From f082f7f183cc8e9e4a7ce2374e74acb009c3c3be Mon Sep 17 00:00:00 2001
From: sayanc82 <sayan.chowdhury@stability.ai>
Date: Mon, 14 Jul 2025 16:49:47 -0400
Subject: [PATCH] feat(generate_image_stability.py): Add output file and image
 prompt to return

---
 .gitignore                                    |  3 +
 src/strands_tools/generate_image_stability.py | 10 +++
 tests/test_generate_image_stability.py        | 75 +++++++++++++++++++
 3 files changed, 88 insertions(+)

diff --git a/.gitignore b/.gitignore
index 010f51c1..f8b91770 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,6 @@ errors/
 repl_state/
 venv/
 *.egg-info
+*.iml
+CLAUDE.md
+.idea/
diff --git a/src/strands_tools/generate_image_stability.py b/src/strands_tools/generate_image_stability.py
index ebcccf17..7a5a5b5a 100644
--- a/src/strands_tools/generate_image_stability.py
+++ b/src/strands_tools/generate_image_stability.py
@@ -447,6 +447,16 @@ def generate_image_stability(tool: ToolUse, **kwargs: Any) -> ToolResult:
                         f"Generated image using {model_id}. Finish reason: {finish_reason}"
                         f"{' ' + save_info if save_info else ''}"
                     ),
+                    "json": {
+                        "image_prompt": prompt,
+                        "output_filename": filename,
+                    },
+                },
+                {
+                    "image": {
+                        "format": output_format,
+                        "source": {"bytes": image_bytes},
+                    }
                 },
                 {"image": image_object},
             ],
diff --git a/tests/test_generate_image_stability.py b/tests/test_generate_image_stability.py
index 3b0e39bf..5a6ff783 100644
--- a/tests/test_generate_image_stability.py
+++ b/tests/test_generate_image_stability.py
@@ -524,3 +524,78 @@ def test_tool_spec_exists():
     properties = generate_image_stability.TOOL_SPEC["inputSchema"]["properties"]
     assert "model_id" not in properties
     assert "prompt" in properties
+
+
+def test_generate_image_stability_no_output_dir(mock_env_api_key, mock_requests):
+    """Test that output_filename is None when STABILITY_OUTPUT_DIR is not set."""
+    mock_post, mock_response = mock_requests
+
+    # Ensure STABILITY_OUTPUT_DIR is not set
+    if "STABILITY_OUTPUT_DIR" in os.environ:
+        del os.environ["STABILITY_OUTPUT_DIR"]
+
+    tool_use = {
+        "toolUseId": "test-tool-use-id",
+        "input": {
+            "prompt": "A test image",
+        },
+    }
+
+    result = generate_image_stability.generate_image_stability(tool=tool_use)
+
+    # Check that the result is successful
+    assert result["status"] == "success"
+    assert result["toolUseId"] == "test-tool-use-id"
+
+    # Check that output_filename is None in the JSON response
+    json_content = result["content"][0]["json"]
+    assert json_content["output_filename"] is None
+
+    # Check that the text response doesn't include file save info
+    text_content = result["content"][0]["text"]
+    assert "Image saved to" not in text_content
+
+
+def test_generate_image_stability_with_output_dir(mock_env_api_key, mock_requests, tmp_path):
+    """Test that output_filename is set correctly when STABILITY_OUTPUT_DIR is set."""
+    mock_post, mock_response = mock_requests
+
+    # Set STABILITY_OUTPUT_DIR to a temporary directory
+    output_dir = str(tmp_path)
+    os.environ["STABILITY_OUTPUT_DIR"] = output_dir
+
+    tool_use = {
+        "toolUseId": "test-tool-use-id",
+        "input": {
+            "prompt": "A test image for file output",
+        },
+    }
+
+    result = generate_image_stability.generate_image_stability(tool=tool_use)
+
+    # Check that the result is successful
+    assert result["status"] == "success"
+    assert result["toolUseId"] == "test-tool-use-id"
+
+    # Check that output_filename is set in the JSON response
+    json_content = result["content"][0]["json"]
+    assert json_content["output_filename"] is not None
+    assert json_content["output_filename"].startswith(output_dir)
+    assert json_content["output_filename"].endswith(".png")  # default format
+
+    # Check that the text response includes file save info
+    text_content = result["content"][0]["text"]
+    assert "Image saved to" in text_content
+    assert output_dir in text_content
+
+    # Verify the file actually exists
+    filename = json_content["output_filename"]
+    assert os.path.exists(filename)
+
+    # Verify the file content matches the mock response
+    with open(filename, "rb") as f:
+        file_content = f.read()
+    assert file_content == b"mock_image_data"
+
+    # Clean up environment variable
+    del os.environ["STABILITY_OUTPUT_DIR"]