sweet-search 0.0.1 → 2.3.0

package/core/graph/graph-expansion.js

@@ -0,0 +1,839 @@
+/**
+ * Graph Expansion Module
+ *
+ * Expands search results by following relationship edges in the entity graph.
+ * Given top-k search results, performs 1-hop (or 2-hop) traversal to include
+ * related code chunks (imports, extends, implements, uses, calls).
+ *
+ * This helps with multi-hop coding questions where the answer spans
+ * multiple related entities.
+ */
+
+// Default edge types to follow during expansion
+const DEFAULT_EDGE_TYPES = new Set(['imports', 'extends', 'implements', 'uses', 'calls']);
+
+// --- Token Estimation Helpers ---
+
+// Language-specific tokens-per-line averages (from CodeSearchNet analysis)
+const TOKENS_PER_LINE = {
+  java: 15, kotlin: 14, swift: 13,
+  go: 12, c: 12, cpp: 12, php: 11,
+  javascript: 10, typescript: 10, jsx: 10, tsx: 10,
+  ruby: 9, python: 8,
+};
+
+// Map file extensions to language keys
+const EXT_TO_LANG = {
+  js: 'javascript', ts: 'typescript', py: 'python', rb: 'ruby',
+  kt: 'kotlin', cc: 'cpp', cxx: 'cpp', h: 'c', hpp: 'cpp', m: 'c',
+};
+
+/**
+ * Estimate token count from text using whitespace splitting.
+ * ±10-15% of real BPE counts, <0.1ms for typical chunks.
+ */
+function estimateTokenCount(text) {
+  if (!text) return 0;
+  return (text.match(/\S+/g) || []).length;
+}
+
+/**
+ * Fallback token estimate using language-specific multipliers.
+ * Much better than flat ×10 for mixed-language codebases.
+ */
+function fallbackTokenEstimate(result) {
+  const ext = (result.file_path || result.file || result.metadata?.file || result.metadata?.path || '')
+    .split('.').pop()?.toLowerCase();
+  const lang = result.metadata?.language || EXT_TO_LANG[ext] || ext;
+  const perLine = TOKENS_PER_LINE[lang] || 10;
+  const startLine = result.start_line || result.startLine || 0;
+  const endLine = result.end_line || result.endLine || startLine;
+  return Math.max(1, (endLine - startLine + 1)) * perLine;
+}
+
+/**
+ * Batch-load chunk texts from codebase.db vectors table.
+ * Accepts either a CodebaseRepository or a raw better-sqlite3 Database (legacy).
+ * @param {import('../infrastructure/codebase-repository.js').CodebaseRepository|import('better-sqlite3').Database} codebaseDbOrRepo
+ * @param {string[]} ids - Vector IDs to look up
+ * @returns {Map<string, string>} id → text
+ */
+export function loadChunkTexts(codebaseDbOrRepo, ids) {
+  if (!codebaseDbOrRepo || !ids || ids.length === 0) return new Map();
+  // Repository path (preferred)
+  if (typeof codebaseDbOrRepo.getChunkTexts === 'function') {
+    return codebaseDbOrRepo.getChunkTexts(ids);
+  }
+  // Legacy raw-DB path (backward compat)
+  try {
+    const ph = ids.map(() => '?').join(',');
+    const rows = codebaseDbOrRepo.prepare(
+      `SELECT id, text FROM vectors WHERE id IN (${ph})`
+    ).all(...ids);
+    return new Map(rows.map(r => [r.id, r.text]));
+  } catch {
+    return new Map();
+  }
+}
+
+/**
+ * Compute accurate token estimates for a mixed set of results.
+ * Original results (from HNSW) use codebaseDb text; expanded results
+ * use readFileLines (injected to keep this module import-free).
+ *
+ * @param {Array} results
+ * @param {Object} options
+ * @param {import('better-sqlite3').Database} [options.codebaseDb]
+ * @param {Function} [options.readFileLines] - (filePath, startLine, endLine) => string|null
+ * @returns {Map<number, number>} index → token count
+ */
+export function computeTokenEstimates(results, options = {}) {
+  const { codebaseDb, readFileLines } = options;
+  const estimates = new Map();
+
+  const originalIds = [];
+  const originalIndexes = [];
+
+  for (let i = 0; i < results.length; i++) {
+    const r = results[i];
+    if (!r.is_expanded && r.id) {
+      originalIds.push(r.id);
+      originalIndexes.push(i);
+    } else if (r.is_expanded && readFileLines) {
+      const filePath = r.file_path || r.file;
+      const startLine = r.start_line || r.startLine;
+      const endLine = r.end_line || r.endLine;
+      if (filePath && startLine) {
+        const text = readFileLines(filePath, startLine, endLine);
+        if (text) {
+          estimates.set(i, estimateTokenCount(text));
+        }
+      }
+    }
+  }
+
+  // Batch-load original chunk texts from codebase.db
+  const textMap = loadChunkTexts(codebaseDb, originalIds);
+  for (let j = 0; j < originalIds.length; j++) {
+    const text = textMap.get(originalIds[j]);
+    if (text) {
+      estimates.set(originalIndexes[j], estimateTokenCount(text));
+    }
+  }
+
+  return estimates;
+}
+
+// Score decay per hop (graph-expanded results are less relevant than direct hits)
+const HOP_DECAY = 0.6;
+const HOP2_DECAY = 0.35;
+
+// Edge priority scores for adaptive 2-hop ranking
+const EDGE_PRIORITY = {
+  extends: 4,
+  implements: 4,
+  imports: 3,
+  calls: 2,
+  uses: 1,
+};
+
+// PathRAG/LEGO-GraphRAG SOTA scoring constants for adaptive 2-hop
+const BASE_ALPHA = 0.55;
+const EDGE_ALPHA_BONUS = {
+  extends: 0.25,     // effective alpha = 0.80
+  implements: 0.25,  // effective alpha = 0.80
+  imports: 0.10,     // effective alpha = 0.65
+  calls: 0.05,       // effective alpha = 0.60
+  uses: 0.00,        // effective alpha = 0.55
+};
+const FLOW_THRESHOLD = 0.05;
+
+// Structural entity type boosts for reranking
+const TYPE_BOOST = {
+  class: 1.3,
+  function: 1.2,
+  method: 1.2,
+  interface: 1.3,
+  struct: 1.2,
+};
+
+function clampSemanticWeight(value) {
+  if (!Number.isFinite(value)) return 0.4;
+  return Math.max(0, Math.min(1, value));
+}
+
+function normalizeMinMax(values) {
+  if (values.length === 0) return [];
+  const min = Math.min(...values);
+  const max = Math.max(...values);
+  if (max === min) return values.map(() => 0.5);
+  return values.map(v => (v - min) / (max - min));
+}
+
+function blendScores(graphScore, cosineSim, weight) {
+  return (1 - weight) * graphScore + weight * cosineSim;
+}
+
+/**
+ * Expand search results using the entity relationship graph.
+ *
+ * @param {import('better-sqlite3').Database} db - The code graph database
+ * @param {Array} results - Initial search results with entity IDs
+ * @param {Object} options
+ * @param {string} options.expandMode - 'none' | '1hop' | '2hop'
+ * @param {number} options.maxExpanded - Max expanded results to add
+ * @param {number} options.tokenBudget - Max total tokens in expanded set
+ * @param {Set<string>} options.edgeTypes - Relationship types to follow
+ * @returns {Array} Expanded and reranked results
+ */
+export function expandResults(db, results, options = {}) {
+  const {
+    expandMode = '1hop',
+    maxExpanded = 10,
+    tokenBudget = 8000,
+    edgeTypes = DEFAULT_EDGE_TYPES,
+    adaptiveHop2 = false,
+    hop2TokenBudget = 4000,
+    expandedBudget,
+    queryInt8 = null,
+    hnswIndex = null,
+    semanticWeight = 0.4,
+    cosineSimilarity = null,
+    codebaseDb = null,
+    readFileLines = null,
+  } = options;
+  const clampedSemanticWeight = clampSemanticWeight(semanticWeight);
+
+  if (expandMode === 'none' || results.length === 0) return results;
+
+  // Collect entity IDs from results
+  const seedIds = collectSeedIds(db, results);
+  if (seedIds.size === 0) return results;
+
+  // 1-hop expansion: find neighbors via forward + reverse edges
+  const expanded = expandOneHop(db, seedIds, edgeTypes);
+
+  // 2-hop expansion (if requested)
+  if (expandMode === '2hop' && expanded.size > 0) {
+    if (adaptiveHop2) {
+      expandSecondHopAdaptive(db, seedIds, expanded, edgeTypes, {
+        maxHop2: maxExpanded,
+        tokenBudget: hop2TokenBudget,
+        queryInt8,
+        hnswIndex,
+        semanticWeight: clampedSemanticWeight,
+        cosineSimilarity,
+      });
+    } else {
+      expandSecondHop(db, seedIds, expanded, edgeTypes, {
+        queryInt8,
+        hnswIndex,
+        semanticWeight: clampedSemanticWeight,
+        cosineSimilarity,
+      });
+    }
+  }
+
+  if (expanded.size === 0) return results;
+
+  // Look up entity details for expanded IDs, respecting maxExpanded
+  const expandedIds = [...expanded.keys()].slice(0, maxExpanded);
+  const expandedResults = lookupEntities(db, expandedIds, expanded);
+
+  // Score expanded results relative to original results
+  const maxOriginalScore = Math.max(...results.map(r => r.score || 0), 1);
+  for (const er of expandedResults) {
+    const hops = er.expansion?.hops || 1;
+    const decay = er.expansion?.decay || (hops === 1 ? HOP_DECAY : HOP2_DECAY);
+    er.score = maxOriginalScore * decay;
+  }
+
+  // Rerank expanded results using composite scoring (file proximity + entity type + semantic)
+  rerankExpanded(expandedResults, results, {
+    queryInt8,
+    hnswIndex,
+    semanticWeight: clampedSemanticWeight,
+    cosineSimilarity,
+  });
+
+  // Apply token budget
+  const { results: budgeted, stats: budgetStats } = applyTokenBudget(
+    [...results, ...expandedResults], tokenBudget,
+    { expandedBudget, codebaseDb, readFileLines }
+  );
+
+  budgeted._budgetStats = budgetStats;
+  return budgeted;
+}
+
+/**
+ * Collect entity IDs from search results.
+ * Tries entity_id from metadata, then falls back to file_path + line range matching.
+ *
+ * @param {import('better-sqlite3').Database} db
+ * @param {Array} results
+ * @returns {Set<string>}
+ */
+function collectSeedIds(db, results) {
+  const seedIds = new Set();
+
+  for (const r of results) {
+    if (r.entity_id) seedIds.add(r.entity_id);
+    else if (r.metadata?.entity_id) seedIds.add(r.metadata.entity_id);
+    else if (r.id) seedIds.add(r.id);
+  }
+
+  if (seedIds.size > 0) return seedIds;
+
+  // Fallback: match results to entities by file_path + line range
+  let entityLookup;
+  try {
+    entityLookup = db.prepare(`
+      SELECT id, file_path, start_line, end_line
+      FROM entities WHERE stale_since IS NULL
+    `).all();
+  } catch {
+    return seedIds;
+  }
+
+  for (const r of results) {
+    const filePath = r.file_path || r.file || r.metadata?.file || r.metadata?.path;
+    const lineStart = r.start_line || r.startLine || r.metadata?.line_start || r.metadata?.startLine;
+    if (!filePath) continue;
+
+    for (const e of entityLookup) {
+      if (e.file_path === filePath && e.start_line != null && lineStart != null &&
+          e.start_line <= lineStart && e.end_line >= lineStart) {
+        seedIds.add(e.id);
+        break;
+      }
+    }
+  }
+
+  return seedIds;
+}
+
+/**
+ * Perform 1-hop graph expansion from seed entity IDs.
+ *
+ * @param {import('better-sqlite3').Database} db
+ * @param {Set<string>} seedIds
+ * @param {Set<string>} edgeTypes
+ * @returns {Map<string, {via: string, direction: string, score: number, hops?: number}>}
+ */
+export function expandOneHop(db, seedIds, edgeTypes) {
+  const expanded = new Map();
+  const seedArray = [...seedIds];
+  const placeholders = seedArray.map(() => '?').join(',');
+
+  // Forward edges: seed -> neighbor
+  let forwardRels;
+  try {
+    forwardRels = db.prepare(`
+      SELECT DISTINCT target_id, type FROM relationships
+      WHERE source_id IN (${placeholders}) AND target_id IS NOT NULL
+    `).all(...seedArray);
+  } catch {
+    forwardRels = [];
+  }
+
+  // Reverse edges: neighbor -> seed
+  let reverseRels;
+  try {
+    reverseRels = db.prepare(`
+      SELECT DISTINCT source_id, type FROM relationships
+      WHERE target_id IN (${placeholders}) AND source_id IS NOT NULL
+    `).all(...seedArray);
+  } catch {
+    reverseRels = [];
+  }
+
+  for (const { rels, idField, direction } of [
+    { rels: forwardRels, idField: 'target_id', direction: 'forward' },
+    { rels: reverseRels, idField: 'source_id', direction: 'reverse' },
+  ]) {
+    for (const rel of rels) {
+      const neighborId = rel[idField];
+      if (edgeTypes.has(rel.type) && !seedIds.has(neighborId)) {
+        const effectiveAlpha = BASE_ALPHA + (EDGE_ALPHA_BONUS[rel.type] || 0);
+        const score = effectiveAlpha;
+        const existing = expanded.get(neighborId);
+        if (!existing || score > existing.score) {
+          expanded.set(neighborId, { via: rel.type, direction, score });
+        }
+      }
+    }
+  }
+
+  return expanded;
+}
+
+/**
+ * Perform 2nd-hop expansion from the 1-hop neighbors.
+ *
+ * @param {import('better-sqlite3').Database} db
+ * @param {Set<string>} seedIds - Original seeds
+ * @param {Map<string, Object>} expanded - 1-hop expansion map (mutated in place)
+ * @param {Set<string>} edgeTypes
+ */
+export function expandSecondHop(db, seedIds, expanded, edgeTypes, options = {}) {
+  const {
+    queryInt8 = null,
+    hnswIndex = null,
+    semanticWeight = 0.4,
+    cosineSimilarity = null,
+  } = options;
+  const semanticEnabled = !!(queryInt8 && hnswIndex && cosineSimilarity && semanticWeight > 0);
+
+  const hop1Ids = [...expanded.keys()];
+  if (hop1Ids.length === 0) return;
+
+  const ph = hop1Ids.map(() => '?').join(',');
+
+  let hop2Forward;
+  try {
+    hop2Forward = db.prepare(`
+      SELECT source_id, target_id, type FROM relationships
+      WHERE source_id IN (${ph}) AND target_id IS NOT NULL
+    `).all(...hop1Ids);
+  } catch {
+    return;
+  }
+
+  if (!semanticEnabled) {
+    for (const rel of hop2Forward) {
+      if (edgeTypes.has(rel.type) && !seedIds.has(rel.target_id) && !expanded.has(rel.target_id)) {
+        expanded.set(rel.target_id, { via: rel.type, direction: 'forward', hops: 2 });
+      }
+    }
+    return;
+  }
+
+  const excluded = new Set([...seedIds, ...expanded.keys()]);
+  const candidates = [];
+  for (const rel of hop2Forward) {
+    if (!edgeTypes.has(rel.type) || excluded.has(rel.target_id)) continue;
+
+    const hop1Entry = expanded.get(rel.source_id);
+    const hop1Score = hop1Entry?.score ?? 1;  // identity: preserves old edgePriority × weight
+    const graphScore = hop1Score * (EDGE_PRIORITY[rel.type] || 1) * (rel.weight || 1.0);
+    let normSim = null;
+    const entityInt8 = hnswIndex.getInt8Vector(rel.target_id);
+    if (entityInt8) {
+      const cosSim = cosineSimilarity(queryInt8, entityInt8);
+      normSim = (cosSim + 1) / 2;
+    }
+    candidates.push({ rel, graphScore, normSim });
+  }
+
+  if (candidates.length === 0) return;
+  const normalizedGraphScores = normalizeMinMax(candidates.map(c => c.graphScore));
+  const bestByTarget = new Map();
+
+  for (let i = 0; i < candidates.length; i++) {
+    const c = candidates[i];
+    const normGraph = normalizedGraphScores[i];
+    let score = normGraph;
+    if (c.normSim != null) {
+      score = blendScores(normGraph, c.normSim, semanticWeight);
+    }
+
+    const prev = bestByTarget.get(c.rel.target_id);
+    if (!prev || score > prev.score) {
+      bestByTarget.set(c.rel.target_id, { rel: c.rel, score });
+    }
+  }
+
+  const ranked = [...bestByTarget.values()].sort((a, b) => b.score - a.score);
+  for (const c of ranked) {
+    expanded.set(c.rel.target_id, {
+      via: c.rel.type,
+      direction: 'forward',
+      hops: 2,
+    });
+  }
+}
+
+/**
+ * Perform adaptive 2nd-hop expansion with per-edge-type alpha decay,
+ * degree normalization, and flow-based early stopping (PathRAG-style).
+ *
+ * @param {import('better-sqlite3').Database} db
+ * @param {Set<string>} seedIds - Original seed entity IDs
+ * @param {Map<string, Object>} hop1Expanded - 1-hop expansion map (mutated in place)
+ * @param {Set<string>} edgeTypes - Allowed edge types
+ * @param {Object} options
+ * @param {number} options.maxHop2 - Max 2-hop entities to add
+ * @param {number} options.tokenBudget - Token budget for 2-hop expansion
+ * @returns {{ added: number, budgetUsed: number, candidates: number }}
+ */
+export function expandSecondHopAdaptive(db, seedIds, hop1Expanded, edgeTypes, options = {}) {
+  const {
+    maxHop2 = 5,
+    tokenBudget = 4000,
+    queryInt8 = null,
+    hnswIndex = null,
+    semanticWeight = 0.4,
+    cosineSimilarity = null,
+  } = options;
+  const semanticEnabled = !!(queryInt8 && hnswIndex && cosineSimilarity && semanticWeight > 0);
+
+  const hop1Ids = [...hop1Expanded.keys()];
+  if (hop1Ids.length === 0) return { added: 0, budgetUsed: 0, candidates: 0 };
+
+  const ph = hop1Ids.map(() => '?').join(',');
+
+  // Query out-degrees for hop-1 nodes, filtered to active edge types only.
+  // Counting all edge types would over-penalize nodes with many irrelevant edges.
+  // Safety: edgeTypes is always code-controlled (DEFAULT_EDGE_TYPES or intent policy
+  // constants). Not parameterized because better-sqlite3 doesn't support mixing
+  // positional params across two IN clauses cleanly. Never pass user input here.
+  const typeList = [...edgeTypes].map(t => `'${t}'`).join(',');
+  let degreeMap;
+  try {
+    const degRows = db.prepare(`
+      SELECT source_id, COUNT(*) as deg FROM relationships
+      WHERE source_id IN (${ph}) AND type IN (${typeList})
+      GROUP BY source_id
+    `).all(...hop1Ids);
+    degreeMap = new Map(degRows.map(r => [r.source_id, r.deg]));
+  } catch {
+    degreeMap = new Map();
+  }
+
+  // Query candidate 2-hop targets with source, weights, and line ranges
+  let rawCandidates;
+  try {
+    rawCandidates = db.prepare(`
+      SELECT r.source_id, r.target_id, r.type, r.weight, e.file_path, e.start_line, e.end_line
+      FROM relationships r
+      JOIN entities e ON e.id = r.target_id AND e.stale_since IS NULL
+      WHERE r.source_id IN (${ph}) AND r.target_id IS NOT NULL
+    `).all(...hop1Ids);
+  } catch {
+    return { added: 0, budgetUsed: 0, candidates: 0 };
+  }
+
+  // Filter by edge types and score all paths.
+  const excluded = new Set([...seedIds, ...hop1Expanded.keys()]);
+  const vectorCache = semanticEnabled ? new Map() : null;
+  const scoredCandidates = [];
+
+  for (const c of rawCandidates) {
+    if (!edgeTypes.has(c.type) || excluded.has(c.target_id)) continue;
+
+    const effectiveAlpha = BASE_ALPHA + (EDGE_ALPHA_BONUS[c.type] || 0);
+    const edgePriority = EDGE_PRIORITY[c.type] || 1;
+    const weight = c.weight || 1.0;
+    const outDegree = degreeMap.get(c.source_id) || 1;
+    const hop1Entry = hop1Expanded.get(c.source_id);
+    const hop1Score = hop1Entry?.score ?? effectiveAlpha;
+    const graphScore = (hop1Score * effectiveAlpha * weight * edgePriority) / Math.sqrt(outDegree);
+
+    let normSim = null;
+    if (semanticEnabled) {
+      if (!vectorCache.has(c.target_id)) {
+        vectorCache.set(c.target_id, hnswIndex.getInt8Vector(c.target_id));
+      }
+      const entityInt8 = vectorCache.get(c.target_id);
+      if (entityInt8) {
+        const cosSim = cosineSimilarity(queryInt8, entityInt8);
+        normSim = (cosSim + 1) / 2;
+      }
+    }
+
+    const estimatedTokens = fallbackTokenEstimate({
+      file_path: c.file_path,
+      start_line: c.start_line,
+      end_line: c.end_line,
+    });
+
+    scoredCandidates.push({
+      target_id: c.target_id,
+      source_id: c.source_id,
+      type: c.type,
+      graphScore,
+      normSim,
+      estimatedTokens,
+      effectiveAlpha,
+      outDegree,
+    });
+  }
+
+  if (scoredCandidates.length === 0) return { added: 0, budgetUsed: 0, candidates: 0 };
+
+  const normalizedGraphScores = semanticEnabled
+    ? normalizeMinMax(scoredCandidates.map(c => c.graphScore))
+    : [];
+
+  // Multiple hop-1 sources may reach the same target — keep the highest score.
+  const bestByTarget = new Map(); // target_id -> best scored entry
+  for (let i = 0; i < scoredCandidates.length; i++) {
+    const c = scoredCandidates[i];
+    let score = c.graphScore;
+
+    if (semanticEnabled) {
+      const normGraph = normalizedGraphScores[i];
+      score = normGraph;
+      if (c.normSim != null) {
+        score = blendScores(normGraph, c.normSim, semanticWeight);
+      }
+    }
+
+    // PathRAG-style early stopping
+    if (score < FLOW_THRESHOLD) continue;
+
+    const prev = bestByTarget.get(c.target_id);
+    if (prev && prev.score >= score) continue;
+    bestByTarget.set(c.target_id, { ...c, score });
+  }
+
+  const scored = [...bestByTarget.values()];
+
+  // Sort by composite score descending
+  scored.sort((a, b) => b.score - a.score);
+
+  // Greedily select candidates within token budget and maxHop2 limit
+  let budgetUsed = 0;
+  let count = 0;
+
+  for (const s of scored) {
+    if (count >= maxHop2) break;
+    if (budgetUsed + s.estimatedTokens > tokenBudget && count > 0) break;
+
+    const decay = s.effectiveAlpha * s.effectiveAlpha;
+
+    hop1Expanded.set(s.target_id, {
+      via: s.type,
+      direction: 'forward',
+      hops: 2,
+      adaptiveScore: s.score,
+      decay,
+      sourceOutDegree: s.outDegree,
+    });
+
+    budgetUsed += s.estimatedTokens;
+    count++;
+  }
+
+  return { added: count, budgetUsed, candidates: scored.length };
+}
+
+/**
+ * Look up entity details for expanded IDs.
+ *
+ * @param {import('better-sqlite3').Database} db
+ * @param {string[]} expandedIds
+ * @param {Map<string, Object>} expansionMeta
+ * @returns {Array}
+ */
+function lookupEntities(db, expandedIds, expansionMeta) {
+  if (expandedIds.length === 0) return [];
+
+  const ph = expandedIds.map(() => '?').join(',');
+  let entities;
+  try {
+    entities = db.prepare(`
+      SELECT id, file_path, type, name, signature, start_line, end_line
+      FROM entities WHERE id IN (${ph}) AND stale_since IS NULL
+    `).all(...expandedIds);
+  } catch {
+    return [];
+  }
+
+  return entities.map(e => ({
+    entity_id: e.id,
+    id: e.id,
+    file_path: e.file_path,
+    file: e.file_path,
+    name: e.name,
+    type: e.type,
+    signature: e.signature,
+    startLine: e.start_line,
+    endLine: e.end_line,
+    start_line: e.start_line,
+    end_line: e.end_line,
+    expansion: expansionMeta.get(e.id),
+    score: 0,
+    is_expanded: true,
+  }));
+}
+
+/**
+ * Rerank expanded results using a composite score that combines the
+ * decay-based score with file proximity and entity type relevance.
+ *
+ * Score factors:
+ * 1. Decay score (already assigned by the caller)
+ * 2. File proximity: entities in the same file as a seed result get a 1.5x boost
+ * 3. Entity type relevance: structural types (class, interface, function, method, struct)
+ *    receive a multiplicative boost (1.2-1.3x)
+ *
+ * Mutates `expandedResults` in place: updates `.score` and re-sorts descending.
+ *
+ * @param {Array} expandedResults - Expanded results with scores already assigned
+ * @param {Array} seedResults - Original seed results (used to determine file proximity)
+ * @returns {Array} The same array, sorted by reranked score descending
+ */
+export function rerankExpanded(expandedResults, seedResults, options = {}) {
+  const {
+    queryInt8 = null,
+    hnswIndex = null,
+    semanticWeight = 0.4,
+    cosineSimilarity = null,
+  } = options;
+  const clampedSemanticWeight = clampSemanticWeight(semanticWeight);
+  const semanticEnabled = !!(queryInt8 && hnswIndex && cosineSimilarity && clampedSemanticWeight > 0);
+
+  if (expandedResults.length === 0) return expandedResults;
+
+  const seedFiles = new Set(
+    seedResults.map(r => r.file_path || r.file || r.metadata?.path).filter(Boolean)
+  );
+
+  const baseScores = [];
+
+  for (const er of expandedResults) {
+    let rerankScore = er.score || 0;
+
+    // File proximity boost: entities in same file as seeds are more relevant
+    const erFile = er.file_path || er.file || er.metadata?.path;
+    if (erFile && seedFiles.has(erFile)) {
+      rerankScore *= 1.5;
+    }
+
+    // Entity type relevance: structural entities are more valuable
+    const entType = er.type || er.metadata?.chunk_type;
+    if (TYPE_BOOST[entType]) {
+      rerankScore *= TYPE_BOOST[entType];
+    }
+
+    baseScores.push(rerankScore);
+  }
+
+  if (!semanticEnabled) {
+    for (let i = 0; i < expandedResults.length; i++) {
+      expandedResults[i].score = baseScores[i];
+    }
+  } else {
+    const normalizedGraphScores = normalizeMinMax(baseScores);
+    for (let i = 0; i < expandedResults.length; i++) {
+      const er = expandedResults[i];
+      const normGraph = normalizedGraphScores[i];
+      let rerankScore = normGraph;
+      const entityId = er.entity_id || er.id;
+      const entityInt8 = hnswIndex.getInt8Vector(entityId);
+      if (entityInt8) {
+        const cosSim = cosineSimilarity(queryInt8, entityInt8);
+        const normSim = (cosSim + 1) / 2;
+        rerankScore = blendScores(normGraph, normSim, clampedSemanticWeight);
+      }
+      er.score = rerankScore;
+    }
+  }
+
+  // Re-sort by reranked score descending
+  expandedResults.sort((a, b) => (b.score || 0) - (a.score || 0));
+  return expandedResults;
+}
+
+/**
+ * Apply token budget to limit total result set size.
+ * Uses accurate token counts from chunk text when available,
+ * falls back to language-specific per-line multipliers.
+ *
+ * @param {Array} results
+ * @param {number} budget - Total token budget
+ * @param {Object} [options]
+ * @param {number} [options.expandedBudget] - Separate budget for expanded results
+ * @param {import('better-sqlite3').Database} [options.codebaseDb] - For chunk text lookup
+ * @param {Function} [options.readFileLines] - For expanded result text lookup
+ * @returns {{ results: Array, stats: { original: Object, hop1: Object, hop2: Object, total: Object } }}
+ */
+export function applyTokenBudget(results, budget, options = {}) {
+  const { expandedBudget, codebaseDb, readFileLines } = options;
+
+  // Pre-compute accurate token estimates when data sources are available
+  const tokenEstimates = (codebaseDb || readFileLines)
+    ? computeTokenEstimates(results, { codebaseDb, readFileLines })
+    : new Map();
+
+  let totalTokens = 0;
+  let expandedTokens = 0;
+  const budgeted = [];
+  const stats = {
+    original: { count: 0, tokens: 0 },
+    hop1: { count: 0, tokens: 0 },
+    hop2: { count: 0, tokens: 0 },
+    total: { count: 0, tokens: 0 },
+  };
+
+  for (let i = 0; i < results.length; i++) {
+    const r = results[i];
+    // Use accurate estimate if available, otherwise fall back to per-language heuristic
+    const accurate = tokenEstimates.get(i);
+    const estimatedTokens = (accurate != null && accurate > 0) ? accurate : fallbackTokenEstimate(r);
+    const isExpanded = !!r.is_expanded;
+    const hops = r.expansion?.hops || (isExpanded ? 1 : 0);
+
+    // Check total budget (always include at least one result)
+    if (totalTokens + estimatedTokens > budget && budgeted.length > 0) break;
+
+    // Check expanded-specific budget
+    if (isExpanded && expandedBudget != null && expandedTokens + estimatedTokens > expandedBudget) {
+      continue;
+    }
+
+    totalTokens += estimatedTokens;
+    if (isExpanded) expandedTokens += estimatedTokens;
+    budgeted.push(r);
+
+    // Track per-category stats
+    if (!isExpanded) {
+      stats.original.count++;
+      stats.original.tokens += estimatedTokens;
+    } else if (hops === 2) {
+      stats.hop2.count++;
+      stats.hop2.tokens += estimatedTokens;
+    } else {
+      stats.hop1.count++;
+      stats.hop1.tokens += estimatedTokens;
+    }
+  }
+
+  stats.total = { count: budgeted.length, tokens: totalTokens };
+  return { results: budgeted, stats };
+}
+
+/**
+ * Get expansion statistics for a set of entities.
+ *
+ * @param {import('better-sqlite3').Database} db
+ * @param {string[]} entityIds
+ * @returns {{ total: number, byType: Record<string, number> }}
+ */
+export function getExpansionStats(db, entityIds) {
+  if (!entityIds || entityIds.length === 0) return { total: 0, byType: {} };
+
+  const ph = entityIds.map(() => '?').join(',');
+  let rels;
+  try {
+    rels = db.prepare(`
+      SELECT type, COUNT(*) as count FROM relationships
+      WHERE (source_id IN (${ph}) OR target_id IN (${ph}))
+      AND source_id IS NOT NULL AND target_id IS NOT NULL
+      GROUP BY type
+    `).all(...entityIds, ...entityIds);
+  } catch {
+    return { total: 0, byType: {} };
+  }
+
+  const byType = {};
+  let total = 0;
+  for (const r of rels) {
+    byType[r.type] = r.count;
+    total += r.count;
+  }
+
+  return { total, byType };
+}