Added LiteLLM to the stack

2025-08-18 09:40:50 +00:00
parent 0648c1968c
commit d220b04e32
2682 changed files with 533609 additions and 1 deletions
--- a/Development/litellm/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py
+++ b/Development/litellm/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py
@@ -0,0 +1,75 @@
+import pytest
+from litellm.llms.bedrock.image.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig
+from litellm.types.utils import ImageResponse
+
+def test_transform_request_body_text_to_image():
+    params = {
+        "imageGenerationConfig": {
+            "cfgScale": 7,
+            "seed": 42,
+            "quality": "standard",
+            "width": 512,
+            "height": 512,
+            "numberOfImages": 1,
+            "textToImageParams": {
+                "negativeText": "blurry"
+            }
+        }
+    }
+    req = AmazonNovaCanvasConfig.transform_request_body("cat", params.copy())
+    assert isinstance(req, dict)
+    assert "textToImageParams" in req
+    assert req["textToImageParams"]["text"] == "cat"
+    assert req["imageGenerationConfig"]["width"] == 512
+
+def test_transform_request_body_color_guided():
+    params = {
+        "taskType": "COLOR_GUIDED_GENERATION",
+        "imageGenerationConfig": {
+            "cfgScale": 7,
+            "seed": 42,
+            "quality": "standard",
+            "width": 512,
+            "height": 512,
+            "numberOfImages": 1,
+            "colorGuidedGenerationParams": {
+                "colors": ["#FFFFFF"],
+                "referenceImage": "img",
+                "negativeText": "blurry"
+            }
+        }
+    }
+    req = AmazonNovaCanvasConfig.transform_request_body("cat", params.copy())
+    assert "colorGuidedGenerationParams" in req
+    assert req["colorGuidedGenerationParams"]["text"] == "cat"
+    assert req["imageGenerationConfig"]["width"] == 512
+
+def test_transform_request_body_inpainting():
+    params = {
+        "taskType": "INPAINTING",
+        "imageGenerationConfig": {
+            "cfgScale": 7,
+            "seed": 42,
+            "quality": "standard",
+            "width": 512,
+            "height": 512,
+            "numberOfImages": 1,
+            "inpaintingParams": {
+                "maskImage": "mask",
+                "inputImage": "input",
+                "negativeText": "blurry"
+            }
+        }
+    }
+    req = AmazonNovaCanvasConfig.transform_request_body("cat", params.copy())
+    assert "inpaintingParams" in req
+    assert req["inpaintingParams"]["text"] == "cat"
+    assert req["imageGenerationConfig"]["width"] == 512
+
+def test_transform_response_dict_to_openai_response():
+    response_dict = {"images": ["b64img1", "b64img2"]}
+    model_response = ImageResponse()
+    result = AmazonNovaCanvasConfig.transform_response_dict_to_openai_response(model_response, response_dict)
+    assert hasattr(result, "data")
+    assert len(result.data) == 2
+    assert result.data[0].b64_json == "b64img1"
--- a/Development/litellm/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py
+++ b/Development/litellm/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py
@@ -0,0 +1,20 @@
+import json
+import os
+import sys
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+from unittest.mock import MagicMock, patch
+
+from litellm.llms.bedrock.image.amazon_stability3_transformation import (
+    AmazonStability3Config,
+)
+
+
+def test_stability_image_core_is_v3_model():
+    model = "stability.stable-image-core-v1:1"
+    assert AmazonStability3Config._is_stability_3_model(model)
--- a/Development/litellm/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py
+++ b/Development/litellm/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py
@@ -0,0 +1,130 @@
+import json
+import os
+import sys
+from unittest.mock import Mock, patch
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../../.."))  # Adds the parent directory to the system path
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
+
+# Mock response for Bedrock image generation
+mock_image_response = {
+    "images": ["base64_encoded_image_data"],
+    "error": None
+}
+
+class TestBedrockImageGeneration:
+    def test_image_generation_with_api_key_bearer_token(self):
+        """Test image generation with bearer token authentication"""
+        litellm.set_verbose = True
+        test_api_key = "test-bearer-token-12345"
+        model = "bedrock/stability.sd3-large-v1:0"
+        prompt = "A cute baby sea otter"
+
+        with patch("litellm.llms.bedrock.image.image_handler.BedrockImageGeneration.image_generation") as mock_bedrock_image_gen:
+            # Setup mock response
+            mock_image_response_obj = litellm.ImageResponse()
+            mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}]
+            mock_bedrock_image_gen.return_value = mock_image_response_obj
+
+            response = litellm.image_generation(
+                model=model,
+                prompt=prompt,
+                aws_region_name="us-west-2",
+                api_key=test_api_key
+            )
+
+            assert response is not None
+            assert len(response.data) > 0
+            
+            mock_bedrock_image_gen.assert_called_once()
+            for call in mock_bedrock_image_gen.call_args_list:
+                if "headers" in call.kwargs:
+                    headers = call.kwargs["headers"]
+                    if "Authorization" in headers and headers["Authorization"] == f"Bearer {test_api_key}":
+                        break
+
+    def test_image_generation_with_env_variable_bearer_token(self, monkeypatch):
+        """Test image generation with bearer token from environment variable"""
+        litellm.set_verbose = True
+        test_api_key = "env-bearer-token-12345"
+        model = "bedrock/stability.sd3-large-v1:0"
+        prompt = "A cute baby sea otter"
+        
+        # Mock the environment variable
+        with patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": test_api_key}), \
+             patch("litellm.llms.bedrock.image.image_handler.BedrockImageGeneration.image_generation") as mock_bedrock_image_gen:
+            
+            mock_image_response_obj = litellm.ImageResponse()
+            mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}]
+            mock_bedrock_image_gen.return_value = mock_image_response_obj
+
+            response = litellm.image_generation(
+                model=model,
+                prompt=prompt,
+                aws_region_name="us-west-2"
+            )
+
+            assert response is not None
+            assert len(response.data) > 0
+            
+            mock_bedrock_image_gen.assert_called_once()
+            for call in mock_bedrock_image_gen.call_args_list:
+                if "headers" in call.kwargs:
+                    headers = call.kwargs["headers"]
+                    if "Authorization" in headers and headers["Authorization"] == f"Bearer {test_api_key}":
+                        break
+
+    @pytest.mark.asyncio
+    async def test_async_image_generation_with_bearer_token(self):
+        """Test async image generation with bearer token authentication"""
+        litellm.set_verbose = True
+        test_api_key = "async-bearer-token-12345"
+        model = "bedrock/stability.sd3-large-v1:0"
+        prompt = "A cute baby sea otter"
+
+        with patch("litellm.llms.bedrock.image.image_handler.BedrockImageGeneration.async_image_generation") as mock_async_bedrock_image_gen:
+            mock_image_response_obj = litellm.ImageResponse()
+            mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}]
+            mock_async_bedrock_image_gen.return_value = mock_image_response_obj
+
+            # Call async image generation with api_key parameter
+            response = await litellm.aimage_generation(
+                model=model,
+                prompt=prompt,
+                aws_region_name="us-west-2",
+                api_key=test_api_key
+            )
+
+            assert response is not None
+            assert len(response.data) > 0
+            
+            mock_async_bedrock_image_gen.assert_called_once()
+            for call in mock_async_bedrock_image_gen.call_args_list:
+                if "headers" in call.kwargs:
+                    headers = call.kwargs["headers"]
+                    if "Authorization" in headers and headers["Authorization"] == f"Bearer {test_api_key}":
+                        break
+
+    def test_image_generation_with_sigv4(self):
+        """Test image generation falls back to SigV4 auth when no bearer token is provided"""
+        litellm.set_verbose = True
+        model = "bedrock/stability.sd3-large-v1:0"
+        prompt = "A cute baby sea otter"
+
+        with patch("litellm.llms.bedrock.image.image_handler.BedrockImageGeneration.image_generation") as mock_bedrock_image_gen:
+            mock_image_response_obj = litellm.ImageResponse()
+            mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}]
+            mock_bedrock_image_gen.return_value = mock_image_response_obj
+
+            response = litellm.image_generation(
+                model=model,
+                prompt=prompt,
+                aws_region_name="us-west-2"
+            )
+            
+            assert response is not None
+            assert len(response.data) > 0
+            mock_bedrock_image_gen.assert_called_once()