Add Support for Customizing default_headers in Azure OpenAI (#1925)

tests/embeddings/test_azure_openai_embeddings.py

@@ -1,7 +1,9 @@
-import pytest
 from unittest.mock import Mock, patch
-from mem0.embeddings.azure_openai import AzureOpenAIEmbedding
+
+import pytest
+
 from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.azure_openai import AzureOpenAIEmbedding
 
 
 @pytest.fixture
@@ -29,18 +31,26 @@ def test_embed_text(mock_openai_client):
     assert embedding == [0.1, 0.2, 0.3]
 
 
-def test_embed_text_with_newlines(mock_openai_client):
-    config = BaseEmbedderConfig(model="text-embedding-ada-002")
-    embedder = AzureOpenAIEmbedding(config)
-
-    mock_embedding_response = Mock()
-    mock_embedding_response.data = [Mock(embedding=[0.4, 0.5, 0.6])]
-    mock_openai_client.embeddings.create.return_value = mock_embedding_response
-
-    text = "Hello,\nthis is a test\nwith newlines."
-    embedding = embedder.embed(text)
-
-    mock_openai_client.embeddings.create.assert_called_once_with(
-        input=["Hello, this is a test with newlines."], model="text-embedding-ada-002"
+@pytest.mark.parametrize(
+    "default_headers, expected_header",
+    [
+        (None, None),
+        ({"Test": "test_value"}, "test_value"),
+        ({}, None)
+    ],
+)
+def test_embed_text_with_default_headers(default_headers, expected_header):
+    config = BaseEmbedderConfig(
+        model="text-embedding-ada-002",
+        azure_kwargs={
+            "api_key": "test",
+            "api_version": "test_version",
+            "azure_endpoint": "test_endpoint",
+            "azuer_deployment": "test_deployment",
+            "default_headers": default_headers
+        }
     )
-    assert embedding == [0.4, 0.5, 0.6]
+    embedder = AzureOpenAIEmbedding(config)
+    assert embedder.client.api_key == "test"
+    assert embedder.client._api_version == "test_version"
+    assert embedder.client.default_headers.get("Test") == expected_header