Add Grok media providers and improve selector routing

This commit is contained in:
calesthio
2026-04-05 15:31:37 -07:00
parent 6a6d456e50
commit 7ca04e66d8
15 changed files with 1251 additions and 28 deletions
+290
View File
@@ -0,0 +1,290 @@
"""xAI Grok video generation."""
from __future__ import annotations
import base64
import mimetypes
import os
import time
from pathlib import Path
from typing import Any
from tools.base_tool import (
BaseTool,
Determinism,
ExecutionMode,
ResourceProfile,
RetryPolicy,
ToolResult,
ToolRuntime,
ToolStability,
ToolStatus,
ToolTier,
)
def _file_to_data_uri(path_str: str) -> str:
path = Path(path_str)
if not path.exists():
raise FileNotFoundError(f"Input file not found: {path}")
mime_type, _ = mimetypes.guess_type(path.name)
if not mime_type:
mime_type = "application/octet-stream"
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
return f"data:{mime_type};base64,{encoded}"
def _normalize_media_ref(url_value: str | None, path_value: str | None) -> dict[str, str] | None:
if url_value:
return {"url": url_value}
if path_value:
return {"url": _file_to_data_uri(path_value)}
return None
class GrokVideo(BaseTool):
name = "grok_video"
version = "0.1.0"
tier = ToolTier.GENERATE
capability = "video_generation"
provider = "grok"
stability = ToolStability.BETA
execution_mode = ExecutionMode.SYNC
determinism = Determinism.STOCHASTIC
runtime = ToolRuntime.API
dependencies = []
install_instructions = (
"Set XAI_API_KEY to your xAI API key.\n"
" Get one from the xAI developer console"
)
agent_skills = ["grok-media", "ai-video-gen"]
capabilities = ["text_to_video", "image_to_video", "reference_to_video"]
supports = {
"text_to_video": True,
"image_to_video": True,
"reference_to_video": True,
"reference_image": True,
"multiple_reference_images": True,
}
best_for = [
"reference-conditioned video generation",
"product placement or character-consistent motion clips",
"xAI-native image-guided and text-guided short videos",
]
not_good_for = ["offline generation", "very long clips"]
fallback_tools = ["veo_video", "runway_video", "kling_video", "minimax_video"]
input_schema = {
"type": "object",
"required": ["prompt"],
"properties": {
"prompt": {"type": "string"},
"operation": {
"type": "string",
"enum": ["text_to_video", "image_to_video", "reference_to_video"],
"default": "text_to_video",
},
"model": {
"type": "string",
"enum": ["grok-imagine-video"],
"default": "grok-imagine-video",
},
"duration": {
"type": "integer",
"minimum": 2,
"maximum": 10,
"default": 5,
},
"aspect_ratio": {
"type": "string",
"enum": ["16:9", "9:16", "1:1"],
"default": "16:9",
},
"resolution": {
"type": "string",
"enum": ["480p", "720p"],
"default": "720p",
},
"image_url": {"type": "string", "description": "Reference image URL for image_to_video"},
"image_path": {"type": "string", "description": "Local reference image path for image_to_video"},
"reference_image_urls": {
"type": "array",
"items": {"type": "string"},
"description": "Reference image URLs for reference_to_video",
},
"reference_image_paths": {
"type": "array",
"items": {"type": "string"},
"description": "Local reference image paths for reference_to_video",
},
"output_path": {"type": "string"},
"poll_interval_seconds": {"type": "integer", "minimum": 2, "default": 5},
"timeout_seconds": {"type": "integer", "minimum": 30, "default": 900},
},
}
resource_profile = ResourceProfile(
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True
)
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
idempotency_key_fields = ["prompt", "operation", "model", "duration", "aspect_ratio", "resolution"]
side_effects = ["writes video file to output_path", "calls xAI video API"]
user_visible_verification = ["Watch generated clip for motion quality and prompt fidelity"]
def get_status(self) -> ToolStatus:
if os.environ.get("XAI_API_KEY"):
return ToolStatus.AVAILABLE
return ToolStatus.UNAVAILABLE
@staticmethod
def _normalize_resolution(value: str | None) -> str:
if value == "540p":
return "480p"
return value or "720p"
@staticmethod
def _input_image_count(inputs: dict[str, Any]) -> int:
count = 0
if inputs.get("image_url") or inputs.get("image_path"):
count += 1
count += len(inputs.get("reference_image_urls") or [])
count += len(inputs.get("reference_image_paths") or [])
return count
def estimate_cost(self, inputs: dict[str, Any]) -> float:
duration = int(inputs.get("duration", 5))
resolution = self._normalize_resolution(inputs.get("resolution"))
base_per_second = 0.07 if resolution == "720p" else 0.05
input_image_cost = self._input_image_count(inputs) * 0.002
# xAI currently publishes Grok Imagine Video at $0.05/sec for 480p,
# $0.07/sec for 720p, plus $0.002 per input image.
return base_per_second * duration + input_image_cost
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
duration = int(inputs.get("duration", 5))
return 90.0 + duration * 8.0
def _build_payload(self, inputs: dict[str, Any]) -> dict[str, Any]:
operation = inputs.get("operation", "text_to_video")
payload: dict[str, Any] = {
"model": inputs.get("model", "grok-imagine-video"),
"prompt": inputs["prompt"],
}
if operation != "reference_to_video":
payload["duration"] = int(inputs.get("duration", 5))
if inputs.get("aspect_ratio"):
payload["aspect_ratio"] = inputs["aspect_ratio"]
if inputs.get("resolution"):
payload["resolution"] = self._normalize_resolution(inputs["resolution"])
if operation == "image_to_video":
image = _normalize_media_ref(inputs.get("image_url"), inputs.get("image_path"))
if not image:
raise ValueError("image_to_video requires image_url or image_path")
payload["image"] = image
elif operation == "reference_to_video":
refs = [{"url": url} for url in (inputs.get("reference_image_urls") or [])]
refs.extend(
{"url": _file_to_data_uri(path)}
for path in (inputs.get("reference_image_paths") or [])
)
if not refs:
raise ValueError(
"reference_to_video requires reference_image_urls or reference_image_paths"
)
payload["reference_images"] = refs
payload["duration"] = int(inputs.get("duration", 5))
if inputs.get("aspect_ratio"):
payload["aspect_ratio"] = inputs["aspect_ratio"]
if inputs.get("resolution"):
payload["resolution"] = self._normalize_resolution(inputs["resolution"])
return payload
def execute(self, inputs: dict[str, Any]) -> ToolResult:
api_key = os.environ.get("XAI_API_KEY")
if not api_key:
return ToolResult(
success=False,
error="XAI_API_KEY not set. " + self.install_instructions,
)
import requests
from tools.video._shared import probe_output
start = time.time()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
try:
payload = self._build_payload(inputs)
response = requests.post(
"https://api.x.ai/v1/videos/generations",
headers=headers,
json=payload,
timeout=60,
)
response.raise_for_status()
request_id = response.json()["request_id"]
timeout_seconds = int(inputs.get("timeout_seconds", 900))
poll_interval = int(inputs.get("poll_interval_seconds", 5))
deadline = time.time() + timeout_seconds
result_data: dict[str, Any] | None = None
while time.time() < deadline:
result = requests.get(
f"https://api.x.ai/v1/videos/{request_id}",
headers={"Authorization": headers["Authorization"]},
timeout=30,
)
result.raise_for_status()
result_data = result.json()
status = result_data.get("status")
if status == "done":
break
if status in {"failed", "expired"}:
detail = result_data.get("error") or result_data.get("message") or status
return ToolResult(success=False, error=f"Grok video generation {status}: {detail}")
time.sleep(poll_interval)
if not result_data or result_data.get("status") != "done":
return ToolResult(success=False, error="Grok video generation timed out")
video_url = (result_data.get("video") or {}).get("url")
if not video_url:
return ToolResult(success=False, error="xAI video output missing url")
download = requests.get(video_url, timeout=300)
download.raise_for_status()
output_path = Path(inputs.get("output_path", "grok_video_output.mp4"))
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_bytes(download.content)
except Exception as e:
return ToolResult(success=False, error=f"Grok video generation failed: {e}")
probed = probe_output(output_path)
return ToolResult(
success=True,
data={
"provider": "grok",
"model": payload["model"],
"prompt": inputs["prompt"],
"operation": inputs.get("operation", "text_to_video"),
"request_id": request_id,
"output": str(output_path),
"output_path": str(output_path),
"format": "mp4",
**probed,
},
artifacts=[str(output_path)],
cost_usd=self.estimate_cost(inputs),
duration_seconds=round(time.time() - start, 2),
model=payload["model"],
)
+92 -6
View File
@@ -49,7 +49,11 @@ class VideoSelector(BaseTool):
"default": "auto",
},
"allowed_providers": {"type": "array", "items": {"type": "string"}},
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video", "rank"], "default": "text_to_video"},
"operation": {
"type": "string",
"enum": ["text_to_video", "image_to_video", "reference_to_video", "rank"],
"default": "text_to_video",
},
"aspect_ratio": {
"type": "string",
"enum": ["16:9", "9:16", "1:1"],
@@ -68,10 +72,24 @@ class VideoSelector(BaseTool):
"type": "string",
"description": "URL of a reference image for image_to_video.",
},
"reference_image_urls": {
"type": "array",
"items": {"type": "string"},
"description": "Reference image URLs for providers that support reference-conditioned video.",
},
"reference_image_paths": {
"type": "array",
"items": {"type": "string"},
"description": "Local reference image paths for providers that support reference-conditioned video.",
},
"image_url": {
"type": "string",
"description": "Alias for reference_image_url (used by some providers like Kling via fal.ai).",
},
"resolution": {
"type": "string",
"description": "Resolution hint for providers that support named output resolutions.",
},
"output_path": {"type": "string"},
},
}
@@ -103,23 +121,23 @@ class VideoSelector(BaseTool):
return ToolStatus.UNAVAILABLE
def estimate_cost(self, inputs: dict[str, object]) -> float:
candidates = self._providers()
candidates = self._filter_candidates(inputs, self._providers())
if not candidates:
return 0.0
tool, _ = self._select_best_tool(inputs, candidates, inputs.get("task_context", {}))
tool, _ = self._select_best_tool(inputs, candidates, self._prepare_task_context(inputs))
return tool.estimate_cost(inputs) if tool else 0.0
def estimate_runtime(self, inputs: dict[str, object]) -> float:
candidates = self._providers()
if not candidates:
return 0.0
tool, _ = self._select_best_tool(inputs, candidates, inputs.get("task_context", {}))
tool, _ = self._select_best_tool(inputs, candidates, self._prepare_task_context(inputs))
return tool.estimate_runtime(inputs) if tool else 0.0
def execute(self, inputs: dict[str, object]) -> ToolResult:
from lib.scoring import rank_providers
task_context = inputs.get("task_context", {})
task_context = self._prepare_task_context(inputs)
candidates = self._providers()
# Rank mode — return scored provider rankings without generating
@@ -128,8 +146,9 @@ class VideoSelector(BaseTool):
return ToolResult(
success=True,
data={
"rankings": [r.to_dict() for r in rankings],
"rankings": self._serialize_rankings(candidates, rankings),
"explanation": "\n".join(r.explain() for r in rankings[:5]),
"normalized_task_context": task_context,
},
)
@@ -159,9 +178,11 @@ class VideoSelector(BaseTool):
result = tool.execute(adapted)
if result.success:
result.data.setdefault("selected_tool", tool.name)
result.data["selected_provider"] = tool.provider
result.data["selection_reason"] = score.explain() if score else f"Selected {tool.provider} ({tool.name})"
if score:
result.data["provider_score"] = score.to_dict()
result.data.update(self._tool_context_payload(tool))
result.data["alternatives_considered"] = [
t.name for t in candidates
if t.name != tool.name and t.get_status().value == "available"
@@ -185,6 +206,7 @@ class VideoSelector(BaseTool):
allowed = set(inputs.get("allowed_providers") or [])
if allowed:
candidates = [tool for tool in candidates if tool.provider in allowed]
candidates = self._filter_candidates(inputs, candidates)
env_hint = os.environ.get("VIDEO_GEN_LOCAL_MODEL", "").lower()
env_map = {
@@ -219,3 +241,67 @@ class VideoSelector(BaseTool):
return tool_by_provider[score.provider], score
return None, None
def _prepare_task_context(self, inputs: dict[str, object]) -> dict[str, object]:
from lib.scoring import normalize_task_context
return normalize_task_context(
inputs.get("task_context", {}),
prompt=str(inputs.get("prompt", "")),
capability=self.capability,
operation=str(inputs.get("operation", "text_to_video")),
)
@staticmethod
def _tool_context_payload(tool: BaseTool) -> dict[str, object]:
info = tool.get_info()
return {
"selected_tool_agent_skills": info.get("agent_skills", []),
"required_agent_skills": info.get("agent_skills", []),
"selected_tool_usage_location": info.get("usage_location"),
"selected_tool_best_for": info.get("best_for", []),
}
def _serialize_rankings(self, candidates: list[BaseTool], rankings: list[object]) -> list[dict[str, object]]:
tool_by_name = {tool.name: tool for tool in candidates}
serialized: list[dict[str, object]] = []
for score in rankings:
item = score.to_dict()
tool = tool_by_name.get(score.tool_name)
if tool:
info = tool.get_info()
item["agent_skills"] = info.get("agent_skills", [])
item["usage_location"] = info.get("usage_location")
item["best_for"] = info.get("best_for", [])
item["supports"] = info.get("supports", {})
item["status"] = str(tool.get_status())
serialized.append(item)
return serialized
def _filter_candidates(
self,
inputs: dict[str, object],
candidates: list[BaseTool],
) -> list[BaseTool]:
operation = inputs.get("operation", "text_to_video")
if operation == "rank":
return candidates
filtered: list[BaseTool] = []
for tool in candidates:
supports = getattr(tool, "supports", {})
props = getattr(tool, "input_schema", {}).get("properties", {})
if operation == "image_to_video":
if supports.get("image_to_video") or "image_url" in props or "reference_image_url" in props:
filtered.append(tool)
continue
if operation == "reference_to_video":
if supports.get("reference_to_video") or "reference_image_urls" in props:
filtered.append(tool)
continue
filtered.append(tool)
return filtered or candidates