# Copyright 2026 Umbra Tools (umbratools.com) # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import logging from typing import Any from umbraaudio.analysis.clarity import assess_utterance_clarity from umbraaudio.analysis.llm_attributor import LLMAttributor from umbraaudio.config import load_config from umbraaudio.database import Database logger = logging.getLogger(__name__) async def run_llm_attribution( project_dir: str = ".", chapter_ids: list[int] | None = None, ) -> dict[str, Any]: """ Run multi-pass LLM speaker attribution on chapters. This replaces the mechanical _attribute_speakers with LLM-based attribution that includes context analysis, confidence scoring, and exception detection. """ from pathlib import Path config = load_config(Path(project_dir)) db = Database(config.db_path) await db.connect() try: pid = await db.active_project_id() # Get chapters to process all_chapters = await db.get_chapters(pid) if chapter_ids: chapters = [ch for ch in all_chapters if ch["id"] in chapter_ids] else: chapters = all_chapters # Initialize LLM attributor attributor = LLMAttributor(config.llm, data_dir=config.data_dir) results = { "chapters_processed": 0, "total_utterances": 0, "attributed": 0, "exceptions_found": 0, "quality_gate_passed": True, } for ch in chapters: chapter_id = ch["id"] chapter_title = ch.get("title") or f"Chapter {ch['chapter_index'] + 1}" logger.info(f"Running LLM attribution for chapter {chapter_id}: {chapter_title}") # Get utterances from database utterances = await db.get_utterances(chapter_id) if not utterances: logger.warning(f"Chapter {chapter_id}: no utterances found") continue # Prepare chapter data for context analysis chapter_data = { "id": chapter_id, "title": chapter_title, "source_text": ch.get("source_text", ""), "utterances": [dict(u) for u in utterances], } # Run multi-pass LLM attribution updated_utterances = attributor.attribute_speakers( utterances=[dict(u) for u in utterances], chapter_title=chapter_title, chapter_id=chapter_id, chapter_data=chapter_data, ) # Deterministic clarity + review flags (evidence kind, pronoun # uniqueness, QA issues, confidence thresholds) alias_map: dict[str, str] = {} try: import json as _json from pathlib import Path as _Path sp = _Path(config.data_dir) / "speakers.json" if sp.exists(): with open(sp) as f: for s in _json.load(f).get("speakers", []): primary = s.get("primary_name", "") if primary: alias_map[primary.lower()] = primary for a in s.get("aliases", []): if a: alias_map[a.lower()] = primary except Exception: pass clarity_aware = assess_utterance_clarity( [dict(u) for u in updated_utterances], chapter_id, alias_map, ) # Update database with new attributions for u in clarity_aware: if u.get("type") == "dialogue" and "speaker_name" in u: await db.update_utterance_speaker( u["id"], u["speaker_name"], u.get("attribution_confidence", 0.5), u.get("attribution_reasoning", ""), u.get("emotion"), u.get("emotion_confidence"), evidence_kind=u.get("attribution_evidence_kind"), tag_name=u.get("attribution_tag_name"), uncertain=1 if u.get("attribution_uncertain") else 0, needs_review=1 if u.get("needs_review") else 0, review_reason=u.get("review_reason") or "", ) # Populate exception reviews from clarity flags await _populate_exceptions(db, pid, chapter_id, clarity_aware) results["chapters_processed"] += 1 results["total_utterances"] += len([u for u in updated_utterances if u.get("type") == "dialogue"]) results["attributed"] += len([u for u in updated_utterances if u.get("type") == "dialogue" and u.get("speaker_name") != "narrator"]) # Check quality gate quality_gate = await _check_quality_gate(db) results["quality_gate_passed"] = quality_gate["passed"] results["exception_rate"] = quality_gate["exception_rate"] results["exceptions_found"] = quality_gate["total_exceptions"] logger.info(f"LLM attribution complete: {results}") return results finally: await db.close() async def _populate_exceptions( db: Database, project_id: int, chapter_id: int, utterances: list[dict[str, Any]] ) -> int: """Populate exception_reviews table from clarity review flags.""" import json total_added = 0 for u in utterances: if u.get("type") != "dialogue": continue if not u.get("needs_review"): continue # Check if already in exception_reviews cursor = await db.conn.execute( "SELECT id FROM exception_reviews WHERE utterance_id = ?", (u["id"],) ) existing = await cursor.fetchone() if existing: continue reason = u.get("review_reason") or "" evidence_kind = u.get("attribution_evidence_kind") or "" if "pronoun" in reason: ambiguity_type = "pronoun" elif "role" in reason: ambiguity_type = "unnamed_role" elif evidence_kind in ("none", "address_only"): ambiguity_type = "ambiguous_tag" elif u.get("speaker_name") == "narrator": ambiguity_type = "narrator_dialogue" else: ambiguity_type = "low_confidence" await db.conn.execute(""" INSERT INTO exception_reviews ( project_id, chapter_id, utterance_id, utterance_index, current_speaker, attributed_speaker, confidence, ambiguity_type, evidence, reasoning, alternatives, text_preview, status ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending') """, ( project_id, chapter_id, u["id"], u["utterance_index"], u.get("speaker_name", ""), u.get("speaker_name", ""), u.get("attribution_confidence", 0.0), ambiguity_type, json.dumps([reason] if reason else ["Needs review"]), reason, json.dumps([]), (u.get("text") or "")[:200], )) total_added += 1 return total_added async def _check_quality_gate(db: Database) -> dict[str, Any]: """Check if exception rate exceeds 1% quality gate.""" cursor = await db.conn.execute( "SELECT COUNT(*) as cnt FROM utterances WHERE type = 'dialogue'" ) total_utterances = await cursor.fetchone() cursor = await db.conn.execute( "SELECT COUNT(*) as cnt FROM exception_reviews WHERE status = 'pending'" ) total_exceptions = await cursor.fetchone() total_utt = total_utterances["cnt"] if total_utterances else 0 total_exc = total_exceptions["cnt"] if total_exceptions else 0 exception_rate = (total_exc / total_utt * 100) if total_utt > 0 else 0 return { "passed": exception_rate <= 1.0, "exception_rate": exception_rate, "total_exceptions": total_exc, "total_utterances": total_utt, "threshold": 1.0, } # CLI entry point if __name__ == "__main__": import asyncio import sys chapter_ids = None if len(sys.argv) > 1: chapter_ids = [int(x) for x in sys.argv[1].split(",")] result = asyncio.run(run_llm_attribution(chapter_ids=chapter_ids)) print(result)