@nomos-arc/arc 0.1.0

package/src/search/query-engine.ts

@@ -0,0 +1,225 @@
+import fs from 'node:fs/promises';
+import path from 'node:path';
+import type { Logger } from 'winston';
+import { NomosError } from '../core/errors.js';
+import type { IndexMetadata, NomosConfig, ProjectMap, SearchResult } from '../types/index.js';
+import { Embedder } from './embedder.js';
+import { GraphEnricher } from './graph-enricher.js';
+import { LexicalSearcher } from './lexical-searcher.js';
+import { VectorStore } from './vector-store.js';
+import { AuthManager } from '../core/auth/manager.js';
+
+// ─── QueryEngine ──────────────────────────────────────────────────────────────
+
+export class QueryEngine {
+  private _embedder: Embedder | null = null; // lazy — not created until first query
+  private readonly store: VectorStore;
+  private readonly enricher: GraphEnricher;
+  private readonly lexical: LexicalSearcher;
+  private readonly metaPath: string;
+  private readonly projectMapPath: string;
+
+  constructor(
+    private readonly projectRoot: string,
+    private readonly config: NomosConfig,
+    private readonly logger: Logger,
+    private readonly authManager?: AuthManager | null,
+  ) {
+    this.store = new VectorStore(
+      config.search.vector_store_path,
+      logger,
+      config.search.embedding_dimensions,
+    );
+    this.metaPath = path.join(config.search.vector_store_path, 'index-meta.json');
+    this.projectMapPath = path.join(projectRoot, config.graph.output_dir, 'project_map.json');
+    this.enricher = new GraphEnricher(this.projectMapPath, logger);
+    this.lexical = new LexicalSearcher();
+  }
+
+  /**
+   * Execute a hybrid search query.
+   *
+   * Pipeline:
+   * 1. If vector index exists → embed query → vector search → graph enrich
+   * 2. Find files NOT in the vector index (structural-only) → lexical search
+   * 3. Merge results (rank fusion: vector results ranked higher) → deduplicate → return
+   *
+   * If NO vector index exists, falls back entirely to lexical search against project_map.
+   */
+  async search(
+    query: string,
+    options?: { topK?: number; threshold?: number },
+  ): Promise<SearchResult[]> {
+    const topK = options?.topK ?? this.config.search.default_top_k;
+    const threshold = options?.threshold ?? this.config.search.default_threshold;
+
+    // Step 1: Validate query
+    if (!query.trim()) {
+      throw new NomosError('search_query_failed', 'Query must be a non-empty string.');
+    }
+
+    // Step 2: Try to load vector index metadata
+    let meta: IndexMetadata | null = null;
+    try {
+      const raw = await fs.readFile(this.metaPath, 'utf-8');
+      meta = JSON.parse(raw) as IndexMetadata;
+    } catch {
+      // No vector index — will use lexical-only search
+    }
+
+    // Step 3: Load project map (required for both lexical and graph enrichment)
+    let projectMap: ProjectMap;
+    try {
+      const raw = await fs.readFile(this.projectMapPath, 'utf-8');
+      projectMap = JSON.parse(raw) as ProjectMap;
+    } catch {
+      throw new NomosError(
+        'search_index_not_found',
+        'No project map found. Run: arc map',
+      );
+    }
+
+    let vectorResults: SearchResult[] = [];
+    let indexedFilePaths = new Set<string>();
+
+    // Step 4: Vector search (if index available)
+    if (meta !== null) {
+      if (meta.status === 'in_progress') {
+        this.logger.warn(
+          '[nomos:search:warn] Index is incomplete. Results may be partial.',
+        );
+      }
+      await this.warnIfIndexStale(meta);
+
+      try {
+        await this.store.init();
+        const queryVector = await (await this.getEmbedder()).embedOne(query.trim());
+        const rawResults = await this.store.query(queryVector, topK, threshold);
+        vectorResults = await this.enricher.enrich(rawResults);
+        indexedFilePaths = new Set(Object.keys(meta.files));
+      } catch (err) {
+        // Vector search failed — log and continue with lexical-only
+        this.logger.warn(
+          `[nomos:search:warn] Vector search failed, falling back to lexical search: ${
+            err instanceof Error ? err.message : String(err)
+          }`,
+        );
+      }
+    } else {
+      this.logger.info(
+        '[nomos:search:info] No vector index found. Using lexical search only.',
+      );
+    }
+
+    // Step 5: Lexical search for files NOT in the vector index
+    const nonIndexedFiles = new Set<string>();
+    for (const fp of Object.keys(projectMap.files)) {
+      if (!indexedFilePaths.has(fp)) {
+        nonIndexedFiles.add(fp);
+      }
+    }
+
+    let lexicalResults: SearchResult[] = [];
+    if (nonIndexedFiles.size > 0) {
+      lexicalResults = this.lexical.search(
+        query, projectMap, nonIndexedFiles, topK, threshold,
+      );
+      if (lexicalResults.length > 0) {
+        this.logger.info(
+          `[nomos:search:info] Lexical fallback found ${lexicalResults.length} results from ${nonIndexedFiles.size} non-indexed files.`,
+        );
+      }
+    }
+
+    // Step 6: Merge results (vector results naturally score higher)
+    const merged = [...vectorResults, ...lexicalResults];
+
+    // Step 7: De-duplicate
+    const deduped = this.deduplicate(merged);
+
+    // Step 8: Sort by similarity_score descending
+    deduped.sort((a, b) => b.similarity_score - a.similarity_score);
+
+    // Step 9: Return top-K
+    return deduped.slice(0, topK);
+  }
+
+  /** Lazy Embedder accessor — throws at construction time only if actually needed. */
+  private async getEmbedder(): Promise<Embedder> {
+    if (!this._embedder) {
+      this._embedder = await Embedder.create(
+        this.config.search, this.logger, this.authManager,
+      );
+    }
+    return this._embedder;
+  }
+
+  // ─── Private helpers ────────────────────────────────────────────────────────
+
+  /**
+   * De-duplication [TRAP-3]:
+   * For each file_path that has BOTH a 'file' type result AND one or more 'symbol'
+   * type results:
+   *   - Compute score gap = abs(file_result.similarity_score - max(symbol_scores))
+   *   - If gap <= 0.05: REMOVE the file-level result (keep only symbol results)
+   *   - If gap > 0.05: keep both (sufficiently distinct in relevance)
+   */
+  private deduplicate(results: SearchResult[]): SearchResult[] {
+    // Group by file_path
+    const byFile = new Map<string, { file: SearchResult | null; symbols: SearchResult[] }>();
+
+    for (const r of results) {
+      if (!byFile.has(r.file_path)) {
+        byFile.set(r.file_path, { file: null, symbols: [] });
+      }
+      const group = byFile.get(r.file_path)!;
+      if (r.type === 'file') {
+        group.file = r;
+      } else {
+        group.symbols.push(r);
+      }
+    }
+
+    const output: SearchResult[] = [];
+
+    for (const group of byFile.values()) {
+      if (group.file && group.symbols.length > 0) {
+        const maxSymbolScore = Math.max(...group.symbols.map((s) => s.similarity_score));
+        const gap = Math.abs(group.file.similarity_score - maxSymbolScore);
+
+        if (gap <= 0.05) {
+          // TRAP-3: file result is within 0.05 of best symbol — remove file-level result
+          output.push(...group.symbols);
+        } else {
+          // gap > 0.05: both are sufficiently distinct, keep both
+          output.push(group.file, ...group.symbols);
+        }
+      } else {
+        // Only file-level, or only symbol-level — keep everything
+        if (group.file) output.push(group.file);
+        output.push(...group.symbols);
+      }
+    }
+
+    return output;
+  }
+
+  /**
+   * Warn if the vector index is older than the project map [STALE INDEX WARNING].
+   */
+  private async warnIfIndexStale(meta: IndexMetadata): Promise<void> {
+    try {
+      const raw = await fs.readFile(this.projectMapPath, 'utf-8');
+      const projectMap = JSON.parse(raw) as { generated_at: string };
+      const indexTime = new Date(meta.last_full_index).getTime();
+      const mapTime = new Date(projectMap.generated_at).getTime();
+      if (indexTime < mapTime) {
+        this.logger.warn(
+          '[nomos:search:warn] Index is older than project map. Consider running: arc index --incremental',
+        );
+      }
+    } catch {
+      // project_map.json may not exist — non-fatal, skip warning
+    }
+  }
+}

package/src/search/vector-store.ts

@@ -0,0 +1,311 @@
+import * as lancedb from '@lancedb/lancedb';
+import { makeArrowTable } from '@lancedb/lancedb';
+import { Field, FixedSizeList, Float32, Int32, Schema, Utf8 } from 'apache-arrow';
+import type { Connection } from '@lancedb/lancedb';
+import type { Logger } from 'winston';
+import { NomosError } from '../core/errors.js';
+import type { VectorRecord } from '../types/index.js';
+
+const LIVE_TABLE = 'nomos_vectors';
+const STAGING_TABLE = 'nomos_vectors_staging';
+
+/**
+ * Explicit Arrow schema for VectorRecord rows.
+ * Required because LanceDB cannot infer types for nullable (null-valued) fields.
+ * vector_dims is parameterized — must match embedding_dimensions in config.
+ */
+function makeVectorSchema(vectorDims: number): Schema {
+  return new Schema([
+    new Field('id', new Utf8(), false),
+    new Field('type', new Utf8(), false),
+    new Field('vector', new FixedSizeList(vectorDims, new Field('item', new Float32(), true)), false),
+    new Field('file_path', new Utf8(), false),
+    new Field('module', new Utf8(), false),
+    new Field('purpose', new Utf8(), false),
+    new Field('symbol_name', new Utf8(), true),   // nullable
+    new Field('symbol_type', new Utf8(), true),   // nullable
+    new Field('line_start', new Int32(), true),   // nullable
+    new Field('line_end', new Int32(), true),     // nullable
+    new Field('parent_file_id', new Utf8(), true), // nullable
+    new Field('graph_depth', new Int32(), false),
+    new Field('dependents_count', new Int32(), false),
+    new Field('last_indexed', new Utf8(), false),
+    new Field('content_hash', new Utf8(), false),
+  ]);
+}
+
+type PlainRow = Record<string, unknown>;
+type RawQueryRow = PlainRow & { _distance: number };
+
+export class VectorStore {
+  private db: Connection | null = null;
+  private hasMergeInsert: boolean = false; // [TRAP-1] detected at init
+  private vectorDims: number;
+
+  constructor(
+    private readonly storePath: string,
+    private readonly logger: Logger,
+    vectorDims: number = 768,
+  ) {
+    this.vectorDims = vectorDims;
+  }
+
+  /**
+   * Initialize the DB connection. Detects mergeInsert capability. [GAP-2] wrapped in try-catch.
+   */
+  async init(): Promise<void> {
+    try {
+      this.db = await lancedb.connect(this.storePath);
+    } catch (err) {
+      throw new NomosError(
+        'search_index_corrupted',
+        `Failed to open vector store at ${this.storePath}. ` +
+          `The index may be corrupted. Run: arc index --force\n` +
+          `Original error: ${err instanceof Error ? err.message : String(err)}`,
+      );
+    }
+
+    // [TRAP-1] Detect mergeInsert capability
+    try {
+      const names = await this.db.tableNames();
+      if (names.includes(LIVE_TABLE)) {
+        const table = await this.db.openTable(LIVE_TABLE);
+        this.hasMergeInsert = typeof table.mergeInsert === 'function';
+      }
+    } catch {
+      this.hasMergeInsert = false;
+    }
+  }
+
+  /**
+   * Upsert a batch of vector records into the LIVE table.
+   * Uses mergeInsert if available; otherwise overwrites by id.
+   * Called per-batch during incremental indexing [BLOCKER-4].
+   */
+  async upsert(records: VectorRecord[]): Promise<void> {
+    this.assertConnected();
+    const schema = makeVectorSchema(this.vectorDims);
+    const names = await this.db!.tableNames();
+
+    if (!names.includes(LIVE_TABLE)) {
+      // Table doesn't exist yet — create with explicit schema
+      const arrowTable = makeArrowTable(this.toPlainRows(records), { schema });
+      await this.db!.createTable(LIVE_TABLE, arrowTable, { mode: 'overwrite' });
+      this.hasMergeInsert =
+        typeof (await this.db!.openTable(LIVE_TABLE)).mergeInsert === 'function';
+      return;
+    }
+
+    const table = await this.db!.openTable(LIVE_TABLE);
+
+    if (this.hasMergeInsert) {
+      const arrowTable = makeArrowTable(this.toPlainRows(records), { schema });
+      await table
+        .mergeInsert('id')
+        .whenMatchedUpdateAll()
+        .whenNotMatchedInsertAll()
+        .execute(arrowTable);
+    } else {
+      // [TRAP-1] Fallback: single-transaction overwrite — read existing, merge by id, overwrite.
+      // No delete-then-add: avoids a window where records are missing.
+      const existing = (await table.query().toArray()) as PlainRow[];
+      const merged = new Map(existing.map((r) => [r['id'] as string, r]));
+      for (const rec of this.toPlainRows(records)) {
+        merged.set(rec['id'] as string, rec);
+      }
+      const arrowTable = makeArrowTable(Array.from(merged.values()), { schema });
+      await this.db!.createTable(LIVE_TABLE, arrowTable, { mode: 'overwrite' });
+    }
+  }
+
+  /**
+   * Write records to the STAGING table for full re-index [BLOCKER-1].
+   * Does NOT touch the live table. Called per-batch.
+   */
+  async upsertToStaging(records: VectorRecord[]): Promise<void> {
+    this.assertConnected();
+    const schema = makeVectorSchema(this.vectorDims);
+    const names = await this.db!.tableNames();
+    const arrowTable = makeArrowTable(this.toPlainRows(records), { schema });
+
+    if (!names.includes(STAGING_TABLE)) {
+      await this.db!.createTable(STAGING_TABLE, arrowTable, { mode: 'overwrite' });
+    } else {
+      const staging = await this.db!.openTable(STAGING_TABLE);
+      await staging.add(arrowTable);
+    }
+  }
+
+  /**
+   * Atomic table swap: drop live table, rename staging → live [BLOCKER-1].
+   * LanceDB lacks a native rename; workaround: read staging → create live → drop staging.
+   * If staging doesn't exist, throws immediately — live table is never touched.
+   */
+  async promoteStagingToLive(): Promise<void> {
+    this.assertConnected();
+    const names = await this.db!.tableNames();
+
+    if (!names.includes(STAGING_TABLE)) {
+      throw new NomosError(
+        'search_index_failed',
+        'Staging table does not exist. Index may have failed before any data was written.',
+      );
+    }
+
+    // Drop old live table if it exists — committed to the swap from this point
+    if (names.includes(LIVE_TABLE)) {
+      await this.db!.dropTable(LIVE_TABLE);
+    }
+
+    // Read all staging data as an Arrow Table and re-create as the live table.
+    // Use toArrow() to preserve the exact Arrow schema from the staging table —
+    // avoids re-inferring types for Arrow-native vector fields.
+    const staging = await this.db!.openTable(STAGING_TABLE);
+    const stagingArrow = await staging.query().toArrow();
+    await this.db!.createTable(LIVE_TABLE, stagingArrow, { mode: 'overwrite' });
+    await this.db!.dropTable(STAGING_TABLE);
+
+    // Re-detect mergeInsert capability on new live table
+    try {
+      const live = await this.db!.openTable(LIVE_TABLE);
+      this.hasMergeInsert = typeof live.mergeInsert === 'function';
+    } catch {
+      this.hasMergeInsert = false;
+    }
+  }
+
+  /**
+   * Drop the staging table if it exists (cleanup after failed/interrupted index). [GAP-3]
+   */
+  async cleanupStaging(): Promise<void> {
+    this.assertConnected();
+    try {
+      const names = await this.db!.tableNames();
+      if (names.includes(STAGING_TABLE)) {
+        await this.db!.dropTable(STAGING_TABLE);
+        this.logger.info('[nomos:search] Cleaned up orphaned staging table from prior run.');
+      }
+    } catch (err) {
+      // Non-fatal: log and continue
+      this.logger.warn(
+        `[nomos:search:warn] Failed to clean up staging table: ${err instanceof Error ? err.message : String(err)}`,
+      );
+    }
+  }
+
+  /**
+   * Query the LIVE table with a vector. Returns top-K results above threshold.
+   *
+   * Distance type: `dot` (mathematically equivalent to cosine for L2-normalized vectors,
+   * which is what Gemini embeddings produce). Cosine SIMD path requires AVX2 which is not
+   * universally available; dot product path is hardware-portable. [S-3]
+   *
+   * Similarity conversion: similarity = 1 - _distance, then clamped to [0, 1].
+   * Result records do NOT contain the raw vector field. [S-5]
+   */
+  async query(
+    vector: Float32Array,
+    topK: number,
+    threshold: number,
+  ): Promise<Array<Omit<VectorRecord, 'vector'> & { similarity_score: number }>> {
+    this.assertConnected();
+    const names = await this.db!.tableNames();
+    if (!names.includes(LIVE_TABLE)) {
+      throw new NomosError(
+        'search_index_not_found',
+        `Vector index not found at ${this.storePath}. Run: arc index`,
+      );
+    }
+
+    const table = await this.db!.openTable(LIVE_TABLE);
+    const raw = (await table
+      .query()
+      .nearestTo(Array.from(vector))
+      .distanceType('dot') // equivalent to cosine for normalized vectors; hardware-portable
+      .limit(topK * 2) // over-fetch to allow post-filter
+      .toArray()) as RawQueryRow[];
+
+    return raw
+      .map((r) => {
+        const similarity = 1 - r._distance;
+        // [S-3] Clamp similarity ∈ [0, 1]
+        const similarity_score = Math.max(0, Math.min(1, similarity));
+
+        // Strip _distance and vector from output [S-5]
+        const { _distance, vector: _vec, ...rest } = r as RawQueryRow & { vector?: unknown };
+        void _distance;
+        void _vec;
+        return { ...(rest as Omit<VectorRecord, 'vector'>), similarity_score };
+      })
+      .filter((r) => r.similarity_score >= threshold)
+      .slice(0, topK);
+  }
+
+  /**
+   * Delete all records whose file_path matches any of the given paths.
+   * Used for incremental re-indexing: delete stale → upsert fresh.
+   */
+  async deleteByFilePaths(filePaths: string[]): Promise<void> {
+    this.assertConnected();
+    const names = await this.db!.tableNames();
+    if (!names.includes(LIVE_TABLE)) {
+      return; // Nothing to delete
+    }
+
+    if (filePaths.length === 0) return;
+
+    const table = await this.db!.openTable(LIVE_TABLE);
+    // Build SQL-safe IN clause — escape single quotes in paths
+    const escaped = filePaths.map((p) => `'${p.replace(/'/g, "''")}'`).join(', ');
+    await table.delete(`file_path IN (${escaped})`);
+  }
+
+  /**
+   * Return total record count in live table.
+   */
+  async count(): Promise<number> {
+    this.assertConnected();
+    const names = await this.db!.tableNames();
+    if (!names.includes(LIVE_TABLE)) {
+      return 0;
+    }
+    const table = await this.db!.openTable(LIVE_TABLE);
+    return await table.countRows();
+  }
+
+  // ─── Private helpers ────────────────────────────────────────────────────────
+
+  private assertConnected(): void {
+    if (!this.db) {
+      throw new NomosError(
+        'search_index_corrupted',
+        'VectorStore.init() must be called before any operation.',
+      );
+    }
+  }
+
+  /**
+   * Convert VectorRecord[] to plain row objects for LanceDB.
+   * vector: Float32Array → number[] (LanceDB handles typed arrays but plain arrays are safer).
+   * Null fields preserved as-is — schema declares them nullable.
+   */
+  private toPlainRows(records: VectorRecord[]): PlainRow[] {
+    return records.map((r) => ({
+      id: r.id,
+      type: r.type,
+      vector: Array.from(r.vector),
+      file_path: r.file_path,
+      module: r.module,
+      purpose: r.purpose,
+      symbol_name: r.symbol_name,
+      symbol_type: r.symbol_type,
+      line_start: r.line_start,
+      line_end: r.line_end,
+      parent_file_id: r.parent_file_id,
+      graph_depth: r.graph_depth,
+      dependents_count: r.dependents_count,
+      last_indexed: r.last_indexed,
+      content_hash: r.content_hash,
+    }));
+  }
+}