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.
This commit is contained in:
2026-07-03 13:16:05 +04:00
parent a0b3a1865a
commit 6437d29914
5 changed files with 102 additions and 30 deletions
@@ -116,6 +116,15 @@ class LlamaCppEmbedder(
return arrayResult
}
// llama.cpp --pooling mean: [{"embedding": [[...]]}] (nested list-of-one-vector).
val nestedResult = runCatching {
json.decodeFromString<List<EmbeddingNestedArrayResponse>>(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)}"
@@ -19,6 +19,13 @@ data class EmbeddingArrayResponse(
val embedding: List<Float>,
)
// 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<List<Float>>,
)
@Serializable
data class OpenAiEmbeddingData(
val embedding: List<Float>,
@@ -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 ->
@@ -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"}")
}
}
}
@@ -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)}