[Bug Fix] LLM Translation - Allow using dynamic api_key for image generation requests (#14007)

* fix - allow using dyanmic api key for img gen

* test_aiml_image_generation_with_dynamic_api_key

litellm/images/main.py

@@ -1,7 +1,7 @@
 import asyncio
 import contextvars
 from functools import partial
-from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast, overload, List
+from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, cast, overload
 
 import httpx
 
@@ -347,6 +347,7 @@ def image_generation(  # noqa: PLR0915
                 raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
             
             return llm_http_handler.image_generation_handler(
+                api_key=api_key,
                 model=model,
                 prompt=prompt,
                 image_generation_provider_config=image_generation_config,

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2681,6 +2681,7 @@ class BaseLLMHTTPHandler:
         _is_async: bool = False,
         fake_stream: bool = False,
         litellm_metadata: Optional[Dict[str, Any]] = None,
+        api_key: Optional[str] = None,
     ) -> Union[
         ImageResponse,
         Coroutine[Any, Any, ImageResponse],
@@ -2705,6 +2706,7 @@ class BaseLLMHTTPHandler:
                 client=client if isinstance(client, AsyncHTTPHandler) else None,
                 fake_stream=fake_stream,
                 litellm_metadata=litellm_metadata,
+                api_key=api_key,
             )
 
         if client is None or not isinstance(client, HTTPHandler):
@@ -2715,7 +2717,7 @@ class BaseLLMHTTPHandler:
             sync_httpx_client = client
 
         headers = image_generation_provider_config.validate_environment(
-            api_key=litellm_params.get("api_key", None),
+            api_key=api_key,
             headers=image_generation_optional_request_params.get("extra_headers", {})
             or {},
             model=model,
@@ -2798,6 +2800,7 @@ class BaseLLMHTTPHandler:
         client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
         fake_stream: bool = False,
         litellm_metadata: Optional[Dict[str, Any]] = None,
+        api_key: Optional[str] = None,
     ) -> ImageResponse:
         """
         Async version of the image generation handler.
@@ -2812,7 +2815,7 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         headers = image_generation_provider_config.validate_environment(
-            api_key=litellm_params.get("api_key", None),
+            api_key=api_key,
             headers=image_generation_optional_request_params.get("extra_headers", {})
             or {},
             model=model,

tests/image_gen_tests/test_image_generation.py

@@ -330,3 +330,78 @@ async def test_gpt_image_1_with_input_fidelity():
         assert captured_kwargs["quality"] == "medium"
         assert captured_kwargs["size"] == "1024x1024"
 
+
+@pytest.mark.asyncio
+async def test_aiml_image_generation_with_dynamic_api_key():
+    """
+    Test that when api_key is passed as a dynamic parameter to aimage_generation,
+    it gets properly used for AIML provider authentication instead of falling back
+    to environment variables.
+    
+    This test validates the fix for ensuring dynamic API keys are respected
+    when making image generation requests to the AIML provider.
+    """
+    from unittest.mock import AsyncMock, patch, MagicMock
+    import httpx
+    
+    # Mock AIML response
+    mock_aiml_response = {
+        "created": 1703658209,
+        "data": [
+            {
+                "url": "https://example.com/generated_image.png"
+            }
+        ]
+    }
+    
+    # Track captured arguments
+    captured_headers = None
+    captured_url = None
+    captured_json_data = None
+    
+    def capture_post_call(*args, **kwargs):
+        nonlocal captured_headers, captured_url, captured_json_data
+        captured_url = kwargs.get('url') or (args[0] if args else None)
+        captured_headers = kwargs.get('headers', {})
+        captured_json_data = kwargs.get('json', {})
+        
+        # Create a mock response
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = mock_aiml_response
+        mock_response.text = json.dumps(mock_aiml_response)
+        return mock_response
+    
+    # Mock the HTTP client that actually makes the request (sync version for image generation)
+    with patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post') as mock_post:
+        mock_post.side_effect = capture_post_call
+        
+        # Test with dynamic api_key
+        test_api_key = "test-dynamic-api-key-12345"
+        
+        response = await litellm.aimage_generation(
+            prompt="A cute baby sea otter",
+            model="aiml/flux-pro/v1.1",
+            api_key=test_api_key,  # This should be used instead of env vars
+        )
+
+        # Validate the response (mocked response processing might not populate data correctly)
+        assert response is not None
+        
+        # The most important validations: API key and endpoint usage
+        # These prove that the dynamic API key was properly used
+        assert captured_headers is not None
+        assert "Authorization" in captured_headers
+        assert captured_headers["Authorization"] == f"Bearer {test_api_key}"
+        print("TESTCAPTURED HEADERS", captured_headers)
+        # Validate the correct AIML endpoint was called
+        assert captured_url is not None
+        assert "api.aimlapi.com" in captured_url
+        assert "/v1/images/generations" in captured_url
+        
+        # Validate the request data
+        assert captured_json_data is not None
+        assert captured_json_data["prompt"] == "A cute baby sea otter"
+        assert captured_json_data["model"] == "flux-pro/v1.1"
+
+