diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 0574a12a..2c49bf49 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -193,7 +193,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 @@ -214,6 +214,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 new file mode 100644 index 00000000..4932a41d --- /dev/null +++ b/tests/tools/test_seedream_image.py @@ -0,0 +1,299 @@ +"""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, 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)) + + 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: + @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) + 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)) + + 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 new file mode 100644 index 00000000..64acc8d8 --- /dev/null +++ b/tools/graphics/seedream_image.py @@ -0,0 +1,275 @@ +"""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 = "bytedance" + stability = ToolStability.EXPERIMENTAL + execution_mode = ExecutionMode.ASYNC + determinism = Determinism.STOCHASTIC + runtime = ToolRuntime.API + + 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 = ["visual-style"] + + capabilities = [ + "generate_image", + "text_to_image", + "structured_designs", + "dense_layouts", + "multi_language_text", + ] + supports = { + "text_rendering": True, + "color_palette": True, + "custom_size": True, + "structured_designs": True, + "dense_layouts": True, + "multi_language_text": True, + } + best_for = [ + "raster brand and campaign assets", + "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": "integer", + "minimum": 1, + "maximum": 4, + "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.", + }, + "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", + "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"] + + 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) + + @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 + + 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"] + 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": num_images, + "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") + + ext = inputs.get("output_format", "jpeg") + expected_paths = self._output_paths( + inputs.get("output_path"), len(images), ext + ) + output_paths = [] + 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() + + 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", + )