Spaces:
Running
Running
File size: 2,755 Bytes
f1f74fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """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")
|