mARTin-B78
37e2622d4e
Verify / recast now run span-natively (no text re-emission): a second differently-windowed pass marks lines agreed/disputed and records the other answer, without changing any speaker (~9 min vs ~49 min). Measured on the real pipeline, overruling disputed lines with a reasoning adjudication made them worse (66.5% -> 61.4%), so that step was removed; agreed lines were 83% right vs 48.5% for disputed ones, which is the signal used for confidence and review. Auto-repair no longer runs after a span cast; segments carry a spanId. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
168 lines
8.2 KiB
Python
168 lines
8.2 KiB
Python
"""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 = ('<think>D1 is Linosch, {"id":"D1","speaker":"Wrong"}</think>\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()
|