diff --git a/litellm/llms/AzureOpenAI/azure.py b/litellm/llms/AzureOpenAI/azure.py index 39dea14e25f9..c9bbf2516e0d 100644 --- a/litellm/llms/AzureOpenAI/azure.py +++ b/litellm/llms/AzureOpenAI/azure.py @@ -817,6 +817,7 @@ async def aembedding( logging_obj: LiteLLMLoggingObj, api_key: Optional[str] = None, client: Optional[AsyncAzureOpenAI] = None, + extra_headers: Optional[dict] = None, timeout=None, ): response = None @@ -826,7 +827,7 @@ async def aembedding( else: openai_aclient = client raw_response = await openai_aclient.embeddings.with_raw_response.create( - **data, timeout=timeout + **data, timeout=timeout, extra_headers=extra_headers, ) headers = dict(raw_response.headers) response = raw_response.parse() @@ -867,6 +868,7 @@ def embedding( optional_params: dict, api_key: Optional[str] = None, azure_ad_token: Optional[str] = None, + extra_headers: Optional[dict] = None, client=None, aembedding=None, ) -> litellm.EmbeddingResponse: @@ -923,13 +925,14 @@ def embedding( azure_client_params=azure_client_params, timeout=timeout, client=client, + extra_headers=extra_headers, ) if client is None: azure_client = AzureOpenAI(**azure_client_params) # type: ignore else: azure_client = client ## COMPLETION CALL - raw_response = azure_client.embeddings.with_raw_response.create(**data, timeout=timeout) # type: ignore + raw_response = azure_client.embeddings.with_raw_response.create(**data, timeout=timeout, extra_headers=extra_headers) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() ## LOGGING diff --git a/litellm/main.py b/litellm/main.py index 8334f35d7bf5..8e9be7884676 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3387,6 +3387,7 @@ def embedding( # noqa: PLR0915 optional_params=optional_params, client=client, aembedding=aembedding, + extra_headers=extra_headers, ) elif ( model in litellm.open_ai_embedding_models