🔧 Refactor code interpreter API to use langchain_core and langchain_openai libraries

Author: shroominic

src/codeinterpreterapi/__init__.py

@@ -1,10 +1,9 @@
+from . import _patch_parser  # noqa
+
 from codeinterpreterapi.config import settings
 from codeinterpreterapi.schema import File
 from codeinterpreterapi.session import CodeInterpreterSession
 
-from ._patch_parser import patch
-
-patch()
 
 __all__ = [
     "CodeInterpreterSession",

src/codeinterpreterapi/_patch_parser.py

@@ -3,16 +3,13 @@
 from json import JSONDecodeError
 from typing import List, Union
 
+from langchain.agents.agent import AgentOutputParser
+from langchain.agents.openai_functions_agent import base
 from langchain_core.agents import AgentAction, AgentActionMessageLog, AgentFinish
 from langchain_core.exceptions import OutputParserException
-from langchain_core.messages import (
-    AIMessage,
-    BaseMessage,
-)
+from langchain_core.messages import AIMessage, BaseMessage
 from langchain_core.outputs import ChatGeneration, Generation
 
-from langchain.agents.agent import AgentOutputParser
-
 
 class OpenAIFunctionsAgentOutputParser(AgentOutputParser):
     """Parses a message into agent action/finish.
@@ -102,8 +99,5 @@ def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
         raise ValueError("Can only parse messages")
 
 
-def patch() -> None:
-    """Patch the parser."""
-    from langchain.agents import openai_functions_agent
-
-    openai_functions_agent.OpenAIFunctionsAgentOutputParser = OpenAIFunctionsAgentOutputParser  # type: ignore
+# patch
+base.OpenAIFunctionsAgentOutputParser = OpenAIFunctionsAgentOutputParser  # type: ignore

src/codeinterpreterapi/chains/extract_code.py

@@ -1,5 +1,4 @@
-from langchain.base_language import BaseLanguageModel
-from langchain.chat_models.anthropic import ChatAnthropic
+from langchain_core.language_models import BaseLanguageModel
 
 
 def extract_python_code(
@@ -19,7 +18,9 @@ async def aextract_python_code(
 
 
 async def test() -> None:
-    llm = ChatAnthropic(model="claude-1.3")  # type: ignore
+    from langchain_openai import ChatOpenAI
+
+    llm = ChatOpenAI()
 
     code = """
         import matplotlib.pyplot as plt

src/codeinterpreterapi/chains/modifications_check.py

@@ -1,26 +1,25 @@
 import json
 from typing import List, Optional
 
-from langchain.base_language import BaseLanguageModel
-from langchain.chat_models.anthropic import ChatAnthropic
+from langchain_core.language_models import BaseLanguageModel
 
 from codeinterpreterapi.prompts import determine_modifications_prompt
 
 
 def get_file_modifications(
     code: str,
     llm: BaseLanguageModel,
-    retry: int = 2,
+    retry: int = 4,
 ) -> Optional[List[str]]:
     if retry < 1:
         return None
 
     prompt = determine_modifications_prompt.format(code=code)
 
-    result = llm.predict(prompt, stop="```")
+    result = llm.invoke(prompt)
 
     try:
-        result = json.loads(result)
+        result = json.loads(result.content)
     except json.JSONDecodeError:
         result = ""
     if not result or not isinstance(result, dict) or "modifications" not in result:
@@ -31,17 +30,17 @@ def get_file_modifications(
 async def aget_file_modifications(
     code: str,
     llm: BaseLanguageModel,
-    retry: int = 2,
+    retry: int = 4,
 ) -> Optional[List[str]]:
     if retry < 1:
         return None
 
     prompt = determine_modifications_prompt.format(code=code)
 
-    result = await llm.apredict(prompt, stop="```")
+    result = await llm.ainvoke(prompt)
 
     try:
-        result = json.loads(result)
+        result = json.loads(result.content)
     except json.JSONDecodeError:
         result = ""
     if not result or not isinstance(result, dict) or "modifications" not in result:
@@ -50,7 +49,9 @@ async def aget_file_modifications(
 
 
 async def test() -> None:
-    llm = ChatAnthropic(model="claude-2")  # type: ignore
+    from langchain_openai import ChatOpenAI
+
+    llm = ChatOpenAI()
 
     code = """
         import matplotlib.pyplot as plt

src/codeinterpreterapi/chains/rm_dl_link.py

@@ -1,6 +1,7 @@
-from langchain.base_language import BaseLanguageModel
-from langchain.chat_models.openai import ChatOpenAI
-from langchain.schema import AIMessage, OutputParserException
+from langchain_core.exceptions import OutputParserException
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.messages import AIMessage
+from langchain_openai import ChatOpenAI
 
 from codeinterpreterapi.prompts import remove_dl_link_prompt
 
@@ -12,7 +13,7 @@ def remove_download_link(
     messages = remove_dl_link_prompt.format_prompt(
         input_response=input_response
     ).to_messages()
-    message = llm.predict_messages(messages)
+    message = llm.invoke(messages)
 
     if not isinstance(message, AIMessage):
         raise OutputParserException("Expected an AIMessage")
@@ -28,7 +29,7 @@ async def aremove_download_link(
     messages = remove_dl_link_prompt.format_prompt(
         input_response=input_response
     ).to_messages()
-    message = await llm.apredict_messages(messages)
+    message = await llm.ainvoke(messages)
 
     if not isinstance(message, AIMessage):
         raise OutputParserException("Expected an AIMessage")

src/codeinterpreterapi/chat_history.py

@@ -2,12 +2,11 @@
 import json
 from typing import List
 
-from codeboxapi import CodeBox  # type: ignore
-from langchain.schema import BaseChatMessageHistory
-from langchain.schema.messages import BaseMessage, messages_from_dict, messages_to_dict
+from codeboxapi import CodeBox
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.messages import BaseMessage, messages_from_dict, messages_to_dict
 
 
-# TODO: This is probably not efficient, but it works for now.
 class CodeBoxChatMessageHistory(BaseChatMessageHistory):
     """
     Chat message history that stores history inside the codebox.

src/codeinterpreterapi/config.py

@@ -1,14 +1,10 @@
 from typing import Optional
 
-from dotenv import load_dotenv
-from langchain.pydantic_v1 import BaseSettings, SecretStr
-from langchain.schema import SystemMessage
+from langchain_core.messages import SystemMessage
+from langchain_core.pydantic_v1 import BaseSettings, SecretStr
 
 from codeinterpreterapi.prompts import code_interpreter_system_message
 
-# .env file
-load_dotenv(dotenv_path="./.env")
-
 
 class CodeInterpreterAPISettings(BaseSettings):
     """
@@ -18,8 +14,8 @@ class CodeInterpreterAPISettings(BaseSettings):
     DEBUG: bool = False
 
     # Models
-    OPENAI_API_KEY: Optional[str] = None
-    AZURE_API_KEY: Optional[str] = None
+    OPENAI_API_KEY: Optional[SecretStr] = None
+    AZURE_OPENAI_API_KEY: Optional[SecretStr] = None
     AZURE_API_BASE: Optional[str] = None
     AZURE_API_VERSION: Optional[str] = None
     AZURE_DEPLOYMENT_NAME: Optional[str] = None
@@ -46,5 +42,9 @@ class CodeInterpreterAPISettings(BaseSettings):
     # deprecated
     VERBOSE: bool = DEBUG
 
+    class Config:
+        env_file = "./.env"
+        extra = "ignore"
+
 
 settings = CodeInterpreterAPISettings()

src/codeinterpreterapi/prompts/modifications_check.py

@@ -1,4 +1,4 @@
-from langchain.prompts import PromptTemplate
+from langchain_core.prompts import PromptTemplate
 
 determine_modifications_prompt = PromptTemplate(
     input_variables=["code"],

src/codeinterpreterapi/prompts/remove_dl_link.py

@@ -1,5 +1,5 @@
-from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
-from langchain.schema import AIMessage, HumanMessage, SystemMessage
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_core.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
 
 remove_dl_link_prompt = ChatPromptTemplate(
     input_variables=["input_response"],

src/codeinterpreterapi/prompts/system_message.py

@@ -1,4 +1,4 @@
-from langchain.schema import SystemMessage
+from langchain_core.messages import SystemMessage
 
 system_message = SystemMessage(
     content="""