f964631d02
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>
110 lines
3.4 KiB
Python
110 lines
3.4 KiB
Python
"""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
|