feat(chat): forward document_ids through retrieval for multi-document chat#669
Merged
param20h merged 1 commit intoJun 23, 2026
Merged
Conversation
Contributor
Author
param20h
approved these changes
Jun 23, 2026
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
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.
🔗 Related Issue
Closes #649
📝 What does this PR do?
document_idsalready existed as a field onChatRequest, and both storage layers already supported it —query_chunks(vectorstore) filters with ChromaDB's$in, andquery_bm25queries per-document indexes. Butretrieve()never forwardeddocument_idsto either, and the chat routes never passed it down, so asking a question across multiple documents silently behaved exactly like no document filter at all. This PR connects that existing capability to the request flow.retrieve()acceptsdocument_idsand forwards it to bothquery_chunksandquery_bm25(and its trace metadata factory).PDFSearchTooland the agent (get_agent_executor,generate_answer,generate_answer_stream) threaddocument_idsthrough to retrieval./askand/ask/streamforwarddocument_idsand validate every requested document: 404 if any are missing or not owned by the user, 400 if any are still processing — mirroring the existing single-document guard.document_ids, so a multi-document query doesn't collide with a single-document query (or a different set of documents) and return a stale answer.document_idkeeps precedence overdocument_idseverywhere, matching the existing vectorstore logic.🗂️ Type of Change
🧪 How was this tested?
Added
tests/test_multi_document_chat.py(7 tests):document_idsreaches both the vector and BM25 calls inretrieve(); single-document requests still leavedocument_idsasNone; the comparison guidance appears in the agent prompt only when more than one document is selected; and the route guard returns 404 for missing/unowned documents and 400 for not-ready ones. Full suite passes (239 tests). I also updated one existing mock (fake_retrieveintest_rag_tools.py) to accept thedocument_idskwarg the tool now forwards — a signature-only change, assertions unchanged.I didn't test via live API calls — the agent path needs a HuggingFace model, so verification is through the unit and route tests rather than a real generation.
The storage-layer support for
document_ids(vectorstore$infilter, per-document BM25) was already present ondevfrom earlier work — it just wasn't reachable becauseretrieve()and the routes didn't forward the field. This PR is the wiring plus the ownership guard, cache-key fix, and comparison prompt, not the storage filters themselves.One thing worth a look: the cache-key change. Previously the cache keyed on
document_idonly, so a multi-document query (wheredocument_idisNone) would have keyed on an empty string and could return a cached answer from an unrelated query. I prefixed multi-doc cache keys withmulti:plus the sorted document IDs so they're distinct and order-independent. Happy to adjust if you'd prefer a different cache strategy.✅ Self-Review Checklist
dev, notmainmainbranch or any HuggingFace deployment config