fix(core): reuse semantic query vectors and bound inference batches

docs/configuration.md

@@ -148,6 +148,10 @@ mutable and therefore unverified.
 Transformers pipelines are shared per exact model identity and disposed when the final
 `RagmirClient` owner closes. Cache retirement waits for active inference leases, so a model switch
 or shutdown does not dispose a session still serving a request.
+Repeated single-query embeddings use a project-scoped, in-memory LRU cache tied to that model
+identity, capped at 128 entries and 65,536 vector values. The cache retains hashed query keys and
+vectors, never search results; each search still reads the current index generation. Every semantic
+model call respects `embeddingBatchSize`, including callers that supply a larger input array.
 
 ## Local extractors
 

packages/ragmir-core/src/embeddings.test.ts

@@ -3,6 +3,7 @@ import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"
 import os from "node:os"
 import path from "node:path"
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
+import { RagmirClient } from "./client.js"
 import {
   clearTransformersCache,
   disposeTransformersCache,
@@ -105,6 +106,208 @@ describe("local hash embeddings", () => {
 })
 
 describe("embedding model adapters", () => {
+  it("reuses a single-query embedding without caching document vectors or sharing mutable arrays", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [0.25, 0.75]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+
+    const first = await embedText("same evidence", config)
+    first[0] = 99
+    expect(await embedText("same evidence", config)).toEqual([0.25, 0.75])
+    await embedTexts(["same evidence"], config)
+    await embedTexts(["same evidence"], config)
+
+    expect(extractor).toHaveBeenCalledTimes(3)
+  })
+
+  it("isolates query vectors by project root and exact model identity", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [1, 0]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+    await embedText("same evidence", config)
+    await embedText("same evidence", { ...config, projectRoot: "/tmp/another-project" })
+    await embedText("same evidence", { ...config, embeddingModelRevision: "new-revision" })
+    await embedText("same evidence", {
+      ...config,
+      embeddingModelDigest: `sha256:${"a".repeat(64)}`,
+    })
+
+    expect(extractor).toHaveBeenCalledTimes(4)
+    expect(transformersMock.pipeline).toHaveBeenCalledTimes(3)
+  })
+
+  it("evicts the least recently used query when the entry limit is reached", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [1, 0]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+    for (let index = 0; index < 128; index += 1) await embedText(`query ${index}`, config)
+    await embedText("query 0", config)
+    await embedText("query 128", config)
+    await embedText("query 0", config)
+    expect(extractor).toHaveBeenCalledTimes(129)
+
+    await embedText("query 1", config)
+    expect(extractor).toHaveBeenCalledTimes(130)
+  })
+
+  it("bounds cached vector values independently from the entry count", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => Array.from({ length: 40_000 }, () => 1)),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+    await embedText("first", config)
+    await embedText("second", config)
+    await embedText("second", config)
+    expect(extractor).toHaveBeenCalledTimes(2)
+
+    await embedText("first", config)
+    expect(extractor).toHaveBeenCalledTimes(3)
+  })
+
+  it("does not cache a vector larger than the total cache budget", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => Array.from({ length: 65_537 }, () => 1)),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+    await embedText("oversized", config)
+    await embedText("oversized", config)
+
+    expect(extractor).toHaveBeenCalledTimes(2)
+  })
+
+  it("retries failed inference and drops query vectors when the model is disposed", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [1, 0]),
+    }))
+    extractor.mockRejectedValueOnce(new Error("inference failed"))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+
+    await expect(embedText("same evidence", config)).rejects.toThrow("inference failed")
+    await expect(embedText("same evidence", config)).resolves.toEqual([1, 0])
+    await embedText("same evidence", config)
+    expect(extractor).toHaveBeenCalledTimes(2)
+    await disposeTransformersCache()
+    await embedText("same evidence", config)
+    expect(extractor).toHaveBeenCalledTimes(3)
+  })
+
+  it("rejects an aborted cached query without running inference", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [1, 0]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig()
+    await embedText("same evidence", config)
+    const controller = new AbortController()
+    controller.abort()
+
+    await expect(embedText("same evidence", config, controller.signal)).rejects.toMatchObject({
+      code: "ABORTED",
+    })
+    expect(extractor).toHaveBeenCalledTimes(1)
+  })
+
+  it("refreshes search evidence after ingestion while reusing the query vector", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [1, 0]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const modelConfig = await transformerArtifactConfig()
+    const root = path.dirname(modelConfig.embeddingModelPath)
+    const config = testConfig(root, {
+      embeddingProvider: "transformers",
+      embeddingModel: modelConfig.embeddingModel,
+      embeddingModelRevision: modelConfig.embeddingModelRevision,
+      embeddingModelPath: modelConfig.embeddingModelPath,
+    })
+    await mkdir(config.rawDir, { recursive: true })
+    const sourcePath = path.join(config.rawDir, "decision.md")
+    await writeFile(sourcePath, "Production approval requires the original evidence.\n")
+    const reader = await RagmirClient.createWithConfig(config)
+    let writer: RagmirClient | undefined
+    try {
+      await reader.ingest()
+      const first = await reader.search("production approval")
+      const warm = await reader.search("production approval")
+      expect(warm).toEqual(first)
+
+      await writeFile(sourcePath, "Production approval now requires signed refreshed evidence.\n")
+      writer = await RagmirClient.createWithConfig(config)
+      await writer.ingest({ rebuild: true })
+      await writer.close()
+      const inferenceCalls = extractor.mock.calls.length
+      const refreshed = await reader.search("production approval")
+
+      expect(refreshed[0]?.text).toContain("signed refreshed evidence")
+      expect(refreshed[0]?.evidence.id).not.toBe(first[0]?.evidence.id)
+      expect(extractor).toHaveBeenCalledTimes(inferenceCalls)
+    } finally {
+      await writer?.close()
+      await reader.close()
+    }
+  })
+
+  it("bounds every semantic model call and preserves input order across batches", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map((text) => [Number(text.match(/\d+$/u)?.[0]), 0]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+    const config = transformerConfig({ embeddingBatchSize: 2 })
+
+    expect(await embedTexts(["row 0", "row 1", "row 2", "row 3", "row 4"], config)).toEqual([
+      [0, 0],
+      [1, 0],
+      [2, 0],
+      [3, 0],
+      [4, 0],
+    ])
+    expect(extractor.mock.calls.map(([texts]) => texts.length)).toEqual([2, 2, 1])
+  })
+
+  it("enforces the hard batch ceiling even for a directly supplied config", async () => {
+    const extractor = vi.fn(async (texts: string[]) => ({
+      tolist: () => texts.map(() => [1, 0]),
+    }))
+    transformersMock.pipeline.mockResolvedValue(extractor)
+
+    await embedTexts(
+      Array.from({ length: 129 }, () => "evidence"),
+      transformerConfig({ embeddingBatchSize: 1_000 }),
+    )
+    expect(extractor.mock.calls.map(([texts]) => texts.length)).toEqual([128, 1])
+  })
+
+  it("stops between semantic batches and releases the model after cancellation", async () => {
+    const controller = new AbortController()
+    const dispose = vi.fn(async () => undefined)
+    const extractor = vi.fn(async (texts: string[]) => {
+      controller.abort()
+      return { tolist: () => texts.map(() => [1, 0]) }
+    })
+    transformersMock.pipeline.mockResolvedValue(Object.assign(extractor, { dispose }))
+
+    await expect(
+      embedTexts(
+        ["first", "second", "third"],
+        transformerConfig({ embeddingBatchSize: 2 }),
+        "document",
+        controller.signal,
+      ),
+    ).rejects.toMatchObject({ code: "ABORTED" })
+    expect(extractor).toHaveBeenCalledTimes(1)
+    await disposeTransformersCache()
+    expect(dispose).toHaveBeenCalledOnce()
+  })
+
   it("adds the asymmetric E5 retrieval prefixes", () => {
     const model = "intfloat/multilingual-e5-small"
 

packages/ragmir-core/src/embeddings.ts

@@ -2,6 +2,7 @@ import { createHash } from "node:crypto"
 import { createReadStream } from "node:fs"
 import { mkdir, readdir, stat } from "node:fs/promises"
 import path from "node:path"
+import { MAX_EMBEDDING_BATCH_SIZE } from "./defaults.js"
 import { throwIfAborted } from "./operation.js"
 import { tokenize } from "./text.js"
 import type { Config, IngestionEmbeddingModelState } from "./types.js"
@@ -19,6 +20,8 @@ const MAX_LOCAL_HASH_FEATURE_CACHE_ENTRIES = 16_384
  * pipeline is disposed as soon as its final lease is released.
  */
 const MAX_TRANSFORMERS_PIPELINES = 1
+const MAX_QUERY_EMBEDDING_CACHE_ENTRIES = 128
+const MAX_QUERY_EMBEDDING_CACHE_VALUES = 65_536
 const transformersPipelines = new Map<string, TransformersPipelineEntry>()
 const transformersModelOwners = new Map<string, number>()
 
@@ -38,13 +41,20 @@ interface TransformersPipelineEntry {
   retired: boolean
   idleWaiters: Array<() => void>
   disposal: Promise<void> | undefined
+  queryEmbeddings: QueryEmbeddingCache
 }
 
 interface TransformersPipelineLease {
   extractor: TransformersExtractor
+  queryEmbeddings: QueryEmbeddingCache
   release(): Promise<void>
 }
 
+interface QueryEmbeddingCache {
+  vectors: Map<string, number[]>
+  values: number
+}
+
 interface LocalHashFeatureLocation {
   index: number
   sign: 1 | -1
@@ -77,7 +87,7 @@ export async function embedTexts(
       return texts.map((text) => localHashEmbedding(text, featureCache))
     }
 
-    const embeddings = await embedWithTransformers(texts, config, inputType)
+    const embeddings = await embedWithTransformers(texts, config, inputType, signal)
     throwIfAborted(signal)
     return embeddings
   })
@@ -151,25 +161,74 @@ async function embedWithTransformers(
   texts: string[],
   config: Config,
   inputType: EmbeddingInputType,
+  signal?: AbortSignal,
 ): Promise<number[][]> {
   const lease = await acquireTransformersPipeline(config)
   try {
-    const preparedTexts = texts.map((text) =>
-      prepareEmbeddingText(text, config.embeddingModel, inputType),
-    )
-    const output = await lease.extractor(preparedTexts, { pooling: "mean", normalize: true })
-    const rows = tensorToEmbeddingRows(output)
-
-    if (rows.length !== texts.length) {
-      throw new Error(`Expected ${texts.length} embeddings, received ${rows.length}.`)
+    throwIfAborted(signal)
+    // Cache only single-query vectors. Search still reads and ranks the current index snapshot.
+    const cacheKey =
+      inputType === "query" && texts.length === 1
+        ? createHash("sha256")
+            .update(JSON.stringify([config.projectRoot, texts[0]]))
+            .digest("hex")
+        : null
+    if (cacheKey !== null) {
+      const cached = cachedQueryEmbedding(lease.queryEmbeddings, cacheKey)
+      if (cached) return [cached]
     }
 
-    return rows
+    const embeddings: number[][] = []
+    const batchSize = Math.max(1, Math.min(config.embeddingBatchSize, MAX_EMBEDDING_BATCH_SIZE))
+    for (let offset = 0; offset < texts.length; offset += batchSize) {
+      throwIfAborted(signal)
+      const preparedTexts = texts
+        .slice(offset, offset + batchSize)
+        .map((text) => prepareEmbeddingText(text, config.embeddingModel, inputType))
+      const output = await lease.extractor(preparedTexts, { pooling: "mean", normalize: true })
+      throwIfAborted(signal)
+      const rows = tensorToEmbeddingRows(output)
+      if (rows.length !== preparedTexts.length) {
+        throw new Error(`Expected ${preparedTexts.length} embeddings, received ${rows.length}.`)
+      }
+      embeddings.push(...rows)
+    }
+    const queryEmbedding = embeddings[0]
+    if (cacheKey !== null && queryEmbedding) {
+      cacheQueryEmbedding(lease.queryEmbeddings, cacheKey, queryEmbedding)
+    }
+    return embeddings
   } finally {
     await lease.release()
   }
 }
 
+function cachedQueryEmbedding(cache: QueryEmbeddingCache, key: string): number[] | undefined {
+  const vector = cache.vectors.get(key)
+  if (!vector) return undefined
+  cache.vectors.delete(key)
+  cache.vectors.set(key, vector)
+  return [...vector]
+}
+
+function cacheQueryEmbedding(cache: QueryEmbeddingCache, key: string, vector: number[]): void {
+  if (vector.length > MAX_QUERY_EMBEDDING_CACHE_VALUES) return
+  const previous = cache.vectors.get(key)
+  cache.values -= previous?.length ?? 0
+  cache.vectors.delete(key)
+  while (
+    cache.vectors.size >= MAX_QUERY_EMBEDDING_CACHE_ENTRIES ||
+    cache.values + vector.length > MAX_QUERY_EMBEDDING_CACHE_VALUES
+  ) {
+    const oldest = cache.vectors.entries().next().value
+    if (!oldest) break
+    cache.vectors.delete(oldest[0])
+    cache.values -= oldest[1].length
+  }
+  cache.vectors.set(key, [...vector])
+  cache.values += vector.length
+}
+
 export function prepareEmbeddingText(
   text: string,
   model: string,
@@ -198,6 +257,7 @@ async function acquireTransformersPipeline(config: Config): Promise<Transformers
       retired: false,
       idleWaiters: [],
       disposal: undefined,
+      queryEmbeddings: { vectors: new Map(), values: 0 },
     }
     transformersPipelines.set(key, entry)
     void entry.creation.catch(() => {
@@ -214,6 +274,7 @@ async function acquireTransformersPipeline(config: Config): Promise<Transformers
     let released = false
     return {
       extractor,
+      queryEmbeddings: entry.queryEmbeddings,
       async release() {
         if (released) {
           return