Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions tests/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,26 @@ async def test_ask_returns_only_permitted_documents(client):
cited = {c["doc_id"] for c in body["citations"]}
assert "ceo-private" not in cited, "LEAK: CEO notes surfaced over the API"

async def test_dropped_citations_are_surfaced_in_ask_response(client):
"""The model cites source [9] when only one source (index 1) was actually given.

The dropped citation should be visible in the API response, not silently discarded.
"""
await _seed(client)
from vaultrag.generate import FakeLLM
client._transport.app.state.llm = FakeLLM(
'{"answer": "the bonus is 10%", "cited": [1, 9], "conflict": false}'
)
r = await client.post(
"/ask",
json={"question": "quarterly bonus payout policy"},
headers={"X-User-Id": "alice"},
)
assert r.status_code == 200
body = r.json()
assert body["answered"] is True
assert body["dropped_citations"] == [9]


async def test_user_with_no_access_gets_a_refusal_not_a_leak(client):
"""bob is in sales. Neither seeded document is his, so he should be told nothing was found,
Expand Down
4 changes: 4 additions & 0 deletions vaultrag/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,10 @@ async def _ask(args) -> int:
console.print(f"[yellow]stale[/] {s.doc_id} last updated {s.age_days}d ago (owner: {s.owner})")

answer = generate(llm, args.question, hits)

if answer.dropped_citations:
console.print(f"[yellow]unverified citations dropped[/] {answer.dropped_citations}")

console.print(f"\n[bold]{answer.text}[/]")
if answer.citations:
console.print("[dim]sources: " + ", ".join(c.doc_id for c in answer.citations) + "[/]")
Expand Down
2 changes: 2 additions & 0 deletions vaultrag/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ class AskResponse(BaseModel):
citations: list[CitationOut] = []
conflict: bool = False
refusal_reason: str | None = None
dropped_citations: list[int] = []
query_id: int


Expand Down Expand Up @@ -161,6 +162,7 @@ async def ask(req: AskRequest, user_id: str = Depends(current_user)) -> AskRespo
answered=answer.answered,
conflict=answer.conflict,
refusal_reason=answer.refusal_reason,
dropped_citations=answer.dropped_citations,
query_id=query_id,
citations=[
CitationOut(doc_id=c.doc_id, title=c.title, chunk_id=c.chunk_id, url=c.url, owner=c.owner)
Expand Down
Loading