trellis 1.0.8 → 2.0.5

package/src/embeddings/search.ts

@@ -0,0 +1,305 @@
+/**
+ * Semantic Search Integration
+ *
+ * Connects the TrellisVcsEngine to the embedding system.
+ * Provides reindex (full rebuild) and search (query → ranked results).
+ * The embedder function is pluggable for testing with mock vectors.
+ *
+ * @see TRL-20
+ */
+
+import { join } from 'path';
+import { readFileSync, existsSync } from 'fs';
+import { VectorStore } from './store.js';
+import { embed } from './model.js';
+import {
+  chunkIssue,
+  chunkMilestone,
+  chunkDecision,
+  chunkMarkdown,
+  chunkCodeEntities,
+  chunkFile,
+} from './chunker.js';
+import type {
+  ChunkMeta,
+  EmbeddingRecord,
+  SearchOptions,
+  SearchResult,
+} from './types.js';
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+/** Minimal engine interface — avoids importing the full engine for testability */
+export interface SearchableEngine {
+  getRootPath(): string;
+  trackedFiles(): Array<{ path: string; contentHash?: string }>;
+  listIssues(filters?: any): Array<{
+    id: string;
+    title?: string;
+    description?: string;
+  }>;
+  listMilestones(): Array<{ id: string; message?: string }>;
+  queryDecisions?(): Array<{
+    id: string;
+    toolName: string;
+    rationale?: string;
+    context?: string;
+    outputSummary?: string;
+  }>;
+  parseFile?(filePath: string): any;
+}
+
+/** Embedder function type — maps text → vector. Pluggable for testing. */
+export type Embedder = (text: string) => Promise<Float32Array>;
+
+// ---------------------------------------------------------------------------
+// EmbeddingManager
+// ---------------------------------------------------------------------------
+
+export class EmbeddingManager {
+  private store: VectorStore;
+  private embedFn: Embedder;
+
+  constructor(dbPath: string, embedFn?: Embedder) {
+    this.store = new VectorStore(dbPath);
+    this.embedFn = embedFn ?? embed;
+  }
+
+  /**
+   * Full reindex: clear store, re-chunk all entities, embed, and insert.
+   */
+  async reindex(engine: SearchableEngine): Promise<{ chunks: number }> {
+    this.store.clear();
+
+    const allChunks: ChunkMeta[] = [];
+
+    // 1. Issues
+    const issues = engine.listIssues();
+    for (const issue of issues) {
+      allChunks.push(...chunkIssue(issue));
+    }
+
+    // 2. Milestones
+    const milestones = engine.listMilestones();
+    for (const ms of milestones) {
+      allChunks.push(...chunkMilestone(ms));
+    }
+
+    // 3. Decisions
+    if (engine.queryDecisions) {
+      const decisions = engine.queryDecisions();
+      for (const dec of decisions) {
+        allChunks.push(...chunkDecision(dec));
+      }
+    }
+
+    // 4. Files (markdown, summaries)
+    const rootPath = engine.getRootPath();
+    const trackedFiles = engine.trackedFiles();
+    for (const tf of trackedFiles) {
+      try {
+        const absPath = join(rootPath, tf.path);
+        if (!existsSync(absPath)) continue;
+        const content = readFileSync(absPath, 'utf-8');
+        allChunks.push(...chunkFile(tf.path, content));
+      } catch {}
+    }
+
+    // 5. Code entities (from parsed files)
+    if (engine.parseFile) {
+      for (const tf of trackedFiles) {
+        const ext = tf.path.split('.').pop()?.toLowerCase() ?? '';
+        if (
+          ![
+            'ts',
+            'js',
+            'tsx',
+            'jsx',
+            'py',
+            'go',
+            'rs',
+            'rb',
+            'java',
+            'cs',
+          ].includes(ext)
+        ) {
+          continue;
+        }
+        try {
+          const parsed = engine.parseFile(tf.path);
+          if (parsed && Array.isArray(parsed.entities)) {
+            const declarations = parsed.entities.map((e: any) => ({
+              id: e.id ?? e.name,
+              name: e.name,
+              kind: e.kind,
+              signature: e.signature ?? e.rawText?.split('\n')[0] ?? '',
+              docComment: e.docComment,
+            }));
+            allChunks.push(...chunkCodeEntities(tf.path, declarations));
+          }
+        } catch {}
+      }
+    }
+
+    // Embed and insert all chunks
+    const records: EmbeddingRecord[] = [];
+    for (const chunk of allChunks) {
+      try {
+        const vector = await this.embedFn(chunk.content);
+        records.push({ ...chunk, embedding: vector });
+      } catch {}
+    }
+
+    this.store.upsertBatch(records);
+
+    return { chunks: records.length };
+  }
+
+  /**
+   * Incrementally index a single file (on file change).
+   */
+  async indexFile(
+    filePath: string,
+    content: string,
+    engine?: SearchableEngine,
+  ): Promise<number> {
+    // Remove old chunks for this file
+    this.store.deleteByFile(filePath);
+
+    const chunks = chunkFile(filePath, content);
+
+    // Also index code entities if engine is available
+    if (engine?.parseFile) {
+      const ext = filePath.split('.').pop()?.toLowerCase() ?? '';
+      if (
+        [
+          'ts',
+          'js',
+          'tsx',
+          'jsx',
+          'py',
+          'go',
+          'rs',
+          'rb',
+          'java',
+          'cs',
+        ].includes(ext)
+      ) {
+        try {
+          const parsed = engine.parseFile(filePath);
+          if (parsed && Array.isArray(parsed.entities)) {
+            const declarations = parsed.entities.map((e: any) => ({
+              id: e.id ?? e.name,
+              name: e.name,
+              kind: e.kind,
+              signature: e.signature ?? e.rawText?.split('\n')[0] ?? '',
+              docComment: e.docComment,
+            }));
+            chunks.push(...chunkCodeEntities(filePath, declarations));
+          }
+        } catch {}
+      }
+    }
+
+    const records: EmbeddingRecord[] = [];
+    for (const chunk of chunks) {
+      try {
+        const vector = await this.embedFn(chunk.content);
+        records.push({ ...chunk, embedding: vector });
+      } catch {}
+    }
+
+    if (records.length > 0) {
+      this.store.upsertBatch(records);
+    }
+
+    return records.length;
+  }
+
+  /**
+   * Index an issue (on create/update).
+   */
+  async indexIssue(issue: {
+    id: string;
+    title?: string;
+    description?: string;
+  }): Promise<number> {
+    this.store.deleteByEntity(`issue:${issue.id}`);
+
+    const chunks = chunkIssue(issue);
+    const records: EmbeddingRecord[] = [];
+
+    for (const chunk of chunks) {
+      try {
+        const vector = await this.embedFn(chunk.content);
+        records.push({ ...chunk, embedding: vector });
+      } catch {}
+    }
+
+    if (records.length > 0) {
+      this.store.upsertBatch(records);
+    }
+
+    return records.length;
+  }
+
+  /**
+   * Index a milestone (on create).
+   */
+  async indexMilestone(milestone: {
+    id: string;
+    message?: string;
+  }): Promise<number> {
+    this.store.deleteByEntity(`milestone:${milestone.id}`);
+
+    const chunks = chunkMilestone(milestone);
+    const records: EmbeddingRecord[] = [];
+
+    for (const chunk of chunks) {
+      try {
+        const vector = await this.embedFn(chunk.content);
+        records.push({ ...chunk, embedding: vector });
+      } catch {}
+    }
+
+    if (records.length > 0) {
+      this.store.upsertBatch(records);
+    }
+
+    return records.length;
+  }
+
+  /**
+   * Semantic search: embed query → vector search → ranked results.
+   */
+  async search(query: string, opts?: SearchOptions): Promise<SearchResult[]> {
+    const queryVector = await this.embedFn(query);
+    return this.store.search(queryVector, opts);
+  }
+
+  /**
+   * Remove all data for a file.
+   */
+  removeFile(filePath: string): void {
+    this.store.deleteByFile(filePath);
+  }
+
+  /**
+   * Get store statistics.
+   */
+  stats(): { total: number; byType: Record<string, number> } {
+    return {
+      total: this.store.count(),
+      byType: this.store.countByType(),
+    };
+  }
+
+  /**
+   * Close the store.
+   */
+  close(): void {
+    this.store.close();
+  }
+}

package/src/embeddings/store.ts

@@ -0,0 +1,313 @@
+/**
+ * Embedding Vector Store
+ *
+ * Persistent storage for embedding vectors using bun:sqlite.
+ * Vectors are stored as Float32Array blobs; cosine similarity
+ * is computed in JavaScript for cross-platform portability.
+ *
+ * @see TRL-18
+ */
+
+import { Database } from 'bun:sqlite';
+import type {
+  ChunkMeta,
+  ChunkType,
+  EmbeddingRecord,
+  SearchOptions,
+  SearchResult,
+} from './types.js';
+
+// ---------------------------------------------------------------------------
+// Schema
+// ---------------------------------------------------------------------------
+
+const SCHEMA_SQL = `
+CREATE TABLE IF NOT EXISTS chunks (
+  id TEXT PRIMARY KEY,
+  entity_id TEXT NOT NULL,
+  content TEXT NOT NULL,
+  chunk_type TEXT NOT NULL,
+  file_path TEXT,
+  updated_at TEXT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS vectors (
+  id TEXT PRIMARY KEY,
+  embedding BLOB NOT NULL,
+  FOREIGN KEY (id) REFERENCES chunks(id) ON DELETE CASCADE
+);
+
+CREATE INDEX IF NOT EXISTS idx_chunks_entity ON chunks(entity_id);
+CREATE INDEX IF NOT EXISTS idx_chunks_type ON chunks(chunk_type);
+CREATE INDEX IF NOT EXISTS idx_chunks_file ON chunks(file_path);
+`;
+
+// ---------------------------------------------------------------------------
+// Vector Store
+// ---------------------------------------------------------------------------
+
+export class VectorStore {
+  private db: Database;
+
+  constructor(dbPath: string) {
+    this.db = new Database(dbPath);
+    this.db.exec('PRAGMA journal_mode=WAL;');
+    this.db.exec('PRAGMA foreign_keys=ON;');
+    this.db.exec(SCHEMA_SQL);
+  }
+
+  /**
+   * Insert or update a chunk with its embedding vector.
+   */
+  upsert(record: EmbeddingRecord): void {
+    const insertChunk = this.db.prepare(`
+      INSERT OR REPLACE INTO chunks (id, entity_id, content, chunk_type, file_path, updated_at)
+      VALUES ($id, $entityId, $content, $chunkType, $filePath, $updatedAt)
+    `);
+    const insertVector = this.db.prepare(`
+      INSERT OR REPLACE INTO vectors (id, embedding)
+      VALUES ($id, $embedding)
+    `);
+
+    const embeddingBlob = Buffer.from(record.embedding.buffer);
+
+    this.db.transaction(() => {
+      insertChunk.run({
+        $id: record.id,
+        $entityId: record.entityId,
+        $content: record.content,
+        $chunkType: record.chunkType,
+        $filePath: record.filePath ?? null,
+        $updatedAt: record.updatedAt,
+      });
+      insertVector.run({
+        $id: record.id,
+        $embedding: embeddingBlob,
+      });
+    })();
+  }
+
+  /**
+   * Batch upsert multiple records.
+   */
+  upsertBatch(records: EmbeddingRecord[]): void {
+    if (records.length === 0) return;
+
+    const insertChunk = this.db.prepare(`
+      INSERT OR REPLACE INTO chunks (id, entity_id, content, chunk_type, file_path, updated_at)
+      VALUES ($id, $entityId, $content, $chunkType, $filePath, $updatedAt)
+    `);
+    const insertVector = this.db.prepare(`
+      INSERT OR REPLACE INTO vectors (id, embedding)
+      VALUES ($id, $embedding)
+    `);
+
+    this.db.transaction(() => {
+      for (const record of records) {
+        const embeddingBlob = Buffer.from(record.embedding.buffer);
+        insertChunk.run({
+          $id: record.id,
+          $entityId: record.entityId,
+          $content: record.content,
+          $chunkType: record.chunkType,
+          $filePath: record.filePath ?? null,
+          $updatedAt: record.updatedAt,
+        });
+        insertVector.run({
+          $id: record.id,
+          $embedding: embeddingBlob,
+        });
+      }
+    })();
+  }
+
+  /**
+   * Delete a chunk and its vector by ID.
+   */
+  delete(id: string): void {
+    this.db.prepare('DELETE FROM vectors WHERE id = ?').run(id);
+    this.db.prepare('DELETE FROM chunks WHERE id = ?').run(id);
+  }
+
+  /**
+   * Delete all chunks for an entity.
+   */
+  deleteByEntity(entityId: string): void {
+    const ids = this.db
+      .prepare('SELECT id FROM chunks WHERE entity_id = ?')
+      .all(entityId) as Array<{ id: string }>;
+
+    if (ids.length === 0) return;
+
+    this.db.transaction(() => {
+      for (const { id } of ids) {
+        this.db.prepare('DELETE FROM vectors WHERE id = ?').run(id);
+        this.db.prepare('DELETE FROM chunks WHERE id = ?').run(id);
+      }
+    })();
+  }
+
+  /**
+   * Delete all chunks associated with a file path.
+   */
+  deleteByFile(filePath: string): void {
+    const ids = this.db
+      .prepare('SELECT id FROM chunks WHERE file_path = ?')
+      .all(filePath) as Array<{ id: string }>;
+
+    if (ids.length === 0) return;
+
+    this.db.transaction(() => {
+      for (const { id } of ids) {
+        this.db.prepare('DELETE FROM vectors WHERE id = ?').run(id);
+        this.db.prepare('DELETE FROM chunks WHERE id = ?').run(id);
+      }
+    })();
+  }
+
+  /**
+   * Get a chunk by ID (without vector).
+   */
+  getChunk(id: string): ChunkMeta | null {
+    const row = this.db
+      .prepare('SELECT * FROM chunks WHERE id = ?')
+      .get(id) as any;
+    if (!row) return null;
+    return rowToChunkMeta(row);
+  }
+
+  /**
+   * Search for chunks similar to the query vector.
+   * Uses brute-force cosine similarity scan.
+   */
+  search(queryVector: Float32Array, opts: SearchOptions = {}): SearchResult[] {
+    const limit = opts.limit ?? 10;
+    const minScore = opts.minScore ?? 0.0;
+
+    // Build SQL filter
+    const conditions: string[] = [];
+    const params: Record<string, unknown> = {};
+
+    if (opts.types && opts.types.length > 0) {
+      const placeholders = opts.types.map((_, i) => `$type${i}`).join(', ');
+      conditions.push(`c.chunk_type IN (${placeholders})`);
+      opts.types.forEach((t, i) => {
+        params[`$type${i}`] = t;
+      });
+    }
+
+    if (opts.filePrefix) {
+      conditions.push('c.file_path LIKE $filePrefix');
+      params.$filePrefix = `${opts.filePrefix}%`;
+    }
+
+    const where =
+      conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+
+    const sql = `
+      SELECT c.id, c.entity_id, c.content, c.chunk_type, c.file_path, c.updated_at,
+             v.embedding
+      FROM chunks c
+      JOIN vectors v ON c.id = v.id
+      ${where}
+    `;
+
+    const rows = this.db.prepare(sql).all(params) as any[];
+
+    // Compute cosine similarity for each row
+    const scored: SearchResult[] = [];
+    for (const row of rows) {
+      const storedVec = new Float32Array(
+        (row.embedding as Buffer).buffer,
+        (row.embedding as Buffer).byteOffset,
+        (row.embedding as Buffer).byteLength / 4,
+      );
+      const score = cosineSimilarity(queryVector, storedVec);
+      if (score >= minScore) {
+        scored.push({
+          chunk: rowToChunkMeta(row),
+          score,
+        });
+      }
+    }
+
+    // Sort by score descending and limit
+    scored.sort((a, b) => b.score - a.score);
+    return scored.slice(0, limit);
+  }
+
+  /**
+   * Get total count of chunks in the store.
+   */
+  count(): number {
+    const row = this.db
+      .prepare('SELECT COUNT(*) as cnt FROM chunks')
+      .get() as any;
+    return row?.cnt ?? 0;
+  }
+
+  /**
+   * Get count by chunk type.
+   */
+  countByType(): Record<string, number> {
+    const rows = this.db
+      .prepare(
+        'SELECT chunk_type, COUNT(*) as cnt FROM chunks GROUP BY chunk_type',
+      )
+      .all() as any[];
+    const result: Record<string, number> = {};
+    for (const row of rows) {
+      result[row.chunk_type] = row.cnt;
+    }
+    return result;
+  }
+
+  /**
+   * Clear all data from the store.
+   */
+  clear(): void {
+    this.db.exec('DELETE FROM vectors');
+    this.db.exec('DELETE FROM chunks');
+  }
+
+  /**
+   * Close the database connection.
+   */
+  close(): void {
+    this.db.close();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function rowToChunkMeta(row: any): ChunkMeta {
+  return {
+    id: row.id,
+    entityId: row.entity_id,
+    content: row.content,
+    chunkType: row.chunk_type as ChunkType,
+    filePath: row.file_path ?? undefined,
+    updatedAt: row.updated_at,
+  };
+}
+
+/**
+ * Compute cosine similarity between two vectors.
+ * Both vectors should already be normalized (output of mean pooling + normalize).
+ * For normalized vectors, cosine similarity = dot product.
+ */
+export function cosineSimilarity(a: Float32Array, b: Float32Array): number {
+  if (a.length !== b.length) return 0;
+  let dot = 0;
+  let normA = 0;
+  let normB = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+    normA += a[i] * a[i];
+    normB += b[i] * b[i];
+  }
+  const denom = Math.sqrt(normA) * Math.sqrt(normB);
+  return denom === 0 ? 0 : dot / denom;
+}

package/src/embeddings/types.ts

@@ -0,0 +1,85 @@
+/**
+ * Embedding Types
+ *
+ * Types for the semantic embedding and vector search system.
+ *
+ * @see TRL-18
+ */
+
+// ---------------------------------------------------------------------------
+// Chunk types — what gets embedded
+// ---------------------------------------------------------------------------
+
+export type ChunkType =
+  | 'issue_title'
+  | 'issue_desc'
+  | 'milestone_msg'
+  | 'decision_rationale'
+  | 'summary_md'
+  | 'code_entity'
+  | 'doc_comment'
+  | 'markdown';
+
+export interface ChunkMeta {
+  /** Unique chunk ID, e.g. "issue:TRL-5:title", "file:src/engine.ts:chunk:0" */
+  id: string;
+  /** EAV entity ID this chunk belongs to */
+  entityId: string;
+  /** Original text content */
+  content: string;
+  /** Chunk classification */
+  chunkType: ChunkType;
+  /** Source file path (nullable for non-file entities) */
+  filePath?: string;
+  /** When this chunk was last updated */
+  updatedAt: string;
+}
+
+// ---------------------------------------------------------------------------
+// Embedding record — chunk + vector
+// ---------------------------------------------------------------------------
+
+export interface EmbeddingRecord extends ChunkMeta {
+  /** 384-dimensional embedding vector */
+  embedding: Float32Array;
+}
+
+// ---------------------------------------------------------------------------
+// Search
+// ---------------------------------------------------------------------------
+
+export interface SearchResult {
+  /** Chunk metadata */
+  chunk: ChunkMeta;
+  /** Cosine similarity score (0..1) */
+  score: number;
+}
+
+export interface SearchOptions {
+  /** Max results to return (default: 10) */
+  limit?: number;
+  /** Filter by chunk type(s) */
+  types?: ChunkType[];
+  /** Filter by file path prefix */
+  filePrefix?: string;
+  /** Minimum similarity threshold (default: 0.0) */
+  minScore?: number;
+}
+
+// ---------------------------------------------------------------------------
+// Embedding model config
+// ---------------------------------------------------------------------------
+
+export interface EmbeddingModelConfig {
+  /** Model name for @xenova/transformers (default: "Xenova/all-MiniLM-L6-v2") */
+  modelName: string;
+  /** Embedding dimension (default: 384) */
+  dimension: number;
+  /** Cache directory for model files */
+  cacheDir?: string;
+}
+
+export const DEFAULT_MODEL_CONFIG: EmbeddingModelConfig = {
+  modelName: 'Xenova/all-MiniLM-L6-v2',
+  dimension: 384,
+};