feat: Added support for default_headers for azure_openai. (#2917)

src/phoenix/experimental/evals/models/openai.py

@@ -8,6 +8,7 @@
     Callable,
     Dict,
     List,
+    Mapping,
     Optional,
     Tuple,
     Union,
@@ -108,6 +109,8 @@ class OpenAIModel(BaseEvalModel):
     azure_deployment: Optional[str] = field(default=None)
     azure_ad_token: Optional[str] = field(default=None)
     azure_ad_token_provider: Optional[Callable[[], str]] = field(default=None)
+    default_headers: Optional[Mapping[str, str]] = field(default=None)
+    """Default headers required by AzureOpenAI"""
 
     # Deprecated fields
     model_name: Optional[str] = field(default=None)
@@ -193,6 +196,7 @@ def _init_open_ai(self) -> None:
                 azure_ad_token_provider=azure_options.azure_ad_token_provider,
                 api_key=self.api_key,
                 organization=self.organization,
+                default_headers=self.default_headers,
             )
             self._async_client = self._openai.AsyncAzureOpenAI(
                 azure_endpoint=azure_options.azure_endpoint,
@@ -202,6 +206,7 @@ def _init_open_ai(self) -> None:
                 azure_ad_token_provider=azure_options.azure_ad_token_provider,
                 api_key=self.api_key,
                 organization=self.organization,
+                default_headers=self.default_headers,
             )
             # return early since we don't need to check the model
             return

tests/experimental/evals/models/test_openai.py

@@ -30,6 +30,31 @@ def test_azure_openai_model(monkeypatch):
     assert isinstance(model._client, AzureOpenAI)
 
 
+def test_azure_openai_model_added_custom_header(monkeypatch):
+    monkeypatch.setenv(OPENAI_API_KEY_ENVVAR_NAME, "sk-0123456789")
+    with patch.object(OpenAIModel, "_init_tiktoken", return_value=None):
+        header_key = "header"
+        header_value = "my-example-header-value"
+        default_headers = {header_key: header_value}
+        model = OpenAIModel(
+            model="gpt-4-turbo-preview",
+            api_version="2023-07-01-preview",
+            azure_endpoint="https://example-endpoint.openai.azure.com",
+            default_headers=default_headers,
+        )
+
+    assert isinstance(model._client, AzureOpenAI)
+    # check if custom header is added to headers
+    assert (
+        header_key in model._client.default_headers
+        and model._client.default_headers.get(header_key) == header_value
+    )
+    assert (
+        header_key in model._async_client.default_headers
+        and model._async_client.default_headers.get(header_key) == header_value
+    )
+
+
 def test_azure_fails_when_missing_options(monkeypatch):
     monkeypatch.setenv(OPENAI_API_KEY_ENVVAR_NAME, "sk-0123456789")
     # Test missing api_version