From f35ce02475a0b341c21531ded111b4bb1993c1c7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Aug 2025 10:51:24 -0700 Subject: [PATCH] [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 | 3 +- litellm/llms/custom_httpx/llm_http_handler.py | 7 +- .../image_gen_tests/test_image_generation.py | 75 +++++++++++++++++++ 3 files changed, 82 insertions(+), 3 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 4e4dfa752f..4993a48c72 100644 --- a/litellm/images/main.py +++ b/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, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d404077a5b..2faea53901 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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, diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 043a0a88fd..9aa8152716 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/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" + +