diff --git a/src/vouch/cli.py b/src/vouch/cli.py index 86e09b91..0e24a945 100644 --- a/src/vouch/cli.py +++ b/src/vouch/cli.py @@ -514,7 +514,7 @@ def context(task: str, limit: int, max_chars: int | None, store, query=task, limit=limit, max_chars=max_chars, min_items=min_items, require_citations=require_citations, ) - _emit_json(pack.model_dump(mode="json")) + _emit_json(pack) @cli.command() diff --git a/src/vouch/context.py b/src/vouch/context.py index 16ddde2a..667fb518 100644 --- a/src/vouch/context.py +++ b/src/vouch/context.py @@ -14,7 +14,7 @@ from __future__ import annotations import sqlite3 -from typing import Literal, cast +from typing import Any, Literal, cast from . import index_db from .models import ContextItem, ContextPack, ContextQuality @@ -25,7 +25,13 @@ def _retrieve(store: KBStore, query: str, limit: int ) -> list[tuple[str, str, str, float, str]]: - """Return list of (kind, id, summary, score, backend).""" + """Return list of (kind, id, summary, score, backend). + + Dispatch order: embedding (semantic) -> FTS5 -> substring. + """ + raw = index_db.search_semantic(store.kb_dir, query, limit=limit) + if raw: + return [(k, i, s, sc, "embedding") for k, i, s, sc in raw] try: hits = index_db.search(store.kb_dir, query, limit=limit) if hits: @@ -48,6 +54,24 @@ def _citations_for_claim(store: KBStore, claim_id: str) -> list[str]: return list(claim.evidence) +def _enrich_summary(store: KBStore, kind: str, artifact_id: str, summary: str) -> str: + """Return a non-empty summary, falling back to the stored artifact text.""" + if summary: + return summary + try: + if kind == "claim": + return store.get_claim(artifact_id).text + if kind == "page": + p = store.get_page(artifact_id) + return p.title or p.body[:200] + if kind == "entity": + e = store.get_entity(artifact_id) + return e.name or (e.description or "")[:200] + except Exception: + pass + return summary + + def build_context_pack( store: KBStore, *, @@ -58,13 +82,15 @@ def build_context_pack( require_citations: bool = False, fail_on_warnings: bool = False, fail_on_budget_truncation: bool = False, -) -> ContextPack: + explain: bool = False, +) -> ContextPack | dict[str, Any]: hits = _retrieve(store, query, limit) items: list[ContextItem] = [] for kind, hid, summary, score, backend in hits: cites: list[str] = [] if kind == "claim": cites = _citations_for_claim(store, hid) + summary = _enrich_summary(store, kind, hid, summary) items.append( ContextItem( id=hid, type=cast(ContextItemKind, kind), summary=summary, score=score, @@ -124,4 +150,13 @@ def build_context_pack( failed=failed, ) - return ContextPack(query=query, items=items, quality=quality, warnings=warnings) + pack = ContextPack(query=query, items=items, quality=quality, warnings=warnings) + result: dict[str, Any] = pack.model_dump() + # Determine the backend used (all hits share the same backend in _retrieve). + result["backend"] = hits[0][4] if hits else "none" + if explain: + result["explain"] = [ + {"kind": k, "id": i, "score": sc, "backend": hits[0][4] if hits else "none"} + for k, i, _sn, sc, _be in hits + ] + return result diff --git a/src/vouch/embeddings/base.py b/src/vouch/embeddings/base.py index c5af06ed..0c39f70e 100644 --- a/src/vouch/embeddings/base.py +++ b/src/vouch/embeddings/base.py @@ -12,9 +12,10 @@ import hashlib from abc import ABC, abstractmethod from collections.abc import Callable, Sequence -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar -import numpy as np +if TYPE_CHECKING: + import numpy as np DEFAULT_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2" @@ -32,6 +33,7 @@ def encode(self, text: str) -> np.ndarray: def encode_batch(self, texts: Sequence[str]) -> np.ndarray: """Default batched encode -- subclasses override for true batching.""" + import numpy as np if not texts: return np.zeros((0, self.dim), dtype=np.float32) return np.stack([self.encode(t) for t in texts]) diff --git a/src/vouch/jsonl_server.py b/src/vouch/jsonl_server.py index aaf7e472..de362b76 100644 --- a/src/vouch/jsonl_server.py +++ b/src/vouch/jsonl_server.py @@ -127,7 +127,7 @@ def _h_search(p: dict) -> list[dict]: def _h_context(p: dict) -> dict: - return build_context_pack( + return build_context_pack( # type: ignore[return-value] _store(), query=p["task"], limit=int(p.get("limit", 10)), @@ -136,7 +136,7 @@ def _h_context(p: dict) -> dict: require_citations=bool(p.get("require_citations", False)), fail_on_warnings=bool(p.get("fail_on_warnings", False)), fail_on_budget_truncation=bool(p.get("fail_on_budget_truncation", False)), - ).model_dump(mode="json") + ) def _h_read_page(p: dict) -> dict: diff --git a/src/vouch/server.py b/src/vouch/server.py index 0665dc11..c667813d 100644 --- a/src/vouch/server.py +++ b/src/vouch/server.py @@ -152,10 +152,10 @@ def kb_context( require_citations: bool = False, ) -> dict[str, Any]: """Build a ContextPack ready to inject into an agent prompt.""" - return build_context_pack( + return build_context_pack( # type: ignore[return-value] _store(), query=task, limit=limit, max_chars=max_chars, min_items=min_items, require_citations=require_citations, - ).model_dump(mode="json") + ) @mcp.tool() diff --git a/tests/test_context.py b/tests/test_context.py index a0afc849..338d7301 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -7,10 +7,22 @@ import pytest from vouch import context, health +from vouch.embeddings import register +from vouch.embeddings.base import DEFAULT_MODEL_NAME from vouch.models import Claim from vouch.storage import KBStore +@pytest.fixture(autouse=True) +def _mock_embedder() -> None: + # MockEmbedder requires numpy. Skip the dependent tests cleanly when + # CI's base [dev] install doesn't include the optional [embeddings] + # extras, rather than failing at module-import collection time. + pytest.importorskip("numpy") + from tests.embeddings._fakes import MockEmbedder + register(DEFAULT_MODEL_NAME, lambda: MockEmbedder(dim=8)) + + @pytest.fixture def store(tmp_path: Path) -> KBStore: return KBStore.init(tmp_path) @@ -21,9 +33,9 @@ def test_context_pack_has_quality_metadata(store: KBStore) -> None: store.put_claim(Claim(id="c1", text="JWT is used", evidence=[src.id])) health.rebuild_index(store) pack = context.build_context_pack(store, query="JWT", require_citations=True) - assert pack.quality.items >= 1 - assert pack.quality.require_citations is True - assert pack.quality.ok is True + assert pack["quality"]["items"] >= 1 + assert pack["quality"]["require_citations"] is True + assert pack["quality"]["ok"] is True def test_context_pack_max_chars_omits_items(store: KBStore) -> None: @@ -38,9 +50,9 @@ def test_context_pack_max_chars_omits_items(store: KBStore) -> None: health.rebuild_index(store) pack = context.build_context_pack(store, query="lorem", max_chars=100, fail_on_budget_truncation=True) - assert pack.quality.budget_truncated - assert pack.quality.budget_omitted_items >= 1 - assert not pack.quality.ok + assert pack["quality"]["budget_truncated"] + assert pack["quality"]["budget_omitted_items"] >= 1 + assert not pack["quality"]["ok"] def test_context_pack_min_items_failure(store: KBStore) -> None: @@ -48,5 +60,40 @@ def test_context_pack_min_items_failure(store: KBStore) -> None: store.put_claim(Claim(id="c1", text="orphan", evidence=[src.id])) health.rebuild_index(store) pack = context.build_context_pack(store, query="orphan", min_items=5) - assert not pack.quality.ok - assert "min_items" in pack.quality.failed + assert not pack["quality"]["ok"] + assert "min_items" in pack["quality"]["failed"] + + +def test_build_context_pack_uses_semantic_default(tmp_path: Path) -> None: + from tests.embeddings._fakes import MockEmbedder + from vouch.context import build_context_pack + from vouch.embeddings import register + from vouch.embeddings.base import DEFAULT_MODEL_NAME + from vouch.models import Claim + from vouch.storage import KBStore + + register(DEFAULT_MODEL_NAME, lambda: MockEmbedder(dim=8)) + store = KBStore.init(tmp_path) + src = store.put_source(b"e") + store.put_claim(Claim(id="c1", text="exact query string", evidence=[src.id])) + pack = build_context_pack(store, query="exact query string", limit=5) + assert any(item["id"] == "c1" for item in pack.get("items", [])) + + +def test_build_context_pack_explain_flag_returns_score_breakdown( + tmp_path: Path, +) -> None: + from tests.embeddings._fakes import MockEmbedder + from vouch.context import build_context_pack + from vouch.embeddings import register + from vouch.embeddings.base import DEFAULT_MODEL_NAME + from vouch.models import Claim + from vouch.storage import KBStore + + register(DEFAULT_MODEL_NAME, lambda: MockEmbedder(dim=8)) + store = KBStore.init(tmp_path) + src = store.put_source(b"e") + store.put_claim(Claim(id="c1", text="hello", evidence=[src.id])) + pack = build_context_pack(store, query="hello", limit=5, explain=True) + assert "explain" in pack + assert any("backend" in row for row in pack["explain"])