diff --git a/CLAUDE.md b/CLAUDE.md index 07590973b..e180b0db2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -133,10 +133,16 @@ 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 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: `{"sources": [filtered_source_list]}`. +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/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/templates/raycluster.yaml b/infra/charts/openrag-stack/templates/raycluster.yaml index af2ed6a80..b8474cba9 100644 --- a/infra/charts/openrag-stack/templates/raycluster.yaml +++ b/infra/charts/openrag-stack/templates/raycluster.yaml @@ -34,6 +34,13 @@ 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. + # 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." else 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/infra/scripts/entrypoint.sh b/infra/scripts/entrypoint.sh index ea30bd732..38298cc98 100644 --- a/infra/scripts/entrypoint.sh +++ b/infra/scripts/entrypoint.sh @@ -30,6 +30,18 @@ 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. +# 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=() if [[ -n "${SHARED_ENV}" ]]; then ENV_ARGS+=("--env-file=${SHARED_ENV}") diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index 81f93eeb7..a31bcd328 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -446,7 +446,18 @@ def check_tokens_limit( **Response:** Returns OpenAI-compatible response with additional `extra` field containing: -- `sources`: Array of source documents with metadata and URLs +- `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. @@ -552,7 +563,18 @@ async def stream_response(): **Response:** Returns OpenAI-compatible response with additional `extra` field containing: -- `sources`: Array of source documents with metadata and URLs +- `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/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/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/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index ff7b52833..7a5a693a5 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,11 +112,31 @@ async def stream_with_source_filtering( model_name: str, buffer_size: int | None = None, *, - allow_uncited_sources: bool = False, 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 — 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 upstream closing the connection without one, or the upstream generator @@ -263,8 +283,15 @@ async def stream_with_source_filtering( else: final_clean, citations = pending, None - filtered = filter_sources_by_citations(sources, citations, allow_uncited=allow_uncited_sources) - extra_payload = {"sources": filtered} + filtered = filter_sources_by_citations(sources, citations) + extra_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: + 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 7c9287b23..882c8eeff 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 @@ -471,7 +477,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,9 +493,20 @@ 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] + + # 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) @@ -538,7 +555,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 +575,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 +584,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 +604,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 @@ -648,14 +669,26 @@ async def chat( ) -> dict: """Non-streaming chat completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} + 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): - 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) - structured_output = _allows_uncited_sources(payload) + # `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"]) chunk = await llm.chat(payload["messages"], **_sampling(payload)) @@ -670,13 +703,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, - allow_uncited=structured_output, - ) - } + _build_extra_payload(sources, citations, all_sources, include_all_retrieved=include_all_retrieved) ) return chunk @@ -690,14 +717,23 @@ async def chat_stream( ) -> AsyncIterator[str]: """Streaming chat completion → SSE strings with filtered sources.""" metadata = payload.get("metadata") or {} + 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): - 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) - structured_output = _allows_uncited_sources(payload) + 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"]) llm_stream = llm.stream_chat(payload["messages"], **_sampling(payload)) @@ -705,7 +741,8 @@ async def chat_stream( llm_stream, sources, model_name, - allow_uncited_sources=structured_output, + all_sources=all_sources, + include_all_retrieved=include_all_retrieved, citation_protocol_active=citation_protocol_active and not structured_output, ): yield sse_line @@ -718,14 +755,17 @@ async def complete( prepare_sources: PrepareSources, ) -> dict: """Non-streaming text completion → finalized OpenAI dict.""" + metadata = payload.get("metadata") or {} + 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: - 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, []) - structured_output = _allows_uncited_sources(payload) + 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")) text = resp.get("choices", [{}])[0].get("text", "") or "" @@ -738,13 +778,7 @@ async def complete( clean, citations = text, None resp["choices"][0]["text"] = clean resp["extra"] = json.dumps( - { - "sources": filter_sources_by_citations( - sources, - citations, - allow_uncited=structured_output, - ) - } + _build_extra_payload(sources, citations, all_sources, include_all_retrieved=include_all_retrieved) ) return resp @@ -789,10 +823,33 @@ 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"} +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/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/tests/unit/core/utils/test_source_filtering.py b/tests/unit/core/utils/test_source_filtering.py index 877587542..e7d4f6b50 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): @@ -298,6 +294,49 @@ 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_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]"), + _make_finish(), + 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 + 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 @@ -391,10 +430,13 @@ 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_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,24 +444,10 @@ 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 + # 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): @@ -434,12 +462,13 @@ 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, ) ) 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): @@ -514,8 +543,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..b2f0d3a77 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 @@ -452,14 +452,125 @@ async def test_chat_direct_mode_preserves_literal_terminal_sources_marker(): async def test_chat_with_partition_retrieves_and_filters_sources(): svc = _svc(llm=FakeLLM(chat_responses=["answer [Sources: 1]"])) sources = [{"source_type": "document", "n": 1}, {"source_type": "document", "n": 2}] + out = await svc.chat( + partitions=["p"], + 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["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", ) - filtered = json.loads(out["extra"])["sources"] - assert filtered == [{"source_type": "document", "n": 1}] # only cited source 1 + + 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 +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": {"include_all_retrieved_sources": True}, + }, + 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": {"include_all_retrieved_sources": True}, + }, + 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", "metadata": {"include_all_retrieved_sources": True}}, + 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 @@ -487,7 +598,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 +680,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 +692,16 @@ async def test_chat_without_citation_does_not_attribute_retrieved_sources(): model_name="m", ) - assert json.loads(out["extra"])["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 + # 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 @@ -595,7 +716,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 @@ -694,14 +817,21 @@ 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], 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 @@ -733,7 +863,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 @@ -759,7 +889,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": {}} ) @@ -966,7 +1096,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"] == [] @@ -1006,13 +1136,37 @@ 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 +@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()) @@ -1046,6 +1200,43 @@ 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, "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", + ) + + 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 # --------------------------------------------------------------------------- # @@ -1226,7 +1417,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": {}} ) 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" },