diff --git a/packages/kb-retrieval/src/kb_retrieval/embed.py b/packages/kb-retrieval/src/kb_retrieval/embed.py index 996ecb0..a2aad28 100644 --- a/packages/kb-retrieval/src/kb_retrieval/embed.py +++ b/packages/kb-retrieval/src/kb_retrieval/embed.py @@ -7,6 +7,12 @@ client without the service image pulling in `jobs/`'s `anthropic` dependency and `check_ollama_health` is new here (not moved) -- a generic `/api/tags` probe kb-query's health-check + circuit breaker (plan §2 decision 2) builds on, kept in this module so the HTTP probing logic lives in exactly one place. + +`embed_batch` -- module 5, faza mailowa, plan Krok 1 (docs/kb/modules/05-faza-mailowa-plan.md, +§4, decision 5): `/api/embed` with `input` as a list, batch 64 measured live at ~8-18 ms/chunk +vs ~150-200 ms/chunk sequential through `embed_chunk`'s `/api/embeddings` (plan §1.4). Used by +`jobs/mail-body-ingest` only -- `embed_chunk` stays the single-text path for kb-query (one query += one text) and the paperless cyclic ingest. """ from __future__ import annotations @@ -16,6 +22,16 @@ import aiohttp DEFAULT_OLLAMA_URL = "http://localhost:11434" DEFAULT_MODEL = "bge-m3" +EXPECTED_DIM = 1024 + + +class EmbeddingDimensionError(RuntimeError): + """Ollama returned a vector of the wrong dimension for the target `document_chunk` schema. + + Abort-run class, mirroring `documents_ingest.chunk_embed.EmbeddingDimensionError` (a separate + class in a separate module -- `embed_chunk`'s single-vector path still raises its own local + one; this one is for `embed_batch` callers, currently only `jobs/mail-body-ingest`). + """ def _vector_literal(embedding: list[float]) -> str: @@ -42,6 +58,42 @@ async def embed_chunk( return embedding, elapsed +async def embed_batch( + session: aiohttp.ClientSession, base_url: str, model: str, texts: list[str] +) -> tuple[list[list[float]], float]: + """POST /api/embed on Ollama with a batch of texts. Returns (embeddings, elapsed_seconds). + + Sequential batches (no HTTP parallelism) -- the GPU saturates on the batch itself (plan + §1.4/§4 decision 5), so parallel requests would only add overhead and unpredictable error + ordering in the caller's stats bilans. + + Raises `EmbeddingDimensionError` (abort-run) if the response's embedding count doesn't match + `len(texts)`, or if any single embedding's dimension isn't `EXPECTED_DIM` -- never silently + indexes a vector that doesn't match the `document_chunk.embedding VECTOR(1024)` column. + Same `aiohttp.ClientError` propagation as `embed_chunk` -- no built-in timeout/retry. + """ + t0 = time.monotonic() + async with session.post(f"{base_url}/api/embed", json={"model": model, "input": texts}) as resp: + resp.raise_for_status() + data = await resp.json() + elapsed = time.monotonic() - t0 + + embeddings = data.get("embeddings") + if embeddings is None: + raise ValueError(f"ollama response missing 'embeddings': {data!r}") + if len(embeddings) != len(texts): + raise EmbeddingDimensionError( + f"ollama returned {len(embeddings)} embeddings for {len(texts)} input texts" + ) + for i, embedding in enumerate(embeddings): + if len(embedding) != EXPECTED_DIM: + raise EmbeddingDimensionError( + f"ollama model={model!r} batch item {i} returned dim={len(embedding)}, " + f"expected {EXPECTED_DIM} (document_chunk.embedding is VECTOR({EXPECTED_DIM}))" + ) + return embeddings, elapsed + + async def check_ollama_health( session: aiohttp.ClientSession, base_url: str, timeout_s: float ) -> bool: diff --git a/packages/kb-retrieval/tests/test_embed.py b/packages/kb-retrieval/tests/test_embed.py index 0526881..adf7cf7 100644 --- a/packages/kb-retrieval/tests/test_embed.py +++ b/packages/kb-retrieval/tests/test_embed.py @@ -3,7 +3,13 @@ from __future__ import annotations import pytest -from kb_retrieval.embed import _vector_literal, check_ollama_health, embed_chunk +from kb_retrieval.embed import ( + EmbeddingDimensionError, + _vector_literal, + check_ollama_health, + embed_batch, + embed_chunk, +) class TestVectorLiteral: @@ -67,6 +73,60 @@ class TestEmbedChunk: await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text") +class _FakeBatchResponse(_FakeEmbedResponse): + pass + + +class _FakeBatchOllamaSession: + """Serves a fixed set of embeddings for every /api/embed POST.""" + + def __init__(self, dim=1024, n_override=None, bad_index=None): + self._dim = dim + self._n_override = n_override + self._bad_index = bad_index + self.requests: list[dict] = [] + + def post(self, url, json): + self.requests.append({"url": url, "json": json}) + texts = json["input"] + n = self._n_override if self._n_override is not None else len(texts) + embeddings = [[0.01] * self._dim for _ in range(n)] + if self._bad_index is not None and self._bad_index < len(embeddings): + embeddings[self._bad_index] = [0.01] * (self._dim - 1) + return _FakeBatchResponse({"embeddings": embeddings}) + + +class TestEmbedBatch: + async def test_returns_embeddings_and_elapsed(self): + session = _FakeBatchOllamaSession(dim=1024) + texts = ["hello", "world", "third"] + embeddings, elapsed = await embed_batch(session, "http://fake-ollama", "bge-m3", texts) + assert len(embeddings) == 3 + assert all(len(e) == 1024 for e in embeddings) + assert elapsed >= 0 + assert session.requests == [ + {"url": "http://fake-ollama/api/embed", "json": {"model": "bge-m3", "input": texts}} + ] + + async def test_length_mismatch_raises_dimension_error(self): + session = _FakeBatchOllamaSession(dim=1024, n_override=2) + with pytest.raises(EmbeddingDimensionError): + await embed_batch(session, "http://fake-ollama", "bge-m3", ["a", "b", "c"]) + + async def test_wrong_dimension_on_one_item_raises(self): + session = _FakeBatchOllamaSession(dim=1024, bad_index=1) + with pytest.raises(EmbeddingDimensionError): + await embed_batch(session, "http://fake-ollama", "bge-m3", ["a", "b", "c"]) + + async def test_missing_embeddings_key_raises_value_error(self): + class _Session: + def post(self, url, json): + return _FakeBatchResponse({}) + + with pytest.raises(ValueError): + await embed_batch(_Session(), "http://fake-ollama", "bge-m3", ["a"]) + + class _FakeTagsResponse: def __init__(self, status): self.status = status