Add Grok media providers and improve selector routing
This commit is contained in:
@@ -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"],
|
||||
)
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user