kiro-memory 1.9.0 → 3.0.0

package/plugin/dist/services/search/VectorSearch.js

@@ -220,14 +220,48 @@ ${data.stack}` : ` ${data.message}`;
 var logger = new Logger();
 
 // src/services/search/EmbeddingService.ts
+var MODEL_CONFIGS = {
+  "all-MiniLM-L6-v2": {
+    modelId: "Xenova/all-MiniLM-L6-v2",
+    dimensions: 384
+  },
+  "jina-code-v2": {
+    modelId: "jinaai/jina-embeddings-v2-base-code",
+    dimensions: 768
+  },
+  "bge-small-en": {
+    modelId: "BAAI/bge-small-en-v1.5",
+    dimensions: 384
+  }
+};
+var FASTEMBED_COMPATIBLE_MODELS = /* @__PURE__ */ new Set(["all-MiniLM-L6-v2", "bge-small-en"]);
 var EmbeddingService = class {
   provider = null;
   model = null;
   initialized = false;
   initializing = null;
+  config;
+  configName;
+  constructor() {
+    const envModel = process.env.KIRO_MEMORY_EMBEDDING_MODEL || "all-MiniLM-L6-v2";
+    this.configName = envModel;
+    if (MODEL_CONFIGS[envModel]) {
+      this.config = MODEL_CONFIGS[envModel];
+    } else if (envModel.includes("/")) {
+      const dimensions = parseInt(process.env.KIRO_MEMORY_EMBEDDING_DIMENSIONS || "384", 10);
+      this.config = {
+        modelId: envModel,
+        dimensions: isNaN(dimensions) ? 384 : dimensions
+      };
+    } else {
+      logger.warn("EMBEDDING", `Unknown model name '${envModel}', falling back to 'all-MiniLM-L6-v2'`);
+      this.configName = "all-MiniLM-L6-v2";
+      this.config = MODEL_CONFIGS["all-MiniLM-L6-v2"];
+    }
+  }
   /**
-   * Inizializza il servizio di embedding.
-   * Tenta fastembed, poi @huggingface/transformers, poi fallback a null.
+   * Initialize the embedding service.
+   * Tries fastembed (when compatible), then @huggingface/transformers, then falls back to null.
    */
   async initialize() {
     if (this.initialized) return this.provider !== null;
@@ -238,45 +272,48 @@ var EmbeddingService = class {
     return result;
   }
   async _doInitialize() {
-    try {
-      const fastembed = await import("fastembed");
-      const EmbeddingModel = fastembed.EmbeddingModel || fastembed.default?.EmbeddingModel;
-      const FlagEmbedding = fastembed.FlagEmbedding || fastembed.default?.FlagEmbedding;
-      if (FlagEmbedding && EmbeddingModel) {
-        this.model = await FlagEmbedding.init({
-          model: EmbeddingModel.BGESmallENV15
-        });
-        this.provider = "fastembed";
-        this.initialized = true;
-        logger.info("EMBEDDING", "Inizializzato con fastembed (BGE-small-en-v1.5)");
-        return true;
+    const fastembedCompatible = FASTEMBED_COMPATIBLE_MODELS.has(this.configName);
+    if (fastembedCompatible) {
+      try {
+        const fastembed = await import("fastembed");
+        const EmbeddingModel = fastembed.EmbeddingModel || fastembed.default?.EmbeddingModel;
+        const FlagEmbedding = fastembed.FlagEmbedding || fastembed.default?.FlagEmbedding;
+        if (FlagEmbedding && EmbeddingModel) {
+          this.model = await FlagEmbedding.init({
+            model: EmbeddingModel.BGESmallENV15
+          });
+          this.provider = "fastembed";
+          this.initialized = true;
+          logger.info("EMBEDDING", `Initialized with fastembed (BGE-small-en-v1.5) for model '${this.configName}'`);
+          return true;
+        }
+      } catch (error) {
+        logger.debug("EMBEDDING", `fastembed not available: ${error}`);
       }
-    } catch (error) {
-      logger.debug("EMBEDDING", `fastembed non disponibile: ${error}`);
     }
     try {
       const transformers = await import("@huggingface/transformers");
       const pipeline = transformers.pipeline || transformers.default?.pipeline;
       if (pipeline) {
-        this.model = await pipeline("feature-extraction", "Xenova/all-MiniLM-L6-v2", {
+        this.model = await pipeline("feature-extraction", this.config.modelId, {
           quantized: true
         });
         this.provider = "transformers";
         this.initialized = true;
-        logger.info("EMBEDDING", "Inizializzato con @huggingface/transformers (all-MiniLM-L6-v2)");
+        logger.info("EMBEDDING", `Initialized with @huggingface/transformers (${this.config.modelId})`);
         return true;
       }
     } catch (error) {
-      logger.debug("EMBEDDING", `@huggingface/transformers non disponibile: ${error}`);
+      logger.debug("EMBEDDING", `@huggingface/transformers not available: ${error}`);
     }
     this.provider = null;
     this.initialized = true;
-    logger.warn("EMBEDDING", "Nessun provider embedding disponibile, ricerca semantica disabilitata");
+    logger.warn("EMBEDDING", "No embedding provider available, semantic search disabled");
     return false;
   }
   /**
-   * Genera embedding per un singolo testo.
-   * Ritorna Float32Array con 384 dimensioni, o null se non disponibile.
+   * Generate embedding for a single text.
+   * Returns Float32Array with configured dimensions, or null if not available.
    */
   async embed(text) {
     if (!this.initialized) await this.initialize();
@@ -289,46 +326,118 @@ var EmbeddingService = class {
         return await this._embedTransformers(truncated);
       }
     } catch (error) {
-      logger.error("EMBEDDING", `Errore generazione embedding: ${error}`);
+      logger.error("EMBEDDING", `Error generating embedding: ${error}`);
     }
     return null;
   }
   /**
-   * Genera embeddings in batch.
+   * Generate embeddings in batch.
+   * Uses native batch support when available (fastembed, transformers),
+   * falls back to serial processing on batch failure.
    */
   async embedBatch(texts) {
     if (!this.initialized) await this.initialize();
     if (!this.provider || !this.model) return texts.map(() => null);
-    const results = [];
-    for (const text of texts) {
-      try {
-        const embedding = await this.embed(text);
-        results.push(embedding);
-      } catch {
-        results.push(null);
+    if (texts.length === 0) return [];
+    const truncated = texts.map((t) => t.substring(0, 2e3));
+    try {
+      if (this.provider === "fastembed") {
+        return await this._embedBatchFastembed(truncated);
+      } else if (this.provider === "transformers") {
+        return await this._embedBatchTransformers(truncated);
       }
+    } catch (error) {
+      logger.warn("EMBEDDING", `Batch embedding failed, falling back to serial: ${error}`);
     }
-    return results;
+    return this._embedBatchSerial(truncated);
   }
   /**
-   * Verifica se il servizio è disponibile.
+   * Check if the service is available.
    */
   isAvailable() {
     return this.initialized && this.provider !== null;
   }
   /**
-   * Nome del provider attivo.
+   * Name of the active provider.
    */
   getProvider() {
     return this.provider;
   }
   /**
-   * Dimensioni del vettore embedding.
+   * Embedding vector dimensions for the active model configuration.
    */
   getDimensions() {
-    return 384;
+    return this.config.dimensions;
   }
-  // --- Provider specifici ---
+  /**
+   * Human-readable model name used as identifier in the observation_embeddings table.
+   * Returns the short name (e.g., 'all-MiniLM-L6-v2') or the full HF model ID for custom models.
+   */
+  getModelName() {
+    return this.configName;
+  }
+  // --- Batch implementations ---
+  /**
+   * Native batch embedding with fastembed.
+   * FlagEmbedding.embed() accepts string[] and returns an async iterable of batches.
+   */
+  async _embedBatchFastembed(texts) {
+    const results = [];
+    const embeddings = this.model.embed(texts, texts.length);
+    for await (const batch of embeddings) {
+      if (batch) {
+        for (const vec of batch) {
+          results.push(vec instanceof Float32Array ? vec : new Float32Array(vec));
+        }
+      }
+    }
+    while (results.length < texts.length) {
+      results.push(null);
+    }
+    return results;
+  }
+  /**
+   * Batch embedding with @huggingface/transformers pipeline.
+   * The pipeline accepts string[] and returns a Tensor with shape [N, dims].
+   */
+  async _embedBatchTransformers(texts) {
+    const output = await this.model(texts, {
+      pooling: "mean",
+      normalize: true
+    });
+    if (!output?.data) {
+      return texts.map(() => null);
+    }
+    const dims = this.getDimensions();
+    const data = output.data instanceof Float32Array ? output.data : new Float32Array(output.data);
+    const results = [];
+    for (let i = 0; i < texts.length; i++) {
+      const offset = i * dims;
+      if (offset + dims <= data.length) {
+        results.push(data.slice(offset, offset + dims));
+      } else {
+        results.push(null);
+      }
+    }
+    return results;
+  }
+  /**
+   * Serial fallback: embed texts one at a time.
+   * Used when native batch fails.
+   */
+  async _embedBatchSerial(texts) {
+    const results = [];
+    for (const text of texts) {
+      try {
+        const embedding = await this.embed(text);
+        results.push(embedding);
+      } catch {
+        results.push(null);
+      }
+    }
+    return results;
+  }
+  // --- Single-text provider implementations ---
   async _embedFastembed(text) {
     const embeddings = this.model.embed([text], 1);
     for await (const batch of embeddings) {
@@ -359,17 +468,21 @@ function getEmbeddingService() {
 }
 
 // src/services/search/VectorSearch.ts
+var DEFAULT_MAX_CANDIDATES = 2e3;
 function cosineSimilarity(a, b) {
-  if (a.length !== b.length) return 0;
+  const len = a.length;
+  if (len !== b.length) return 0;
   let dotProduct = 0;
   let normA = 0;
   let normB = 0;
-  for (let i = 0; i < a.length; i++) {
-    dotProduct += a[i] * b[i];
-    normA += a[i] * a[i];
-    normB += b[i] * b[i];
-  }
-  const denominator = Math.sqrt(normA) * Math.sqrt(normB);
+  for (let i = 0; i < len; i++) {
+    const ai = a[i];
+    const bi = b[i];
+    dotProduct += ai * bi;
+    normA += ai * ai;
+    normB += bi * bi;
+  }
+  const denominator = Math.sqrt(normA * normB);
   if (denominator === 0) return 0;
   return dotProduct / denominator;
 }
@@ -382,23 +495,36 @@ function bufferToFloat32(buf) {
 }
 var VectorSearch = class {
   /**
-   * Ricerca semantica: calcola cosine similarity tra query e tutti gli embeddings.
+   * Semantic search with SQL pre-filtering for scalability.
+   *
+   * 2-phase strategy:
+   * 1. SQL pre-filters by project + sorts by recency (loads max N candidates)
+   * 2. JS computes cosine similarity only on filtered candidates
+   *
+   * With 50k observations and maxCandidates=2000, loads only ~4% of data.
    */
   async search(db, queryEmbedding, options = {}) {
     const limit = options.limit || 10;
     const threshold = options.threshold || 0.3;
+    const maxCandidates = options.maxCandidates || DEFAULT_MAX_CANDIDATES;
     try {
-      let sql = `
+      const conditions = [];
+      const params = [];
+      if (options.project) {
+        conditions.push("o.project = ?");
+        params.push(options.project);
+      }
+      const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
+      const sql = `
         SELECT e.observation_id, e.embedding,
                o.title, o.text, o.type, o.project, o.created_at, o.created_at_epoch
         FROM observation_embeddings e
         JOIN observations o ON o.id = e.observation_id
+        ${whereClause}
+        ORDER BY o.created_at_epoch DESC
+        LIMIT ?
       `;
-      const params = [];
-      if (options.project) {
-        sql += " WHERE o.project = ?";
-        params.push(options.project);
-      }
+      params.push(maxCandidates);
       const rows = db.query(sql).all(...params);
       const scored = [];
       for (const row of rows) {
@@ -419,14 +545,15 @@ var VectorSearch = class {
         }
       }
       scored.sort((a, b) => b.similarity - a.similarity);
+      logger.debug("VECTOR", `Search: ${rows.length} candidates \u2192 ${scored.length} above threshold \u2192 ${Math.min(scored.length, limit)} results`);
       return scored.slice(0, limit);
     } catch (error) {
-      logger.error("VECTOR", `Errore ricerca vettoriale: ${error}`);
+      logger.error("VECTOR", `Vector search error: ${error}`);
       return [];
     }
   }
   /**
-   * Salva embedding per un'osservazione.
+   * Store embedding for an observation.
    */
   async storeEmbedding(db, observationId, embedding, model) {
     try {
@@ -442,18 +569,18 @@ var VectorSearch = class {
         embedding.length,
         (/* @__PURE__ */ new Date()).toISOString()
       );
-      logger.debug("VECTOR", `Embedding salvato per osservazione ${observationId}`);
+      logger.debug("VECTOR", `Embedding saved for observation ${observationId}`);
     } catch (error) {
-      logger.error("VECTOR", `Errore salvataggio embedding: ${error}`);
+      logger.error("VECTOR", `Error saving embedding: ${error}`);
     }
   }
   /**
-   * Genera embeddings per osservazioni che non li hanno ancora.
+   * Generate embeddings for observations that don't have them yet.
    */
   async backfillEmbeddings(db, batchSize = 50) {
     const embeddingService2 = getEmbeddingService();
     if (!await embeddingService2.initialize()) {
-      logger.warn("VECTOR", "Embedding service non disponibile, backfill saltato");
+      logger.warn("VECTOR", "Embedding service not available, backfill skipped");
       return 0;
     }
     const rows = db.query(`
@@ -466,7 +593,7 @@ var VectorSearch = class {
     `).all(batchSize);
     if (rows.length === 0) return 0;
     let count = 0;
-    const model = embeddingService2.getProvider() || "unknown";
+    const model = embeddingService2.getModelName();
     for (const row of rows) {
       const parts = [row.title];
       if (row.text) parts.push(row.text);
@@ -479,11 +606,11 @@ var VectorSearch = class {
         count++;
       }
     }
-    logger.info("VECTOR", `Backfill completato: ${count}/${rows.length} embeddings generati`);
+    logger.info("VECTOR", `Backfill completed: ${count}/${rows.length} embeddings generated`);
     return count;
   }
   /**
-   * Statistiche sugli embeddings.
+   * Embedding statistics.
    */
   getStats(db) {
     try {