HaBraid/src/search/hybrid.ts

/**
 * Hybrid search combining BM25 (FTS5) and Vector similarity via RRF.
 *
 * Implements Reciprocal Rank Fusion to merge results from two search methods:
 * - BM25: keyword matching via SQLite FTS5
 * - Vector: semantic similarity via fastembed embeddings
 */

import type Database from "better-sqlite3";

import { searchItems } from "../db/items.js";
import { searchByVector } from "./vector.js";
import { embed } from "./embedder.js";
import { enrichSearchResults, type EnrichedSearchResult } from "./context.js";
import type { Item, ItemSource, SearchResult, WikiCategory } from "../types.js";

/** RRF constant k — prevents top ranks from dominating. */
const RRF_K = 60;

/**
 * Source-based boost multipliers for RRF scoring.
 * Manual/wiki content is curated and higher quality — boost it.
 * MemPalace raw drawers are noisy — dampen them.
 */
const SOURCE_BOOST: Record<string, number> = {
  manual: 1.5,
  wiki: 1.3,
  mempalace: 0.7,
};

/** Category-based boost — curated wiki pages get extra weight. */
const CATEGORY_BOOST: Record<string, number> = {
  projects: 1.4,
  decisions: 1.3,
  guides: 1.2,
  topics: 1.1,
  people: 1.1,
  infrastructure: 1.1,
};

/**
 * Compute a combined source + category boost multiplier for an item.
 */
function getBoost(source: string, category: string | null): number {
  const s = SOURCE_BOOST[source] ?? 1.0;
  const c = category ? (CATEGORY_BOOST[category] ?? 1.0) : 1.0;
  return s * c;
}

/** Search mode options. */
export type SearchMode = "keyword" | "semantic" | "hybrid";

/**
 * Result of a hybrid search.
 */
export interface HybridSearchResult {
  item: Item;
  score: number;
  snippet: string;
  sources: string[];
}

/**
 * Performs hybrid search combining BM25 and vector search with RRF fusion.
 *
 * @param db Connected database.
 * @param query Search query string.
 * @param mode Search mode: keyword, semantic, or hybrid.
 * @param limit Max results.
 * @param dbPath Path to the database file (for HNSW index).
 * @param enrichContext When true, returns enriched results with context tree data.
 * @returns Ranked search results, optionally enriched with context.
 */
export async function hybridSearch(
  db: Database.Database,
  query: string,
  mode: SearchMode = "hybrid",
  limit: number = 20,
  dbPath?: string,
  enrichContext: boolean = false,
): Promise<{
  results: HybridSearchResult[];
  enrichedResults?: EnrichedSearchResult[];
  mode: SearchMode;
  fallback?: string;
}> {
  const candidateMultiplier = 3;
  const candidateLimit = limit * candidateMultiplier;

  // 1. BM25 search (if keyword or hybrid)
  let bm25Results: SearchResult[] = [];
  if (mode === "keyword" || mode === "hybrid") {
    try {
      bm25Results = searchItems(db, { query, limit: candidateLimit });
    } catch {
      // FTS5 might fail on special characters
      bm25Results = [];
    }
  }

  // 2. Vector search (if semantic or hybrid)
  let vectorResults: Array<{ itemId: string; score: number }> = [];
  let usedMode = mode;
  let fallback: string | undefined;

  if (mode === "semantic" || mode === "hybrid") {
    try {
      const queryVector = await embed(query);
      vectorResults = searchByVector(db, queryVector, candidateLimit, dbPath);
    } catch {
      // Embedding failed — fall back to keyword
      if (mode === "semantic") {
        usedMode = "keyword";
        fallback = "Embedding unavailable, falling back to keyword search";
        bm25Results = searchItems(db, { query, limit: candidateLimit });
      }
    }

    if (vectorResults.length === 0) {
      usedMode = "keyword" as SearchMode;
      fallback = "No vectors indexed, using keyword search only";
      if (bm25Results.length === 0) {
        bm25Results = searchItems(db, { query, limit: candidateLimit });
      }
    }
  }

  // 3. If only one method has results, return directly
  if (bm25Results.length === 0 && vectorResults.length === 0) {
    return { results: [], mode: usedMode, fallback };
  }

  if (vectorResults.length === 0) {
    // BM25 only — apply source/category boost
    const results = bm25Results.slice(0, limit).map((r) => ({
      item: r.item,
      score: r.rank * getBoost(r.item.source, r.item.category),
      snippet: r.snippet,
      sources: ["bm25"],
    }));
    if (enrichContext) {
      const enrichedResults = enrichSearchResults(db, bm25Results.slice(0, limit));
      return { results, enrichedResults, mode: usedMode, fallback };
    }
    return { results, mode: usedMode, fallback };
  }

  if (bm25Results.length === 0) {
    // Vector only — need to load items, apply source/category boost, then re-sort
    const boostedResults = vectorResults.map((vr) => {
      const item = loadItemById(db, vr.itemId);
      if (!item) return null;
      return {
        item,
        score: vr.score * getBoost(item.source, item.category),
        snippet: item.content.slice(0, 200),
        sources: ["vector"],
      };
    }).filter((r): r is HybridSearchResult => r !== null);

    // Re-sort by boosted score descending
    boostedResults.sort((a, b) => b.score - a.score);

    // Normalize
    const maxScore = boostedResults[0]?.score ?? 1;
    const results: HybridSearchResult[] = boostedResults.slice(0, limit).map((r) => ({
      ...r,
      score: Math.round((r.score / maxScore) * 100) / 100,
    }));

    if (enrichContext) {
      const searchResults = boostedResults.slice(0, limit).map((r) => ({
        item: r.item, rank: r.score, snippet: r.snippet,
      }));
      const enrichedResults = enrichSearchResults(db, searchResults);
      return { results, enrichedResults, mode: usedMode, fallback };
    }
    return { results, mode: usedMode, fallback };
  }

  // 4. RRF fusion with source/category boost
  const rrfScores = new Map<string, { score: number; item: Item; snippet: string; sources: Set<string> }>();

  // BM25 ranks
  for (let rank = 0; rank < bm25Results.length; rank++) {
    const r = bm25Results[rank];
    const id = r.item.id;
    const boost = getBoost(r.item.source, r.item.category);
    const rrf = boost * (1 / (RRF_K + rank + 1));
    const existing = rrfScores.get(id);
    if (existing) {
      existing.score += rrf;
      existing.sources.add("bm25");
    } else {
      rrfScores.set(id, { score: rrf, item: r.item, snippet: r.snippet, sources: new Set(["bm25"]) });
    }
  }

  // Vector ranks
  for (let rank = 0; rank < vectorResults.length; rank++) {
    const vr = vectorResults[rank];
    const item = loadItemById(db, vr.itemId);
    if (!item) continue;

    const boost = getBoost(item.source, item.category);
    const rrf = boost * (1 / (RRF_K + rank + 1));
    const existing = rrfScores.get(vr.itemId);
    if (existing) {
      existing.score += rrf;
      existing.sources.add("vector");
    } else {
      rrfScores.set(vr.itemId, { score: rrf, item, snippet: item.content.slice(0, 200), sources: new Set(["vector"]) });
    }
  }

  // Sort by RRF score descending
  const sorted = [...rrfScores.values()].sort((a, b) => b.score - a.score);

  // Normalize scores to 0-1 range
  const maxScore = sorted[0]?.score ?? 1;
  const results: HybridSearchResult[] = sorted.slice(0, limit).map((r) => ({
    item: r.item,
    score: Math.round((r.score / maxScore) * 100) / 100,
    snippet: r.snippet,
    sources: [...r.sources],
  }));

  // Optional context enrichment
  if (enrichContext) {
    const searchResults: SearchResult[] = sorted.slice(0, limit).map((r) => ({
      item: r.item,
      rank: r.score,
      snippet: r.snippet,
    }));
    const enrichedResults = enrichSearchResults(db, searchResults);
    return { results, enrichedResults, mode: usedMode, fallback };
  }

  return { results, mode: usedMode, fallback };
}

/**
 * Loads a single item by ID.
 */
function loadItemById(db: Database.Database, itemId: string): Item | null {
  const row = db.prepare("SELECT * FROM items WHERE id = ?").get(itemId) as Record<string, unknown> | undefined;
  if (!row) return null;

  return {
    id: String(row.id),
    title: String(row.title),
    content: String(row.content),
    source: String(row.source) as ItemSource,
    category: row.category ? String(row.category) as WikiCategory : null,
    tags: JSON.parse(String(row.tags ?? "[]")),
    createdAt: String(row.created_at),
    updatedAt: String(row.updated_at),
    metadata: JSON.parse(String(row.metadata ?? "{}")),
  };
}