diff --git a/jobs/documents-ingest/eval/retrieval_eval.py b/jobs/documents-ingest/eval/retrieval_eval.py index 588afb2..b2e126c 100644 --- a/jobs/documents-ingest/eval/retrieval_eval.py +++ b/jobs/documents-ingest/eval/retrieval_eval.py @@ -3,7 +3,7 @@ Read-only integration script (NOT collected by pytest -- it hits the live kb-postgres DB and the live Ollama instance, exactly like the plan asked for a separate eval script rather than a mocked test). Runs every query in `queries.yaml` through both the flat baseline -(`documents_ingest.retrieval.flat_query`) and the cascade (`cascade_query`), for a sweep of N +(`kb_retrieval.retrieval.flat_query`) and the cascade (`cascade_query`), for a sweep of N values, and checks the plan's three gate criteria: 1. every query the flat path hits (top-1 dist < 0.45) must still be a hit in the cascade @@ -34,9 +34,11 @@ import aiohttp import asyncpg import yaml +_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) +sys.path.insert(0, str(_REPO_ROOT / "packages" / "kb-retrieval" / "src")) -from documents_ingest.retrieval import ( # noqa: E402 +from kb_retrieval.retrieval import ( # noqa: E402 DEFAULT_EMBED_MODEL, DEFAULT_K, DEFAULT_N, diff --git a/jobs/documents-ingest/pyproject.toml b/jobs/documents-ingest/pyproject.toml index 379ecd7..cfbfbcf 100644 --- a/jobs/documents-ingest/pyproject.toml +++ b/jobs/documents-ingest/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "structlog>=24.1", "aiohttp>=3.9", "kb-mail", + "kb-retrieval", "PyYAML>=6.0", "anthropic>=0.40", ] diff --git a/jobs/documents-ingest/src/documents_ingest/chunk_embed.py b/jobs/documents-ingest/src/documents_ingest/chunk_embed.py index 99d84b6..645ffd7 100644 --- a/jobs/documents-ingest/src/documents_ingest/chunk_embed.py +++ b/jobs/documents-ingest/src/documents_ingest/chunk_embed.py @@ -14,8 +14,14 @@ Runs on SOLARIA (needs Ollama on localhost) against kb-postgres@PIHA over Tailsc Install (from repo root): pip install -e packages/kb-mail/ + pip install -e packages/kb-retrieval/ pip install -e jobs/documents-ingest/ +`embed_chunk`/`_vector_literal`/`DEFAULT_MODEL`/`DEFAULT_OLLAMA_URL` moved to +`kb_retrieval.embed` in module 5 phase 4 (docs/kb/modules/05-faza4-plan.md, §3, decision 1) so +`kb-query` (Docker service) can share the same client without pulling in this job's `anthropic` +dependency; re-exported here unchanged so nothing importing them from this module breaks. + Usage: # Dry run (default) — chunk and count, no Ollama calls, no DB writes: documents-ingest-embed --dsn postgresql://kb:@piha:5433/kb @@ -49,17 +55,15 @@ import json import os import re import sys -import time from typing import Optional import aiohttp import asyncpg import structlog +from kb_retrieval.embed import DEFAULT_MODEL, DEFAULT_OLLAMA_URL, _vector_literal, embed_chunk _log = structlog.get_logger(__name__) -DEFAULT_OLLAMA_URL = "http://localhost:11434" -DEFAULT_MODEL = "bge-m3" EXPECTED_DIM = 1024 # Plan decision 3: ~600 tok/chunk, ~150 tok overlap. No local bge-m3 tokenizer available @@ -204,11 +208,6 @@ def _decode_jsonb(value: object) -> object: return value -def _vector_literal(embedding: list[float]) -> str: - """Render a Python float list as a pgvector input literal, e.g. '[0.1,0.2,...]'.""" - return "[" + ",".join(repr(v) for v in embedding) + "]" - - async def fetch_documents(conn: asyncpg.Connection, limit: Optional[int], offset: Optional[int]) -> list: """`source='paperless'` envelopes, ordered by id for stable --limit/--offset slicing.""" query = "SELECT id, entities FROM envelope WHERE source = 'paperless' ORDER BY id" @@ -230,22 +229,6 @@ async def fetch_existing_chunk_keys(conn: asyncpg.Connection, model: str) -> set return {(r["envelope_id"], r["chunk_index"]) for r in rows} -async def embed_chunk( - session: aiohttp.ClientSession, base_url: str, model: str, text: str -) -> tuple[list[float], float]: - """POST /api/embeddings on Ollama for one chunk. Returns (embedding, elapsed_seconds).""" - t0 = time.monotonic() - async with session.post(f"{base_url}/api/embeddings", json={"model": model, "prompt": text}) as resp: - resp.raise_for_status() - data = await resp.json() - elapsed = time.monotonic() - t0 - - embedding = data.get("embedding") - if not embedding: - raise ValueError(f"ollama response missing 'embedding': {data!r}") - return embedding, elapsed - - async def insert_chunk( conn: asyncpg.Connection, envelope_id: str, chunk_index: int, text: str, embedding: Optional[list[float]], model: str, excluded_reason: Optional[str] = None, diff --git a/jobs/documents-ingest/src/documents_ingest/retrieval.py b/jobs/documents-ingest/src/documents_ingest/retrieval.py index f4320aa..8934076 100644 --- a/jobs/documents-ingest/src/documents_ingest/retrieval.py +++ b/jobs/documents-ingest/src/documents_ingest/retrieval.py @@ -1,135 +1,28 @@ -"""Retrieval module -- module 5, phase 3, plan step 4 (docs/kb/modules/05-faza3-plan.md, §6). - -Two retrieval paths over the same corpus, sharing one query embedding (bge-m3, via Ollama): - -- `flat_retrieve`: baseline -- ranks every active `document_chunk` row directly. This - formalizes the pilot's ad hoc `/tmp/kbq.sh` query (docs/kb/eval/retrieval-pilot-2026-07-16.md) - into a tested, versioned module instead of a script living only in a session transcript. -- `cascade_retrieve`: pre-filters to the top-N `document_summary` matches for one configured - `model` (plan §2 decision 3, resolved 2026-07-17 as D3: `claude-haiku-4-5` is the compilation - track; `gemma3:12b` stays in reserve as the local track) before ranking `document_chunk` - within just those envelopes. At the 186-document pilot scale this doesn't speed anything up - -- it is an architecture test for the mail-scale corpus (225k envelopes, plan §1.1) where a - flat chunk scan stops being cheap. `eval/retrieval_eval.py` runs the quality gate (plan §6.2) - that decides whether it becomes the default path. - -`cascade_query` / `flat_query` are the intended clean entry points for phase 4's kb-query: -query_text -> chunk hits with `dist` (cosine distance) and `source` ("flat" | "cascade"). -`envelope`, `document_chunk`, and `document_summary` are read-only here -- this module only -ever `SELECT`s. +"""Re-export shim -- the retrieval module moved to `packages/kb-retrieval/` in module 5 phase +4 (docs/kb/modules/05-faza4-plan.md, §3, decision 1) so both this job and `kb-query` (Docker +service) share one tested module. Kept here unchanged so nothing importing +`documents_ingest.retrieval` breaks; new code should import `kb_retrieval.retrieval` directly. """ from __future__ import annotations -import aiohttp -import asyncpg +from kb_retrieval.retrieval import ( + DEFAULT_EMBED_MODEL, + DEFAULT_K, + DEFAULT_N, + DEFAULT_SUMMARY_MODEL, + cascade_query, + cascade_retrieve, + flat_query, + flat_retrieve, +) -from documents_ingest.chunk_embed import _vector_literal, embed_chunk - -DEFAULT_SUMMARY_MODEL = "claude-haiku-4-5" # plan §2 D3 resolution 2026-07-17: compilation track -DEFAULT_EMBED_MODEL = "bge-m3" -DEFAULT_N = 10 # plan §6.1 start value -DEFAULT_K = 5 # plan §6.1 start value - - -async def flat_retrieve(conn: asyncpg.Connection, query_vector: str, k: int = DEFAULT_K) -> list[dict]: - """Baseline: rank every active chunk directly against the query embedding, no pre-filter.""" - rows = await conn.fetch( - "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist " - "FROM document_chunk WHERE excluded_reason IS NULL AND embedding IS NOT NULL " - "ORDER BY embedding <=> $1::vector LIMIT $2", - query_vector, k, - ) - return [ - { - "envelope_id": r["envelope_id"], - "chunk_index": r["chunk_index"], - "text": r["text"], - "dist": r["dist"], - "source": "flat", - } - for r in rows - ] - - -async def cascade_retrieve( - conn: asyncpg.Connection, - query_vector: str, - summary_model: str = DEFAULT_SUMMARY_MODEL, - n: int = DEFAULT_N, - k: int = DEFAULT_K, -) -> dict: - """Stage 1: top-N `document_summary` envelopes for `summary_model`. Stage 2: top-k - `document_chunk` ranked within just those envelopes. - - `n` exceeding the number of summarized envelopes is not an error -- plain SQL `LIMIT` - semantics just return all of them. An empty stage 1 (no summaries for this model, or a - corpus that hasn't been summarized yet) short-circuits before stage 2 runs at all: a - cascade can never rank chunks in envelopes it didn't pre-filter into, so there is nothing - for stage 2 to query. - """ - stage1 = await conn.fetch( - "SELECT envelope_id, embedding <=> $1::vector AS dist FROM document_summary " - "WHERE model = $2 AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3", - query_vector, summary_model, n, - ) - stage1_summaries = [{"envelope_id": r["envelope_id"], "dist": r["dist"]} for r in stage1] - - envelope_ids = [s["envelope_id"] for s in stage1_summaries] - if not envelope_ids: - return {"stage1_summaries": stage1_summaries, "chunks": []} - - stage2 = await conn.fetch( - "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist " - "FROM document_chunk WHERE envelope_id = ANY($2::text[]) AND excluded_reason IS NULL " - "AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3", - query_vector, envelope_ids, k, - ) - chunks = [ - { - "envelope_id": r["envelope_id"], - "chunk_index": r["chunk_index"], - "text": r["text"], - "dist": r["dist"], - "source": "cascade", - } - for r in stage2 - ] - return {"stage1_summaries": stage1_summaries, "chunks": chunks} - - -async def flat_query( - conn: asyncpg.Connection, - session: aiohttp.ClientSession, - ollama_url: str, - query_text: str, - embed_model: str = DEFAULT_EMBED_MODEL, - k: int = DEFAULT_K, -) -> dict: - """query_text -> flat chunk hits. One Ollama embed call, one SQL query.""" - embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) - chunks = await flat_retrieve(conn, _vector_literal(embedding), k=k) - return {"query": query_text, "k": k, "chunks": chunks} - - -async def cascade_query( - conn: asyncpg.Connection, - session: aiohttp.ClientSession, - ollama_url: str, - query_text: str, - summary_model: str = DEFAULT_SUMMARY_MODEL, - embed_model: str = DEFAULT_EMBED_MODEL, - n: int = DEFAULT_N, - k: int = DEFAULT_K, -) -> dict: - """query_text -> cascade chunk hits, plus the stage-1 envelope pre-filter (needed for the - gate's diagnosis order, plan §6.2: "N too small" is diagnosed by looking at stage 1). - - Same query embedding as `flat_query`'s single Ollama call -- the cascade's only added - cost over the flat path is one extra SQL query (stage 1), never an extra embed. - """ - embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) - result = await cascade_retrieve(conn, _vector_literal(embedding), summary_model, n=n, k=k) - result["query"] = query_text - result["n"] = n - result["k"] = k - return result +__all__ = [ + "DEFAULT_EMBED_MODEL", + "DEFAULT_K", + "DEFAULT_N", + "DEFAULT_SUMMARY_MODEL", + "cascade_query", + "cascade_retrieve", + "flat_query", + "flat_retrieve", +] diff --git a/jobs/documents-ingest/src/documents_ingest/summarize.py b/jobs/documents-ingest/src/documents_ingest/summarize.py index f41530c..c039ace 100644 --- a/jobs/documents-ingest/src/documents_ingest/summarize.py +++ b/jobs/documents-ingest/src/documents_ingest/summarize.py @@ -10,7 +10,7 @@ Pipeline: `document_chunk.text WHERE excluded_reason IS NULL ORDER BY chunk_inde Two backends write the same table under different `model` values — `UNIQUE (envelope_id, model)` exists precisely so both A/B tracks coexist (plan §2 decision 3). A second, separate mode (`--embed-summaries`) embeds existing summaries with bge-m3 via Ollama, reusing -`chunk_embed.embed_chunk` 1:1 — writing and embedding are split so either can be re-run +`kb_retrieval.embed.embed_chunk` 1:1 — writing and embedding are split so either can be re-run without redoing the other (SOLARIA asleep blocks embedding, not writing; API downtime blocks writing, not embedding). @@ -55,12 +55,8 @@ import structlog import yaml from anthropic import AsyncAnthropic -from documents_ingest.chunk_embed import ( - EmbeddingDimensionError, - EXPECTED_DIM, - _vector_literal, - embed_chunk, -) +from documents_ingest.chunk_embed import EmbeddingDimensionError, EXPECTED_DIM +from kb_retrieval.embed import _vector_literal, embed_chunk _log = structlog.get_logger(__name__) diff --git a/packages/kb-retrieval/pyproject.toml b/packages/kb-retrieval/pyproject.toml new file mode 100644 index 0000000..b433d5e --- /dev/null +++ b/packages/kb-retrieval/pyproject.toml @@ -0,0 +1,25 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "kb-retrieval" +version = "0.1.0" +requires-python = ">=3.11" +dependencies = [ + "asyncpg>=0.29", + "aiohttp>=3.9", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.1", + "pytest-asyncio>=0.23", +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] diff --git a/packages/kb-retrieval/src/kb_retrieval/__init__.py b/packages/kb-retrieval/src/kb_retrieval/__init__.py new file mode 100644 index 0000000..d9c745c --- /dev/null +++ b/packages/kb-retrieval/src/kb_retrieval/__init__.py @@ -0,0 +1,26 @@ +from .embed import DEFAULT_MODEL, DEFAULT_OLLAMA_URL, check_ollama_health, embed_chunk +from .retrieval import ( + DEFAULT_EMBED_MODEL, + DEFAULT_K, + DEFAULT_N, + DEFAULT_SUMMARY_MODEL, + cascade_query, + cascade_retrieve, + flat_query, + flat_retrieve, +) + +__all__ = [ + "DEFAULT_MODEL", + "DEFAULT_OLLAMA_URL", + "check_ollama_health", + "embed_chunk", + "DEFAULT_EMBED_MODEL", + "DEFAULT_K", + "DEFAULT_N", + "DEFAULT_SUMMARY_MODEL", + "cascade_query", + "cascade_retrieve", + "flat_query", + "flat_retrieve", +] diff --git a/packages/kb-retrieval/src/kb_retrieval/embed.py b/packages/kb-retrieval/src/kb_retrieval/embed.py new file mode 100644 index 0000000..996ecb0 --- /dev/null +++ b/packages/kb-retrieval/src/kb_retrieval/embed.py @@ -0,0 +1,57 @@ +"""Ollama embedding client -- module 5, phase 4, plan step 0 (docs/kb/modules/05-faza4-plan.md, +§3, decision 1). Moved 1:1 out of `jobs/documents_ingest/chunk_embed.py` (`embed_chunk`, +`_vector_literal`, `DEFAULT_MODEL`, `DEFAULT_OLLAMA_URL`) so both a venv-based host job +(`documents-ingest`) and a long-lived Docker service (`kb-query`) can depend on the same tested +client without the service image pulling in `jobs/`'s `anthropic` dependency and CLI scripts. + +`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. +""" +from __future__ import annotations + +import time + +import aiohttp + +DEFAULT_OLLAMA_URL = "http://localhost:11434" +DEFAULT_MODEL = "bge-m3" + + +def _vector_literal(embedding: list[float]) -> str: + """Render a Python float list as a pgvector input literal, e.g. '[0.1,0.2,...]'.""" + return "[" + ",".join(repr(v) for v in embedding) + "]" + + +async def embed_chunk( + session: aiohttp.ClientSession, base_url: str, model: str, text: str +) -> tuple[list[float], float]: + """POST /api/embeddings on Ollama for one chunk. Returns (embedding, elapsed_seconds). + No built-in timeout/retry -- inherits whatever `aiohttp.ClientSession(timeout=...)` the + caller constructed; `raise_for_status()` propagates `aiohttp.ClientError` when Ollama is + unreachable.""" + t0 = time.monotonic() + async with session.post(f"{base_url}/api/embeddings", json={"model": model, "prompt": text}) as resp: + resp.raise_for_status() + data = await resp.json() + elapsed = time.monotonic() - t0 + + embedding = data.get("embedding") + if not embedding: + raise ValueError(f"ollama response missing 'embedding': {data!r}") + return embedding, elapsed + + +async def check_ollama_health( + session: aiohttp.ClientSession, base_url: str, timeout_s: float +) -> bool: + """Probe `GET {base_url}/api/tags`. True iff it answers within `timeout_s` with a non-error + status; any exception (timeout, connection refused, 4xx/5xx) is treated as `down`, never + raised -- callers use this for liveness decisions, not error propagation.""" + try: + async with session.get( + f"{base_url}/api/tags", timeout=aiohttp.ClientTimeout(total=timeout_s) + ) as resp: + return resp.status < 400 + except (aiohttp.ClientError, TimeoutError): + return False diff --git a/packages/kb-retrieval/src/kb_retrieval/retrieval.py b/packages/kb-retrieval/src/kb_retrieval/retrieval.py new file mode 100644 index 0000000..ee61f33 --- /dev/null +++ b/packages/kb-retrieval/src/kb_retrieval/retrieval.py @@ -0,0 +1,138 @@ +"""Retrieval module -- module 5, phase 3, plan step 4 (docs/kb/modules/05-faza3-plan.md, §6). +Moved 1:1 into `packages/kb-retrieval` in phase 4 (docs/kb/modules/05-faza4-plan.md, §3, +decision 1) so both `documents-ingest` (venv job) and `kb-query` (Docker service) share one +tested module instead of the service image needing to pull in all of `jobs/documents-ingest`. + +Two retrieval paths over the same corpus, sharing one query embedding (bge-m3, via Ollama): + +- `flat_retrieve`: baseline -- ranks every active `document_chunk` row directly. This + formalizes the pilot's ad hoc `/tmp/kbq.sh` query (docs/kb/eval/retrieval-pilot-2026-07-16.md) + into a tested, versioned module instead of a script living only in a session transcript. +- `cascade_retrieve`: pre-filters to the top-N `document_summary` matches for one configured + `model` (plan §2 decision 3, resolved 2026-07-17 as D3: `claude-haiku-4-5` is the compilation + track; `gemma3:12b` stays in reserve as the local track) before ranking `document_chunk` + within just those envelopes. At the 186-document pilot scale this doesn't speed anything up + -- it is an architecture test for the mail-scale corpus (225k envelopes, plan §1.1) where a + flat chunk scan stops being cheap. `eval/retrieval_eval.py` runs the quality gate (plan §6.2) + that decides whether it becomes the default path. + +`cascade_query` / `flat_query` are the intended clean entry points for phase 4's kb-query: +query_text -> chunk hits with `dist` (cosine distance) and `source` ("flat" | "cascade"). +`envelope`, `document_chunk`, and `document_summary` are read-only here -- this module only +ever `SELECT`s. +""" +from __future__ import annotations + +import aiohttp +import asyncpg + +from kb_retrieval.embed import _vector_literal, embed_chunk + +DEFAULT_SUMMARY_MODEL = "claude-haiku-4-5" # plan §2 D3 resolution 2026-07-17: compilation track +DEFAULT_EMBED_MODEL = "bge-m3" +DEFAULT_N = 10 # plan §6.1 start value +DEFAULT_K = 5 # plan §6.1 start value + + +async def flat_retrieve(conn: asyncpg.Connection, query_vector: str, k: int = DEFAULT_K) -> list[dict]: + """Baseline: rank every active chunk directly against the query embedding, no pre-filter.""" + rows = await conn.fetch( + "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist " + "FROM document_chunk WHERE excluded_reason IS NULL AND embedding IS NOT NULL " + "ORDER BY embedding <=> $1::vector LIMIT $2", + query_vector, k, + ) + return [ + { + "envelope_id": r["envelope_id"], + "chunk_index": r["chunk_index"], + "text": r["text"], + "dist": r["dist"], + "source": "flat", + } + for r in rows + ] + + +async def cascade_retrieve( + conn: asyncpg.Connection, + query_vector: str, + summary_model: str = DEFAULT_SUMMARY_MODEL, + n: int = DEFAULT_N, + k: int = DEFAULT_K, +) -> dict: + """Stage 1: top-N `document_summary` envelopes for `summary_model`. Stage 2: top-k + `document_chunk` ranked within just those envelopes. + + `n` exceeding the number of summarized envelopes is not an error -- plain SQL `LIMIT` + semantics just return all of them. An empty stage 1 (no summaries for this model, or a + corpus that hasn't been summarized yet) short-circuits before stage 2 runs at all: a + cascade can never rank chunks in envelopes it didn't pre-filter into, so there is nothing + for stage 2 to query. + """ + stage1 = await conn.fetch( + "SELECT envelope_id, embedding <=> $1::vector AS dist FROM document_summary " + "WHERE model = $2 AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3", + query_vector, summary_model, n, + ) + stage1_summaries = [{"envelope_id": r["envelope_id"], "dist": r["dist"]} for r in stage1] + + envelope_ids = [s["envelope_id"] for s in stage1_summaries] + if not envelope_ids: + return {"stage1_summaries": stage1_summaries, "chunks": []} + + stage2 = await conn.fetch( + "SELECT envelope_id, chunk_index, text, embedding <=> $1::vector AS dist " + "FROM document_chunk WHERE envelope_id = ANY($2::text[]) AND excluded_reason IS NULL " + "AND embedding IS NOT NULL ORDER BY embedding <=> $1::vector LIMIT $3", + query_vector, envelope_ids, k, + ) + chunks = [ + { + "envelope_id": r["envelope_id"], + "chunk_index": r["chunk_index"], + "text": r["text"], + "dist": r["dist"], + "source": "cascade", + } + for r in stage2 + ] + return {"stage1_summaries": stage1_summaries, "chunks": chunks} + + +async def flat_query( + conn: asyncpg.Connection, + session: aiohttp.ClientSession, + ollama_url: str, + query_text: str, + embed_model: str = DEFAULT_EMBED_MODEL, + k: int = DEFAULT_K, +) -> dict: + """query_text -> flat chunk hits. One Ollama embed call, one SQL query.""" + embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) + chunks = await flat_retrieve(conn, _vector_literal(embedding), k=k) + return {"query": query_text, "k": k, "chunks": chunks} + + +async def cascade_query( + conn: asyncpg.Connection, + session: aiohttp.ClientSession, + ollama_url: str, + query_text: str, + summary_model: str = DEFAULT_SUMMARY_MODEL, + embed_model: str = DEFAULT_EMBED_MODEL, + n: int = DEFAULT_N, + k: int = DEFAULT_K, +) -> dict: + """query_text -> cascade chunk hits, plus the stage-1 envelope pre-filter (needed for the + gate's diagnosis order, plan §6.2: "N too small" is diagnosed by looking at stage 1). + + Same query embedding as `flat_query`'s single Ollama call -- the cascade's only added + cost over the flat path is one extra SQL query (stage 1), never an extra embed. + """ + embedding, _elapsed = await embed_chunk(session, ollama_url, embed_model, query_text) + result = await cascade_retrieve(conn, _vector_literal(embedding), summary_model, n=n, k=k) + result["query"] = query_text + result["n"] = n + result["k"] = k + return result diff --git a/packages/kb-retrieval/tests/test_embed.py b/packages/kb-retrieval/tests/test_embed.py new file mode 100644 index 0000000..0526881 --- /dev/null +++ b/packages/kb-retrieval/tests/test_embed.py @@ -0,0 +1,109 @@ +"""Unit tests for the Ollama embedding client -- no real HTTP, no real Ollama.""" +from __future__ import annotations + +import pytest + +from kb_retrieval.embed import _vector_literal, check_ollama_health, embed_chunk + + +class TestVectorLiteral: + def test_formats_as_bracketed_csv(self): + assert _vector_literal([0.1, 0.2, -0.3]) == "[0.1,0.2,-0.3]" + + +class _FakeEmbedResponse: + def __init__(self, payload, status=200): + self._payload = payload + self._status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def json(self): + return self._payload + + def raise_for_status(self): + if self._status >= 400: + raise RuntimeError(f"HTTP {self._status}") + + +class _FakeOllamaSession: + """Serves a fixed embedding vector for every /api/embeddings POST, or errors by prompt.""" + + def __init__(self, dim=1024, fail_for=None): + self._dim = dim + self._fail_for = fail_for or set() + self.requests: list[dict] = [] + + def post(self, url, json): + self.requests.append({"url": url, "json": json}) + if json["prompt"] in self._fail_for: + return _FakeEmbedResponse({}, status=500) + return _FakeEmbedResponse({"embedding": [0.01] * self._dim}) + + +class TestEmbedChunk: + async def test_returns_embedding_and_elapsed(self): + session = _FakeOllamaSession(dim=1024) + embedding, elapsed = await embed_chunk(session, "http://fake-ollama", "bge-m3", "hello world") + assert len(embedding) == 1024 + assert elapsed >= 0 + assert session.requests == [ + {"url": "http://fake-ollama/api/embeddings", "json": {"model": "bge-m3", "prompt": "hello world"}} + ] + + async def test_missing_embedding_key_raises(self): + class _EmptyResponse(_FakeEmbedResponse): + pass + + class _Session: + def post(self, url, json): + return _EmptyResponse({}) + + with pytest.raises(ValueError): + await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text") + + +class _FakeTagsResponse: + def __init__(self, status): + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + +class _FakeHealthSession: + def __init__(self, status=200, raises=None): + self._status = status + self._raises = raises + + def get(self, url, timeout=None): + if self._raises is not None: + raise self._raises + return _FakeTagsResponse(self._status) + + +class TestCheckOllamaHealth: + async def test_up_on_2xx(self): + session = _FakeHealthSession(status=200) + assert await check_ollama_health(session, "http://fake-ollama", 0.5) is True + + async def test_down_on_error_status(self): + session = _FakeHealthSession(status=500) + assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False + + async def test_down_on_connection_error(self): + import aiohttp + + session = _FakeHealthSession(raises=aiohttp.ClientConnectionError("refused")) + assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False + + async def test_down_on_timeout(self): + session = _FakeHealthSession(raises=TimeoutError()) + assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False diff --git a/jobs/documents-ingest/tests/test_retrieval.py b/packages/kb-retrieval/tests/test_retrieval.py similarity index 99% rename from jobs/documents-ingest/tests/test_retrieval.py rename to packages/kb-retrieval/tests/test_retrieval.py index 8cccaaf..ae21819 100644 --- a/jobs/documents-ingest/tests/test_retrieval.py +++ b/packages/kb-retrieval/tests/test_retrieval.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest -from documents_ingest.retrieval import ( +from kb_retrieval.retrieval import ( DEFAULT_K, DEFAULT_N, cascade_query,