Initial release — OpenMontage: the first open-source agentic video production system
11 production pipelines, 47 tools, 124 agent skills. Supports cloud APIs (fal.ai, OpenAI, ElevenLabs, Suno, HeyGen, Runway) and free local providers (diffusers, Piper TTS, WAN 2.1, Hunyuan, CogVideo). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Graphics tools for image, diagram, and animation generation."""
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Code snippet renderer for overlay images.
|
||||
|
||||
Generates styled code screenshots using Pygments for syntax
|
||||
highlighting and Pillow for rendering. No external services required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Theme presets mapping to Pygments styles and background colors
|
||||
THEMES = {
|
||||
"monokai": {
|
||||
"pygments_style": "monokai",
|
||||
"bg_color": "#272822",
|
||||
"text_color": "#f8f8f2",
|
||||
"border_color": "#3e3d32",
|
||||
},
|
||||
"github_dark": {
|
||||
"pygments_style": "github-dark",
|
||||
"bg_color": "#0d1117",
|
||||
"text_color": "#c9d1d9",
|
||||
"border_color": "#30363d",
|
||||
},
|
||||
"dracula": {
|
||||
"pygments_style": "dracula",
|
||||
"bg_color": "#282a36",
|
||||
"text_color": "#f8f8f2",
|
||||
"border_color": "#44475a",
|
||||
},
|
||||
"one_dark": {
|
||||
"pygments_style": "one-dark",
|
||||
"bg_color": "#282c34",
|
||||
"text_color": "#abb2bf",
|
||||
"border_color": "#3e4452",
|
||||
},
|
||||
"solarized_dark": {
|
||||
"pygments_style": "solarized-dark",
|
||||
"bg_color": "#002b36",
|
||||
"text_color": "#839496",
|
||||
"border_color": "#073642",
|
||||
},
|
||||
"light": {
|
||||
"pygments_style": "default",
|
||||
"bg_color": "#ffffff",
|
||||
"text_color": "#333333",
|
||||
"border_color": "#e1e4e8",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class CodeSnippet(BaseTool):
|
||||
name = "code_snippet"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "graphics"
|
||||
provider = "pygments"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["python:pygments", "python:PIL"]
|
||||
install_instructions = "pip install Pygments Pillow"
|
||||
agent_skills = []
|
||||
|
||||
capabilities = [
|
||||
"render_code_image",
|
||||
"syntax_highlight",
|
||||
"themed_code_card",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["code"],
|
||||
"properties": {
|
||||
"code": {"type": "string"},
|
||||
"language": {"type": "string", "default": "python"},
|
||||
"theme": {
|
||||
"type": "string",
|
||||
"enum": list(THEMES.keys()),
|
||||
"default": "monokai",
|
||||
},
|
||||
"font_size": {"type": "integer", "default": 20},
|
||||
"padding": {"type": "integer", "default": 40},
|
||||
"border_radius": {"type": "integer", "default": 12},
|
||||
"line_numbers": {"type": "boolean", "default": True},
|
||||
"title": {"type": "string", "description": "Optional title bar text"},
|
||||
"output_path": {"type": "string"},
|
||||
"width": {"type": "integer", "description": "Force specific width"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50)
|
||||
idempotency_key_fields = ["code", "language", "theme", "font_size"]
|
||||
side_effects = ["writes image to output_path"]
|
||||
user_visible_verification = [
|
||||
"Verify code is readable and syntax highlighting is correct",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import pygments # noqa: F401
|
||||
from PIL import Image # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
try:
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from pygments import highlight
|
||||
from pygments.lexers import get_lexer_by_name, guess_lexer
|
||||
from pygments.formatters import ImageFormatter
|
||||
except ImportError:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pygments and Pillow required. Run: pip install Pygments Pillow",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
code = inputs["code"]
|
||||
language = inputs.get("language", "python")
|
||||
theme_name = inputs.get("theme", "monokai")
|
||||
font_size = inputs.get("font_size", 20)
|
||||
padding = inputs.get("padding", 40)
|
||||
line_numbers = inputs.get("line_numbers", True)
|
||||
title = inputs.get("title")
|
||||
output_path = Path(inputs.get("output_path", "code_snippet.png"))
|
||||
|
||||
theme = THEMES.get(theme_name, THEMES["monokai"])
|
||||
|
||||
try:
|
||||
lexer = get_lexer_by_name(language)
|
||||
except Exception:
|
||||
lexer = guess_lexer(code)
|
||||
|
||||
# Use Pygments ImageFormatter for rendering
|
||||
formatter = ImageFormatter(
|
||||
style=theme["pygments_style"],
|
||||
font_size=font_size,
|
||||
line_numbers=line_numbers,
|
||||
image_pad=padding,
|
||||
line_number_bg=theme["bg_color"],
|
||||
line_number_fg="#6272a4",
|
||||
)
|
||||
|
||||
# Render to bytes
|
||||
image_bytes = highlight(code, lexer, formatter)
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_bytes)
|
||||
|
||||
# Add title bar if requested
|
||||
if title:
|
||||
self._add_title_bar(output_path, title, theme, font_size)
|
||||
|
||||
elapsed = time.time() - start
|
||||
img = Image.open(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"output": str(output_path),
|
||||
"language": language,
|
||||
"theme": theme_name,
|
||||
"width": img.width,
|
||||
"height": img.height,
|
||||
"line_count": code.count("\n") + 1,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _add_title_bar(
|
||||
self, image_path: Path, title: str, theme: dict, font_size: int
|
||||
) -> None:
|
||||
"""Add a title bar to the top of the code image."""
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
img = Image.open(image_path)
|
||||
bar_height = font_size + 20
|
||||
|
||||
new_img = Image.new("RGB", (img.width, img.height + bar_height), theme["bg_color"])
|
||||
|
||||
# Draw title bar
|
||||
draw = ImageDraw.Draw(new_img)
|
||||
draw.rectangle(
|
||||
[(0, 0), (img.width, bar_height)],
|
||||
fill=theme["border_color"],
|
||||
)
|
||||
|
||||
# Draw window dots
|
||||
dot_y = bar_height // 2
|
||||
for i, color in enumerate(["#ff5f56", "#ffbd2e", "#27c93f"]):
|
||||
draw.ellipse(
|
||||
[(15 + i * 22, dot_y - 6), (15 + i * 22 + 12, dot_y + 6)],
|
||||
fill=color,
|
||||
)
|
||||
|
||||
# Draw title text
|
||||
try:
|
||||
font = ImageFont.truetype("arial.ttf", font_size - 4)
|
||||
except (IOError, OSError):
|
||||
font = ImageFont.load_default()
|
||||
|
||||
bbox = draw.textbbox((0, 0), title, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_x = (img.width - text_width) // 2
|
||||
draw.text((text_x, 8), title, fill=theme["text_color"], font=font)
|
||||
|
||||
# Paste original image below title bar
|
||||
new_img.paste(img, (0, bar_height))
|
||||
new_img.save(image_path)
|
||||
|
||||
@staticmethod
|
||||
def list_themes() -> dict[str, str]:
|
||||
return {name: f"Background: {t['bg_color']}" for name, t in THEMES.items()}
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Diagram generation tool using Mermaid CLI, Cairo/Pillow, or Graphviz.
|
||||
|
||||
Generates technical diagrams from text descriptions. Supports Mermaid
|
||||
syntax (flowcharts, sequence diagrams, etc.) and simple box/arrow
|
||||
diagrams via Pillow as fallback.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class DiagramGen(BaseTool):
|
||||
name = "diagram_gen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "graphics"
|
||||
provider = "mermaid"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = [] # checked dynamically
|
||||
install_instructions = (
|
||||
"For Mermaid diagrams:\n"
|
||||
" npm install -g @mermaid-js/mermaid-cli\n"
|
||||
"For Pillow-based diagrams (fallback):\n"
|
||||
" pip install Pillow"
|
||||
)
|
||||
agent_skills = ["beautiful-mermaid", "d3-viz"]
|
||||
|
||||
capabilities = [
|
||||
"generate_mermaid",
|
||||
"generate_flowchart",
|
||||
"generate_box_diagram",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["diagram_type"],
|
||||
"properties": {
|
||||
"diagram_type": {
|
||||
"type": "string",
|
||||
"enum": ["mermaid", "flowchart", "boxes"],
|
||||
},
|
||||
"definition": {
|
||||
"type": "string",
|
||||
"description": "Mermaid syntax or diagram description",
|
||||
},
|
||||
"boxes": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string"},
|
||||
"color": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"description": "Box definitions for box diagram type",
|
||||
},
|
||||
"connections": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": {"type": "integer"},
|
||||
"to": {"type": "integer"},
|
||||
"label": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"title": {"type": "string"},
|
||||
"theme": {
|
||||
"type": "string",
|
||||
"enum": ["dark", "light", "neutral"],
|
||||
"default": "dark",
|
||||
},
|
||||
"width": {"type": "integer", "default": 1200},
|
||||
"height": {"type": "integer", "default": 800},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50)
|
||||
idempotency_key_fields = ["diagram_type", "definition", "boxes"]
|
||||
side_effects = ["writes diagram image to output_path"]
|
||||
user_visible_verification = [
|
||||
"Verify diagram accurately represents the described structure",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._has_mermaid() or self._has_pillow():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def _has_mermaid(self) -> bool:
|
||||
return shutil.which("mmdc") is not None
|
||||
|
||||
def _has_pillow(self) -> bool:
|
||||
try:
|
||||
from PIL import Image # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
diagram_type = inputs["diagram_type"]
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if diagram_type == "mermaid":
|
||||
result = self._render_mermaid(inputs)
|
||||
elif diagram_type in ("flowchart", "boxes"):
|
||||
result = self._render_boxes(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown diagram type: {diagram_type}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Diagram generation failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _render_mermaid(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
definition = inputs.get("definition", "")
|
||||
if not definition:
|
||||
return ToolResult(success=False, error="Mermaid definition required")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "diagram.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
theme = inputs.get("theme", "dark")
|
||||
|
||||
if self._has_mermaid():
|
||||
# Write temp mermaid file
|
||||
temp_mmd = output_path.with_suffix(".mmd")
|
||||
temp_mmd.write_text(definition, encoding="utf-8")
|
||||
|
||||
mermaid_config = {"theme": theme}
|
||||
config_path = output_path.with_suffix(".mermaid.json")
|
||||
config_path.write_text(json.dumps(mermaid_config), encoding="utf-8")
|
||||
|
||||
cmd = [
|
||||
"mmdc",
|
||||
"-i", str(temp_mmd),
|
||||
"-o", str(output_path),
|
||||
"-c", str(config_path),
|
||||
"-b", "transparent",
|
||||
"-w", str(inputs.get("width", 1200)),
|
||||
]
|
||||
|
||||
try:
|
||||
self.run_command(cmd, timeout=30)
|
||||
finally:
|
||||
temp_mmd.unlink(missing_ok=True)
|
||||
config_path.unlink(missing_ok=True)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method": "mermaid-cli",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
else:
|
||||
# Fallback: render mermaid text as a styled text card
|
||||
return self._render_text_card(definition, inputs)
|
||||
|
||||
def _render_boxes(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Render a box-and-arrow diagram using Pillow."""
|
||||
if not self._has_pillow():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pillow required for box diagrams. Run: pip install Pillow",
|
||||
)
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
boxes = inputs.get("boxes", [])
|
||||
connections = inputs.get("connections", [])
|
||||
title = inputs.get("title", "")
|
||||
theme = inputs.get("theme", "dark")
|
||||
width = inputs.get("width", 1200)
|
||||
height = inputs.get("height", 800)
|
||||
output_path = Path(inputs.get("output_path", "diagram.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Theme colors
|
||||
if theme == "dark":
|
||||
bg, text_color, box_default, line_color = "#1e1e2e", "#cdd6f4", "#45475a", "#89b4fa"
|
||||
elif theme == "light":
|
||||
bg, text_color, box_default, line_color = "#ffffff", "#333333", "#e1e4e8", "#0366d6"
|
||||
else:
|
||||
bg, text_color, box_default, line_color = "#2d2d2d", "#d4d4d4", "#404040", "#569cd6"
|
||||
|
||||
img = Image.new("RGB", (width, height), bg)
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype("arial.ttf", 18)
|
||||
title_font = ImageFont.truetype("arial.ttf", 24)
|
||||
except (IOError, OSError):
|
||||
font = ImageFont.load_default()
|
||||
title_font = font
|
||||
|
||||
# Draw title
|
||||
y_offset = 20
|
||||
if title:
|
||||
bbox = draw.textbbox((0, 0), title, font=title_font)
|
||||
tw = bbox[2] - bbox[0]
|
||||
draw.text(((width - tw) // 2, y_offset), title, fill=text_color, font=title_font)
|
||||
y_offset += 50
|
||||
|
||||
# Layout boxes in a grid
|
||||
if not boxes:
|
||||
boxes = [{"label": "Empty"}]
|
||||
|
||||
cols = min(len(boxes), 4)
|
||||
rows = (len(boxes) + cols - 1) // cols
|
||||
box_w = min(200, (width - 80) // cols - 20)
|
||||
box_h = 60
|
||||
x_gap = (width - cols * box_w) // (cols + 1)
|
||||
y_gap = max(40, (height - y_offset - rows * box_h) // (rows + 1))
|
||||
|
||||
box_positions = []
|
||||
for i, box in enumerate(boxes):
|
||||
col = i % cols
|
||||
row = i // cols
|
||||
x = x_gap + col * (box_w + x_gap)
|
||||
y = y_offset + y_gap + row * (box_h + y_gap)
|
||||
|
||||
fill = box.get("color", box_default)
|
||||
draw.rounded_rectangle(
|
||||
[(x, y), (x + box_w, y + box_h)],
|
||||
radius=8,
|
||||
fill=fill,
|
||||
outline=line_color,
|
||||
width=2,
|
||||
)
|
||||
|
||||
label = box.get("label", f"Box {i}")
|
||||
bbox = draw.textbbox((0, 0), label, font=font)
|
||||
lw = bbox[2] - bbox[0]
|
||||
lh = bbox[3] - bbox[1]
|
||||
draw.text(
|
||||
(x + (box_w - lw) // 2, y + (box_h - lh) // 2),
|
||||
label, fill=text_color, font=font,
|
||||
)
|
||||
|
||||
box_positions.append((x, y, x + box_w, y + box_h))
|
||||
|
||||
# Draw connections
|
||||
for conn in connections:
|
||||
fi = conn.get("from", 0)
|
||||
ti = conn.get("to", 0)
|
||||
if fi >= len(box_positions) or ti >= len(box_positions):
|
||||
continue
|
||||
|
||||
fx1, fy1, fx2, fy2 = box_positions[fi]
|
||||
tx1, ty1, tx2, ty2 = box_positions[ti]
|
||||
|
||||
start_x = (fx1 + fx2) // 2
|
||||
start_y = fy2
|
||||
end_x = (tx1 + tx2) // 2
|
||||
end_y = ty1
|
||||
|
||||
draw.line([(start_x, start_y), (end_x, end_y)], fill=line_color, width=2)
|
||||
|
||||
# Arrow head
|
||||
arrow_size = 8
|
||||
draw.polygon(
|
||||
[(end_x, end_y), (end_x - arrow_size, end_y - arrow_size * 2), (end_x + arrow_size, end_y - arrow_size * 2)],
|
||||
fill=line_color,
|
||||
)
|
||||
|
||||
# Connection label
|
||||
conn_label = conn.get("label")
|
||||
if conn_label:
|
||||
mid_x = (start_x + end_x) // 2
|
||||
mid_y = (start_y + end_y) // 2
|
||||
draw.text((mid_x + 5, mid_y - 10), conn_label, fill=text_color, font=font)
|
||||
|
||||
img.save(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method": "pillow",
|
||||
"output": str(output_path),
|
||||
"box_count": len(boxes),
|
||||
"connection_count": len(connections),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _render_text_card(self, text: str, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Fallback: render text as a styled card image."""
|
||||
if not self._has_pillow():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pillow required. Run: pip install Pillow",
|
||||
)
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
output_path = Path(inputs.get("output_path", "diagram.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
width = inputs.get("width", 800)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype("consola.ttf", 16)
|
||||
except (IOError, OSError):
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# Calculate needed height
|
||||
lines = text.split("\n")
|
||||
line_height = 22
|
||||
height = max(200, len(lines) * line_height + 80)
|
||||
|
||||
img = Image.new("RGB", (width, height), "#1e1e2e")
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
y = 40
|
||||
for line in lines:
|
||||
draw.text((40, y), line, fill="#cdd6f4", font=font)
|
||||
y += line_height
|
||||
|
||||
img.save(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method": "text_card",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
@@ -0,0 +1,164 @@
|
||||
"""FLUX image generation via fal.ai API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class FluxImage(BaseTool):
|
||||
name = "flux_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "flux"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically via env var
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = ["flux-best-practices", "bfl-api"]
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"negative_prompt": True,
|
||||
"seed": True,
|
||||
"custom_size": True,
|
||||
}
|
||||
best_for = [
|
||||
"photorealistic images",
|
||||
"general-purpose image generation",
|
||||
"high quality at low cost (~$0.03/image)",
|
||||
]
|
||||
not_good_for = ["text rendering in images", "offline generation"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"negative_prompt": {"type": "string", "default": ""},
|
||||
"width": {"type": "integer", "default": 1024},
|
||||
"height": {"type": "integer", "default": 1024},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["flux-pro/v1.1", "flux/dev", "flux-pro"],
|
||||
"default": "flux-pro/v1.1",
|
||||
},
|
||||
"seed": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"guidance_scale": {"type": "number"},
|
||||
"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", "width", "height", "seed", "model"]
|
||||
side_effects = ["writes image file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "flux-pro/v1.1")
|
||||
if "pro" in model:
|
||||
return 0.05
|
||||
return 0.03 # dev tier
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No fal.ai API key found. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
model = inputs.get("model", "flux-pro/v1.1")
|
||||
prompt = inputs["prompt"]
|
||||
width = inputs.get("width", 1024)
|
||||
height = inputs.get("height", 1024)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"image_size": {"width": width, "height": height},
|
||||
}
|
||||
if inputs.get("seed") is not None:
|
||||
payload["seed"] = inputs["seed"]
|
||||
if inputs.get("num_inference_steps"):
|
||||
payload["num_inference_steps"] = inputs["num_inference_steps"]
|
||||
if inputs.get("guidance_scale"):
|
||||
payload["guidance_scale"] = inputs["guidance_scale"]
|
||||
if inputs.get("negative_prompt"):
|
||||
payload["negative_prompt"] = inputs["negative_prompt"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
image_url = data["images"][0]["url"]
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"FLUX generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "flux",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
"seed": data.get("seed"),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
seed=data.get("seed"),
|
||||
model=f"fal-ai/{model}",
|
||||
)
|
||||
@@ -0,0 +1,269 @@
|
||||
"""Image generation tool for diagrams, overlays, and illustrations.
|
||||
|
||||
.. deprecated::
|
||||
Use ``image_selector`` instead. This monolithic tool has been replaced by
|
||||
the selector/provider pattern: ``image_selector`` routes to per-provider
|
||||
tools (flux_image, openai_image, recraft_image, local_diffusion,
|
||||
pexels_image, pixabay_image). This file is kept for backwards
|
||||
compatibility and will be removed in a future release.
|
||||
|
||||
Supports cloud API providers (FLUX via fal.ai/Replicate, OpenAI DALL-E)
|
||||
and local Stable Diffusion via diffusers. Reports unavailable with
|
||||
install instructions when no provider is configured.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class ImageGen(BaseTool):
|
||||
name = "image_gen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "image_generation"
|
||||
provider = "multi"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.HYBRID # API (DALL-E/FLUX) or local (diffusers)
|
||||
|
||||
dependencies = [] # checked dynamically based on provider
|
||||
install_instructions = (
|
||||
"Set one of these environment variables:\n"
|
||||
" OPENAI_API_KEY — for DALL-E 3\n"
|
||||
" FAL_KEY — for FLUX via fal.ai\n"
|
||||
"Or install diffusers for local generation:\n"
|
||||
" pip install diffusers transformers accelerate torch"
|
||||
)
|
||||
agent_skills = ["flux-best-practices", "bfl-api"]
|
||||
|
||||
capabilities = [
|
||||
"generate_image",
|
||||
"generate_diagram_overlay",
|
||||
"generate_illustration",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"negative_prompt": {"type": "string", "default": ""},
|
||||
"width": {"type": "integer", "default": 1024},
|
||||
"height": {"type": "integer", "default": 1024},
|
||||
"provider": {
|
||||
"type": "string",
|
||||
"enum": ["openai", "flux", "local"],
|
||||
"description": "Auto-detected if not specified",
|
||||
},
|
||||
"model": {"type": "string"},
|
||||
"seed": {"type": "integer"},
|
||||
"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", "width", "height", "seed"]
|
||||
side_effects = ["writes image file to output_path", "calls external API"]
|
||||
user_visible_verification = [
|
||||
"Inspect generated image for relevance and quality",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
provider = self._detect_provider()
|
||||
if provider:
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def _detect_provider(self) -> Optional[str]:
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
return "openai"
|
||||
if os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY"):
|
||||
return "flux"
|
||||
try:
|
||||
import diffusers # noqa: F401
|
||||
return "local"
|
||||
except ImportError:
|
||||
pass
|
||||
return None
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
provider = inputs.get("provider") or self._detect_provider()
|
||||
if provider == "openai":
|
||||
return 0.04 # DALL-E 3 standard
|
||||
if provider == "flux":
|
||||
return 0.03
|
||||
return 0.0 # local
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
provider = inputs.get("provider") or self._detect_provider()
|
||||
if not provider:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No image generation provider available. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if provider == "openai":
|
||||
result = self._generate_openai(inputs)
|
||||
elif provider == "flux":
|
||||
result = self._generate_flux(inputs)
|
||||
elif provider == "local":
|
||||
result = self._generate_local(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown provider: {provider}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Generation failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
def _generate_openai(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
from openai import OpenAI
|
||||
import base64
|
||||
|
||||
client = OpenAI()
|
||||
prompt = inputs["prompt"]
|
||||
size = f"{inputs.get('width', 1024)}x{inputs.get('height', 1024)}"
|
||||
model = inputs.get("model", "dall-e-3")
|
||||
|
||||
response = client.images.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
n=1,
|
||||
response_format="b64_json",
|
||||
)
|
||||
|
||||
image_data = base64.b64decode(response.data[0].b64_json)
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_data)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "openai",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=model,
|
||||
)
|
||||
|
||||
def _generate_flux(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
import requests
|
||||
|
||||
api_key = os.environ.get("FAL_KEY") or os.environ["FAL_AI_API_KEY"]
|
||||
prompt = inputs["prompt"]
|
||||
width = inputs.get("width", 1024)
|
||||
height = inputs.get("height", 1024)
|
||||
seed = inputs.get("seed")
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"image_size": {"width": width, "height": height},
|
||||
}
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
|
||||
response = requests.post(
|
||||
"https://fal.run/fal-ai/flux/dev",
|
||||
headers={"Authorization": f"Key {api_key}", "Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
image_url = data["images"][0]["url"]
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "flux",
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
"seed": data.get("seed"),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
seed=data.get("seed"),
|
||||
model="flux-dev",
|
||||
)
|
||||
|
||||
def _generate_local(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
prompt = inputs["prompt"]
|
||||
negative = inputs.get("negative_prompt", "")
|
||||
width = inputs.get("width", 512)
|
||||
height = inputs.get("height", 512)
|
||||
seed = inputs.get("seed")
|
||||
model_id = inputs.get("model", "stabilityai/stable-diffusion-2-1-base")
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.float16 if device == "cuda" else torch.float32
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
|
||||
pipe = pipe.to(device)
|
||||
|
||||
generator = None
|
||||
if seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
image = pipe(
|
||||
prompt,
|
||||
negative_prompt=negative,
|
||||
width=width,
|
||||
height=height,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
image.save(str(output_path))
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "local",
|
||||
"model": model_id,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
seed=seed,
|
||||
model=model_id,
|
||||
)
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Capability-level image selector that routes between generation and stock providers.
|
||||
|
||||
Provider discovery is automatic — any BaseTool with capability="image_generation"
|
||||
is picked up from the registry. Adding a new image provider requires only creating
|
||||
the tool file in tools/graphics/; no changes to this selector are needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import BaseTool, ToolResult, ToolRuntime, ToolStability, ToolStatus, ToolTier
|
||||
|
||||
|
||||
class ImageSelector(BaseTool):
|
||||
name = "image_selector"
|
||||
version = "0.2.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "selector"
|
||||
stability = ToolStability.BETA
|
||||
runtime = ToolRuntime.HYBRID
|
||||
agent_skills = ["flux-best-practices", "bfl-api"]
|
||||
|
||||
capabilities = [
|
||||
"generate_image", "search_image", "download_image",
|
||||
"provider_selection", "text_to_image", "stock_image",
|
||||
]
|
||||
supports = {
|
||||
"user_preference_routing": True,
|
||||
"offline_fallback": True,
|
||||
"stock_fallback": True,
|
||||
}
|
||||
best_for = [
|
||||
"preflight routing — pick the best image provider for the task",
|
||||
"switching between generated and stock images",
|
||||
"automatic fallback when preferred provider is unavailable",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": "Image description (used as prompt for generation or query for stock)",
|
||||
},
|
||||
"negative_prompt": {
|
||||
"type": "string",
|
||||
"description": "What to avoid in the generated image. Passed to providers that support it.",
|
||||
},
|
||||
"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)"},
|
||||
"preferred_provider": {
|
||||
"type": "string",
|
||||
"description": "Provider name or 'auto'. Valid values are discovered at runtime from the registry.",
|
||||
"default": "auto",
|
||||
},
|
||||
"allowed_providers": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
def _providers(self) -> list[BaseTool]:
|
||||
"""Auto-discover image generation providers from the registry."""
|
||||
from tools.tool_registry import registry
|
||||
registry.ensure_discovered()
|
||||
return [t for t in registry.get_by_capability("image_generation")
|
||||
if t.name != self.name]
|
||||
|
||||
@property
|
||||
def fallback_tools(self) -> list[str]:
|
||||
"""Dynamically built from discovered providers."""
|
||||
return [t.name for t in self._providers()]
|
||||
|
||||
@property
|
||||
def provider_matrix(self) -> dict[str, dict[str, str]]:
|
||||
"""Built at runtime from each provider's best_for field."""
|
||||
matrix = {}
|
||||
for tool in self._providers():
|
||||
strength = ", ".join(tool.best_for) if tool.best_for else tool.name
|
||||
matrix[tool.provider] = {"tool": tool.name, "strength": strength}
|
||||
return matrix
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if any(tool.get_status() == ToolStatus.AVAILABLE for tool in self._providers()):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
tool = self._select_tool(inputs)
|
||||
return tool.estimate_cost(inputs) if tool else 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
tool = self._select_tool(inputs)
|
||||
if tool is None:
|
||||
return ToolResult(success=False, error="No image provider available.")
|
||||
|
||||
# Adapt input keys: stock tools use 'query' while generators use 'prompt'
|
||||
adapted = dict(inputs)
|
||||
if hasattr(tool, 'input_schema'):
|
||||
props = tool.input_schema.get("properties", {})
|
||||
if "query" in props and "query" not in adapted:
|
||||
adapted["query"] = adapted.get("prompt", "")
|
||||
|
||||
# Strip selector-only keys that downstream tools don't understand
|
||||
adapted.pop("preferred_provider", None)
|
||||
adapted.pop("allowed_providers", None)
|
||||
|
||||
# Pass through generation params only to tools that accept them
|
||||
if hasattr(tool, 'input_schema'):
|
||||
props = tool.input_schema.get("properties", {})
|
||||
for passthrough_key in ("negative_prompt", "width", "height", "seed"):
|
||||
if passthrough_key in adapted and passthrough_key not in props:
|
||||
adapted.pop(passthrough_key)
|
||||
|
||||
result = tool.execute(adapted)
|
||||
if result.success:
|
||||
result.data.setdefault("selected_tool", tool.name)
|
||||
return result
|
||||
|
||||
def _select_tool(self, inputs: dict[str, Any]) -> BaseTool | None:
|
||||
preferred = inputs.get("preferred_provider", "auto")
|
||||
allowed = set(inputs.get("allowed_providers") or [])
|
||||
candidates = self._providers()
|
||||
if allowed:
|
||||
candidates = [tool for tool in candidates if tool.provider in allowed]
|
||||
|
||||
if preferred != "auto":
|
||||
ordered = [tool for tool in candidates if tool.provider == preferred]
|
||||
ordered.extend([tool for tool in candidates if tool.provider != preferred])
|
||||
else:
|
||||
ordered = candidates
|
||||
|
||||
for tool in ordered:
|
||||
if tool.get_status() == ToolStatus.AVAILABLE:
|
||||
return tool
|
||||
return None
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Local Stable Diffusion image generation via diffusers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class LocalDiffusion(BaseTool):
|
||||
name = "local_diffusion"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "local_diffusion"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = [] # checked dynamically
|
||||
install_instructions = (
|
||||
"Install diffusers for local Stable Diffusion:\n"
|
||||
" pip install diffusers transformers accelerate torch"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"negative_prompt": True,
|
||||
"seed": True,
|
||||
"offline": True,
|
||||
"custom_size": True,
|
||||
}
|
||||
best_for = [
|
||||
"offline/air-gapped generation",
|
||||
"free image generation (no API cost)",
|
||||
"privacy-sensitive workflows",
|
||||
]
|
||||
not_good_for = [
|
||||
"CPU-only machines (very slow)",
|
||||
"highest quality output (API models are better)",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"negative_prompt": {"type": "string", "default": ""},
|
||||
"width": {"type": "integer", "default": 512},
|
||||
"height": {"type": "integer", "default": 512},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"default": "stabilityai/stable-diffusion-2-1-base",
|
||||
},
|
||||
"seed": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer", "default": 30},
|
||||
"guidance_scale": {"type": "number", "default": 7.5},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=8000, vram_mb=4000, disk_mb=5000, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1)
|
||||
idempotency_key_fields = ["prompt", "width", "height", "seed", "model"]
|
||||
side_effects = ["writes image file to output_path", "may download model weights on first run"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import diffusers # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return 30.0 # ~30s on a mid-range GPU
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="diffusers not installed. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
start = time.time()
|
||||
prompt = inputs["prompt"]
|
||||
negative = inputs.get("negative_prompt", "")
|
||||
width = inputs.get("width", 512)
|
||||
height = inputs.get("height", 512)
|
||||
seed = inputs.get("seed")
|
||||
model_id = inputs.get("model", "stabilityai/stable-diffusion-2-1-base")
|
||||
steps = inputs.get("num_inference_steps", 30)
|
||||
guidance = inputs.get("guidance_scale", 7.5)
|
||||
|
||||
try:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.float16 if device == "cuda" else torch.float32
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
|
||||
pipe = pipe.to(device)
|
||||
|
||||
generator = None
|
||||
if seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
image = pipe(
|
||||
prompt,
|
||||
negative_prompt=negative,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
image.save(str(output_path))
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Local diffusion generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "local_diffusion",
|
||||
"model": model_id,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
seed=seed,
|
||||
model=model_id,
|
||||
)
|
||||
@@ -0,0 +1,373 @@
|
||||
"""Mathematical animation tool via ManimCE.
|
||||
|
||||
Generates animated math/science/explainer videos from Python scene code
|
||||
using the Manim Community Edition engine. Free, local, no API key required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Quality presets mapping to Manim CLI flags
|
||||
QUALITY_PRESETS = {
|
||||
"low": {"flag": "-ql", "resolution": "854x480", "fps": 15},
|
||||
"medium": {"flag": "-qm", "resolution": "1280x720", "fps": 30},
|
||||
"high": {"flag": "-qh", "resolution": "1920x1080", "fps": 60},
|
||||
"4k": {"flag": "-qk", "resolution": "3840x2160", "fps": 60},
|
||||
"preview": {"flag": "-ql --format gif", "resolution": "854x480", "fps": 15},
|
||||
}
|
||||
|
||||
|
||||
class MathAnimate(BaseTool):
|
||||
name = "math_animate"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "graphics"
|
||||
provider = "manim"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL
|
||||
|
||||
dependencies = ["cmd:manim"]
|
||||
install_instructions = (
|
||||
"Install ManimCE:\n"
|
||||
" pip install manim\n"
|
||||
" manim checkhealth\n"
|
||||
"Requires: Python 3.8+, FFmpeg, LaTeX (optional, for math formulas)\n"
|
||||
" Windows: choco install miktex ffmpeg\n"
|
||||
" macOS: brew install mactex ffmpeg\n"
|
||||
" Linux: sudo apt install texlive-full ffmpeg"
|
||||
)
|
||||
agent_skills = ["manimce-best-practices", "manim-composer"]
|
||||
|
||||
capabilities = [
|
||||
"render_scene",
|
||||
"render_from_code",
|
||||
"render_from_template",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["scene_code"],
|
||||
"properties": {
|
||||
"scene_code": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Python code defining a Manim scene. Must contain a class "
|
||||
"inheriting from Scene with a construct() method. "
|
||||
"Import 'from manim import *' is auto-added if missing."
|
||||
),
|
||||
},
|
||||
"scene_name": {
|
||||
"type": "string",
|
||||
"description": "Name of the Scene class to render. Auto-detected if only one scene.",
|
||||
},
|
||||
"quality": {
|
||||
"type": "string",
|
||||
"enum": list(QUALITY_PRESETS.keys()),
|
||||
"default": "medium",
|
||||
"description": "Render quality preset",
|
||||
},
|
||||
"format": {
|
||||
"type": "string",
|
||||
"enum": ["mp4", "gif", "png", "webm"],
|
||||
"default": "mp4",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
"transparent": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Render with transparent background (PNG sequence or WebM)",
|
||||
},
|
||||
"background_color": {
|
||||
"type": "string",
|
||||
"description": "Background hex color (e.g., '#1a1a2e'). Default: Manim default (black).",
|
||||
},
|
||||
"extra_args": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Additional Manim CLI arguments",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=500, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["timeout"])
|
||||
idempotency_key_fields = ["scene_code", "scene_name", "quality"]
|
||||
side_effects = ["writes video/image file to output_path", "creates temp files"]
|
||||
user_visible_verification = [
|
||||
"Watch the animation for correctness and visual quality",
|
||||
"Verify math formulas render correctly (requires LaTeX)",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if shutil.which("manim"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # local, free
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
quality = inputs.get("quality", "medium")
|
||||
# Rough estimates based on scene complexity (assuming ~10s scene)
|
||||
estimates = {
|
||||
"low": 5.0,
|
||||
"medium": 15.0,
|
||||
"high": 45.0,
|
||||
"4k": 120.0,
|
||||
"preview": 3.0,
|
||||
}
|
||||
return estimates.get(quality, 15.0)
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if not shutil.which("manim"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Manim not found. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
result = self._render(inputs)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Manim render failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _render(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
scene_code = inputs["scene_code"]
|
||||
scene_name = inputs.get("scene_name")
|
||||
quality = inputs.get("quality", "medium")
|
||||
output_format = inputs.get("format", "mp4")
|
||||
output_path = inputs.get("output_path")
|
||||
transparent = inputs.get("transparent", False)
|
||||
bg_color = inputs.get("background_color")
|
||||
extra_args = inputs.get("extra_args", [])
|
||||
|
||||
# Ensure import statement
|
||||
if "from manim import" not in scene_code:
|
||||
scene_code = "from manim import *\n\n" + scene_code
|
||||
|
||||
# Auto-detect scene name if not provided
|
||||
if not scene_name:
|
||||
scene_name = self._detect_scene_name(scene_code)
|
||||
if not scene_name:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Could not detect Scene class name. Provide scene_name explicitly.",
|
||||
)
|
||||
|
||||
# Write scene code to temp file
|
||||
work_dir = Path(tempfile.mkdtemp(prefix="manim_"))
|
||||
scene_file = work_dir / "scene.py"
|
||||
scene_file.write_text(scene_code, encoding="utf-8")
|
||||
|
||||
# Build Manim CLI command
|
||||
cmd = ["manim"]
|
||||
|
||||
# Quality flag
|
||||
preset = QUALITY_PRESETS.get(quality, QUALITY_PRESETS["medium"])
|
||||
for flag_part in preset["flag"].split():
|
||||
cmd.append(flag_part)
|
||||
|
||||
# Format
|
||||
if output_format == "gif":
|
||||
cmd.append("--format")
|
||||
cmd.append("gif")
|
||||
elif output_format == "webm":
|
||||
cmd.append("--format")
|
||||
cmd.append("webm")
|
||||
elif output_format == "png":
|
||||
cmd.append("-s") # save last frame as PNG
|
||||
|
||||
# Transparent background
|
||||
if transparent:
|
||||
cmd.append("--transparent")
|
||||
|
||||
# Background color
|
||||
if bg_color:
|
||||
cmd.extend(["--background_color", bg_color])
|
||||
|
||||
# Disable window preview (headless rendering)
|
||||
cmd.append("--disable_caching")
|
||||
|
||||
# Extra args
|
||||
cmd.extend(extra_args)
|
||||
|
||||
# Scene file and class name
|
||||
cmd.append(str(scene_file))
|
||||
cmd.append(scene_name)
|
||||
|
||||
# Execute Manim
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300, # 5 min timeout
|
||||
cwd=str(work_dir),
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._cleanup(work_dir)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Manim render timed out after 300s. Try 'low' or 'preview' quality.",
|
||||
)
|
||||
|
||||
if proc.returncode != 0:
|
||||
error_msg = proc.stderr or proc.stdout or "Unknown error"
|
||||
# Extract the most useful part of the error
|
||||
lines = error_msg.strip().split("\n")
|
||||
# Look for the actual error (skip Manim header/progress)
|
||||
error_lines = [l for l in lines if "Error" in l or "error" in l or "Traceback" in l]
|
||||
if error_lines:
|
||||
error_msg = "\n".join(lines[lines.index(error_lines[0]):])
|
||||
self._cleanup(work_dir)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Manim render failed:\n{error_msg}",
|
||||
data={"full_stderr": proc.stderr, "full_stdout": proc.stdout},
|
||||
)
|
||||
|
||||
# Find the output file
|
||||
rendered_file = self._find_output(work_dir, scene_name, output_format)
|
||||
if not rendered_file:
|
||||
self._cleanup(work_dir)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Render succeeded but output file not found. Manim output:\n{proc.stdout}",
|
||||
)
|
||||
|
||||
# Move to desired output path
|
||||
if output_path:
|
||||
final_path = Path(output_path)
|
||||
else:
|
||||
ext = rendered_file.suffix
|
||||
final_path = Path(f"manim_{scene_name}{ext}")
|
||||
|
||||
final_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(str(rendered_file), str(final_path))
|
||||
|
||||
# Get video info
|
||||
video_info = self._probe_output(final_path)
|
||||
|
||||
# Cleanup temp directory
|
||||
self._cleanup(work_dir)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"scene_name": scene_name,
|
||||
"quality": quality,
|
||||
"format": output_format,
|
||||
"output": str(final_path),
|
||||
"resolution": preset["resolution"],
|
||||
"fps": preset["fps"],
|
||||
**video_info,
|
||||
},
|
||||
artifacts=[str(final_path)],
|
||||
)
|
||||
|
||||
def _detect_scene_name(self, code: str) -> Optional[str]:
|
||||
"""Extract Scene subclass name from code."""
|
||||
import re
|
||||
|
||||
# Match class definitions that inherit from Scene or its variants
|
||||
pattern = r"class\s+(\w+)\s*\(\s*(?:Scene|ThreeDScene|MovingCameraScene|ZoomedScene)\s*\)"
|
||||
matches = re.findall(pattern, code)
|
||||
if len(matches) == 1:
|
||||
return matches[0]
|
||||
if len(matches) > 1:
|
||||
# Return the last one (convention: main scene is last)
|
||||
return matches[-1]
|
||||
return None
|
||||
|
||||
def _find_output(self, work_dir: Path, scene_name: str, fmt: str) -> Optional[Path]:
|
||||
"""Find Manim's output file in the media directory."""
|
||||
media_dir = work_dir / "media"
|
||||
if not media_dir.exists():
|
||||
return None
|
||||
|
||||
# Manim outputs to media/videos/<scene_file>/<quality>/<SceneName>.<ext>
|
||||
# or media/images/<scene_file>/<SceneName>.<ext> for -s flag
|
||||
ext_map = {"mp4": ".mp4", "gif": ".gif", "webm": ".webm", "png": ".png"}
|
||||
target_ext = ext_map.get(fmt, ".mp4")
|
||||
|
||||
# Search recursively for the output file
|
||||
for path in media_dir.rglob(f"{scene_name}{target_ext}"):
|
||||
return path
|
||||
|
||||
# Fallback: any file with the right extension
|
||||
for path in media_dir.rglob(f"*{target_ext}"):
|
||||
return path
|
||||
|
||||
return None
|
||||
|
||||
def _probe_output(self, path: Path) -> dict[str, Any]:
|
||||
"""Get basic info about the rendered file."""
|
||||
info: dict[str, Any] = {"file_size_bytes": path.stat().st_size}
|
||||
|
||||
if not shutil.which("ffprobe"):
|
||||
return info
|
||||
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format", "-show_streams",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
if proc.returncode == 0:
|
||||
import json
|
||||
probe = json.loads(proc.stdout)
|
||||
fmt = probe.get("format", {})
|
||||
info["duration_seconds"] = float(fmt.get("duration", 0))
|
||||
info["file_size_mb"] = round(path.stat().st_size / (1024 * 1024), 2)
|
||||
for stream in probe.get("streams", []):
|
||||
if stream.get("codec_type") == "video":
|
||||
info["video_width"] = int(stream.get("width", 0))
|
||||
info["video_height"] = int(stream.get("height", 0))
|
||||
info["video_codec"] = stream.get("codec_name", "")
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return info
|
||||
|
||||
def _cleanup(self, work_dir: Path) -> None:
|
||||
"""Remove temp working directory."""
|
||||
try:
|
||||
shutil.rmtree(str(work_dir), ignore_errors=True)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,176 @@
|
||||
"""OpenAI GPT Image generation (gpt-image-1 / DALL-E 3)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class OpenAIImage(BaseTool):
|
||||
name = "openai_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "openai"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically
|
||||
install_instructions = (
|
||||
"Set OPENAI_API_KEY to your OpenAI API key.\n"
|
||||
" pip install openai"
|
||||
)
|
||||
agent_skills = ["flux-best-practices"] # general image gen knowledge
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"complex_instructions": True,
|
||||
"text_in_image": True,
|
||||
"multiple_outputs": True,
|
||||
}
|
||||
best_for = [
|
||||
"complex multi-element compositions",
|
||||
"images with text/labels",
|
||||
"following detailed instructions accurately",
|
||||
]
|
||||
not_good_for = ["offline generation", "budget-constrained projects at high quality"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["gpt-image-1", "dall-e-3"],
|
||||
"default": "gpt-image-1",
|
||||
},
|
||||
"size": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"1024x1024", "1536x1024", "1024x1536", "auto",
|
||||
"1024x1792", "1792x1024", # dall-e-3 only
|
||||
],
|
||||
"default": "1024x1024",
|
||||
},
|
||||
"quality": {
|
||||
"type": "string",
|
||||
"enum": ["low", "medium", "high", "auto", "standard", "hd"],
|
||||
"default": "high",
|
||||
},
|
||||
"output_format": {
|
||||
"type": "string",
|
||||
"enum": ["png", "jpeg", "webp"],
|
||||
"default": "png",
|
||||
},
|
||||
"n": {"type": "integer", "default": 1, "minimum": 1, "maximum": 4},
|
||||
"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", "size", "quality", "model"]
|
||||
side_effects = ["writes image file to output_path", "calls OpenAI API"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "gpt-image-1")
|
||||
quality = inputs.get("quality", "high")
|
||||
n = inputs.get("n", 1)
|
||||
if model == "gpt-image-1":
|
||||
cost_map = {"low": 0.011, "medium": 0.042, "high": 0.167, "auto": 0.042}
|
||||
return cost_map.get(quality, 0.042) * n
|
||||
# dall-e-3 fallback pricing
|
||||
quality_map = {"standard": 0.04, "hd": 0.08}
|
||||
return quality_map.get(quality, 0.04) * n
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="OPENAI_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
start = time.time()
|
||||
client = OpenAI()
|
||||
model = inputs.get("model", "gpt-image-1")
|
||||
prompt = inputs["prompt"]
|
||||
size = inputs.get("size", "1024x1024")
|
||||
n = inputs.get("n", 1)
|
||||
|
||||
try:
|
||||
if model == "gpt-image-1":
|
||||
quality = inputs.get("quality", "high")
|
||||
output_format = inputs.get("output_format", "png")
|
||||
response = client.images.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
quality=quality,
|
||||
output_format=output_format,
|
||||
n=n,
|
||||
)
|
||||
else:
|
||||
# dall-e-3 path
|
||||
quality = inputs.get("quality", "standard")
|
||||
if quality in ("low", "medium", "high", "auto"):
|
||||
quality = "standard" # map to dall-e-3 quality options
|
||||
response = client.images.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
quality=quality,
|
||||
n=1, # dall-e-3 only supports n=1
|
||||
response_format="b64_json",
|
||||
)
|
||||
|
||||
image_data = base64.b64decode(response.data[0].b64_json)
|
||||
ext = inputs.get("output_format", "png")
|
||||
output_path = Path(inputs.get("output_path", f"generated_image.{ext}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_data)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"OpenAI image generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "openai",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=model,
|
||||
)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Stock image acquisition from Pexels API (free)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class PexelsImage(BaseTool):
|
||||
name = "pexels_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.SOURCE
|
||||
capability = "image_generation"
|
||||
provider = "pexels"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set PEXELS_API_KEY to your Pexels API key.\n"
|
||||
" Get one free at https://www.pexels.com/api/"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["search_image", "download_image", "stock_image"]
|
||||
supports = {
|
||||
"orientation_filter": True,
|
||||
"size_filter": True,
|
||||
"color_filter": True,
|
||||
"locale": True,
|
||||
"free_commercial_use": True,
|
||||
}
|
||||
best_for = [
|
||||
"real-world photography (cities, nature, people, objects)",
|
||||
"establishing shots and B-roll stills",
|
||||
"free stock images — no cost, no attribution required",
|
||||
]
|
||||
not_good_for = [
|
||||
"custom/specific compositions",
|
||||
"abstract or stylized graphics",
|
||||
"offline use",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search term"},
|
||||
"orientation": {
|
||||
"type": "string",
|
||||
"enum": ["landscape", "portrait", "square"],
|
||||
},
|
||||
"size": {
|
||||
"type": "string",
|
||||
"enum": ["large", "medium", "small"],
|
||||
"description": "large=24MP+, medium=12MP+, small=4MP+",
|
||||
},
|
||||
"color": {
|
||||
"type": "string",
|
||||
"description": "Hex without # (e.g. FF0000) or color name (red, blue, etc.)",
|
||||
},
|
||||
"per_page": {"type": "integer", "default": 5, "minimum": 1, "maximum": 80},
|
||||
"page": {"type": "integer", "default": 1},
|
||||
"download_size": {
|
||||
"type": "string",
|
||||
"enum": ["original", "large2x", "large", "medium"],
|
||||
"default": "large2x",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["query", "orientation", "size", "color", "page"]
|
||||
side_effects = ["writes image file to output_path", "calls Pexels API"]
|
||||
user_visible_verification = ["Check that downloaded image matches the intended scene"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("PEXELS_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # Pexels is free
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("PEXELS_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="PEXELS_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
query = inputs["query"]
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"query": query,
|
||||
"per_page": inputs.get("per_page", 5),
|
||||
"page": inputs.get("page", 1),
|
||||
}
|
||||
if inputs.get("orientation"):
|
||||
params["orientation"] = inputs["orientation"]
|
||||
if inputs.get("size"):
|
||||
params["size"] = inputs["size"]
|
||||
if inputs.get("color"):
|
||||
params["color"] = inputs["color"]
|
||||
|
||||
try:
|
||||
search_response = requests.get(
|
||||
"https://api.pexels.com/v1/search",
|
||||
headers={"Authorization": api_key},
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
search_response.raise_for_status()
|
||||
data = search_response.json()
|
||||
|
||||
photos = data.get("photos", [])
|
||||
if not photos:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No images found for query: {query}",
|
||||
data={"total_results": data.get("total_results", 0)},
|
||||
)
|
||||
|
||||
# Pick the first result (agent can refine query if needed)
|
||||
photo = photos[0]
|
||||
download_size = inputs.get("download_size", "large2x")
|
||||
image_url = photo["src"].get(download_size, photo["src"]["large2x"])
|
||||
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"pexels_{photo['id']}.jpg"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Pexels image search failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "pexels",
|
||||
"photo_id": photo["id"],
|
||||
"photographer": photo.get("photographer", "Unknown"),
|
||||
"photographer_url": photo.get("photographer_url", ""),
|
||||
"alt": photo.get("alt", ""),
|
||||
"width": photo.get("width"),
|
||||
"height": photo.get("height"),
|
||||
"query": query,
|
||||
"output": str(output_path),
|
||||
"total_results": data.get("total_results", 0),
|
||||
"results_returned": len(photos),
|
||||
"license": "Pexels License (free, no attribution required)",
|
||||
"pexels_url": photo.get("url", ""),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
)
|
||||
@@ -0,0 +1,196 @@
|
||||
"""Stock image acquisition from Pixabay API (free)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class PixabayImage(BaseTool):
|
||||
name = "pixabay_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.SOURCE
|
||||
capability = "image_generation"
|
||||
provider = "pixabay"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set PIXABAY_API_KEY to your Pixabay API key.\n"
|
||||
" Get one free at https://pixabay.com/api/docs/"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["search_image", "download_image", "stock_image"]
|
||||
supports = {
|
||||
"orientation_filter": True,
|
||||
"category_filter": True,
|
||||
"color_filter": True,
|
||||
"image_type_filter": True,
|
||||
"editors_choice": True,
|
||||
"free_commercial_use": True,
|
||||
}
|
||||
best_for = [
|
||||
"large royalty-free library (5M+ images)",
|
||||
"category-based filtering (nature, business, science, etc.)",
|
||||
"free stock images — no cost, no attribution required",
|
||||
]
|
||||
not_good_for = [
|
||||
"full-resolution originals (standard API limited to 1280px)",
|
||||
"custom compositions",
|
||||
"offline use",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search term (max 100 chars)"},
|
||||
"image_type": {
|
||||
"type": "string",
|
||||
"enum": ["all", "photo", "illustration", "vector"],
|
||||
"default": "all",
|
||||
},
|
||||
"orientation": {
|
||||
"type": "string",
|
||||
"enum": ["all", "horizontal", "vertical"],
|
||||
"default": "all",
|
||||
},
|
||||
"category": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"backgrounds", "fashion", "nature", "science", "education",
|
||||
"feelings", "health", "people", "religion", "places",
|
||||
"animals", "industry", "computer", "food", "sports",
|
||||
"transportation", "travel", "buildings", "business", "music",
|
||||
],
|
||||
},
|
||||
"colors": {
|
||||
"type": "string",
|
||||
"description": "Comma-separated: grayscale, transparent, red, orange, yellow, green, turquoise, blue, lilac, pink, white, gray, black, brown",
|
||||
},
|
||||
"editors_choice": {"type": "boolean", "default": False},
|
||||
"safesearch": {"type": "boolean", "default": True},
|
||||
"per_page": {"type": "integer", "default": 5, "minimum": 3, "maximum": 200},
|
||||
"page": {"type": "integer", "default": 1},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["query", "image_type", "orientation", "category", "page"]
|
||||
side_effects = ["writes image file to output_path", "calls Pixabay API"]
|
||||
user_visible_verification = ["Check that downloaded image matches the intended scene"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("PIXABAY_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # Pixabay is free
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("PIXABAY_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="PIXABAY_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
query = inputs["query"]
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"key": api_key,
|
||||
"q": query,
|
||||
"per_page": inputs.get("per_page", 5),
|
||||
"page": inputs.get("page", 1),
|
||||
"safesearch": str(inputs.get("safesearch", True)).lower(),
|
||||
}
|
||||
if inputs.get("image_type") and inputs["image_type"] != "all":
|
||||
params["image_type"] = inputs["image_type"]
|
||||
if inputs.get("orientation") and inputs["orientation"] != "all":
|
||||
params["orientation"] = inputs["orientation"]
|
||||
if inputs.get("category"):
|
||||
params["category"] = inputs["category"]
|
||||
if inputs.get("colors"):
|
||||
params["colors"] = inputs["colors"]
|
||||
if inputs.get("editors_choice"):
|
||||
params["editors_choice"] = "true"
|
||||
|
||||
try:
|
||||
search_response = requests.get(
|
||||
"https://pixabay.com/api/",
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
search_response.raise_for_status()
|
||||
data = search_response.json()
|
||||
|
||||
hits = data.get("hits", [])
|
||||
if not hits:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No images found for query: {query}",
|
||||
data={"total_results": data.get("total", 0)},
|
||||
)
|
||||
|
||||
hit = hits[0]
|
||||
# largeImageURL is the best available at standard API tier (1280px)
|
||||
image_url = hit.get("largeImageURL", hit.get("webformatURL"))
|
||||
|
||||
# Download immediately — Pixabay URLs contain embedded tokens that expire
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"pixabay_{hit['id']}.jpg"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Pixabay image search failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "pixabay",
|
||||
"image_id": hit["id"],
|
||||
"user": hit.get("user", "Unknown"),
|
||||
"tags": hit.get("tags", ""),
|
||||
"image_width": hit.get("imageWidth"),
|
||||
"image_height": hit.get("imageHeight"),
|
||||
"query": query,
|
||||
"output": str(output_path),
|
||||
"total_results": data.get("total", 0),
|
||||
"results_returned": len(hits),
|
||||
"license": "Pixabay Content License (free, no attribution required)",
|
||||
"page_url": hit.get("pageURL", ""),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Recraft V4 image generation via fal.ai API.
|
||||
|
||||
Best for logos, brand assets, SVG vectors, and images with accurate text rendering.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class RecraftImage(BaseTool):
|
||||
name = "recraft_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "recraft"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = [
|
||||
"generate_image",
|
||||
"generate_logo",
|
||||
"generate_vector",
|
||||
"text_to_image",
|
||||
]
|
||||
supports = {
|
||||
"svg_output": True,
|
||||
"text_rendering": True,
|
||||
"color_palette": True,
|
||||
"custom_size": True,
|
||||
}
|
||||
best_for = [
|
||||
"logos and brand assets",
|
||||
"SVG vector output",
|
||||
"images with accurate text rendering",
|
||||
"clean professional graphics",
|
||||
]
|
||||
not_good_for = ["photorealistic images", "offline generation"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["v4", "v4-pro"],
|
||||
"default": "v4",
|
||||
},
|
||||
"image_size": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"square", "square_hd",
|
||||
"landscape_4_3", "landscape_16_9",
|
||||
"portrait_4_3", "portrait_16_9",
|
||||
],
|
||||
"default": "square_hd",
|
||||
},
|
||||
"style": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"any", "realistic_image", "digital_illustration",
|
||||
"vector_illustration", "icon",
|
||||
],
|
||||
"default": "any",
|
||||
},
|
||||
"colors": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Color palette as hex strings, e.g. ['#FF5733', '#2E86C1']",
|
||||
},
|
||||
"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", "model", "style", "image_size"]
|
||||
side_effects = ["writes image file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = ["Inspect generated image for brand accuracy and text readability"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "v4")
|
||||
if model == "v4-pro":
|
||||
return 0.25
|
||||
return 0.04
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="FAL_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
model = inputs.get("model", "v4")
|
||||
prompt = inputs["prompt"]
|
||||
|
||||
model_path = f"recraft/{model}/text-to-image"
|
||||
if model == "v4-pro":
|
||||
model_path = "recraft/v4/pro/text-to-image"
|
||||
elif model == "v4":
|
||||
model_path = "recraft/v4/text-to-image"
|
||||
|
||||
payload: dict[str, Any] = {"prompt": prompt}
|
||||
if inputs.get("image_size"):
|
||||
payload["image_size"] = inputs["image_size"]
|
||||
if inputs.get("style"):
|
||||
payload["style"] = inputs["style"]
|
||||
if inputs.get("colors"):
|
||||
payload["colors"] = inputs["colors"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model_path}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
image_url = data["images"][0]["url"]
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
ext = "svg" if inputs.get("style") == "vector_illustration" else "png"
|
||||
output_path = Path(inputs.get("output_path", f"generated_image.{ext}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Recraft generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "recraft",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=f"fal-ai/{model_path}",
|
||||
)
|
||||
Reference in New Issue
Block a user