fix: prevent abstention control-token leaks

agent.py

@@ -188,6 +188,12 @@ def answer_question(
             retrieved_source_ids=retrieved_source_ids,
             abstained=True,
         )
+    if INSUFFICIENT_EVIDENCE_TOKEN in normalized_answer:
+        return AnswerResult(
+            answer=CITATION_VALIDATION_MESSAGE,
+            sources=(),
+            retrieved_source_ids=retrieved_source_ids,
+        )
     validated = _validate_and_number_citations(normalized_answer, matches)
     if validated is None:
         return AnswerResult(

tests/test_agent.py

@@ -164,14 +164,41 @@ def test_model_can_abstain_when_retrieved_reviews_are_insufficient(self) -> None
         result = answer_question(
             "Is parking available?",
             vector_store=FakeStore([(document, 0.5)]),
-            model=FakeModel("INSUFFICIENT_EVIDENCE"),
+            model=FakeModel("\n  INSUFFICIENT_EVIDENCE  \n"),
         )
 
         self.assertEqual(result.answer, NO_MATCH_MESSAGE)
         self.assertEqual(result.sources, ())
         self.assertEqual(result.retrieved_source_ids, ("review-1",))
         self.assertTrue(result.abstained)
 
+    def test_rejects_answer_mixed_with_insufficient_evidence_token(self) -> None:
+        document = Document(
+            page_content="The crust was perfectly crispy.",
+            metadata={"source_id": "review-1"},
+            id="review-1",
+        )
+
+        responses = {
+            "followed": "Guests praise the crispy crust [1].\n\nINSUFFICIENT_EVIDENCE",
+            "preceded": "INSUFFICIENT_EVIDENCE\nGuests praise the crispy crust [1].",
+            "embedded": ("The raw marker INSUFFICIENT_EVIDENCE must not be shown [1]."),
+        }
+
+        for position, response in responses.items():
+            with self.subTest(position=position):
+                result = answer_question(
+                    "What do guests say about the crust?",
+                    vector_store=FakeStore([(document, 0.5)]),
+                    model=FakeModel(response),
+                )
+
+                self.assertEqual(result.answer, CITATION_VALIDATION_MESSAGE)
+                self.assertNotIn("INSUFFICIENT_EVIDENCE", result.answer)
+                self.assertEqual(result.sources, ())
+                self.assertEqual(result.retrieved_source_ids, ("review-1",))
+                self.assertFalse(result.abstained)
+
     def test_does_not_call_model_when_filters_match_no_reviews(self) -> None:
         model = FakeModel()
 

tests/test_dashboard.py

@@ -3,9 +3,39 @@
 import unittest
 from unittest.mock import patch
 
+from langchain_core.documents import Document
 from streamlit.testing.v1 import AppTest
 
+import agent
 import ollama_health
+import vector
+
+
+class MixedTokenModel:
+    def invoke(self, _: str) -> str:
+        return "Guests praise the crispy crust [1].\n\nINSUFFICIENT_EVIDENCE"
+
+
+class DashboardStore:
+    def __init__(self) -> None:
+        self.document = Document(
+            page_content="The crust was perfectly crispy.",
+            metadata={
+                "source_id": "review-1",
+                "rating": 5,
+                "date": "2024-01-10",
+            },
+            id="review-1",
+        )
+
+    def get(self, **_: object) -> dict[str, list[str]]:
+        return {"ids": ["review-1"]}
+
+    def similarity_search_with_score(
+        self, _: str, *, k: int, filter: dict | None = None
+    ) -> list[tuple[Document, float]]:
+        del filter
+        return [(self.document, 0.1)][:k]
 
 
 class DashboardRenderTest(unittest.TestCase):
@@ -74,6 +104,30 @@ def test_adapts_uploaded_platform_schema(self) -> None:
             ["Sentiment", "Restaurant", "Country or region"],
         )
 
+    def test_does_not_render_mixed_insufficient_evidence_token(self) -> None:
+        health = ollama_health.OllamaHealth(
+            True,
+            ("llama3.2:latest", "mxbai-embed-large:latest"),
+            (),
+        )
+        host = "http://dashboard-token-regression:11434"
+        with (
+            patch.object(ollama_health, "DEFAULT_OLLAMA_HOST", host),
+            patch.object(ollama_health, "check_ollama", return_value=health),
+            patch.object(vector, "create_vector_store", return_value=DashboardStore()),
+            patch.object(agent, "create_chat_model", return_value=MixedTokenModel()),
+        ):
+            application = AppTest.from_file("dashboard.py").run(timeout=30)
+            application.chat_input[0].set_value(
+                "What do guests say about the crust?"
+            ).run(timeout=30)
+
+        self.assertEqual(list(application.exception), [])
+        rendered_markdown = "\n".join(item.value for item in application.markdown)
+        self.assertIn("I could not produce an answer with citations", rendered_markdown)
+        self.assertNotIn("INSUFFICIENT_EVIDENCE", rendered_markdown)
+        self.assertNotIn("#### Evidence", rendered_markdown)
+
 
 if __name__ == "__main__":
     unittest.main()