fix(azure): return response headers for sync embedding calls

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:

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(

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