diff --git a/agent.py b/agent.py index 304000b..d5715ed 100644 --- a/agent.py +++ b/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( diff --git a/tests/test_agent.py b/tests/test_agent.py index 455475c..e53334b 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -164,7 +164,7 @@ 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) @@ -172,6 +172,33 @@ def test_model_can_abstain_when_retrieved_reviews_are_insufficient(self) -> None 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() diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index f9eeadd..f1014dd 100644 --- a/tests/test_dashboard.py +++ b/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()