fix(caching_handler.py): fix embedding str caching result (#10700)

* fix(caching_handler.py): fix embedding str caching result

Fixes issue where str caching results were not being correctly assembled on str input

* feat(azure/image_generation): Support dropping response_format for azure gpt-image-1

Fixes LIT-118

* test(test_utils.py): add unit testing

* test: rename file to avoid testing conflict

litellm/caching/caching_handler.py

@@ -515,9 +515,11 @@ class LLMCachingHandler:
             )
         )
         cached_result: Optional[Any] = None
-        if call_type == CallTypes.aembedding.value and isinstance(
-            new_kwargs["input"], list
-        ):
+        if call_type == CallTypes.aembedding.value:
+            if isinstance(new_kwargs["input"], str):
+                new_kwargs["input"] = [new_kwargs["input"]]
+            elif not isinstance(new_kwargs["input"], list):
+                raise ValueError("input must be a string or a list")
             tasks = []
             for idx, i in enumerate(new_kwargs["input"]):
                 preset_cache_key = litellm.cache.get_cache_key(
@@ -700,6 +702,7 @@ class LLMCachingHandler:
         Raises:
             None
         """
+
         if litellm.cache is None:
             return
 

litellm/llms/azure/image_generation/__init__.py

@@ -0,0 +1,22 @@
+from litellm.llms.base_llm.image_generation.transformation import (
+    BaseImageGenerationConfig,
+)
+
+from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig
+from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig
+from .gpt_transformation import AzureGPTImageGenerationConfig
+
+__all__ = [
+    "AzureDallE2ImageGenerationConfig",
+    "AzureDallE3ImageGenerationConfig",
+    "AzureGPTImageGenerationConfig",
+]
+
+
+def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig:
+    if model.startswith("dall-e-2") or model == "":  # empty model is dall-e-2
+        return AzureDallE2ImageGenerationConfig()
+    elif model.startswith("dall-e-3"):
+        return AzureDallE3ImageGenerationConfig()
+    else:
+        return AzureGPTImageGenerationConfig()

litellm/llms/azure/image_generation/dall_e_2_transformation.py

@@ -0,0 +1,9 @@
+from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig
+
+
+class AzureDallE2ImageGenerationConfig(DallE2ImageGenerationConfig):
+    """
+    Azure dall-e-2 image generation config
+    """
+
+    pass

litellm/llms/azure/image_generation/dall_e_3_transformation.py

@@ -0,0 +1,9 @@
+from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig
+
+
+class AzureDallE3ImageGenerationConfig(DallE3ImageGenerationConfig):
+    """
+    Azure dall-e-3 image generation config
+    """
+
+    pass

litellm/llms/azure/image_generation/gpt_transformation.py

@@ -0,0 +1,9 @@
+from litellm.llms.openai.image_generation import GPTImageGenerationConfig
+
+
+class AzureGPTImageGenerationConfig(GPTImageGenerationConfig):
+    """
+    Azure gpt-image-1 image generation config
+    """
+
+    pass

litellm/llms/openai/image_generation/__init__.py

@@ -14,7 +14,7 @@ __all__ = [
 
 
 def get_openai_image_generation_config(model: str) -> BaseImageGenerationConfig:
-    if model.startswith("dall-e-2"):
+    if model.startswith("dall-e-2") or model == "":  # empty model is dall-e-2
         return DallE2ImageGenerationConfig()
     elif model.startswith("dall-e-3"):
         return DallE3ImageGenerationConfig()

litellm/utils.py

@@ -6569,6 +6569,12 @@ class ProviderConfigManager:
             )
 
             return get_openai_image_generation_config(model)
+        elif LlmProviders.AZURE == provider:
+            from litellm.llms.azure.image_generation import (
+                get_azure_image_generation_config,
+            )
+
+            return get_azure_image_generation_config(model)
         return None
 
 

tests/litellm/test_utils.py

@@ -0,0 +1,30 @@
+import json
+import os
+import sys
+
+import httpx
+import pytest
+import respx
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+from litellm.utils import get_optional_params_image_gen
+
+
+def test_get_optional_params_image_gen():
+    from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig
+
+    provider_config = AzureGPTImageGenerationConfig()
+    optional_params = get_optional_params_image_gen(
+        model="gpt-image-1",
+        response_format="b64_json",
+        n=3,
+        custom_llm_provider="azure",
+        drop_params=True,
+        provider_config=provider_config,
+    )
+    assert optional_params is not None
+    assert "response_format" not in optional_params
+    assert optional_params["n"] == 3

tests/local_testing/test_caching.py

@@ -483,6 +483,23 @@ async def test_embedding_caching_individual_items_and_then_list():
     )
     assert embedding4.usage.prompt_tokens > embedding3.usage.prompt_tokens
 
+@pytest.mark.asyncio
+async def test_embedding_caching_individual_items():
+    litellm.cache = Cache()
+    text_to_embed = "hello"
+    embedding1 = await aembedding(
+        model="text-embedding-ada-002", input=text_to_embed, caching=True
+    )
+
+    await asyncio.sleep(1)
+
+    embedding3 = await aembedding(
+        model="text-embedding-ada-002", input=text_to_embed, caching=True
+    )
+    final_prompt_tokens = embedding3.usage.prompt_tokens
+    assert embedding3["data"][0]["embedding"] == embedding1["data"][0]["embedding"]
+    assert embedding3._hidden_params["cache_hit"] == True
+    assert embedding3.usage.prompt_tokens != 0 
 
 
 def test_embedding_caching_azure():