"""Span-based speaker attribution: request validation, prompt, parsing, route.""" from __future__ import annotations import os import tempfile import unittest from unittest import mock os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-spans-test-config-")) from fastapi import FastAPI from fastapi.testclient import TestClient from core.attribution_spans import ( SpanRequestError, build_messages, max_tokens_for, parse_response, validate_only_ids, validate_spans, ) import routes.conversation as conversation SPANS = [ {"id": "N0", "kind": "narration", "text": "Der Zwerg schneuzte sich."}, {"id": "D1", "kind": "dialogue", "text": "Na los, steh hier nicht weiter rum."}, {"id": "N2", "kind": "narration", "text": "sagte Linosch."}, {"id": "D3", "kind": "dialogue", "text": "Vielen Dank!"}, ] class ValidateSpansTests(unittest.TestCase): def test_accepts_a_well_formed_window(self): self.assertEqual([s["id"] for s in validate_spans(SPANS)], ["N0", "D1", "N2", "D3"]) def test_rejects_bad_shapes(self): bad = [ None, [], "x", [1], [{"id": "X1", "kind": "dialogue", "text": "a"}], # id pattern [{"id": "D1", "kind": "narration", "text": "a"}], # kind/id mismatch [{"id": "D1", "kind": "dialogue", "text": " "}], # empty text [{"id": "D1", "kind": "dialogue", "text": "a"}, {"id": "D1", "kind": "dialogue", "text": "b"}], [{"id": "N1", "kind": "narration", "text": "only narration"}], # nothing to attribute [{"id": "D%d" % i, "kind": "dialogue", "text": "a"} for i in range(401)], [{"id": "D1", "kind": "dialogue", "text": "x" * 20_001}], ] for case in bad: with self.subTest(case=str(case)[:60]): with self.assertRaises(SpanRequestError): validate_spans(case) class PromptTests(unittest.TestCase): def test_prompt_lists_every_line_and_the_ids_to_attribute(self): system, user = build_messages(SPANS, ["Marcian"], "Marcian: Hallo", "Danach ging er.", "German") self.assertIn("Do NOT rewrite", system) self.assertIn("German", system) self.assertIn('[D1] "Na los, steh hier nicht weiter rum."', user) self.assertIn("[N2] sagte Linosch.", user) self.assertIn("Known characters so far: Marcian", user) self.assertIn("Recent dialogue", user) self.assertIn("Attribute these lines: D1, D3", user) def test_only_ids_restricts_the_request_for_a_retry(self): _, user = build_messages(SPANS, only_ids=["D3"]) self.assertIn("Attribute these lines: D3", user) self.assertNotIn("D1, D3", user) def test_token_budget_grows_with_dialogue_and_is_capped(self): many = lambda n: [{"id": "D%d" % i, "kind": "dialogue", "text": "a"} for i in range(n)] self.assertLess(max_tokens_for(SPANS), max_tokens_for(many(30))) self.assertLessEqual(max_tokens_for(many(400)), 4096) class AttributeOnlyTests(unittest.TestCase): def test_attribute_only_keeps_valid_dialogue_ids_in_order(self): spans = validate_spans(SPANS) self.assertEqual(validate_only_ids(["D3", "D1", "D3", "N0", "D9", 5], spans), ["D3", "D1"]) self.assertEqual(validate_only_ids("D1", spans), []) def test_prompt_restricts_the_request_to_those_lines(self): _, user = build_messages(SPANS, only_ids=["D1"]) self.assertIn("Attribute these lines: D1", user) self.assertNotIn("D3", user.split("Attribute these lines:")[1]) class ParseTests(unittest.TestCase): def test_parses_clean_json_and_normalises_names(self): raw = '{"speakers":[{"id":"D1","speaker":" Linosch ","emotion":"gereizt"},{"id":"D3","speaker":"unbekannt","emotion":""}]}' out = parse_response(raw, validate_spans(SPANS)) self.assertEqual(out["speakers"]["D1"], {"speaker": "Linosch", "emotion": "gereizt"}) self.assertEqual(out["speakers"]["D3"]["speaker"], "Unknown") self.assertEqual(out["missing"], []) self.assertEqual(out["characters"], ["Linosch"]) def test_ignores_reasoning_fences_and_unknown_ids(self): raw = ('D1 is Linosch, {"id":"D1","speaker":"Wrong"}\n```json\n' '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":""},{"id":"D9","speaker":"Ghost"}]}\n```') out = parse_response(raw, validate_spans(SPANS)) self.assertEqual(out["speakers"]["D1"]["speaker"], "Linosch") self.assertNotIn("D9", out["speakers"]) self.assertEqual(out["missing"], ["D3"]) def test_salvages_a_truncated_answer(self): raw = '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":"x"},{"id":"D3","speak' out = parse_response(raw, validate_spans(SPANS)) self.assertEqual(list(out["speakers"]), ["D1"]) self.assertEqual(out["missing"], ["D3"]) def test_narrator_answer_for_a_non_speech_line_is_kept_out_of_the_cast(self): raw = '{"speakers":[{"id":"D1","speaker":"Erzähler"},{"id":"D3","speaker":"Marcian"}]}' out = parse_response(raw, validate_spans(SPANS)) self.assertEqual(out["speakers"]["D1"]["speaker"], "Narrator") self.assertEqual(out["characters"], ["Marcian"]) class RouteTests(unittest.TestCase): def setUp(self): app = FastAPI() app.include_router(conversation.router) self.client = TestClient(app) def test_rejects_an_invalid_request_before_calling_the_llm(self): with mock.patch.object(conversation, "_span_llm_answer") as llm: r = self.client.post("/api/attribute-spans", json={"spans": [{"id": "N1", "kind": "narration", "text": "x"}]}) self.assertEqual(r.status_code, 400) llm.assert_not_called() def test_returns_speakers_and_never_the_text(self): answer = '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":"gereizt"},{"id":"D3","speaker":"Alrik","emotion":""}]}' with mock.patch.object(conversation, "_span_llm_answer", return_value=answer) as llm: r = self.client.post("/api/attribute-spans", json={"spans": SPANS, "known_characters": ["Alrik"]}) self.assertEqual(r.status_code, 200) body = r.json() self.assertEqual(body["speakers"]["D3"]["speaker"], "Alrik") self.assertEqual(body["missing"], []) self.assertNotIn("Zwerg", r.text) llm.assert_called_once() def test_retries_only_the_lines_the_model_skipped(self): first = '{"speakers":[{"id":"D1","speaker":"Linosch","emotion":""}]}' second = '{"speakers":[{"id":"D3","speaker":"Alrik","emotion":""}]}' with mock.patch.object(conversation, "_span_llm_answer", side_effect=[first, second]) as llm: r = self.client.post("/api/attribute-spans", json={"spans": SPANS}) body = r.json() self.assertEqual(llm.call_count, 2) self.assertEqual({k: v["speaker"] for k, v in body["speakers"].items()}, {"D1": "Linosch", "D3": "Alrik"}) self.assertEqual(body["missing"], []) retry_prompt = llm.call_args_list[1].args[0]["messages"][1]["content"] self.assertIn("Attribute these lines: D3", retry_prompt) def test_attribute_only_limits_the_request_and_the_missing_accounting(self): answer = '{"speakers":[{"id":"D3","speaker":"Alrik","emotion":""}]}' with mock.patch.object(conversation, "_span_llm_answer", return_value=answer) as llm: r = self.client.post("/api/attribute-spans", json={"spans": SPANS, "attribute_only": ["D3"]}) body = r.json() self.assertEqual(body["speakers"]["D3"]["speaker"], "Alrik") self.assertEqual(body["missing"], []) llm.assert_called_once() self.assertIn("Attribute these lines: D3", llm.call_args.args[0]["messages"][1]["content"]) def test_reports_failure_instead_of_erroring_when_the_llm_is_down(self): with mock.patch.object(conversation, "_span_llm_answer", side_effect=RuntimeError("LLM HTTP 503")): r = self.client.post("/api/attribute-spans", json={"spans": SPANS}) self.assertEqual(r.status_code, 200) body = r.json() self.assertTrue(body["failed"]) self.assertEqual(body["missing"], ["D1", "D3"]) if __name__ == "__main__": unittest.main()