spb/note-spboucher-ai
Public
1/**2 * AI link suggestions. Candidates = top-k pages by cosine similarity between page chunks (pgvector HNSW, one SQL3 * query), excluding self, already-linked pages and pairs that already carry an ai_suggested edge (any status — a4 * rejected pair is never re-suggested). A quick model then confirms each pair (structured JSON) and writes a5 * one-sentence rationale + confidence; related pairs are stored as knowledge_edges kind=ai_suggested status=suggested.6 */7import { sql as dsql } from "drizzle-orm";8import { db, schema } from "@/db";9import { route, EMBEDDING_MODEL } from "@/lib/ai/router";10import { callAI } from "@/lib/ai/providers";11import { recordUsage } from "@/lib/ai/usage";12import { INJECTION_GUARD } from "@/lib/ai/prompts";1314const rows = async <T = any>(q: ReturnType<typeof dsql>): Promise<T[]> => (await db.execute(q)) as unknown as T[];15export interface Candidate { id: string; title: string; type: string; section: string; excerpt: string; sim: number; hits: number }16export interface SuggestResult { pageId: string; candidates: number; suggested: number; declined: number; model?: string; costUsd?: number | null; skipped?: string }1718/** Nearest pages by chunk-embedding cosine similarity (max over chunk pairs), excluding self / linked / already-judged pairs. */19export async function nearestPages(wsId: string, pageId: string, k = 8, minSim = 0.25): Promise<Candidate[]> {20 return rows<Candidate>(dsql`21 with src as (22 select e.embedding as v from embeddings e join chunks c on c.id = e.chunk_id23 where c.workspace_id = ${wsId} and c.page_id = ${pageId} and e.model = ${EMBEDDING_MODEL} order by c.ordinal limit 24),24 nn as (25 select n.page_id, n.sim from src s cross join lateral (26 select c.page_id, 1 - (e.embedding <=> s.v) as sim from embeddings e join chunks c on c.id = e.chunk_id27 where c.workspace_id = ${wsId} and c.page_id is not null and c.page_id <> ${pageId} and e.model = ${EMBEDDING_MODEL}28 order by e.embedding <=> s.v limit 20) n)29 select p.id, p.title, p.type, p.section, left(p.plain_text, 1200) as excerpt, max(nn.sim)::float as sim, count(*)::int as hits30 from nn join pages p on p.id = nn.page_id31 where p.deleted_at is null32 and not exists (select 1 from links l where (l.from_page_id = ${pageId} and l.to_page_id = p.id) or (l.from_page_id = p.id and l.to_page_id = ${pageId}))33 and not exists (select 1 from knowledge_edges x where x.kind = 'ai_suggested' and ((x.from_id = ${pageId} and x.to_id = p.id) or (x.from_id = p.id and x.to_id = ${pageId})))34 group by p.id having max(nn.sim) >= ${minSim}35 order by sim desc limit ${k}`);36}3738const SCHEMA = { name: "link_suggestions", schema: { type: "object", additionalProperties: false, required: ["results"], properties: { results: { type: "array", items: { type: "object", additionalProperties: false, required: ["id", "related", "confidence", "rationale"], properties: { id: { type: "string", description: "candidate id, copied verbatim" }, related: { type: "boolean", description: "true only if the two pages are genuinely about the same topic, project, entity or question and a reader of one would benefit from the other" }, confidence: { type: "number", description: "0 to 1" }, rationale: { type: "string", description: "one sentence, in the language of the pages, naming the concrete shared topic" } } } } } } };39const SYSTEM = `You curate the knowledge graph of a private research workspace. You receive a SOURCE page and candidate pages that are semantically close by embedding. For each candidate decide whether a link between the two pages is genuinely useful: they discuss the same specific topic, project, entity, dataset, course, paper or question — not merely the same broad field or a shared writing style. Be strict: generic similarity (both are "notes", both mention finance) is NOT a link. Return one result per candidate id, in the same order. ${INJECTION_GUARD}`;40const clamp01 = (x: unknown) => Math.max(0, Math.min(1, Number(x) || 0));4142/** Suggest links for one page. k = max candidates sent to the model. */43export async function suggestLinksForPage(wsId: string, pageId: string, opts: { k?: number; minSim?: number } = {}): Promise<SuggestResult> {44 const [src] = await rows<{ id: string; title: string; type: string; section: string; excerpt: string }>(dsql`select id, title, type, section, left(plain_text, 2000) as excerpt from pages where id = ${pageId} and workspace_id = ${wsId} and deleted_at is null`);45 if (!src) return { pageId, candidates: 0, suggested: 0, declined: 0, skipped: "page not found" };46 const cands = await nearestPages(wsId, pageId, opts.k ?? 8, opts.minSim ?? 0.25);47 if (!cands.length) return { pageId, candidates: 0, suggested: 0, declined: 0, skipped: "no candidates" };48 const r = await route({ task: "quick", schema: true });49 const user = [`<workspace>`, `SOURCE PAGE (id ${src.id}) — ${src.title} [${src.type}/${src.section}]`, src.excerpt || "(empty)", "", ...cands.flatMap((c) => [`CANDIDATE id=${c.id} — ${c.title} [${c.type}/${c.section}] similarity=${c.sim.toFixed(2)}`, c.excerpt || "(empty)", ""]), `</workspace>`, `Judge each of the ${cands.length} candidates. Candidate ids: ${cands.map((c) => c.id).join(", ")}.`].join("\n");50 const t0 = performance.now();51 let parsed: { results: { id: string; related: boolean; confidence: number; rationale: string }[] };52 try {53 const res = await callAI({ model: r.model.id, provider: r.provider, system: SYSTEM, messages: [{ role: "user", content: user }], jsonSchema: SCHEMA, effort: r.effort, maxTokens: 2000, timeoutMs: 120_000 });54 await recordUsage({ workspaceId: wsId, provider: r.provider, model: r.model.id, task: "suggest-links", usage: res.usage, costUsd: res.costUsd, latencyMs: res.totalMs, ttftMs: res.ttftMs, ok: true, refId: pageId });55 try { parsed = JSON.parse(res.text); } catch { throw new Error("model returned invalid JSON"); }56 const byId = new Map(cands.map((c) => [c.id, c])); let suggested = 0, declined = 0;57 for (const x of parsed.results ?? []) {58 const c = byId.get(String(x.id)); if (!c) continue; byId.delete(c.id);59 const related = Boolean(x.related) && clamp01(x.confidence) >= 0.35;60 await db.insert(schema.knowledgeEdges).values({ workspaceId: wsId, fromId: pageId, toId: c.id, fromType: "page", toType: "page", kind: "ai_suggested", status: related ? "suggested" : "rejected", confidence: clamp01(x.confidence), rationale: String(x.rationale ?? "").slice(0, 600), model: r.model.id, decidedAt: null }).onConflictDoNothing();61 if (related) suggested++; else declined++;62 }63 return { pageId, candidates: cands.length, suggested, declined, model: r.model.id, costUsd: res.costUsd };64 } catch (e: any) {65 await recordUsage({ workspaceId: wsId, provider: r.provider, model: r.model.id, task: "suggest-links", ok: false, error: String(e?.message ?? e).slice(0, 300), latencyMs: performance.now() - t0, refId: pageId });66 throw e;67 }68}6970/** Batched workspace pass: every page that has embeddings, most recently edited first (bounded per run). */71export async function suggestLinksForWorkspace(wsId: string, opts: { limit?: number; k?: number } = {}) {72 const pages = await rows<{ id: string }>(dsql`select p.id from pages p where p.workspace_id = ${wsId} and p.deleted_at is null and exists (select 1 from chunks c join embeddings e on e.chunk_id = c.id where c.page_id = p.id and e.model = ${EMBEDDING_MODEL}) order by p.updated_at desc limit ${opts.limit ?? 150}`);73 const out = { pages: pages.length, suggested: 0, declined: 0, failed: 0 };74 for (const p of pages) { try { const r = await suggestLinksForPage(wsId, p.id, { k: opts.k ?? 5 }); out.suggested += r.suggested; out.declined += r.declined; } catch (e) { out.failed++; console.error("[suggest-links]", p.id, e); } }75 return out;76}77