Review a RAG Answer Function as a Pull Request, Then Write Tests for Its Bugs

You are reviewing a pull request for a question-answering service built on retrieval-augmented generation (RAG) over documents that carry per-user read permissions. The pull request adds the function below. It retrieves documents relevant to a question, keeps the ones the user may read, asks an internal model for an answer with citations, falls back to an external model if the internal one times out, writes an audit record and returns the answer. ```python async def answer(question, user): docs = await retrieve(question, top_k=8) visible = [doc for doc in docs if can_read(user, doc)] try: result = await internal_model(question, visible) if not result.citations: return {"status": "abstained"} except TimeoutError: result = await external_model(question, docs) if not result.citations: return {"status": "abstained"} await audit.write({ "user_id": user.id, "question": question, "answer": result.text, }) return { "status": "answered", "text": result.text, "citations": result.citations, } ``` The report gives the code without indentation; it is restored here in its most natural reading, with the second citation check inside the `except` block. Treat this as a code review: find the problems in the change and discuss them, then write test cases based on the problems you found. The tests do not need exact syntax, but each one should be concrete enough to run after minor fixes. ### Constraints and Clarifications - Assume `retrieve(question, top_k)` returns up to `top_k` documents ranked by relevance to the question, without regard to who is asking. - `can_read(user, doc)` is intended as the authoritative check of whether `user` may read `doc`. - Assume `internal_model(question, docs)` runs inside the organization and `external_model(question, docs)` is hosted outside it, for example by a third-party provider. Both generate an answer grounded in `docs` and return an object with `text` and `citations`, the documents the answer cites. - `audit.write(record)` persists an audit record. - The helpers' code is not shown. Reason from their names and call sites, and say which of your findings depend on an assumption about them. ### Clarifying Questions - Is `can_read` a plain function or a coroutine function? - Is the external model approved to receive the user's question and the contents of internal documents, and for which kinds of data? - Must every request, including abstentions and failures, produce an audit record, and should a failed audit write block the response? - Does `internal_model` enforce its own deadline and raise `TimeoutError`, or is the caller expected to? What is the endpoint's overall latency budget? - What does a citation contain (a document ID, a title, a snippet), and is it shown to the user as is? ### Part 1 — Review the change List the problems you would raise in review, from most to least severe. For each one, explain what can go wrong, when it would happen in production, and what change you would ask the author to make. ```hint Compare the two paths The happy path and the fallback path look alike. Read them side by side, argument by argument. ``` ```hint Question each helper's contract Ask what the code assumes about `can_read`, about what `result.citations` contains and about `audit.write`, and what happens when each assumption is false. ``` ```hint Walk the failure paths For every awaited call, ask what happens if it raises something other than the exception that is caught, or never returns at all. ``` #### What This Part Should Cover - Authorization: whether every model call and everything returned to the user respects the user's read permissions on every path - Failure handling: which failures trigger the fallback, which escape to the caller, and what the caller then sees - Grounding and abstention: whether the citation check shows that the answer rests on documents the user may read - Data governance and auditability: what leaves the organization, and whether the audit record lets someone reconstruct what happened ### Part 2 — Write test cases Based on the problems you found, write test cases for this function. For each test give the setup, the call and the assertions, and say whether it would fail against the code as submitted. ```hint Control the collaborators Every outside call in the function is a name a test can replace with a fake whose return value or exception you choose. ``` ```hint Make each problem observable For each problem from Part 1, choose inputs for which the submitted code and the fixed code produce different results that a test can see, including what each fake was called with. ``` #### What This Part Should Cover - Isolation of the function from its collaborators with fakes or mocks, including timeouts and exceptions - One test per Part 1 finding that fails on the submitted code and passes once it is fixed, plus regression tests for behavior that is already correct - Assertions on what each collaborator received and what was audited, not only on the returned value - Tests that depend on an unresolved policy, marked as such ### What a Strong Answer Covers - Findings ranked so that security and data-governance risks come before efficiency and style points - A concrete requested change for each finding, with the assumption or open policy question it depends on - Reasoning about how often each failure path runs in production, not only whether it can run - Deterministic, fast tests that fake the models, the permission check and the audit log instead of calling real services - Awareness of what unit tests cannot show about a model-backed service, and what would cover it ### Follow-up Questions - How would you move the permission check into retrieval itself, and what has to happen when a document's permissions change after it was indexed? - A document the user is allowed to read contains text telling the model to ignore its instructions. What can go wrong in this function, and how would you test for it? - If policy forbids sending internal documents to the external model, what should the function do when the internal model times out, and how do your tests change? - How would you test the timeout and fallback logic without making the suite slow or flaky?