cognitive-core 0.0.2 → 0.1.0

package/src/embeddings/bm25.ts

@@ -0,0 +1,337 @@
+/**
+ * BM25 Index
+ *
+ * A text-based similarity search index using BM25 (Best Match 25) algorithm.
+ * Used as a fallback when no embedding provider is configured.
+ *
+ * BM25 is a ranking function used in information retrieval that considers
+ * term frequency, document length, and inverse document frequency.
+ */
+
+export type MemoryItemType = 'experience' | 'strategy' | 'concept' | 'skill';
+
+interface Document {
+  id: string;
+  tokens: string[];
+  type: MemoryItemType;
+  content: string;
+}
+
+/**
+ * BM25 text similarity search index
+ */
+export class BM25Index {
+  private documents: Map<string, Document>;
+  private idf: Map<string, number>;
+  private avgDocLength: number;
+  private k1: number;
+  private b: number;
+  private dirty: boolean;
+
+  constructor(options?: { k1?: number; b?: number }) {
+    this.documents = new Map();
+    this.idf = new Map();
+    this.avgDocLength = 0;
+    // k1 controls term frequency saturation (typical: 1.2-2.0)
+    this.k1 = options?.k1 ?? 1.5;
+    // b controls document length normalization (typical: 0.75)
+    this.b = options?.b ?? 0.75;
+    this.dirty = false;
+  }
+
+  /**
+   * Add a document to the index
+   */
+  add(id: string, content: string, type: MemoryItemType): void {
+    const tokens = this.tokenize(content);
+    this.documents.set(id, { id, tokens, type, content });
+    this.dirty = true;
+  }
+
+  /**
+   * Remove a document from the index
+   */
+  remove(id: string): boolean {
+    const removed = this.documents.delete(id);
+    if (removed) {
+      this.dirty = true;
+    }
+    return removed;
+  }
+
+  /**
+   * Check if a document exists
+   */
+  has(id: string): boolean {
+    return this.documents.has(id);
+  }
+
+  /**
+   * Get document count
+   */
+  get size(): number {
+    return this.documents.size;
+  }
+
+  /**
+   * Search for similar documents
+   */
+  search(
+    query: string,
+    options: {
+      k?: number;
+      type?: MemoryItemType;
+      threshold?: number;
+    } = {}
+  ): Array<{ id: string; score: number; content: string }> {
+    // Rebuild IDF if needed
+    if (this.dirty) {
+      this.rebuildIdf();
+    }
+
+    const queryTokens = this.tokenize(query);
+
+    // Return empty for empty query
+    if (queryTokens.length === 0) {
+      return [];
+    }
+
+    const results: Array<{ id: string; score: number; content: string }> = [];
+
+    for (const doc of this.documents.values()) {
+      // Skip if type filter doesn't match
+      if (options.type && doc.type !== options.type) continue;
+
+      const score = this.computeScore(queryTokens, doc.tokens);
+
+      // Normalize score to 0-1 range (approximate)
+      const normalizedScore = this.normalizeScore(score, queryTokens.length);
+
+      if (normalizedScore >= (options.threshold ?? 0)) {
+        results.push({
+          id: doc.id,
+          score: normalizedScore,
+          content: doc.content,
+        });
+      }
+    }
+
+    return results.sort((a, b) => b.score - a.score).slice(0, options.k ?? 10);
+  }
+
+  /**
+   * Tokenize text for indexing and search
+   */
+  private tokenize(text: string): string[] {
+    return text
+      .toLowerCase()
+      .replace(/[^\w\s]/g, ' ')
+      .split(/\s+/)
+      .filter((t) => t.length > 1)
+      .filter((t) => !STOP_WORDS.has(t));
+  }
+
+  /**
+   * Rebuild IDF values for all terms
+   */
+  private rebuildIdf(): void {
+    this.idf.clear();
+    const N = this.documents.size;
+
+    if (N === 0) {
+      this.avgDocLength = 0;
+      this.dirty = false;
+      return;
+    }
+
+    const termDocCounts = new Map<string, number>();
+    let totalLength = 0;
+
+    for (const doc of this.documents.values()) {
+      totalLength += doc.tokens.length;
+      const uniqueTerms = new Set(doc.tokens);
+      for (const term of uniqueTerms) {
+        termDocCounts.set(term, (termDocCounts.get(term) ?? 0) + 1);
+      }
+    }
+
+    this.avgDocLength = totalLength / N;
+
+    // Calculate IDF for each term
+    for (const [term, df] of termDocCounts) {
+      // IDF formula: log((N - df + 0.5) / (df + 0.5) + 1)
+      this.idf.set(term, Math.log((N - df + 0.5) / (df + 0.5) + 1));
+    }
+
+    this.dirty = false;
+  }
+
+  /**
+   * Compute BM25 score for a document
+   */
+  private computeScore(queryTokens: string[], docTokens: string[]): number {
+    const docLength = docTokens.length;
+
+    // Build term frequency map for document
+    const termFreq = new Map<string, number>();
+    for (const token of docTokens) {
+      termFreq.set(token, (termFreq.get(token) ?? 0) + 1);
+    }
+
+    let score = 0;
+
+    for (const term of queryTokens) {
+      const tf = termFreq.get(term) ?? 0;
+      const idf = this.idf.get(term) ?? 0;
+
+      if (tf === 0 || idf === 0) continue;
+
+      // BM25 formula
+      const numerator = tf * (this.k1 + 1);
+      const denominator =
+        tf +
+        this.k1 * (1 - this.b + this.b * (docLength / this.avgDocLength));
+
+      score += idf * (numerator / denominator);
+    }
+
+    return score;
+  }
+
+  /**
+   * Normalize BM25 score to approximately 0-1 range
+   * This is an approximation since BM25 scores are unbounded
+   */
+  private normalizeScore(score: number, queryLength: number): number {
+    if (score <= 0 || queryLength === 0) return 0;
+
+    // Approximate maximum possible score per term
+    // Assuming perfect TF and high IDF
+    const maxScorePerTerm = 2.5; // Typical upper bound
+    const maxPossibleScore = queryLength * maxScorePerTerm;
+
+    // Sigmoid-like normalization
+    const normalized = score / (score + maxPossibleScore * 0.5);
+
+    return Math.min(1, Math.max(0, normalized));
+  }
+
+  /**
+   * Serialize index to JSON (for persistence)
+   */
+  toJSON(): {
+    documents: Array<{
+      id: string;
+      tokens: string[];
+      type: MemoryItemType;
+      content: string;
+    }>;
+    k1: number;
+    b: number;
+  } {
+    return {
+      documents: Array.from(this.documents.values()),
+      k1: this.k1,
+      b: this.b,
+    };
+  }
+
+  /**
+   * Restore index from JSON
+   */
+  static fromJSON(data: {
+    documents: Array<{
+      id: string;
+      tokens: string[];
+      type: MemoryItemType;
+      content: string;
+    }>;
+    k1: number;
+    b: number;
+  }): BM25Index {
+    const index = new BM25Index({ k1: data.k1, b: data.b });
+    for (const doc of data.documents) {
+      index.documents.set(doc.id, doc);
+    }
+    index.dirty = true;
+    return index;
+  }
+}
+
+/**
+ * Common English stop words to filter out
+ */
+const STOP_WORDS = new Set([
+  'a',
+  'an',
+  'and',
+  'are',
+  'as',
+  'at',
+  'be',
+  'by',
+  'for',
+  'from',
+  'has',
+  'he',
+  'in',
+  'is',
+  'it',
+  'its',
+  'of',
+  'on',
+  'or',
+  'that',
+  'the',
+  'to',
+  'was',
+  'were',
+  'will',
+  'with',
+  'this',
+  'but',
+  'they',
+  'have',
+  'had',
+  'what',
+  'when',
+  'where',
+  'who',
+  'which',
+  'why',
+  'how',
+  'all',
+  'each',
+  'every',
+  'both',
+  'few',
+  'more',
+  'most',
+  'other',
+  'some',
+  'such',
+  'no',
+  'nor',
+  'not',
+  'only',
+  'own',
+  'same',
+  'so',
+  'than',
+  'too',
+  'very',
+  'can',
+  'just',
+  'should',
+  'now',
+]);
+
+/**
+ * Create a new BM25 index
+ */
+export function createBM25Index(options?: {
+  k1?: number;
+  b?: number;
+}): BM25Index {
+  return new BM25Index(options);
+}

package/src/embeddings/index.ts

@@ -0,0 +1,39 @@
+/**
+ * Embeddings Module
+ *
+ * Provides embedding generation, storage, and similarity search capabilities.
+ * Supports multiple embedding providers (OpenAI, Voyage, local) with automatic
+ * fallback to BM25 text search when no provider is configured.
+ */
+
+// Provider interface and implementations
+export {
+  type EmbeddingProvider,
+  type EmbeddingProviderConfig,
+  OpenAIEmbeddingProvider,
+  VoyageEmbeddingProvider,
+  LocalEmbeddingProvider,
+  createEmbeddingProvider,
+} from './provider.js';
+
+// BM25 text search (fallback)
+export {
+  BM25Index,
+  createBM25Index,
+  type MemoryItemType,
+} from './bm25.js';
+
+// Vector store using sqlite-vec
+export {
+  SqliteVectorStore,
+  createVectorStore,
+  type VectorSearchResult,
+} from './vector-store.js';
+
+// Embedding manager (main interface)
+export {
+  EmbeddingManager,
+  createEmbeddingManager,
+  type EmbeddingManagerOptions,
+  type SearchResult,
+} from './manager.js';

package/src/embeddings/manager.ts

@@ -0,0 +1,288 @@
+/**
+ * Embedding Manager
+ *
+ * Orchestrates embedding generation, storage, and search.
+ * Provides automatic fallback to BM25 when no embedding provider is configured.
+ */
+
+import type { EmbeddingProvider } from './provider.js';
+import { SqliteVectorStore } from './vector-store.js';
+import { BM25Index, type MemoryItemType } from './bm25.js';
+
+export { type MemoryItemType } from './bm25.js';
+
+export interface SearchResult {
+  id: string;
+  score: number;
+  content: string;
+}
+
+export interface EmbeddingManagerOptions {
+  /** Embedding provider (null for BM25-only mode) */
+  provider?: EmbeddingProvider | null;
+  /** Path to sqlite-vec database file */
+  vectorStorePath?: string;
+  /** Embedding dimension (required if provider is set) */
+  dimension?: number;
+  /** BM25 parameters */
+  bm25?: {
+    k1?: number;
+    b?: number;
+  };
+}
+
+/**
+ * Manages embeddings with automatic fallback to BM25
+ */
+export class EmbeddingManager {
+  private provider: EmbeddingProvider | null;
+  private vectorStore: SqliteVectorStore | null;
+  private bm25Index: BM25Index;
+  private useVectorSearch: boolean;
+
+  constructor(options: EmbeddingManagerOptions) {
+    this.provider = options.provider ?? null;
+    this.useVectorSearch = this.provider !== null && !!options.vectorStorePath;
+
+    // Initialize vector store if we have a provider and path
+    if (this.useVectorSearch && options.vectorStorePath) {
+      const dimension = options.dimension ?? this.provider!.dimension;
+      this.vectorStore = new SqliteVectorStore({
+        dbPath: options.vectorStorePath,
+        dimension,
+      });
+    } else {
+      this.vectorStore = null;
+    }
+
+    // Always maintain BM25 as fallback
+    this.bm25Index = new BM25Index(options.bm25);
+  }
+
+  /**
+   * Check if vector search is available
+   */
+  get hasVectorSearch(): boolean {
+    return this.useVectorSearch && this.vectorStore !== null;
+  }
+
+  /**
+   * Get the embedding dimension (if available)
+   */
+  get dimension(): number | null {
+    return this.provider?.dimension ?? null;
+  }
+
+  /**
+   * Store an item with its embedding
+   * Returns the embedding if generated, undefined otherwise
+   */
+  async store(
+    id: string,
+    content: string,
+    type: MemoryItemType
+  ): Promise<number[] | undefined> {
+    // Always add to BM25 index for fallback
+    this.bm25Index.add(id, content, type);
+
+    // Generate and store embedding if provider available
+    if (this.useVectorSearch && this.provider && this.vectorStore) {
+      try {
+        const embedding = await this.provider.embed(content);
+        this.vectorStore.insert(id, embedding, type, content);
+        return embedding;
+      } catch (error) {
+        console.error(`Failed to generate embedding for ${id}:`, error);
+        // Fall through to return undefined - BM25 is still available
+      }
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Store multiple items with embeddings (batch)
+   * More efficient than calling store() multiple times
+   */
+  async storeBatch(
+    items: Array<{ id: string; content: string; type: MemoryItemType }>
+  ): Promise<Map<string, number[] | undefined>> {
+    const results = new Map<string, number[] | undefined>();
+
+    // Add all to BM25 index
+    for (const item of items) {
+      this.bm25Index.add(item.id, item.content, item.type);
+    }
+
+    // Batch embed if provider available
+    if (this.useVectorSearch && this.provider && this.vectorStore) {
+      try {
+        const contents = items.map((i) => i.content);
+        const embeddings = await this.provider.embedBatch(contents);
+
+        const batchItems = items.map((item, i) => ({
+          id: item.id,
+          embedding: embeddings[i],
+          type: item.type,
+          content: item.content,
+        }));
+
+        this.vectorStore.insertBatch(batchItems);
+
+        for (let i = 0; i < items.length; i++) {
+          results.set(items[i].id, embeddings[i]);
+        }
+      } catch (error) {
+        console.error('Failed to generate batch embeddings:', error);
+        // Fall through - items are still in BM25 index
+        for (const item of items) {
+          results.set(item.id, undefined);
+        }
+      }
+    } else {
+      for (const item of items) {
+        results.set(item.id, undefined);
+      }
+    }
+
+    return results;
+  }
+
+  /**
+   * Search for similar items
+   * Uses vector search if available, falls back to BM25
+   */
+  async search(
+    query: string,
+    options: {
+      k?: number;
+      type?: MemoryItemType;
+      threshold?: number;
+    } = {}
+  ): Promise<SearchResult[]> {
+    // Use vector search if available
+    if (this.useVectorSearch && this.provider && this.vectorStore) {
+      try {
+        const queryEmbedding = await this.provider.embed(query);
+        const results = this.vectorStore.search(queryEmbedding, options);
+        return results.map((r) => ({
+          id: r.id,
+          score: r.score,
+          content: r.content,
+        }));
+      } catch (error) {
+        console.error('Vector search failed, falling back to BM25:', error);
+        // Fall through to BM25
+      }
+    }
+
+    // Fallback to BM25
+    return this.bm25Index.search(query, options);
+  }
+
+  /**
+   * Search using a pre-computed embedding
+   * Only works if vector store is available
+   */
+  searchByEmbedding(
+    embedding: number[],
+    options: {
+      k?: number;
+      type?: MemoryItemType;
+      threshold?: number;
+    } = {}
+  ): SearchResult[] | null {
+    if (!this.vectorStore) {
+      return null;
+    }
+
+    const results = this.vectorStore.search(embedding, options);
+    return results.map((r) => ({
+      id: r.id,
+      score: r.score,
+      content: r.content,
+    }));
+  }
+
+  /**
+   * Generate an embedding without storing it
+   */
+  async embed(text: string): Promise<number[] | null> {
+    if (!this.provider) {
+      return null;
+    }
+
+    try {
+      return await this.provider.embed(text);
+    } catch (error) {
+      console.error('Failed to generate embedding:', error);
+      return null;
+    }
+  }
+
+  /**
+   * Generate embeddings for multiple texts without storing
+   */
+  async embedBatch(texts: string[]): Promise<number[][] | null> {
+    if (!this.provider) {
+      return null;
+    }
+
+    try {
+      return await this.provider.embedBatch(texts);
+    } catch (error) {
+      console.error('Failed to generate batch embeddings:', error);
+      return null;
+    }
+  }
+
+  /**
+   * Remove an item from all indexes
+   */
+  remove(id: string): void {
+    this.bm25Index.remove(id);
+    this.vectorStore?.delete(id);
+  }
+
+  /**
+   * Check if an item exists
+   */
+  has(id: string): boolean {
+    return this.bm25Index.has(id);
+  }
+
+  /**
+   * Get statistics about the indexes
+   */
+  getStats(): {
+    bm25Count: number;
+    vectorCount: number | null;
+    hasVectorSearch: boolean;
+    dimension: number | null;
+    providerName: string | null;
+  } {
+    return {
+      bm25Count: this.bm25Index.size,
+      vectorCount: this.vectorStore?.count ?? null,
+      hasVectorSearch: this.hasVectorSearch,
+      dimension: this.dimension,
+      providerName: this.provider?.name ?? null,
+    };
+  }
+
+  /**
+   * Close all resources
+   */
+  close(): void {
+    this.vectorStore?.close();
+  }
+}
+
+/**
+ * Create an embedding manager
+ */
+export function createEmbeddingManager(
+  options: EmbeddingManagerOptions
+): EmbeddingManager {
+  return new EmbeddingManager(options);
+}