Momento/memento-note/lib/ai/tools/note-crud.tool.ts

/**
 * Note CRUD Tools
 * note_create, note_read, note_update
 */

import { tool } from 'ai'
import { z } from 'zod'
import { toolRegistry } from './registry'
import { prisma } from '@/lib/prisma'
import { markdownToHtml } from '@/lib/markdown-to-html'

// --- note_read ---
toolRegistry.register({
  name: 'note_read',
  description: 'Read a specific note by its ID. Returns the full note content.',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: 'Read a specific note by ID. Returns the full content.',
      inputSchema: z.object({
        noteId: z.string().describe('The ID of the note to read'),
      }),
      execute: async ({ noteId }) => {
        try {
          const note = await prisma.note.findFirst({
            where: { id: noteId, userId: ctx.userId },
            select: { id: true, title: true, content: true, isMarkdown: true, createdAt: true, updatedAt: true },
          })
          if (!note) return { error: 'Note not found' }
          return note
        } catch (e: any) {
          return { error: `Read note failed: ${e.message}` }
        }
      },
    }),
})

// --- note_create ---
toolRegistry.register({
  name: 'note_create',
  description: 'Create a new note with a title and content.',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: 'Create a new note.',
      inputSchema: z.object({
        title: z.string().describe('Title for the note'),
        content: z.string().describe('Content of the note (markdown supported)'),
        notebookId: z.string().optional().describe('Optional notebook ID to place the note in'),
        images: z.array(z.string()).optional().describe('Optional array of local image URL paths to attach to the note (e.g. ["/uploads/notes/abc.jpg"])'),
      }),
      execute: async ({ title, content, notebookId, images }) => {
        try {
          const htmlContent = markdownToHtml(content)
          const note = await prisma.note.create({
            data: {
              title,
              content: htmlContent,
              type: 'richtext',
              isMarkdown: false,
              autoGenerated: true,
              userId: ctx.userId,
              notebookId: notebookId || null,
              images: images && images.length > 0 ? JSON.stringify(images) : null,
            },
            select: { id: true, title: true },
          })
          return { success: true, noteId: note.id, title: note.title }
        } catch (e: any) {
          return { error: `Create note failed: ${e.message}` }
        }
      },
    }),
})

// --- note_update ---
toolRegistry.register({
  name: 'note_update',
  description: 'Update an existing note\'s content.',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: 'Update an existing note.',
      inputSchema: z.object({
        noteId: z.string().describe('The ID of the note to update'),
        title: z.string().optional().describe('New title (optional)'),
        content: z.string().optional().describe('New content (optional)'),
      }),
      execute: async ({ noteId, title, content }) => {
        try {
          const existing = await prisma.note.findFirst({
            where: { id: noteId, userId: ctx.userId },
          })
          if (!existing) return { error: 'Note not found' }

          const data: Record<string, any> = {}
          if (title !== undefined) data.title = title
          if (content !== undefined) data.content = content

          await prisma.note.update({ where: { id: noteId }, data })
          return { success: true, noteId }
        } catch (e: any) {
          return { error: `Update note failed: ${e.message}` }
        }
      },
    }),
})