memorylink 1.0.0

package/src/search/vector-search.ts

@@ -0,0 +1,608 @@
+/**
+ * MemoryLink - Vector Search
+ * Semantic search using embeddings for better recall
+ * 
+ * v1.1: Vector search implementation
+ * 
+ * Supports:
+ * - OpenAI embeddings (if API key provided)
+ * - Local embeddings (optional, via transformers.js)
+ * - In-memory vector storage (for small datasets)
+ * - Optional persistent storage (SQLite + FAISS)
+ */
+
+import * as fs from 'fs';
+import * as path from 'path';
+import {
+  findRoot,
+  isInitialized,
+  getFilePath,
+  readJsonlEntries,
+  isJsonlFormat,
+  warn,
+  info,
+} from '../utils.js';
+import type {
+  MemoryEntry,
+  LearningEntry,
+  ChangeEntry,
+  BaseEntry,
+} from '../types.js';
+import {
+  isLocalEmbeddingsAvailable,
+  generateEmbeddingLocal,
+  generateEmbeddingsLocal,
+} from './local-embeddings.js';
+
+/**
+ * Vector embedding (1536 dimensions for OpenAI text-embedding-3-small, 384 for local models)
+ */
+export type Embedding = number[];
+
+/**
+ * Entry with embedding
+ */
+interface EmbeddedEntry {
+  id: string;
+  entry: BaseEntry;
+  embedding: Embedding;
+  text: string;
+}
+
+/**
+ * Embedding provider type
+ */
+type EmbeddingProvider = 'openai' | 'local';
+
+/**
+ * In-memory vector store with persistent storage support
+ */
+class VectorStore {
+  private entries: EmbeddedEntry[] = [];
+  private initialized = false;
+  private indexDir?: string;
+  private provider: EmbeddingProvider = 'openai';
+
+  /**
+   * Set index directory for persistent storage
+   */
+  setIndexDir(dir: string): void {
+    this.indexDir = dir;
+    // Ensure directory exists
+    if (!fs.existsSync(dir)) {
+      fs.mkdirSync(dir, { recursive: true });
+    }
+  }
+
+  /**
+   * Set embedding provider
+   */
+  setProvider(provider: EmbeddingProvider): void {
+    this.provider = provider;
+  }
+
+  /**
+   * Get embedding provider
+   */
+  getProvider(): EmbeddingProvider {
+    return this.provider;
+  }
+
+  /**
+   * Add entry with embedding
+   */
+  add(entry: BaseEntry, embedding: Embedding, text: string): void {
+    this.entries.push({
+      id: entry.id,
+      entry,
+      embedding,
+      text,
+    });
+    this.initialized = true;
+
+    // Persist to disk if index directory is set
+    if (this.indexDir) {
+      this.saveToDisk();
+    }
+  }
+
+  /**
+   * Save embeddings to disk (JSONL format)
+   */
+  private saveToDisk(): void {
+    if (!this.indexDir) return;
+
+    const embeddingsPath = path.join(this.indexDir, 'embeddings.jsonl');
+    const metadataPath = path.join(this.indexDir, 'metadata.json');
+
+    // Write embeddings (append mode for incremental updates)
+    const lines = this.entries.map(e => JSON.stringify({
+      id: e.id,
+      embedding: e.embedding,
+      text: e.text,
+      entry: e.entry,
+    }));
+    fs.writeFileSync(embeddingsPath, lines.join('\n') + '\n');
+
+    // Write metadata
+    const metadata = {
+      version: '1.0',
+      provider: this.provider,
+      count: this.entries.length,
+      lastUpdated: new Date().toISOString(),
+    };
+    fs.writeFileSync(metadataPath, JSON.stringify(metadata, null, 2));
+  }
+
+  /**
+   * Load embeddings from disk
+   */
+  loadFromDisk(): boolean {
+    if (!this.indexDir) return false;
+
+    const embeddingsPath = path.join(this.indexDir, 'embeddings.jsonl');
+    const metadataPath = path.join(this.indexDir, 'metadata.json');
+
+    if (!fs.existsSync(embeddingsPath)) {
+      return false;
+    }
+
+    try {
+      // Read metadata
+      if (fs.existsSync(metadataPath)) {
+        const metadata = JSON.parse(fs.readFileSync(metadataPath, 'utf-8'));
+        this.provider = metadata.provider || 'openai';
+      }
+
+      // Read embeddings
+      const lines = fs.readFileSync(embeddingsPath, 'utf-8').split('\n').filter(l => l.trim());
+      this.entries = [];
+      
+      for (const line of lines) {
+        try {
+          const data = JSON.parse(line);
+          this.entries.push({
+            id: data.id,
+            entry: data.entry,
+            embedding: data.embedding,
+            text: data.text,
+          });
+        } catch (e) {
+          // Skip invalid lines
+        }
+      }
+
+      this.initialized = this.entries.length > 0;
+      return this.initialized;
+    } catch (error) {
+      warn(`Failed to load embeddings from disk: ${error instanceof Error ? error.message : 'Unknown error'}`);
+      return false;
+    }
+  }
+
+  /**
+   * Search for similar entries
+   */
+  search(queryEmbedding: Embedding, limit: number = 10): Array<{ entry: BaseEntry; score: number; text: string }> {
+    if (!this.initialized || this.entries.length === 0) {
+      return [];
+    }
+
+    // Calculate cosine similarity for each entry
+    const scored = this.entries.map(item => ({
+      entry: item.entry,
+      text: item.text,
+      score: cosineSimilarity(queryEmbedding, item.embedding),
+    }));
+
+    // Sort by score (descending) and return top N
+    return scored
+      .sort((a, b) => b.score - a.score)
+      .slice(0, limit);
+  }
+
+  /**
+   * Get all entries count
+   */
+  size(): number {
+    return this.entries.length;
+  }
+
+  /**
+   * Clear all entries
+   */
+  clear(): void {
+    this.entries = [];
+    this.initialized = false;
+
+    // Clear disk storage
+    if (this.indexDir) {
+      const embeddingsPath = path.join(this.indexDir, 'embeddings.jsonl');
+      const metadataPath = path.join(this.indexDir, 'metadata.json');
+      if (fs.existsSync(embeddingsPath)) {
+        fs.unlinkSync(embeddingsPath);
+      }
+      if (fs.existsSync(metadataPath)) {
+        fs.unlinkSync(metadataPath);
+      }
+    }
+  }
+
+  /**
+   * Check if initialized
+   */
+  isInitialized(): boolean {
+    return this.initialized;
+  }
+}
+
+/**
+ * Calculate cosine similarity between two vectors
+ */
+function cosineSimilarity(vec1: Embedding, vec2: Embedding): number {
+  if (vec1.length !== vec2.length) {
+    return 0;
+  }
+
+  let dotProduct = 0;
+  let norm1 = 0;
+  let norm2 = 0;
+
+  for (let i = 0; i < vec1.length; i++) {
+    dotProduct += vec1[i] * vec2[i];
+    norm1 += vec1[i] * vec1[i];
+    norm2 += vec2[i] * vec2[i];
+  }
+
+  if (norm1 === 0 || norm2 === 0) {
+    return 0;
+  }
+
+  return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
+}
+
+/**
+ * Extract text from entry for embedding
+ */
+function extractTextForEmbedding(entry: BaseEntry): string {
+  let text = '';
+
+  // Memory entry
+  if ('details' in entry) {
+    const memEntry = entry as MemoryEntry;
+    text += memEntry.details || '';
+    if (memEntry.code) {
+      text += ' ' + memEntry.code;
+    }
+    if (memEntry.context) {
+      text += ' ' + memEntry.context;
+    }
+  }
+
+  // Learning entry
+  if ('problem' in entry) {
+    const learnEntry = entry as LearningEntry;
+    text += (learnEntry.problem || '') + ' ' + (learnEntry.solution || '');
+  }
+
+  // Change entry
+  if ('file' in entry) {
+    const changeEntry = entry as ChangeEntry;
+    text += (changeEntry.file || '') + ' ' + (changeEntry.reason || '');
+  }
+
+  // Tags
+  if (entry.tags && entry.tags.length > 0) {
+    text += ' ' + entry.tags.join(' ');
+  }
+
+  return text.trim();
+}
+
+/**
+ * Generate embeddings using OpenAI API
+ */
+async function generateEmbeddingOpenAI(text: string): Promise<Embedding> {
+  const apiKey = process.env.OPENAI_API_KEY;
+  if (!apiKey) {
+    throw new Error('OPENAI_API_KEY environment variable not set. Set it to use vector search.');
+  }
+
+  try {
+    // Dynamic import to avoid requiring openai if not used
+    const { OpenAI } = await import('openai');
+    const openai = new OpenAI({ apiKey });
+
+    const response = await openai.embeddings.create({
+      model: 'text-embedding-3-small',
+      input: text,
+    });
+
+    return response.data[0].embedding;
+  } catch (error) {
+    throw new Error(`Failed to generate embedding: ${error instanceof Error ? error.message : 'Unknown error'}`);
+  }
+}
+
+/**
+ * Generate embeddings for multiple texts (batch)
+ */
+async function generateEmbeddingsOpenAI(texts: string[]): Promise<Embedding[]> {
+  const apiKey = process.env.OPENAI_API_KEY;
+  if (!apiKey) {
+    throw new Error('OPENAI_API_KEY environment variable not set. Set it to use vector search.');
+  }
+
+  try {
+    const { OpenAI } = await import('openai');
+    const openai = new OpenAI({ apiKey });
+
+    const response = await openai.embeddings.create({
+      model: 'text-embedding-3-small',
+      input: texts,
+    });
+
+    return response.data.map(item => item.embedding);
+  } catch (error) {
+    throw new Error(`Failed to generate embeddings: ${error instanceof Error ? error.message : 'Unknown error'}`);
+  }
+}
+
+/**
+ * Retry wrapper with exponential backoff for embedding calls
+ */
+async function withRetry<T>(fn: () => Promise<T>, label: string, maxRetries = 3): Promise<T> {
+  let attempt = 0;
+  let lastErr: unknown;
+  while (attempt < maxRetries) {
+    try {
+      return await fn();
+    } catch (err) {
+      lastErr = err;
+      attempt += 1;
+      const delay = Math.min(1000 * 2 ** (attempt - 1), 5000);
+      warn(`Vector search: ${label} failed (attempt ${attempt}/${maxRetries}): ${err instanceof Error ? err.message : String(err)}. Retrying in ${delay}ms...`);
+      await new Promise(res => setTimeout(res, delay));
+    }
+  }
+  throw new Error(`${label} failed after ${maxRetries} attempts: ${lastErr instanceof Error ? lastErr.message : String(lastErr)}`);
+}
+
+// Global vector store instance
+let vectorStore: VectorStore | null = null;
+
+/**
+ * Initialize vector store and index existing entries
+ */
+export async function initializeVectorStore(root?: string, provider?: EmbeddingProvider): Promise<void> {
+  const projectRoot = root || findRoot();
+  if (!projectRoot || !isInitialized(projectRoot)) {
+    throw new Error('MemoryLink is not initialized. Run: memorylink init');
+  }
+
+  // Determine provider
+  let embeddingProvider: EmbeddingProvider = provider || 'openai';
+  
+  if (!provider) {
+    // Auto-detect: prefer local if available, otherwise OpenAI
+    if (isLocalEmbeddingsAvailable()) {
+      embeddingProvider = 'local';
+      info('Using local embeddings (offline mode)');
+    } else if (process.env.OPENAI_API_KEY) {
+      embeddingProvider = 'openai';
+      info('Using OpenAI embeddings');
+    } else {
+      throw new Error('No embedding provider available. Install @xenova/transformers for local embeddings or set OPENAI_API_KEY for OpenAI embeddings.');
+    }
+  }
+
+  // Validate provider
+  if (embeddingProvider === 'openai' && !process.env.OPENAI_API_KEY) {
+    throw new Error('OPENAI_API_KEY environment variable not set. Set it to use OpenAI embeddings, or install @xenova/transformers for local embeddings.');
+  }
+
+  if (embeddingProvider === 'local' && !isLocalEmbeddingsAvailable()) {
+    throw new Error('@xenova/transformers not installed. Install with: npm install @xenova/transformers');
+  }
+
+  vectorStore = new VectorStore();
+  vectorStore.setProvider(embeddingProvider);
+
+  // Set up persistent storage
+  const indexDir = path.join(projectRoot, '.memorylink', 'index');
+  vectorStore.setIndexDir(indexDir);
+
+  // Try to load from disk first
+  if (vectorStore.loadFromDisk()) {
+    const stats = getVectorStoreStats();
+    info(`Loaded ${stats.size} embeddings from disk`);
+    return; // Already indexed
+  }
+
+  const memoryPath = getFilePath('MEMORY', projectRoot);
+  const learningsPath = getFilePath('LEARNINGS', projectRoot);
+  const changesPath = getFilePath('CHANGES', projectRoot);
+
+  // Read all entries
+  const allEntries: BaseEntry[] = [];
+
+  if (isJsonlFormat(memoryPath)) {
+    const memoryEntries = readJsonlEntries<MemoryEntry>(memoryPath)
+      .filter(e => e.status === 'ACTIVE');
+    allEntries.push(...memoryEntries);
+  }
+
+  if (isJsonlFormat(learningsPath)) {
+    const learningEntries = readJsonlEntries<LearningEntry>(learningsPath)
+      .filter(e => e.status === 'ACTIVE');
+    allEntries.push(...learningEntries);
+  }
+
+  if (isJsonlFormat(changesPath)) {
+    const changeEntries = readJsonlEntries<ChangeEntry>(changesPath);
+    allEntries.push(...changeEntries);
+  }
+
+  if (allEntries.length === 0) {
+    return; // No entries to index
+  }
+
+  // Generate embeddings in batches
+  const batchSize = embeddingProvider === 'openai' ? 100 : 10; // Smaller batches for local
+  const texts = allEntries.map(extractTextForEmbedding);
+
+  for (let i = 0; i < texts.length; i += batchSize) {
+    const batch = texts.slice(i, i + batchSize);
+    const batchEntries = allEntries.slice(i, i + batchSize);
+
+    try {
+      let embeddings: Embedding[];
+      
+      if (embeddingProvider === 'openai') {
+        embeddings = await withRetry(() => generateEmbeddingsOpenAI(batch), 'generateEmbeddings');
+      } else {
+        // Local embeddings (no retry needed, but can add if needed)
+        embeddings = await generateEmbeddingsLocal(batch);
+      }
+      
+      // Add to vector store
+      for (let j = 0; j < batch.length; j++) {
+        vectorStore.add(batchEntries[j], embeddings[j], batch[j]);
+      }
+    } catch (error) {
+      warn(`Failed to index batch ${i / batchSize + 1}: ${error instanceof Error ? error.message : 'Unknown error'}`);
+      // Continue with next batch
+    }
+  }
+}
+
+/**
+ * Search using vector embeddings
+ */
+export async function vectorSearch(
+  query: string,
+  options: {
+    limit?: number;
+    threshold?: number;
+    root?: string;
+  } = {}
+): Promise<Array<{ entry: BaseEntry; score: number; text: string }>> {
+  const { limit = 10, threshold = 0.5, root } = options;
+
+  // Check if vector store is initialized
+  if (!vectorStore || !vectorStore.isInitialized()) {
+    // Try to load from disk first
+    if (root) {
+      const loaded = await loadVectorStoreFromDisk(root);
+      if (!loaded) {
+        // If not on disk, initialize fresh
+        await initializeVectorStore(root);
+      }
+    } else {
+      throw new Error('Vector store not initialized. Run: memorylink index --vector');
+    }
+  }
+
+  if (!vectorStore || vectorStore.size() === 0) {
+    return [];
+  }
+
+  // Generate query embedding based on provider
+  let queryEmbedding: Embedding;
+  const provider = vectorStore.getProvider();
+
+  if (provider === 'openai') {
+    if (!process.env.OPENAI_API_KEY) {
+      throw new Error('OPENAI_API_KEY environment variable not set. Set it to use OpenAI embeddings.');
+    }
+    queryEmbedding = await withRetry(() => generateEmbeddingOpenAI(query), 'generateEmbedding');
+  } else {
+    queryEmbedding = await generateEmbeddingLocal(query);
+  }
+
+  // Search
+  const results = vectorStore.search(queryEmbedding, limit * 2); // Get more results to filter by threshold
+
+  // Filter by threshold
+  return results.filter(result => result.score >= threshold).slice(0, limit);
+}
+
+/**
+ * Check if vector store is initialized (including from disk)
+ */
+export function isVectorStoreInitialized(root?: string): boolean {
+  if (vectorStore && vectorStore.isInitialized()) {
+    return true;
+  }
+
+  // Check if index exists on disk
+  if (root) {
+    const projectRoot = root || findRoot();
+    if (projectRoot && isInitialized(projectRoot)) {
+      const indexDir = path.join(projectRoot, '.memorylink', 'index');
+      const embeddingsPath = path.join(indexDir, 'embeddings.jsonl');
+      return fs.existsSync(embeddingsPath);
+    }
+  }
+
+  return false;
+}
+
+/**
+ * Load vector store from disk if available
+ */
+export async function loadVectorStoreFromDisk(root?: string): Promise<boolean> {
+  const projectRoot = root || findRoot();
+  if (!projectRoot || !isInitialized(projectRoot)) {
+    return false;
+  }
+
+  const indexDir = path.join(projectRoot, '.memorylink', 'index');
+  const embeddingsPath = path.join(indexDir, 'embeddings.jsonl');
+  const metadataPath = path.join(indexDir, 'metadata.json');
+
+  if (!fs.existsSync(embeddingsPath)) {
+    return false;
+  }
+
+  try {
+    // Read metadata to determine provider
+    let provider: EmbeddingProvider = 'openai';
+    if (fs.existsSync(metadataPath)) {
+      const metadata = JSON.parse(fs.readFileSync(metadataPath, 'utf-8'));
+      provider = metadata.provider || 'openai';
+    }
+
+    vectorStore = new VectorStore();
+    vectorStore.setProvider(provider);
+    vectorStore.setIndexDir(indexDir);
+    
+    return vectorStore.loadFromDisk();
+  } catch (error) {
+    warn(`Failed to load vector store from disk: ${error instanceof Error ? error.message : 'Unknown error'}`);
+    return false;
+  }
+}
+
+/**
+ * Get vector store statistics
+ */
+export function getVectorStoreStats(): { size: number; initialized: boolean } {
+  if (!vectorStore) {
+    return { size: 0, initialized: false };
+  }
+  return {
+    size: vectorStore.size(),
+    initialized: vectorStore.isInitialized(),
+  };
+}
+
+/**
+ * Clear vector store (for re-indexing)
+ */
+export function clearVectorStore(): void {
+  if (vectorStore) {
+    vectorStore.clear();
+  }
+}
+