Video Captioning Team

from functools import partial
from os import environ
from pathlib import Path
from sys import argv

from agno.agent import Agent
from agno.models.openai import OpenAIResponses
from agno.team import Team
from agno.tools.moviepy_video import MoviePyVideoTools
from openai import OpenAI


def transcribe_audio_to_srt(audio_path: str, output_path: str) -> str:
    """Transcribe an audio file to a timestamped SRT file."""
    with open(audio_path, "rb") as audio_file:
        captions = OpenAI().audio.transcriptions.create(
            model="whisper-1", file=audio_file, response_format="srt"
        )
    Path(output_path).write_text(captions, encoding="utf-8")
    return output_path


caption_tools = MoviePyVideoTools(
    enable_process_video=False,
    enable_generate_captions=False,
    enable_embed_captions=True,
)
font_path = Path(environ["CAPTION_FONT_PATH"]).expanduser().resolve()
if not font_path.is_file():
    raise SystemExit(f"Font not found: {font_path}")
caption_tools.create_caption_clips = partial(
    caption_tools.create_caption_clips, font=str(font_path)
)

video_processor = Agent(
    name="Video Processor",
    role="Extract audio from the requested local video",
    model=OpenAIResponses(id="gpt-5.2"),
    tools=[MoviePyVideoTools(
        enable_process_video=True,
        enable_generate_captions=False,
        enable_embed_captions=False,
    )],
    instructions="Use extract_audio with the exact video and audio paths supplied.",
)
caption_generator = Agent(
    name="Caption Generator",
    role="Transcribe audio to timed captions and embed them in the video",
    model=OpenAIResponses(id="gpt-5.2"),
    tools=[transcribe_audio_to_srt, caption_tools],
    instructions=[
        "Use transcribe_audio_to_srt to obtain the timestamped SRT file.",
        "Pass that file directly to embed_captions without rewriting its timing.",
        "Follow the requested file paths exactly.",
    ],
)
caption_team = Team(
    name="Video Caption Team",
    members=[video_processor, caption_generator],
    model=OpenAIResponses(id="gpt-5.2"),
    instructions=[
        "First ask Video Processor to extract the audio.",
        "Then ask Caption Generator to transcribe it and embed the resulting SRT captions.",
        "Pass all exact file paths to the members and return the output video path.",
        "If a tool fails, report the error instead of claiming the video is complete.",
    ],
    markdown=True,
)

if __name__ == "__main__":
    if len(argv) != 2:
        raise SystemExit("Usage: python video_caption_generation.py /path/to/input.mp4")
    video_path = Path(argv[1]).expanduser().resolve()
    if not video_path.is_file():
        raise SystemExit(f"Video not found: {video_path}")
    audio_path = video_path.with_suffix(".wav")
    srt_path = video_path.with_suffix(".srt")
    output_path = video_path.with_name(f"{video_path.stem}_captioned.mp4")
    caption_team.print_response(
        f"""Generate captions using these exact paths.
Video: {video_path}
Audio: {audio_path}
SRT: {srt_path}
Captioned video: {output_path}"""
    )