@claude-flow/embeddings/src/bm25.ts

V3 Embedding Service - OpenAI, Transformers.js, Agentic-Flow (ONNX), Mock providers with hyperbolic embeddings, normalization, and chunking

/**
 * ADR-121 Phase 17 — Okapi BM25 (sparse lexical retrieval).
 *
 * Dense vector retrieval (cosine over learned embeddings) is what
 * Phases 9-16 ship. It's powerful but has a known failure mode:
 * rare or technical terms that the embedding model never saw at
 * training time get averaged out in the dense vector, so docs that
 * literally contain "CVE-2024-12345" or "sha256:abc…" rank poorly
 * even when they're the exact match.
 *
 * The standard production fix is **hybrid retrieval**: combine a
 * sparse lexical signal (BM25 — what classical search engines use)
 * with the dense signal (cosine), then RRF-fuse the two ranked
 * lists. The sparse signal catches the rare-term docs the dense
 * signal underweights; the dense signal catches the semantic-only
 * matches the sparse signal misses.
 *
 * This module ships pure-function BM25 (Okapi variant — the modern
 * production default since Robertson-Walker 1994) with no external
 * deps. Composes with the Phase 11 RRF primitive for hybrid retrieval.
 *
 *   score(q, d) = Σ_{t in q} IDF(t) · (tf(t,d) · (k1 + 1)) / (tf(t,d) + k1·(1 - b + b·|d|/avgdl))
 *
 *   IDF(t) = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
 *
 * Parameters:
 *   k1 ∈ [1.2, 2.0] — term frequency saturation (default 1.5)
 *   b  ∈ [0.0, 1.0] — document-length normalization (default 0.75)
 */

export interface Bm25Options {
  /** Term frequency saturation. Default 1.5 (Okapi standard). */
  readonly k1?: number;
  /** Length normalization. Default 0.75. */
  readonly b?: number;
}

export interface Bm25Document {
  readonly id: string;
  readonly text: string;
}

export interface Bm25Index {
  readonly docs: ReadonlyArray<{ id: string; tokens: string[]; length: number }>;
  /** Document frequency: token → number of docs containing it. */
  readonly df: ReadonlyMap<string, number>;
  /** Number of documents. */
  readonly N: number;
  /** Average document length. */
  readonly avgdl: number;
  readonly k1: number;
  readonly b: number;
}

export interface Bm25Hit {
  readonly id: string;
  readonly score: number;
}

/**
 * Default tokenizer: lowercase, split on non-alphanumeric, drop
 * empty strings, drop tokens of length < 2 (stopword-ish noise).
 *
 * Production callers can supply their own tokenizer (e.g. stemmer,
 * stop-word filter, n-grams) by tokenizing themselves and passing
 * the pre-tokenized form via `buildBm25IndexFromTokens`.
 */
export function tokenize(text: string): string[] {
  if (typeof text !== 'string' || text.length === 0) return [];
  return text
    .toLowerCase()
    .split(/[^a-z0-9]+/i)
    .filter(t => t.length >= 2);
}

/**
 * Build a BM25 index over a corpus. O(N · avg-doc-length).
 *
 * The index is immutable — to add docs, build a new index. For
 * streaming corpora the right move is to maintain the (df, N, avgdl,
 * docs) structures explicitly and rebuild incrementally; this
 * module keeps it simple.
 */
export function buildBm25Index(corpus: ReadonlyArray<Bm25Document>, options: Bm25Options = {}): Bm25Index {
  const k1 = options.k1 ?? 1.5;
  const b = options.b ?? 0.75;
  if (k1 <= 0) throw new Error('k1 must be > 0');
  if (b < 0 || b > 1) throw new Error('b must be in [0, 1]');

  const docs: Array<{ id: string; tokens: string[]; length: number }> = [];
  const df = new Map<string, number>();
  let totalLen = 0;

  for (const d of corpus) {
    const tokens = tokenize(d.text);
    docs.push({ id: d.id, tokens, length: tokens.length });
    totalLen += tokens.length;
    // Count distinct tokens per doc for df.
    const seen = new Set<string>();
    for (const t of tokens) {
      if (!seen.has(t)) {
        seen.add(t);
        df.set(t, (df.get(t) ?? 0) + 1);
      }
    }
  }

  return {
    docs,
    df,
    N: docs.length,
    avgdl: docs.length > 0 ? totalLen / docs.length : 0,
    k1,
    b,
  };
}

/**
 * Score a single document against a query under the index. Used
 * by bm25TopK; exported for callers who want to score a specific
 * (query, doc) pair.
 */
export function bm25Score(
  queryTokens: ReadonlyArray<string>,
  doc: { tokens: string[]; length: number },
  index: Bm25Index,
): number {
  const { df, N, avgdl, k1, b } = index;
  // Count term frequencies in the doc once.
  const tfMap = new Map<string, number>();
  for (const t of doc.tokens) {
    tfMap.set(t, (tfMap.get(t) ?? 0) + 1);
  }
  let score = 0;
  for (const qt of queryTokens) {
    const tf = tfMap.get(qt) ?? 0;
    if (tf === 0) continue;
    const dfT = df.get(qt) ?? 0;
    const idf = Math.log(((N - dfT + 0.5) / (dfT + 0.5)) + 1);
    const numerator = tf * (k1 + 1);
    const denominator = tf + k1 * (1 - b + b * (doc.length / Math.max(avgdl, 1e-9)));
    score += idf * (numerator / denominator);
  }
  return score;
}

/**
 * Top-k BM25 search. Returns docs with score > 0 (docs that share
 * at least one token with the query), sorted descending. Ties
 * resolve by id ascending for determinism.
 *
 * Time complexity: O(N · |q|) where N = doc count, |q| = unique
 * query tokens.
 */
export function bm25TopK(
  query: string,
  index: Bm25Index,
  k: number,
): Bm25Hit[] {
  if (!Number.isInteger(k) || k <= 0) return [];
  const queryTokens = Array.from(new Set(tokenize(query)));
  if (queryTokens.length === 0) return [];

  const scored: Bm25Hit[] = [];
  for (const d of index.docs) {
    const score = bm25Score(queryTokens, d, index);
    if (score > 0) scored.push({ id: d.id, score });
  }
  scored.sort((a, b) => {
    if (b.score !== a.score) return b.score - a.score;
    return a.id < b.id ? -1 : a.id > b.id ? 1 : 0;
  });
  return scored.slice(0, k);
}

/**
 * Compute the inverse document frequency for a single token. Returns
 * 0 for tokens not in the corpus. Useful for the adaptive router's
 * rare-token-density signal.
 */
export function idfOf(token: string, index: Bm25Index): number {
  const dfT = index.df.get(token) ?? 0;
  if (dfT === 0) return 0;
  return Math.log(((index.N - dfT + 0.5) / (dfT + 0.5)) + 1);
}

/**
 * Rare-token density of a query relative to a BM25 index. Returns
 * the **mean IDF** of unique tokens in the query that appear in the
 * corpus at least once (tokens not in the corpus would have IDF 0
 * and would dilute the signal — they don't represent "rare
 * technical terms our retrieval can act on", they represent
 * "tokens our corpus doesn't have").
 *
 * Used by the adaptive router as the hybrid-retrieval signal: high
 * mean IDF means the query has rare/technical terms that dense
 * cosine retrieval tends to underweight.
 */
export function queryMeanIdf(query: string, index: Bm25Index): number {
  const queryTokens = Array.from(new Set(tokenize(query)));
  if (queryTokens.length === 0) return 0;
  let total = 0;
  let counted = 0;
  for (const t of queryTokens) {
    const idf = idfOf(t, index);
    if (idf > 0) {
      total += idf;
      counted++;
    }
  }
  return counted === 0 ? 0 : total / counted;
}