feat: retrieval and citation-grounded RAG query - #4
Merged
Conversation
cosine_top_k(session, *, query_vec, k) returns ChunkScore(chunk, score) tuples ranked by cosine similarity descending. Score is 1 - distance so 1.0 == perfect match. NULL embeddings are excluded; ties are broken deterministically by chunk id. Caller embeds the query themselves so the same EmbeddingProvider used at ingest time is also used at query time (and so this module stays a pure SQL adapter).
- base.py: LLMClient Protocol with .model_name and .complete(system, user, max_tokens, temperature) → LLMResponse(text, model, stop_reason). Single-turn semantics; richer features can extend later without breaking M3. - fake.py: deterministic canned-response client for tests. Optional response_factory callable lets tests vary output across calls (e.g. different citations per question). - claude.py: thin httpx POST against Anthropic /v1/messages, no SDK dep. Concatenates text content blocks; surfaces stop_reason. Uses anthropic-version: 2023-06-01. - __init__.py: get_llm(settings) factory dispatches anthropic|fake. - config.py: adds llm_provider, claude_model, llm_temperature (default 0.0 per CLAUDE.md 'pin temperatures for LLM calls used in eval'), llm_max_tokens. - .env.example: documents the new LLM_* vars.
answer_query(session, *, query, embedder, llm, settings) returns a RagAnswer with status='answered' (with citations) or status='refused' (with a reason). Refusal triggers: - 'empty_query' — whitespace-only input. - 'no_support' — no retrieved chunks, or top score < retrieval_min_score. - 'uncited' — LLM produced text but no [chunk:N] referenced a retrieved id. The system prompt instructs the model to cite using [chunk:N] markers and to refuse verbatim with REFUSAL_TEXT if context is insufficient. The parser maps citations back to retrieved chunks; citations to chunk ids we did not supply are dropped.
Adds backend/app/routers/query.py with QueryRequest, CitationOut, and QueryResponse Pydantic models. The handler resolves session, embedder, and LLM via FastAPI Depends so tests can override each independently (TestClient + dependency_overrides). main.py includes the query router; module also imports cleanly without a DB connection, and OpenAPI/docs routes still work alongside /health and /query.
…moke - test_retrieval.py (5 tests): cosine_top_k orders by descending similarity using crafted unit vectors with known strengths; respects k limit; excludes NULL-embedding chunks (property-based — robust to other DB state); self-similarity = 1.0; rejects invalid inputs. - test_rag.py (7 tests): happy path returns at least one citation that maps to a retrieved chunk; refusal when top score < retrieval_min_score; refusal when corpus is empty; refusal when LLM does not cite ([chunk:N] absent); refusal when LLM cites only unknown chunk ids; mixed valid+bogus citation parsing keeps unique valid ones in first-seen order; refusal on whitespace-only query without touching the session. - test_query_router.py (5 tests): TestClient with dependency_overrides for session, embedder, and llm. Validates request schema (empty body, oversized query → 422); happy path round-trip; refusal on empty corpus; refusal on uncited LLM response. All RAG/router tests use FakeLLM and an aligned-query embedder; no live API calls in the test suite. Total project tests: 74.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4060f83002
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+92
to
+95
| if hit is None: | ||
| # The model cited a chunk we did not supply. Drop it; the caller's | ||
| # uncited-answer check below will refuse if no valid citations remain. | ||
| continue |
There was a problem hiding this comment.
When the LLM output contains at least one valid retrieved citation plus one fabricated citation, this branch silently drops the fabricated id while answer_query still returns status="answered" with the original answer text, including the bogus [chunk:N] marker. In that mixed-citation scenario the caller sees an unsupported source reference that is absent from the structured citations, so the citation-or-refuse guardrail is bypassed; the response should be refused or the invalid marker otherwise handled before returning the answer.
Useful? React with 👍 / 👎.
8 tasks
div0rce
added a commit
that referenced
this pull request
May 29, 2026
* docs(progress): record M3 merged (PR #4) and mark M4 in progress * feat(extraction): ExtractedField[T] wrapper, schema registry, invoice schema ExtractedField[T] (PEP 695 generic Pydantic v2 model, extra='forbid') forces every extracted field to carry value, confidence in [0,1], and a source_chunk_id ≥ 1. InvoicePayload composes four ExtractedField[T] fields (invoice_number, vendor, issue_date, total_due) and matches the synthetic invoice corpus. The registry exposes get_schema(name) / list_schemas() so the orchestrator and the POST /extract router can resolve schemas by string name. Adding a schema is a two-line edit: define the model and register it. * feat(extract): schema-constrained extraction orchestrator with strict validation extract_document(session, *, document_id, schema_name, llm, settings) returns an ExtractionResult with status='ok' (persisted extraction_id, flat payload + per-field confidence + per-field citations) or status='failed' (with a typed reason and detail string). Failure reasons are stable strings: - document_not_found / no_chunks — preconditions. - unknown_schema — registry lookup miss. - parse_error — output is not valid JSON (markdown fences are tolerated; nothing else). - schema_invalid — JSON does not match the Pydantic schema (missing field, wrong type, confidence out of [0,1], extra keys, etc.). - invalid_citation — a source_chunk_id was not in the supplied context. Same hard-fail posture as the M3 RAG layer. Failures do NOT persist an extraction row. Only ok results land in the extractions table via extractions_repo.create(). * feat(extract): POST /extract router with strict request/response schemas ExtractRequest validates document_id (>=1) and schema_name (1-128 chars). Resolves session and LLMClient via FastAPI Depends so tests can override each via app.dependency_overrides. The handler is intentionally thin — it delegates to backend.app.extract.extract_document and converts the ExtractionResult to ExtractResponse. On status='ok' the session is committed so the new extractions row is durable; failure paths issue no writes. main.py now includes both the M3 query router and the M4 extract router. * test(m4): valid extraction, malformed/schema-invalid/invalid-citation refusals, persistence test_extract.py (12 tests): - happy path: extraction validates, persists, payload + per-field confidence + per-field citations round-trip from the repo, model_name='fake-llm' captured. - markdown-fenced JSON parses successfully. - parse_error: non-JSON output and JSON-but-not-an-object (array). - schema_invalid: missing required field, confidence out of [0,1], extra unknown keys (extra='forbid'). - invalid_citation: source_chunk_id pointing at a chunk not in the supplied context. - document_not_found: unknown document_id. - no_chunks: document with zero chunks. - unknown_schema: unregistered schema_name. - failures do not persist (parametrized over parse_error and schema_invalid). - InvoicePayload sanity round-trip. test_extract_router.py (6 tests): - request validation (empty body, document_id=0). - happy path: 200 with status='ok', extraction_id, payload, field_confidence, field_citations. - failed surfaces: malformed LLM output (status='failed', reason='parse_error'), unknown document (document_not_found), unknown schema (unknown_schema). All tests use FakeLLM only — no live API calls in CI. Total project tests: 96. * docs(progress): mark M4 complete on branch with DoD verification and decision log * fix: enforce invoice issue date schema
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Milestone
M3 — Retrieval + citation-grounded RAG
Summary
End-to-end RAG: pgvector cosine top-k retrieval over chunks, an
LLMClientProtocolwith
FakeLLM(used in tests) andClaudeClient(httpx-based;/v1/messages), ananswer_queryorchestrator that enforces citation-or-refuse, aPOST /queryFastAPI endpoint with Pydantic request/response schemas, and 17 new tests covering
retrieval ordering, refusal triggers, and citation→chunk mapping correctness — all
running offline against
FakeLLMandFakeEmbedder.Definition of Done
make checkpasses (ruff + ruff-format + mypy strict + pytest)and the
/queryrouterwith PR feat: document ingestion and embedding pipeline #3 link; 4 new decision-log entries
logging) — citation-or-refuse implemented and tested; PII redaction, confidence
gating, and audit logging arrive in M5/M7
M3 DoD verification (from MILESTONES.md)
POST /queryreturns answer + citations for an in-corpus question (manualcheck with real key). End-to-end
POST /queryreturns{status:"answered", answer:"... [chunk:N]", citations:[{chunk_id,document_id,score,text}], reason:null}via the FakeLLM smoke test (
test_query_router::test_post_query_happy_path_returns_answer_with_citations).The path is the same with
LLM_PROVIDER=anthropicand a real key —answer_querydelegates to
LLMClient, andClaudeClientis a drop-in.mapping correctness.
test_retrieval.py): crafted unit vectors with knownstrengthyield strict descending similarity; k limit honoured; NULL embeddings excluded;
self-similarity = 1.0.
no_supportwhen top score <RETRIEVAL_MIN_SCOREand when thecorpus is empty;
uncitedwhen the LLM produced no[chunk:N]markers;uncitedwhen the LLM cited only chunk ids that were not in the retrieved set;
empty_queryon whitespace-only input.test_rag.py::test_answer_query_citation_to_chunk_mapping_is_correct):a mixed valid+bogus+duplicate citation string parses to unique valid citations in
first-seen order, with
score,document_id, andtextresolved from theretrieved set.
Notes / decisions
See PROGRESS.md "Decision log" — four M3 entries:
[chunk:N]whereNis the integer chunk id.empty_query,no_support,uncited) soM7 audit and M9 eval can bucket failures without re-running the pipeline.
LLMClient.complete(system, user, max_tokens, temperature)is single-turn bydesign; streaming and tool use can extend later.
llm_temperaturedefaults to0.0per CLAUDE.md.Other deliberate choices:
/queryrouter resolvesSession, embedder, and LLM via three separateDependscallables so tests can override each independently with FastAPI'sapp.dependency_overrides.ingest-time and query-time embedders can be unified by the FastAPI dependency.
OpenAIEmbedder/ClaudeClientboth go throughhttpxdirectly; no SDK deps.Test coverage
test_retrieval.py(5): cosine ordering, k limit, NULL exclusion, self-similarity,invalid-input rejection.
test_rag.py(7): happy path, no_support refusal, empty-corpus refusal, uncitedrefusal, unknown-chunk-id refusal, citation-mapping correctness, empty-query refusal.
test_query_router.py(5): request validation (empty body, oversized query),happy path round-trip, empty-corpus refusal, uncited refusal.
Total project tests: 74 (was 57 in M2).
Reminder
Please squash-merge this PR (Golden Rule #2: one milestone = one squash-merge
commit on
main).