1fcea6ed7d
All checks were successful
Deploy to Production / Build and Deploy (push) Successful in 7s
- Add brainstorm feature with collaborative canvas, AI idea generation, live cursors, playback, and export - Add PDF upload/extraction/ingestion pipeline with pgvector document search (RAG) - Add document Q&A overlay with streaming chat and PDF preview - Add note attachments UI with status polling, grid layout, and auto-scroll - Add task extraction AI tool and agent executor improvements - Fix NoteEmbedding missing updatedAt column, re-index 66 notes with 1536-dim embeddings - Fix brainstorm 'Create Note' button: add success toast and redirect to created note - Fix memory echo notification infinite polling - Fix chat route to always include document_search tool - Add brainstorm i18n keys across all 14 locales - Add socket server for real-time brainstorm collaboration - Add hierarchical notebook selector and organize notebook dialog improvements - Add sidebar brainstorm section with session management - Update prisma schema with brainstorm tables, attachments, and document chunks
74 lines
2.7 KiB
TypeScript
74 lines
2.7 KiB
TypeScript
import { tool } from 'ai'
|
|
import { z } from 'zod'
|
|
import { toolRegistry } from './registry'
|
|
import { embeddingService } from '@/lib/ai/services/embedding.service'
|
|
import prisma from '@/lib/prisma'
|
|
|
|
toolRegistry.register({
|
|
name: 'document_search',
|
|
description: 'Search within PDF documents attached to notes. Returns relevant passages with page numbers and source document info.',
|
|
isInternal: true,
|
|
buildTool: (ctx) =>
|
|
tool({
|
|
description: `Search within PDF documents attached to the user's notes.
|
|
Returns matching passages with page numbers, chunk content, and the source note/document info.
|
|
Use this when the user asks about specific documents, PDFs, or attached files.`,
|
|
inputSchema: z.object({
|
|
query: z.string().describe('The search query to find relevant passages in documents'),
|
|
noteId: z.string().optional().describe('Optional: restrict search to attachments of a specific note'),
|
|
limit: z.number().optional().describe('Max results to return (default 5)').default(5),
|
|
}),
|
|
execute: async ({ query, noteId, limit = 5 }) => {
|
|
try {
|
|
const queryEmbedding = await embeddingService.generateEmbedding(query)
|
|
const vectorStr = embeddingService.toVectorString(queryEmbedding.embedding)
|
|
|
|
let noteFilter = ''
|
|
const params: any[] = [vectorStr, limit, ctx.userId]
|
|
|
|
if (noteId) {
|
|
noteFilter = `AND na."noteId" = $4`
|
|
params.push(noteId)
|
|
}
|
|
|
|
const results = await prisma.$queryRawUnsafe(
|
|
`SELECT
|
|
dc.id as "chunkId",
|
|
dc.content,
|
|
dc."pageNumber",
|
|
dc."chunkIndex",
|
|
na.id as "attachmentId",
|
|
na."fileName",
|
|
na."pageCount",
|
|
na."noteId",
|
|
n.title as "noteTitle"
|
|
FROM "DocumentChunk" dc
|
|
JOIN "NoteAttachment" na ON na.id = dc."attachmentId"
|
|
JOIN "Note" n ON n.id = na."noteId"
|
|
WHERE dc."embedding" IS NOT NULL
|
|
AND na.status = 'ready'
|
|
AND n."trashedAt" IS NULL
|
|
AND n."userId" = $3
|
|
${noteFilter}
|
|
ORDER BY dc."embedding" <=> $1::vector
|
|
LIMIT $2`,
|
|
...params
|
|
) as any[]
|
|
|
|
if (!results.length) return { results: [], message: 'No matching documents found' }
|
|
|
|
return results.map(r => ({
|
|
content: r.content.substring(0, 600),
|
|
pageNumber: r.pageNumber,
|
|
chunkIndex: r.chunkIndex,
|
|
fileName: r.fileName,
|
|
noteId: r.noteId,
|
|
noteTitle: r.noteTitle || 'Untitled',
|
|
}))
|
|
} catch (e: any) {
|
|
return { error: `Document search failed: ${e.message}` }
|
|
}
|
|
},
|
|
}),
|
|
})
|