SPB Git
6.0 KB · 51 lines typescript
1import "server-only";2import { sqlPublic as sql } from "@/db";3import { embed, embeddingsEnabled } from "@/lib/ai/embeddings";4export type SearchMode = "hybrid" | "lexical" | "semantic";5export type SearchHit = { slug: string; identifier: string; title: string; subtitle: string | null; year: number | null; authors: string[]; score: number; signals: string[]; snippets: { text: string; page: number | null; kind: "chunk" | "abstract" }[] };6const vec = (v: number[]) => `[${v.map((x) => x.toFixed(6)).join(",")}]`;78/** Hybrid retrieval: lexical (tsvector over metadata + full-text chunks, trigram typo tolerance) fused with semantic (pgvector) by reciprocal rank fusion. */9export async function search(q: string, mode: SearchMode = "hybrid", limit = 20): Promise<{ hits: SearchHit[]; mode: SearchMode; semanticUsed: boolean; tookMs: number }> {10  const t0 = Date.now(); const query = q.trim().slice(0, 300); if (!query) return { hits: [], mode, semanticUsed: false, tookMs: 0 };11  const rr = new Map<string, { score: number; signals: Set<string> }>();12  const add = (id: string, rank: number, signal: string, w = 1) => { const e = rr.get(id) ?? { score: 0, signals: new Set() }; e.score += w / (60 + rank); e.signals.add(signal); rr.set(id, e); };13  let semanticUsed = false;14  if (mode !== "semantic") {15    const meta = await sql`SELECT p.id, ts_rank_cd(p.fts, websearch_to_tsquery('english', f_unaccent(${query}))) AS r, similarity(f_unaccent(lower(${query})), p.search_blob) AS s,16        EXISTS (SELECT 1 FROM paper_authors pa JOIN authors a ON a.id = pa.author_id WHERE pa.paper_id = p.id AND f_unaccent(lower(a.display_name)) LIKE '%' || f_unaccent(lower(${query})) || '%') AS by_author,17        p.identifier ILIKE ${"%" + query.replace(/\s+/g, "") + "%"} AS by_id18      FROM public_papers p WHERE p.status = 'published' AND (p.fts @@ websearch_to_tsquery('english', f_unaccent(${query})) OR p.search_blob % f_unaccent(lower(${query})) OR word_similarity(f_unaccent(lower(${query})), p.search_blob) > 0.4519        OR EXISTS (SELECT 1 FROM paper_authors pa JOIN authors a ON a.id = pa.author_id WHERE pa.paper_id = p.id AND f_unaccent(lower(a.display_name)) LIKE '%' || f_unaccent(lower(${query})) || '%') OR p.identifier ILIKE ${"%" + query.replace(/\s+/g, "") + "%"})20      ORDER BY (CASE WHEN p.identifier ILIKE ${"%" + query.replace(/\s+/g, "") + "%"} THEN 10 ELSE 0 END) + ts_rank_cd(p.fts, websearch_to_tsquery('english', f_unaccent(${query}))) * 10 + similarity(f_unaccent(lower(${query})), p.search_blob) DESC LIMIT 40`;21    meta.forEach((m, i) => add(m.id, i, m.by_id ? "identifier" : m.by_author ? "author" : Number(m.r) > 0 ? "metadata" : "fuzzy-title", 1.2));22    const chunks = await sql`SELECT DISTINCT ON (c.paper_id) c.paper_id AS id, c.page_from, ts_headline('english', c.text, websearch_to_tsquery('english', f_unaccent(${query})), 'MaxWords=40, MinWords=20, StartSel=<mark>, StopSel=</mark>, MaxFragments=2, FragmentDelimiter=" … "') AS snippet, ts_rank_cd(c.fts, websearch_to_tsquery('english', f_unaccent(${query}))) AS r23      FROM public_paper_chunks c WHERE c.fts @@ websearch_to_tsquery('english', f_unaccent(${query})) ORDER BY c.paper_id, r DESC`;24    const sorted = [...chunks].sort((a, b) => Number(b.r) - Number(a.r)).slice(0, 40);25    sorted.forEach((c, i) => { add(c.id, i, "full-text", 1); const e = rr.get(c.id)!; (e as unknown as { snip?: { text: string; page: number | null } }).snip = { text: c.snippet, page: c.page_from }; });26  }27  if (mode !== "lexical" && embeddingsEnabled()) {28    try {29      const e = await embed([query]);30      if (e) {31        semanticUsed = true; const v = vec(e[0]);32        const sem = await sql`SELECT id, 1 - (embedding <=> ${v}::vector) AS sim FROM public_papers WHERE status = 'published' AND embedding IS NOT NULL ORDER BY embedding <=> ${v}::vector LIMIT 20`;33        sem.forEach((s, i) => { if (Number(s.sim) > 0.25) add(s.id, i, "semantic", 1); });34        const semc = await sql`SELECT DISTINCT ON (paper_id) paper_id AS id, text, page_from, 1 - (embedding <=> ${v}::vector) AS sim FROM public_paper_chunks WHERE embedding IS NOT NULL ORDER BY paper_id, embedding <=> ${v}::vector`;35        [...semc].sort((a, b) => Number(b.sim) - Number(a.sim)).slice(0, 20).forEach((c, i) => { if (Number(c.sim) > 0.35) { add(c.id, i, "semantic-passage", 0.8); const e2 = rr.get(c.id) as unknown as { snip?: { text: string; page: number | null } }; if (!e2.snip) e2.snip = { text: String(c.text).slice(0, 320) + "…", page: c.page_from }; } });36      }37    } catch (err) { console.warn("semantic search unavailable:", (err as Error).message); }38  }39  const ids = [...rr.entries()].sort((a, b) => b[1].score - a[1].score).slice(0, limit).map(([id]) => id);40  if (!ids.length) return { hits: [], mode, semanticUsed, tookMs: Date.now() - t0 };41  const rows = await sql`SELECT p.id, p.slug, p.identifier, p.title, p.subtitle, p.year, p.abstract, (SELECT array_agg(a.display_name ORDER BY pa.position) FROM paper_authors pa JOIN authors a ON a.id = pa.author_id WHERE pa.paper_id = p.id) AS authors FROM public_papers p WHERE p.id = ANY(${ids})`;42  const by = new Map(rows.map((r) => [r.id, r]));43  const hits: SearchHit[] = ids.map((id) => { const r = by.get(id)!; const e = rr.get(id)! as { score: number; signals: Set<string>; snip?: { text: string; page: number | null } };44    const snippets = e.snip ? [{ text: e.snip.text, page: e.snip.page, kind: "chunk" as const }] : [{ text: String((r.abstract as string[])[0] ?? "").slice(0, 300) + "…", page: null, kind: "abstract" as const }];45    return { slug: r.slug, identifier: r.identifier, title: r.title, subtitle: r.subtitle, year: r.year, authors: r.authors ?? [], score: Number(e.score.toFixed(4)), signals: [...e.signals], snippets }; });46  return { hits, mode, semanticUsed, tookMs: Date.now() - t0 };47}48export async function logSearch(query: string, mode: string, results: number, locale: string, tookMs: number) {49  try { await sql`INSERT INTO search_log (query, mode, results, locale, took_ms) VALUES (${query.slice(0, 300)}, ${mode}, ${results}, ${locale}, ${tookMs})`; } catch { /* ignore */ }50}51
Files · ba542514eac709a3e76ac3b3f88404c61356d594
Menu

Theme