Merge branch 'main' into add_docker

shroominic · Aug 8, 2023 · 97ac6e6 · 97ac6e6
2 parents 8f2d89c + 0f676c4
commit 97ac6e6
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 10 deletions.
diff --git a/.env.example b/.env.example
@@ -4,7 +4,11 @@ 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=
 
 STORAGE_TYPE="file" # "file" or "redis"
 

diff --git a/README.md b/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

diff --git a/codeinterpreterapi/session.py b/codeinterpreterapi/session.py
@@ -15,7 +15,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.memory.chat_message_histories import (
@@ -103,13 +103,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: