diff --git a/EvoScientist/EvoScientist.py b/EvoScientist/EvoScientist.py index 643e173..dabb7f0 100644 --- a/EvoScientist/EvoScientist.py +++ b/EvoScientist/EvoScientist.py @@ -430,7 +430,14 @@ def _build_base_kwargs( base_backend, base_middleware, *, cfg=None, chat_model=None, workspace_dir=None ): """Build agent kwargs *without* MCP (fast, no subprocess spawning).""" - from .tools import skill_manager, tavily_search, think_tool + from .tools import ( + edit_image, + generate_image, + refresh_image_tool_descriptions, + skill_manager, + tavily_search, + think_tool, + ) from .utils import load_subagents cfg = cfg if cfg is not None else _ensure_config() @@ -438,7 +445,8 @@ def _build_base_kwargs( tool_registry = {"think_tool": think_tool} if os.environ.get("TAVILY_API_KEY"): tool_registry["tavily_search"] = tavily_search - base_tools = [think_tool] + refresh_image_tool_descriptions() + base_tools = [think_tool, generate_image, edit_image] if cfg.workspace_isolation != "required": base_tools.append(skill_manager) @@ -485,7 +493,14 @@ def load_mcp_and_build_kwargs( chat_model: Explicit chat model to bind instead of ``_ensure_chat_model()`` (which would write module globals). """ - from .tools import skill_manager, tavily_search, think_tool + from .tools import ( + edit_image, + generate_image, + refresh_image_tool_descriptions, + skill_manager, + tavily_search, + think_tool, + ) from .utils import load_subagents cfg = cfg if cfg is not None else _ensure_config() @@ -503,7 +518,8 @@ def load_mcp_and_build_kwargs( tool_registry = {"think_tool": think_tool} if os.environ.get("TAVILY_API_KEY"): tool_registry["tavily_search"] = tavily_search - base_tools = [think_tool] + refresh_image_tool_descriptions() + base_tools = [think_tool, generate_image, edit_image] if cfg.workspace_isolation != "required": base_tools.append(skill_manager) diff --git a/EvoScientist/tools/__init__.py b/EvoScientist/tools/__init__.py index 5fe21d4..ac697ed 100644 --- a/EvoScientist/tools/__init__.py +++ b/EvoScientist/tools/__init__.py @@ -4,12 +4,16 @@ External imports like ``from EvoScientist.tools import tavily_search`` continue to work unchanged thanks to these re-exports. """ +from .image import edit_image, generate_image, refresh_image_tool_descriptions from .search import fetch_webpage_content, tavily_search from .skill_manager import skill_manager from .think import think_tool __all__ = [ + "edit_image", "fetch_webpage_content", + "generate_image", + "refresh_image_tool_descriptions", "skill_manager", "tavily_search", "think_tool", diff --git a/EvoScientist/tools/image.py b/EvoScientist/tools/image.py new file mode 100644 index 0000000..1ee889b --- /dev/null +++ b/EvoScientist/tools/image.py @@ -0,0 +1,132 @@ +"""Agent tools for image generation and editing. + +Thin wrappers over ``image_gen.service``. They never see API keys: entries +resolve ``${ENV_VAR}`` references server-side at call time. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from langchain_core.tools import tool + + +def _workspace() -> Path: + from EvoScientist.paths import _active_workspace + + return Path(_active_workspace).resolve() + + +def _settings_path() -> Path: + from EvoScientist.config.settings import get_config_path + + return get_config_path() + + +# Late-bound indirections so tests can patch without importing the service's +# adapter stack. +async def _generate_for_workspace(workspace: Path, **kwargs): + from EvoScientist.image_gen import service + + return await service.generate_for_workspace(workspace, **kwargs) + + +async def _edit_for_workspace(workspace: Path, **kwargs): + from EvoScientist.image_gen import service + + return await service.edit_for_workspace(workspace, **kwargs) + + +def _json_result(payload: dict) -> str: + return json.dumps(payload, ensure_ascii=False) + + +def _available_models_hint() -> str: + try: + from EvoScientist.image_gen.config import load_image_generation_settings + + settings = load_image_generation_settings(config_path=_settings_path()) + names = [entry.display_name() for entry in settings.models] + if not names: + return "" + default = settings.default_model or settings.models[0].id + return ( + f" Available image models: {', '.join(names)}. Default: {default}." + ) + except Exception: + return "" + + +def refresh_image_tool_descriptions() -> None: + """Refresh tool descriptions with the current image model list. + + Called on every agent build so config.yaml edits take effect without a + restart. Descriptions must never include credentials. + """ + hint = _available_models_hint() + generate_image.description = ( + "Generate one or more images with a dedicated image model and save " + f"them under artifacts/.{hint}" + ) + edit_image.description = ( + "Edit an existing workspace image with a dedicated image model and " + f"save the result under artifacts/.{hint}" + ) + + +@tool +async def generate_image( + prompt: str, + model: str | None = None, + size: str = "1024x1024", + quality: str = "auto", + background: str = "auto", + output_path: str | None = None, + n: int = 1, +) -> str: + """Generate one or more images with a dedicated image model and save them under artifacts/.""" + try: + result = await _generate_for_workspace( + _workspace(), + prompt=prompt, + model=model, + size=size, + quality=quality, + background=background, + output_path=output_path, + n=n, + ) + return _json_result(result) + except Exception as exc: + return _json_result({"ok": False, "error": str(exc) or exc.__class__.__name__}) + + +@tool +async def edit_image( + image_path: str, + prompt: str, + model: str | None = None, + mask_path: str | None = None, + size: str = "1024x1024", + quality: str = "auto", + output_path: str | None = None, +) -> str: + """Edit an existing workspace image with a dedicated image model and save the result under artifacts/.""" + try: + result = await _edit_for_workspace( + _workspace(), + image_path=image_path, + prompt=prompt, + model=model, + mask_path=mask_path, + size=size, + quality=quality, + output_path=output_path, + ) + return _json_result(result) + except Exception as exc: + return _json_result({"ok": False, "error": str(exc) or exc.__class__.__name__}) + + +refresh_image_tool_descriptions() diff --git a/tests/test_image_tools.py b/tests/test_image_tools.py new file mode 100644 index 0000000..9527557 --- /dev/null +++ b/tests/test_image_tools.py @@ -0,0 +1,93 @@ +"""Tests for the image generation agent tools.""" + +import json + +import pytest + +from EvoScientist.tools import image as image_tools + + +@pytest.fixture(autouse=True) +def fake_service(monkeypatch, tmp_path): + """Patch the service entry points the tools call.""" + async def fake_generate(workspace, **kwargs): + return { + "ok": True, + "path": "artifacts/generated_1.png", + "paths": ["artifacts/generated_1.png"], + "mime_type": "image/png", + "model": "gpt-image-2", + "size": kwargs.get("size", "1024x1024"), + } + + async def fake_edit(workspace, **kwargs): + return { + "ok": True, + "path": "artifacts/edited_1.png", + "paths": ["artifacts/edited_1.png"], + "mime_type": "image/png", + "model": "gpt-image-2", + "size": kwargs.get("size", "1024x1024"), + } + + monkeypatch.setattr(image_tools, "_generate_for_workspace", fake_generate) + monkeypatch.setattr(image_tools, "_edit_for_workspace", fake_edit) + monkeypatch.setattr(image_tools, "_workspace", lambda: tmp_path) + + +async def test_generate_image_success(): + result = json.loads( + await image_tools.generate_image.ainvoke({"prompt": "a cat"}) + ) + assert result["ok"] is True + assert result["paths"] == ["artifacts/generated_1.png"] + + +async def test_edit_image_success(): + result = json.loads( + await image_tools.edit_image.ainvoke( + {"image_path": "artifacts/src.png", "prompt": "make blue"} + ) + ) + assert result["ok"] is True + assert result["paths"] == ["artifacts/edited_1.png"] + + +async def test_error_becomes_json(monkeypatch): + from EvoScientist.image_gen.adapters.base import ImageGenError + + async def boom(workspace, **kwargs): + raise ImageGenError("no image models configured") + + monkeypatch.setattr(image_tools, "_generate_for_workspace", boom) + result = json.loads(await image_tools.generate_image.ainvoke({"prompt": "x"})) + assert result["ok"] is False + assert "no image models" in result["error"] + + +def test_refresh_descriptions_includes_models(monkeypatch, tmp_path): + config = tmp_path / "config.yaml" + config.write_text( + """ +image_generation: + default_model: gpt-image-2 + models: + - id: gpt-image-2 + name: GPT Image 2 + provider: openai + api_key: sk-x +""", + encoding="utf-8", + ) + monkeypatch.setattr(image_tools, "_settings_path", lambda: config) + image_tools.refresh_image_tool_descriptions() + assert "GPT Image 2" in image_tools.generate_image.description + assert "Default: gpt-image-2" in image_tools.generate_image.description + assert "sk-x" not in image_tools.generate_image.description + + +def test_tools_package_exports(): + from EvoScientist import tools + + assert tools.generate_image is image_tools.generate_image + assert tools.edit_image is image_tools.edit_image