diff --git a/litellm/llms/AzureOpenAI/azure.py b/litellm/llms/AzureOpenAI/azure.py index d7a528c737..1a9f1b769d 100644 --- a/litellm/llms/AzureOpenAI/azure.py +++ b/litellm/llms/AzureOpenAI/azure.py @@ -1053,7 +1053,7 @@ class AzureChatCompletion(BaseLLM): response_object=stringified_response, model_response_object=model_response, hidden_params={"headers": headers}, - _response_headers=headers, + _response_headers=process_azure_headers(headers), response_type="embedding", ) except Exception as e: @@ -1142,6 +1142,7 @@ class AzureChatCompletion(BaseLLM): azure_client = client ## COMPLETION CALL raw_response = azure_client.embeddings.with_raw_response.create(**data, timeout=timeout) # type: ignore + headers = dict(raw_response.headers) response = raw_response.parse() ## LOGGING logging_obj.post_call( @@ -1151,7 +1152,7 @@ class AzureChatCompletion(BaseLLM): original_response=response, ) - return convert_to_model_response_object(response_object=response.model_dump(), model_response_object=model_response, response_type="embedding") # type: ignore + return convert_to_model_response_object(response_object=response.model_dump(), model_response_object=model_response, response_type="embedding", _response_headers=process_azure_headers(headers)) # type: ignore except AzureOpenAIError as e: raise e except Exception as e: diff --git a/litellm/utils.py b/litellm/utils.py index 8b257cc42a..524abe3635 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1422,9 +1422,11 @@ def client(original_function): or isinstance(result, EmbeddingResponse) or isinstance(result, TranscriptionResponse) ): - result._response_ms = ( - end_time - start_time - ).total_seconds() * 1000 # return response latency in ms like openai + setattr( + result, + "_response_ms", + (end_time - start_time).total_seconds() * 1000, + ) # return response latency in ms like openai ### POST-CALL RULES ### post_call_processing( diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 143784e88d..949052ab73 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -1034,3 +1034,25 @@ async def test_hf_embedddings_with_optional_params(sync_mode): assert json_data["options"]["wait_for_model"] is True assert json_data["parameters"]["top_p"] == 10 assert json_data["parameters"]["top_k"] == 10 + + +@pytest.mark.parametrize( + "model", + [ + "text-embedding-ada-002", + "azure/azure-embedding-model", + ], +) +def test_embedding_response_ratelimit_headers(model): + response = embedding( + model=model, + input=["Hello world"], + ) + hidden_params = response._hidden_params + additional_headers = hidden_params.get("additional_headers", {}) + + print(additional_headers) + assert "x-ratelimit-remaining-requests" in additional_headers + assert int(additional_headers["x-ratelimit-remaining-requests"]) > 0 + assert "x-ratelimit-remaining-tokens" in additional_headers + assert int(additional_headers["x-ratelimit-remaining-tokens"]) > 0