diff --git a/agent.py b/agent.py index d5715ed..be67dfc 100644 --- a/agent.py +++ b/agent.py @@ -83,6 +83,19 @@ def _format_context(matches: list[ReviewMatch]) -> str: return "\n\n".join(sections) +def _remove_standalone_control_token(answer: str) -> str | None: + if INSUFFICIENT_EVIDENCE_TOKEN not in answer: + return answer + retained_lines: list[str] = [] + for line in answer.splitlines(): + if line.strip() == INSUFFICIENT_EVIDENCE_TOKEN: + continue + if INSUFFICIENT_EVIDENCE_TOKEN in line: + return None + retained_lines.append(line) + return "\n".join(retained_lines).strip() + + def _validate_and_number_citations( answer: str, matches: list[ReviewMatch], @@ -188,7 +201,8 @@ def answer_question( retrieved_source_ids=retrieved_source_ids, abstained=True, ) - if INSUFFICIENT_EVIDENCE_TOKEN in normalized_answer: + normalized_answer = _remove_standalone_control_token(normalized_answer) + if normalized_answer is None: return AnswerResult( answer=CITATION_VALIDATION_MESSAGE, sources=(), diff --git a/tests/test_agent.py b/tests/test_agent.py index e53334b..a29f859 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -172,7 +172,9 @@ 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: + def test_accepts_cited_answer_with_standalone_insufficient_token_line( + self, + ) -> None: document = Document( page_content="The crust was perfectly crispy.", metadata={"source_id": "review-1"}, @@ -182,7 +184,9 @@ def test_rejects_answer_mixed_with_insufficient_evidence_token(self) -> None: 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]."), + "multiline_whitespace": ( + "\n INSUFFICIENT_EVIDENCE \n\nGuests praise the crispy crust [1].\n" + ), } for position, response in responses.items(): @@ -193,12 +197,50 @@ def test_rejects_answer_mixed_with_insufficient_evidence_token(self) -> None: model=FakeModel(response), ) - self.assertEqual(result.answer, CITATION_VALIDATION_MESSAGE) + self.assertEqual(result.answer, "Guests praise the crispy crust [1].") self.assertNotIn("INSUFFICIENT_EVIDENCE", result.answer) - self.assertEqual(result.sources, ()) + self.assertEqual(len(result.sources), 1) + self.assertEqual(result.sources[0].document.id, "review-1") self.assertEqual(result.retrieved_source_ids, ("review-1",)) self.assertFalse(result.abstained) + def test_rejects_insufficient_evidence_token_embedded_in_prose(self) -> None: + document = Document( + page_content="The crust was perfectly crispy.", + metadata={"source_id": "review-1"}, + id="review-1", + ) + + result = answer_question( + "What do guests say about the crust?", + vector_store=FakeStore([(document, 0.5)]), + model=FakeModel( + "The raw marker INSUFFICIENT_EVIDENCE must not be shown [1]." + ), + ) + + self.assertEqual(result.answer, CITATION_VALIDATION_MESSAGE) + self.assertNotIn("INSUFFICIENT_EVIDENCE", result.answer) + self.assertEqual(result.sources, ()) + self.assertFalse(result.abstained) + + def test_rejects_uncited_answer_after_removing_control_token_line(self) -> None: + document = Document( + page_content="The crust was perfectly crispy.", + metadata={"source_id": "review-1"}, + id="review-1", + ) + + result = answer_question( + "What do guests say about the crust?", + vector_store=FakeStore([(document, 0.5)]), + model=FakeModel("Guests praise the crust.\nINSUFFICIENT_EVIDENCE"), + ) + + self.assertEqual(result.answer, CITATION_VALIDATION_MESSAGE) + self.assertEqual(result.sources, ()) + 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 f1014dd..209c3cb 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -124,9 +124,9 @@ def test_does_not_render_mixed_insufficient_evidence_token(self) -> None: 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.assertIn("Guests praise the crispy crust [1].", rendered_markdown) self.assertNotIn("INSUFFICIENT_EVIDENCE", rendered_markdown) - self.assertNotIn("#### Evidence", rendered_markdown) + self.assertIn("#### Evidence", rendered_markdown) if __name__ == "__main__":