npm install node-llama-cpp @xenova/transformers
# Or: use @tylermarchetti/flashrank for a pre-packaged cross-encoder

// ✅ Cross-encoder reranker wrapper
// file: packages/db/src/retrieval/CrossEncoderReranker.ts

import { pipeline } from '@xenova/transformers'

export interface RankCandidate {
  id: string
  content: string
  heading: string
  documentId: string
  tokenCount: number
}

export interface RankedResult extends RankCandidate {
  rerankerScore: number
}

export class CrossEncoderReranker {
  private model: Awaited<ReturnType<typeof pipeline>> | null = null

  async load(): Promise<void> {
    // Loads cross-encoder/ms-marco-MiniLM-L-6-v2 (ONNX quantised, ~23MB)
    this.model = await pipeline(
      'text-classification',
      'Xenova/ms-marco-MiniLM-L-6-v2',
      { quantized: true },
    )
  }

  async rerank(
    query: string,
    candidates: RankCandidate[],
    topK = 5,
  ): Promise<RankedResult[]> {
    if (!this.model) throw new Error('Reranker not loaded. Call load() first.')
    if (candidates.length === 0) return []

    const start = performance.now()

    // Score all candidates in parallel (ONNX is CPU-bound, not GPU)
    const scores = await Promise.all(
      candidates.map(async (candidate) => {
        const result = await this.model!([
          { text: query, text_pair: candidate.content },
        ])
        return { ...candidate, rerankerScore: result[0].score as number }
      })
    )

    const elapsed = performance.now() - start
    if (elapsed > 150) {
      console.warn(`[Reranker] Scored ${candidates.length} candidates in ${elapsed.toFixed(0)}ms — approaching latency budget`)
    }

    // Sort descending by reranker score and return top K
    return scores
      .sort((a, b) => b.rerankerScore - a.rerankerScore)
      .slice(0, topK)
  }
}

// Singleton — load once at startup
export const reranker = new CrossEncoderReranker()

// ✅ Full two-stage retrieval pipeline
// file: packages/db/src/retrieval/RetrievalPipeline.ts

import { hybridSearch } from './HybridSearch'
import { reranker } from './CrossEncoderReranker'

export interface RetrievalOptions {
  candidateCount?: number   // Stage 1 candidates (default: 50)
  finalCount?: number       // Stage 2 reranked results (default: 5)
  efSearch?: number         // HNSW ef_search parameter (default: 40)
  rerankerTimeout?: number  // Max ms for reranker (default: 150)
}

export async function retrieve(
  orgId: string,
  query: string,
  options: RetrievalOptions = {},
): Promise<RankedResult[]> {
  const {
    candidateCount = 50,
    finalCount = 5,
    rerankerTimeout = 150,
  } = options

  // Stage 1: Hybrid retrieval
  const candidates = await hybridSearch(orgId, query, candidateCount)

  if (candidates.length === 0) return []

  // Stage 2: Cross-encoder reranking with timeout guard
  const rerankedPromise = reranker.rerank(query, candidates, finalCount)
  const timeoutPromise = new Promise<null>(resolve =>
    setTimeout(() => resolve(null), rerankerTimeout)
  )

  const result = await Promise.race([rerankedPromise, timeoutPromise])

  if (result === null) {
    // Reranker timed out — fall back to hybrid RRF order
    console.warn('[RetrievalPipeline] Reranker timeout — using hybrid RRF order')
    return candidates.slice(0, finalCount).map(c => ({ ...c, rerankerScore: 0 }))
  }

  return result
}