From 6437d29914316a71e13881b8fc9f82913b3d4bf2 Mon Sep 17 00:00:00 2001 From: kami Date: Fri, 3 Jul 2026 13:16:05 +0400 Subject: [PATCH] fix(router): make the real embedder + TurboVec L3 actually work end-to-end - LlamaCppEmbedder: parse llama.cpp --pooling mean nested response [[...]] (was silently failing every embed with 'unexpected response shape') - qa-stack: embedder -b/-ub 8192 so docs >512 tokens don't 500; fix default path - TurboVecL3MemoryStore: send init handshake on startup (was 'Index not initialized' on every add/query) - turbovec_sidecar.py: rewrite add/search/save/load against the real turbovec API (batched ndarray, L2-normalize so IP=cosine, prepare() before search, classmethod load). The sidecar was scaffolding never run against the real lib. --- .../inference/llama/cpp/LlamaCppEmbedder.kt | 9 +++ .../llama/cpp/LlamaCppEmbeddingApiModels.kt | 7 ++ .../llama/cpp/LlamaCppEmbedderTest.kt | 22 ++++++ .../router/turbovec/TurboVecL3MemoryStore.kt | 17 +++- .../main/resources/python/turbovec_sidecar.py | 77 ++++++++++++------- 5 files changed, 102 insertions(+), 30 deletions(-) diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedder.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedder.kt index e621eda2..9e581029 100644 --- a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedder.kt +++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedder.kt @@ -116,6 +116,15 @@ class LlamaCppEmbedder( return arrayResult } + // llama.cpp --pooling mean: [{"embedding": [[...]]}] (nested list-of-one-vector). + val nestedResult = runCatching { + json.decodeFromString>(responseBody) + .firstOrNull()?.embedding?.firstOrNull().orEmpty() + }.getOrNull() + if (nestedResult != null && nestedResult.isNotEmpty()) { + return nestedResult + } + throw IllegalStateException( "Unexpected embedding response shape from $url. " + "Response body (first 200 chars): ${responseBody.take(200)}" diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbeddingApiModels.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbeddingApiModels.kt index d90707e6..f8742e3b 100644 --- a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbeddingApiModels.kt +++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbeddingApiModels.kt @@ -19,6 +19,13 @@ data class EmbeddingArrayResponse( val embedding: List, ) +// llama.cpp /embedding with --pooling mean returns a list-of-one-vector per item: +// [{"index":0,"embedding":[[floats...]]}]. embedding is nested, not flat. +@Serializable +data class EmbeddingNestedArrayResponse( + val embedding: List>, +) + @Serializable data class OpenAiEmbeddingData( val embedding: List, diff --git a/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedderTest.kt b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedderTest.kt index 60856bce..f1f1b8df 100644 --- a/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedderTest.kt +++ b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppEmbedderTest.kt @@ -79,6 +79,28 @@ class LlamaCppEmbedderTest { assert(kotlin.math.abs(result[2] - 0.7f) < 0.0001f) } + @Test + fun `embed parses llama-cpp nested pooling-mean array format`(): Unit = runBlocking { + // llama.cpp /embedding with --pooling mean: [{"index":0,"embedding":[[...]]}] + val mockEngine = MockEngine { + respond( + content = """[{"index":0,"embedding":[[0.5,0.6,0.7]]}]""", + headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()) + ) + } + val embedder = LlamaCppEmbedder( + url = "http://localhost:11000", + dimension = 3, + httpClient = HttpClient(mockEngine), + ) + + val result = embedder.embed("test text") + + assert(result.size == 3) + assert(kotlin.math.abs(result[0] - 0.5f) < 0.0001f) + assert(kotlin.math.abs(result[2] - 0.7f) < 0.0001f) + } + @Test fun `embed throws on dimension mismatch`(): Unit = runBlocking { val mockEngine = MockEngine { request -> diff --git a/infrastructure/router/turbovec/src/main/kotlin/com/correx/infrastructure/router/turbovec/TurboVecL3MemoryStore.kt b/infrastructure/router/turbovec/src/main/kotlin/com/correx/infrastructure/router/turbovec/TurboVecL3MemoryStore.kt index fafe978d..eda5eeda 100644 --- a/infrastructure/router/turbovec/src/main/kotlin/com/correx/infrastructure/router/turbovec/TurboVecL3MemoryStore.kt +++ b/infrastructure/router/turbovec/src/main/kotlin/com/correx/infrastructure/router/turbovec/TurboVecL3MemoryStore.kt @@ -120,6 +120,21 @@ class TurboVecL3MemoryStore(private val config: TurboVecSidecarConfig) : Rehydra val pb = ProcessBuilder(config.pythonExecutable, config.scriptPath.toString()) pb.redirectError(logFile) - process = pb.start() + val started = pb.start() + process = started + + // The sidecar rejects add/search until the index is created. Send init once on startup; + // without this every store/query fails with "Index not initialized". + val initLine = json.encodeToString( + SidecarRequest.serializer(), + SidecarRequest(op = "init", dim = config.dim, bitWidth = config.bitWidth) + ) + "\n" + started.outputStream.write(initLine.toByteArray()) + started.outputStream.flush() + val response = started.inputStream.bufferedReader().readLine() + ?.let { json.decodeFromString(SidecarResponse.serializer(), it) } + if (response?.ok != true) { + throw RuntimeException("TurboVec init failed: ${response?.error ?: "no response"}") + } } } diff --git a/infrastructure/router/turbovec/src/main/resources/python/turbovec_sidecar.py b/infrastructure/router/turbovec/src/main/resources/python/turbovec_sidecar.py index dcdcc62f..eceabf4d 100644 --- a/infrastructure/router/turbovec/src/main/resources/python/turbovec_sidecar.py +++ b/infrastructure/router/turbovec/src/main/resources/python/turbovec_sidecar.py @@ -19,17 +19,17 @@ Requires: pip install turbovec import sys import json +import numpy as np from turbovec import TurboQuantIndex -def cosine_similarity(a, b): - """Compute cosine similarity between two vectors.""" - dot_product = sum(x * y for x, y in zip(a, b)) - norm_a = sum(x * x for x in a) ** 0.5 - norm_b = sum(x * x for x in b) ** 0.5 - if norm_a == 0 or norm_b == 0: - return 0.0 - return dot_product / (norm_a * norm_b) +def _normalize(vector): + """L2-normalize to a (1, dim) float32 row so TurboVec's inner-product search + behaves as cosine similarity (matching the in-memory L3 store).""" + a = np.asarray(vector, dtype=np.float32).reshape(1, -1) + n = np.linalg.norm(a, axis=-1, keepdims=True) + n[n == 0] = 1.0 + return (a / n).astype(np.float32) class TurboVecSidecar: @@ -37,57 +37,70 @@ class TurboVecSidecar: self.index = None self.id_to_idx = {} # id (str) -> index position self.idx_to_id = [] # index position -> id (str) + self.dim = None + self.bit_width = 4 + self.dirty = False # True when vectors added since last prepare() def init(self, dim, bit_width): """Initialize a new TurboQuantIndex.""" try: self.index = TurboQuantIndex(dim=dim, bit_width=bit_width) + self.dim = dim + self.bit_width = bit_width self.id_to_idx = {} self.idx_to_id = [] + self.dirty = False return {"ok": True} except Exception as e: return {"ok": False, "error": str(e)} def add(self, entry_id, vector): - """Add a vector to the index.""" + """Add one vector. TurboVec expects a (1, dim) float32 batch, normalized.""" try: if self.index is None: return {"ok": False, "error": "Index not initialized"} - idx = len(self.idx_to_id) - self.index.add(vector) - self.id_to_idx[entry_id] = idx + self.index.add(_normalize(vector)) + self.id_to_idx[entry_id] = len(self.idx_to_id) self.idx_to_id.append(entry_id) + self.dirty = True return {"ok": True} except Exception as e: return {"ok": False, "error": str(e)} def search(self, vector, k): - """Search for top-k nearest neighbors.""" + """Top-k inner-product (≈cosine) search via the quantized index.""" try: if self.index is None: return {"ok": False, "error": "Index not initialized"} - - # Compute distances to all indexed vectors using cosine similarity - hits = [] - for idx, stored_vector in enumerate(self.index.vectors): - score = cosine_similarity(vector, stored_vector) - entry_id = self.idx_to_id[idx] - hits.append({"id": entry_id, "score": score}) - - # Sort by score descending, take top k - hits = sorted(hits, key=lambda x: x["score"], reverse=True)[:k] + if len(self.idx_to_id) == 0: + return {"ok": True, "hits": []} + if self.dirty: + self.index.prepare() # warm caches; required after adds + self.dirty = False + scores, indices = self.index.search(_normalize(vector), k) + hits = [ + {"id": self.idx_to_id[int(i)], "score": float(s)} + for s, i in zip(scores[0], indices[0]) + ] return {"ok": True, "hits": hits} except Exception as e: return {"ok": False, "error": str(e)} def save(self, path): - """Persist the index to disk.""" + """Persist the index + id mapping to disk.""" try: if self.index is None: return {"ok": False, "error": "Index not initialized"} + if self.dirty: + self.index.prepare() + self.dirty = False self.index.write(path) - # Also save metadata - metadata = {"id_to_idx": self.id_to_idx, "idx_to_id": self.idx_to_id} + metadata = { + "id_to_idx": self.id_to_idx, + "idx_to_id": self.idx_to_id, + "dim": self.dim, + "bit_width": self.bit_width, + } with open(f"{path}.meta.json", "w") as f: json.dump(metadata, f) return {"ok": True} @@ -95,10 +108,16 @@ class TurboVecSidecar: return {"ok": False, "error": str(e)} def load(self, path): - """Load the index from disk.""" + """Load a previously-saved index + id mapping.""" try: - # SCAFFOLDING: full load not yet implemented - # For now, return success but do not actually load + self.index = TurboQuantIndex.load(path) + with open(f"{path}.meta.json") as f: + metadata = json.load(f) + self.id_to_idx = metadata.get("id_to_idx", {}) + self.idx_to_id = metadata.get("idx_to_id", []) + self.dim = metadata.get("dim") + self.bit_width = metadata.get("bit_width", 4) + self.dirty = True # re-prepare on first search after load return {"ok": True} except Exception as e: return {"ok": False, "error": str(e)}