src/ctp_slack_bot/models/webvtt.py

from datetime import datetime, timedelta
from io import BytesIO
from more_itertools import windowed
from pydantic import BaseModel, ConfigDict, Field, field_validator
from types import MappingProxyType
from typing import Any, ClassVar, Literal, Mapping, Optional, Self
from webvtt import Caption, WebVTT

from ctp_slack_bot.utils import to_deep_immutable
from .base import Chunk, Content


class WebVTTFrame(BaseModel):
    """Represents a WebVTT frame"""

    model_config = ConfigDict(frozen=True)

    _SPEAKER_SPEECH_TEXT_SEPARATOR: ClassVar[str] = ": "

    identifier: str
    start: timedelta
    end: timedelta
    speaker: Optional[str] = None
    speech: str

    @classmethod
    def from_webvtt_caption(cls, caption: Caption, index: int) -> Self:
        identifier = caption.identifier if caption.identifier else str(index)
        start = timedelta(**caption.start_time.__dict__)
        end = timedelta(**caption.end_time.__dict__)
        match caption.text.split(cls._SPEAKER_SPEECH_TEXT_SEPARATOR, 1):
            case [speaker, speech]:
                return cls(identifier=identifier, start=start, end=end, speaker=speaker, speech=speech)
            case [speech]:
                return cls(identifier=identifier, start=start, end=end, speech=speech)


class WebVTTContent(Content):
    """Represents parsed WebVTT content."""

    CHUNK_FRAMES_OVERLAP: ClassVar[int] = 1
    CHUNK_FRAMES_WINDOW: ClassVar[int] = 5

    id: str
    metadata: Mapping[str, Any] = Field(default_factory=lambda: MappingProxyType({}))
    start_time: Optional[datetime]
    frames: tuple[WebVTTFrame, ...]

    def get_id(self: Self) -> str:
        return self.id

    def get_chunks(self: Self) -> tuple[Chunk]:
        windows = (tuple(filter(None, window))
                   for window
                   in windowed(self.frames, self.CHUNK_FRAMES_WINDOW, step=self.CHUNK_FRAMES_WINDOW-self.CHUNK_FRAMES_OVERLAP))
        return tuple(Chunk(text="\n\n".join(": ".join(filter(None, (frame.speaker, frame.speech)))
                                            for frame
                                            in frames),
                           parent_id=self.get_id(),
                           chunk_id=f"{frames[0].identifier}-{frames[-1].identifier}",
                           metadata={
                               "start": self.start_time + frames[0].start if self.start_time else None,
                               "end": self.start_time + frames[-1].end if self.start_time else None,
                               "speakers": (frame.speaker for frame in frames if frame.speaker)
                           })
                     for frames
                     in windows)

    def get_metadata(self: Self) -> Mapping[str, Any]:
        return MappingProxyType(self.metadata)