diff --git a/EvoScientist/image_gen/__init__.py b/EvoScientist/image_gen/__init__.py new file mode 100644 index 0000000..1f45452 --- /dev/null +++ b/EvoScientist/image_gen/__init__.py @@ -0,0 +1,15 @@ +"""Image generation: dedicated image models, adapters, and agent tools.""" + +from .config import ( + ImageGenerationSettings, + ImageModelEntry, + is_image_generation_model, + load_image_generation_settings, +) + +__all__ = [ + "ImageGenerationSettings", + "ImageModelEntry", + "is_image_generation_model", + "load_image_generation_settings", +] diff --git a/EvoScientist/image_gen/config.py b/EvoScientist/image_gen/config.py new file mode 100644 index 0000000..98ea282 --- /dev/null +++ b/EvoScientist/image_gen/config.py @@ -0,0 +1,120 @@ +"""Configuration for dedicated image-generation models. + +Image generation models are service models, not chat models. They live in a +separate ``image_generation`` section of ``config.yaml`` so image-only models +such as ``gpt-image-2`` are never offered in the chat model selector. + +Secrets are stored as ``${ENV_VAR}`` references and resolved only at call +time, server-side. Tool signatures, skill docs, logs and error messages must +never carry a resolved key. +""" + +from __future__ import annotations + +import os +import re +from pathlib import Path +from typing import Any, Literal + +import yaml +from pydantic import BaseModel, Field, field_validator + +IMAGE_GENERATION_SECTION = "image_generation" +DEFAULT_IMAGE_GENERATION_TIMEOUT_SECONDS = 120.0 + +_IMAGE_MODEL_PREFIXES = ("gpt-image", "chatgpt-image", "dall-e", "dalle") +_IMAGE_MODEL_MARKERS = ("imagen", "wanx", "seedream") + +_ENV_REF = re.compile(r"^\$\{([A-Za-z_][A-Za-z0-9_]*)\}$") + + +def _resolve_env_ref(value: str) -> str: + match = _ENV_REF.match(value.strip()) + if match: + return os.environ.get(match.group(1), "") + return value + + +class ImageModelEntry(BaseModel): + """A single dedicated image-generation model.""" + + id: str = Field(..., description="Model ID sent to the image API") + name: str = "" + provider: Literal["openai", "gemini"] = "openai" + api_key: str = Field("", description="API key or ${ENV_VAR} reference") + base_url: str = Field("", description="API base URL or ${ENV_VAR} reference") + supports_generation: bool = True + supports_edit: bool = True + default_size: str = "1024x1024" + default_quality: str = "auto" + params: dict[str, Any] = Field(default_factory=dict) + + @field_validator("id") + @classmethod + def id_not_empty(cls, value: str) -> str: + value = value.strip() + if not value: + raise ValueError("image model id must not be empty") + return value + + def resolved_api_key(self) -> str: + return _resolve_env_ref(self.api_key) + + def resolved_base_url(self) -> str: + return _resolve_env_ref(self.base_url) + + def display_name(self) -> str: + return self.name or self.id + + +class ImageGenerationSettings(BaseModel): + """The ``image_generation`` section of config.yaml.""" + + default_model: str = "" + timeout_seconds: float = DEFAULT_IMAGE_GENERATION_TIMEOUT_SECONDS + models: list[ImageModelEntry] = Field(default_factory=list) + + @field_validator("timeout_seconds") + @classmethod + def timeout_must_be_positive(cls, value: float) -> float: + if value <= 0: + raise ValueError("timeout_seconds must be greater than 0") + return value + + def find_model(self, model_id_or_name: str) -> ImageModelEntry | None: + for entry in self.models: + if entry.id == model_id_or_name or entry.name == model_id_or_name: + return entry + return None + + +def is_image_generation_model(model_ref: str | None) -> bool: + """Return True when a model ID is known to be image-generation only.""" + if not model_ref: + return False + value = str(model_ref).strip().lower() + if "/" in value: + value = value.rsplit("/", 1)[1] + return value.startswith(_IMAGE_MODEL_PREFIXES) or any( + marker in value for marker in _IMAGE_MODEL_MARKERS + ) + + +def load_image_generation_settings( + *, config_path: Path | None = None +) -> ImageGenerationSettings: + """Load the image_generation section; empty settings when absent.""" + if config_path is None: + from EvoScientist.config.settings import get_config_path + + config_path = get_config_path() + if not config_path.exists(): + return ImageGenerationSettings() + try: + data = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {} + except yaml.YAMLError: + return ImageGenerationSettings() + section = data.get(IMAGE_GENERATION_SECTION) + if not isinstance(section, dict): + return ImageGenerationSettings() + return ImageGenerationSettings.model_validate(section) diff --git a/tests/test_image_gen_config.py b/tests/test_image_gen_config.py new file mode 100644 index 0000000..a3e95d2 --- /dev/null +++ b/tests/test_image_gen_config.py @@ -0,0 +1,120 @@ +"""Tests for image_gen.config.""" + +import os + +from EvoScientist.image_gen.config import ( + ImageGenerationSettings, + ImageModelEntry, + is_image_generation_model, + load_image_generation_settings, +) + + +def _write_config(tmp_path, body: str): + path = tmp_path / "config.yaml" + path.write_text(body, encoding="utf-8") + return path + + +def test_load_full_section(tmp_path, monkeypatch): + monkeypatch.setenv("TEST_IMG_KEY", "sk-secret-123") + path = _write_config( + tmp_path, + """ +image_generation: + default_model: gpt-image-2 + timeout_seconds: 90 + models: + - id: gpt-image-2 + name: GPT Image 2 + provider: openai + api_key: ${TEST_IMG_KEY} + base_url: https://api.openai.com + default_size: "1024x1024" + - id: imagen-4.0-generate-001 + provider: gemini + api_key: ${TEST_IMG_KEY} + supports_edit: false +""", + ) + settings = load_image_generation_settings(config_path=path) + assert settings.default_model == "gpt-image-2" + assert settings.timeout_seconds == 90 + assert len(settings.models) == 2 + openai_entry = settings.models[0] + assert openai_entry.resolved_api_key() == "sk-secret-123" + assert openai_entry.resolved_base_url() == "https://api.openai.com" + assert openai_entry.display_name() == "GPT Image 2" + gemini_entry = settings.models[1] + assert gemini_entry.provider == "gemini" + assert gemini_entry.supports_edit is False + assert gemini_entry.display_name() == "imagen-4.0-generate-001" + + +def test_missing_file_returns_empty_settings(tmp_path): + settings = load_image_generation_settings(config_path=tmp_path / "nope.yaml") + assert settings == ImageGenerationSettings() + + +def test_missing_section_returns_empty_settings(tmp_path): + path = _write_config(tmp_path, "webui_port: 4716\n") + settings = load_image_generation_settings(config_path=path) + assert settings.models == [] + + +def test_unresolved_env_ref_yields_empty(tmp_path, monkeypatch): + monkeypatch.delenv("NOPE_NOT_SET", raising=False) + path = _write_config( + tmp_path, + """ +image_generation: + models: + - id: gpt-image-2 + provider: openai + api_key: ${NOPE_NOT_SET} +""", + ) + settings = load_image_generation_settings(config_path=path) + assert settings.models[0].resolved_api_key() == "" + + +def test_plain_value_passthrough(tmp_path): + path = _write_config( + tmp_path, + """ +image_generation: + models: + - id: gpt-image-2 + provider: openai + api_key: sk-literal +""", + ) + settings = load_image_generation_settings(config_path=path) + assert settings.models[0].resolved_api_key() == "sk-literal" + + +def test_is_image_generation_model(): + assert is_image_generation_model("gpt-image-2") is True + assert is_image_generation_model("dall-e-3") is True + assert is_image_generation_model("openai/gpt-image-1") is True + assert is_image_generation_model("imagen-4.0-generate-001") is True + assert is_image_generation_model("doubao-seedream-4-0") is True + assert is_image_generation_model("wanx2.1-t2i-turbo") is True + assert is_image_generation_model("glm-5.2") is False + assert is_image_generation_model("gpt-4o") is False + assert is_image_generation_model(None) is False + assert is_image_generation_model("") is False + + +def test_evoscientist_config_ignores_section(tmp_path, monkeypatch): + """The platform config loader filters unknown keys, so the image section + never leaks through EvoScientistConfig / its API surface.""" + from EvoScientist.config import settings as config_settings + + path = _write_config( + tmp_path, + "image_generation:\n models:\n - id: gpt-image-2\n api_key: sk-x\n", + ) + monkeypatch.setattr(config_settings, "get_config_path", lambda: path) + cfg = config_settings.load_config() + assert not hasattr(cfg, "image_generation")