Merge pull request #68 from somethingwentwell/add-azure-openai

Add azure openai

Author: shroominic

.env.example

@@ -4,3 +4,8 @@ OPENAI_API_KEY=
 # CODEBOX_API_KEY=
 # (set True to enable logging)
 VERBOSE=False
+# (optional, required for Azure OpenAI)
+# OPENAI_API_TYPE=azure
+# OPENAI_API_VERSION=2023-07-01-preview
+# OPENAI_API_BASE=
+# DEPLOYMENT_NAME=

README.md

@@ -27,7 +27,23 @@ For deployments, you can use `pip install codeinterpreterapi` instead which does
 
 ## Usage
 
-Make sure to set the `OPENAI_API_KEY` environment variable (or use a `.env` file)
+To configure OpenAI and Azure OpenAI, ensure that you set the appropriate environment variables (or use a .env file):
+
+For OpenAI, set the OPENAI_API_KEY environment variable:
+```
+export OPENAI_API_KEY=your_openai_api_key
+```
+
+For Azure OpenAI, set the following environment variables:
+```
+export OPENAI_API_TYPE=azure
+export OPENAI_API_VERSION=your_api_version
+export OPENAI_API_BASE=your_api_base
+export OPENAI_API_KEY=your_azure_openai_api_key
+export DEPLOYMENT_NAME=your_deployment_name
+```
+
+Remember to replace the placeholders with your actual API keys and other required information.
 
 ```python
 from codeinterpreterapi import CodeInterpreterSession

codeinterpreterapi/session.py

@@ -14,7 +14,7 @@
     ConversationalAgent,
     ConversationalChatAgent,
 )
-from langchain.chat_models import ChatAnthropic, ChatOpenAI
+from langchain.chat_models import AzureChatOpenAI, ChatAnthropic, ChatOpenAI
 from langchain.chat_models.base import BaseChatModel
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts.chat import MessagesPlaceholder
@@ -87,13 +87,33 @@ def _choose_llm(
                     "OpenAI API key missing. Set OPENAI_API_KEY env variable "
                     "or pass `openai_api_key` to session."
                 )
-            return ChatOpenAI(
-                temperature=0.03,
-                model=model,
-                openai_api_key=openai_api_key,
-                max_retries=3,
-                request_timeout=60 * 3,
-            )  # type: ignore
+            openai_api_version = getenv("OPENAI_API_VERSION")
+            openai_api_base = getenv("OPENAI_API_BASE")
+            deployment_name = getenv("DEPLOYMENT_NAME")
+            openapi_type = getenv("OPENAI_API_TYPE")
+            if (
+                openapi_type == "azure"
+                and openai_api_version
+                and openai_api_base
+                and deployment_name
+            ):
+                return AzureChatOpenAI(
+                    temperature=0.03,
+                    openai_api_base=openai_api_base,
+                    openai_api_version=openai_api_version,
+                    deployment_name=deployment_name,
+                    openai_api_key=openai_api_key,
+                    max_retries=3,
+                    request_timeout=60 * 3,
+                )
+            else:
+                return ChatOpenAI(
+                    temperature=0.03,
+                    model=model,
+                    openai_api_key=openai_api_key,
+                    max_retries=3,
+                    request_timeout=60 * 3,
+                )
         elif "claude" in model:
             return ChatAnthropic(model=model)
         else: