Merge branch 'main' into ducanhdt:add_docker

Author: shroominic

.pre-commit-config.yaml

@@ -24,4 +24,4 @@ repos:
     rev: 6.1.0
     hooks:
     -   id: flake8
-        args: [--max-line-length=88]
+        args: [--max-line-length=120]

codeinterpreterapi/agents/functions_agent.py

@@ -1,28 +1,32 @@
-"""
-Module implements an agent that uses OpenAI's APIs function enabled API.
-
-This file is a modified version of the original file
-from langchain/agents/openai_functions_agent/base.py.
-Credits go to the original authors :)
-"""
-
+"""Module implements an agent that uses OpenAI's APIs function enabled API."""
 import json
 from dataclasses import dataclass
 from json import JSONDecodeError
 from typing import Any, List, Optional, Sequence, Tuple, Union
 
 from langchain.agents import BaseSingleActionAgent
-from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.callbacks.manager import Callbacks
 from langchain.chat_models.openai import ChatOpenAI
-from langchain.prompts.chat import (BaseMessagePromptTemplate,
-                                    ChatPromptTemplate,
-                                    HumanMessagePromptTemplate,
-                                    MessagesPlaceholder)
-from langchain.schema import (AgentAction, AgentFinish, AIMessage, BaseMessage,
-                              BasePromptTemplate, FunctionMessage,
-                              OutputParserException, SystemMessage)
+from langchain.prompts.chat import (
+    BaseMessagePromptTemplate,
+    ChatPromptTemplate,
+    HumanMessagePromptTemplate,
+    MessagesPlaceholder,
+)
+from langchain.schema import (
+    AgentAction,
+    AgentFinish,
+    BasePromptTemplate,
+    OutputParserException,
+)
+from langchain.schema.language_model import BaseLanguageModel
+from langchain.schema.messages import (
+    AIMessage,
+    BaseMessage,
+    FunctionMessage,
+    SystemMessage,
+)
 from langchain.tools import BaseTool
 from langchain.tools.convert_to_openai import format_tool_to_openai_function
 from pydantic import root_validator
@@ -95,17 +99,14 @@ def _format_intermediate_steps(
     return messages
 
 
-async def _parse_ai_message(
-    message: BaseMessage, llm: BaseLanguageModel
-) -> Union[AgentAction, AgentFinish]:
+def _parse_ai_message(message: BaseMessage) -> Union[AgentAction, AgentFinish]:
     """Parse an AI message."""
     if not isinstance(message, AIMessage):
         raise TypeError(f"Expected an AI message got {type(message)}")
 
     function_call = message.additional_kwargs.get("function_call", {})
 
     if function_call:
-        function_call = message.additional_kwargs["function_call"]
         function_name = function_call["name"]
         try:
             _tool_input = json.loads(function_call["arguments"])
@@ -189,8 +190,42 @@ def input_keys(self) -> List[str]:
     def functions(self) -> List[dict]:
         return [dict(format_tool_to_openai_function(t)) for t in self.tools]
 
-    def plan(self):
-        raise NotImplementedError
+    def plan(
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        with_functions: bool = True,
+        **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        agent_scratchpad = _format_intermediate_steps(intermediate_steps)
+        selected_inputs = {
+            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
+        }
+        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
+        prompt = self.prompt.format_prompt(**full_inputs)
+        messages = prompt.to_messages()
+        if with_functions:
+            predicted_message = self.llm.predict_messages(
+                messages,
+                functions=self.functions,
+                callbacks=callbacks,
+            )
+        else:
+            predicted_message = self.llm.predict_messages(
+                messages,
+                callbacks=callbacks,
+            )
+        agent_decision = _parse_ai_message(predicted_message)
+        return agent_decision
 
     async def aplan(
         self,
@@ -218,9 +253,38 @@ async def aplan(
         predicted_message = await self.llm.apredict_messages(
             messages, functions=self.functions, callbacks=callbacks
         )
-        agent_decision = await _parse_ai_message(predicted_message, self.llm)
+        agent_decision = _parse_ai_message(predicted_message)
         return agent_decision
 
+    def return_stopped_response(
+        self,
+        early_stopping_method: str,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        **kwargs: Any,
+    ) -> AgentFinish:
+        """Return response when agent has been stopped due to max iterations."""
+        if early_stopping_method == "force":
+            # `force` just returns a constant string
+            return AgentFinish(
+                {"output": "Agent stopped due to iteration limit or time limit."}, ""
+            )
+        elif early_stopping_method == "generate":
+            # Generate does one final forward pass
+            agent_decision = self.plan(
+                intermediate_steps, with_functions=False, **kwargs
+            )
+            if type(agent_decision) == AgentFinish:  # noqa: E721
+                return agent_decision
+            else:
+                raise ValueError(
+                    f"got AgentAction with no functions provided: {agent_decision}"
+                )
+        else:
+            raise ValueError(
+                "early_stopping_method should be one of `force` or `generate`, "
+                f"got {early_stopping_method}"
+            )
+
     @classmethod
     def create_prompt(
         cls,
@@ -275,7 +339,7 @@ def from_llm_and_tools(
             extra_prompt_messages=extra_prompt_messages,
             system_message=system_message,
         )
-        return cls(
+        return cls(  # type: ignore
             llm=llm,
             prompt=prompt,
             tools=tools,

codeinterpreterapi/chains/__init__.py

@@ -1,9 +1,11 @@
 from .extract_code import extract_python_code
-from .modifications_check import get_file_modifications
-from .rm_dl_link import remove_download_link
+from .modifications_check import aget_file_modifications, get_file_modifications
+from .rm_dl_link import aremove_download_link, remove_download_link
 
 __all__ = [
     "extract_python_code",
     "get_file_modifications",
+    "aget_file_modifications",
     "remove_download_link",
+    "aremove_download_link",
 ]

codeinterpreterapi/chains/modifications_check.py

@@ -7,7 +7,28 @@
 from codeinterpreterapi.prompts import determine_modifications_prompt
 
 
-async def get_file_modifications(
+def get_file_modifications(
+    code: str,
+    llm: BaseLanguageModel,
+    retry: int = 2,
+) -> Optional[List[str]]:
+    if retry < 1:
+        return None
+
+    prompt = determine_modifications_prompt.format(code=code)
+
+    result = llm.predict(prompt, stop="```")
+
+    try:
+        result = json.loads(result)
+    except json.JSONDecodeError:
+        result = ""
+    if not result or not isinstance(result, dict) or "modifications" not in result:
+        return get_file_modifications(code, llm, retry=retry - 1)
+    return result["modifications"]
+
+
+async def aget_file_modifications(
     code: str,
     llm: BaseLanguageModel,
     retry: int = 2,
@@ -24,12 +45,12 @@ async def get_file_modifications(
     except json.JSONDecodeError:
         result = ""
     if not result or not isinstance(result, dict) or "modifications" not in result:
-        return await get_file_modifications(code, llm, retry=retry - 1)
+        return await aget_file_modifications(code, llm, retry=retry - 1)
     return result["modifications"]
 
 
 async def test():
-    llm = ChatAnthropic(model="claude-1.3")  # type: ignore
+    llm = ChatAnthropic(model="claude-2")  # type: ignore
 
     code = """
         import matplotlib.pyplot as plt
@@ -45,7 +66,7 @@ async def test():
         plt.show()
         """
 
-    print(await get_file_modifications(code, llm))
+    print(get_file_modifications(code, llm))
 
 
 if __name__ == "__main__":

codeinterpreterapi/chains/rm_dl_link.py

@@ -5,7 +5,22 @@
 from codeinterpreterapi.prompts import remove_dl_link_prompt
 
 
-async def remove_download_link(
+def remove_download_link(
+    input_response: str,
+    llm: BaseLanguageModel,
+) -> str:
+    messages = remove_dl_link_prompt.format_prompt(
+        input_response=input_response
+    ).to_messages()
+    message = llm.predict_messages(messages)
+
+    if not isinstance(message, AIMessage):
+        raise OutputParserException("Expected an AIMessage")
+
+    return message.content
+
+
+async def aremove_download_link(
     input_response: str,
     llm: BaseLanguageModel,
 ) -> str:
@@ -20,21 +35,19 @@ async def remove_download_link(
     return message.content
 
 
-async def test():
+def test():
     llm = ChatOpenAI(model="gpt-3.5-turbo-0613")  # type: ignore
 
     example = (
         "I have created the plot to your dataset.\n\n"
         "Link to the file [here](sandbox:/plot.png)."
     )
-    print(await remove_download_link(example, llm))
+    print(remove_download_link(example, llm))
 
 
 if __name__ == "__main__":
-    import asyncio
-
     from dotenv import load_dotenv
 
     load_dotenv()
 
-    asyncio.run(test())
+    test()

codeinterpreterapi/utils/parser.py renamed to codeinterpreterapi/parser.py

@@ -39,8 +39,7 @@ def _type(self) -> str:
 
 class CodeChatAgentOutputParser(AgentOutputParser):
     def get_format_instructions(self) -> str:
-        from langchain.agents.conversational_chat.prompt import \
-            FORMAT_INSTRUCTIONS
+        from langchain.agents.conversational_chat.prompt import FORMAT_INSTRUCTIONS
 
         return FORMAT_INSTRUCTIONS
 

codeinterpreterapi/prompts/remove_dl_link.py

@@ -1,5 +1,4 @@
-from langchain.prompts.chat import (ChatPromptTemplate,
-                                    HumanMessagePromptTemplate)
+from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
 from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 remove_dl_link_prompt = ChatPromptTemplate(

codeinterpreterapi/schema/__init__.py

@@ -1,11 +1,13 @@
 from .file import File
 from .input import CodeInput, FileInput
 from .response import CodeInterpreterResponse, UserRequest
+from .status import SessionStatus
 
 __all__ = [
     "CodeInterpreterResponse",
     "CodeInput",
     "File",
     "FileInput",
     "UserRequest",
+    "SessionStatus",
 ]

codeinterpreterapi/schema/response.py

@@ -14,9 +14,15 @@ def __repr__(self):
 
 
 class CodeInterpreterResponse(AIMessage):
+    """
+    Response from the code interpreter agent.
+
+    files: list of files to be sent to the user (File )
+    code_log: list[tuple[str, str]] = []
+    """
+
     files: list[File] = []
-    # final_code: str = ""  TODO: implement
-    # final_output: str = ""  TODO: implement
+    code_log: list[tuple[str, str]] = []
 
     def show(self):
         print("AI: ", self.content)

codeinterpreterapi/schema/status.py

@@ -0,0 +1,10 @@
+from codeboxapi.schema import CodeBoxStatus  # type: ignore
+
+
+class SessionStatus(CodeBoxStatus):
+    @classmethod
+    def from_codebox_status(cls, cbs: CodeBoxStatus) -> "SessionStatus":
+        return cls(status=cbs.status)
+
+    def __repr__(self):
+        return f"<SessionStatus status={self.status}>"