bluera-knowledge 0.19.7 → 0.21.0

package/dist/{chunk-JPJI3VMA.js → chunk-ZR23KJPJ.js}

@@ -4,7 +4,7 @@ import {
 } from "./chunk-CLIMKLTW.js";
 import {
   parseIgnorePatternsForScanning
-} from "./chunk-HXBIIMYL.js";
+} from "./chunk-H25AEF47.js";
 import {
   __require
 } from "./chunk-DGUM43GV.js";
@@ -2063,8 +2063,14 @@ var DEFAULT_CONFIG = {
   version: 1,
   dataDir: ".bluera/bluera-knowledge/data",
   embedding: {
-    model: "Xenova/all-MiniLM-L6-v2",
-    batchSize: 32
+    model: "Xenova/bge-small-en-v1.5",
+    batchSize: 32,
+    dtype: "fp32",
+    pooling: "mean",
+    normalize: true,
+    queryPrefix: "Represent this sentence for searching relevant passages: ",
+    docPrefix: "",
+    maxInFlightBatches: 1
   },
   indexing: {
     concurrency: 4,
@@ -3749,15 +3755,13 @@ function detectContentType(results) {
 }
 var SearchService = class {
   lanceStore;
-  embeddingEngine;
   codeUnitService;
   codeGraphService;
   graphCache;
   searchConfig;
   unsubscribeCacheInvalidation;
-  constructor(lanceStore, embeddingEngine, codeGraphService, searchConfig) {
+  constructor(lanceStore, codeGraphService, searchConfig) {
     this.lanceStore = lanceStore;
-    this.embeddingEngine = embeddingEngine;
     this.codeUnitService = new CodeUnitService();
     this.codeGraphService = codeGraphService;
     this.graphCache = /* @__PURE__ */ new Map();
@@ -3972,17 +3976,17 @@ var SearchService = class {
   /**
    * Fetch raw vector search results without normalization.
    * Returns results with raw cosine similarity scores [0-1].
+   * Uses LanceDB's embedding function for query embedding,
+   * ensuring consistent query/document embedding through a single code path.
    */
   async vectorSearchRaw(query, stores, limit) {
-    const queryVector = await this.embeddingEngine.embed(query);
     const results = [];
     for (const storeId of stores) {
-      const hits = await this.lanceStore.search(storeId, queryVector, limit);
+      const hits = await this.lanceStore.searchText(storeId, query, limit);
       results.push(
         ...hits.map((r) => ({
           id: r.id,
           score: r.score,
-          // Raw cosine similarity (1 - distance)
           content: r.content,
           metadata: r.metadata
         }))
@@ -4816,6 +4820,9 @@ function extractRepoName(url) {
   return name;
 }
 
+// src/types/store.ts
+var CURRENT_SCHEMA_VERSION = 2;
+
 // src/services/store.service.ts
 async function fileExists4(path4) {
   try {
@@ -4830,12 +4837,21 @@ var StoreService = class {
   definitionService;
   gitignoreService;
   projectRoot;
+  embeddingModelId;
   registry = { stores: [] };
   constructor(dataDir, options) {
     this.dataDir = dataDir;
-    this.definitionService = options?.definitionService ?? void 0;
-    this.gitignoreService = options?.gitignoreService ?? void 0;
-    this.projectRoot = options?.projectRoot ?? void 0;
+    this.definitionService = options.definitionService ?? void 0;
+    this.gitignoreService = options.gitignoreService ?? void 0;
+    this.projectRoot = options.projectRoot ?? void 0;
+    this.embeddingModelId = options.embeddingModelId;
+  }
+  /**
+   * Get the current embedding model ID used for new stores.
+   * Used by model compatibility validation.
+   */
+  getCurrentModelId() {
+    return this.embeddingModelId;
   }
   async initialize() {
     await mkdir5(this.dataDir, { recursive: true });
@@ -4974,7 +4990,9 @@ var StoreService = class {
           tags: input.tags,
           status: "ready",
           createdAt: now,
-          updatedAt: now
+          updatedAt: now,
+          schemaVersion: CURRENT_SCHEMA_VERSION,
+          modelId: this.embeddingModelId
         };
         break;
       }
@@ -5019,7 +5037,9 @@ var StoreService = class {
           tags: input.tags,
           status: "ready",
           createdAt: now,
-          updatedAt: now
+          updatedAt: now,
+          schemaVersion: CURRENT_SCHEMA_VERSION,
+          modelId: this.embeddingModelId
         };
         break;
       }
@@ -5040,7 +5060,9 @@ var StoreService = class {
           tags: input.tags,
           status: "ready",
           createdAt: now,
-          updatedAt: now
+          updatedAt: now,
+          schemaVersion: CURRENT_SCHEMA_VERSION,
+          modelId: this.embeddingModelId
         };
         break;
       default: {
@@ -5418,6 +5440,16 @@ import { homedir as homedir2 } from "os";
 import { join as join11 } from "path";
 import { pipeline, env } from "@huggingface/transformers";
 env.cacheDir = join11(homedir2(), ".cache", "huggingface-transformers");
+var DEFAULT_EMBEDDING_CONFIG = {
+  model: "Xenova/bge-small-en-v1.5",
+  batchSize: 32,
+  dtype: "fp32",
+  pooling: "mean",
+  normalize: true,
+  queryPrefix: "Represent this sentence for searching relevant passages: ",
+  docPrefix: "",
+  maxInFlightBatches: 1
+};
 var EmbeddingEngine = class {
   extractor = null;
   initPromise = null;
@@ -5425,11 +5457,9 @@ var EmbeddingEngine = class {
   _dimensions = null;
   // eslint-disable-next-line @typescript-eslint/prefer-readonly -- mutated in dispose()
   disposed = false;
-  modelName;
-  batchSize;
-  constructor(modelName = "Xenova/all-MiniLM-L6-v2", batchSize = 32) {
-    this.modelName = modelName;
-    this.batchSize = batchSize;
+  config;
+  constructor(config = DEFAULT_EMBEDDING_CONFIG) {
+    this.config = config;
   }
   /**
    * Guard against use-after-dispose
@@ -5448,8 +5478,8 @@ var EmbeddingEngine = class {
     if (this.extractor !== null) return;
     this.initPromise ??= (async () => {
       try {
-        this.extractor = await pipeline("feature-extraction", this.modelName, {
-          dtype: "fp32"
+        this.extractor = await pipeline("feature-extraction", this.config.model, {
+          dtype: this.config.dtype
         });
       } catch (error) {
         this.initPromise = null;
@@ -5458,7 +5488,22 @@ var EmbeddingEngine = class {
     })();
     await this.initPromise;
   }
-  async embed(text) {
+  /**
+   * Embed a search query. Applies queryPrefix for asymmetric models.
+   */
+  async embedQuery(text) {
+    return this.embedText(this.config.queryPrefix + text);
+  }
+  /**
+   * Embed a document for indexing. Applies docPrefix for asymmetric models.
+   */
+  async embedDocument(text) {
+    return this.embedText(this.config.docPrefix + text);
+  }
+  /**
+   * Internal: embed text without prefix.
+   */
+  async embedText(text) {
     this.assertNotDisposed();
     if (this.extractor === null) {
       await this.initialize();
@@ -5467,13 +5512,17 @@ var EmbeddingEngine = class {
       throw new Error("Failed to initialize embedding model");
     }
     const output = await this.extractor(text, {
-      pooling: "mean",
-      normalize: true
+      pooling: this.config.pooling,
+      normalize: this.config.normalize
     });
-    const result = Array.from(output.data, (v) => Number(v));
-    this._dimensions ??= result.length;
-    return result;
+    const dim = output.dims[output.dims.length - 1] ?? 0;
+    this._dimensions ??= dim;
+    return Float32Array.from(output.data);
   }
+  /**
+   * Embed a batch of documents with optional parallelism.
+   * When maxInFlightBatches > 1, processes multiple batches concurrently.
+   */
   async embedBatch(texts) {
     this.assertNotDisposed();
     if (this.extractor === null) {
@@ -5482,26 +5531,79 @@ var EmbeddingEngine = class {
     if (this.extractor === null) {
       throw new Error("Failed to initialize embedding model");
     }
+    const batches = [];
+    for (let i = 0; i < texts.length; i += this.config.batchSize) {
+      batches.push(texts.slice(i, i + this.config.batchSize));
+    }
+    if (batches.length === 0) {
+      return [];
+    }
+    if (this.config.maxInFlightBatches <= 1) {
+      return this.embedBatchesSequential(batches);
+    } else {
+      return this.embedBatchesConcurrent(batches);
+    }
+  }
+  /**
+   * Process batches sequentially (original behavior).
+   */
+  async embedBatchesSequential(batches) {
     const results = [];
-    for (let i = 0; i < texts.length; i += this.batchSize) {
-      const batch = texts.slice(i, i + this.batchSize);
-      const output = await this.extractor(batch, {
-        pooling: "mean",
-        normalize: true
-      });
-      const dim = output.dims[output.dims.length - 1] ?? 0;
-      for (let b = 0; b < batch.length; b++) {
-        const start = b * dim;
-        const end = start + dim;
-        results.push(Array.from(output.data.slice(start, end), (v) => Number(v)));
-      }
-      this._dimensions ??= dim;
-      if (i + this.batchSize < texts.length) {
+    for (let i = 0; i < batches.length; i++) {
+      const batch = batches[i];
+      if (batch === void 0) continue;
+      const batchResults = await this.processSingleBatch(batch);
+      results.push(...batchResults);
+      if (i < batches.length - 1) {
         await new Promise((resolve4) => setImmediate(resolve4));
       }
     }
     return results;
   }
+  /**
+   * Process batches with controlled concurrency.
+   */
+  async embedBatchesConcurrent(batches) {
+    const results = new Array(batches.length);
+    let inFlight = 0;
+    const maxConcurrent = this.config.maxInFlightBatches;
+    await Promise.all(
+      batches.map(async (batch, idx) => {
+        while (inFlight >= maxConcurrent) {
+          await new Promise((resolve4) => setImmediate(resolve4));
+        }
+        inFlight++;
+        try {
+          results[idx] = await this.processSingleBatch(batch);
+        } finally {
+          inFlight--;
+        }
+      })
+    );
+    return results.flat();
+  }
+  /**
+   * Process a single batch and return embeddings.
+   */
+  async processSingleBatch(batch) {
+    if (this.extractor === null) {
+      throw new Error("Extractor not initialized");
+    }
+    const prefixedBatch = batch.map((text) => this.config.docPrefix + text);
+    const output = await this.extractor(prefixedBatch, {
+      pooling: this.config.pooling,
+      normalize: this.config.normalize
+    });
+    const dim = output.dims[output.dims.length - 1] ?? 0;
+    const batchResults = [];
+    for (let b = 0; b < batch.length; b++) {
+      const start = b * dim;
+      const end = start + dim;
+      batchResults.push(Float32Array.from(output.data.slice(start, end)));
+    }
+    this._dimensions ??= dim;
+    return batchResults;
+  }
   /**
    * Get cached embedding dimensions. Throws if embed() hasn't been called yet.
    * Use ensureDimensions() if you need to guarantee dimensions are available.
@@ -5512,13 +5614,38 @@ var EmbeddingEngine = class {
     }
     return this._dimensions;
   }
+  /**
+   * Check if the embedding pipeline is initialized.
+   */
+  isInitialized() {
+    return this.extractor !== null;
+  }
+  /**
+   * Check if this engine has been disposed.
+   */
+  isDisposed() {
+    return this.disposed;
+  }
+  /**
+   * Reset the engine to uninitialized state, allowing reuse after disposal.
+   * If currently initialized, disposes the pipeline first.
+   */
+  async reset() {
+    if (this.extractor !== null) {
+      await this.extractor.dispose();
+      this.extractor = null;
+    }
+    this.initPromise = null;
+    this._dimensions = null;
+    this.disposed = false;
+  }
   /**
    * Ensure dimensions are available, initializing the model if needed.
    * Returns the embedding dimensions for the current model.
    */
   async ensureDimensions() {
     if (this._dimensions === null) {
-      await this.embed("dimension probe");
+      await this.embedText("dimension probe");
     }
     if (this._dimensions === null) {
       throw new Error("Failed to determine embedding dimensions");
@@ -5543,6 +5670,88 @@ var EmbeddingEngine = class {
 
 // src/db/lance.ts
 import * as lancedb from "@lancedb/lancedb";
+import { LanceSchema } from "@lancedb/lancedb/embedding";
+import { Utf8 } from "apache-arrow";
+
+// src/db/lance-embedding-function.ts
+import { TextEmbeddingFunction, getRegistry } from "@lancedb/lancedb/embedding";
+import { Float32 } from "apache-arrow";
+var HuggingFaceEmbeddingFunction = class extends TextEmbeddingFunction {
+  engine;
+  embeddingConfig;
+  _ndims = null;
+  constructor(optionsRaw) {
+    super();
+    const options = this.resolveVariables(optionsRaw ?? {});
+    this.embeddingConfig = {
+      model: options.model ?? "Xenova/bge-small-en-v1.5",
+      batchSize: options.batchSize ?? 32,
+      dtype: options.dtype ?? "fp32",
+      pooling: options.pooling ?? "mean",
+      normalize: options.normalize ?? true,
+      queryPrefix: options.queryPrefix ?? "",
+      docPrefix: options.docPrefix ?? "",
+      maxInFlightBatches: 1
+      // Single-threaded for LanceDB integration
+    };
+    this.engine = new EmbeddingEngine(this.embeddingConfig);
+  }
+  /**
+   * Initialize the embedding model. Called by LanceDB before embeddings are computed.
+   */
+  async init() {
+    this._ndims = await this.engine.ensureDimensions();
+  }
+  /**
+   * Return embedding dimensions. Must call init() first.
+   */
+  ndims() {
+    if (this._ndims === null) {
+      throw new Error("HuggingFaceEmbeddingFunction not initialized. Call init() first.");
+    }
+    return this._ndims;
+  }
+  /**
+   * Return embedding data type (always Float32 for our models).
+   */
+  embeddingDataType() {
+    return new Float32();
+  }
+  /**
+   * Generate embeddings for a batch of texts (documents).
+   * Called during table.add() operations.
+   */
+  async generateEmbeddings(texts) {
+    return this.engine.embedBatch(texts);
+  }
+  /**
+   * Compute embedding for a single query.
+   * Called during table.search(query) operations.
+   */
+  async computeQueryEmbeddings(data) {
+    const embedding = await this.engine.embedQuery(data);
+    return Array.from(embedding);
+  }
+  /**
+   * Get the model ID for provenance tracking.
+   */
+  getModelId() {
+    return this.embeddingConfig.model;
+  }
+  /**
+   * Get the full embedding config.
+   */
+  getConfig() {
+    return this.embeddingConfig;
+  }
+  /**
+   * Dispose the underlying engine to free resources.
+   */
+  async dispose() {
+    await this.engine.dispose();
+  }
+};
+getRegistry().register("HuggingFaceEmbeddingFunction")(HuggingFaceEmbeddingFunction);
 
 // src/types/document.ts
 import { z as z5 } from "zod";
@@ -5560,15 +5769,51 @@ var DocumentMetadataSchema = z5.object({
 }).loose();
 
 // src/db/lance.ts
+function isSearchHit(value) {
+  if (typeof value !== "object" || value === null) return false;
+  return "id" in value && "content" in value && "metadata" in value && "_distance" in value && typeof value.id === "string" && typeof value.content === "string" && typeof value.metadata === "string" && typeof value._distance === "number";
+}
+function parseDocumentMetadata(jsonStr) {
+  const parsed = DocumentMetadataSchema.parse(JSON.parse(jsonStr));
+  return {
+    ...parsed,
+    storeId: createStoreId(parsed.storeId)
+  };
+}
 var LanceStore = class {
   connection = null;
   tables = /* @__PURE__ */ new Map();
   dataDir;
   // eslint-disable-next-line @typescript-eslint/prefer-readonly -- set via setDimensions()
   _dimensions = null;
+  embeddingFunction = null;
   constructor(dataDir) {
     this.dataDir = dataDir;
   }
+  /**
+   * Set the embedding function for auto-embedding queries.
+   * Must be called before initialize() for new tables.
+   * The embedding function is initialized and its dimensions are used for schema creation.
+   */
+  async setEmbeddingFunction(config) {
+    this.embeddingFunction = new HuggingFaceEmbeddingFunction({
+      model: config.model,
+      batchSize: config.batchSize,
+      dtype: config.dtype,
+      pooling: config.pooling,
+      normalize: config.normalize,
+      queryPrefix: config.queryPrefix,
+      docPrefix: config.docPrefix
+    });
+    await this.embeddingFunction.init();
+    this._dimensions = this.embeddingFunction.ndims();
+  }
+  /**
+   * Check if embedding function is available for auto-embedding queries.
+   */
+  hasEmbeddingFunction() {
+    return this.embeddingFunction !== null;
+  }
   /**
    * Set the embedding dimensions. Must be called before initialize().
    * This allows dimensions to be derived from the embedding model at runtime.
@@ -5579,22 +5824,35 @@ var LanceStore = class {
   }
   async initialize(storeId) {
     if (this._dimensions === null) {
-      throw new Error("Dimensions not set. Call setDimensions() before initialize().");
+      throw new Error(
+        "Dimensions not set. Call setDimensions() or setEmbeddingFunction() before initialize()."
+      );
     }
     this.connection ??= await lancedb.connect(this.dataDir);
     const tableName = this.getTableName(storeId);
     const tableNames = await this.connection.tableNames();
     if (!tableNames.includes(tableName)) {
-      const table = await this.connection.createTable(tableName, [
-        {
-          id: "__init__",
-          content: "",
-          vector: new Array(this._dimensions).fill(0),
-          metadata: "{}"
-        }
-      ]);
-      await table.delete('id = "__init__"');
-      this.tables.set(tableName, table);
+      if (this.embeddingFunction !== null) {
+        const schema = LanceSchema({
+          id: new Utf8(),
+          content: this.embeddingFunction.sourceField(),
+          vector: this.embeddingFunction.vectorField(),
+          metadata: new Utf8()
+        });
+        const table = await this.connection.createEmptyTable(tableName, schema);
+        this.tables.set(tableName, table);
+      } else {
+        const table = await this.connection.createTable(tableName, [
+          {
+            id: "__init__",
+            content: "",
+            vector: new Array(this._dimensions).fill(0),
+            metadata: "{}"
+          }
+        ]);
+        await table.delete('id = "__init__"');
+        this.tables.set(tableName, table);
+      }
     } else {
       const table = await this.connection.openTable(tableName);
       this.tables.set(tableName, table);
@@ -5638,6 +5896,29 @@ var LanceStore = class {
       };
     });
   }
+  /**
+   * Search using a text query with automatic embedding.
+   * Requires setEmbeddingFunction() to have been called.
+   * Uses the embedding function to compute query embeddings consistently with document embeddings.
+   */
+  async searchText(storeId, query, limit) {
+    if (this.embeddingFunction === null) {
+      throw new Error(
+        "Embedding function not set. Call setEmbeddingFunction() before searchText()."
+      );
+    }
+    const queryEmbedding = await this.embeddingFunction.computeQueryEmbeddings(query);
+    const table = await this.getTable(storeId);
+    const searchQuery = table.vectorSearch(queryEmbedding).limit(limit).distanceType("cosine");
+    const rawResults = await searchQuery.toArray();
+    const results = rawResults.filter(isSearchHit);
+    return results.map((r) => ({
+      id: createDocumentId(r.id),
+      content: r.content,
+      score: 1 - r._distance,
+      metadata: parseDocumentMetadata(r.metadata)
+    }));
+  }
   async createFtsIndex(storeId) {
     const table = await this.getTable(storeId);
     await table.createIndex("content", {
@@ -5734,10 +6015,7 @@ var LazyServiceContainer = class {
   get embeddings() {
     if (this._embeddings === null) {
       logger4.debug("Lazy-initializing EmbeddingEngine");
-      this._embeddings = new EmbeddingEngine(
-        this.appConfig.embedding.model,
-        this.appConfig.embedding.batchSize
-      );
+      this._embeddings = new EmbeddingEngine(this.appConfig.embedding);
     }
     return this._embeddings;
   }
@@ -5757,12 +6035,7 @@ var LazyServiceContainer = class {
   get search() {
     if (this._search === null) {
       logger4.debug("Lazy-initializing SearchService");
-      this._search = new SearchService(
-        this.lance,
-        this.embeddings,
-        this.codeGraph,
-        this.appConfig.search
-      );
+      this._search = new SearchService(this.lance, this.codeGraph, this.appConfig.search);
     }
     return this._search;
   }
@@ -5821,7 +6094,8 @@ async function createLazyServices(configPath, dataDir, projectRoot) {
   const storeOptions = {
     definitionService,
     gitignoreService,
-    projectRoot: resolvedProjectRoot
+    projectRoot: resolvedProjectRoot,
+    embeddingModelId: appConfig.embedding.model
   };
   const store = new StoreService(resolvedDataDir, storeOptions);
   await store.initialize();
@@ -5840,21 +6114,23 @@ async function createServices(configPath, dataDir, projectRoot) {
   const pythonBridge = new PythonBridge();
   await pythonBridge.start();
   const lance = new LanceStore(resolvedDataDir);
-  const embeddings = new EmbeddingEngine(appConfig.embedding.model, appConfig.embedding.batchSize);
+  const embeddings = new EmbeddingEngine(appConfig.embedding);
   await embeddings.initialize();
+  await lance.setEmbeddingFunction(appConfig.embedding);
   const resolvedProjectRoot = config.resolveProjectRoot();
   const definitionService = new StoreDefinitionService(resolvedProjectRoot);
   const gitignoreService = new GitignoreService(resolvedProjectRoot);
   const storeOptions = {
     definitionService,
     gitignoreService,
-    projectRoot: resolvedProjectRoot
+    projectRoot: resolvedProjectRoot,
+    embeddingModelId: appConfig.embedding.model
   };
   const store = new StoreService(resolvedDataDir, storeOptions);
   await store.initialize();
   const codeGraph = new CodeGraphService(resolvedDataDir, pythonBridge);
   const manifest = new ManifestService(resolvedDataDir);
-  const search = new SearchService(lance, embeddings, codeGraph, appConfig.search);
+  const search = new SearchService(lance, codeGraph, appConfig.search);
   const index = new IndexService(lance, embeddings, {
     codeGraphService: codeGraph,
     manifestService: manifest,
@@ -5946,4 +6222,4 @@ export {
   createServices,
   destroyServices
 };
-//# sourceMappingURL=chunk-JPJI3VMA.js.map
+//# sourceMappingURL=chunk-ZR23KJPJ.js.map