kiro-memory 1.9.0 → 3.0.0

package/plugin/dist/services/search/EmbeddingService.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;
+  }
+  /**
+   * 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;
   }
-  // --- Provider specifici ---
+  // --- Single-text provider implementations ---
   async _embedFastembed(text) {
     const embeddings = this.model.embed([text], 1);
     for await (const batch of embeddings) {