Reranker chain — Cohere → Workers AI → LLM

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

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

/**
 * Reranker abstraction — pluggable implementations.
 * All use the same interface so search.ts doesn't care which is active.
 */

export interface Reranker {
  /**
   * True when `rerank` returns absolute, calibrated relevance scores on a
   * stable 0..1 scale (Cohere v3.5). The scope gate (ADR-0112 D3) only
   * applies its relevance threshold when this is true — the LLM/noop
   * rerankers return rank-derived scores (top is always ~1.0), which cannot
   * gate relevance. When false, the gate falls back to "any hit at all".
   */
  calibrated: boolean;
  rerank(
    query: string,
    documents: string[],
    topN: number,
  ): Promise<Array<{ index: number; score: number }>>;
}

/**
 * LLM-as-reranker via OpenRouter.
 * Sends query + candidate docs to Haiku and asks it to rank by relevance.
 * Uses the same OPENROUTER_API_KEY as answer generation — no new key needed.
 * ~500ms for 10 docs, ~$0.0003/call with Haiku.
 */
export function createLLMReranker(
  openRouterKey: string,
  model = 'anthropic/claude-haiku-4-5',
): Reranker {
  return {
    // Scores are rank-derived (1 - rank/topN), not absolute relevance.
    calibrated: false,
    async rerank(query, documents, topN) {
      if (documents.length === 0) return [];
      if (documents.length <= topN) {
        return documents.map((_, i) => ({ index: i, score: 1 - i * 0.01 }));
      }

      const docList = documents.map((d, i) => `[${i}] ${d.slice(0, 300)}`).join('\n\n');

      try {
        const res = await fetch('https://openrouter.ai/api/v1/chat/completions', {
          method: 'POST',
          headers: {
            Authorization: `Bearer ${openRouterKey}`,
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            model,
            messages: [
              {
                role: 'system',
                content: `You are a relevance ranker. Given a query and numbered documents, return ONLY a JSON array of the top ${topN} most relevant document indices, ordered by relevance. Example: [3, 0, 7, 1, 5]. No explanation.`,
              },
              {
                role: 'user',
                content: `Query: ${query}\n\nDocuments:\n${docList}`,
              },
            ],
            max_tokens: 100,
            temperature: 0,
          }),
        });

        if (!res.ok) {
          return fallbackOrder(documents, topN);
        }

        const data = (await res.json()) as {
          choices: Array<{ message: { content: string } }>;
        };
        const content = data.choices?.[0]?.message?.content?.trim() ?? '';

        // Parse JSON array of indices
        const match = content.match(/\[[\d,\s]+\]/);
        if (!match) return fallbackOrder(documents, topN);

        const indices: number[] = JSON.parse(match[0]);
        const valid = indices.filter((i) => i >= 0 && i < documents.length).slice(0, topN);

        if (valid.length === 0) return fallbackOrder(documents, topN);

        return valid.map((idx, rank) => ({
          index: idx,
          score: 1 - rank * (1 / topN),
        }));
      } catch {
        return fallbackOrder(documents, topN);
      }
    },
  };
}

/**
 * Cohere Rerank v3.5 — dedicated reranker API, best quality.
 * Use if you have a separate COHERE_API_KEY.
 */
export function createCohereReranker(apiKey: string): Reranker {
  return {
    // `relevance_score` is an absolute 0..1 relevance — safe to threshold.
    calibrated: true,
    async rerank(query, documents, topN) {
      if (documents.length === 0) return [];

      const res = await fetch('https://api.cohere.com/v2/rerank', {
        method: 'POST',
        headers: {
          Authorization: `Bearer ${apiKey}`,
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          model: 'rerank-v3.5',
          query,
          documents,
          top_n: topN,
          return_documents: false,
        }),
      });

      if (!res.ok) {
        return fallbackOrder(documents, topN);
      }

      const data = (await res.json()) as {
        results: Array<{ index: number; relevance_score: number }>;
      };
      return data.results.map((r) => ({
        index: r.index,
        score: r.relevance_score,
      }));
    },
  };
}

/**
 * Cloudflare Workers AI reranker (`@cf/baai/bge-reranker-base`).
 *
 * Runs **in-network** — no external HTTPS hop, unlike Cohere/LLM — so it's the
 * preferred fallback when no Cohere key is present (cheaper + faster than the
 * external LLM reranker, better ordering than noop). The model returns a raw
 * logit per context; we map it to a 0..1 scale with a sigmoid for a stable,
 * comparable score. `id` in the response is the index into the input contexts.
 *
 * Marked `calibrated: false` deliberately: bge-reranker-base is the English-
 * leaning *base* model and the corpus is Thai, so we do NOT let it drive the
 * ADR-0112 scope-gate relevance floor until a Thai golden-set eval validates
 * the score calibration. It only reorders candidates today.
 */
export function createCfReranker(ai: Ai): Reranker {
  return {
    calibrated: false,
    async rerank(query, documents, topN) {
      if (documents.length === 0) return [];
      try {
        // @cloudflare/workers-types documents but does not declare `query` on
        // the bge-reranker input type, so a fresh object literal trips the
        // excess-property check. Build it as a variable (non-fresh objects skip
        // that check) to pass the runtime-required `query` without a cast.
        const input = { query, contexts: documents.map((text) => ({ text })), top_k: topN };
        const out = (await ai.run('@cf/baai/bge-reranker-base', input)) as {
          response?: Array<{ id: number; score: number }>;
        };
        const ranked = out.response;
        if (!ranked || ranked.length === 0) return fallbackOrder(documents, topN);
        return ranked
          .filter((r) => typeof r?.id === 'number' && r.id >= 0 && r.id < documents.length)
          .slice(0, topN)
          .map((r) => ({ index: r.id, score: 1 / (1 + Math.exp(-r.score)) }));
      } catch {
        return fallbackOrder(documents, topN);
      }
    },
  };
}

/** No-op reranker — returns original order. */
export function createNoopReranker(): Reranker {
  return {
    calibrated: false,
    async rerank(_query, documents, topN) {
      return fallbackOrder(documents, topN);
    },
  };
}

/**
 * Pick the active reranker by available credentials/bindings, in priority
 * order: Cohere v3.5 (calibrated, primary) → Workers AI bge-reranker
 * (in-network, no external hop) → LLM-as-reranker (external, only if there's
 * no AI binding) → noop. Centralizes the selection chain that was duplicated
 * across every RAG entry point (chat, answer, conversations, mcp, line).
 */
export function selectReranker(opts: {
  cohereKey?: string;
  ai?: Ai;
  openrouterKey?: string;
}): Reranker {
  if (opts.cohereKey) return createCohereReranker(opts.cohereKey);
  if (opts.ai) return createCfReranker(opts.ai);
  if (opts.openrouterKey) return createLLMReranker(opts.openrouterKey);
  return createNoopReranker();
}

function fallbackOrder(documents: string[], topN: number) {
  return documents.slice(0, topN).map((_, i) => ({ index: i, score: 1 - i * 0.01 }));
}