Extending to Azure OpenAI implementation

evals/completion_fns/openai.py

@@ -6,6 +6,7 @@
 
 from evals.api import CompletionFn, CompletionResult
 from evals.base import CompletionFnSpec
+from evals.openai_client.openai_client_provider import OpenAIClientProvider
 from evals.prompt.base import (
     ChatCompletionPrompt,
     CompletionPrompt,
@@ -116,7 +117,7 @@ def __call__(
         openai_create_prompt: OpenAICreatePrompt = prompt.to_formatted_prompt()
 
         result = openai_completion_create_retrying(
-            OpenAI(api_key=self.api_key, base_url=self.api_base),
+            OpenAIClientProvider(api_key=self.api_key).get_client(),
             model=self.model,
             prompt=openai_create_prompt,
             **{**kwargs, **self.extra_options},
@@ -166,7 +167,7 @@ def __call__(
         openai_create_prompt: OpenAICreateChatPrompt = prompt.to_formatted_prompt()
 
         result = openai_chat_completion_create_retrying(
-            OpenAI(api_key=self.api_key, base_url=self.api_base),
+            OpenAIClientProvider(api_key=self.api_key).get_client(),
             model=self.model,
             messages=openai_create_prompt,
             **{**kwargs, **self.extra_options},

evals/openai_client/openai_client_provider.py

@@ -0,0 +1,28 @@
+import os
+from abc import ABC
+
+from openai import OpenAI
+from openai.lib.azure import AzureOpenAI
+
+
+class OpenAIClientProvider(ABC):
+    def __init__(self, api_key, **kwargs):
+        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        self.azure_endpoint = kwargs.pop("azure_endpoint", None) or os.environ.get(
+            "OPENAI_AZURE_ENDPOINT"
+        )
+        self.api_version = kwargs.pop("opena_api_version", None) or os.environ.get(
+            "OPENAI_API_VERSION"
+        )
+
+        self.kwargs = kwargs
+
+    def get_client(self):
+        if not self.azure_endpoint:
+            return OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+        return AzureOpenAI(
+            api_key=os.environ.get("OPENAI_API_KEY"),
+            azure_endpoint=os.environ.get("OPENAI_AZURE_ENDPOINT"),
+            api_version=os.environ.get("OPENAI_API_VERSION"),
+            azure_deployment=os.environ.get("OPENAI_AZURE_DEPLOYMENT"),
+        )

evals/registry.py

@@ -15,15 +15,16 @@
 
 import openai
 import yaml
-from openai import OpenAI
+
 
 from evals import OpenAIChatCompletionFn, OpenAICompletionFn
 from evals.api import CompletionFn, DummyCompletionFn
 from evals.base import BaseEvalSpec, CompletionFnSpec, EvalSetSpec, EvalSpec
 from evals.elsuite.modelgraded.base import ModelGradedSpec
+from evals.openai_client.openai_client_provider import OpenAIClientProvider
 from evals.utils.misc import make_object
 
-client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+client = OpenAIClientProvider(api_key=os.environ.get("OPENAI_API_KEY")).get_client()
 
 logger = logging.getLogger(__name__)