freedak/OpenMontage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Grok media providers and improve selector routing

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
tools/graphics/grok_image.py
+296
View File
@@ -0,0 +1,296 @@
"""xAI Grok image generation and editing."""
from __future__ import annotations
import base64
import mimetypes
import os
import time
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
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_image_input(url_value: str | None, path_value: str | None) -> dict[str, str] | None:
if url_value:
return {"url": url_value, "type": "image_url"}
if path_value:
return {"url": _file_to_data_uri(path_value), "type": "image_url"}
return None
class GrokImage(BaseTool):
name = "grok_image"
version = "0.1.0"
tier = ToolTier.GENERATE
capability = "image_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"]
capabilities = [
"generate_image",
"edit_image",
"text_to_image",
"image_to_image",
"style_transfer",
]
supports = {
"image_edit": True,
"multiple_outputs": True,
"aspect_ratio": True,
"resolution": True,
"reference_image": True,
"multiple_reference_images": True,
}
best_for = [
"single-image edits and style transfers",
"multi-image compositing into one generated frame",
"general-purpose image generation with aspect ratio control",
]
not_good_for = ["offline generation", "strict seeded reproducibility"]
input_schema = {
"type": "object",
"required": ["prompt"],
"properties": {
"prompt": {"type": "string"},
"generation_mode": {
"type": "string",
"enum": ["generate", "edit"],
"default": "generate",
"description": "Use 'edit' when providing one or more source images.",
},
"model": {
"type": "string",
"enum": ["grok-imagine-image"],
"default": "grok-imagine-image",
},
"aspect_ratio": {"type": "string", "description": "Examples: 1:1, 3:2, 16:9, 9:16"},
"resolution": {
"type": "string",
"enum": ["1k", "2k"],
"description": "xAI image output resolution tier",
},
"n": {
"type": "integer",
"minimum": 1,
"maximum": 10,
"default": 1,
},
"image_url": {"type": "string", "description": "Single source image URL for edit mode"},
"image_path": {"type": "string", "description": "Single local source image path for edit mode"},
"image_urls": {
"type": "array",
"items": {"type": "string"},
"description": "Multiple source image URLs for compositing edits",
},
"image_paths": {
"type": "array",
"items": {"type": "string"},
"description": "Multiple local source image paths for compositing edits",
},
"output_path": {"type": "string"},
},
}
resource_profile = ResourceProfile(
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
)
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
idempotency_key_fields = ["prompt", "generation_mode", "model", "aspect_ratio", "resolution", "n"]
side_effects = ["writes image file(s) to output_path", "calls xAI image API"]
user_visible_verification = ["Inspect generated image(s) for composition quality and edit fidelity"]
def get_status(self) -> ToolStatus:
if os.environ.get("XAI_API_KEY"):
return ToolStatus.AVAILABLE
return ToolStatus.UNAVAILABLE
@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("image_urls") or [])
count += len(inputs.get("image_paths") or [])
return count
def estimate_cost(self, inputs: dict[str, Any]) -> float:
output_count = int(inputs.get("n", 1))
input_count = self._input_image_count(inputs)
# xAI currently publishes Grok Imagine Image at $0.02 per generated
# image plus $0.002 per input image for edits or composites.
return output_count * 0.02 + input_count * 0.002
def _build_payload(self, inputs: dict[str, Any]) -> tuple[str, dict[str, Any]]:
mode = inputs.get("generation_mode", "generate")
payload: dict[str, Any] = {
"model": inputs.get("model", "grok-imagine-image"),
"prompt": inputs["prompt"],
}
if inputs.get("aspect_ratio"):
payload["aspect_ratio"] = inputs["aspect_ratio"]
if inputs.get("resolution"):
payload["resolution"] = inputs["resolution"]
if inputs.get("n"):
payload["n"] = inputs["n"]
primary_image = _normalize_image_input(inputs.get("image_url"), inputs.get("image_path"))
extra_images = [
{"url": url, "type": "image_url"}
for url in (inputs.get("image_urls") or [])
]
extra_images.extend(
{"url": _file_to_data_uri(path), "type": "image_url"}
for path in (inputs.get("image_paths") or [])
)
if primary_image or extra_images:
mode = "edit"
if mode == "edit":
endpoint = "https://api.x.ai/v1/images/edits"
if primary_image and not extra_images:
payload["image"] = primary_image
else:
images = []
if primary_image:
images.append(primary_image)
images.extend(extra_images)
if not images:
raise ValueError(
"Edit mode requires image_url/image_path or image_urls/image_paths"
)
payload["images"] = images
else:
endpoint = "https://api.x.ai/v1/images/generations"
return endpoint, payload
@staticmethod
def _infer_extension(url: str) -> str:
suffix = Path(urlparse(url).path).suffix.lower()
if suffix in {".png", ".jpg", ".jpeg", ".webp"}:
return suffix
return ".png"
@staticmethod
def _output_paths(output_path: str | None, count: int, extension: str) -> list[Path]:
if not output_path:
stem = "grok_image"
return [Path(f"{stem}_{idx + 1}{extension}") for idx in range(count)]
path = Path(output_path)
suffix = path.suffix or extension
if count == 1:
return [path if path.suffix else path.with_suffix(suffix)]
base = path.with_suffix("") if path.suffix else path
return [base.parent / f"{base.name}_{idx + 1}{suffix}" for idx in range(count)]
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
start = time.time()
try:
endpoint, payload = self._build_payload(inputs)
response = requests.post(
endpoint,
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
},
json=payload,
timeout=180,
)
response.raise_for_status()
data = response.json()
items = data.get("data", [])
if not items:
return ToolResult(success=False, error="xAI returned no image outputs")
extension = ".png"
first_url = items[0].get("url")
if first_url:
extension = self._infer_extension(first_url)
output_paths = self._output_paths(inputs.get("output_path"), len(items), extension)
artifacts: list[str] = []
outputs: list[str] = []
for item, output_path in zip(items, output_paths):
output_path.parent.mkdir(parents=True, exist_ok=True)
if item.get("b64_json"):
output_path.write_bytes(base64.b64decode(item["b64_json"]))
else:
image_url = item.get("url")
if not image_url:
return ToolResult(success=False, error="xAI image output missing url")
download = requests.get(image_url, timeout=120)
download.raise_for_status()
output_path.write_bytes(download.content)
artifacts.append(str(output_path))
outputs.append(str(output_path))
except Exception as e:
return ToolResult(success=False, error=f"Grok image generation failed: {e}")
primary_output = outputs[0]
return ToolResult(
success=True,
data={
"provider": "grok",
"model": payload["model"],
"prompt": inputs["prompt"],
"generation_mode": inputs.get("generation_mode", "generate"),
"output": primary_output,
"outputs": outputs,
"images_generated": len(outputs),
},
artifacts=artifacts,
cost_usd=self.estimate_cost(inputs),
duration_seconds=round(time.time() - start, 2),
model=payload["model"],
)
tools/graphics/image_selector.py
+106 -5
View File
@@ -52,6 +52,33 @@ class ImageSelector(BaseTool):
"width": {"type": "integer", "description": "Image width in pixels"},
"height": {"type": "integer", "description": "Image height in pixels"},
"seed": {"type": "integer", "description": "Random seed for reproducibility (generation providers only)"},
"n": {"type": "integer", "description": "Number of image variations to request when supported."},
"aspect_ratio": {
"type": "string",
"description": "Aspect ratio hint for providers that support ratio-based generation.",
},
"resolution": {
"type": "string",
"description": "Resolution tier for providers that support named resolutions.",
},
"generation_mode": {
"type": "string",
"enum": ["generate", "edit"],
"default": "generate",
"description": "Use 'edit' when providing one or more source images.",
},
"image_url": {"type": "string", "description": "Single source image URL for edit-capable providers."},
"image_path": {"type": "string", "description": "Single local source image path for edit-capable providers."},
"image_urls": {
"type": "array",
"items": {"type": "string"},
"description": "Multiple source image URLs for compositing edits.",
},
"image_paths": {
"type": "array",
"items": {"type": "string"},
"description": "Multiple local source image paths for compositing edits.",
},
"preferred_provider": {
"type": "string",
"description": "Provider name or 'auto'. Valid values are discovered at runtime from the registry.",
@@ -101,7 +128,7 @@ class ImageSelector(BaseTool):
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_cost(inputs) if tool else 0.0
def execute(self, inputs: dict[str, Any]) -> ToolResult:
@@ -109,8 +136,8 @@ class ImageSelector(BaseTool):
from lib.scoring import rank_providers
logger = logging.getLogger(__name__)
task_context = inputs.get("task_context", {})
candidates = self._providers()
task_context = self._prepare_task_context(inputs)
candidates = self._filter_candidates(inputs, self._providers())
# Rank mode — return scored provider rankings without generating
if inputs.get("operation") == "rank":
@@ -118,8 +145,9 @@ class ImageSelector(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,
},
)
@@ -143,7 +171,20 @@ class ImageSelector(BaseTool):
if hasattr(tool, 'input_schema'):
props = tool.input_schema.get("properties", {})
stripped = []
for passthrough_key in ("negative_prompt", "width", "height", "seed"):
for passthrough_key in (
"negative_prompt",
"width",
"height",
"seed",
"n",
"aspect_ratio",
"resolution",
"generation_mode",
"image_url",
"image_path",
"image_urls",
"image_paths",
):
if passthrough_key in adapted and passthrough_key not in props:
stripped.append(f"{passthrough_key}={adapted.pop(passthrough_key)}")
if stripped:
@@ -155,9 +196,11 @@ class ImageSelector(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"
@@ -177,6 +220,7 @@ class ImageSelector(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)
rankings = rank_providers(candidates, task_context)
@@ -195,3 +239,60 @@ class ImageSelector(BaseTool):
return tool_by_provider[score_item.provider], score_item
return None, None
def _prepare_task_context(self, inputs: dict[str, Any]) -> dict[str, Any]:
from lib.scoring import normalize_task_context
return normalize_task_context(
inputs.get("task_context", {}),
prompt=inputs.get("prompt", ""),
capability=self.capability,
operation=inputs.get("generation_mode", inputs.get("operation", "generate")),
)
@staticmethod
def _tool_context_payload(tool: BaseTool) -> dict[str, Any]:
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, Any]]:
tool_by_name = {tool.name: tool for tool in candidates}
serialized: list[dict[str, Any]] = []
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, Any], candidates: list[BaseTool]) -> list[BaseTool]:
wants_edit = (
inputs.get("generation_mode") == "edit"
or inputs.get("image_url")
or inputs.get("image_path")
or inputs.get("image_urls")
or inputs.get("image_paths")
)
if not wants_edit:
return candidates
filtered: list[BaseTool] = []
for tool in candidates:
props = getattr(tool, "input_schema", {}).get("properties", {})
supports = getattr(tool, "supports", {})
if supports.get("image_edit") or any(
key in props for key in ("image", "images", "image_url", "image_path", "image_urls", "image_paths")
):
filtered.append(tool)
return filtered or candidates