feat: Add Azure OpenAI support and configuration instructions

README.md

@@ -26,13 +26,27 @@ For technical details and a complete evaluation, see our paper (to appear at FSE
 
 > [!IMPORTANT]
 >
-> ChatDBG currently needs to be connected to an [OpenAI account](https://openai.com/api/). _Your account will need to have a positive balance for this to work_ ([check your balance](https://platform.openai.com/account/usage)). If you have never purchased credits, you will need to purchase at least \$1 in credits (if your API account was created before August 13, 2023) or \$0.50 (if you have a newer API account) in order to have access to GPT-4, which ChatDBG uses. [Get a key here.](https://platform.openai.com/account/api-keys)
+> ChatDBG can be used with either OpenAI or Azure OpenAI.
+>
+> ### OpenAI Configuration
+> ChatDBG needs to be connected to an [OpenAI account](https://openai.com/api/). _Your account will need to have a positive balance for this to work_ ([check your balance](https://platform.openai.com/account/usage)). If you have never purchased credits, you will need to purchase at least \$1 in credits (if your API account was created before August 13, 2023) or \$0.50 (if you have a newer API account) in order to have access to GPT-4, which ChatDBG uses. [Get a key here.](https://platform.openai.com/account/api-keys)
 >
 > Once you have an API key, set it as an environment variable called `OPENAI_API_KEY`.
 >
 > ```bash
 > export OPENAI_API_KEY=<your-api-key>
 > ```
+>
+> ### Azure OpenAI Configuration
+> Alternatively, you can use Azure OpenAI by setting the following environment variables:
+>
+> ```bash
+> export AZURE_API_KEY=<your-azure-api-key>
+> export AZURE_API_BASE=<your-azure-endpoint>  # e.g., https://YOUR_RESOURCE.openai.azure.com
+> export AZURE_API_VERSION=<api-version>       # e.g., 2024-02-15-preview
+> export CHATDBG_MODEL=<your-deployment-name>  # e.g., azure/gpt-4o
+> ```
+>
 
 Install ChatDBG using `pip` (you need to do this whether you are debugging Python, C, or C++ code):
 

src/chatdbg/assistant/assistant.py

@@ -1,4 +1,5 @@
 import json
+import os
 import string
 import textwrap
 import time
@@ -124,7 +125,24 @@ def _check_model(self):
         missing_keys = result["missing_keys"]
         if missing_keys != []:
             _, provider, _, _ = litellm.get_llm_provider(self._model)
-            if provider == "openai":
+            if provider == "azure":
+                missing_azure_vars = [k for k in missing_keys if k.startswith("AZURE_")]
+                raise AssistantError(
+                    textwrap.dedent(
+                        f"""\
+                    You need to set the following Azure OpenAI environment variables:
+                    - AZURE_API_KEY: Your Azure OpenAI API key
+                    - AZURE_API_BASE: Your Azure OpenAI endpoint (e.g., https://YOUR_RESOURCE.openai.azure.com)
+                    - AZURE_API_VERSION: API version (e.g., 2024-02-15-preview)
+                    - CHATDBG_MODEL: The model name, which should be in the format azure/<deployment_name>
+
+                    Missing variables: {', '.join(missing_azure_vars)}
+
+                    For Azure OpenAI, use the deployment name as the model name.
+                    The deployment should be using a compatible model."""
+                    )
+                )
+            elif provider == "openai":
                 raise AssistantError(
                     textwrap.dedent(
                         f"""\
@@ -265,19 +283,31 @@ def _streamed_query(self, prompt: str, user_text):
         return stats
 
     def _stream_completion(self):
-
         self._trim_conversation()
 
-        return litellm.completion(
-            model=self._model,
-            messages=self._conversation,
-            tools=[
+        completion_params = {
+            "model": self._model,
+            "messages": self._conversation,
+            "tools": [
                 {"type": "function", "function": f["schema"]}
                 for f in self._functions.values()
             ],
-            timeout=self._timeout,
-            stream=True,
-        )
+            "timeout": self._timeout,
+            "stream": True,
+        }
+
+        # Add Azure specific configuration if Azure env vars are present and model looks like an Azure deployment
+        if all(
+            k in os.environ
+            for k in ["AZURE_API_KEY", "AZURE_API_BASE", "AZURE_API_VERSION"]
+        ) and self._model.startswith("azure"):
+            completion_params |= {
+                "api_key": os.getenv("AZURE_API_KEY"),
+                "api_base": os.getenv("AZURE_API_BASE"),
+                "api_version": os.getenv("AZURE_API_VERSION"),
+            }
+
+        return litellm.completion(**completion_params)
 
     def _trim_conversation(self):
         old_len = litellm.token_counter(self._model, messages=self._conversation)