UmbraAudio/umbraaudio/analysis/llm_segmenter.py

from __future__ import annotations

import json
import logging
import re
import time
from pathlib import Path
from typing import Any

from openai import OpenAI

from umbraaudio.config import LLMConfig
from umbraaudio.voices.lmstudio_gate import queued_chat_completion

logger = logging.getLogger(__name__)

PROMPT_PATH = Path(__file__).parent / "CHARACTER_SEGMENTATION_PROMPT.md"

CHUNK_SIZE = 15000


class LLMSegmenter:
    def __init__(self, config: LLMConfig, model: str | None = None, data_dir: Path | None = None):
        self.config = config
        self.model = model or config.lmstudio_model
        self.data_dir = data_dir or Path(".umbra")

    def _get_client(self) -> OpenAI:
        return OpenAI(
            base_url=self.config.lmstudio_base_url,
            api_key="lm-studio",
        )

    def _load_prompt(self) -> str:
        return PROMPT_PATH.read_text(encoding="utf-8")

    def _load_character_profiles(self) -> str:
        speakers_path = self.data_dir / "speakers.json"
        if not speakers_path.exists():
            return "No character profiles available."

        with open(speakers_path) as f:
            data = json.load(f)

        speakers = data.get("speakers", [])
        if not speakers:
            return "No character profiles available."

        lines = []
        for s in speakers:
            name = s.get("primary_name", "Unknown")
            demos = s.get("demographics", {})
            dtype = demos.get("type", "Human")
            gender = demos.get("gender", "Unknown")
            traits = ", ".join(s.get("vocal_traits", [])) or "none"
            aliases = ", ".join(s.get("aliases", [])) or "none"

            lines.append(f"- {name} ({dtype}, {gender}): traits=[{traits}], aliases=[{aliases}]")

        return "\n".join(lines)

    def segment_chapter(
        self,
        chapter_text: str,
        chapter_id: int,
    ) -> list[dict[str, Any]]:
        scenes = self._split_into_scenes(chapter_text)
        all_utterances: list[dict[str, Any]] = []

        for scene_idx, scene_text in enumerate(scenes):
            logger.info(
                "Chapter %d, scene %d/%d: %d chars",
                chapter_id,
                scene_idx + 1,
                len(scenes),
                len(scene_text),
            )

            if len(scene_text) <= CHUNK_SIZE:
                utterances = self._call_llm(scene_text, chapter_id)
                all_utterances.extend(utterances)
            else:
                paragraphs = self._split_into_paragraphs(scene_text)
                chunk_utterances: list[dict[str, Any]] = []
                current_chunk = ""

                for para in paragraphs:
                    if len(current_chunk) + len(para) > CHUNK_SIZE and current_chunk:
                        utterances = self._call_llm(current_chunk, chapter_id)
                        chunk_utterances.extend(utterances)
                        current_chunk = ""
                        time.sleep(1)

                    current_chunk += para + "\n\n"

                if current_chunk.strip():
                    utterances = self._call_llm(current_chunk, chapter_id)
                    chunk_utterances.extend(utterances)

                all_utterances.extend(chunk_utterances)

            if scene_idx < len(scenes) - 1:
                all_utterances.append(
                    {
                        "speaker": "scene_break",
                        "text": "<<<>>>",
                        "type": "narration",
                    }
                )

            time.sleep(1)

        all_utterances = self._renumber(all_utterances)

        logger.info(
            "Chapter %d complete: %d utterances",
            chapter_id,
            len(all_utterances),
        )

        return all_utterances

    def _split_into_scenes(self, text: str) -> list[str]:
        parts = re.split(r"\n\s*<<<>>>\s*\n", text)
        return [p.strip() for p in parts if p.strip()]

    def _split_into_paragraphs(self, text: str) -> list[str]:
        return [p.strip() for p in text.split("\n\n") if p.strip()]

    def _call_llm(self, text: str, chapter_id: int) -> list[dict[str, Any]]:
        prompt_template = self._load_prompt()
        profiles = self._load_character_profiles()

        prompt = prompt_template.replace("{chapter_text}", text)
        prompt = prompt.replace("{character_profiles}", profiles)

        client = self._get_client()

        try:
            response = queued_chat_completion(
                client,
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.2,
                max_tokens=16384,
            )
        except Exception as e:
            logger.error("LLM call failed: %s", e)
            return []

        content = response.choices[0].message.content or ""
        content = content.strip()

        if not content:
            logger.warning("Empty LLM response")
            return []

        utterances = self._parse_response(content)
        utterances = self._validate(utterances)

        return utterances

    def _parse_response(self, content: str) -> list[dict[str, Any]]:
        stripped = content.strip()

        if stripped.startswith("```"):
            first_newline = stripped.find("\n")
            if first_newline > 0:
                stripped = stripped[first_newline + 1 :]
            if stripped.endswith("```"):
                stripped = stripped[:-3].strip()

        utterances = self._try_json_parse(stripped)
        if utterances is not None:
            return utterances

        json_start = stripped.find("[")
        json_end = stripped.rfind("]") + 1
        if json_start >= 0 and json_end > json_start:
            candidate = stripped[json_start:json_end]
            utterances = self._try_json_parse(candidate)
            if utterances is not None:
                return utterances

            last_brace = candidate.rfind("}")
            if last_brace > 0:
                candidate = candidate[: last_brace + 1]
                open_brackets = candidate.count("[") - candidate.count("]")
                candidate += "]" * open_brackets
                utterances = self._try_json_parse(candidate)
                if utterances is not None:
                    return utterances

        logger.error("Failed to parse LLM segmentation response")
        return []

    def _try_json_parse(self, text: str) -> list[dict[str, Any]] | None:
        try:
            result = json.loads(text)
            if isinstance(result, list):
                return result
            if isinstance(result, dict) and "utterances" in result:
                return result["utterances"]
            return None
        except json.JSONDecodeError:
            pass

        fixed = text
        for _ in range(10):
            try:
                result = json.loads(fixed)
                if isinstance(result, list):
                    return result
                return None
            except json.JSONDecodeError:
                last_brace = fixed.rfind("}")
                last_bracket = fixed.rfind("]")
                last_comma = fixed.rfind(",")
                cut = max(last_brace, last_bracket, last_comma)
                if cut > 0:
                    fixed = fixed[:cut]
                    if fixed.endswith(","):
                        fixed = fixed[:-1]
                else:
                    break

        return None

    def _validate(self, utterances: list[dict[str, Any]]) -> list[dict[str, Any]]:
        valid = []
        for u in utterances:
            if not isinstance(u, dict):
                continue
            speaker = u.get("speaker", "narrator")
            text = u.get("text", "").strip()
            if not text:
                continue
            valid.append(
                {
                    "speaker": speaker,
                    "text": text,
                    "type": self._utterance_type(speaker),
                }
            )
        return valid

    def _utterance_type(self, speaker: str) -> str:
        if speaker in ("narrator", "scene_break", "title"):
            return "narration"
        return "dialogue"

    def _renumber(self, utterances: list[dict[str, Any]]) -> list[dict[str, Any]]:
        for i, u in enumerate(utterances):
            u["index"] = i
        return utterances