Heal LLM-simplified citations to canonical reference

app/backend/approaches/approach.py

@@ -830,6 +830,33 @@ def get_image_citation(self, sourcepage: Optional[str], image_url: str):
         image_filename = image_url.split("/")[-1]
         return f"{sourcepage_citation}({image_filename})"
 
+    @staticmethod
+    def heal_citation(citation: str, valid_citations: Optional[list[str]]) -> str:
+        """Map a possibly-simplified citation back to a valid source citation.
+
+        LLMs sometimes drop leading special characters (for example a "- " prefix or
+        parentheses) when emitting a citation for a filename. When that happens the
+        citation no longer matches the indexed ``sourcefile``, so the document link
+        returns 403/404. If the citation is not an exact match for a known source,
+        fall back to the unique valid citation that ends with the same string.
+        """
+        if not citation or not valid_citations or citation in valid_citations:
+            return citation
+        matches = [valid for valid in valid_citations if valid != citation and valid.endswith(citation)]
+        if len(matches) == 1:
+            return matches[0]
+        return citation
+
+    def heal_citations(self, answer: str, valid_citations: Optional[list[str]]) -> str:
+        """Heal every ``[citation]`` token in an answer against the valid citations."""
+        if not answer or not valid_citations:
+            return answer
+
+        def _sub(match: re.Match) -> str:
+            return f"[{self.heal_citation(match.group(1), valid_citations)}]"
+
+        return re.sub(r"\[([^\[\]]+)\]", _sub, answer)
+
     async def download_blob_as_base64(self, blob_url: str, user_oid: Optional[str] = None) -> Optional[str]:
         """
         Downloads a blob from either Azure Blob Storage or Azure Data Lake Storage and returns it as a base64 encoded string.

app/backend/approaches/chatreadretrieveread.py

@@ -126,6 +126,8 @@ async def run_without_streaming(
         if overrides.get("suggest_followup_questions"):
             content, followup_questions = self.extract_followup_questions(content)
             extra_info.followup_questions = followup_questions
+        # Repair citations the model may have simplified so they match the indexed sources
+        content = self.heal_citations(content, extra_info.data_points.citations)
         if self.include_token_usage and extra_info.thoughts and response.usage:
             extra_info.thoughts[-1].update_token_usage(response.usage)
         chat_app_response = {
@@ -166,6 +168,9 @@ async def run_with_streaming(
                 content, followup_questions = self.extract_followup_questions(content)
                 extra_info.followup_questions = followup_questions
 
+            # Repair citations the model may have simplified so they match the indexed sources
+            content = self.heal_citations(content, extra_info.data_points.citations)
+
             if self.include_token_usage and extra_info.thoughts and result.usage:
                 extra_info.thoughts[-1].update_token_usage(result.usage)
 
@@ -178,24 +183,57 @@ async def run_with_streaming(
         # Handle streaming Response events
         stream = cast(AsyncStream[ResponseStreamEvent], result)
 
+        # Buffer text inside a "[citation]" token so it can be repaired before being emitted.
+        # The model sometimes simplifies citations (e.g. dropping a "- " filename prefix),
+        # which would otherwise break the document link the frontend builds from the token.
+        valid_citations = extra_info.data_points.citations
+        citation_buffer: Optional[str] = None
+
+        def heal_streamed_text(text: str) -> str:
+            nonlocal citation_buffer
+            output = ""
+            for char in text:
+                if citation_buffer is not None:
+                    citation_buffer += char
+                    if char == "]":
+                        inner = citation_buffer[1:-1]
+                        output += f"[{self.heal_citation(inner, valid_citations)}]"
+                        citation_buffer = None
+                elif char == "[":
+                    citation_buffer = "["
+                else:
+                    output += char
+            return output
+
         async for event in stream:
             if isinstance(event, ResponseTextDeltaEvent):
                 delta_content: str = event.delta or ""
                 if overrides.get("suggest_followup_questions") and "<<" in delta_content:
                     followup_questions_started = True
-                    earlier_content = delta_content[: delta_content.index("<<")]
+                    earlier_content = heal_streamed_text(delta_content[: delta_content.index("<<")])
+                    # Flush any unterminated citation token before the followup section begins
+                    if citation_buffer:
+                        earlier_content += citation_buffer
+                        citation_buffer = None
                     if earlier_content:
                         yield {"type": "response.output_text.delta", "delta": earlier_content}
                     followup_content += delta_content[delta_content.index("<<") :]
                 elif followup_questions_started:
                     followup_content += delta_content
                 else:
-                    yield {"type": "response.output_text.delta", "delta": delta_content}
+                    healed_content = heal_streamed_text(delta_content)
+                    if healed_content:
+                        yield {"type": "response.output_text.delta", "delta": healed_content}
             elif isinstance(event, ResponseCompletedEvent):
                 if event.response.usage and extra_info.thoughts and self.include_token_usage:
                     extra_info.thoughts[-1].update_token_usage(event.response.usage)
                     yield {"type": "response.context", "context": extra_info, "session_state": session_state}
 
+        # Flush any unterminated citation token left in the buffer at the end of the stream
+        if citation_buffer:
+            yield {"type": "response.output_text.delta", "delta": citation_buffer}
+            citation_buffer = None
+
         if followup_content:
             _, followup_questions = self.extract_followup_questions(followup_content)
             extra_info.followup_questions = followup_questions

app/frontend/src/components/Answer/AnswerParser.tsx

@@ -117,14 +117,18 @@ const collectCitations = (answer: ChatAppResponse, isStreaming: boolean): { frag
             return;
         }
 
-        const isValidCitation = possibleCitations.some(citation => citation.endsWith(part));
-        if (!isValidCitation) {
+        // The LLM sometimes simplifies citations by dropping leading special characters,
+        // e.g. rendering "- PyCon US 2025.pdf#page=1" as "PyCon US 2025.pdf#page=1". Resolve
+        // the cited text back to the canonical citation so links and lookups use the indexed
+        // value instead of the simplified text (which would 404/403 against the content route).
+        const canonicalCitation = possibleCitations.find(citation => citation === part) ?? possibleCitations.find(citation => citation.endsWith(part));
+        if (!canonicalCitation) {
             fragments.push({ type: "text", value: `[${part}]` });
             return;
         }
 
         // Resolve SharePoint filename to URL if applicable
-        const resolvedReference = resolveSharePointUrl(part);
+        const resolvedReference = resolveSharePointUrl(canonicalCitation);
 
         // Check if this resolved reference already exists
         const existing = citationMap.get(resolvedReference);
@@ -133,7 +137,7 @@ const collectCitations = (answer: ChatAppResponse, isStreaming: boolean): { frag
             return;
         }
 
-        const backendDetail = citationActivityDetails?.[part];
+        const backendDetail = citationActivityDetails?.[canonicalCitation];
         const activityId = backendDetail?.id;
         const stepMeta = activityId ? activitySteps[String(activityId)] : undefined;
 

tests/test_chatapproach.py

@@ -8,6 +8,7 @@
     Response,
     ResponseFunctionToolCall,
     ResponseOutputMessage,
+    ResponseTextDeltaEvent,
     ResponseUsage,
 )
 from openai.types.responses.response_usage import (
@@ -624,6 +625,108 @@ async def test_run_until_final_call_rejects_web_streaming(chat_approach):
         )
 
 
+def test_heal_citation_returns_exact_match_unchanged(chat_approach):
+    valid_citations = ["PyCon US 2025.pdf#page=1"]
+    assert chat_approach.heal_citation("PyCon US 2025.pdf#page=1", valid_citations) == "PyCon US 2025.pdf#page=1"
+
+
+def test_heal_citation_recovers_simplified_prefix(chat_approach):
+    # The model dropped the leading "- " prefix that exists in the indexed sourcefile
+    valid_citations = ["- PyCon US 2025.pdf#page=1", "- PyCon US 2025.pdf#page=2"]
+    assert chat_approach.heal_citation("PyCon US 2025.pdf#page=1", valid_citations) == "- PyCon US 2025.pdf#page=1"
+
+
+def test_heal_citation_ambiguous_match_returns_original(chat_approach):
+    # Two valid citations end with the same string, so we cannot safely pick one
+    valid_citations = ["- report.pdf#page=1", "-- report.pdf#page=1"]
+    assert chat_approach.heal_citation("report.pdf#page=1", valid_citations) == "report.pdf#page=1"
+
+
+def test_heal_citation_no_match_returns_original(chat_approach):
+    valid_citations = ["other.pdf#page=1"]
+    assert chat_approach.heal_citation("missing.pdf#page=1", valid_citations) == "missing.pdf#page=1"
+
+
+def test_heal_citations_repairs_tokens_in_answer(chat_approach):
+    valid_citations = ["- PyCon US 2025.pdf#page=1", "(1) notes.pdf#page=3"]
+    answer = "See [PyCon US 2025.pdf#page=1] and [notes.pdf#page=3] for details."
+    healed = chat_approach.heal_citations(answer, valid_citations)
+    assert healed == "See [- PyCon US 2025.pdf#page=1] and [(1) notes.pdf#page=3] for details."
+
+
+def test_heal_citations_leaves_non_citation_brackets_unchanged(chat_approach):
+    valid_citations = ["- PyCon US 2025.pdf#page=1"]
+    answer = "A list item [not a citation] stays as is."
+    assert chat_approach.heal_citations(answer, valid_citations) == answer
+
+
+def test_heal_citations_noop_without_valid_citations(chat_approach):
+    answer = "Nothing to heal [foo.pdf#page=1]."
+    assert chat_approach.heal_citations(answer, []) == answer
+
+
+@pytest.mark.asyncio
+async def test_run_with_streaming_heals_citations_split_across_deltas(chat_approach, monkeypatch):
+    extra_info = ExtraInfo(
+        data_points=DataPoints(text=[], images=[], citations=["- PyCon US 2025.pdf#page=1"]),
+        thoughts=[ThoughtStep("Final", None, props={})],
+    )
+
+    seq = 0
+
+    def make_delta(delta: str) -> ResponseTextDeltaEvent:
+        nonlocal seq
+        seq += 1
+        return ResponseTextDeltaEvent(
+            content_index=0,
+            delta=delta,
+            item_id="item-0",
+            logprobs=[],
+            output_index=0,
+            sequence_number=seq,
+            type="response.output_text.delta",
+        )
+
+    class FakeEventStream:
+        def __init__(self, events):
+            self._events = events
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            if not self._events:
+                raise StopAsyncIteration
+            return self._events.pop(0)
+
+    # The "[PyCon US 2025.pdf#page=1]" citation token is split across several deltas
+    events = [
+        make_delta("Answer "),
+        make_delta("[PyCon US 2025"),
+        make_delta(".pdf#page=1]"),
+        make_delta(" done."),
+    ]
+
+    async def fake_completion():
+        return FakeEventStream(events)
+
+    async def fake_run_until_final_call(messages, overrides, auth_claims, should_stream):
+        return extra_info, fake_completion()
+
+    monkeypatch.setattr(chat_approach, "run_until_final_call", fake_run_until_final_call)
+
+    deltas = []
+    async for event in chat_approach.run_with_streaming(
+        messages=[{"role": "user", "content": "Hello"}],
+        overrides={},
+        auth_claims={},
+        session_state="state",
+    ):
+        if event["type"] == "response.output_text.delta":
+            deltas.append(event["delta"])
+
+    assert "".join(deltas) == "Answer [- PyCon US 2025.pdf#page=1] done."
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "auth_claims,expected",