Add Azure OpenAI

shroominic · Aug 4, 2023 · 3960c0b · 3960c0b
1 parent a91cd26
commit 3960c0b
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 27 deletions.
diff --git a/.env.example b/.env.example
@@ -4,3 +4,8 @@ OPENAI_API_KEY=
 # CODEBOX_API_KEY=
 # (set True to enable logging)
 VERBOSE=False 
+# (optional, required for production)
+OPENAI_API_TYPE=azure
+OPENAI_API_VERSION=2023-07-01-preview
+OPENAI_API_BASE=
+DEPLOYMENT_NAME=
diff --git a/codeinterpreterapi/session.py b/codeinterpreterapi/session.py
@@ -5,7 +5,7 @@
 from codeboxapi import CodeBox  # type: ignore
 from codeboxapi.schema import CodeBoxOutput  # type: ignore
 from langchain.tools import StructuredTool, BaseTool
-from langchain.chat_models import ChatOpenAI, ChatAnthropic
+from langchain.chat_models import ChatOpenAI, ChatAnthropic, AzureChatOpenAI
 from langchain.chat_models.base import BaseChatModel
 from langchain.schema.language_model import BaseLanguageModel
 from langchain.prompts.chat import MessagesPlaceholder
@@ -72,33 +72,50 @@ def _tools(
             ),
         ]
 
-    def _choose_llm(
-        self,
-        model: str = "gpt-4",
-        openai_api_key: Optional[str] = None,
-        **kwargs
-    ) -> BaseChatModel:
-        if "gpt" in model:
-            openai_api_key = (
-                openai_api_key 
-                or settings.OPENAI_API_KEY 
-                or getenv("OPENAI_API_KEY", None)
-            )
-            if openai_api_key is None:
-                raise ValueError(
-                    "OpenAI API key missing. Set OPENAI_API_KEY env variable or pass `openai_api_key` to session."
+    class CodeInterpreterSession:
+        # ...
+        def _choose_llm(
+            self,
+            model: str = "gpt-4",
+            openai_api_key: Optional[str] = None,
+            **kwargs
+        ) -> BaseChatModel:
+            if "gpt" in model:
+                openai_api_key = (
+                    openai_api_key 
+                    or settings.OPENAI_API_KEY 
+                    or getenv("OPENAI_API_KEY", None)
                 )
-            return ChatOpenAI(
-                temperature=0.03,
-                model=model,
-                openai_api_key=openai_api_key,
-                max_retries=3,
-                request_timeout=60 * 3,
-            )  # type: ignore
-        elif "claude" in model:
-            return ChatAnthropic(model=model)
-        else:
-            raise ValueError(f"Unknown model: {model} (expected gpt or claude model)")
+                if openai_api_key is None:
+                    raise ValueError(
+                        "OpenAI API key missing. Set OPENAI_API_KEY env variable or pass `openai_api_key` to session."
+                    )
+                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:
+                raise ValueError(f"Unknown model: {model} (expected gpt or claude model)")
 
     def _choose_agent(self) -> BaseSingleActionAgent:
         return (