spb/chat-spboucher-ai

1import "server-only";2import { sql, vec } from "../db";3import { openaiEmbed } from "../ai/openai";4import { recordUsage } from "../usage";5import { chunkText } from "../storage";6import type { ToolContext } from "./index";78/** Index an attachment's extracted text into document_chunks (embeddings + FTS). */9export async function indexAttachment(attachmentId: string) {10  const a = await sql<{ extracted_text: string | null }[]>`SELECT extracted_text FROM attachments WHERE id = ${attachmentId}`;11  const text = a[0]?.extracted_text;12  if (!text || text.length < 40) return 0;13  await sql`DELETE FROM document_chunks WHERE attachment_id = ${attachmentId}`;14  const chunks = chunkText(text).slice(0, 400); // cap ~500k chars per document15  for (let i = 0; i < chunks.length; i += 64) {16    const batch = chunks.slice(i, i + 64);17    let vectors: number[][] | null = null;18    try { const r = await openaiEmbed(batch); vectors = r.vectors; await recordUsage({ provider: "openai", modelId: "text-embedding-3-small", kind: "embedding", usage: { input: r.tokens, output: 0 }, costEstimated: r.tokens * 0.02 / 1e6 }); } catch { vectors = null; }19    for (let j = 0; j < batch.length; j++) {20      const v = vectors?.[j];21      if (v) await sql`INSERT INTO document_chunks (attachment_id, idx, content, embedding) VALUES (${attachmentId}, ${i + j}, ${batch[j]}, ${vec(v)}::vector)`;22      else await sql`INSERT INTO document_chunks (attachment_id, idx, content) VALUES (${attachmentId}, ${i + j}, ${batch[j]})`;23    }24  }25  return chunks.length;26}2728export async function searchDocuments(query: string, ctx: Pick<ToolContext, "conversationId" | "projectId">, limit = 6) {29  let qv: number[] | null = null;30  try { qv = (await openaiEmbed([query])).vectors[0]; } catch { qv = null; }31  const rows = await sql<{ id: number; content: string; filename: string; page: number | null; attachment_id: string; score: number }[]>`32    SELECT c.id, c.content, c.page, c.attachment_id, a.filename,33      (${qv ? sql`1 - (c.embedding <=> ${vec(qv)}::vector)` : sql`0`}) * 0.7 + ts_rank(c.tsv, plainto_tsquery('simple', ${query})) * 0.3 AS score34    FROM document_chunks c JOIN attachments a ON a.id = c.attachment_id35    WHERE (a.conversation_id = ${ctx.conversationId} OR (${ctx.projectId ?? null}::text IS NOT NULL AND a.project_id = ${ctx.projectId ?? null}))36    ORDER BY score DESC LIMIT ${limit}`;37  return { query, passages: rows.map(r => ({ file: r.filename, attachmentId: r.attachment_id, page: r.page, score: Number(r.score.toFixed(3)), text: r.content })) };38}39