analyze.py

from __future__ import annotations

import json
import os
from typing import Any

from .chunking import merge_highlight_candidates, select_top_highlights
from .prompts import (
    SYSTEM_PROMPT,
    build_chapter_prompt,
    build_highlight_prompt,
    build_overall_summary_prompt,
)
from .schemas import AnalysisBundle, ChapterSummary, HighlightCandidate, TranscriptChunk


class Analyzer:
    def __init__(self, *, model: str | None = None) -> None:
        api_key = os.environ.get("OPENAI_API_KEY")
        if not api_key:
            raise RuntimeError("OPENAI_API_KEY is not set.")
        try:
            from openai import OpenAI
        except ImportError as exc:
            raise RuntimeError("openai is not installed. Install dependencies from requirements.txt first.") from exc

        self.client = OpenAI(api_key=api_key)
        self.model = model or os.environ.get("OPENAI_MODEL", "gpt-4o-mini")

    def analyze(
        self,
        chunks: list[TranscriptChunk],
        *,
        highlight_count: int,
        min_clip_seconds: int,
        max_clip_seconds: int,
    ) -> AnalysisBundle:
        if not chunks:
            return AnalysisBundle(chapters=[], highlights=[], overall_summary="No transcript content was available.")

        chapters = self.create_chapters(chunks)
        highlight_pool: list[HighlightCandidate] = []
        for chunk in chunks:
            highlight_pool.extend(
                self.score_highlights(
                    chunk,
                    min_clip_seconds=min_clip_seconds,
                    max_clip_seconds=max_clip_seconds,
                )
            )
        merged = merge_highlight_candidates(highlight_pool)
        highlights = select_top_highlights(merged, limit=highlight_count)
        summary = self.create_overall_summary(chapters, highlight_count=len(highlights))
        return AnalysisBundle(chapters=chapters, highlights=highlights, overall_summary=summary)

    def create_chapters(self, chunks: list[TranscriptChunk]) -> list[ChapterSummary]:
        payload = self._response_json(
            name="chapter_summary",
            schema=_chapter_schema(),
            prompt=build_chapter_prompt(chunks),
        )
        chapters: list[ChapterSummary] = []
        for item in payload.get("chapters", []):
            chapters.append(
                ChapterSummary(
                    title=item["title"].strip(),
                    start=float(item["start"]),
                    end=float(item["end"]),
                    summary=item["summary"].strip(),
                )
            )
        return sorted(chapters, key=lambda item: (item.start, item.end))

    def score_highlights(
        self,
        chunk: TranscriptChunk,
        *,
        min_clip_seconds: int,
        max_clip_seconds: int,
    ) -> list[HighlightCandidate]:
        payload = self._response_json(
            name="highlight_candidates",
            schema=_highlight_schema(),
            prompt=build_highlight_prompt(
                chunk,
                min_clip_seconds=min_clip_seconds,
                max_clip_seconds=max_clip_seconds,
            ),
        )
        items: list[HighlightCandidate] = []
        for index, raw in enumerate(payload.get("highlights", []), start=1):
            start = max(chunk.start, float(raw["start"]))
            end = min(chunk.end, float(raw["end"]))
            if end <= start:
                continue
            items.append(
                HighlightCandidate(
                    id=f"{chunk.chunk_id}-h{index}",
                    start=start,
                    end=end,
                    duration=end - start,
                    score=float(raw["score"]),
                    title=raw["title"].strip(),
                    reason=raw["reason"].strip(),
                    excerpt=raw["excerpt"].strip(),
                    source_chunk_ids=[chunk.chunk_id],
                )
            )
        return items

    def create_overall_summary(self, chapters: list[ChapterSummary], *, highlight_count: int) -> str:
        response = self.client.responses.create(
            model=self.model,
            instructions=SYSTEM_PROMPT,
            input=build_overall_summary_prompt(chapters, highlight_count),
        )
        return (response.output_text or "").strip()

    def _response_json(self, *, name: str, schema: dict[str, Any], prompt: str) -> dict[str, Any]:
        response = self.client.responses.create(
            model=self.model,
            instructions=SYSTEM_PROMPT,
            input=prompt,
            text={
                "format": {
                    "type": "json_schema",
                    "name": name,
                    "schema": schema,
                    "strict": True,
                }
            },
        )
        content = (response.output_text or "").strip()
        if not content:
            raise RuntimeError(f"The model returned an empty response for {name}.")
        return json.loads(content)


def _chapter_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "chapters": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "title": {"type": "string"},
                        "start": {"type": "number"},
                        "end": {"type": "number"},
                        "summary": {"type": "string"},
                    },
                    "required": ["title", "start", "end", "summary"],
                    "additionalProperties": False,
                },
            }
        },
        "required": ["chapters"],
        "additionalProperties": False,
    }


def _highlight_schema() -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "highlights": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "start": {"type": "number"},
                        "end": {"type": "number"},
                        "score": {"type": "number"},
                        "title": {"type": "string"},
                        "reason": {"type": "string"},
                        "excerpt": {"type": "string"},
                    },
                    "required": ["start", "end", "score", "title", "reason", "excerpt"],
                    "additionalProperties": False,
                },
            }
        },
        "required": ["highlights"],
        "additionalProperties": False,
    }