234 lines
8.3 KiB
Python
234 lines
8.3 KiB
Python
# 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)
|