From 6f32189093976dc71c8d8a1669d87992bc0aa64d Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 9 May 2025 23:37:02 -0700 Subject: [PATCH] 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 | 9 ++++-- .../llms/azure/image_generation/__init__.py | 22 ++++++++++++++ .../dall_e_2_transformation.py | 9 ++++++ .../dall_e_3_transformation.py | 9 ++++++ .../image_generation/gpt_transformation.py | 9 ++++++ .../llms/openai/image_generation/__init__.py | 2 +- litellm/utils.py | 6 ++++ ...s.py => test_litellm_proxy_types_utils.py} | 0 tests/litellm/test_utils.py | 30 +++++++++++++++++++ tests/local_testing/test_caching.py | 17 +++++++++++ 10 files changed, 109 insertions(+), 4 deletions(-) create mode 100644 litellm/llms/azure/image_generation/__init__.py create mode 100644 litellm/llms/azure/image_generation/dall_e_2_transformation.py create mode 100644 litellm/llms/azure/image_generation/dall_e_3_transformation.py create mode 100644 litellm/llms/azure/image_generation/gpt_transformation.py rename tests/litellm/proxy/types_utils/{test_utils.py => test_litellm_proxy_types_utils.py} (100%) create mode 100644 tests/litellm/test_utils.py diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index afe96c7f46..6b41c1ff40 100644 --- a/litellm/caching/caching_handler.py +++ b/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 diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py new file mode 100644 index 0000000000..d51a9229bd --- /dev/null +++ b/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() diff --git a/litellm/llms/azure/image_generation/dall_e_2_transformation.py b/litellm/llms/azure/image_generation/dall_e_2_transformation.py new file mode 100644 index 0000000000..3fe702f57f --- /dev/null +++ b/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 diff --git a/litellm/llms/azure/image_generation/dall_e_3_transformation.py b/litellm/llms/azure/image_generation/dall_e_3_transformation.py new file mode 100644 index 0000000000..5e0bfcd108 --- /dev/null +++ b/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 diff --git a/litellm/llms/azure/image_generation/gpt_transformation.py b/litellm/llms/azure/image_generation/gpt_transformation.py new file mode 100644 index 0000000000..1f5f65f693 --- /dev/null +++ b/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 diff --git a/litellm/llms/openai/image_generation/__init__.py b/litellm/llms/openai/image_generation/__init__.py index 7e6ca6e7ae..eb2a0576b6 100644 --- a/litellm/llms/openai/image_generation/__init__.py +++ b/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() diff --git a/litellm/utils.py b/litellm/utils.py index 376eb7899f..45da62631a 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/litellm/proxy/types_utils/test_utils.py b/tests/litellm/proxy/types_utils/test_litellm_proxy_types_utils.py similarity index 100% rename from tests/litellm/proxy/types_utils/test_utils.py rename to tests/litellm/proxy/types_utils/test_litellm_proxy_types_utils.py diff --git a/tests/litellm/test_utils.py b/tests/litellm/test_utils.py new file mode 100644 index 0000000000..6d76439b94 --- /dev/null +++ b/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 diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 0acfecfe33..9de9ff3178 100644 --- a/tests/local_testing/test_caching.py +++ b/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():