Merge pull request #6261 from MervinPraison/claude/issue-6260-20261105

fix: ImageAgent results gain .save() (no manual base64 decoding)

src/praisonai-agents/praisonaiagents/agent/_image_result.py

@@ -0,0 +1,69 @@
+"""Savable wrapper for ImageAgent responses (keeps the raw response API intact)."""
+import base64
+from typing import Any, List
+
+_MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024
+
+
+class ImageResult:
+    """Wraps a provider image response and adds ``save()``.
+
+    Attribute and item access fall through to the wrapped response, so
+    ``result.data[0].b64_json`` and ``result["data"]`` keep working.
+    """
+
+    def __init__(self, response: Any):
+        self._response = response
+
+    def __getattr__(self, name: str) -> Any:
+        if "_response" not in self.__dict__:
+            raise AttributeError(name)
+        return getattr(self.__dict__["_response"], name)
+
+    def __getitem__(self, key: Any) -> Any:
+        return self._response[key]
+
+    def __contains__(self, key: Any) -> bool:
+        return key in self._response
+
+    def __iter__(self):
+        return iter(self._response)
+
+    def __repr__(self) -> str:
+        return f"ImageResult({self._response!r})"
+
+    @property
+    def images(self) -> List[Any]:
+        """List of image objects returned by the provider (one per ``n``)."""
+        data = self._response.get("data") if isinstance(self._response, dict) else getattr(self._response, "data", None)
+        return list(data or [])
+
+    def save(self, path: str, index: int = 0) -> str:
+        """Write image ``index`` to ``path`` (base64 is decoded, http(s) URLs are downloaded)."""
+        images = self.images
+        if not images:
+            raise ValueError("No image data in response")
+        image = images[index]
+        get = image.get if isinstance(image, dict) else (lambda k: getattr(image, k, None))
+        b64, url = get("b64_json"), get("url")
+        if b64:
+            data = base64.b64decode(b64)
+        elif url:
+            data = _download(url)
+        else:
+            raise ValueError("Image has neither b64_json nor url")
+        with open(path, "wb") as f:
+            f.write(data)
+        return path
+
+
+def _download(url: str) -> bytes:
+    import urllib.request
+
+    if not url.startswith(("https://", "http://")):
+        raise ValueError("Image URL must be http(s)")
+    with urllib.request.urlopen(url, timeout=30) as resp:  # noqa: S310 - scheme checked above
+        data = resp.read(_MAX_DOWNLOAD_BYTES + 1)
+    if len(data) > _MAX_DOWNLOAD_BYTES:
+        raise ValueError("Image exceeds maximum download size")
+    return data

src/praisonai-agents/praisonaiagents/agent/image_agent.py

@@ -6,6 +6,7 @@ including support for different image models, sizes, and quality settings.
 import asyncio
 from typing import Optional, Any, Dict, Union, List
 from ..agent.agent import Agent
+from ._image_result import ImageResult
 from pydantic import BaseModel, Field
 import logging
 from praisonaiagents._logging import get_logger
@@ -153,7 +154,7 @@ class ImageAgent(Agent):
         model_info = self.llm_model
         return model_info.model if hasattr(model_info, 'model') else str(model_info)
 
-    def generate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
+    def generate_image(self, prompt: str, **kwargs) -> ImageResult:
         """Generate an image based on the provided prompt."""
         # Merge default config with any provided kwargs
         config = self.image_config.dict(exclude_none=True)
@@ -208,27 +209,27 @@ class ImageAgent(Agent):
                 
                 # Mark task as complete
                 progress.update(task, completed=True)
-                return response
+                return ImageResult(response)
             except Exception as e:
                 error_msg = f"Error generating image: {str(e)}"
                 if self.verbose:
                     self.console.print(f"[red]{error_msg}[/red]")
                 logging.error(error_msg)
                 raise
-    async def agenerate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
+    async def agenerate_image(self, prompt: str, **kwargs) -> ImageResult:
         """Offload the synchronous generation path to keep the event loop responsive."""
         return await asyncio.to_thread(self.generate_image, prompt, **kwargs)
     
     # Aliases for consistency with other agents
-    def generate(self, prompt: str, **kwargs) -> Dict[str, Any]:
+    def generate(self, prompt: str, **kwargs) -> ImageResult:
         """Alias for generate_image() - for consistency with VideoAgent/AudioAgent."""
         return self.generate_image(prompt, **kwargs)
     
-    async def agenerate(self, prompt: str, **kwargs) -> Dict[str, Any]:
+    async def agenerate(self, prompt: str, **kwargs) -> ImageResult:
         """Async alias for generate_image()."""
         return await self.agenerate_image(prompt, **kwargs)
 
-    def chat(self, prompt: str, **kwargs) -> Dict[str, Any]:
+    def chat(self, prompt: str, **kwargs) -> Union[ImageResult, Dict[str, Any]]:
         """Generate an image from the prompt."""
         try:
             result = self.generate_image(prompt, **kwargs)
@@ -249,7 +250,7 @@ class ImageAgent(Agent):
         output_pydantic: Optional[Any] = None,
         reasoning_steps: bool = False,
         **kwargs
-    ) -> Union[str, Dict[str, Any]]:
+    ) -> Union[ImageResult, Dict[str, Any]]:
         """Async chat method for image generation."""
         try:
             image_result = await self.agenerate_image(prompt, **kwargs)
@@ -270,7 +271,7 @@ class ImageAgent(Agent):
         n: int = 1,
         size: Optional[str] = None,
         **kwargs
-    ) -> Dict[str, Any]:
+    ) -> ImageResult:
         """
         Edit an existing image with a prompt.
         
@@ -288,7 +289,7 @@ class ImageAgent(Agent):
         Example:
             ```python
             agent = ImageAgent(llm="openai/gpt-image-2.5-flare")
-            result = agent.edit("photo.png", "Add a sunset in the background")
+            agent.edit("photo.png", "Add a sunset in the background").save("out.png")
             ```
         """
         try:
@@ -318,7 +319,7 @@ class ImageAgent(Agent):
             if self.verbose:
                 self.console.print(f"[cyan]Editing image with {model_name}...[/cyan]")
             
-            response = litellm.image_edit(**config)
+            response = ImageResult(litellm.image_edit(**config))
             
             if self.verbose:
                 self.console.print(f"[green]✓ Image edited successfully[/green]")
@@ -338,7 +339,7 @@ class ImageAgent(Agent):
         n: int = 1,
         size: Optional[str] = None,
         **kwargs
-    ) -> Dict[str, Any]:
+    ) -> ImageResult:
         """Offload the synchronous editing path to keep the event loop responsive."""
         return await asyncio.to_thread(self.edit, image, prompt, mask, n, size, **kwargs)
 
@@ -348,7 +349,7 @@ class ImageAgent(Agent):
         n: int = 1,
         size: Optional[str] = None,
         **kwargs
-    ) -> Dict[str, Any]:
+    ) -> ImageResult:
         """
         Generate variations of an existing image.
         
@@ -390,7 +391,7 @@ class ImageAgent(Agent):
             if self.verbose:
                 self.console.print(f"[cyan]Generating variations with {model_name}...[/cyan]")
             
-            response = litellm.image_variation(**config)
+            response = ImageResult(litellm.image_variation(**config))
             
             if self.verbose:
                 self.console.print(f"[green]✓ Variations generated successfully[/green]")
@@ -408,6 +409,6 @@ class ImageAgent(Agent):
         n: int = 1,
         size: Optional[str] = None,
         **kwargs
-    ) -> Dict[str, Any]:
+    ) -> ImageResult:
         """Offload the synchronous variation path to keep the event loop responsive."""
         return await asyncio.to_thread(self.variation, image, n, size, **kwargs)

src/praisonai-agents/praisonaiagents/workflows/workflows.py

@@ -3005,17 +3005,16 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
             result = agent.generate(prompt)
             # Extract URL from result and store in variables for next step
             if hasattr(result, 'data') and result.data:
-                image_url = result.data[0].url if hasattr(result.data[0], 'url') else None
+                image_url = getattr(result.data[0], 'url', None)
                 if image_url:
                     variables['_last_image_url'] = image_url
-                    # Save to file if output path specified
-                    if output_path:
-                        try:
-                            import urllib.request
-                            urllib.request.urlretrieve(image_url, output_path)
-                            variables['_last_image_file'] = output_path
-                        except Exception as e:
-                            logger.warning(f"Failed to save image to {output_path}: {e}")
+            # Save to file if output path specified (handles b64 and URL responses)
+            if output_path:
+                try:
+                    result.save(output_path)
+                    variables['_last_image_file'] = output_path
+                except Exception as e:
+                    logger.warning(f"Failed to save image to {output_path}: {e}")
             return result
         
         # OCRAgent - extract with retry logic for external API failures

src/praisonai-agents/tests/unit/test_image_agent_save.py

@@ -0,0 +1,158 @@
+"""ImageAgent results expose .save() while keeping the raw response API (#6260)."""
+
+import base64
+import sys
+import types
+
+import pytest
+
+from praisonaiagents import ImageAgent
+
+GIF = b"GIF89a\x01\x00\x01\x00fake"
+SVG = b'<svg xmlns="http://www.w3.org/2000/svg"></svg>'
+
+
+def _response(*payloads):
+    return types.SimpleNamespace(
+        data=[types.SimpleNamespace(b64_json=base64.b64encode(p).decode(), url=None) for p in payloads],
+        usage={"total_tokens": 5},
+    )
+
+
+@pytest.fixture
+def agent(monkeypatch):
+    agent = ImageAgent(llm="openai/gpt-image-2.6-flare", verbose=False)
+    sent = {}
+
+    def fake(name, payloads):
+        def call(**kwargs):
+            sent[name] = kwargs
+            return _response(*payloads)
+        return call
+
+    monkeypatch.setattr(agent, "_litellm", fake("generate", [GIF]))
+    monkeypatch.setitem(sys.modules, "litellm", types.SimpleNamespace(
+        image_edit=fake("edit", [SVG, GIF]),
+        image_variation=fake("variation", [GIF]),
+        get_llm_provider=lambda model: (model, "openai", None, None),
+    ))
+    agent.sent = sent
+    return agent
+
+
+def test_generate_save_writes_decoded_bytes(agent, tmp_path):
+    out = agent.generate("A cat waving", output_format="gif", frame_count=8).save(str(tmp_path / "cat.gif"))
+    assert open(out, "rb").read() == GIF
+    assert agent.sent["generate"]["output_format"] == "gif"
+    assert agent.sent["generate"]["frame_count"] == 8
+
+
+def test_edit_save_index_and_raw_access(agent, tmp_path):
+    src = tmp_path / "photo.png"
+    src.write_bytes(b"png")
+    result = agent.edit(str(src), "make the cup blue", n=2)
+    assert agent.sent["edit"]["n"] == 2
+    assert len(result.images) == 2
+    assert result.usage == {"total_tokens": 5}
+    assert base64.b64decode(result.data[1].b64_json) == GIF
+    result.save(str(tmp_path / "a.svg"))
+    assert (tmp_path / "a.svg").read_bytes() == SVG
+
+
+def test_variation_and_dict_response(agent, tmp_path, monkeypatch):
+    src = tmp_path / "photo.png"
+    src.write_bytes(b"png")
+    assert agent.variation(str(src)).save(str(tmp_path / "v.gif"))
+    monkeypatch.setattr(agent, "_litellm", lambda **kw: {"data": [{"b64_json": base64.b64encode(SVG).decode()}]})
+    result = agent.generate("logo")
+    assert result["data"][0]["b64_json"]
+    result.save(str(tmp_path / "logo.svg"))
+    assert (tmp_path / "logo.svg").read_bytes() == SVG
+
+
+def test_save_without_image_data_raises(agent, tmp_path, monkeypatch):
+    monkeypatch.setattr(agent, "_litellm", lambda **kw: {"data": []})
+    with pytest.raises(ValueError):
+        agent.generate("nothing").save(str(tmp_path / "x.png"))
+
+
+def test_save_nonzero_index(agent, tmp_path):
+    src = tmp_path / "photo.png"
+    src.write_bytes(b"png")
+    result = agent.edit(str(src), "two takes", n=2)
+    assert result.save(str(tmp_path / "b.gif"), index=1) == str(tmp_path / "b.gif")
+    assert (tmp_path / "b.gif").read_bytes() == GIF
+
+
+class _FakeHTTP:
+    def __init__(self, body):
+        self.body = body
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+    def read(self, n=-1):
+        return self.body if n < 0 else self.body[:n]
+
+
+def _url_response(url):
+    return {"data": [{"b64_json": None, "url": url}]}
+
+
+def test_save_downloads_url(agent, tmp_path, monkeypatch):
+    import urllib.request
+
+    seen = {}
+
+    def fake_urlopen(url, timeout=None):
+        seen.update(url=url, timeout=timeout)
+        return _FakeHTTP(GIF)
+
+    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
+    monkeypatch.setattr(agent, "_litellm", lambda **kw: _url_response("https://example.com/i.gif"))
+    agent.generate("cat").save(str(tmp_path / "u.gif"))
+    assert (tmp_path / "u.gif").read_bytes() == GIF
+    assert seen == {"url": "https://example.com/i.gif", "timeout": 30}
+
+
+def test_save_rejects_non_http_url(agent, tmp_path, monkeypatch):
+    monkeypatch.setattr(agent, "_litellm", lambda **kw: _url_response("file:///etc/passwd"))
+    with pytest.raises(ValueError, match="http"):
+        agent.generate("x").save(str(tmp_path / "x.png"))
+    assert not (tmp_path / "x.png").exists()
+
+
+def test_save_rejects_oversized_download(agent, tmp_path, monkeypatch):
+    import urllib.request
+    from praisonaiagents.agent import _image_result
+
+    monkeypatch.setattr(_image_result, "_MAX_DOWNLOAD_BYTES", 4)
+    monkeypatch.setattr(urllib.request, "urlopen", lambda url, timeout=None: _FakeHTTP(b"123456"))
+    monkeypatch.setattr(agent, "_litellm", lambda **kw: _url_response("https://example.com/big.png"))
+    with pytest.raises(ValueError, match="maximum"):
+        agent.generate("x").save(str(tmp_path / "big.png"))
+    assert not (tmp_path / "big.png").exists()
+
+
+def test_container_protocol_passes_through(agent):
+    from litellm.types.utils import ImageResponse
+
+    raw = ImageResponse(data=[{"b64_json": base64.b64encode(GIF).decode()}])
+    agent._litellm = lambda **kw: raw
+    result = agent.generate("cat")
+    assert "data" in result
+    assert dict(result).keys() == dict(raw).keys()
+
+
+def test_workflow_output_saves_b64_image(agent, tmp_path):
+    from praisonaiagents.workflows.workflows import AgentFlow
+
+    flow = object.__new__(AgentFlow)
+    out = str(tmp_path / "wf.gif")
+    variables = {"output": out}
+    flow._execute_specialized_agent(agent, "ImageAgent", "a cat", None, variables, False)
+    assert (tmp_path / "wf.gif").read_bytes() == GIF
+    assert variables["_last_image_file"] == out