UmbraAudio/scripts/run_llm_attribution.py

from __future__ import annotations

import logging
from typing import Any

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:
        # Get chapters to process
        all_chapters = await db.get_chapters(1)
        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,
            )

            # Update database with new attributions
            for u in updated_utterances:
                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"),
                    )

            # Populate exception reviews from confidence scores
            await _populate_exceptions(db, chapter_id, updated_utterances)

            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, chapter_id: int, utterances: list[dict[str, Any]]) -> int:
    """Populate exception_reviews table from low-confidence attributions."""
    import json

    total_added = 0

    for u in utterances:
        if u.get("type") != "dialogue":
            continue

        confidence = u.get("attribution_confidence", 1.0)
        if confidence >= 0.7:
            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

        # Determine ambiguity type
        ambiguity_type = "low_confidence"
        reasoning = u.get("attribution_reasoning", "")
        if confidence == 0.0:
            ambiguity_type = "no_confidence"
        elif u.get("speaker_name") == "narrator":
            ambiguity_type = "narrator_dialogue"
        elif "pronoun" in reasoning.lower():
            ambiguity_type = "pronoun"
        elif "role" in reasoning.lower() or "unnamed" in reasoning.lower():
            ambiguity_type = "unnamed_role"
        elif "retroactive" in reasoning.lower():
            ambiguity_type = "retroactive_name"
        elif "generic" in reasoning.lower():
            ambiguity_type = "generic_role"

        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')
        """, (
            1, chapter_id, u["id"], u["utterance_index"],
            u.get("speaker_name", ""), u.get("speaker_name", ""),
            confidence,
            ambiguity_type,
            json.dumps([reasoning] if reasoning else ["Low confidence attribution"]),
            reasoning,
            json.dumps([]),
            u["text"][: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)