From 2ac186690bd11f527f14ffe611d213a893c60444 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Sun, 2 Aug 2026 21:27:39 +0000 Subject: [PATCH 01/14] fix(rag): keep sources when the LLM omits the citation tag No [Sources: ...] tag means the model didn't report which sources it used, not that it used none. filter_sources_by_citations now keeps all retrieved sources in that case instead of hiding them, since answers were frequently coming back with no cited sources at all. --- openrag/core/utils/source_filtering.py | 19 +++++------ .../services/orchestrators/query_service.py | 29 ++++------------ .../unit/core/utils/test_source_filtering.py | 34 ++++--------------- .../orchestrators/test_query_service.py | 11 +++--- 4 files changed, 27 insertions(+), 66 deletions(-) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index ff7b52833..7f8fb462c 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -79,15 +79,15 @@ def extract_and_strip_sources_block( return cleaned, set() -def filter_sources_by_citations( - sources: list, - citations: set[int] | None, - *, - allow_uncited: bool = False, -) -> list: - """Keep only sources whose 1-based index was cited.""" +def filter_sources_by_citations(sources: list, citations: set[int] | None) -> list: + """Keep only sources whose 1-based index was cited. + + No tag at all (``citations is None``) means the model didn't report which + sources it used, not that it used none — the answer may still be grounded + in them, so keep everything rather than silently dropping real sources. + """ if citations is None: - return sources if allow_uncited else [] + return sources if not citations: return [] return [source for i, source in enumerate(sources, start=1) if i in citations] @@ -112,7 +112,6 @@ async def stream_with_source_filtering( model_name: str, buffer_size: int | None = None, *, - allow_uncited_sources: bool = False, citation_protocol_active: bool = True, ): """Process an LLM SSE stream and, when active, strip source tags. @@ -263,7 +262,7 @@ async def stream_with_source_filtering( else: final_clean, citations = pending, None - filtered = filter_sources_by_citations(sources, citations, allow_uncited=allow_uncited_sources) + filtered = filter_sources_by_citations(sources, citations) extra_payload = {"sources": filtered} if not saw_done: extra_payload["truncated"] = True diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 7c9287b23..92cedfe24 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -655,7 +655,7 @@ async def chat( else: payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) - structured_output = _allows_uncited_sources(payload) + structured_output = _is_structured_output(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) chunk = await llm.chat(payload["messages"], **_sampling(payload)) @@ -669,15 +669,7 @@ async def chat( else: clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean - chunk["extra"] = json.dumps( - { - "sources": filter_sources_by_citations( - sources, - citations, - allow_uncited=structured_output, - ) - } - ) + chunk["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) return chunk async def chat_stream( @@ -697,7 +689,7 @@ async def chat_stream( else: payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) - structured_output = _allows_uncited_sources(payload) + structured_output = _is_structured_output(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) llm_stream = llm.stream_chat(payload["messages"], **_sampling(payload)) @@ -705,7 +697,6 @@ async def chat_stream( llm_stream, sources, model_name, - allow_uncited_sources=structured_output, citation_protocol_active=citation_protocol_active and not structured_output, ): yield sse_line @@ -725,7 +716,7 @@ async def complete( else: payload, docs = await self._prepare_completions(partitions, payload, llm) sources = prepare_sources(docs, []) - structured_output = _allows_uncited_sources(payload) + structured_output = _is_structured_output(payload) resp = await llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) text = resp.get("choices", [{}])[0].get("text", "") or "" @@ -737,15 +728,7 @@ async def complete( else: clean, citations = text, None resp["choices"][0]["text"] = clean - resp["extra"] = json.dumps( - { - "sources": filter_sources_by_citations( - sources, - citations, - allow_uncited=structured_output, - ) - } - ) + resp["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) return resp @@ -789,7 +772,7 @@ def _sampling(payload: dict, key: str = "messages") -> dict: return {k: v for k, v in payload.items() if k not in drop} -def _allows_uncited_sources(payload: dict) -> bool: +def _is_structured_output(payload: dict) -> bool: """Structured output cannot carry the plain-text citation marker.""" response_format = payload.get("response_format") return isinstance(response_format, dict) and response_format.get("type") in {"json_object", "json_schema"} diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index 877587542..22eed3247 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -160,14 +160,10 @@ def test_basic_filtering(self): result = filter_sources_by_citations(sources, {1, 3, 5}) assert result == ["a", "c", "e"] - def test_none_citations_returns_empty(self): + def test_none_citations_returns_all_sources(self): + """No tag at all means the model didn't report citations, not that it used none.""" sources = ["a", "b", "c"] result = filter_sources_by_citations(sources, None) - assert result == [] - - def test_none_citations_can_be_allowed_for_structured_output(self): - sources = ["a", "b", "c"] - result = filter_sources_by_citations(sources, None, allow_uncited=True) assert result == ["a", "b", "c"] def test_empty_citations_returns_empty(self): @@ -393,8 +389,8 @@ async def test_case2_llm_says_sources_none(self): assert _parse_finish_sources(result) == [] @pytest.mark.asyncio - async def test_case3_llm_no_tag_returns_no_sources(self): - """Case 3: LLM omits tag entirely → no source is attributed.""" + async def test_case3_llm_no_tag_returns_all_sources(self): + """Case 3: LLM omits tag entirely → treated as unreported, not uncited; all sources kept.""" lines = [ _make_chunk("Answer without any sources tag."), _make_finish(), @@ -402,23 +398,6 @@ async def test_case3_llm_no_tag_returns_no_sources(self): ] result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) assert _collect_content(result) == "Answer without any sources tag." - assert _parse_finish_sources(result) == [] - - @pytest.mark.asyncio - async def test_no_tag_keeps_sources_when_uncited_output_is_allowed(self): - lines = [ - _make_chunk('{"answer": "structured"}'), - _make_finish(), - DONE_LINE, - ] - result = await _collect( - stream_with_source_filtering( - _fake_stream(lines), - self.SOURCES, - "test-model", - allow_uncited_sources=True, - ) - ) assert _parse_finish_sources(result) == self.SOURCES @pytest.mark.asyncio @@ -434,7 +413,6 @@ async def test_structured_output_preserves_source_like_json_values(self): _fake_stream(lines), self.SOURCES, "test-model", - allow_uncited_sources=True, citation_protocol_active=False, ) ) @@ -514,8 +492,8 @@ async def test_inline_prose_tag_preserved_in_stream(self): result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) content = _collect_content(result) assert content == "Use the format [Sources: 1, 3] at the very end of your response." - # No line-terminal tag means no source was actually cited. - assert _parse_finish_sources(result) == [] + # No line-terminal tag means the model didn't report citations — kept, not dropped. + assert _parse_finish_sources(result) == self.SOURCES @pytest.mark.asyncio async def test_mid_response_tag_stripped_plus_trailing_tag(self): diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 3fce8a7ae..ae0e27d74 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -407,7 +407,7 @@ async def _spy(**kwargs): out = await svc.chat( partitions=None, payload={"messages": [{"role": "user", "content": "hi"}], "metadata": {}}, - prepare_sources=lambda d, w: [{"source_type": "document"}], + prepare_sources=lambda d, w: [{"source_type": "document"}] if d or w else [], model_name="m1", ) assert called["n"] == 0 # no retrieval in direct mode @@ -487,7 +487,7 @@ async def test_chat_conversational_request_skips_partition_retrieval(): out = await svc.chat( partitions=["p"], payload={"messages": [{"role": "user", "content": "How can you help me?"}], "metadata": {}}, - prepare_sources=lambda d, w: [{"source_type": "document"}], + prepare_sources=lambda d, w: [{"source_type": "document"}] if d or w else [], model_name="m", ) @@ -569,7 +569,8 @@ async def test_chat_inconsistent_classifier_result_prefers_supplied_query(): @pytest.mark.asyncio -async def test_chat_without_citation_does_not_attribute_retrieved_sources(): +async def test_chat_without_citation_keeps_retrieved_sources(): + """No tag at all means the model didn't report citations, not that the answer is unsourced.""" svc = _svc(llm=FakeLLM(chat_responses=["A general answer with no citation marker."])) sources = [{"source_type": "document", "filename": "unrelated.pdf"}] @@ -580,7 +581,7 @@ async def test_chat_without_citation_does_not_attribute_retrieved_sources(): model_name="m", ) - assert json.loads(out["extra"])["sources"] == [] + assert json.loads(out["extra"])["sources"] == sources @pytest.mark.asyncio @@ -966,7 +967,7 @@ async def test_complete_direct_mode_preserves_literal_source_marker(): out = await svc.complete( partitions=None, payload={"prompt": "do x"}, - prepare_sources=lambda d, w: [{"x": 1}], + prepare_sources=lambda d, w: [{"x": 1}] if d or w else [], ) assert out["choices"][0]["text"] == answer assert json.loads(out["extra"])["sources"] == [] From 7d45146bc21fcbc0759c65b2e58ca8f35fc30010 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Sun, 2 Aug 2026 21:44:27 +0000 Subject: [PATCH 02/14] feat(rag): expose all retrieved sources alongside cited ones Adds extra.retrieved_sources, carrying every source retrieval produced regardless of citation filtering, so clients can see what was searched even when extra.sources ends up narrower than the full retrieval set. --- CLAUDE.md | 4 ++-- openrag/core/utils/source_filtering.py | 2 +- openrag/services/orchestrators/query_service.py | 8 ++++++-- tests/unit/core/utils/test_source_filtering.py | 12 ++++++++++++ .../services/orchestrators/test_query_service.py | 5 +++-- 5 files changed, 24 insertions(+), 7 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 07590973b..54c9c415a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -133,10 +133,10 @@ The RAG pipeline filters out false-positive sources by having the LLM self-repor 1. `format_context()` (`openrag/core/prompts/chat_prompt_builder.py`) numbers each source (`[Source 1]`, `[Source 2]`, ...) in the context and returns `(formatted_text, included_indices)` — the indices track which docs fit within the token budget 2. Prompt templates (`openrag/prompts/templates/*.txt`) instruct the LLM to append `[Sources: 1, 3, 5]` at the end of its response 3. `extract_and_strip_sources_block()` (`openrag/core/utils/source_filtering.py`) strips this tag from the response before sending to the client -4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources (falls back to all sources if none match) +4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every retrieved source is kept instead (a missing tag means the model didn't report citations, not that it used none) 5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client -The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list]}`. +The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "retrieved_sources": [all_retrieved_sources]}`. `retrieved_sources` always carries every source retrieval produced, unfiltered by citation — useful for clients that want to show what was searched regardless of what the model cited. ### API Routers (`openrag/api/routers/`) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index 7f8fb462c..289091809 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -263,7 +263,7 @@ async def stream_with_source_filtering( final_clean, citations = pending, None filtered = filter_sources_by_citations(sources, citations) - extra_payload = {"sources": filtered} + extra_payload = {"sources": filtered, "retrieved_sources": sources} if not saw_done: extra_payload["truncated"] = True logger.warning( diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 92cedfe24..3eba62bd7 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -669,7 +669,9 @@ async def chat( else: clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean - chunk["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + chunk["extra"] = json.dumps( + {"sources": filter_sources_by_citations(sources, citations), "retrieved_sources": sources} + ) return chunk async def chat_stream( @@ -728,7 +730,9 @@ async def complete( else: clean, citations = text, None resp["choices"][0]["text"] = clean - resp["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)}) + resp["extra"] = json.dumps( + {"sources": filter_sources_by_citations(sources, citations), "retrieved_sources": sources} + ) return resp diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index 22eed3247..21da4e22e 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -294,6 +294,18 @@ async def test_case1_llm_cites_specific_sources(self): assert _collect_content(result) == "Here is the answer." assert _parse_finish_sources(result) == [{"file": "a.pdf"}, {"file": "c.pdf"}] + @pytest.mark.asyncio + async def test_retrieved_sources_includes_uncited_ones(self): + """`retrieved_sources` always carries every candidate, unfiltered by citation.""" + lines = [ + _make_chunk("Here is the answer."), + _make_chunk("\n[Sources: 1, 3]"), + _make_finish(), + DONE_LINE, + ] + result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) + assert _parse_finish_extra(result)["retrieved_sources"] == self.SOURCES + @pytest.mark.asyncio async def test_content_and_finish_reason_in_same_chunk_keeps_last_token(self): """A provider that packs the final token and finish_reason into one chunk diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index ae0e27d74..5d36d0d52 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -458,8 +458,9 @@ async def test_chat_with_partition_retrieves_and_filters_sources(): prepare_sources=lambda d, w: sources, model_name="m", ) - filtered = json.loads(out["extra"])["sources"] - assert filtered == [{"source_type": "document", "n": 1}] # only cited source 1 + extra = json.loads(out["extra"]) + assert extra["sources"] == [{"source_type": "document", "n": 1}] # only cited source 1 + assert extra["retrieved_sources"] == sources # unfiltered, everything retrieved @pytest.mark.asyncio From 18d35c7f772659223aea0c8011966f490e24ac05 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Sun, 2 Aug 2026 22:01:36 +0000 Subject: [PATCH 03/14] chore(release): bump version to 2.1.1 --- infra/charts/openrag-stack/Chart.yaml | 4 ++-- infra/charts/openrag-stack/values.yaml | 6 +++--- infra/compose/docker-compose.yaml | 4 ++-- pyproject.toml | 2 +- uv.lock | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/infra/charts/openrag-stack/Chart.yaml b/infra/charts/openrag-stack/Chart.yaml index 557bbbd0a..deea0a96a 100644 --- a/infra/charts/openrag-stack/Chart.yaml +++ b/infra/charts/openrag-stack/Chart.yaml @@ -3,8 +3,8 @@ name: openrag-stack description: A Helm chart for Kubernetes type: application -version: 0.6.1 -appVersion: "2.1.0" +version: 0.6.2 +appVersion: "2.1.1" maintainers: - name: linagora diff --git a/infra/charts/openrag-stack/values.yaml b/infra/charts/openrag-stack/values.yaml index 23e52e495..59d8c08d5 100644 --- a/infra/charts/openrag-stack/values.yaml +++ b/infra/charts/openrag-stack/values.yaml @@ -99,7 +99,7 @@ ray: registry: "ghcr.io" repository: "linagora/openrag-ray" # Pin to a release tag (ideally a digest) for reproducible deploys. - tag: "v2.1.0" + tag: "v2.1.1" resources: head: requests: @@ -383,7 +383,7 @@ adminUi: repository: "linagoraai/openrag-admin-ui" # Pin to a release tag (ideally a digest) for reproducible deploys. Must be a # build from infra/docker/ui.Dockerfile (nginx-unprivileged, listens :8080). - tag: "v2.1.0" + tag: "v2.1.1" pullPolicy: IfNotPresent replicaCount: 1 service: @@ -440,7 +440,7 @@ openrag: registry: "" repository: "linagoraai/openrag" # Pin to a release tag (ideally a digest) for reproducible deploys. - tag: "v2.1.0" + tag: "v2.1.1" pullPolicy: IfNotPresent service: type: ClusterIP diff --git a/infra/compose/docker-compose.yaml b/infra/compose/docker-compose.yaml index 38d1bb3a4..5a6d1a1b9 100644 --- a/infra/compose/docker-compose.yaml +++ b/infra/compose/docker-compose.yaml @@ -18,7 +18,7 @@ x-openrag-env: &openrag_env FONT_PATH: ${FONT_PATH:-/app/data/fonts/GoNotoCurrent-Regular.ttf} x-openrag: &openrag_template - image: linagoraai/openrag:v2.1.0 + image: linagoraai/openrag:v2.1.1 # Start as root so entrypoint.sh can grant GID-0 write on the bind-mounted # writable dirs (data/, logs/, the HF cache) — which a non-root container # can't write when Docker auto-creates them root-owned — then it immediately @@ -113,7 +113,7 @@ x-vllm: &vllm_template services: # ── Admin UI (React SPA + nginx, same-origin reverse proxy to the API) ── admin-ui: - image: linagoraai/openrag-admin-ui:v2.1.0 + image: linagoraai/openrag-admin-ui:v2.1.1 build: context: ../.. dockerfile: infra/docker/ui.Dockerfile diff --git a/pyproject.toml b/pyproject.toml index a312607a9..c0db9d14f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "openrag" -version = "2.1.0" +version = "2.1.1" description = "Add your description here" readme = "README.md" requires-python = ">=3.12" diff --git a/uv.lock b/uv.lock index 5dfa8b8f1..9b10a8e6f 100644 --- a/uv.lock +++ b/uv.lock @@ -2713,7 +2713,7 @@ wheels = [ [[package]] name = "openrag" -version = "2.1.0" +version = "2.1.1" source = { editable = "." } dependencies = [ { name = "aiobreaker" }, From f4bd28a6aff3bf26f73552d10162b098a604dfb7 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Sun, 2 Aug 2026 22:03:01 +0000 Subject: [PATCH 04/14] refactor(rag): rename extra field to all_retrieved_sources Clearer than retrieved_sources for a field meant for debugging and RAG evaluation. --- CLAUDE.md | 2 +- openrag/core/utils/source_filtering.py | 2 +- openrag/services/orchestrators/query_service.py | 4 ++-- tests/unit/core/utils/test_source_filtering.py | 6 +++--- tests/unit/services/orchestrators/test_query_service.py | 2 +- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 54c9c415a..e30dae1d8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -136,7 +136,7 @@ The RAG pipeline filters out false-positive sources by having the LLM self-repor 4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every retrieved source is kept instead (a missing tag means the model didn't report citations, not that it used none) 5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client -The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "retrieved_sources": [all_retrieved_sources]}`. `retrieved_sources` always carries every source retrieval produced, unfiltered by citation — useful for clients that want to show what was searched regardless of what the model cited. +The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "all_retrieved_sources": [all_retrieved_sources]}`. `all_retrieved_sources` always carries every source retrieval produced, unfiltered by citation — useful for debugging and RAG evaluation, regardless of what the model cited. ### API Routers (`openrag/api/routers/`) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index 289091809..12f453a37 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -263,7 +263,7 @@ async def stream_with_source_filtering( final_clean, citations = pending, None filtered = filter_sources_by_citations(sources, citations) - extra_payload = {"sources": filtered, "retrieved_sources": sources} + extra_payload = {"sources": filtered, "all_retrieved_sources": sources} if not saw_done: extra_payload["truncated"] = True logger.warning( diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 3eba62bd7..2f4bc80e7 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -670,7 +670,7 @@ async def chat( clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean chunk["extra"] = json.dumps( - {"sources": filter_sources_by_citations(sources, citations), "retrieved_sources": sources} + {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": sources} ) return chunk @@ -731,7 +731,7 @@ async def complete( clean, citations = text, None resp["choices"][0]["text"] = clean resp["extra"] = json.dumps( - {"sources": filter_sources_by_citations(sources, citations), "retrieved_sources": sources} + {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": sources} ) return resp diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index 21da4e22e..f1f44a83f 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -295,8 +295,8 @@ async def test_case1_llm_cites_specific_sources(self): assert _parse_finish_sources(result) == [{"file": "a.pdf"}, {"file": "c.pdf"}] @pytest.mark.asyncio - async def test_retrieved_sources_includes_uncited_ones(self): - """`retrieved_sources` always carries every candidate, unfiltered by citation.""" + async def test_all_retrieved_sources_includes_uncited_ones(self): + """`all_retrieved_sources` always carries every candidate, unfiltered by citation.""" lines = [ _make_chunk("Here is the answer."), _make_chunk("\n[Sources: 1, 3]"), @@ -304,7 +304,7 @@ async def test_retrieved_sources_includes_uncited_ones(self): DONE_LINE, ] result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) - assert _parse_finish_extra(result)["retrieved_sources"] == self.SOURCES + assert _parse_finish_extra(result)["all_retrieved_sources"] == self.SOURCES @pytest.mark.asyncio async def test_content_and_finish_reason_in_same_chunk_keeps_last_token(self): diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 5d36d0d52..986d075e3 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -460,7 +460,7 @@ async def test_chat_with_partition_retrieves_and_filters_sources(): ) extra = json.loads(out["extra"]) assert extra["sources"] == [{"source_type": "document", "n": 1}] # only cited source 1 - assert extra["retrieved_sources"] == sources # unfiltered, everything retrieved + assert extra["all_retrieved_sources"] == sources # unfiltered, everything retrieved @pytest.mark.asyncio From f3c5b2f839ad3b9589c8ad55a4e276c8206b5923 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 08:59:07 +0000 Subject: [PATCH 05/14] fix(rag): capture all_retrieved_sources before context-budget truncation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit all_retrieved_sources was built from the same docs/web_results already truncated by format_context()/format_web_context() to fit the prompt's token budget, so it silently dropped anything retrieval returned but couldn't fit — defeating its purpose as the complete set for debugging and RAG evaluation. Snapshot the full retrieval set before that truncation and thread it through chat, chat_stream, and complete. --- CLAUDE.md | 2 +- openrag/core/utils/source_filtering.py | 12 ++- .../services/orchestrators/query_service.py | 53 ++++++++--- .../orchestrators/test_query_service.py | 93 ++++++++++++++++++- 4 files changed, 142 insertions(+), 18 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index e30dae1d8..a5c99f79f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -136,7 +136,7 @@ The RAG pipeline filters out false-positive sources by having the LLM self-repor 4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every retrieved source is kept instead (a missing tag means the model didn't report citations, not that it used none) 5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client -The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "all_retrieved_sources": [all_retrieved_sources]}`. `all_retrieved_sources` always carries every source retrieval produced, unfiltered by citation — useful for debugging and RAG evaluation, regardless of what the model cited. +The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "all_retrieved_sources": [all_retrieved_sources]}`. `all_retrieved_sources` always carries the complete retrieval set — captured before the context-token-budget truncation in `format_context()`/`format_web_context()`, so it also includes documents/web results that didn't fit in the prompt — unfiltered by citation, useful for debugging and RAG evaluation regardless of what the model cited. ### API Routers (`openrag/api/routers/`) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index 12f453a37..e41cc3c45 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -113,9 +113,16 @@ async def stream_with_source_filtering( buffer_size: int | None = None, *, citation_protocol_active: bool = True, + all_sources: list | None = None, ): """Process an LLM SSE stream and, when active, strip source tags. + ``sources`` is the prompt-visible (context-budget-truncated) list used to + resolve citation indices; ``all_sources`` — the complete pre-truncation + retrieval set — is reported separately as ``extra.all_retrieved_sources`` + and defaults to ``sources`` when the caller has nothing more complete to + offer. + The terminal flush (tail content + ``extra.sources``) runs exactly once after the loop on *every* termination path — a clean ``data: [DONE]``, the upstream closing the connection without one, or the upstream generator @@ -263,7 +270,10 @@ async def stream_with_source_filtering( final_clean, citations = pending, None filtered = filter_sources_by_citations(sources, citations) - extra_payload = {"sources": filtered, "all_retrieved_sources": sources} + extra_payload = { + "sources": filtered, + "all_retrieved_sources": all_sources if all_sources is not None else sources, + } if not saw_done: extra_payload["truncated"] = True logger.warning( diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 2f4bc80e7..e62bfdeab 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -471,7 +471,7 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L context="", current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), ) - return payload, [], [], True + return payload, [], [], [], [], True queries = SearchQueries(query_list=[Query(query=messages[-1]["content"])]) web_results: list = [] @@ -487,12 +487,19 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L chunks = [] if not chunks and not web_results and partition is None: - return payload, [], [], False + return payload, [], [], [], [], False docs = [c.to_langchain() for c in chunks] if use_map_reduce and docs: docs = await self._map_reduce(" ".join(q.query for q in queries.query_list), docs) + # Full retrieval set, captured before the token-budget selection below + # drops anything that didn't fit in the prompt — kept separately for + # `all_retrieved_sources` (debugging/eval), while `docs`/`web_results` + # stay budget-truncated to match the citation indices the LLM sees (#847). + retrieved_docs = docs + retrieved_web_results = web_results + web_formatted, web_source_numbers, web_tokens = "", [], 0 web_start_index = 1 if web_results: @@ -538,7 +545,7 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L }, ) payload["messages"] = new_messages - return payload, docs, web_results, True + return payload, docs, web_results, retrieved_docs, retrieved_web_results, True async def _gather_rag_and_web(self, queries, partition, top_k, filter_params): # Fuse the doc branch through retrieve_multi so a partition's rrf_k drives @@ -558,6 +565,7 @@ async def _prepare_completions(self, partition: list[str], payload: dict, llm: L # partition= is ours: the retrieval preset's query_contextualizer is # resolved per partition. The skip below is from #807. queries = await self.generate_query([{"role": "user", "content": prompt}], llm=llm, partition=partition) + retrieved_docs: list = [] if not queries.query_list: if not queries.requires_retrieval: docs, context = [], "" @@ -566,6 +574,9 @@ async def _prepare_completions(self, partition: list[str], payload: dict, llm: L if queries.query_list: chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) docs = [c.to_langchain() for c in chunks] + # Full retrieval set before the token-budget selection below, kept + # separately for `all_retrieved_sources` (#847). + retrieved_docs = docs context, included = format_context( [doc.page_content for doc in docs], max_context_tokens=self._max_context_tokens, @@ -583,7 +594,7 @@ async def _prepare_completions(self, partition: list[str], payload: dict, llm: L current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), ) payload["prompt"] = f"{instructions}\n\n# User request\n{prompt}" - return payload, docs + return payload, docs, retrieved_docs # ------------------------------------------------------------------ # Message sanitization @@ -651,10 +662,18 @@ async def chat( llm = self._resolve_llm(partitions) citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): - docs, web_results = [], [] + docs, web_results, retrieved_docs, retrieved_web_results = [], [], [], [] else: - payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) + ( + payload, + docs, + web_results, + retrieved_docs, + retrieved_web_results, + citation_protocol_active, + ) = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) + all_sources = prepare_sources(retrieved_docs, retrieved_web_results) structured_output = _is_structured_output(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) @@ -670,7 +689,7 @@ async def chat( clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean chunk["extra"] = json.dumps( - {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": sources} + {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": all_sources} ) return chunk @@ -687,10 +706,18 @@ async def chat_stream( llm = self._resolve_llm(partitions) citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): - docs, web_results = [], [] + docs, web_results, retrieved_docs, retrieved_web_results = [], [], [], [] else: - payload, docs, web_results, citation_protocol_active = await self._prepare_chat(partitions, payload, llm) + ( + payload, + docs, + web_results, + retrieved_docs, + retrieved_web_results, + citation_protocol_active, + ) = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) + all_sources = prepare_sources(retrieved_docs, retrieved_web_results) structured_output = _is_structured_output(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) @@ -699,6 +726,7 @@ async def chat_stream( llm_stream, sources, model_name, + all_sources=all_sources, citation_protocol_active=citation_protocol_active and not structured_output, ): yield sse_line @@ -714,10 +742,11 @@ async def complete( llm = self._resolve_llm(partitions) citation_protocol_active = partitions is not None if partitions is None: - docs = [] + docs, retrieved_docs = [], [] else: - payload, docs = await self._prepare_completions(partitions, payload, llm) + payload, docs, retrieved_docs = await self._prepare_completions(partitions, payload, llm) sources = prepare_sources(docs, []) + all_sources = prepare_sources(retrieved_docs, []) structured_output = _is_structured_output(payload) resp = await llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) @@ -731,7 +760,7 @@ async def complete( clean, citations = text, None resp["choices"][0]["text"] = clean resp["extra"] = json.dumps( - {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": sources} + {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": all_sources} ) return resp diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 986d075e3..dfa07a242 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -463,6 +463,84 @@ async def test_chat_with_partition_retrieves_and_filters_sources(): assert extra["all_retrieved_sources"] == sources # unfiltered, everything retrieved +@pytest.mark.asyncio +async def test_chat_all_retrieved_sources_survives_context_budget_truncation(): + """#847 review: all_retrieved_sources must reflect the complete retrieval + set, not just the docs that fit the prompt's context-token budget — a + doc dropped only for lack of room must still show up there.""" + chunks = [ + Chunk(id="c1", text="short", metadata={"_id": "c1"}), + Chunk(id="c2", text="this one does not fit the token budget", metadata={"_id": "c2"}), + ] + svc = _svc(retrieval=FakeRetrieval(chunks=chunks), llm=FakeLLM(chat_responses=["answer [Sources: 1]"])) + svc._max_context_tokens = qs.get_num_tokens()("[Source 1]\nshort") + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id")} for doc in d], + model_name="m", + ) + + extra = json.loads(out["extra"]) + assert extra["sources"] == [{"id": "c1"}] # only the doc that fit the prompt and was cited + assert extra["all_retrieved_sources"] == [{"id": "c1"}, {"id": "c2"}] # both, unfiltered + + +@pytest.mark.asyncio +async def test_chat_stream_all_retrieved_sources_survives_context_budget_truncation(): + chunks = [ + Chunk(id="c1", text="short", metadata={"_id": "c1"}), + Chunk(id="c2", text="this one does not fit the token budget", metadata={"_id": "c2"}), + ] + stream_lines = [ + 'data: {"choices":[{"delta":{"content":"answer [Sources: 1]"},"finish_reason":null}]}\n\n', + 'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\n', + "data: [DONE]\n\n", + ] + svc = _svc(retrieval=FakeRetrieval(chunks=chunks), llm=FakeLLM(stream_lines=stream_lines)) + svc._max_context_tokens = qs.get_num_tokens()("[Source 1]\nshort") + + lines = [ + line + async for line in svc.chat_stream( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id")} for doc in d], + model_name="m", + ) + ] + chunks_out = [ + json.loads(line[len("data: ") :]) + for line in lines + if line.startswith("data: ") and line.strip() != "data: [DONE]" + ] + extra = next(json.loads(c["extra"]) for c in reversed(chunks_out) if c.get("extra") not in (None, "{}")) + + assert extra["sources"] == [{"id": "c1"}] + assert extra["all_retrieved_sources"] == [{"id": "c1"}, {"id": "c2"}] + + +@pytest.mark.asyncio +async def test_complete_all_retrieved_sources_survives_context_budget_truncation(): + chunks = [ + Chunk(id="c1", text="short", metadata={"_id": "c1"}), + Chunk(id="c2", text="this one does not fit the token budget", metadata={"_id": "c2"}), + ] + svc = _svc(retrieval=FakeRetrieval(chunks=chunks), llm=FakeLLM(gen_text="answer\n[Sources: 1]")) + svc._max_context_tokens = qs.get_num_tokens()("[Source 1]\nshort") + + out = await svc.complete( + partitions=["p"], + payload={"prompt": "q"}, + prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id")} for doc in d], + ) + + extra = json.loads(out["extra"]) + assert extra["sources"] == [{"id": "c1"}] + assert extra["all_retrieved_sources"] == [{"id": "c1"}, {"id": "c2"}] + + @pytest.mark.asyncio async def test_chat_recovers_context_markers_as_citations(): svc = _svc(llm=FakeLLM(chat_responses=["First claim [Source 2]. Second claim [Source 1][Source 2]."])) @@ -703,7 +781,14 @@ async def test_structured_websearch_returns_only_sources_included_in_context(): model_name="m", ) - assert json.loads(out["extra"])["sources"] == [{"url": "https://example.test/included"}] + extra = json.loads(out["extra"]) + assert extra["sources"] == [{"url": "https://example.test/included"}] + # #847 review: excluded (didn't fit the web token budget) still shows up + # in all_retrieved_sources. + assert extra["all_retrieved_sources"] == [ + {"url": "https://example.test/included"}, + {"url": "https://example.test/excluded"}, + ] @pytest.mark.asyncio @@ -735,7 +820,7 @@ async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): retrieval = FakeRetrieval() web_result = SimpleNamespace(url="https://ex.com", title="T", content="web body", snippet="") svc = _svc(retrieval=retrieval, web=FakeWeb(results=[web_result])) - _payload, _docs, web, _citation_protocol_active = await svc._prepare_chat( + _payload, _docs, web, _retrieved_docs, _retrieved_web, _citation_protocol_active = await svc._prepare_chat( ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {"websearch": True}} ) assert len(retrieval.retrieve_multi_calls) == 1 # doc branch fused via the rrf_k-aware retrieve_multi @@ -761,7 +846,7 @@ async def resolve_prompt(self, prompt_type, names=None): marker = MarkerPromptService() svc._prompt_service = marker - payload, _docs, _web, _citation_protocol_active = await svc._prepare_chat( + payload, _docs, _web, _retrieved_docs, _retrieved_web, _citation_protocol_active = await svc._prepare_chat( ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}} ) @@ -1228,7 +1313,7 @@ async def resolve_prompt(self, prompt_type, names=None): "p": SimpleNamespace(generation_prompt_names={"sys_prompt": "chatty"}, chat_history_depth=4) } - out, docs, web, _ = await svc._prepare_chat( + out, docs, web, _retrieved_docs, _retrieved_web, _ = await svc._prepare_chat( ["p"], {"messages": [{"role": "user", "content": "hello!"}], "metadata": {}} ) From f782b071abc4bfec94586b1a4720d60a921cbe66 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 09:00:39 +0000 Subject: [PATCH 06/14] feat(rag): report whether the model actually emitted a citations tag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With no [Sources: ...] tag, `sources` falls back to keeping every retrieved source — identical, from the client's view, to the model explicitly citing all of them. Add `citations_reported` (true only when a tag, even an empty/none one, was found) so clients can tell "cited everything" apart from "didn't report citations at all". --- CLAUDE.md | 2 +- openrag/core/utils/source_filtering.py | 10 +++++++- .../services/orchestrators/query_service.py | 12 ++++++++-- .../unit/core/utils/test_source_filtering.py | 24 +++++++++++++++++++ .../orchestrators/test_query_service.py | 14 ++++++++--- 5 files changed, 55 insertions(+), 7 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index a5c99f79f..50d0e33b4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -136,7 +136,7 @@ The RAG pipeline filters out false-positive sources by having the LLM self-repor 4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every retrieved source is kept instead (a missing tag means the model didn't report citations, not that it used none) 5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client -The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "all_retrieved_sources": [all_retrieved_sources]}`. `all_retrieved_sources` always carries the complete retrieval set — captured before the context-token-budget truncation in `format_context()`/`format_web_context()`, so it also includes documents/web results that didn't fit in the prompt — unfiltered by citation, useful for debugging and RAG evaluation regardless of what the model cited. +The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "all_retrieved_sources": [all_retrieved_sources], "citations_reported": bool}`. `all_retrieved_sources` always carries the complete retrieval set — captured before the context-token-budget truncation in `format_context()`/`format_web_context()`, so it also includes documents/web results that didn't fit in the prompt — unfiltered by citation, useful for debugging and RAG evaluation regardless of what the model cited. `citations_reported` is `true` only when the model actually emitted a `[Sources: ...]` tag (even an empty/`none` one); it's `false` when the tag was missing entirely, which is the only case where `sources` falls back to keeping everything — this lets a client tell "the model cited every source" apart from "the model didn't report citations at all". ### API Routers (`openrag/api/routers/`) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index e41cc3c45..7eefb5e64 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -121,7 +121,11 @@ async def stream_with_source_filtering( resolve citation indices; ``all_sources`` — the complete pre-truncation retrieval set — is reported separately as ``extra.all_retrieved_sources`` and defaults to ``sources`` when the caller has nothing more complete to - offer. + offer. ``extra.citations_reported`` distinguishes a genuine "the model + cited every source" from "no ``[Sources: ...]`` tag was found, so + everything was kept by default" — both leave ``sources`` covering the + full prompt-visible list, but only the former means ``citations_reported`` + is ``true``. The terminal flush (tail content + ``extra.sources``) runs exactly once after the loop on *every* termination path — a clean ``data: [DONE]``, the @@ -273,6 +277,10 @@ async def stream_with_source_filtering( extra_payload = { "sources": filtered, "all_retrieved_sources": all_sources if all_sources is not None else sources, + # Disambiguates "the model cited exactly these" from "no [Sources: ...] + # tag was found, so everything was kept" — both cases can otherwise + # leave `sources` covering every retrieved source. + "citations_reported": citations is not None, } if not saw_done: extra_payload["truncated"] = True diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index e62bfdeab..085fab734 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -689,7 +689,11 @@ async def chat( clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean chunk["extra"] = json.dumps( - {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": all_sources} + { + "sources": filter_sources_by_citations(sources, citations), + "all_retrieved_sources": all_sources, + "citations_reported": citations is not None, + } ) return chunk @@ -760,7 +764,11 @@ async def complete( clean, citations = text, None resp["choices"][0]["text"] = clean resp["extra"] = json.dumps( - {"sources": filter_sources_by_citations(sources, citations), "all_retrieved_sources": all_sources} + { + "sources": filter_sources_by_citations(sources, citations), + "all_retrieved_sources": all_sources, + "citations_reported": citations is not None, + } ) return resp diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index f1f44a83f..f3ce0c778 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -306,6 +306,22 @@ async def test_all_retrieved_sources_includes_uncited_ones(self): result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) assert _parse_finish_extra(result)["all_retrieved_sources"] == self.SOURCES + @pytest.mark.asyncio + async def test_citations_reported_true_when_model_cites_every_source(self): + """Citing every source and reporting no tag both leave `sources` equal + to the full list — `citations_reported` is the only way to tell them + apart (#847 review).""" + lines = [ + _make_chunk("Here is the answer."), + _make_chunk("\n[Sources: 1, 2, 3]"), + _make_finish(), + DONE_LINE, + ] + result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) + extra = _parse_finish_extra(result) + assert extra["sources"] == self.SOURCES + assert extra["citations_reported"] is True + @pytest.mark.asyncio async def test_content_and_finish_reason_in_same_chunk_keeps_last_token(self): """A provider that packs the final token and finish_reason into one chunk @@ -399,6 +415,9 @@ async def test_case2_llm_says_sources_none(self): result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) assert _collect_content(result) == "I cannot find this in the documents." assert _parse_finish_sources(result) == [] + # Explicit "[Sources: none]" is a reported (empty) citation set, not a + # missing tag — distinguishable from case 3 below via citations_reported. + assert _parse_finish_extra(result)["citations_reported"] is True @pytest.mark.asyncio async def test_case3_llm_no_tag_returns_all_sources(self): @@ -411,6 +430,9 @@ async def test_case3_llm_no_tag_returns_all_sources(self): result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) assert _collect_content(result) == "Answer without any sources tag." assert _parse_finish_sources(result) == self.SOURCES + # No tag at all → citations_reported is False even though `sources` + # ends up covering everything, same as an explicit "cite all" case. + assert _parse_finish_extra(result)["citations_reported"] is False @pytest.mark.asyncio async def test_structured_output_preserves_source_like_json_values(self): @@ -430,6 +452,8 @@ async def test_structured_output_preserves_source_like_json_values(self): ) assert _collect_content(result) == structured assert _parse_finish_sources(result) == self.SOURCES + # citation_protocol_active=False → citations always None → not reported. + assert _parse_finish_extra(result)["citations_reported"] is False @pytest.mark.asyncio async def test_direct_output_preserves_terminal_source_marker(self): diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index dfa07a242..f81ead806 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -660,7 +660,11 @@ async def test_chat_without_citation_keeps_retrieved_sources(): model_name="m", ) - assert json.loads(out["extra"])["sources"] == sources + extra = json.loads(out["extra"]) + assert extra["sources"] == sources + # No tag at all → not reported, even though `sources` ends up covering + # everything, same as if the model had explicitly cited all of them (#847 review). + assert extra["citations_reported"] is False @pytest.mark.asyncio @@ -675,7 +679,9 @@ async def test_chat_invalid_citation_does_not_fallback_to_unrelated_sources(): model_name="m", ) - assert json.loads(out["extra"])["sources"] == [] + extra = json.loads(out["extra"]) + assert extra["sources"] == [] + assert extra["citations_reported"] is True # a tag was present, just out of range @pytest.mark.asyncio @@ -1093,8 +1099,10 @@ async def test_complete_partition_request_keeps_context_and_filters_citations(): prepare_sources=lambda d, w: sources, ) + extra = json.loads(out["extra"]) assert out["choices"][0]["text"] == "The answer is grounded." - assert json.loads(out["extra"])["sources"] == sources + assert extra["sources"] == sources + assert extra["citations_reported"] is True answer_prompt = llm.generate_calls[0][0] assert "ctx" in answer_prompt assert "What does the report say?" in answer_prompt From c564472c0c365d939ccfee2df23bd1fe73e6a065 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 09:32:45 +0000 Subject: [PATCH 07/14] fix(rag): capture all_retrieved_sources before map-reduce replaces docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Map-reduce reassigns `docs` to LLM-generated summaries before the prompt is built, and the retrieval snapshot for all_retrieved_sources was taken after that reassignment — so on the map-reduce path it held summaries instead of anything retrieval actually returned. Move the snapshot above the map-reduce call, alongside the existing pre-truncation capture. Follow-up to f3c5b2f8, from PR #847 review. --- .../services/orchestrators/query_service.py | 16 +++++---- .../orchestrators/test_query_service.py | 34 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 085fab734..376c441cc 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -490,16 +490,20 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L return payload, [], [], [], [], False docs = [c.to_langchain() for c in chunks] - if use_map_reduce and docs: - docs = await self._map_reduce(" ".join(q.query for q in queries.query_list), docs) - # Full retrieval set, captured before the token-budget selection below - # drops anything that didn't fit in the prompt — kept separately for - # `all_retrieved_sources` (debugging/eval), while `docs`/`web_results` - # stay budget-truncated to match the citation indices the LLM sees (#847). + # Full retrieval set, captured right after retrieval — before map-reduce + # replaces `docs` with LLM-generated summaries, and before the + # token-budget selection below drops anything that didn't fit in the + # prompt. Kept separately for `all_retrieved_sources` (debugging/eval), + # while `docs`/`web_results` stay map-reduced and budget-truncated to + # match what the LLM actually saw and the citation indices it cites + # (#847 review — the map-reduce gap was called out in a follow-up pass). retrieved_docs = docs retrieved_web_results = web_results + if use_map_reduce and docs: + docs = await self._map_reduce(" ".join(q.query for q in queries.query_list), docs) + web_formatted, web_source_numbers, web_tokens = "", [], 0 web_start_index = 1 if web_results: diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index f81ead806..ed31a67e0 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -1141,6 +1141,40 @@ async def test_map_reduce_keeps_relevant_drops_irrelevant(): assert out[0].page_content == "kept" +@pytest.mark.asyncio +async def test_chat_all_retrieved_sources_survives_map_reduce_replacement(): + """#847 follow-up review: map-reduce replaces `docs` with LLM-generated + summaries before the prompt is built. all_retrieved_sources must still + reflect what retrieval actually returned, not those summaries.""" + chunks = [ + Chunk(id="c1", text="original text one", metadata={"_id": "c1"}), + Chunk(id="c2", text="original text two", metadata={"_id": "c2"}), + ] + rel1 = json.dumps({"relevancy": True, "summary": "summary one"}) + rel2 = json.dumps({"relevancy": True, "summary": "summary two"}) + answer = "Grounded answer. [Sources: 1, 2]" + svc = _svc(retrieval=FakeRetrieval(chunks=chunks), llm=FakeLLM(chat_responses=[rel1, rel2, answer])) + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {"use_map_reduce": True}}, + prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id"), "text": doc.page_content} for doc in d], + model_name="m", + ) + + extra = json.loads(out["extra"]) + # What the LLM actually saw and cited: the map-reduce summaries. + assert extra["sources"] == [ + {"id": "c1", "text": "summary one"}, + {"id": "c2", "text": "summary two"}, + ] + # The real retrieval, unreplaced by summarization. + assert extra["all_retrieved_sources"] == [ + {"id": "c1", "text": "original text one"}, + {"id": "c2", "text": "original text two"}, + ] + + # --------------------------------------------------------------------------- # # helpers # --------------------------------------------------------------------------- # From 9429a7b967f50cf15454a710589d5feed741605b Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 09:32:48 +0000 Subject: [PATCH 08/14] docs(api): document all_retrieved_sources and citations_reported The OpenAPI-facing description of the extra response field (the public contract clients actually read) still only mentioned sources. CLAUDE.md was updated for these fields but the router docstrings weren't. --- openrag/api/routers/user/chat.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index 81f93eeb7..41f7db1ea 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -446,7 +446,12 @@ def check_tokens_limit( **Response:** Returns OpenAI-compatible response with additional `extra` field containing: -- `sources`: Array of source documents with metadata and URLs +- `sources`: Array of source documents actually cited by the model (or every + retrieved source if it didn't report citations — see `citations_reported`) +- `all_retrieved_sources`: Array of every source retrieval returned, unfiltered + by citation — for debugging and evaluation +- `citations_reported`: `true` only if the model emitted a citations tag + (even an empty one); `false` means `sources` fell back to keeping everything **Streaming:** Set `stream: true` for Server-Sent Events (SSE) streaming responses. @@ -552,7 +557,12 @@ async def stream_response(): **Response:** Returns OpenAI-compatible response with additional `extra` field containing: -- `sources`: Array of source documents with metadata and URLs +- `sources`: Array of source documents actually cited by the model (or every + retrieved source if it didn't report citations — see `citations_reported`) +- `all_retrieved_sources`: Array of every source retrieval returned, unfiltered + by citation — for debugging and evaluation +- `citations_reported`: `true` only if the model emitted a citations tag + (even an empty one); `false` means `sources` fell back to keeping everything **Note:** Streaming is not supported for this endpoint. """, From 87c07b07ab087295360cd951194c58606f06a3a3 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 09:49:42 +0000 Subject: [PATCH 09/14] feat(rag): split cited/presented sources and gate the full retrieval dump MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses andyne13's design note on #847: `sources` conflated "what was cited" with "everything, because there was no tag" — no distinct field existed for what was actually shown to the LLM. Add `presented_sources` (everything shown, pre-citation-filter) and `cited_sources` (strictly what was cited, never falling back like `sources` does), so a client can render "cited" vs "consulted" and Chainlit can adopt `cited_sources` with a `presented_sources` fallback when it's empty. `sources` is left untouched for backward compatibility with existing clients (e.g. Twake). Also gate `all_retrieved_sources` behind a new `metadata.include_all_retrieved_sources` request flag (default off): dumping the full, uncapped retrieval set on every response is debug/eval telemetry that most callers don't need on the hot path. Added the `metadata` field to OpenAICompletionRequest, which was missing it entirely — completions couldn't reach any metadata flag before this, including the pre-existing spoken_style_answer. --- CLAUDE.md | 8 ++- openrag/api/routers/user/chat.py | 28 +++++++--- openrag/api/schemas/user/chat.py | 13 ++++- openrag/core/utils/source_filtering.py | 34 +++++++----- .../services/orchestrators/query_service.py | 49 ++++++++++++----- .../unit/core/utils/test_source_filtering.py | 19 ++++++- .../orchestrators/test_query_service.py | 54 ++++++++++++++++--- 7 files changed, 161 insertions(+), 44 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 50d0e33b4..46075236d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -136,7 +136,13 @@ The RAG pipeline filters out false-positive sources by having the LLM self-repor 4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every retrieved source is kept instead (a missing tag means the model didn't report citations, not that it used none) 5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client -The `extra` field in API responses is a JSON string: `{"sources": [filtered_source_list], "all_retrieved_sources": [all_retrieved_sources], "citations_reported": bool}`. `all_retrieved_sources` always carries the complete retrieval set — captured before the context-token-budget truncation in `format_context()`/`format_web_context()`, so it also includes documents/web results that didn't fit in the prompt — unfiltered by citation, useful for debugging and RAG evaluation regardless of what the model cited. `citations_reported` is `true` only when the model actually emitted a `[Sources: ...]` tag (even an empty/`none` one); it's `false` when the tag was missing entirely, which is the only case where `sources` falls back to keeping everything — this lets a client tell "the model cited every source" apart from "the model didn't report citations at all". +The `extra` field in API responses is a JSON string with these keys: + +- `sources` — legacy field, kept as-is for existing clients (e.g. Twake): cited sources, or every presented source as a fallback when no `[Sources: ...]` tag was found. +- `presented_sources` — every source actually shown to the LLM (after `format_context()`/`format_web_context()` truncation), regardless of citation. Always present; a client can fall back to this ("sources consulted") when nothing was cited. +- `cited_sources` — strictly what the model cited via the tag; unlike `sources`, this never falls back to "everything" — it's `[]` whenever no tag was found. Chainlit is expected to move to this field, falling back to `presented_sources` in its UI when `cited_sources` is empty. +- `citations_reported` (bool) — `true` only when the model actually emitted a `[Sources: ...]` tag (even an empty/`none` one); `false` when the tag was missing entirely, which is the only case where `sources` falls back to keeping everything. Lets a client tell "the model cited every source" apart from "the model didn't report citations at all". +- `all_retrieved_sources` — the complete retrieval set, captured before the context-token-budget truncation, so it also includes documents/web results that didn't fit in the prompt (and, on the map-reduce path, the original retrieved docs rather than the LLM-generated summaries). Only included when the request sets `metadata.include_all_retrieved_sources: true` — it's debug/eval telemetry, gated off by default since retrieval is uncapped up to `retriever.top_k` while the context budget only fits a handful of documents. ### API Routers (`openrag/api/routers/`) diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index 41f7db1ea..a31bcd328 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -446,12 +446,18 @@ def check_tokens_limit( **Response:** Returns OpenAI-compatible response with additional `extra` field containing: -- `sources`: Array of source documents actually cited by the model (or every - retrieved source if it didn't report citations — see `citations_reported`) -- `all_retrieved_sources`: Array of every source retrieval returned, unfiltered - by citation — for debugging and evaluation +- `sources`: Legacy field, kept for backward compatibility. Cited sources, or + every presented source as a fallback when the model didn't report citations +- `presented_sources`: Array of every source actually shown to the model + (after context-budget truncation), regardless of what it cited +- `cited_sources`: Array of only the sources the model explicitly cited; empty + whenever no citations tag was found (never falls back like `sources` does) - `citations_reported`: `true` only if the model emitted a citations tag (even an empty one); `false` means `sources` fell back to keeping everything +- `all_retrieved_sources`: Array of every source retrieval returned, unfiltered + by citation or context-budget truncation — only included when the request's + `metadata.include_all_retrieved_sources` is `true` (off by default; this is + debug/evaluation telemetry and can be large) **Streaming:** Set `stream: true` for Server-Sent Events (SSE) streaming responses. @@ -557,12 +563,18 @@ async def stream_response(): **Response:** Returns OpenAI-compatible response with additional `extra` field containing: -- `sources`: Array of source documents actually cited by the model (or every - retrieved source if it didn't report citations — see `citations_reported`) -- `all_retrieved_sources`: Array of every source retrieval returned, unfiltered - by citation — for debugging and evaluation +- `sources`: Legacy field, kept for backward compatibility. Cited sources, or + every presented source as a fallback when the model didn't report citations +- `presented_sources`: Array of every source actually shown to the model + (after context-budget truncation), regardless of what it cited +- `cited_sources`: Array of only the sources the model explicitly cited; empty + whenever no citations tag was found (never falls back like `sources` does) - `citations_reported`: `true` only if the model emitted a citations tag (even an empty one); `false` means `sources` fell back to keeping everything +- `all_retrieved_sources`: Array of every source retrieval returned, unfiltered + by citation or context-budget truncation — only included when the request's + `metadata.include_all_retrieved_sources` is `true` (off by default; this is + debug/evaluation telemetry and can be large) **Note:** Streaming is not supported for this endpoint. """, diff --git a/openrag/api/schemas/user/chat.py b/openrag/api/schemas/user/chat.py index 4ea5ab7e0..82034dc3a 100644 --- a/openrag/api/schemas/user/chat.py +++ b/openrag/api/schemas/user/chat.py @@ -43,8 +43,10 @@ class OpenAIChatCompletionRequest(BaseModel): "spoken_style_answer": False, "websearch": False, "llm_override": None, + "include_all_retrieved_sources": False, }, - description="Extra custom parameters. Supports 'llm_override' object with an optional 'model' to override the downstream model name. The LLM endpoint and credentials are fixed by server configuration and cannot be overridden by the client.", + description="Extra custom parameters. Supports 'llm_override' object with an optional 'model' to override the downstream model name. The LLM endpoint and credentials are fixed by server configuration and cannot be overridden by the client. " + "'include_all_retrieved_sources' (default false) adds the full, unfiltered retrieval set to the response's extra.all_retrieved_sources — off by default since it can be large; opt in only for debugging/evaluation.", ) @model_validator(mode="after") @@ -84,3 +86,12 @@ class OpenAICompletionRequest(BaseModel): stream: bool | None = Field(False) temperature: float | None = Field(0.3) top_p: float | None = Field(1.0) + metadata: dict[str, Any] | None = Field( + { + "spoken_style_answer": False, + "llm_override": None, + "include_all_retrieved_sources": False, + }, + description="Extra custom parameters. Supports 'llm_override' object with an optional 'model' to override the downstream model name. The LLM endpoint and credentials are fixed by server configuration and cannot be overridden by the client. " + "'include_all_retrieved_sources' (default false) adds the full, unfiltered retrieval set to the response's extra.all_retrieved_sources — off by default since it can be large; opt in only for debugging/evaluation.", + ) diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index 7eefb5e64..7a5a693a5 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -114,18 +114,28 @@ async def stream_with_source_filtering( *, citation_protocol_active: bool = True, all_sources: list | None = None, + include_all_retrieved: bool = False, ): """Process an LLM SSE stream and, when active, strip source tags. ``sources`` is the prompt-visible (context-budget-truncated) list used to - resolve citation indices; ``all_sources`` — the complete pre-truncation - retrieval set — is reported separately as ``extra.all_retrieved_sources`` - and defaults to ``sources`` when the caller has nothing more complete to - offer. ``extra.citations_reported`` distinguishes a genuine "the model - cited every source" from "no ``[Sources: ...]`` tag was found, so - everything was kept by default" — both leave ``sources`` covering the - full prompt-visible list, but only the former means ``citations_reported`` - is ``true``. + resolve citation indices — reported as-is via ``extra.sources`` (legacy, + kept for existing clients: cited sources, or every presented source as a + fallback when no tag was found) and ``extra.presented_sources`` (always + the raw pre-filter list, so a client can render "sources consulted" when + nothing was cited). ``extra.cited_sources`` is the strict version: only + what the model actually cited, empty whenever no tag was found — never + falling back to "everything" the way ``sources`` does. + ``extra.citations_reported`` disambiguates those two empty/full states + on the wire: ``true`` only when a ``[Sources: ...]`` tag (even an empty + one) was actually found. + + ``all_sources`` — the complete pre-truncation retrieval set — is reported + as ``extra.all_retrieved_sources`` only when ``include_all_retrieved`` is + true (it defaults to ``sources`` when the caller has nothing more + complete to offer). It's gated because a full retrieval dump on every + response is debug/eval telemetry, not something most callers need on the + hot path. The terminal flush (tail content + ``extra.sources``) runs exactly once after the loop on *every* termination path — a clean ``data: [DONE]``, the @@ -276,12 +286,12 @@ async def stream_with_source_filtering( filtered = filter_sources_by_citations(sources, citations) extra_payload = { "sources": filtered, - "all_retrieved_sources": all_sources if all_sources is not None else sources, - # Disambiguates "the model cited exactly these" from "no [Sources: ...] - # tag was found, so everything was kept" — both cases can otherwise - # leave `sources` covering every retrieved source. + "presented_sources": sources, + "cited_sources": filtered if citations is not None else [], "citations_reported": citations is not None, } + if include_all_retrieved: + extra_payload["all_retrieved_sources"] = all_sources if all_sources is not None else sources if not saw_done: extra_payload["truncated"] = True logger.warning( diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 376c441cc..734c0aaee 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -663,6 +663,7 @@ async def chat( ) -> dict: """Non-streaming chat completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} + include_all_retrieved = bool(metadata.get("include_all_retrieved_sources", False)) llm = self._resolve_llm(partitions) citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): @@ -677,7 +678,10 @@ async def chat( citation_protocol_active, ) = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) - all_sources = prepare_sources(retrieved_docs, retrieved_web_results) + # `all_retrieved_sources` is debug/eval telemetry, not needed by most + # callers — skip building it (and calling prepare_sources on the full, + # uncapped retrieval set) unless the caller opted in (#847 review). + all_sources = prepare_sources(retrieved_docs, retrieved_web_results) if include_all_retrieved else None structured_output = _is_structured_output(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) @@ -693,11 +697,7 @@ async def chat( clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean chunk["extra"] = json.dumps( - { - "sources": filter_sources_by_citations(sources, citations), - "all_retrieved_sources": all_sources, - "citations_reported": citations is not None, - } + _build_extra_payload(sources, citations, all_sources, include_all_retrieved=include_all_retrieved) ) return chunk @@ -711,6 +711,7 @@ async def chat_stream( ) -> AsyncIterator[str]: """Streaming chat completion → SSE strings with filtered sources.""" metadata = payload.get("metadata") or {} + include_all_retrieved = bool(metadata.get("include_all_retrieved_sources", False)) llm = self._resolve_llm(partitions) citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): @@ -725,7 +726,7 @@ async def chat_stream( citation_protocol_active, ) = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) - all_sources = prepare_sources(retrieved_docs, retrieved_web_results) + all_sources = prepare_sources(retrieved_docs, retrieved_web_results) if include_all_retrieved else None structured_output = _is_structured_output(payload) payload["messages"] = self._sanitize_messages(payload["messages"]) @@ -735,6 +736,7 @@ async def chat_stream( sources, model_name, all_sources=all_sources, + include_all_retrieved=include_all_retrieved, citation_protocol_active=citation_protocol_active and not structured_output, ): yield sse_line @@ -747,6 +749,8 @@ async def complete( prepare_sources: PrepareSources, ) -> dict: """Non-streaming text completion → finalized OpenAI dict.""" + metadata = payload.get("metadata") or {} + include_all_retrieved = bool(metadata.get("include_all_retrieved_sources", False)) llm = self._resolve_llm(partitions) citation_protocol_active = partitions is not None if partitions is None: @@ -754,7 +758,7 @@ async def complete( else: payload, docs, retrieved_docs = await self._prepare_completions(partitions, payload, llm) sources = prepare_sources(docs, []) - all_sources = prepare_sources(retrieved_docs, []) + all_sources = prepare_sources(retrieved_docs, []) if include_all_retrieved else None structured_output = _is_structured_output(payload) resp = await llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) @@ -768,11 +772,7 @@ async def complete( clean, citations = text, None resp["choices"][0]["text"] = clean resp["extra"] = json.dumps( - { - "sources": filter_sources_by_citations(sources, citations), - "all_retrieved_sources": all_sources, - "citations_reported": citations is not None, - } + _build_extra_payload(sources, citations, all_sources, include_all_retrieved=include_all_retrieved) ) return resp @@ -823,4 +823,27 @@ def _is_structured_output(payload: dict) -> bool: return isinstance(response_format, dict) and response_format.get("type") in {"json_object", "json_schema"} +def _build_extra_payload( + sources: list, + citations: set[int] | None, + all_sources: list | None, + *, + include_all_retrieved: bool, +) -> dict: + """Shared ``extra`` shape for ``chat``/``complete`` (mirrors + ``stream_with_source_filtering``'s payload, minus the streaming-only + ``truncated`` flag) so the three response paths can't drift apart. + """ + filtered = filter_sources_by_citations(sources, citations) + payload = { + "sources": filtered, + "presented_sources": sources, + "cited_sources": filtered if citations is not None else [], + "citations_reported": citations is not None, + } + if include_all_retrieved: + payload["all_retrieved_sources"] = all_sources + return payload + + __all__ = ["QueryService", "RAGMODE"] diff --git a/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index f3ce0c778..e7d4f6b50 100644 --- a/tests/unit/core/utils/test_source_filtering.py +++ b/tests/unit/core/utils/test_source_filtering.py @@ -295,8 +295,9 @@ async def test_case1_llm_cites_specific_sources(self): assert _parse_finish_sources(result) == [{"file": "a.pdf"}, {"file": "c.pdf"}] @pytest.mark.asyncio - async def test_all_retrieved_sources_includes_uncited_ones(self): - """`all_retrieved_sources` always carries every candidate, unfiltered by citation.""" + async def test_all_retrieved_sources_omitted_unless_opted_in(self): + """all_retrieved_sources is debug/eval telemetry — absent by default, + only included when the caller passes include_all_retrieved=True.""" lines = [ _make_chunk("Here is the answer."), _make_chunk("\n[Sources: 1, 3]"), @@ -304,6 +305,20 @@ async def test_all_retrieved_sources_includes_uncited_ones(self): DONE_LINE, ] result = await _collect(stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model")) + assert "all_retrieved_sources" not in _parse_finish_extra(result) + + @pytest.mark.asyncio + async def test_all_retrieved_sources_includes_uncited_ones_when_opted_in(self): + """`all_retrieved_sources` always carries every candidate, unfiltered by citation.""" + lines = [ + _make_chunk("Here is the answer."), + _make_chunk("\n[Sources: 1, 3]"), + _make_finish(), + DONE_LINE, + ] + result = await _collect( + stream_with_source_filtering(_fake_stream(lines), self.SOURCES, "test-model", include_all_retrieved=True) + ) assert _parse_finish_extra(result)["all_retrieved_sources"] == self.SOURCES @pytest.mark.asyncio diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index ed31a67e0..073e333cf 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -454,13 +454,39 @@ async def test_chat_with_partition_retrieves_and_filters_sources(): sources = [{"source_type": "document", "n": 1}, {"source_type": "document", "n": 2}] out = await svc.chat( partitions=["p"], - payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"include_all_retrieved_sources": True}, + }, prepare_sources=lambda d, w: sources, model_name="m", ) extra = json.loads(out["extra"]) assert extra["sources"] == [{"source_type": "document", "n": 1}] # only cited source 1 - assert extra["all_retrieved_sources"] == sources # unfiltered, everything retrieved + assert extra["presented_sources"] == sources # everything shown to the model + assert extra["cited_sources"] == [{"source_type": "document", "n": 1}] # strictly what was cited + assert extra["all_retrieved_sources"] == sources # opted in: unfiltered, everything retrieved + + +@pytest.mark.asyncio +async def test_chat_all_retrieved_sources_omitted_by_default(): + """#847 follow-up: all_retrieved_sources is debug/eval telemetry, gated + behind metadata.include_all_retrieved_sources — absent unless requested.""" + svc = _svc(llm=FakeLLM(chat_responses=["answer [Sources: 1]"])) + sources = [{"source_type": "document", "n": 1}] + + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: sources, + model_name="m", + ) + + extra = json.loads(out["extra"]) + assert "all_retrieved_sources" not in extra + assert extra["sources"] == sources + assert extra["presented_sources"] == sources + assert extra["cited_sources"] == sources @pytest.mark.asyncio @@ -477,7 +503,10 @@ async def test_chat_all_retrieved_sources_survives_context_budget_truncation(): out = await svc.chat( partitions=["p"], - payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"include_all_retrieved_sources": True}, + }, prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id")} for doc in d], model_name="m", ) @@ -505,7 +534,10 @@ async def test_chat_stream_all_retrieved_sources_survives_context_budget_truncat line async for line in svc.chat_stream( partitions=["p"], - payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"include_all_retrieved_sources": True}, + }, prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id")} for doc in d], model_name="m", ) @@ -532,7 +564,7 @@ async def test_complete_all_retrieved_sources_survives_context_budget_truncation out = await svc.complete( partitions=["p"], - payload={"prompt": "q"}, + payload={"prompt": "q", "metadata": {"include_all_retrieved_sources": True}}, prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id")} for doc in d], ) @@ -665,6 +697,11 @@ async def test_chat_without_citation_keeps_retrieved_sources(): # No tag at all → not reported, even though `sources` ends up covering # everything, same as if the model had explicitly cited all of them (#847 review). assert extra["citations_reported"] is False + # presented_sources always shows what the model saw; cited_sources — unlike + # legacy `sources` — never falls back and stays empty when nothing was + # actually reported cited. + assert extra["presented_sources"] == sources + assert extra["cited_sources"] == [] @pytest.mark.asyncio @@ -780,7 +817,7 @@ async def test_structured_websearch_returns_only_sources_included_in_context(): partitions=None, payload={ "messages": [{"role": "user", "content": "Question"}], - "metadata": {"websearch": True}, + "metadata": {"websearch": True, "include_all_retrieved_sources": True}, "response_format": {"type": "json_object"}, }, prepare_sources=lambda _docs, results: [{"url": result.url} for result in results], @@ -1157,7 +1194,10 @@ async def test_chat_all_retrieved_sources_survives_map_reduce_replacement(): out = await svc.chat( partitions=["p"], - payload={"messages": [{"role": "user", "content": "q"}], "metadata": {"use_map_reduce": True}}, + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"use_map_reduce": True, "include_all_retrieved_sources": True}, + }, prepare_sources=lambda d, w: [{"id": doc.metadata.get("_id"), "text": doc.page_content} for doc in d], model_name="m", ) From 4c71a7e58e151dc97a1f80ba1c1859182dfb268e Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 09:55:59 +0000 Subject: [PATCH 10/14] feat(chainlit): adopt cited_sources, falling back to presented_sources Chainlit was still reading the legacy extra.sources field. Switch to cited_sources (strictly what the model cited) with a presented_sources fallback when nothing was cited, per the plan discussed on PR #847. Also add the complete()-path no-citation test flagged as a coverage gap in that review's nits. --- openrag/app_front.py | 7 +++++- .../orchestrators/test_query_service.py | 22 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/openrag/app_front.py b/openrag/app_front.py index 7b8efd788..5b33a6357 100644 --- a/openrag/app_front.py +++ b/openrag/app_front.py @@ -690,7 +690,12 @@ async def on_message(message: cl.Message): async for chunk in stream: if chunk.extra: extra = json.loads(chunk.extra) - if "sources" in extra: + if "cited_sources" in extra: + # Strictly what the model cited; fall back to + # everything it was shown when nothing was cited + # (e.g. it didn't report citations at all). + sources = extra["cited_sources"] or extra.get("presented_sources") + elif "sources" in extra: sources = extra["sources"] if chunk.choices and chunk.choices[0].delta.content: diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 073e333cf..b2f0d3a77 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -1145,6 +1145,28 @@ async def test_complete_partition_request_keeps_context_and_filters_citations(): assert "What does the report say?" in answer_prompt +@pytest.mark.asyncio +async def test_complete_without_citation_keeps_retrieved_sources(): + """complete()'s equivalent of test_chat_without_citation_keeps_retrieved_sources + (#847 review: test coverage asymmetry between chat and complete).""" + llm = FakeLLM(gen_text="A general answer with no citation marker.") + svc = _svc(llm=llm) + sources = [{"source_type": "document", "filename": "unrelated.pdf"}] + + out = await svc.complete( + partitions=["p"], + payload={"prompt": "What does the report say?"}, + prepare_sources=lambda d, w: sources, + ) + + extra = json.loads(out["extra"]) + assert out["choices"][0]["text"] == "A general answer with no citation marker." + assert extra["sources"] == sources + assert extra["citations_reported"] is False + assert extra["presented_sources"] == sources + assert extra["cited_sources"] == [] + + @pytest.mark.asyncio async def test_chat_stream_yields_sse_and_done(): svc = _svc(llm=FakeLLM()) From 60887453db14a4cd641997a0e18ea3afa9f8baa6 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 10:29:11 +0000 Subject: [PATCH 11/14] docs: note the reranker.top_k / context-budget mismatch in code Point future readers at #851 from the two spots where the gap actually lives: the token-budget sizing in query_service.py, and the pipeline code where reranker_top_k is read but never applied as a final cutoff. --- openrag/core/retrieval/pipeline.py | 7 +++++++ openrag/services/orchestrators/query_service.py | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/openrag/core/retrieval/pipeline.py b/openrag/core/retrieval/pipeline.py index 47574e1b7..ee9687368 100644 --- a/openrag/core/retrieval/pipeline.py +++ b/openrag/core/retrieval/pipeline.py @@ -130,6 +130,13 @@ async def retrieve_docs( if self.reranker_enabled: chunks = await _rerank_chunks(self.reranker, query.query, chunks) + # `reranker_top_k` is NOT applied here as a final cutoff — only + # `top_k` (an explicit caller-supplied value, e.g. map-reduce's + # max_total_documents) truncates. On the common no-`top_k` chat path + # this returns everything reranked (up to the retriever's own + # top_k), which callers sizing a token budget off reranker_top_k + # should not assume is bounded by it. Tracked separately: + # https://github.com/linagora/openrag/issues/851 if top_k is not None: chunks = chunks[:top_k] return chunks diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 734c0aaee..8133f9028 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -143,6 +143,12 @@ def __init__( config.rag.chat_history_depth if config.rag.chat_history_depth >= 1 else self._CHAT_HISTORY_DEPTH_DEFAULT ) self._max_contextualized_query_len = config.rag.max_contextualized_query_len + # Sized on the assumption that retrieval returns ~reranker.top_k chunks, + # but reranker_top_k is never actually applied as a cutoff in + # RetrieverPipeline.retrieve_docs() on the no-map-reduce path — retrieval + # can return up to retriever.top_k candidates, so this budget (not + # reranker.top_k) is what actually determines how many reach the prompt. + # Tracked separately: https://github.com/linagora/openrag/issues/851 self._max_context_tokens = config.reranker.top_k * config.chunker.chunk_size mr = config.map_reduce From e40cd0740a6ee608979ac7081f96fc9acdcefba9 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Fri, 31 Jul 2026 12:45:17 +0000 Subject: [PATCH 12/14] fix(docker): make venv writes group-writable to survive APP_UID drift The openrag_venv named volume persists across container recreations, but the openrag image's non-root APP_UID can differ between a locally built image (defaults to the host UID) and a pulled/CI-built image (defaults to 10001) even though both share the same image tag. The default umask (022) makes `uv sync` create new venv entries writable only by the exact UID that wrote them, so a later sync under a different APP_UID fails to remove/replace files (e.g. the editable install's __editable__*.pth on every version bump) with "Permission denied". Force umask 002 before running uv so venv entries stay group (GID 0) writable, letting any APP_UID that shares that group resync. --- infra/scripts/entrypoint.sh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/infra/scripts/entrypoint.sh b/infra/scripts/entrypoint.sh index ea30bd732..3df3fbc12 100644 --- a/infra/scripts/entrypoint.sh +++ b/infra/scripts/entrypoint.sh @@ -30,6 +30,15 @@ if [ "$(id -u)" = "0" ]; then exec setpriv --reuid "$APP_UID" --regid 0 --clear-groups /app/entrypoint.sh "$@" fi +# The persisted openrag_venv volume can be re-synced by different APP_UIDs +# across the container's lifetime (e.g. a locally-built image bakes in the +# host UID while a pulled/CI-built image bakes in the Dockerfile default) — +# both share GID 0. Default umask (022) makes `uv sync` create new venv +# entries owner-writable only, so a later sync under a different UID can't +# remove/replace them ("Permission denied" on __editable__*.pth). Force +# group-writable new files/dirs so any GID-0 UID can always resync. +umask 002 + ENV_ARGS=() if [[ -n "${SHARED_ENV}" ]]; then ENV_ARGS+=("--env-file=${SHARED_ENV}") From 14e377f72dc38a62161d5af610c1a9457497b9db Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 08:43:48 +0000 Subject: [PATCH 13/14] fix(helm): apply umask 002 to Ray init-container venv sync The Ray init container syncs the shared venv volume directly with uv sync and didn't inherit entrypoint.sh's umask 002, so a later sync under a different APP_UID could still fail to replace owner-only editable venv entries left by this path. --- infra/charts/openrag-stack/templates/raycluster.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/infra/charts/openrag-stack/templates/raycluster.yaml b/infra/charts/openrag-stack/templates/raycluster.yaml index af2ed6a80..a63a01068 100644 --- a/infra/charts/openrag-stack/templates/raycluster.yaml +++ b/infra/charts/openrag-stack/templates/raycluster.yaml @@ -34,6 +34,10 @@ spec: - sh - -c - | + # Keep venv writes group-writable (GID 0) so a later sync under a + # different APP_UID can still remove/replace owner-only entries + # left by a previous sync. See infra/scripts/entrypoint.sh. + umask 002 if [ -f /app/.venv/.ready ]; then echo "Existing env detected, skipping install." else From 5eba9fc9b89d5197b288c69c784f5bfd1673a7c0 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Mon, 24 Aug 2026 12:27:48 +0000 Subject: [PATCH 14/14] fix: address CodeRabbit review findings on PR #848 - Fix CLAUDE.md wording: filter_sources_by_citations() falls back to the presented source list, not the full retrieved set. - Note in entrypoint.sh/raycluster.yaml that the umask 002 fix only covers files written after upgrading; an existing openrag_venv volume/PVC still needs to be recreated once. - Require a real JSON boolean for metadata.include_all_retrieved_sources instead of any truthy value (e.g. the string "false" was enabling it). --- CLAUDE.md | 2 +- infra/charts/openrag-stack/templates/raycluster.yaml | 3 +++ infra/scripts/entrypoint.sh | 3 +++ openrag/services/orchestrators/query_service.py | 6 +++--- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 46075236d..e180b0db2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -133,7 +133,7 @@ The RAG pipeline filters out false-positive sources by having the LLM self-repor 1. `format_context()` (`openrag/core/prompts/chat_prompt_builder.py`) numbers each source (`[Source 1]`, `[Source 2]`, ...) in the context and returns `(formatted_text, included_indices)` — the indices track which docs fit within the token budget 2. Prompt templates (`openrag/prompts/templates/*.txt`) instruct the LLM to append `[Sources: 1, 3, 5]` at the end of its response 3. `extract_and_strip_sources_block()` (`openrag/core/utils/source_filtering.py`) strips this tag from the response before sending to the client -4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every retrieved source is kept instead (a missing tag means the model didn't report citations, not that it used none) +4. `filter_sources_by_citations()` (`openrag/core/utils/source_filtering.py`) filters the source metadata to only include cited sources; if no `[Sources: ...]` tag is found at all, every presented source is kept instead (a missing tag means the model didn't report citations, not that it used none) 5. For streaming, the OpenAI router buffers the last 100 chars to catch the sources tag before it reaches the client The `extra` field in API responses is a JSON string with these keys: diff --git a/infra/charts/openrag-stack/templates/raycluster.yaml b/infra/charts/openrag-stack/templates/raycluster.yaml index a63a01068..b8474cba9 100644 --- a/infra/charts/openrag-stack/templates/raycluster.yaml +++ b/infra/charts/openrag-stack/templates/raycluster.yaml @@ -37,6 +37,9 @@ spec: # Keep venv writes group-writable (GID 0) so a later sync under a # different APP_UID can still remove/replace owner-only entries # left by a previous sync. See infra/scripts/entrypoint.sh. + # NOTE: only fixes files written from here on — a PVC provisioned + # before this fix still has owner-only entries and needs to be + # recreated once after upgrading. umask 002 if [ -f /app/.venv/.ready ]; then echo "Existing env detected, skipping install." diff --git a/infra/scripts/entrypoint.sh b/infra/scripts/entrypoint.sh index 3df3fbc12..38298cc98 100644 --- a/infra/scripts/entrypoint.sh +++ b/infra/scripts/entrypoint.sh @@ -37,6 +37,9 @@ fi # entries owner-writable only, so a later sync under a different UID can't # remove/replace them ("Permission denied" on __editable__*.pth). Force # group-writable new files/dirs so any GID-0 UID can always resync. +# NOTE: this only fixes files written from here on — an openrag_venv volume +# that already existed before this fix still has owner-only entries and needs +# to be recreated once after upgrading (delete the volume / let it resync). umask 002 ENV_ARGS=() diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 8133f9028..882c8eeff 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -669,7 +669,7 @@ async def chat( ) -> dict: """Non-streaming chat completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} - include_all_retrieved = bool(metadata.get("include_all_retrieved_sources", False)) + include_all_retrieved = metadata.get("include_all_retrieved_sources") is True llm = self._resolve_llm(partitions) citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): @@ -717,7 +717,7 @@ async def chat_stream( ) -> AsyncIterator[str]: """Streaming chat completion → SSE strings with filtered sources.""" metadata = payload.get("metadata") or {} - include_all_retrieved = bool(metadata.get("include_all_retrieved_sources", False)) + include_all_retrieved = metadata.get("include_all_retrieved_sources") is True llm = self._resolve_llm(partitions) citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): @@ -756,7 +756,7 @@ async def complete( ) -> dict: """Non-streaming text completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} - include_all_retrieved = bool(metadata.get("include_all_retrieved_sources", False)) + include_all_retrieved = metadata.get("include_all_retrieved_sources") is True llm = self._resolve_llm(partitions) citation_protocol_active = partitions is not None if partitions is None: