from __future__ import annotations import mimetypes from pathlib import Path from typing import Any from agent.prompts import normalize_sdk_package from agent.runtime_limits import BatchRuntimeLimits from agent.tools.apply_patch import ApplyPatchFreeformTool, ApplyPatchJsonTool from agent.tools.base import ( BaseDeclarativeTool, BaseToolInvocation, ToolResult, ToolSchema, make_tool_schema, ) from agent.tools.compile_model import CompileModelTool from agent.tools.edit_code import ReplaceTool from agent.tools.find_examples import FindExamplesTool from agent.tools.probe_model import ProbeModelTool from agent.tools.read_file import ReadFileTool from agent.tools.registry import ToolRegistry from agent.tools.write_code import WriteFileTool from articraft.values import ProviderName, normalize_provider_name SUPPORTED_IMAGE_MIME_TYPES_BY_PROVIDER: dict[str, set[str]] = { ProviderName.ANTHROPIC.value: { "image/png", "image/jpeg", "image/webp", "image/gif", }, ProviderName.DASHSCOPE.value: { "image/png", "image/jpeg", "image/webp", "image/gif", }, ProviderName.OPENAI.value: { "image/png", "image/jpeg", "image/webp", "image/gif", }, ProviderName.CODEX_CLI.value: { "image/png", "image/jpeg", "image/webp", "image/gif", }, ProviderName.GEMINI.value: { "image/png", "image/jpeg", "image/webp", "image/heic", "image/heif", }, ProviderName.OPENROUTER.value: { "image/png", "image/jpeg", "image/webp", "image/gif", }, } _FIRST_TURN_RUNTIME_GUIDANCE_SHARED = ( "\n" "- Read the current `model.py` before editing.\n" "- Start with a realism-first structure plan. Use one coherent scaffold when the real object needs layered bodies, hollow forms, mechanisms, or repeated features; otherwise make small focused edits.\n" "- Treat visual realism as part of the deliverable: make the object read clearly as the requested thing, with believable proportions, silhouette, colors/materials, and major visible surface treatment.\n" "- Run `compile_model` to check your latest revision.\n" "- If compile is clean and the model already satisfies the realism/mechanism brief, conclude.\n" "" ) def build_tool_registry( provider: str, *, sdk_package: str = "sdk", runtime_limits: BatchRuntimeLimits | None = None, ) -> ToolRegistry: provider_norm = normalize_provider_name(provider) package = normalize_sdk_package(sdk_package) if provider_norm is ProviderName.OPENAI: tools: list[BaseDeclarativeTool] = [ ReadFileTool(), ApplyPatchFreeformTool(), CompileModelTool(), ProbeModelTool(sdk_package=package, runtime_limits=runtime_limits), ] elif provider_norm is ProviderName.CODEX_CLI: tools = [ ReadFileTool(), ApplyPatchJsonTool(), ReplaceTool(), WriteFileTool(), CompileModelTool(), ProbeModelTool(sdk_package=package, runtime_limits=runtime_limits), ] else: tools = [ ReadFileTool(), ReplaceTool(), WriteFileTool(), CompileModelTool(), ProbeModelTool(sdk_package=package, runtime_limits=runtime_limits), ] tools.append( FindExamplesTool(sdk_package=package, include_paths=provider_norm is ProviderName.OPENAI) ) return ToolRegistry(tools) def build_first_turn_runtime_guidance(_provider: str) -> str: return _FIRST_TURN_RUNTIME_GUIDANCE_SHARED def prepend_runtime_guidance( user_content: Any, *, runtime_guidance_text: str | None = None, ) -> Any: guidance = (runtime_guidance_text or "").strip() if not guidance: return user_content if isinstance(user_content, str): if not user_content.strip(): return guidance return f"{guidance}\n\n{user_content}" if not isinstance(user_content, list): return user_content return [{"type": "input_text", "text": guidance}, *user_content] def build_first_turn_messages( user_content: Any, *, sdk_docs_context: str, provider: str, ) -> list[dict[str, Any]]: messages: list[dict[str, Any]] = [] if sdk_docs_context: messages.append({"role": "user", "content": sdk_docs_context}) messages.append( { "role": "user", "content": prepend_runtime_guidance( user_content, runtime_guidance_text=build_first_turn_runtime_guidance(provider), ), } ) return messages def build_initial_user_content( text_prompt: str, *, image_path: Path | None = None, image_detail: str = "high", runtime_guidance_text: str | None = None, ) -> Any: content: Any if image_path is None: content = text_prompt else: content = [ {"type": "input_text", "text": text_prompt}, { "type": "input_image", "image_path": str(image_path), "detail": image_detail, }, ] return prepend_runtime_guidance(content, runtime_guidance_text=runtime_guidance_text) def resolve_image_path( image_arg: str | None, *, provider: str | None = None, ) -> Path | None: if not image_arg: return None path = Path(image_arg).expanduser().resolve() if not path.exists(): raise FileNotFoundError(f"Image file not found: {path}") if not path.is_file(): raise ValueError(f"Image path is not a file: {path}") mime_type, _ = mimetypes.guess_type(path.name) provider_norm = normalize_provider_name(provider) supported_mime_types = SUPPORTED_IMAGE_MIME_TYPES_BY_PROVIDER.get(provider_norm.value) if supported_mime_types is None: raise ValueError(f"Unsupported provider for image validation: {provider}") if mime_type not in supported_mime_types: raise ValueError( f"Unsupported image type for {provider_norm.value}: {path.name} ({mime_type or 'unknown'})" ) size_bytes = path.stat().st_size if provider_norm is ProviderName.GEMINI: if size_bytes >= 20 * 1024 * 1024: raise ValueError( f"Image file exceeds Gemini inline request limit: {path} " "(must stay under 20 MB including prompt text)" ) elif size_bytes > 50 * 1024 * 1024: raise ValueError(f"Image file exceeds 50 MB request limit: {path}") return path __all__ = [ "BaseDeclarativeTool", "BaseToolInvocation", "ToolResult", "ToolSchema", "make_tool_schema", "ApplyPatchFreeformTool", "ApplyPatchJsonTool", "CompileModelTool", "FindExamplesTool", "ProbeModelTool", "ReadFileTool", "ReplaceTool", "WriteFileTool", "ToolRegistry", "SUPPORTED_IMAGE_MIME_TYPES_BY_PROVIDER", "build_tool_registry", "build_first_turn_runtime_guidance", "prepend_runtime_guidance", "build_first_turn_messages", "build_initial_user_content", "resolve_image_path", ]