feat: handle LLM [Sources: none] to filter out unused sources

openrag/components/test_source_filtering.py

@@ -1,6 +1,13 @@
 """Tests for source citation extraction and filtering utilities."""
 
-from components.utils import extract_and_strip_sources_block, filter_sources_by_citations
+import json
+
+import pytest
+from components.utils import (
+    extract_and_strip_sources_block,
+    filter_sources_by_citations,
+    stream_with_source_filtering,
+)
 
 
 class TestExtractAndStripSourcesBlock:
@@ -26,7 +33,7 @@ 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()
+        assert citations is None
 
     def test_sources_with_trailing_whitespace(self):
         text = "Answer text\n[Sources: 1, 3]   "
@@ -49,13 +56,13 @@ def test_multiline_answer(self):
     def test_empty_string(self):
         clean, citations = extract_and_strip_sources_block("")
         assert clean == ""
-        assert citations == set()
+        assert citations is None
 
     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()
+        assert citations is None
 
     def test_brackets_around_numbers_only(self):
         text = "Answer text\nSources: [1, 3]"
@@ -75,18 +82,41 @@ def test_singular_no_brackets(self):
         assert clean == "Answer text"
         assert citations == {2}
 
+    def test_sources_none(self):
+        text = "Answer text\n[Sources: none]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == set()
+
+    def test_sources_none_no_brackets(self):
+        text = "Answer text\nSources: none"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == set()
+
+    def test_sources_none_capitalized(self):
+        text = "Answer text\n[Sources: None]"
+        clean, citations = extract_and_strip_sources_block(text)
+        assert clean == "Answer text"
+        assert citations == set()
+
 
 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):
+    def test_none_citations_returns_all(self):
         sources = ["a", "b", "c"]
-        result = filter_sources_by_citations(sources, set())
+        result = filter_sources_by_citations(sources, None)
         assert result == ["a", "b", "c"]
 
+    def test_empty_citations_returns_empty(self):
+        sources = ["a", "b", "c"]
+        result = filter_sources_by_citations(sources, set())
+        assert result == []
+
     def test_out_of_range_citations_fallback(self):
         sources = ["a", "b", "c"]
         result = filter_sources_by_citations(sources, {99})
@@ -120,3 +150,94 @@ 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"}]
+
+
+# --- helpers for streaming tests ---
+
+
+def _make_chunk(content: str, chunk_id: str = "chatcmpl-1") -> str:
+    """Build an SSE line with a content delta."""
+    return "data: " + json.dumps({"id": chunk_id, "choices": [{"delta": {"content": content}, "finish_reason": None}]})
+
+
+def _make_finish(chunk_id: str = "chatcmpl-1") -> str:
+    """Build an SSE line with finish_reason='stop'."""
+    return "data: " + json.dumps({"id": chunk_id, "choices": [{"delta": {}, "finish_reason": "stop"}]})
+
+
+DONE_LINE = "data: [DONE]"
+
+
+async def _fake_stream(lines: list[str]):
+    for line in lines:
+        yield line
+
+
+async def _collect(async_gen) -> list[str]:
+    return [line async for line in async_gen]
+
+
+def _parse_finish_sources(sse_lines: list[str]) -> list:
+    """Extract the sources list from the finish chunk (second-to-last line before [DONE])."""
+    for line in reversed(sse_lines):
+        if line.startswith("data: ") and line.strip() != "data: [DONE]":
+            data = json.loads(line[len("data: ") :])
+            extra = data.get("extra")
+            if extra and extra != "{}":
+                return json.loads(extra).get("sources", [])
+    return []
+
+
+def _collect_content(sse_lines: list[str]) -> str:
+    """Concatenate all content deltas from SSE lines."""
+    parts = []
+    for line in sse_lines:
+        if not line.startswith("data: ") or line.strip() == "data: [DONE]":
+            continue
+        data = json.loads(line[len("data: ") :])
+        content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
+        if content:
+            parts.append(content)
+    return "".join(parts)
+
+
+class TestStreamWithSourceFiltering:
+    SOURCES = [{"file": "a.pdf"}, {"file": "b.pdf"}, {"file": "c.pdf"}]
+
+    @pytest.mark.asyncio
+    async def test_case1_llm_cites_specific_sources(self):
+        """Case 1: LLM cites [Sources: 1, 3] → only cited sources returned."""
+        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 _collect_content(result) == "Here is the answer."
+        assert _parse_finish_sources(result) == [{"file": "a.pdf"}, {"file": "c.pdf"}]
+
+    @pytest.mark.asyncio
+    async def test_case2_llm_says_sources_none(self):
+        """Case 2: LLM says [Sources: none] → no sources returned."""
+        lines = [
+            _make_chunk("I cannot find this in the documents."),
+            _make_chunk("\n[Sources: none]"),
+            _make_finish(),
+            DONE_LINE,
+        ]
+        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) == []
+
+    @pytest.mark.asyncio
+    async def test_case3_llm_no_tag_fallback_all(self):
+        """Case 3: LLM omits tag entirely → fallback to all sources."""
+        lines = [
+            _make_chunk("Answer without any sources tag."),
+            _make_finish(),
+            DONE_LINE,
+        ]
+        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

openrag/components/utils.py

@@ -120,24 +120,52 @@ def format_context(
     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).
+_SOURCES_NONE_RE = re.compile(r"\n?\[?Sources?\]?\s*:\s*\[?\s*none\s*\]?\s*$", re.IGNORECASE)
+_SOURCES_NUMS_RE = re.compile(r"\n?\[?Sources?\]?\s*:\s*\[?([\d,\s]+)\]?\s*$")
 
-    Handles LLM output variations: [Sources: 1, 3], Sources: [1, 3], Sources: 1, 3
+
+def extract_and_strip_sources_block(text: str) -> tuple[str, set[int] | None]:
+    """Extract [Sources: ...] block from end of response. Return (clean_text, citations).
+
+    Returns:
+        (clean_text, citations) where citations is:
+        - set of ints: LLM cited specific sources
+        - empty set:   LLM explicitly said [Sources: none]
+        - None:        LLM didn't include any sources tag
     """
-    pattern = r"\n?\[?Sources?\]?\s*:\s*\[?([\d,\s]+)\]?\s*$"
-    match = re.search(pattern, text)
+    # Check for explicit "none" first
+    match_none = _SOURCES_NONE_RE.search(text)
+    if match_none:
+        clean_text = text[: match_none.start()].rstrip()
+        logger.debug("LLM explicitly reported no sources used")
+        return clean_text, set()
+
+    # Check for numbered citations
+    match = _SOURCES_NUMS_RE.search(text)
     if not match:
-        return text, set()
+        tail = text[-150:] if len(text) > 150 else text
+        logger.debug("No [Sources: ...] tag found in LLM response", tail=repr(tail))
+        return text, None
+
     citations = {int(n.strip()) for n in match.group(1).split(",") if n.strip().isdigit()}
+    logger.debug(
+        "Extracted source citations from LLM response", citations=sorted(citations), matched=repr(match.group(0))
+    )
     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:
+def filter_sources_by_citations(sources: list, citations: set[int] | None) -> list:
+    """Keep only sources whose 1-based index was cited.
+
+    - citations is None:      LLM didn't include tag → fallback to all sources
+    - citations is empty set:  LLM said [Sources: none] → return no sources
+    - citations has values:    filter to cited sources only
+    """
+    if citations is None:
         return sources
+    if not citations:
+        return []
     filtered = [s for i, s in enumerate(sources, start=1) if i in citations]
     return filtered if filtered else sources
 

prompts/example1/spoken_style_answer_tmpl.txt

@@ -7,9 +7,10 @@ Your goal is to give short (1-2 sentences), clear, and accurate explanations, ba
    * Answer strictly from the information in `Context`.
    * Do not guess, infer, or use outside knowledge.
    * If the Context lacks enough information, say so briefly and ask the user for more details.
-   * At the very end of your response, on a new line, list which source numbers you used: [Sources: 1, 3]
-   * Only list sources that actually contributed to your answer
-   * This sources line must be the very last line of your response
+   * At the very end of your response, you MUST always include a sources line on a new line
+   * If sources contributed to your answer, use this format: [Sources: 1, 3]
+   * If no sources were useful, write exactly: [Sources: none]
+   * This sources line must be the very last line of your response, always
 
 2. Match the user’s language
    * Speak in the same language as the user’s question.

prompts/example1/sys_prompt_tmpl.txt

@@ -9,9 +9,10 @@ Prioritize **clarity, accuracy, and completeness** in your responses.
    * **Never infer**, assume, or rely on any external knowledge.
    * If the context is **insufficient**, **invite the user** to clarify their query or provide additional keywords.
    * Do not cite sources or file names within your answer text
-   * At the very end of your response, on a new line, list which source numbers you used in this exact format: [Sources: 1, 3, 5]
-   * Only list sources that actually contributed to your answer
-   * This sources line must be the very last line of your response
+   * At the very end of your response, you MUST always include a sources line on a new line
+   * If sources contributed to your answer, use this format: [Sources: 1, 3, 5]
+   * If no sources were useful, write exactly: [Sources: none]
+   * This sources line must be the very last line of your response, always
 
 2. Language Consistency
    * Always respond **in the same language** as the user’s query.