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>
This commit is contained in:
oskar
parent
88553d1631
commit
f964631d02
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ dependencies = [
|
|||
"structlog>=24.1",
|
||||
"aiohttp>=3.9",
|
||||
"kb-mail",
|
||||
"kb-retrieval",
|
||||
"PyYAML>=6.0",
|
||||
"anthropic>=0.40",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 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,
|
||||
from kb_retrieval.retrieval import (
|
||||
DEFAULT_EMBED_MODEL,
|
||||
DEFAULT_K,
|
||||
DEFAULT_N,
|
||||
DEFAULT_SUMMARY_MODEL,
|
||||
cascade_query,
|
||||
cascade_retrieve,
|
||||
flat_query,
|
||||
flat_retrieve,
|
||||
)
|
||||
return [
|
||||
{
|
||||
"envelope_id": r["envelope_id"],
|
||||
"chunk_index": r["chunk_index"],
|
||||
"text": r["text"],
|
||||
"dist": r["dist"],
|
||||
"source": "flat",
|
||||
}
|
||||
for r in rows
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_EMBED_MODEL",
|
||||
"DEFAULT_K",
|
||||
"DEFAULT_N",
|
||||
"DEFAULT_SUMMARY_MODEL",
|
||||
"cascade_query",
|
||||
"cascade_retrieve",
|
||||
"flat_query",
|
||||
"flat_retrieve",
|
||||
]
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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__)
|
||||
|
||||
|
|
|
|||
25
packages/kb-retrieval/pyproject.toml
Normal file
25
packages/kb-retrieval/pyproject.toml
Normal file
|
|
@ -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"]
|
||||
26
packages/kb-retrieval/src/kb_retrieval/__init__.py
Normal file
26
packages/kb-retrieval/src/kb_retrieval/__init__.py
Normal file
|
|
@ -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",
|
||||
]
|
||||
57
packages/kb-retrieval/src/kb_retrieval/embed.py
Normal file
57
packages/kb-retrieval/src/kb_retrieval/embed.py
Normal file
|
|
@ -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
|
||||
138
packages/kb-retrieval/src/kb_retrieval/retrieval.py
Normal file
138
packages/kb-retrieval/src/kb_retrieval/retrieval.py
Normal file
|
|
@ -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
|
||||
109
packages/kb-retrieval/tests/test_embed.py
Normal file
109
packages/kb-retrieval/tests/test_embed.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
Loading…
Reference in a new issue