@operor/knowledge 0.1.0

package/src/RetrievalPipeline.ts

@@ -0,0 +1,388 @@
+import type { KBSearchResult, KBSearchOptions, KnowledgeStore } from './types.js';
+import type { EmbeddingService } from './EmbeddingService.js';
+import type { QueryRewriter } from './QueryRewriter.js';
+import { normalizeQuery } from './QueryNormalizer.js';
+import { reciprocalRankFusion, weightedScoreFusion } from './RankFusion.js';
+
+export interface RetrievalResult {
+  results: KBSearchResult[];
+  context: string;
+  isFaqMatch: boolean;
+  rewritten?: string;
+  /** Raw FAQ answer extracted from metadata (only set when isFaqMatch is true). */
+  faqAnswer?: string;
+  /** Raw FAQ question extracted from metadata (only set when isFaqMatch is true). */
+  faqQuestion?: string;
+  /** Multiple FAQ matches from compound query splitting. */
+  faqMatches?: Array<{ faqQuestion: string; faqAnswer: string; score: number }>;
+}
+
+export interface RetrievalPipelineOptions {
+  faqThreshold?: number;
+  faqLowThreshold?: number;
+  faqScoreGap?: number;
+  useHybridSearch?: boolean;
+  queryRewriter?: QueryRewriter;
+  rewriteHighThreshold?: number;
+  rewriteLowThreshold?: number;
+  fusionStrategy?: 'rrf' | 'weighted';
+}
+
+/**
+ * Heuristic splitter for compound questions. No LLM call — zero latency cost.
+ * Splits on "?" followed by more text, or " and " when both sides are >3 chars.
+ * Returns the original query in a single-element array if no split detected.
+ * Capped at 4 sub-queries max.
+ */
+export function splitCompoundQuery(query: string): string[] {
+  // Strategy 1: Split on "?" followed by more text
+  const qParts = query.split(/\?\s*/).filter(p => p.trim().length > 3);
+  if (qParts.length > 1) {
+    return qParts.slice(0, 4).map(p => p.trim());
+  }
+
+  // Strategy 2: Split on " and " when both sides are >3 chars
+  // Use word boundary to avoid splitting "android", "band", etc.
+  const andParts = query.split(/\s+and\s+/i).filter(p => p.trim().length > 3);
+  if (andParts.length > 1) {
+    return andParts.slice(0, 4).map(p => p.trim());
+  }
+
+  return [query];
+}
+
+export class RetrievalPipeline {
+  private store: KnowledgeStore;
+  private embedder: EmbeddingService;
+  private faqThreshold: number;
+  private faqLowThreshold: number;
+  private faqScoreGap: number;
+  private useHybridSearch: boolean;
+  private queryRewriter?: QueryRewriter;
+  private rewriteHighThreshold: number;
+  private rewriteLowThreshold: number;
+  private fusionStrategy: 'rrf' | 'weighted';
+
+  constructor(store: KnowledgeStore, embedder: EmbeddingService, faqThreshold?: number);
+  constructor(store: KnowledgeStore, embedder: EmbeddingService, options?: RetrievalPipelineOptions);
+  constructor(
+    store: KnowledgeStore,
+    embedder: EmbeddingService,
+    thresholdOrOptions?: number | RetrievalPipelineOptions,
+  ) {
+    this.store = store;
+    this.embedder = embedder;
+
+    if (typeof thresholdOrOptions === 'number') {
+      this.faqThreshold = thresholdOrOptions;
+      this.faqLowThreshold = 0.70;
+      this.faqScoreGap = 0.15;
+      this.useHybridSearch = true;
+      this.rewriteHighThreshold = 0.70;
+      this.rewriteLowThreshold = 0.50;
+      this.fusionStrategy = 'rrf';
+    } else {
+      const opts = thresholdOrOptions ?? {};
+      this.faqThreshold = opts.faqThreshold ?? 0.85;
+      this.faqLowThreshold = opts.faqLowThreshold ?? 0.70;
+      this.faqScoreGap = opts.faqScoreGap ?? 0.15;
+      this.useHybridSearch = opts.useHybridSearch ?? true;
+      this.queryRewriter = opts.queryRewriter;
+      this.rewriteHighThreshold = opts.rewriteHighThreshold ?? 0.70;
+      this.rewriteLowThreshold = opts.rewriteLowThreshold ?? 0.50;
+      this.fusionStrategy = opts.fusionStrategy ?? 'rrf';
+    }
+  }
+
+  async retrieve(query: string, options?: KBSearchOptions): Promise<RetrievalResult> {
+    const subQueries = splitCompoundQuery(query);
+
+    // Single query — use existing path
+    if (subQueries.length <= 1) {
+      return this.retrieveSingle(query, options);
+    }
+
+    // Compound query — retrieve each sub-query independently
+    const subResults = await Promise.all(
+      subQueries.map(sq => this.retrieveSingle(sq, options))
+    );
+
+    // Collect all FAQ matches, deduplicate by document ID
+    const seen = new Set<string>();
+    const faqMatches: Array<{ faqQuestion: string; faqAnswer: string; score: number }> = [];
+
+    for (const sr of subResults) {
+      if (sr.isFaqMatch && sr.faqAnswer && sr.faqQuestion) {
+        const docId = sr.results[0]?.document?.id;
+        if (docId && !seen.has(docId)) {
+          seen.add(docId);
+          faqMatches.push({
+            faqQuestion: sr.faqQuestion,
+            faqAnswer: sr.faqAnswer,
+            score: sr.results[0]?.score ?? 0,
+          });
+        }
+      }
+    }
+
+    // If 2+ FAQ matches, return combined result
+    if (faqMatches.length >= 2) {
+      const allResults = subResults.flatMap(sr => sr.results);
+      // Deduplicate results by chunk ID
+      const seenChunks = new Set<string>();
+      const dedupedResults = allResults.filter(r => {
+        if (seenChunks.has(r.chunk.id)) return false;
+        seenChunks.add(r.chunk.id);
+        return true;
+      });
+
+      return {
+        results: dedupedResults,
+        context: this.formatContext(dedupedResults),
+        isFaqMatch: true,
+        faqMatches,
+        // Use first match's answer for backward compat
+        faqAnswer: faqMatches[0].faqAnswer,
+        faqQuestion: faqMatches[0].faqQuestion,
+      };
+    }
+
+    // 0-1 FAQ matches — fall back to single-query retrieval with original query
+    return this.retrieveSingle(query, options);
+  }
+
+  private async retrieveSingle(query: string, options?: KBSearchOptions): Promise<RetrievalResult> {
+    // Layer 1: Normalize query (expand abbreviations, lowercase, collapse whitespace)
+    const normalized = normalizeQuery(query);
+    const embedding = await this.embedder.embed(normalized);
+
+    // FAQ fast-path: search FAQ docs first (top 2 for score gap analysis)
+    const faqResults = await this.store.searchByEmbedding(embedding, {
+      ...options,
+      sourceTypes: ['faq'],
+      limit: 2,
+    });
+
+    if (faqResults.length > 0) {
+      // FAQ freshness tiebreak: when top 2 are within 0.02, prefer the newer one
+      let top = faqResults[0];
+      if (faqResults.length > 1) {
+        const scoreDiff = top.score - faqResults[1].score;
+        if (scoreDiff <= 0.02 && (faqResults[1].document.updatedAt ?? 0) > (top.document.updatedAt ?? 0)) {
+          top = faqResults[1];
+        }
+      }
+
+      const faqAnswer = top.chunk.metadata?.answer || top.document.metadata?.answer;
+      const faqQuestion = top.chunk.metadata?.question || top.document.metadata?.question;
+
+      // Layer 2: Score gap analysis
+      // High confidence: score >= 0.85
+      if (top.score >= this.faqThreshold) {
+        return {
+          results: [top],
+          context: this.formatContext([top]),
+          isFaqMatch: true,
+          faqAnswer,
+          faqQuestion,
+        };
+      }
+
+      // Medium confidence with clear standout: score >= 0.70 && gap > 0.15
+      if (top.score >= this.faqLowThreshold) {
+        const gap = faqResults.length > 1 ? top.score - faqResults[1].score : 1;
+        if (gap > this.faqScoreGap) {
+          return {
+            results: [top],
+            context: this.formatContext([top]),
+            isFaqMatch: true,
+            faqAnswer,
+            faqQuestion,
+          };
+        }
+      }
+    }
+
+    // Layer 3: Hybrid search (vector + FTS5 keyword) with RRF
+    const results = await this.hybridSearch(normalized, embedding, options);
+
+    // Layer 4: Conditional LLM query rewrite
+    // Only attempt if: rewriter is configured, top score is in the "uncertain" band,
+    // and there's something in the KB worth re-matching against.
+    const topScore = results.length > 0 ? results[0].score : 0;
+    if (
+      this.queryRewriter &&
+      topScore >= this.rewriteLowThreshold &&
+      topScore < this.rewriteHighThreshold
+    ) {
+      try {
+        const rewriteResult = await this.queryRewriter.rewrite(normalized);
+        const rewrittenEmbedding = await this.embedder.embed(rewriteResult.rewritten);
+
+        // Re-run FAQ fast-path with rewritten query
+        const rewrittenFaqResults = await this.store.searchByEmbedding(rewrittenEmbedding, {
+          ...options,
+          sourceTypes: ['faq'],
+          limit: 2,
+        });
+
+        if (rewrittenFaqResults.length > 0 && rewrittenFaqResults[0].score >= this.faqLowThreshold) {
+          const top = rewrittenFaqResults[0];
+          const gap = rewrittenFaqResults.length > 1 ? top.score - rewrittenFaqResults[1].score : 1;
+          if (top.score >= this.faqThreshold || gap > this.faqScoreGap) {
+            const faqAnswer = top.chunk.metadata?.answer || top.document.metadata?.answer;
+            const faqQuestion = top.chunk.metadata?.question || top.document.metadata?.question;
+            return {
+              results: [top],
+              context: this.formatContext([top]),
+              isFaqMatch: true,
+              rewritten: rewriteResult.rewritten,
+              faqAnswer,
+              faqQuestion,
+            };
+          }
+        }
+
+        // Re-run hybrid search with rewritten query
+        const rewrittenResults = await this.hybridSearch(rewriteResult.rewritten, rewrittenEmbedding, options);
+        if (rewrittenResults.length > 0 && rewrittenResults[0].score > topScore) {
+          return {
+            results: rewrittenResults,
+            context: this.formatContext(rewrittenResults),
+            isFaqMatch: false,
+            rewritten: rewriteResult.rewritten,
+          };
+        }
+      } catch {
+        // Rewrite failed — fall through with original results
+      }
+    }
+
+    return {
+      results,
+      context: this.formatContext(results),
+      isFaqMatch: false,
+    };
+  }
+
+  private async hybridSearch(
+    query: string,
+    embedding: number[],
+    options?: KBSearchOptions,
+  ): Promise<KBSearchResult[]> {
+    const limit = options?.limit || 5;
+
+    // If hybrid search is disabled or store doesn't support keyword search, vector-only
+    if (!this.useHybridSearch || !this.store.searchByKeyword) {
+      const vecResults = await this.store.searchByEmbedding(embedding, {
+        ...options,
+        limit,
+      });
+      return this.applyBoosts(vecResults.slice(0, limit));
+    }
+
+    // Run vector search and FTS5 keyword search in parallel
+    const searchOpts = { ...options, limit: limit * 2 };
+    const [vecResults, ftsResults] = await Promise.all([
+      this.store.searchByEmbedding(embedding, searchOpts),
+      this.store.searchByKeyword(query, searchOpts),
+    ]);
+
+    // If FTS returned nothing, fall back to vector-only
+    if (ftsResults.length === 0) {
+      return this.applyBoosts(vecResults.slice(0, limit));
+    }
+
+    let fusedResults: KBSearchResult[];
+
+    if (this.fusionStrategy === 'weighted') {
+      // Weighted score fusion: combine actual scores
+      const vecItems = vecResults.map(r => ({ id: r.chunk.id, score: r.score }));
+      const ftsItems = ftsResults.map(r => ({ id: r.chunk.id, score: r.score }));
+      const fused = weightedScoreFusion(vecItems, ftsItems);
+
+      const resultMap = new Map<string, KBSearchResult>();
+      for (const r of vecResults) resultMap.set(r.chunk.id, r);
+      for (const r of ftsResults) {
+        if (!resultMap.has(r.chunk.id)) resultMap.set(r.chunk.id, r);
+      }
+
+      fusedResults = [];
+      for (const [chunkId, fusedScore] of fused) {
+        if (fusedResults.length >= limit) break;
+        const result = resultMap.get(chunkId);
+        if (result) {
+          // Use fused score as the result score for weighted strategy
+          fusedResults.push({ ...result, score: fusedScore });
+        }
+      }
+    } else {
+      // RRF: rank-based fusion (default)
+      const vecRanks = new Map<string, number>();
+      vecResults.forEach((r, i) => vecRanks.set(r.chunk.id, i));
+
+      const ftsRanks = new Map<string, number>();
+      ftsResults.forEach((r, i) => ftsRanks.set(r.chunk.id, i));
+
+      const fused = reciprocalRankFusion([vecRanks, ftsRanks]);
+
+      const resultMap = new Map<string, KBSearchResult>();
+      for (const r of vecResults) resultMap.set(r.chunk.id, r);
+      for (const r of ftsResults) {
+        if (!resultMap.has(r.chunk.id)) resultMap.set(r.chunk.id, r);
+      }
+
+      // RRF preserves original vector similarity scores for downstream thresholds
+      fusedResults = [];
+      for (const [chunkId, _rrfScore] of fused) {
+        if (fusedResults.length >= limit) break;
+        const result = resultMap.get(chunkId);
+        if (result) {
+          fusedResults.push(result);
+        }
+      }
+    }
+
+    return this.applyBoosts(fusedResults);
+  }
+
+  /**
+   * Apply freshness and priority boosts to search results, then re-sort.
+   */
+  private applyBoosts(results: KBSearchResult[]): KBSearchResult[] {
+    if (results.length === 0) return results;
+
+    const thirtyDaysAgo = Date.now() - (30 * 24 * 60 * 60 * 1000);
+
+    const boosted = results.map(r => {
+      let score = r.score;
+
+      // Freshness boost: +0.05 for docs updated within 30 days
+      if (r.document.updatedAt && r.document.updatedAt > thirtyDaysAgo) {
+        score += 0.05;
+      }
+
+      // Priority boost
+      const priority = (r.document as any).priority ?? 2;
+      if (priority === 1) score += 0.03;
+      else if (priority === 3) score -= 0.02;
+
+      return { ...r, score };
+    });
+
+    // Re-sort by boosted score
+    boosted.sort((a, b) => b.score - a.score);
+    return boosted;
+  }
+
+  private formatContext(results: KBSearchResult[]): string {
+    if (results.length === 0) return '';
+
+    const sections = results.map((r, i) => {
+      const source = r.document.title || r.document.sourceUrl || r.document.fileName || 'Unknown';
+      return `### Source ${i + 1}: ${source} (score: ${r.score.toFixed(2)})\n${r.chunk.content}`;
+    });
+
+    return `## Knowledge Base Context\n\n${sections.join('\n\n')}`;
+  }
+}