feat(kb-retrieval): add embed_batch for /api/embed (faza mailowa Krok 1)
Batched embedding client (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) for jobs/mail-body-ingest -- embed_chunk stays the single-text path for kb-query and paperless cyclic ingest. Validates the response embedding count and each item's dimension (1024), raising EmbeddingDimensionError (abort-run) on mismatch so a bad batch never silently writes a mismatched vector (docs/kb/modules/05-faza-mailowa-plan.md, §4, decision 5).
This commit is contained in:
parent
348ce10b1d
commit
51998fd891
|
|
@ -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
|
`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
|
health-check + circuit breaker (plan §2 decision 2) builds on, kept in this module so the HTTP
|
||||||
probing logic lives in exactly one place.
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
|
|
@ -16,6 +22,16 @@ import aiohttp
|
||||||
|
|
||||||
DEFAULT_OLLAMA_URL = "http://localhost:11434"
|
DEFAULT_OLLAMA_URL = "http://localhost:11434"
|
||||||
DEFAULT_MODEL = "bge-m3"
|
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:
|
def _vector_literal(embedding: list[float]) -> str:
|
||||||
|
|
@ -42,6 +58,42 @@ async def embed_chunk(
|
||||||
return embedding, elapsed
|
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(
|
async def check_ollama_health(
|
||||||
session: aiohttp.ClientSession, base_url: str, timeout_s: float
|
session: aiohttp.ClientSession, base_url: str, timeout_s: float
|
||||||
) -> bool:
|
) -> bool:
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,13 @@ from __future__ import annotations
|
||||||
|
|
||||||
import pytest
|
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:
|
class TestVectorLiteral:
|
||||||
|
|
@ -67,6 +73,60 @@ class TestEmbedChunk:
|
||||||
await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text")
|
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:
|
class _FakeTagsResponse:
|
||||||
def __init__(self, status):
|
def __init__(self, status):
|
||||||
self.status = status
|
self.status = status
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue