Hybrid search — Vectorize → FTS5 → substring → RRF → rerank → ACL

Source: platform/apps/app/src/lib/rag/search.ts · rendered from main on every deploy — edit in the repo, not here

platform/apps/app/src/lib/rag/search.ts

/**
 * Hybrid search pipeline: Vectorize (semantic) + FTS5 (lexical) + RRF + Rerank.
 * Gracefully degrades when bindings unavailable.
 */
import type { Database } from '@puccha/db';
import { chunks, documents } from '@puccha/db';
import type { ACL, PublicSurface, Role, TenantContext } from '@puccha/types';
import { and, eq, inArray, or, sql } from 'drizzle-orm';
import { containsPattern } from '$lib/db/like.js';
import { citationTitle, sourceUrlFromRef } from './citation-title.js';
import {
  buildEmbeddingCacheKey,
  type EmbeddingCacheLike,
  readEmbedding,
  writeEmbedding,
} from './embedding-cache.js';
import { packSemanticSearch } from './pack-search.js';
import type { Reranker } from './rerank.js';
import { segmentForIndex } from './thai-segment.js';

export interface SearchResult {
  chunkId: string;
  docId: string;
  /**
   * Display title for the LLM context header and the citation. After
   * `hybridSearch` this is the document title (plus ` › section` when the
   * section heading means something on its own) — see citation-title.ts.
   * Inside the search stages it is still the raw chunk/section title.
   */
  title: string;
  /** Document title as stored; absent for pack chunks / deleted docs. */
  docTitle?: string;
  /** Public page URL for crawled / linked documents. */
  url?: string;
  text: string;
  anchor: string;
  score: number;
  source: 'vectorize' | 'fts5' | 'substring';
  /**
   * ADR-0101 §8. Pack provenance — present only when the chunk came
   * from an L1 compliance pack (i.e. via `packSemanticSearch`). The
   * citation renderer keys off this field to show pack badge + version
   * + effective date + deep-link. Tenant chunks leave it undefined.
   *
   * Stored verbatim onto `messages.sourcesJson` so historical citations
   * remain defensibly authoritative after pack versions bump
   * (subscribers may unsubscribe / upgrade — the prior turn's citation
   * must still resolve to what the agent saw at the time).
   */
  pack?: {
    slug: string;
    version: string;
    majorVersion: number;
    effectiveDate: string;
    sourceRef: string;
    sourceUrl: string;
  };
}

export interface SearchOptions {
  db: Database;
  tenant: TenantContext;
  query: string;
  topK?: number;
  vectorize?: VectorizeIndex;
  ai?: Ai;
  reranker?: Reranker;
  /** Filter to specific source IDs (from persona config). Empty = all sources. */
  sourceIds?: string[];
  /**
   * ADR-0183. Restrict retrieval to documents whose ACL is
   * `{ kind: 'public', surfaces: [<this surface>] }` — the ADR-0091 surface
   * gate. Used for the no-token tier on an interlocked tenant: a visitor
   * with no learner session sees the FAQ layer and nothing else. Applied
   * before every phase, so nothing outside the layer reaches the reranker or
   * the model; compliance packs are skipped on this tier too.
   */
  publicSurface?: PublicSurface;
  /**
   * ADR-0039 Stage 3. Restrict to a single workspace within the tenant.
   * Optional during the migration window — when omitted, behaves like the
   * pre-workspaces code (matches all of the tenant's workspaces). Callers
   * with a known workspace context (chat, MCP, playground) should pass
   * `locals.workspace.id`; future stages will tighten this to required.
   */
  workspaceId?: string;
  /**
   * KV namespace for the bge-m3 embedding cache. When passed, semantic
   * search will skip the AI.run() call on a hit. Vectors are tenant-
   * agnostic so the cache is shared globally — same query across tenants
   * pays the AI cost once.
   */
  embeddingCache?: EmbeddingCacheLike;
  /**
   * Cloudflare Worker `ctx.waitUntil` — used to write embedding cache
   * entries without blocking the response. Pass `platform.context.waitUntil`
   * from the route handler.
   */
  waitUntil?: (p: Promise<unknown>) => void;
  /**
   * Optional observability hook — invoked once per call with the
   * embedding-cache outcome. Routes pass a setter that captures the
   * value into a closure so it ends up in query_log via logQuery's
   * embeddingCacheStatus field. NB: only fires when semanticSearch
   * actually runs (skips when vectorize/ai bindings are missing).
   */
  onEmbeddingCacheStatus?: (status: 'hit' | 'miss' | 'skipped') => void;
  /**
   * Fired once per call with whether the reranker actually ran. Rerank
   * only fires when a reranker is present AND there are enough candidates
   * to benefit (>5), so small KBs return un-calibrated RRF/FTS scores.
   * The scope gate uses this to decide whether the top score is on the
   * reranker's calibrated scale — without it, a 3-chunk KB is compared
   * against the Cohere threshold and falsely refused as out-of-scope.
   */
  onReranked?: (ran: boolean) => void;
  /**
   * ADR-0101 §4. Subscribed compliance packs to include alongside the
   * tenant's private KM in the candidate set. Each pack contributes
   * one parallel Vectorize query against `platform:pack:<slug>:v<major>`.
   * Resolution from `personas.packSlugs` → active subscriptions →
   * this list happens in the chat orchestrator; this surface accepts
   * only already-authorized packs (subscription IS the ACL).
   */
  packs?: ReadonlyArray<{ slug: string; majorVersion: number }>;
}

// ── Main entry point ────────────────────────────────────────────────────

// Upper bound on how long we'll wait for the reranker before giving up and
// serving RRF order. Cohere rerank-v3.5 is normally 200–500ms; this only
// trips on the network tail / provider stall, where blocking the whole chat
// response is worse than slightly-degraded ranking for that one turn.
const RERANK_TIMEOUT_MS = 1500;

export async function hybridSearch(opts: SearchOptions): Promise<SearchResult[]> {
  const { db, tenant, query, topK = 5, sourceIds } = opts;
  const candidates: SearchResult[] = [];

  // Resolve allowed docIds ONCE when persona restricts sources (ADR-0023).
  // Passed into each retrieval phase so we don't run expensive queries
  // and then discard 90% of the results post-hoc — which would silently
  // return 0 hits if the top-K happens to be all out-of-scope.
  let allowedDocIds: string[] | null = null;
  if (sourceIds && sourceIds.length > 0) {
    const wsFilter = opts.workspaceId ? eq(documents.workspaceId, opts.workspaceId) : undefined;
    const allowedDocs = await db
      .select({ id: documents.id })
      .from(documents)
      .where(
        wsFilter
          ? and(eq(documents.tenantId, tenant.id), wsFilter, inArray(documents.sourceId, sourceIds))
          : and(eq(documents.tenantId, tenant.id), inArray(documents.sourceId, sourceIds)),
      );
    allowedDocIds = allowedDocs.map((d) => d.id);
    // No allowed docs → short-circuit. Nothing can possibly match.
    if (allowedDocIds.length === 0) return [];
  }

  // ADR-0183 public tier: the surface gate as a doc-id allow-list, resolved
  // once like the persona scope above. `json_each` over `$.surfaces` is the
  // same predicate the help center uses (`helpSurfaceFilter`), so "public on
  // the widget" means exactly one thing across the codebase. Intersects with
  // a persona scope when both apply.
  if (opts.publicSurface) {
    const surfaceRows = await db
      .select({ id: documents.id })
      .from(documents)
      .where(
        and(
          eq(documents.tenantId, tenant.id),
          opts.workspaceId ? eq(documents.workspaceId, opts.workspaceId) : undefined,
          sql`json_extract(${documents.aclJson}, '$.kind') = 'public'
            AND EXISTS (
              SELECT 1 FROM json_each(json_extract(${documents.aclJson}, '$.surfaces'))
              WHERE value = ${opts.publicSurface}
            )`,
        ),
      );
    const publicIds = surfaceRows.map((d) => d.id);
    allowedDocIds = allowedDocIds
      ? allowedDocIds.filter((id) => publicIds.includes(id))
      : publicIds;
    if (allowedDocIds.length === 0) return [];
  }

  // Phase 1: Vectorize semantic search (when available).
  // Tenant chunks and pack chunks query in parallel — same RRF/rerank
  // stage consumes both. Pack results are pre-authorized by subscription
  // (the caller resolved `packs` from active `tenant_pack_subscriptions`),
  // so `allowedDocIds` source-scoping does NOT apply to them.
  // Vectorize filter doesn't support docId IN (...) — oversample topK so
  // post-filter still yields relevant hits after source scoping.
  // Phase 2 (FTS5 lexical) runs in the SAME Promise.all as the vector
  // queries: it depends only on the raw query + ACL filter, never on the
  // semantic results, so awaiting it sequentially just stacked its latency
  // onto the critical path for no reason. Substring fallback (Phase 3) is
  // still sequential — it's gated on the combined candidate count.
  const [semanticResults, packResults, lexicalResults] = await Promise.all([
    semanticSearch({
      ...opts,
      topK: allowedDocIds ? Math.max(50, topK * 10) : 20,
    }),
    opts.packs && opts.packs.length > 0 && !opts.publicSurface
      ? packSemanticSearch({
          query: opts.query,
          topK: opts.topK,
          vectorize: opts.vectorize,
          ai: opts.ai,
          embeddingCache: opts.embeddingCache,
          waitUntil: opts.waitUntil,
          packs: opts.packs,
        })
      : Promise.resolve([] as SearchResult[]),
    fts5Search(db, tenant, query, allowedDocIds, opts.workspaceId),
  ]);
  if (allowedDocIds) {
    const allowedSet = new Set(allowedDocIds);
    candidates.push(...semanticResults.filter((r) => allowedSet.has(r.docId)));
  } else {
    candidates.push(...semanticResults);
  }
  candidates.push(...packResults);
  candidates.push(...lexicalResults);

  // Phase 3: Thai substring fallback (if phases 1+2 returned < 3 results)
  if (candidates.length < 3) {
    // Degrade to no extra candidates rather than failing the request, the
    // same way the FTS5 phase does. A retrieval *fallback* that can take
    // down a chat turn is the wrong shape regardless of which bug is live:
    // ADR-0175 was written after an oversized LIKE pattern did exactly that.
    try {
      const substringResults = await substringSearch(
        db,
        tenant,
        query,
        allowedDocIds,
        opts.workspaceId,
      );
      candidates.push(...substringResults);
    } catch (err) {
      console.error('[SUBSTRING_FALLBACK_FAILED]', err);
    }
  }

  // Deduplicate by chunkId
  const seen = new Set<string>();
  const unique = candidates.filter((c) => {
    if (seen.has(c.chunkId)) return false;
    seen.add(c.chunkId);
    return true;
  });

  // Phase 4: RRF merge (if we have results from multiple sources)
  const sources = new Set(unique.map((r) => r.source));
  let merged: SearchResult[];
  if (sources.size > 1) {
    merged = reciprocalRankFusion(unique);
  } else {
    merged = unique.sort((a, b) => b.score - a.score);
  }

  // Phase 5: Rerank. A calibrated reranker (Cohere) runs at ANY candidate
  // count so the scope gate gets a real relevance floor even on small KBs
  // (ADR-0112 D3) — the score IS the safety signal there. A non-calibrated
  // reranker (LLM/noop) only reorders, so we keep the >5 FinOps guard:
  // reranking ≤5 rank-derived results adds latency without changing scope.
  if (opts.reranker && (opts.reranker.calibrated || merged.length > 5)) {
    const toRerank = merged.slice(0, 25);
    // Time-box the rerank: a stalled provider must not hold the chat
    // response hostage. On timeout (or any error) we fall through to RRF
    // order below. The late rerank promise is swallowed so a rejection
    // arriving after the race doesn't surface as an unhandled rejection.
    const rerankP = opts.reranker.rerank(
      query,
      toRerank.map((c) => `${c.title}\n${c.text}`),
      topK,
    );
    rerankP.catch(() => {});
    let timer: ReturnType<typeof setTimeout> | undefined;
    try {
      const reranked = await Promise.race([
        rerankP,
        new Promise<never>((_, reject) => {
          timer = setTimeout(() => reject(new Error('rerank timeout')), RERANK_TIMEOUT_MS);
        }),
      ]);
      clearTimeout(timer);
      opts.onReranked?.(true);
      return withDocumentTitles(
        opts.db,
        opts.tenant.id,
        reranked.map((r) => ({
          ...toRerank[r.index],
          score: r.score,
        })),
      );
    } catch {
      clearTimeout(timer);
      // Reranker failed or timed out — fall through to RRF order
    }
  }

  opts.onReranked?.(false);
  return withDocumentTitles(opts.db, opts.tenant.id, merged.slice(0, topK));
}

/**
 * Replace raw section titles on the final top-K with `docTitle › section`
 * (or the doc title alone) and attach the page URL. One tenant-scoped
 * `documents` lookup for ≤ topK ids; pack chunks and deleted docs keep
 * their existing title. Runs after rerank so it never affects ranking.
 */
async function withDocumentTitles(
  db: Database,
  tenantId: string,
  results: SearchResult[],
): Promise<SearchResult[]> {
  const ids = [...new Set(results.map((r) => r.docId).filter(Boolean))];
  if (ids.length === 0) return results;
  let rows: Array<{ id: string; title: string; sourceRef: string | null }> = [];
  try {
    rows = await db
      .select({ id: documents.id, title: documents.title, sourceRef: documents.sourceRef })
      .from(documents)
      .where(and(eq(documents.tenantId, tenantId), inArray(documents.id, ids)));
  } catch (err) {
    // Citation polish must never take retrieval down with it.
    console.error('[SEARCH_DOC_TITLES_FAILED]', err);
    return results;
  }
  const byId = new Map(rows.map((r) => [r.id, r]));
  return results.map((r) => {
    const doc = byId.get(r.docId);
    if (!doc) return r;
    return {
      ...r,
      title: citationTitle(doc.title, r.title),
      docTitle: doc.title,
      url: sourceUrlFromRef(doc.sourceRef),
    };
  });
}

// ── Phase 1: Vectorize semantic search ──────────────────────────────────

async function semanticSearch(opts: SearchOptions): Promise<SearchResult[]> {
  if (!opts.vectorize || !opts.ai) return [];

  try {
    // Embedding cache (FinOps): bge-m3 is deterministic, so cache the
    // vector by hashed query and skip the AI.run on a hit. ~50ms +
    // ~$0.0001 saved per repeat. Cache is global (cross-tenant) since
    // the embedding doesn't depend on tenant content.
    let queryVec: number[] | null = null;
    const cacheKey =
      opts.embeddingCache && opts.waitUntil ? await buildEmbeddingCacheKey(opts.query) : null;
    const cacheActive = !!opts.embeddingCache && cacheKey !== null;
    if (opts.embeddingCache && cacheKey) {
      queryVec = await readEmbedding(opts.embeddingCache, cacheKey);
    }
    if (!queryVec) {
      const embedding = (await opts.ai.run('@cf/baai/bge-m3', {
        text: [opts.query],
      })) as { data: number[][] };
      if (!embedding.data?.[0]) return [];
      queryVec = embedding.data[0];
      if (opts.embeddingCache && opts.waitUntil && cacheKey) {
        writeEmbedding(opts.embeddingCache, cacheKey, queryVec, opts.waitUntil);
      }
      opts.onEmbeddingCacheStatus?.(cacheActive ? 'miss' : 'skipped');
    } else {
      opts.onEmbeddingCacheStatus?.('hit');
    }

    const vecResults = await opts.vectorize.query(queryVec, {
      // Oversample when caller asks — source scoping needs more candidates
      // so the post-filter still returns relevant hits.
      topK: opts.topK ?? 20,
      filter: { tenantId: opts.tenant.id },
      returnMetadata: 'all',
    });

    const results: SearchResult[] = [];
    for (const match of vecResults.matches) {
      const meta = match.metadata as Record<string, string> | undefined;
      if (!meta) continue;

      // ADR-0039 Stage 3: workspace post-filter. Vectorize doesn't index
      // workspaceId in its filter expression here (legacy vectors lack it
      // in metadata), so we filter after the query. Missing metadata
      // workspaceId is treated as "default workspace" — correct for
      // pre-Stage-2 vectors during the migration window.
      if (opts.workspaceId && meta.workspaceId && meta.workspaceId !== opts.workspaceId) {
        continue;
      }

      const acl = JSON.parse(meta.acl || '{"kind":"org"}') as ACL;
      if (!checkAcl(acl, opts.tenant.userId, opts.tenant.role)) continue;

      results.push({
        chunkId: meta.chunkId || match.id,
        docId: meta.docId || '',
        title: meta.title || '',
        text: meta.text || '',
        anchor: meta.anchor || '',
        score: match.score ?? 0,
        source: 'vectorize',
      });
    }
    return results;
  } catch {
    return []; // Vectorize unavailable
  }
}

// ── Phase 2: FTS5 lexical search ────────────────────────────────────────

async function fts5Search(
  db: Database,
  tenant: TenantContext,
  query: string,
  allowedDocIds: string[] | null,
  workspaceId: string | undefined,
): Promise<SearchResult[]> {
  try {
    const segmented = segmentForIndex(query);
    // FTS5 MATCH with OR between tokens for recall
    const matchQuery = segmented
      .split(/\s+/)
      .filter((t) => t.length >= 2)
      .join(' OR ');
    if (!matchQuery) return [];

    // chunks_fts itself doesn't carry workspace_id (it indexes denormalized
    // content for FTS), but the JOIN to chunks gives us c.workspace_id —
    // adding the filter there scopes the result set without re-indexing
    // the FTS table. ADR-0039 Stage 3.
    const wsClause = workspaceId ? sql`AND c.workspace_id = ${workspaceId}` : sql``;

    // Source-scope filter pushed into SQL so we don't waste rank budget.
    // Parameterized IN-list via `sql.join` keeps the query bindable.
    const rows = (await (allowedDocIds
      ? db.all(sql`
        SELECT fts.chunk_id, fts.acl_json, fts.rank,
          c.doc_id, c.title, c.text, c.anchor
        FROM chunks_fts fts
        JOIN chunks c ON c.id = fts.chunk_id
        WHERE fts.tenant_id = ${tenant.id}
          ${wsClause}
          AND chunks_fts MATCH ${matchQuery}
          AND c.doc_id IN (${sql.join(
            allowedDocIds.map((id) => sql`${id}`),
            sql`, `,
          )})
        ORDER BY fts.rank
        LIMIT 50
      `)
      : db.all(sql`
        SELECT fts.chunk_id, fts.acl_json, fts.rank,
          c.doc_id, c.title, c.text, c.anchor
        FROM chunks_fts fts
        JOIN chunks c ON c.id = fts.chunk_id
        WHERE fts.tenant_id = ${tenant.id}
          ${wsClause}
          AND chunks_fts MATCH ${matchQuery}
        ORDER BY fts.rank
        LIMIT 20
      `))) as Array<Record<string, unknown>>;

    const results: SearchResult[] = [];
    for (const row of rows) {
      const acl = JSON.parse((row.acl_json as string) || '{"kind":"org"}') as ACL;
      if (!checkAcl(acl, tenant.userId, tenant.role)) continue;

      results.push({
        chunkId: row.chunk_id as string,
        docId: row.doc_id as string,
        title: row.title as string,
        text: row.text as string,
        anchor: row.anchor as string,
        score: Math.abs(row.rank as number),
        source: 'fts5',
      });
    }
    return results;
  } catch {
    return []; // FTS5 table may not exist
  }
}

// ── Phase 3: Thai substring fallback ────────────────────────────────────

async function substringSearch(
  db: Database,
  tenant: TenantContext,
  query: string,
  allowedDocIds: string[] | null,
  workspaceId: string | undefined,
): Promise<SearchResult[]> {
  const segmented = segmentForIndex(query);
  const tokens = segmented.split(/\s+/).filter((t) => t.length >= 2);
  const qLower = query.toLowerCase();

  // Use SQL LIKE to filter at DB level instead of loading all chunks
  // Search for the full query OR the longest token in title/text
  const likePattern = containsPattern(qLower);
  const longestToken = tokens.sort((a, b) => b.length - a.length)[0] ?? qLower;
  const tokenPattern = containsPattern(longestToken);

  const whereConds = [
    eq(chunks.tenantId, tenant.id),
    or(
      sql`lower(${chunks.title}) LIKE ${likePattern}`,
      sql`lower(${chunks.text}) LIKE ${likePattern}`,
      sql`lower(${chunks.title}) LIKE ${tokenPattern}`,
      sql`lower(${chunks.text}) LIKE ${tokenPattern}`,
    ),
  ];
  if (workspaceId) {
    whereConds.push(eq(chunks.workspaceId, workspaceId));
  }
  if (allowedDocIds) {
    whereConds.push(inArray(chunks.docId, allowedDocIds));
  }

  const matchedChunks = await db
    .select({
      id: chunks.id,
      docId: chunks.docId,
      title: chunks.title,
      text: chunks.text,
      anchor: chunks.anchor,
      aclJson: chunks.aclJson,
    })
    .from(chunks)
    .where(and(...whereConds))
    .limit(100);

  const results: SearchResult[] = [];

  for (const chunk of matchedChunks) {
    const h = `${chunk.title} ${chunk.text}`.toLowerCase();
    let score = 0;

    if (h.includes(qLower)) score += 15;

    let hits = 0;
    for (const tok of tokens) {
      if (h.includes(tok)) {
        hits++;
        score += tok.length >= 4 ? 4 : 2;
      }
    }
    if (tokens.length > 0) {
      const coverage = hits / tokens.length;
      if (coverage >= 0.7) score += 5;
      if (coverage >= 0.9) score += 5;
    }

    if (chunk.title.toLowerCase().includes(qLower)) score += 8;

    if (score > 0) {
      const acl = JSON.parse(chunk.aclJson || '{"kind":"org"}') as ACL;
      if (!checkAcl(acl, tenant.userId, tenant.role)) continue;

      results.push({
        chunkId: chunk.id,
        docId: chunk.docId,
        title: chunk.title,
        text: chunk.text,
        anchor: chunk.anchor ?? '',
        score,
        source: 'substring',
      });
    }
  }

  return results.sort((a, b) => b.score - a.score).slice(0, 20);
}

// ── Phase 4: Reciprocal Rank Fusion ─────────────────────────────────────

const RRF_K = 60;

function reciprocalRankFusion(results: SearchResult[]): SearchResult[] {
  // Group by source to get per-source rankings
  const bySource = new Map<string, SearchResult[]>();
  for (const r of results) {
    const list = bySource.get(r.source) || [];
    list.push(r);
    bySource.set(r.source, list);
  }

  // Sort each source list by score descending
  for (const list of bySource.values()) {
    list.sort((a, b) => b.score - a.score);
  }

  // Compute RRF scores
  const rrfScores = new Map<string, { result: SearchResult; score: number }>();

  for (const [, list] of bySource) {
    for (let rank = 0; rank < list.length; rank++) {
      const r = list[rank];
      const rrfScore = 1 / (RRF_K + rank + 1);
      const existing = rrfScores.get(r.chunkId);
      if (existing) {
        existing.score += rrfScore;
      } else {
        rrfScores.set(r.chunkId, { result: r, score: rrfScore });
      }
    }
  }

  return Array.from(rrfScores.values())
    .sort((a, b) => b.score - a.score)
    .map((e) => ({ ...e.result, score: e.score }));
}

// ── ACL check ───────────────────────────────────────────────────────────

function checkAcl(acl: ACL, userId: string | null, role: Role | null): boolean {
  switch (acl.kind) {
    case 'public':
      return true;
    case 'org':
      // org-level content is visible to all users within the tenant context,
      // including anonymous widget users (tenant resolved from /c/{slug}/)
      return true;
    case 'role':
      return role !== null && acl.roles.includes(role);
    case 'user':
      return userId !== null && acl.userIds.includes(userId);
  }
}