cognitive-core 0.0.2 → 0.1.0

package/src/embeddings/provider.ts

@@ -0,0 +1,311 @@
+/**
+ * Embedding Provider Interface and Implementations
+ *
+ * Provides a unified interface for generating embeddings from text.
+ * Supports hosted providers (OpenAI, etc.) and local models.
+ */
+
+/**
+ * Interface for embedding providers (hosted or local)
+ */
+export interface EmbeddingProvider {
+  /** Provider name for logging/debugging */
+  readonly name: string;
+
+  /** Embedding dimension */
+  readonly dimension: number;
+
+  /** Embed a single text */
+  embed(text: string): Promise<number[]>;
+
+  /** Embed multiple texts (batch) */
+  embedBatch(texts: string[]): Promise<number[][]>;
+}
+
+/**
+ * OpenAI embedding provider
+ */
+export class OpenAIEmbeddingProvider implements EmbeddingProvider {
+  readonly name = "openai";
+  readonly dimension: number;
+  private apiKey: string;
+  private model: string;
+  private baseUrl: string;
+
+  constructor(options: {
+    apiKey: string;
+    model?: string;
+    dimension?: number;
+    baseUrl?: string;
+  }) {
+    this.apiKey = options.apiKey;
+    this.model = options.model ?? "text-embedding-3-small";
+    this.dimension = options.dimension ?? 1536;
+    this.baseUrl = options.baseUrl ?? "https://api.openai.com/v1";
+  }
+
+  async embed(text: string): Promise<number[]> {
+    const [embedding] = await this.embedBatch([text]);
+    return embedding;
+  }
+
+  async embedBatch(texts: string[]): Promise<number[][]> {
+    const response = await fetch(`${this.baseUrl}/embeddings`, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model: this.model,
+        input: texts,
+        dimensions: this.dimension,
+      }),
+    });
+
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(`OpenAI embedding error: ${response.status} - ${error}`);
+    }
+
+    const data = (await response.json()) as {
+      data: Array<{ embedding: number[]; index: number }>;
+    };
+
+    // Sort by index to maintain order
+    const sorted = data.data.sort((a, b) => a.index - b.index);
+    return sorted.map((item) => item.embedding);
+  }
+}
+
+/**
+ * Anthropic Voyage embedding provider
+ * Uses Voyage AI embeddings via Anthropic partnership
+ */
+export class VoyageEmbeddingProvider implements EmbeddingProvider {
+  readonly name = "voyage";
+  readonly dimension: number;
+  private apiKey: string;
+  private model: string;
+
+  constructor(options: { apiKey: string; model?: string; dimension?: number }) {
+    this.apiKey = options.apiKey;
+    this.model = options.model ?? "voyage-2";
+    this.dimension = options.dimension ?? 1024;
+  }
+
+  async embed(text: string): Promise<number[]> {
+    const [embedding] = await this.embedBatch([text]);
+    return embedding;
+  }
+
+  async embedBatch(texts: string[]): Promise<number[][]> {
+    const response = await fetch("https://api.voyageai.com/v1/embeddings", {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model: this.model,
+        input: texts,
+      }),
+    });
+
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(`Voyage embedding error: ${response.status} - ${error}`);
+    }
+
+    const data = (await response.json()) as {
+      data: Array<{ embedding: number[]; index: number }>;
+    };
+
+    const sorted = data.data.sort((a, b) => a.index - b.index);
+    return sorted.map((item) => item.embedding);
+  }
+}
+
+/**
+ * Local embedding provider using @huggingface/transformers
+ * Runs transformer models locally in Node.js using ONNX runtime
+ */
+export class LocalEmbeddingProvider implements EmbeddingProvider {
+  readonly name = "local";
+  readonly dimension: number;
+  private modelName: string;
+  private extractor: any = null;
+  private initPromise: Promise<void> | null = null;
+
+  /** Default model - all-MiniLM-L6-v2 is small and fast with good quality */
+  static readonly DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2";
+  /** Dimensions for common models */
+  static readonly MODEL_DIMENSIONS: Record<string, number> = {
+    "sentence-transformers/all-MiniLM-L6-v2": 384,
+    "sentence-transformers/all-MiniLM-L12-v2": 384,
+    "BAAI/bge-small-en-v1.5": 384,
+    "BAAI/bge-base-en-v1.5": 768,
+    "intfloat/e5-small-v2": 384,
+    "intfloat/e5-base-v2": 768,
+    // Keep Xenova aliases for backwards compatibility
+    "Xenova/all-MiniLM-L6-v2": 384,
+    "Xenova/all-MiniLM-L12-v2": 384,
+    "Xenova/bge-small-en-v1.5": 384,
+    "Xenova/bge-base-en-v1.5": 768,
+    "Xenova/e5-small-v2": 384,
+    "Xenova/e5-base-v2": 768,
+  };
+
+  constructor(options?: {
+    modelName?: string;
+    modelPath?: string; // Alias for modelName for backwards compatibility
+    dimension?: number;
+  }) {
+    this.modelName =
+      options?.modelName ??
+      options?.modelPath ??
+      LocalEmbeddingProvider.DEFAULT_MODEL;
+
+    // Auto-detect dimension from known models or use provided/default
+    this.dimension =
+      options?.dimension ??
+      LocalEmbeddingProvider.MODEL_DIMENSIONS[this.modelName] ??
+      384;
+  }
+
+  /** Get the model name/path */
+  get modelPath(): string {
+    return this.modelName;
+  }
+
+  private async initialize(): Promise<void> {
+    if (this.extractor) return;
+
+    // Prevent multiple concurrent initializations
+    if (this.initPromise) {
+      await this.initPromise;
+      return;
+    }
+
+    this.initPromise = this.doInitialize();
+    await this.initPromise;
+  }
+
+  private async doInitialize(): Promise<void> {
+    try {
+      // Dynamic import to avoid issues if transformers.js isn't installed
+      const { pipeline } = await import("@huggingface/transformers");
+
+      // Create feature-extraction pipeline for embeddings
+      this.extractor = await pipeline("feature-extraction", this.modelName, {
+        dtype: "fp32", // Use fp32 for better compatibility
+      });
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      throw new Error(
+        `Failed to initialize local embedding model "${this.modelName}": ${message}`,
+      );
+    }
+  }
+
+  async embed(text: string): Promise<number[]> {
+    await this.initialize();
+
+    // Run the extractor
+    const output = await this.extractor(text, {
+      pooling: "mean", // Mean pooling over tokens
+      normalize: true, // L2 normalize the output
+    });
+
+    // Convert tensor to array
+    return Array.from(output.data as Float32Array);
+  }
+
+  async embedBatch(texts: string[]): Promise<number[][]> {
+    await this.initialize();
+
+    // Process in batch for better efficiency
+    const results: number[][] = [];
+
+    for (const text of texts) {
+      const output = await this.extractor(text, {
+        pooling: "mean",
+        normalize: true,
+      });
+      results.push(Array.from(output.data as Float32Array));
+    }
+
+    return results;
+  }
+
+  /**
+   * Check if a model is available locally (cached)
+   * This can be used to pre-warm the cache
+   */
+  async isModelCached(): Promise<boolean> {
+    try {
+      const { env } = await import("@huggingface/transformers");
+      // Check if model files exist in cache
+      // This is a simplified check - the actual cache structure may vary
+      return env.cacheDir !== undefined;
+    } catch {
+      return false;
+    }
+  }
+}
+
+/**
+ * Configuration for embedding providers
+ */
+export interface EmbeddingProviderConfig {
+  type: "openai" | "voyage" | "local" | "none";
+  apiKey?: string;
+  model?: string;
+  modelPath?: string;
+  dimension?: number;
+  baseUrl?: string;
+}
+
+/**
+ * Create an embedding provider from configuration
+ */
+export function createEmbeddingProvider(
+  config: EmbeddingProviderConfig,
+): EmbeddingProvider | null {
+  switch (config.type) {
+    case "openai":
+      if (!config.apiKey) {
+        throw new Error("OpenAI embedding provider requires apiKey");
+      }
+      return new OpenAIEmbeddingProvider({
+        apiKey: config.apiKey,
+        model: config.model,
+        dimension: config.dimension,
+        baseUrl: config.baseUrl,
+      });
+
+    case "voyage":
+      if (!config.apiKey) {
+        throw new Error("Voyage embedding provider requires apiKey");
+      }
+      return new VoyageEmbeddingProvider({
+        apiKey: config.apiKey,
+        model: config.model,
+        dimension: config.dimension,
+      });
+
+    case "local":
+      // Local provider now has sensible defaults
+      return new LocalEmbeddingProvider({
+        modelName: config.model,
+        modelPath: config.modelPath,
+        dimension: config.dimension,
+      });
+
+    case "none":
+      return null;
+
+    default:
+      return null;
+  }
+}

package/src/embeddings/vector-store.ts

@@ -0,0 +1,353 @@
+/**
+ * Vector Store using sqlite-vec
+ *
+ * Provides efficient vector similarity search using SQLite with the sqlite-vec extension.
+ * sqlite-vec is optimized for in-process vector search without the overhead of external services.
+ */
+
+import Database from 'better-sqlite3';
+import * as sqliteVec from 'sqlite-vec';
+import type { MemoryItemType } from './bm25.js';
+
+export interface VectorSearchResult {
+  id: string;
+  score: number;
+  content: string;
+  type: MemoryItemType;
+}
+
+/**
+ * Vector store using sqlite-vec for efficient similarity search
+ */
+export class SqliteVectorStore {
+  private db: Database.Database;
+  private dimension: number;
+  private tableName: string;
+  private initialized: boolean = false;
+
+  constructor(options: {
+    dbPath: string;
+    dimension: number;
+    tableName?: string;
+  }) {
+    this.dimension = options.dimension;
+    this.tableName = options.tableName ?? 'atlas_embeddings';
+
+    // Create or open the database
+    this.db = new Database(options.dbPath);
+
+    // Load the sqlite-vec extension
+    sqliteVec.load(this.db);
+
+    this.initialize();
+  }
+
+  /**
+   * Initialize the database schema
+   */
+  private initialize(): void {
+    if (this.initialized) return;
+
+    // Create metadata table for storing content and type info
+    this.db.exec(`
+      CREATE TABLE IF NOT EXISTS ${this.tableName}_meta (
+        id TEXT PRIMARY KEY,
+        type TEXT NOT NULL,
+        content TEXT NOT NULL,
+        created_at TEXT DEFAULT CURRENT_TIMESTAMP
+      )
+    `);
+
+    // Create virtual table for vector search using vec0
+    // Note: vec0 uses a separate storage model optimized for vector operations
+    this.db.exec(`
+      CREATE VIRTUAL TABLE IF NOT EXISTS ${this.tableName}_vec USING vec0(
+        id TEXT PRIMARY KEY,
+        embedding FLOAT[${this.dimension}]
+      )
+    `);
+
+    // Create index on type for filtering
+    this.db.exec(`
+      CREATE INDEX IF NOT EXISTS idx_${this.tableName}_type
+      ON ${this.tableName}_meta(type)
+    `);
+
+    this.initialized = true;
+  }
+
+  /**
+   * Insert or update a vector with metadata
+   */
+  insert(
+    id: string,
+    embedding: number[],
+    type: MemoryItemType,
+    content: string
+  ): void {
+    if (embedding.length !== this.dimension) {
+      throw new Error(
+        `Embedding dimension mismatch: expected ${this.dimension}, got ${embedding.length}`
+      );
+    }
+
+    // Use a transaction for atomicity
+    const transaction = this.db.transaction(() => {
+      // Insert or replace metadata
+      const metaStmt = this.db.prepare(`
+        INSERT OR REPLACE INTO ${this.tableName}_meta (id, type, content)
+        VALUES (?, ?, ?)
+      `);
+      metaStmt.run(id, type, content);
+
+      // Insert or replace vector
+      // sqlite-vec expects vectors as JSON arrays or binary blobs
+      const vecStmt = this.db.prepare(`
+        INSERT OR REPLACE INTO ${this.tableName}_vec (id, embedding)
+        VALUES (?, ?)
+      `);
+      vecStmt.run(id, JSON.stringify(embedding));
+    });
+
+    transaction();
+  }
+
+  /**
+   * Insert multiple vectors in batch
+   */
+  insertBatch(
+    items: Array<{
+      id: string;
+      embedding: number[];
+      type: MemoryItemType;
+      content: string;
+    }>
+  ): void {
+    const transaction = this.db.transaction(() => {
+      const metaStmt = this.db.prepare(`
+        INSERT OR REPLACE INTO ${this.tableName}_meta (id, type, content)
+        VALUES (?, ?, ?)
+      `);
+
+      const vecStmt = this.db.prepare(`
+        INSERT OR REPLACE INTO ${this.tableName}_vec (id, embedding)
+        VALUES (?, ?)
+      `);
+
+      for (const item of items) {
+        if (item.embedding.length !== this.dimension) {
+          throw new Error(
+            `Embedding dimension mismatch for ${item.id}: expected ${this.dimension}, got ${item.embedding.length}`
+          );
+        }
+        metaStmt.run(item.id, item.type, item.content);
+        vecStmt.run(item.id, JSON.stringify(item.embedding));
+      }
+    });
+
+    transaction();
+  }
+
+  /**
+   * Search for similar vectors
+   */
+  search(
+    queryEmbedding: number[],
+    options: {
+      k?: number;
+      type?: MemoryItemType;
+      threshold?: number;
+    } = {}
+  ): VectorSearchResult[] {
+    if (queryEmbedding.length !== this.dimension) {
+      throw new Error(
+        `Query embedding dimension mismatch: expected ${this.dimension}, got ${queryEmbedding.length}`
+      );
+    }
+
+    const k = options.k ?? 10;
+    const threshold = options.threshold ?? 0;
+
+    // Build the query based on whether we're filtering by type
+    let query: string;
+    let params: unknown[];
+
+    if (options.type) {
+      // Join with metadata table to filter by type
+      query = `
+        SELECT
+          v.id,
+          m.type,
+          m.content,
+          vec_distance_cosine(v.embedding, ?) as distance
+        FROM ${this.tableName}_vec v
+        JOIN ${this.tableName}_meta m ON v.id = m.id
+        WHERE m.type = ?
+        ORDER BY distance ASC
+        LIMIT ?
+      `;
+      params = [JSON.stringify(queryEmbedding), options.type, k];
+    } else {
+      query = `
+        SELECT
+          v.id,
+          m.type,
+          m.content,
+          vec_distance_cosine(v.embedding, ?) as distance
+        FROM ${this.tableName}_vec v
+        JOIN ${this.tableName}_meta m ON v.id = m.id
+        ORDER BY distance ASC
+        LIMIT ?
+      `;
+      params = [JSON.stringify(queryEmbedding), k];
+    }
+
+    const stmt = this.db.prepare(query);
+    const results = stmt.all(...params) as Array<{
+      id: string;
+      type: MemoryItemType;
+      content: string;
+      distance: number;
+    }>;
+
+    // Convert distance to similarity score (1 - distance for cosine)
+    return results
+      .map((r) => ({
+        id: r.id,
+        type: r.type,
+        content: r.content,
+        score: 1 - r.distance,
+      }))
+      .filter((r) => r.score >= threshold);
+  }
+
+  /**
+   * Get a vector by ID
+   */
+  get(id: string): { embedding: number[]; type: MemoryItemType; content: string } | null {
+    const metaStmt = this.db.prepare(`
+      SELECT type, content FROM ${this.tableName}_meta WHERE id = ?
+    `);
+    const meta = metaStmt.get(id) as { type: MemoryItemType; content: string } | undefined;
+
+    if (!meta) return null;
+
+    const vecStmt = this.db.prepare(`
+      SELECT embedding FROM ${this.tableName}_vec WHERE id = ?
+    `);
+    const vec = vecStmt.get(id) as { embedding: string } | undefined;
+
+    if (!vec) return null;
+
+    return {
+      embedding: JSON.parse(vec.embedding) as number[],
+      type: meta.type,
+      content: meta.content,
+    };
+  }
+
+  /**
+   * Check if an ID exists
+   */
+  has(id: string): boolean {
+    const stmt = this.db.prepare(`
+      SELECT 1 FROM ${this.tableName}_meta WHERE id = ? LIMIT 1
+    `);
+    return stmt.get(id) !== undefined;
+  }
+
+  /**
+   * Delete a vector
+   */
+  delete(id: string): boolean {
+    const transaction = this.db.transaction(() => {
+      const metaStmt = this.db.prepare(
+        `DELETE FROM ${this.tableName}_meta WHERE id = ?`
+      );
+      const vecStmt = this.db.prepare(
+        `DELETE FROM ${this.tableName}_vec WHERE id = ?`
+      );
+
+      const metaResult = metaStmt.run(id);
+      vecStmt.run(id);
+
+      return metaResult.changes > 0;
+    });
+
+    return transaction();
+  }
+
+  /**
+   * Get the number of vectors stored
+   */
+  get count(): number {
+    const stmt = this.db.prepare(
+      `SELECT COUNT(*) as count FROM ${this.tableName}_meta`
+    );
+    const result = stmt.get() as { count: number };
+    return result.count;
+  }
+
+  /**
+   * Get count by type
+   */
+  countByType(type: MemoryItemType): number {
+    const stmt = this.db.prepare(
+      `SELECT COUNT(*) as count FROM ${this.tableName}_meta WHERE type = ?`
+    );
+    const result = stmt.get(type) as { count: number };
+    return result.count;
+  }
+
+  /**
+   * Clear all vectors
+   */
+  clear(): void {
+    const transaction = this.db.transaction(() => {
+      this.db.exec(`DELETE FROM ${this.tableName}_meta`);
+      this.db.exec(`DELETE FROM ${this.tableName}_vec`);
+    });
+
+    transaction();
+  }
+
+  /**
+   * Close the database connection
+   */
+  close(): void {
+    this.db.close();
+  }
+
+  /**
+   * Get database statistics
+   */
+  getStats(): {
+    totalVectors: number;
+    byType: Record<MemoryItemType, number>;
+    dimension: number;
+  } {
+    const types: MemoryItemType[] = ['experience', 'strategy', 'concept', 'skill'];
+    const byType: Record<MemoryItemType, number> = {} as Record<MemoryItemType, number>;
+
+    for (const type of types) {
+      byType[type] = this.countByType(type);
+    }
+
+    return {
+      totalVectors: this.count,
+      byType,
+      dimension: this.dimension,
+    };
+  }
+}
+
+/**
+ * Create a new vector store
+ */
+export function createVectorStore(options: {
+  dbPath: string;
+  dimension: number;
+  tableName?: string;
+}): SqliteVectorStore {
+  return new SqliteVectorStore(options);
+}