🔨parser fix typing

Author: shroominic

codeinterpreterapi/chains/extract_code.py

@@ -2,13 +2,20 @@
 from langchain.chat_models.anthropic import ChatAnthropic
 
 
-# TODO: make async
 def extract_python_code(
     text: str,
     llm: BaseLanguageModel,
     retry: int = 2,
 ):
-    pass
+    pass  # TODO
+
+
+async def aextract_python_code(
+    text: str,
+    llm: BaseLanguageModel,
+    retry: int = 2,
+):
+    pass  # TODO
 
 
 async def test():

codeinterpreterapi/parser.py

@@ -38,6 +38,10 @@ def _type(self) -> str:
 
 
 class CodeChatAgentOutputParser(AgentOutputParser):
+    def __init__(self, llm: BaseChatModel, **kwargs):
+        super().__init__(**kwargs)
+        self.llm = llm
+
     def get_format_instructions(self) -> str:
         from langchain.agents.conversational_chat.prompt import FORMAT_INSTRUCTIONS
 
@@ -46,9 +50,7 @@ def get_format_instructions(self) -> str:
     def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
         raise NotImplementedError
 
-    async def aparse(
-        self, text: str, llm: BaseChatModel
-    ) -> Union[AgentAction, AgentFinish]:
+    async def aparse(self, text: str) -> Union[AgentAction, AgentFinish]:
         try:
             response = parse_json_markdown(text)
             action, action_input = response["action"], response["action_input"]
@@ -58,8 +60,9 @@ async def aparse(
                 return AgentAction(action, action_input, text)
         except Exception:
             if '"action": "python"' in text:
+                print("TODO: Not implemented")
                 # extract python code from text with prompt
-                text = extract_python_code(text, llm=llm) or ""
+                text = extract_python_code(text, llm=self.llm) or ""
                 match = re.search(r"```python\n(.*?)```", text)
                 if match:
                     code = match.group(1).replace("\\n", "; ")

codeinterpreterapi/session.py

@@ -163,7 +163,7 @@ def _choose_agent(self) -> BaseSingleActionAgent:
                 llm=self.llm,
                 tools=self.tools,
                 system_message=code_interpreter_system_message.content,
-                output_parser=CodeChatAgentOutputParser(),
+                output_parser=CodeChatAgentOutputParser(self.llm),
             )
             if isinstance(self.llm, BaseChatModel)
             else ConversationalAgent.from_llm_and_tools(