shared/manim_module.py

"""Python SDK / Module Client for Math Agent integration."""
from __future__ import annotations

import asyncio
import logging
from typing import Any
from uuid import UUID

import httpx

from shared.schemas.math_agent import (
    MathRenderRequest,
    MathRenderResponse,
    VisualizationSpec,
)

logger = logging.getLogger(__name__)


class ManimModule:
    """High-level client for Math Agent to invoke the Manim video generation service."""

    def __init__(
        self,
        base_url: str = "http://localhost:8000",
        internal_token: str | None = None,
        bearer_token: str | None = None,
        timeout: float = 60.0,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.headers: dict[str, str] = {}
        if internal_token:
            self.headers["X-Internal-Token"] = internal_token
        if bearer_token:
            self.headers["Authorization"] = f"Bearer {bearer_token}"
        self.timeout = timeout

    async def generate(
        self,
        spec: VisualizationSpec | dict[str, Any],
        callback_url: str | None = None,
    ) -> MathRenderResponse:
        """Submit a VisualizationSpec to start asynchronous video generation."""
        parsed_spec = spec if isinstance(spec, VisualizationSpec) else VisualizationSpec.model_validate(spec)
        req = MathRenderRequest(spec=parsed_spec, callback_url=callback_url)

        async with httpx.AsyncClient(base_url=self.base_url, headers=self.headers, timeout=self.timeout) as client:
            resp = await client.post("/v1/math/generate", json=req.model_dump(mode="json"))
            resp.raise_for_status()
            return MathRenderResponse.model_validate(resp.json())

    async def get_status(self, job_id: UUID | str) -> MathRenderResponse:
        """Query current status, progress, or completed video URL."""
        async with httpx.AsyncClient(base_url=self.base_url, headers=self.headers, timeout=self.timeout) as client:
            resp = await client.get(f"/v1/math/jobs/{job_id}")
            resp.raise_for_status()
            return MathRenderResponse.model_validate(resp.json())

    async def wait_for_completion(
        self,
        job_id: UUID | str,
        poll_interval: float = 3.0,
        timeout: float = 600.0,
    ) -> MathRenderResponse:
        """Poll until the video rendering is completed or failed."""
        elapsed = 0.0
        while elapsed < timeout:
            status_resp = await self.get_status(job_id)
            if status_resp.status in {"completed", "failed"}:
                return status_resp
            await asyncio.sleep(poll_interval)
            elapsed += poll_interval

        raise TimeoutError(f"Rendering job {job_id} did not finish within {timeout} seconds")