"""AI Mode: retrieval, the citation contract, and thread ownership. Disposable SQLite; the model itself is stubbed, because what is worth testing here is not what a model says but what the server does with it. The safety property — an invented citation cannot survive — has to hold whatever comes back. """ import unittest from unittest.mock import AsyncMock, patch import test_quiz_builder as fixtures from app.models.article import Article, ArticleSectionIndex, QuestionArticleLink from app.models.conversation import Conversation, ConversationMessage from app.models.flashcard import Flashcard, FlashcardDeck from app.models.question import Question from app.routers import ai_mode from app.services import ai_mode_service class CitationContractTests(unittest.TestCase): """The step that makes a hallucinated citation impossible rather than rare.""" def sources(self): return [ {"kind": "article", "ref": "7", "id": 7, "title": "Febrile seizures"}, {"kind": "section", "ref": "7#abc", "id": 7, "section_id": "abc", "title": "Workup"}, ] def test_a_citation_that_was_not_retrieved_is_deleted(self): reply, citations = ai_mode_service.enforce_citations( "Fever first [[article:7]]. Then lumbar puncture [[article:99]].", self.sources()) self.assertIn("[[article:7]]", reply) self.assertNotIn("99", reply) self.assertEqual([c["id"] for c in citations], [7]) def test_deleting_a_marker_does_not_leave_broken_punctuation(self): reply, _ = ai_mode_service.enforce_citations( "This is true [[article:404]].", self.sources()) self.assertEqual(reply, "This is true.") def test_a_url_the_model_invents_is_not_a_citation(self): # Only the marker form counts, so a plausible-looking link cannot smuggle # itself into the citation list. reply, citations = ai_mode_service.enforce_citations( "See https://uptodate.com/febrile-seizures for more.", self.sources()) self.assertEqual(citations, []) self.assertIn("uptodate.com", reply) # left in the prose, cited by nothing def test_the_same_source_cited_twice_is_listed_once(self): _, citations = ai_mode_service.enforce_citations( "One [[article:7]]. Two [[article:7]].", self.sources()) self.assertEqual(len(citations), 1) def test_a_section_citation_keeps_the_section_it_points_at(self): _, citations = ai_mode_service.enforce_citations("Here [[section:7#abc]].", self.sources()) self.assertEqual(citations[0]["section_id"], "abc") def test_with_nothing_close_it_says_so_and_then_helps(self): """Refusing outright reads as a broken assistant, not a careful one. The old behaviour was to say nothing matched and stop. It is honest to name the gap; it is not honest to pretend an unrelated shortlist supports the answer, and it is not useful to withhold one entirely. """ prompt = ai_mode_service.build_prompt([], "open") self.assertIn("Nothing in this learner's library covers their question", prompt) self.assertIn("answer from general knowledge", prompt) self.assertIn("Do not cite anything", prompt) # And nothing it writes can be cited anyway. reply, citations = ai_mode_service.enforce_citations("Anything [[article:1]].", []) self.assertEqual(citations, []) self.assertEqual(reply, "Anything.") def test_something_adjacent_is_named_as_adjacent(self): # `sources()` carries only what citation enforcement needs; a prompt # also prints the text of each source. with_text = [{**s, "text": "Body"} for s in self.sources()] prompt = ai_mode_service.build_prompt(with_text, "adjacent") self.assertIn("closest things", prompt) self.assertIn("[[section:7#abc]]", prompt) # still citable def test_the_three_states_are_chosen_by_the_number_not_the_model(self): sources = self.sources() self.assertEqual(ai_mode_service.answer_mode(0.72, sources), "sourced") self.assertEqual(ai_mode_service.answer_mode(0.52, sources), "adjacent") # 0.49 is where "discuss love" and "photosynthesis" land against this # corpus, alongside "tell me a joke" — noise, not adjacency. self.assertEqual(ai_mode_service.answer_mode(0.49, sources), "open") self.assertEqual(ai_mode_service.answer_mode(0.90, []), "open") # Unmeasurable is not low: retrieval found these by other means, and # dropping every citation because the ruler is missing would be worse. self.assertEqual(ai_mode_service.answer_mode(None, sources), "sourced") def test_a_greeting_is_not_a_query(self): """Whatever it scores, "hello" is not a question about anything. This is the bug in the screenshot: a greeting came back with four confident citations, because reciprocal-rank fusion always returns an order and "ok" happens to embed at 0.51 against a corpus of clinical prose — the wrong side of the adjacency threshold. """ sources = self.sources() for greeting in ["hi", "Hello", "hello!", " thanks ", "ok thanks", "good morning", "Thank you!", "bye", "", "who are you?", "what can you do"]: self.assertTrue(ai_mode_service.is_small_talk(greeting), greeting) self.assertEqual( ai_mode_service.answer_mode(0.9, sources, greeting), "chat", greeting) def test_a_short_clinical_question_is_still_a_query(self): # The guard is a closed vocabulary rather than a length rule, precisely # so that these keep reaching retrieval. sources = self.sources() for query in ["croup dose?", "ok to give ibuprofen at 3 months?", "no stridor now what", "hi flow nasal cannula", "thanks to which vaccine has Hib fallen?"]: self.assertFalse(ai_mode_service.is_small_talk(query), query) self.assertEqual( ai_mode_service.answer_mode(0.72, sources, query), "sourced", query) def test_the_chat_prompt_offers_nothing_and_claims_nothing(self): prompt = ai_mode_service.build_prompt(self.sources(), "chat") # No shortlist reaches a turn that was never a query, so there is # nothing for the model to cite even if it tries. self.assertNotIn("[[section:7#abc]]", prompt) self.assertNotIn("SOURCES", prompt) # And unlike "open", it does not announce a gap nobody asked about. self.assertNotIn("Nothing in this learner's library", prompt) class _AiModeBase(unittest.TestCase): """Fixtures shared by the route and retrieval cases; holds no tests itself.""" def setUp(self): # The daily limit is enforced through the real Redis, so without this a # suite that passes today fails once it has been run sixty times — the # tests would be measuring shared state rather than this code. self._no_limit = patch("app.routers.ai_mode.check_rate_limit", lambda **kwargs: None) self._no_limit.start() self.addCleanup(self._no_limit.stop) self.bank = fixtures.BuilderTests() self.bank.setUp() self.client = self.bank.client self.client.app.include_router(ai_mode.router, prefix='/ai') self.db = self.bank.db self.db.add(Article(id=7, slug='febrile-seizures', title='Febrile seizures', summary='A seizure with fever in a young child', sections=[{"id": "a" * 32, "slug": "workup", "title": "Workup", "content": "LP"}], status='published', user_id=3)) self.db.add(ArticleSectionIndex(id=1, article_id=7, section_id='a' * 32, title='Workup', content='Lumbar puncture in a febrile infant')) self.db.add(FlashcardDeck(id=1, title='Mine', user_id=1)) self.db.flush() self.db.add(Flashcard(id=1, deck_id=1, front='Febrile seizure duration', back='Under 15 minutes')) self.db.query(Question).filter(Question.id == 1).update( {"question_text": "A child with a febrile seizure lasting two minutes"}) self.db.commit() self.bank.user = self.bank.owner def tearDown(self): self.bank.tearDown() def reply_with(self, text): """Stub the model, so the test is about the server's half of the contract.""" return patch("app.services.ai_service.achat", new=AsyncMock(return_value=text)) def ask(self, message, conversation_id=None): if conversation_id is None: conversation_id = self.client.post('/ai/conversations').json()['id'] return conversation_id, self.client.post( f'/ai/conversations/{conversation_id}/messages', json={'message': message}) class AiModeRouteTests(_AiModeBase): def test_an_answer_keeps_only_the_citations_retrieval_found(self): with self.reply_with("Yes [[article:7]], and also [[article:12345]]."): _, response = self.ask('febrile seizure') self.assertEqual(response.status_code, 200, response.text) body = response.json()['message'] self.assertNotIn('12345', body['content']) self.assertEqual([c['id'] for c in body['citations']], [7]) def test_the_first_question_names_the_thread(self): with self.reply_with("An answer."): conversation_id, response = self.ask('What causes febrile seizures?') self.assertEqual(response.json()['title'], 'What causes febrile seizures?') # And a later message does not rename it out from under the learner. with self.reply_with("Another answer."): _, second = self.ask('And the workup?', conversation_id) self.assertEqual(second.json()['title'], 'What causes febrile seizures?') def test_both_turns_are_stored_so_a_thread_can_be_reopened(self): with self.reply_with("Because of fever [[article:7]]."): conversation_id, _ = self.ask('why do febrile seizures happen?') messages = self.client.get(f'/ai/conversations/{conversation_id}').json()['messages'] self.assertEqual([m['role'] for m in messages], ['user', 'assistant']) # The citations are stored as they were filtered, not recomputed later. self.assertEqual(messages[1]['citations'][0]['id'], 7) self.assertEqual(messages[0]['citations'], []) def test_a_thread_belongs_to_one_learner(self): with self.reply_with("An answer."): conversation_id, _ = self.ask('febrile seizure') self.bank.user = self.bank.peer # Not 403: whether somebody else's thread exists is not this user's business. self.assertEqual(self.client.get(f'/ai/conversations/{conversation_id}').status_code, 404) self.assertEqual(self.client.delete(f'/ai/conversations/{conversation_id}').status_code, 404) self.assertEqual(self.client.post( f'/ai/conversations/{conversation_id}/messages', json={'message': 'hello'}).status_code, 404) self.assertEqual(self.client.get('/ai/conversations').json(), []) def test_a_model_failure_is_reported_not_stored(self): conversation_id = self.client.post('/ai/conversations').json()['id'] with patch("app.services.ai_service.achat", new=AsyncMock(side_effect=RuntimeError("down"))): response = self.client.post(f'/ai/conversations/{conversation_id}/messages', json={'message': 'febrile seizure'}) self.assertEqual(response.status_code, 502) # A half-written exchange is worse than none: the question is not kept. self.assertEqual(self.db.query(ConversationMessage).count(), 0) def test_deleting_a_thread_takes_its_messages(self): with self.reply_with("An answer."): conversation_id, _ = self.ask('febrile seizure') self.assertEqual(self.client.delete(f'/ai/conversations/{conversation_id}').status_code, 204) self.assertEqual(self.db.query(Conversation).count(), 0) self.assertEqual(self.db.query(ConversationMessage).count(), 0) class RetrievalTests(_AiModeBase): def test_retrieval_offers_only_what_this_learner_may_see(self): self.db.add(Article(id=8, slug='draft-febrile', title='Febrile draft', summary='Unpublished febrile notes', sections=[], status='draft', user_id=3)) self.db.commit() self.bank.user = self.bank.owner found = ai_mode_service.retrieve(self.db, self.bank.owner, 'febrile') self.assertNotIn(8, [s['id'] for s in found if s['kind'] == 'article']) # Question 3 is another user's private question. self.assertNotIn(3, [s['id'] for s in found if s['kind'] == 'question']) self.bank.user = self.bank.mod found = ai_mode_service.retrieve(self.db, self.bank.mod, 'febrile') self.assertIn(8, [s['id'] for s in found if s['kind'] == 'article']) def test_a_question_source_carries_the_stem_and_not_the_answer(self): found = ai_mode_service.retrieve(self.db, self.bank.owner, 'febrile seizure') questions = [s for s in found if s['kind'] == 'question'] self.assertTrue(questions) for source in questions: self.assertNotIn('Full explanation', source['text']) self.assertNotIn('yes', source['text'].split()) def test_a_curated_link_between_two_hits_lifts_both(self): self.db.add(QuestionArticleLink(question_id=1, article_id=7, section_id=None)) self.db.commit() found = ai_mode_service.retrieve(self.db, self.bank.owner, 'febrile seizure') curated = [s for s in found if s.get('curated')] # An educator tied these two together; both surfacing for one query is # evidence rather than coincidence. self.assertIn(('article', 7), [(s['kind'], s['id']) for s in curated]) self.assertIn(('question', 1), [(s['kind'], s['id']) for s in curated]) def test_the_prompt_carries_the_shortlist_and_the_rules(self): found = ai_mode_service.retrieve(self.db, self.bank.owner, 'febrile') prompt = ai_mode_service.build_prompt(found) self.assertIn('[[article:7]]', prompt) self.assertIn('never cite a marker that is not listed here', prompt) self.assertIn('Never reveal the answer to a practice question', prompt)