diff --git a/README.md b/README.md
index 31aeb866..85813e0c 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,7 @@ python3 -m venv .venv
 source .venv/bin/activate
 pip install -r requirements.txt
 ```
-
+Set up the `OPENAI_API_KEY` environment variable (for OpenAI) or the `AZURE_OPENAI_API_KEY`, `AZURE_API_VERSION` and `AZURE_API_BASE` environment variables (for Azure OpenAI).  
 You can then run the optillm proxy as follows.
 
 ```bash
@@ -77,7 +77,7 @@ response = client.chat.completions.create(
 
 print(response)
 ```
-
+The code above applies to both OpenAI and Azure OpenAI, just remember to populate the `OPENAI_API_KEY` env variable with the proper key.  
 You can control the technique you use for optimization by prepending the slug to the model name `{slug}-model-name`. E.g. in the above code we are using `moa` or
 mixture of agents as the optimization approach. In the proxy logs you will see the following showing the `moa` is been used with the base model as `gpt-4o-mini`.
 
diff --git a/optillm.py b/optillm.py
index c0f987e4..1044b7aa 100644
--- a/optillm.py
+++ b/optillm.py
@@ -3,7 +3,7 @@
 import os
 import secrets
 from flask import Flask, request, jsonify
-from openai import OpenAI
+from openai import AzureOpenAI, OpenAI
 
 # Import approach modules
 from mcts import chat_with_mcts
@@ -26,9 +26,16 @@
 # Initialize Flask app
 app = Flask(__name__)
 
-# OpenAI API configuration
-API_KEY = os.environ.get("OPENAI_API_KEY")
-default_client = OpenAI(api_key=API_KEY)
+# OpenAI or Azure API configuration
+if os.environ.get("OPENAI_API_KEY") != None:
+    API_KEY = os.environ.get("OPENAI_API_KEY")
+    default_client = OpenAI(api_key=API_KEY)
+else:
+    default_client = AzureOpenAI(
+        api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
+        api_version=os.environ.get("AZURE_API_VERSION"),
+        azure_endpoint=os.environ.get("AZURE_API_BASE"),
+)
 
 # Server configuration
 server_config = {