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

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

import { tool } from 'ai'
import { z } from 'zod'
import { Prisma } from '@prisma/client'
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 note = await prisma.note.create({
            data: {
              title,
              content,
              type: 'markdown',
              isMarkdown: true,
              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}` }
        }
      },
    }),
})

// --- note_find_and_update ---
toolRegistry.register({
  name: 'note_find_and_update',
  description: 'Find a note by searching its title/content, then append, prepend, or replace information in it. Use this when the user says "find the note about X and add Y to it".',
  isInternal: true,
  buildTool: (ctx) =>
    tool({
      description: 'Find a note by a search query, then update its content (append/prepend/replace).',
      inputSchema: z.object({
        query: z.string().describe('Search query to find the note (e.g. "bugs and new features")'),
        newContent: z.string().describe('Content to add to the note (markdown supported)'),
        operation: z.enum(['append', 'prepend', 'replace']).default('append').describe('append: add to end, prepend: add to start, replace: overwrite'),
      }),
      execute: async ({ query, newContent, operation }) => {
        try {
          // FTS search for best matching note
          const results = await prisma.$queryRaw<Array<{ id: string; title: string | null; content: string | null }>>(
            Prisma.sql`
              SELECT id, title, content
              FROM "Note"
              WHERE "tsv" @@ plainto_tsquery('simple', ${query})
                AND "trashedAt" IS NULL
                AND "isArchived" = false
                AND "userId" = ${ctx.userId}
              ORDER BY ts_rank("tsv", plainto_tsquery('simple', ${query})) DESC
              LIMIT 1`
          )

          if (!results || results.length === 0) {
            // Fallback: simple title/content ILIKE search
            const fallback = await prisma.note.findFirst({
              where: {
                userId: ctx.userId,
                trashedAt: null,
                isArchived: false,
                OR: [
                  { title: { contains: query, mode: 'insensitive' } },
                  { content: { contains: query, mode: 'insensitive' } },
                ],
              },
              select: { id: true, title: true, content: true },
            })
            if (!fallback) {
              return { error: `No note found matching "${query}". Try a different search term.` }
            }
            results.push(fallback)
          }

          const note = results[0]
          let updatedContent: string

          switch (operation) {
            case 'append':
              updatedContent = note.content ? `${note.content}\n\n${newContent}` : newContent
              break
            case 'prepend':
              updatedContent = note.content ? `${newContent}\n\n${note.content}` : newContent
              break
            case 'replace':
            default:
              updatedContent = newContent
          }

          await prisma.note.update({
            where: { id: note.id },
            data: { content: updatedContent, updatedAt: new Date() },
          })

          return {
            success: true,
            noteId: note.id,
            noteTitle: note.title || 'Untitled',
            operation,
            message: `Successfully updated note "${note.title || 'Untitled'}" (${operation}).`,
          }
        } catch (e: any) {
          return { error: `find_and_update failed: ${e.message}` }
        }
      },
    }),
})