refactor(kb): extract packages/kb-retrieval from documents-ingest

Module 5 phase 4 step 0 (docs/kb/modules/05-faza4-plan.md, §3, decision 1):
kb-query is a long-lived Docker service, documents-ingest is a venv job with
an `anthropic` dependency and CLI scripts it doesn't need. Move
embed_chunk/_vector_literal/cascade_query/flat_query into a shared package
with minimal deps (asyncpg, aiohttp only) so both can depend on the same
tested retrieval code without the service image pulling in the job's extras.

documents_ingest.chunk_embed/retrieval keep thin re-exports so nothing
importing the old paths breaks. Pure refactor: retrieval_eval.py run live
against kb-postgres@PIHA + Ollama@SOLARIA before/after gives byte-identical
`dist`/hit@3/gate results (still PASS) — zero behavior change in the cascade.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

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,

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",
 ]

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:<pw>@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,

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).
+"""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
-Two retrieval paths over the same corpus, sharing one query embedding (bge-m3, via Ollama):
+service) share one tested module. Kept here unchanged so nothing importing
-
+`documents_ingest.retrieval` breaks; new code should import `kb_retrieval.retrieval` directly.
-- `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
+from kb_retrieval.retrieval import (
-import asyncpg
+    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
+__all__ = [
-
+    "DEFAULT_EMBED_MODEL",
-DEFAULT_SUMMARY_MODEL = "claude-haiku-4-5"  # plan §2 D3 resolution 2026-07-17: compilation track
+    "DEFAULT_K",
-DEFAULT_EMBED_MODEL = "bge-m3"
+    "DEFAULT_N",
-DEFAULT_N = 10  # plan §6.1 start value
+    "DEFAULT_SUMMARY_MODEL",
-DEFAULT_K = 5  # plan §6.1 start value
+    "cascade_query",
-
+    "cascade_retrieve",
-
+    "flat_query",
-async def flat_retrieve(conn: asyncpg.Connection, query_vector: str, k: int = DEFAULT_K) -> list[dict]:
+    "flat_retrieve",
-    """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

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 (
+from documents_ingest.chunk_embed import EmbeddingDimensionError, EXPECTED_DIM
-    EmbeddingDimensionError,
+from kb_retrieval.embed import _vector_literal, embed_chunk
-    EXPECTED_DIM,
-    _vector_literal,
-    embed_chunk,
-)
 
 _log = structlog.get_logger(__name__)
 

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"]

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",
+]

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

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

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

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,