feat: filter RAG sources by LLM citations to reduce false positives

CLAUDE.md

@@ -98,6 +98,18 @@ Each file type has a dedicated loader that converts to markdown:
 - PDF/DOCX/PPTX: Extract binary image data from file, pass to VLM directly
 - Markdown: Parse image URLs from text; HTTP URLs require `IMAGE_CAPTIONING_URL=true`
 
+### Source Citation Filtering
+
+The RAG pipeline filters out false-positive sources by having the LLM self-report which sources it actually used:
+
+1. `format_context()` (`openrag/components/utils.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 (`prompts/example1/*.txt`) instruct the LLM to append `[Sources: 1, 3, 5]` at the end of its response
+3. `extract_and_strip_sources_block()` strips this tag from the response before sending to the client
+4. `filter_sources_by_citations()` filters the source metadata to only include cited sources (falls back to all sources if none match)
+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]}`.
+
 ### API Routers (`openrag/routers/`)
 
 - `openai.py` - OpenAI-compatible `/v1/chat/completions` endpoint
@@ -224,7 +236,7 @@ Environment variables override config values (see `.env.example`).
 act -j api-tests -W .github/workflows/api_tests.yml --bind
 ```
 
-**Mock VLLM for CI:** `.github/workflows/api_tests/mock_vllm.py` provides fake embeddings and completions endpoints for testing without a real LLM.
+**Mock VLLM for CI:** `.github/workflows/api_tests/mock_vllm.py` provides fake embeddings and completions endpoints (streaming and non-streaming) for testing without a real LLM. Pydantic request models use `ConfigDict(extra="allow")` to accept vendor-specific fields like `extra_body`.
 
 ## Key Patterns
 

openrag/app_front.py

@@ -233,9 +233,10 @@ async def on_message(message: cl.Message):
             # Stream the response using OpenAI client directly
             stream = await client.chat.completions.create(**data)
             async for chunk in stream:
-                if sources is None:
+                if chunk.extra:
                     extra = json.loads(chunk.extra)
-                    sources = extra["sources"]
+                    if "sources" in extra:
+                        sources = extra["sources"]
 
                 if chunk.choices and chunk.choices[0].delta.content:
                     token = chunk.choices[0].delta.content

openrag/components/pipeline.py

@@ -148,7 +148,8 @@ async def _prepare_for_chat_completion(self, partition: list[str], payload: dict
             docs = await self.map_reduce.map(query=query, chunks=docs)
 
         # 3. Format the retrieved docs
-        context, n_docs = format_context(docs, max_context_tokens=self.max_context_tokens)
+        context, included_indices = format_context(docs, max_context_tokens=self.max_context_tokens)
+        docs = [docs[i] for i in included_indices]
 
         # 4. prepare the output
         messages: list = copy.deepcopy(messages)
@@ -164,7 +165,7 @@ async def _prepare_for_chat_completion(self, partition: list[str], payload: dict
             },
         )
         payload["messages"] = messages
-        return payload, docs[:n_docs]
+        return payload, docs
 
     async def _prepare_for_completions(self, partition: list[str], payload: dict):
         prompt = payload["prompt"]
@@ -175,18 +176,19 @@ async def _prepare_for_completions(self, partition: list[str], payload: dict):
         docs = await self.retriever_pipeline.retrieve_docs(partition=partition, query=query)
 
         # 3. Format the retrieved docs
-        context, n_docs = format_context(docs, max_context_tokens=self.max_context_tokens)
+        context, included_indices = format_context(docs, max_context_tokens=self.max_context_tokens)
+        docs = [docs[i] for i in included_indices]
 
         # 4. prepare the output
         if docs:
             prompt = f"""Given the content
             {context}
             Complete the following prompt: {prompt}
-            """
+            At the very end of your response, on a new line, list which source numbers you used: [Sources: 1, 3]"""
 
         payload["prompt"] = prompt
 
-        return payload, docs[:n_docs]
+        return payload, docs
 
     async def completions(self, partition: list[str], payload: dict):
         if partition is None:

openrag/components/test_source_filtering.py

@@ -0,0 +1,122 @@
+"""Tests for source citation extraction and filtering utilities."""
+
+from components.utils import extract_and_strip_sources_block, filter_sources_by_citations
+
+
+class TestExtractAndStripSourcesBlock:
+    def test_basic_extraction(self):
+        text = "Answer text\n[Sources: 1, 3]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {1, 3}
+
+    def test_single_source(self):
+        text = "Answer text\n[Source: 2]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {2}
+
+    def test_many_sources(self):
+        text = "Answer text\n[Sources: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
+
+    def test_no_sources_block(self):
+        text = "Answer with no block"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer with no block"
+        assert citations == set()
+
+    def test_sources_with_trailing_whitespace(self):
+        text = "Answer text\n[Sources: 1, 3]   "
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {1, 3}
+
+    def test_sources_with_extra_spaces(self):
+        text = "Answer text\n[Sources:  1 ,  3 , 5 ]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {1, 3, 5}
+
+    def test_multiline_answer(self):
+        text = "Line 1\n\nLine 2\n\nLine 3\n[Sources: 2, 4]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Line 1\n\nLine 2\n\nLine 3"
+        assert citations == {2, 4}
+
+    def test_empty_string(self):
+        clean, citations = extract_and_strip_sources_block("")
+        assert clean == ""
+        assert citations == set()
+
+    def test_sources_mid_text_not_stripped(self):
+        text = "Answer [Sources: 1, 2] and more text after"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == text
+        assert citations == set()
+
+    def test_brackets_around_numbers_only(self):
+        text = "Answer text\nSources: [1, 3]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {1, 3}
+
+    def test_no_brackets_at_all(self):
+        text = "Answer text\nSources: 1, 3"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {1, 3}
+
+    def test_singular_no_brackets(self):
+        text = "Answer text\nSource: 2"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == {2}
+
+
+class TestFilterSourcesByCitations:
+    def test_basic_filtering(self):
+        sources = ["a", "b", "c", "d", "e"]
+        result = filter_sources_by_citations(sources, {1, 3, 5})
+        assert result == ["a", "c", "e"]
+
+    def test_empty_citations_returns_all(self):
+        sources = ["a", "b", "c"]
+        result = filter_sources_by_citations(sources, set())
+        assert result == ["a", "b", "c"]
+
+    def test_out_of_range_citations_fallback(self):
+        sources = ["a", "b", "c"]
+        result = filter_sources_by_citations(sources, {99})
+        assert result == ["a", "b", "c"]
+
+    def test_partial_out_of_range(self):
+        sources = ["a", "b", "c"]
+        result = filter_sources_by_citations(sources, {1, 99})
+        assert result == ["a"]
+
+    def test_single_citation(self):
+        sources = ["a", "b", "c"]
+        result = filter_sources_by_citations(sources, {2})
+        assert result == ["b"]
+
+    def test_empty_sources(self):
+        result = filter_sources_by_citations([], {1, 2})
+        assert result == []
+
+    def test_all_cited(self):
+        sources = ["a", "b", "c"]
+        result = filter_sources_by_citations(sources, {1, 2, 3})
+        assert result == ["a", "b", "c"]
+
+    def test_preserves_order(self):
+        sources = ["a", "b", "c", "d"]
+        result = filter_sources_by_citations(sources, {4, 2})
+        assert result == ["b", "d"]
+
+    def test_with_dict_sources(self):
+        sources = [{"file": "a.pdf"}, {"file": "b.pdf"}, {"file": "c.pdf"}]
+        result = filter_sources_by_citations(sources, {1, 3})
+        assert result == [{"file": "a.pdf"}, {"file": "c.pdf"}]

openrag/components/utils.py

@@ -1,5 +1,9 @@
 import asyncio
+import copy
+import json
+import re
 import threading
+from collections import deque
 from typing import ClassVar
 
 import ray
@@ -88,28 +92,138 @@ def get_num_tokens():
     return _cached_length_function
 
 
-def format_context(docs: list[Document], max_context_tokens: int = 4096) -> str:
-    n_docs = 0
+def format_context(
+    docs: list[Document], max_context_tokens: int = 4096, number_sources: bool = True
+) -> tuple[str, list[int]]:
     if not docs:
-        return "No document found from the database", n_docs
+        return "No document found from the database", []
 
     _length_function = get_num_tokens()
 
-    docs_with_tokens = list(map(lambda d: (_length_function(d.page_content), d), docs))  # noqa: C417
-
     reduced_docs = []
-
+    included_indices = []
     total_tokens = 0
-    for n_tokens, doc in docs_with_tokens:
+
+    for i, doc in enumerate(docs):
+        prefix = f"[Source {len(reduced_docs) + 1}]\n" if number_sources else ""
+        n_tokens = _length_function(doc.page_content)
+        if prefix:
+            n_tokens += _length_function(prefix)
         if total_tokens + n_tokens > max_context_tokens:
             break
-        reduced_docs.append(doc.page_content)
-        n_docs += 1
+        reduced_docs.append(f"{prefix}{doc.page_content}")
+        included_indices.append(i)
         total_tokens += n_tokens
 
     sep = "-" * 10 + "\n\n"
     logger.debug("Context formatted", total_tokens=total_tokens, doc_count=len(reduced_docs))
-    return f"{sep}".join(reduced_docs), n_docs
+    return f"{sep}".join(reduced_docs), included_indices
+
+
+def extract_and_strip_sources_block(text: str) -> tuple[str, set[int]]:
+    """Extract [Sources: N, N, ...] block from end of response. Return (clean_text, citations).
+
+    Handles LLM output variations: [Sources: 1, 3], Sources: [1, 3], Sources: 1, 3
+    """
+    pattern = r"\n?\[?Sources?\]?\s*:\s*\[?([\d,\s]+)\]?\s*$"
+    match = re.search(pattern, text)
+    if not match:
+        return text, set()
+    citations = {int(n.strip()) for n in match.group(1).split(",") if n.strip().isdigit()}
+    clean_text = text[: match.start()].rstrip()
+    return clean_text, citations
+
+
+def filter_sources_by_citations(sources: list, citations: set[int]) -> list:
+    """Keep only sources whose 1-based index was cited. Fallback to all if none match."""
+    if not citations:
+        return sources
+    filtered = [s for i, s in enumerate(sources, start=1) if i in citations]
+    return filtered if filtered else sources
+
+
+async def stream_with_source_filtering(
+    llm_stream,
+    sources: list,
+    model_name: str,
+    all_sources_json: str,  # noqa: ARG001 — kept for backward compat; no longer sent to clients
+    buffer_size: int = 100,
+):
+    """Process an LLM SSE stream, stripping the [Sources: ...] tag from content.
+
+    Buffers the last `buffer_size` chars of content to intercept the sources tag
+    before it reaches the client. On stream end, strips the tag and emits a finish
+    chunk with filtered source metadata.
+
+    Yields SSE "data: ..." lines ready to forward to the client.
+    """
+    chunk_buffer: deque[dict] = deque()
+    buffered_content_len = 0
+    last_chunk_template = None
+    last_finish_reason = None
+
+    async for line in llm_stream:
+        if not line.startswith("data:"):
+            continue
+
+        if line.strip() == "data: [DONE]":
+            buffered_text = "".join(
+                (c.get("choices", [{}])[0].get("delta", {}).get("content", "") or "") for c in chunk_buffer
+            )
+            clean_text, citations = extract_and_strip_sources_block(buffered_text)
+
+            remaining = clean_text
+            for chunk in chunk_buffer:
+                original_content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "") or ""
+                if not remaining:
+                    break
+                surviving = remaining[: len(original_content)]
+                remaining = remaining[len(original_content) :]
+                if surviving != original_content:
+                    chunk["choices"][0]["delta"]["content"] = surviving
+                chunk["extra"] = "{}"
+                yield f"data: {json.dumps(chunk)}\n\n"
+
+            if last_chunk_template:
+                filtered = filter_sources_by_citations(sources, citations)
+                finish_chunk = copy.deepcopy(last_chunk_template)
+                finish_chunk["choices"][0]["delta"] = {}
+                finish_chunk["choices"][0]["finish_reason"] = last_finish_reason or "stop"
+                finish_chunk["extra"] = json.dumps({"sources": filtered})
+                yield f"data: {json.dumps(finish_chunk)}\n\n"
+
+            yield "data: [DONE]\n\n"
+
+        else:
+            data_str = line[len("data: ") :]
+            data = json.loads(data_str)
+            data["model"] = model_name
+
+            choice = data.get("choices", [{}])[0]
+            delta = choice.get("delta", {})
+            content = delta.get("content", "") or ""
+            finish_reason = choice.get("finish_reason")
+
+            if finish_reason:
+                # Save finish_reason, don't forward — we emit it at the end
+                last_finish_reason = finish_reason
+                last_chunk_template = data
+            elif content:
+                last_chunk_template = data
+                chunk_buffer.append(data)
+                buffered_content_len += len(content)
+
+                while buffered_content_len > buffer_size:
+                    oldest = chunk_buffer.popleft()
+                    oldest_content = oldest.get("choices", [{}])[0].get("delta", {}).get("content", "") or ""
+                    oldest["extra"] = "{}"
+                    buffered_content_len -= len(oldest_content)
+                    yield f"data: {json.dumps(oldest)}\n\n"
+
+            else:
+                # Forward non-content, non-finish chunks immediately (e.g. role delta)
+                data["extra"] = "{}"
+                yield f"data: {json.dumps(data)}\n\n"
 
 
 # Initialize language detector