From efb86f55dd0d4a92d6e91c5b6dc1f60ee844d5f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=93=9D=E5=8F=8B=E5=92=8C?= <3068150232@qq.com> Date: Sun, 2 Aug 2026 18:23:53 +0800 Subject: [PATCH 1/4] add seedream tool --- tools/graphics/seedream_image.py | 108 +++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 tools/graphics/seedream_image.py diff --git a/tools/graphics/seedream_image.py b/tools/graphics/seedream_image.py new file mode 100644 index 00000000..edaefd2a --- /dev/null +++ b/tools/graphics/seedream_image.py @@ -0,0 +1,108 @@ +"""Seedream V5 image generation via fal.ai API. +deep-thinking prompt understanding, native text in 14 languages, and precise control over dense layouts and structured designs. +""" +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 SeedreamImage(BaseTool): + name = "seedream_image" + version = "0.1.0" + tier = ToolTier.GENERATE + capability = "image_generation" + provider = "seedream" + 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", + "structured_designs", + "dense_layouts", + "multi_language_text", + ] + supports = { + "svg_output": True, + "text_rendering": True, + "color_palette": True, + "custom_size": True, + "structured_designs": True, + "dense_layouts": True, + "multi_language_text": True, + } + best_for = [ + "logos and brand assets", + "SVG vector output", + "images with accurate text rendering", + "structured designs and dense layouts", + "multi-language text rendering (14 languages)", + ] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": {"type": "string"}, + "image_size": { + "type": "string", + "enum": [ + "square", "square_hd", + "landscape_4_3", "landscape_16_9", + "portrait_4_3", "portrait_16_9", + "auto_1K","auto_2K" + ], + "default": "auto_2K", + }, + "num_images": { + "type": "number", + "default": 1, + }, + "output_format": { + "type": "string", + "enum": ["jpeg", "png"], + "description": "Output image format. Use 'jpeg' for smaller file size with lossy compression (suitable for web/preview), or 'png' for lossless quality with transparency support (suitable for design assets and further editing).", + }, + "enable_safety_checker": { + "type": "boolean", + "default": True, + "description": "If set to true, the safety checker will be enabled.", + }, + "enable_safety_checker": { + "type": "boolean", + "description": "IIf True, the media will be returned as a data URI and the output data won't be available in the request history.", + } + }, + } + 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"] From cf3407155cbfd33c66e7237622a61130cde07275 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=93=9D=E5=8F=8B=E5=92=8C?= <3068150232@qq.com> Date: Mon, 3 Aug 2026 00:53:20 +0800 Subject: [PATCH 2/4] add seedream v5 tools for fai.ai --- tools/graphics/seedream_image.py | 164 ++++++++++++++++++++++++++++--- 1 file changed, 153 insertions(+), 11 deletions(-) diff --git a/tools/graphics/seedream_image.py b/tools/graphics/seedream_image.py index edaefd2a..fed26184 100644 --- a/tools/graphics/seedream_image.py +++ b/tools/graphics/seedream_image.py @@ -25,9 +25,9 @@ class SeedreamImage(BaseTool): version = "0.1.0" tier = ToolTier.GENERATE capability = "image_generation" - provider = "seedream" + provider = "bytedance" stability = ToolStability.EXPERIMENTAL - execution_mode = ExecutionMode.SYNC + execution_mode = ExecutionMode.ASYNC determinism = Determinism.STOCHASTIC runtime = ToolRuntime.API @@ -63,7 +63,6 @@ class SeedreamImage(BaseTool): "structured designs and dense layouts", "multi-language text rendering (14 languages)", ] - input_schema = { "type": "object", "required": ["prompt"], @@ -93,16 +92,159 @@ class SeedreamImage(BaseTool): "default": True, "description": "If set to true, the safety checker will be enabled.", }, - "enable_safety_checker": { - "type": "boolean", - "description": "IIf True, the media will be returned as a data URI and the output data won't be available in the request history.", - } - }, - } + "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"] + idempotency_key_fields = ["prompt", "image_size", "output_format"] + side_effects = ["writes image file to output_path", "calls fal.ai queue 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: + image_size = inputs.get("image_size", "auto_2K") + num_images = inputs.get("num_images", 1) + size_price_map = { + "square": 0.0675, + "square_hd": 0.135, + "landscape_4_3": 0.0675, + "landscape_16_9": 0.135, + "portrait_4_3": 0.0675, + "portrait_16_9": 0.135, + "auto_1K": 0.0675, + "auto_2K": 0.135, + } + unit_price = size_price_map.get(image_size, 0.135) + return round(unit_price * num_images, 4) + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + import requests + + api_key = self._get_api_key() + if not api_key: + return ToolResult( + success=False, + error="FAL_KEY not set. " + self.install_instructions, + ) + + start = time.time() + prompt = inputs["prompt"] + submit_url = "https://queue.fal.run/bytedance/seedream/v5/pro/text-to-image" + payload: dict[str, Any] = { + "prompt": prompt, + "image_size": inputs.get("image_size", "auto_2K"), + "output_format": inputs.get("output_format", "jpeg"), + "num_images": inputs.get("num_images", 1), + "enable_safety_checker": inputs.get("enable_safety_checker", True), + } + + try: + headers = { + "Authorization": f"Key {api_key}", + "Content-Type": "application/json", + } + + submit_resp = requests.post( + submit_url, + headers=headers, + json=payload, + timeout=(10, 60), + ) + submit_resp.raise_for_status() + submit_data = submit_resp.json() + request_id = submit_data.get("request_id") + if not request_id: + raise RuntimeError( + "Seedream submit succeeded but did not return request_id" + ) + status_url = ( + f"https://queue.fal.run/bytedance/seedream/requests/" + f"{request_id}/status" + ) + elapsed = 0.0 + while elapsed < 300: + status_resp = requests.get( + status_url, + headers=headers, + timeout=30, + ) + status_resp.raise_for_status() + status_data = status_resp.json() + status = status_data.get("status") + + if status == "COMPLETED": + break + elif status in ("FAILED", "CANCELLED"): + error_msg = status_data.get("error", "Unknown error") + raise RuntimeError(f"Seedream task {status}: {error_msg}") + + time.sleep(10) + elapsed += 10 + + if elapsed >= 300: + raise RuntimeError( + f"Seedream task timed out after {300}s" + ) + + result_resp = requests.get( + f"https://queue.fal.run/bytedance/seedream/requests/" + f"{request_id}", + headers=headers, + timeout=30, + ) + result_resp.raise_for_status() + result_data = result_resp.json() + + images = result_data.get("images", []) + if not images: + raise RuntimeError("Seedream completed but no images returned") + + output_paths = [] + for idx, img in enumerate(images): + image_url = img.get("url") + if not image_url: + continue + image_resp = requests.get(image_url, timeout=60) + image_resp.raise_for_status() + + ext = inputs.get("output_format", "jpeg") + if len(images) > 1: + filename = f"seedream_image_{idx + 1}.{ext}" + else: + filename = f"seedream_image.{ext}" + output_path = Path(inputs.get("output_path", filename)) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_bytes(image_resp.content) + output_paths.append(str(output_path)) + + except Exception as e: + return ToolResult( + success=False, + error=f"Seedream generation failed: {e}", + ) + + return ToolResult( + success=True, + data={ + "provider": "seedream", + "model": "seedream_v5", + "prompt": prompt, + "request_id": request_id, + "image_count": len(output_paths), + "outputs": output_paths, + }, + artifacts=output_paths, + cost_usd=self.estimate_cost(inputs), + duration_seconds=round(time.time() - start, 2), + model="fal-ai/bytedance/seedream/v5", + ) From 63fd646717cd197bcdac6b04aa710a13f68107b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=93=9D=E5=8F=8B=E5=92=8C?= <3068150232@qq.com> Date: Sun, 9 Aug 2026 00:25:44 +0800 Subject: [PATCH 3/4] add seedreamm tools --- tests/tools/test_seedream_image.py | 262 +++++++++++++++++++++++++++++ tools/graphics/seedream_image.py | 4 +- 2 files changed, 264 insertions(+), 2 deletions(-) create mode 100644 tests/tools/test_seedream_image.py diff --git a/tests/tools/test_seedream_image.py b/tests/tools/test_seedream_image.py new file mode 100644 index 00000000..9f95ca0d --- /dev/null +++ b/tests/tools/test_seedream_image.py @@ -0,0 +1,262 @@ +"""Regression tests: seedream_image must return every image it requests and bills for. + +Covers: +- Multi-image output: all requested images must be written and returned +- Cost estimation: billed count matches delivered artifacts +- Single-image output: exact output path preserved +- Async polling: COMPLETED / FAILED / CANCELLED / timeout paths +- API key validation: graceful failure when FAL_KEY is unset +""" + +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent +sys.path.insert(0, str(PROJECT_ROOT)) + + +class _FakeResponse: + def __init__(self, json_data: dict | None = None, status_code: int = 200, content: bytes = b""): + self._json_data = json_data or {} + self.status_code = status_code + self.content = content + + def raise_for_status(self): + if self.status_code >= 400: + import requests + raise requests.HTTPError(response=self) + + def json(self): + return self._json_data + + +def _build_submit_response(request_id: str = "req_123") -> _FakeResponse: + return _FakeResponse({"request_id": request_id}) + + +def _build_status_response(status: str, error: str | None = None) -> _FakeResponse: + data = {"status": status} + if error: + data["error"] = error + return _FakeResponse(data) + + +def _build_result_response(image_urls: list[str]) -> _FakeResponse: + images = [{"url": url} for url in image_urls] + return _FakeResponse({"images": images}) + + +def _build_image_content(index: int) -> bytes: + return f"SEEDREAM_IMAGE_{index}".encode() + + +@pytest.fixture +def seedream_tool(monkeypatch): + monkeypatch.setenv("FAL_KEY", "test-fal-key") + from tools.graphics.seedream_image import SeedreamImage + return SeedreamImage() + + +@pytest.fixture +def mock_requests(monkeypatch): + mock_post = MagicMock() + mock_get = MagicMock() + fake_requests = types.ModuleType("requests") + fake_requests.post = mock_post + fake_requests.get = mock_get + fake_requests.HTTPError = type("HTTPError", (Exception,), {}) + monkeypatch.setitem(sys.modules, "requests", fake_requests) + return mock_post, mock_get + + +def _setup_mock_execution(mock_post, mock_get, num_images: int = 1, status: str = "COMPLETED", + error: str | None = None, extra_gets: list = None): + """Helper to setup common mock execution flow.""" + mock_post.return_value = _build_submit_response() + + side_effects = [_build_status_response(status, error)] + if status == "COMPLETED": + urls = [f"http://img.url/{i}" for i in range(num_images)] + side_effects.append(_build_result_response(urls)) + side_effects.extend([_FakeResponse(content=_build_image_content(i)) for i in range(num_images)]) + elif extra_gets: + side_effects.extend(extra_gets) + + mock_get.side_effect = side_effects + + +# ========== Core Regression Tests ========== + +class TestMultiOutputRegression: + def test_all_requested_images_are_written(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=3) + + result = seedream_tool.execute({ + "prompt": "test", "num_images": 3, + "output_format": "jpeg", "output_path": str(tmp_path / "gen.jpeg"), + }) + + assert result.success + assert result.data["image_count"] == 3 + assert len(result.artifacts) == 3 + + files = sorted(tmp_path.glob("*.jpeg")) + assert len(files) == 3 + contents = {f.read_bytes() for f in files} + assert contents == {b"SEEDREAM_IMAGE_0", b"SEEDREAM_IMAGE_1", b"SEEDREAM_IMAGE_2"} + + def test_artifacts_match_billed_count(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=4) + + inputs = {"prompt": "t", "num_images": 4, "output_path": str(tmp_path / "out.png")} + result = seedream_tool.execute(inputs) + billed = seedream_tool.estimate_cost(inputs) + + assert len(result.artifacts) == 4 + assert billed == pytest.approx(0.135 * 4) + + +class TestSingleOutput: + def test_single_image_keeps_exact_path(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=1) + + out = tmp_path / "single.png" + result = seedream_tool.execute({"prompt": "s", "num_images": 1, "output_path": str(out)}) + + assert result.success + assert result.artifacts == [str(out)] + assert out.read_bytes() == b"SEEDREAM_IMAGE_0" + + +# ========== Cost Estimation (Parameterized) ========== + +class TestCostEstimation: + @pytest.mark.parametrize("size,expected", [ + ("square", 0.0675), ("landscape_4_3", 0.0675), + ("portrait_4_3", 0.0675), ("auto_1K", 0.0675), + ]) + def test_small_size_pricing(self, seedream_tool, size, expected): + cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1}) + assert cost == pytest.approx(expected) + + @pytest.mark.parametrize("size,expected", [ + ("square_hd", 0.135), ("landscape_16_9", 0.135), + ("portrait_16_9", 0.135), ("auto_2K", 0.135), + ]) + def test_large_size_pricing(self, seedream_tool, size, expected): + cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1}) + assert cost == pytest.approx(expected) + + @pytest.mark.parametrize("n",[1, 2, 5, 10]) + def test_cost_scales_with_num_images(self, seedream_tool, n): + cost = seedream_tool.estimate_cost({"image_size": "auto_2K", "num_images": n}) + assert cost == pytest.approx(round(0.135 * n, 4)) + + def test_unknown_size_falls_back_to_high_price(self, seedream_tool): + cost = seedream_tool.estimate_cost({"image_size": "unknown", "num_images": 1}) + assert cost == pytest.approx(0.135) + + def test_default_values(self, seedream_tool): + cost = seedream_tool.estimate_cost({}) + assert cost == pytest.approx(0.135) + + +# ========== Async Polling States ========== + +class TestAsyncPolling: + def test_completed_on_first_poll(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=1) + + result = seedream_tool.execute({"prompt": "q", "output_path": str(tmp_path / "q.png")}) + assert result.success + assert result.data["request_id"] + + @pytest.mark.parametrize("status,error_msg", [ + ("FAILED", "Content policy violation"), + ("CANCELLED", None), + ]) + def test_failed_states_return_error(self, seedream_tool, mock_requests, status, error_msg): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, status=status, error=error_msg) + + result = seedream_tool.execute({"prompt": "bad"}) + assert not result.success + assert status in result.error + + def test_timeout_returns_error(self, seedream_tool, mock_requests): + mock_post, mock_get = mock_requests + mock_post.return_value = _build_submit_response() + mock_get.side_effect = [_build_status_response("IN_PROGRESS")] * 100 + + with patch("tools.graphics.seedream_image.time.sleep"): + result = seedream_tool.execute({"prompt": "timeout"}) + assert not result.success + assert "timed out" in result.error.lower() + + +# ========== Validation & Error Handling ========== + +class TestValidation: + def test_missing_api_key_returns_error(self, monkeypatch): + monkeypatch.delenv("FAL_KEY", raising=False) + monkeypatch.delenv("FAL_AI_API_KEY", raising=False) + from tools.graphics.seedream_image import SeedreamImage + result = SeedreamImage().execute({"prompt": "t"}) + assert not result.success + assert "FAL_KEY" in result.error + + def test_status_available_with_key(self, seedream_tool): + assert seedream_tool.get_status().name == "AVAILABLE" + + def test_status_unavailable_without_key(self, monkeypatch): + monkeypatch.delenv("FAL_KEY", raising=False) + monkeypatch.delenv("FAL_AI_API_KEY", raising=False) + from tools.graphics.seedream_image import SeedreamImage + assert SeedreamImage().get_status().name == "UNAVAILABLE" + + def test_missing_request_id_raises_error(self, seedream_tool, mock_requests): + mock_post, mock_get = mock_requests + mock_post.return_value = _FakeResponse({}) + result = seedream_tool.execute({"prompt": "no id"}) + assert not result.success + assert "request_id" in result.error.lower() + + def test_completed_without_images_raises_error(self, seedream_tool, mock_requests): + mock_post, mock_get = mock_requests + mock_post.return_value = _build_submit_response() + mock_get.side_effect = [ + _build_status_response("COMPLETED"), + _FakeResponse({"images": []}), + ] + result = seedream_tool.execute({"prompt": "empty"}) + assert not result.success + assert "no images" in result.error.lower() + + +# ========== Metadata & Integration ========== + +class TestMetadata: + def test_provider_and_model_info(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=1) + + result = seedream_tool.execute({"prompt": "m", "output_path": str(tmp_path / "m.png")}) + assert result.data["provider"] == "seedream" + assert result.data["model"] == "seedream_v5" + assert result.model == "fal-ai/bytedance/seedream/v5" + + def test_cost_matches_estimate(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=2) + + inputs = {"prompt": "c", "image_size": "square", "num_images": 2, "output_path": str(tmp_path / "c.jpeg")} + result = seedream_tool.execute(inputs) + assert result.cost_usd == pytest.approx(seedream_tool.estimate_cost(inputs)) \ No newline at end of file diff --git a/tools/graphics/seedream_image.py b/tools/graphics/seedream_image.py index fed26184..a39b537c 100644 --- a/tools/graphics/seedream_image.py +++ b/tools/graphics/seedream_image.py @@ -99,7 +99,7 @@ class SeedreamImage(BaseTool): 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", "image_size", "output_format"] + idempotency_key_fields = ["image_size", "output_format"] side_effects = ["writes image file to output_path", "calls fal.ai queue API"] user_visible_verification = ["Inspect generated image for brand accuracy and text readability"] @@ -247,4 +247,4 @@ class SeedreamImage(BaseTool): cost_usd=self.estimate_cost(inputs), duration_seconds=round(time.time() - start, 2), model="fal-ai/bytedance/seedream/v5", - ) + ) \ No newline at end of file From 171866dbbf6cb25e095615542822a7c9d1877766 Mon Sep 17 00:00:00 2001 From: calesthio Date: Thu, 13 Aug 2026 09:12:00 -0700 Subject: [PATCH 4/4] fix: complete Seedream provider integration --- docs/PROVIDERS.md | 4 +- tests/tools/test_seedream_image.py | 41 +++++++++++++++++++- tools/graphics/image_selector.py | 2 + tools/graphics/seedream_image.py | 61 +++++++++++++++++++++--------- 4 files changed, 87 insertions(+), 21 deletions(-) diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 497f78df..515af259 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -145,7 +145,7 @@ The ASR tool (`qwen3-asr-flash-filetrans`) uses an async submit-poll pattern. Au > **Broad single-key coverage.** One API key unlocks image and video providers across multiple models. -**Tools unlocked:** `flux_image`, `recraft_image`, `kling_video`, `veo_video`, `minimax_video` +**Tools unlocked:** `flux_image`, `recraft_image`, `seedream_image`, `kling_video`, `veo_video`, `minimax_video` **Env var:** `FAL_KEY` #### Setup @@ -166,6 +166,8 @@ No subscription — pure pay-as-you-go, no minimum spend. | FLUX Pro v1.1 | $0.05/image | 20 images | | FLUX Dev | $0.03/image | 33 images | | Recraft v3 | ~$0.04/image | 25 images | +| Seedream 5 Pro (up to 1536x1536) | $0.0675/image | ~14 images | +| Seedream 5 Pro (up to 2048x2048) | $0.135/image | ~7 images | **Video generation:** diff --git a/tests/tools/test_seedream_image.py b/tests/tools/test_seedream_image.py index 9f95ca0d..4932a41d 100644 --- a/tests/tools/test_seedream_image.py +++ b/tests/tools/test_seedream_image.py @@ -154,7 +154,7 @@ class TestCostEstimation: cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1}) assert cost == pytest.approx(expected) - @pytest.mark.parametrize("n",[1, 2, 5, 10]) + @pytest.mark.parametrize("n", [1, 2, 3, 4]) def test_cost_scales_with_num_images(self, seedream_tool, n): cost = seedream_tool.estimate_cost({"image_size": "auto_2K", "num_images": n}) assert cost == pytest.approx(round(0.135 * n, 4)) @@ -205,6 +205,16 @@ class TestAsyncPolling: # ========== Validation & Error Handling ========== class TestValidation: + @pytest.mark.parametrize("value", [0, 5, 1.5, True]) + def test_num_images_rejects_invalid_values( + self, seedream_tool, mock_requests, value + ): + mock_post, _ = mock_requests + result = seedream_tool.execute({"prompt": "t", "num_images": value}) + assert not result.success + assert "num_images" in (result.error or "") + mock_post.assert_not_called() + def test_missing_api_key_returns_error(self, monkeypatch): monkeypatch.delenv("FAL_KEY", raising=False) monkeypatch.delenv("FAL_AI_API_KEY", raising=False) @@ -259,4 +269,31 @@ class TestMetadata: inputs = {"prompt": "c", "image_size": "square", "num_images": 2, "output_path": str(tmp_path / "c.jpeg")} result = seedream_tool.execute(inputs) - assert result.cost_usd == pytest.approx(seedream_tool.estimate_cost(inputs)) \ No newline at end of file + assert result.cost_usd == pytest.approx(seedream_tool.estimate_cost(inputs)) + + def test_image_selector_routes_count_and_returns_distinct_artifacts( + self, seedream_tool, tmp_path, mock_requests, monkeypatch + ): + from tools.graphics.image_selector import ImageSelector + + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=2) + selector = ImageSelector() + monkeypatch.setattr(selector, "_providers", lambda: [seedream_tool]) + + result = selector.execute( + { + "prompt": "campaign artwork", + "preferred_provider": "bytedance", + "n": 2, + "output_path": str(tmp_path / "selected.png"), + } + ) + + assert result.success, result.error + assert result.data["selected_tool"] == "seedream_image" + assert len(set(result.artifacts)) == 2 + assert {Path(path).name for path in result.artifacts} == { + "selected_1.png", + "selected_2.png", + } diff --git a/tools/graphics/image_selector.py b/tools/graphics/image_selector.py index b106fbde..c3a3c716 100644 --- a/tools/graphics/image_selector.py +++ b/tools/graphics/image_selector.py @@ -242,6 +242,8 @@ class ImageSelector(BaseTool): props = tool.input_schema.get("properties", {}) if "query" in props and "query" not in adapted: adapted["query"] = adapted.get("prompt", "") + if "n" in adapted and "num_images" in props and "num_images" not in adapted: + adapted["num_images"] = adapted["n"] # Strip selector-only keys that downstream tools don't understand adapted.pop("preferred_provider", None) diff --git a/tools/graphics/seedream_image.py b/tools/graphics/seedream_image.py index a39b537c..64acc8d8 100644 --- a/tools/graphics/seedream_image.py +++ b/tools/graphics/seedream_image.py @@ -31,24 +31,21 @@ class SeedreamImage(BaseTool): determinism = Determinism.STOCHASTIC runtime = ToolRuntime.API - dependencies = [] + dependencies = ["env:FAL_KEY"] install_instructions = ( "Set FAL_KEY to your fal.ai API key.\n" " Get one at https://fal.ai/dashboard/keys" ) - agent_skills = [] + agent_skills = ["visual-style"] capabilities = [ "generate_image", - "generate_logo", - "generate_vector", "text_to_image", "structured_designs", "dense_layouts", "multi_language_text", ] supports = { - "svg_output": True, "text_rendering": True, "color_palette": True, "custom_size": True, @@ -57,8 +54,7 @@ class SeedreamImage(BaseTool): "multi_language_text": True, } best_for = [ - "logos and brand assets", - "SVG vector output", + "raster brand and campaign assets", "images with accurate text rendering", "structured designs and dense layouts", "multi-language text rendering (14 languages)", @@ -79,7 +75,9 @@ class SeedreamImage(BaseTool): "default": "auto_2K", }, "num_images": { - "type": "number", + "type": "integer", + "minimum": 1, + "maximum": 4, "default": 1, }, "output_format": { @@ -99,7 +97,13 @@ class SeedreamImage(BaseTool): 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 = ["image_size", "output_format"] + idempotency_key_fields = [ + "prompt", + "image_size", + "output_format", + "num_images", + "enable_safety_checker", + ] side_effects = ["writes image file to output_path", "calls fal.ai queue API"] user_visible_verification = ["Inspect generated image for brand accuracy and text readability"] @@ -127,6 +131,20 @@ class SeedreamImage(BaseTool): unit_price = size_price_map.get(image_size, 0.135) return round(unit_price * num_images, 4) + @staticmethod + def _output_paths( + output_path: str | None, count: int, output_format: str + ) -> list[Path]: + path = Path(output_path or f"seedream_image.{output_format}") + if not path.suffix: + path = path.with_suffix(f".{output_format}") + if count == 1: + return [path] + return [ + path.with_name(f"{path.stem}_{index}{path.suffix}") + for index in range(1, count + 1) + ] + def execute(self, inputs: dict[str, Any]) -> ToolResult: import requests @@ -139,12 +157,21 @@ class SeedreamImage(BaseTool): start = time.time() prompt = inputs["prompt"] + num_images = inputs.get("num_images", 1) + if isinstance(num_images, bool) or not isinstance(num_images, int): + return ToolResult( + success=False, error="num_images must be an integer from 1 to 4." + ) + if not 1 <= num_images <= 4: + return ToolResult( + success=False, error="num_images must be between 1 and 4." + ) submit_url = "https://queue.fal.run/bytedance/seedream/v5/pro/text-to-image" payload: dict[str, Any] = { "prompt": prompt, "image_size": inputs.get("image_size", "auto_2K"), "output_format": inputs.get("output_format", "jpeg"), - "num_images": inputs.get("num_images", 1), + "num_images": num_images, "enable_safety_checker": inputs.get("enable_safety_checker", True), } @@ -209,20 +236,18 @@ class SeedreamImage(BaseTool): if not images: raise RuntimeError("Seedream completed but no images returned") + ext = inputs.get("output_format", "jpeg") + expected_paths = self._output_paths( + inputs.get("output_path"), len(images), ext + ) output_paths = [] - for idx, img in enumerate(images): + for img, output_path in zip(images, expected_paths): image_url = img.get("url") if not image_url: continue image_resp = requests.get(image_url, timeout=60) image_resp.raise_for_status() - ext = inputs.get("output_format", "jpeg") - if len(images) > 1: - filename = f"seedream_image_{idx + 1}.{ext}" - else: - filename = f"seedream_image.{ext}" - output_path = Path(inputs.get("output_path", filename)) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_bytes(image_resp.content) output_paths.append(str(output_path)) @@ -247,4 +272,4 @@ class SeedreamImage(BaseTool): cost_usd=self.estimate_cost(inputs), duration_seconds=round(time.time() - start, 2), model="fal-ai/bytedance/seedream/v5", - ) \ No newline at end of file + )