diff --git a/tests/test_cli_citations.py b/tests/test_cli_citations.py new file mode 100644 index 0000000..807c327 --- /dev/null +++ b/tests/test_cli_citations.py @@ -0,0 +1,34 @@ +"""CLI rendering regression tests, isolated from database retrieval.""" + +from io import StringIO +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from rich.console import Console + +from vaultrag import cli +from vaultrag.generate import Answer + + +@pytest.mark.parametrize("entry", ["[/]", "[red]bad[/red]"]) +async def test_ask_prints_dropped_citations_literally(monkeypatch, entry): + output = StringIO() + monkeypatch.setattr(cli, "console", Console(file=output, width=120, color_system=None)) + monkeypatch.setattr(cli, "get_embedder", Mock(return_value=Mock(embed=lambda _: [[0.0]]))) + monkeypatch.setattr(cli, "_llm", Mock(return_value=None)) + monkeypatch.setattr(cli.psycopg.AsyncConnection, "connect", AsyncMock(return_value=AsyncMock())) + principal = SimpleNamespace(user_id="alice", principals=["alice"]) + monkeypatch.setattr(cli, "resolve_principal", AsyncMock(return_value=principal)) + hit = SimpleNamespace(doc_id="doc-1", title="Handbook", score=0.5) + monkeypatch.setattr(cli, "search", AsyncMock(return_value=[hit])) + monkeypatch.setattr(cli, "detect_conflicts", Mock(return_value=[])) + monkeypatch.setattr(cli, "detect_stale", Mock(return_value=[])) + monkeypatch.setattr( + cli, + "generate", + Mock(return_value=Answer(text="No source.", answered=False, dropped_citations=[entry])), + ) + + assert await cli._ask(SimpleNamespace(user="alice", question="bonus?", limit=5)) == 0 + assert repr([entry]) in output.getvalue() diff --git a/tests/test_generate.py b/tests/test_generate.py index 33360e3..9d89b57 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -8,9 +8,12 @@ from __future__ import annotations +import json from datetime import datetime -from vaultrag.generate import Answer, FakeLLM, generate +import pytest + +from vaultrag.generate import FakeLLM, generate from vaultrag.retrieval import Hit @@ -112,3 +115,24 @@ def test_json_in_a_code_fence_is_parsed(): ans = generate(llm, "what is the bonus", [_hit()]) assert ans.answered is True assert ans.text == "10% of base [1]." + + +@pytest.mark.parametrize("entry", ["1", "two", 2.5, -1, True, False, None, [], {}]) +def test_malformed_citations_are_recorded_without_coercion(entry): + llm = FakeLLM(json.dumps({"answer": "10% of base [1].", "cited": [entry, 1]})) + ans = generate(llm, "what is the bonus", [_hit(chunk_id=42)]) + + assert ans.answered is True + assert [c.chunk_id for c in ans.citations] == [42] + assert ans.dropped_citations == [entry] + assert type(ans.dropped_citations[0]) is type(entry) + + +def test_only_malformed_citations_refuse_and_preserve_all_entries(): + entries = ["1", "two", 2.5, -1] + llm = FakeLLM(json.dumps({"answer": "10% of base.", "cited": entries})) + ans = generate(llm, "what is the bonus", [_hit()]) + + assert ans.answered is False + assert ans.refusal_reason == "no_verifiable_citation" + assert ans.dropped_citations == entries diff --git a/tests/test_response_models.py b/tests/test_response_models.py new file mode 100644 index 0000000..9199e99 --- /dev/null +++ b/tests/test_response_models.py @@ -0,0 +1,19 @@ +"""Response serialization must preserve malformed citation evidence without coercion.""" + +from vaultrag.main import AskResponse + + +def test_dropped_citation_entries_survive_response_serialization(): + entries = ["1", "two", 2.5, -1, True, False, None, [], {"source": 1}] + response = AskResponse( + answer="No verifiable citation.", + answered=False, + query_id=1, + dropped_citations=entries, + ) + + restored = AskResponse.model_validate_json(response.model_dump_json()) + assert restored.dropped_citations == entries + assert [type(entry) for entry in restored.dropped_citations] == [ + type(entry) for entry in entries + ] diff --git a/vaultrag/cli.py b/vaultrag/cli.py index 863c00c..6ccb1d0 100644 --- a/vaultrag/cli.py +++ b/vaultrag/cli.py @@ -16,6 +16,7 @@ import psycopg from rich.console import Console +from rich.markup import escape from rich.table import Table from .config import get_settings @@ -93,7 +94,9 @@ async def _ask(args) -> int: answer = generate(llm, args.question, hits) if answer.dropped_citations: - console.print(f"[yellow]unverified citations dropped[/] {answer.dropped_citations}") + console.print( + f"[yellow]unverified citations dropped[/] {escape(repr(answer.dropped_citations))}" + ) console.print(f"\n[bold]{answer.text}[/]") if answer.citations: diff --git a/vaultrag/generate.py b/vaultrag/generate.py index de77ebd..89c0533 100644 --- a/vaultrag/generate.py +++ b/vaultrag/generate.py @@ -64,7 +64,8 @@ class Answer: answered: bool = True refusal_reason: str | None = None conflict: bool = False - dropped_citations: list[int] = field(default_factory=list) # model cited these; we couldn't verify + # Preserve raw JSON entries, including malformed values, as evidence of model errors. + dropped_citations: list[object] = field(default_factory=list) _SYSTEM = """You answer questions using ONLY the provided context. @@ -135,9 +136,10 @@ def generate(llm: LLM, question: str, hits: list[Hit]) -> Answer: # Verify: the model can only cite sources we actually handed it. Anything else it invented. citations: list[Citation] = [] - dropped: list[int] = [] + dropped: list[object] = [] for idx in cited_idx: - if 1 <= idx <= len(hits): + # JSON source indices must be integers, not booleans or numeric strings. + if type(idx) is int and 1 <= idx <= len(hits): h = hits[idx - 1] citations.append( Citation(chunk_id=h.chunk_id, doc_id=h.doc_id, title=h.title, url=h.url, owner=h.owner) @@ -163,7 +165,7 @@ def generate(llm: LLM, question: str, hits: list[Hit]) -> Answer: ) -def _parse(raw: str) -> tuple[str, list[int], bool] | None: +def _parse(raw: str) -> tuple[str, list[object], bool] | None: """Pull JSON out of a model response, tolerating code fences and stray prose.""" candidate = raw.strip() @@ -190,8 +192,8 @@ def _parse(raw: str) -> tuple[str, list[int], bool] | None: cited = data.get("cited", []) if not isinstance(cited, list): cited = [] - idx = [int(c) for c in cited if isinstance(c, (int, str)) and str(c).strip().isdigit()] - return (answer, idx, bool(data.get("conflict", False))) + # Verification records invalid entries; filtering here would erase that evidence. + return (answer, cited, bool(data.get("conflict", False))) class FakeLLM: diff --git a/vaultrag/main.py b/vaultrag/main.py index 65b71ac..239ea16 100644 --- a/vaultrag/main.py +++ b/vaultrag/main.py @@ -49,7 +49,7 @@ class AskResponse(BaseModel): citations: list[CitationOut] = [] conflict: bool = False refusal_reason: str | None = None - dropped_citations: list[int] = [] + dropped_citations: list[object] = [] query_id: int