From 318c4781de9f221050dcbbc1a3dd96a52c61ab53 Mon Sep 17 00:00:00 2001
From: adldotori <adldotori@gmail.com>
Date: Mon, 3 Apr 2023 10:54:38 +0000
Subject: [PATCH] feat: logging in parse method

---
 core/agents/callback.py        | 29 +++++++++++------------------
 core/agents/parser.py          | 18 ++++++++++++++++--
 core/tools/terminal/syscall.py |  2 +-
 3 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/core/agents/callback.py b/core/agents/callback.py
index 3f037d4..c18d1f0 100644
--- a/core/agents/callback.py
+++ b/core/agents/callback.py
@@ -1,18 +1,18 @@
 from typing import Any, Dict, List, Optional, Union
 
+from ansi import ANSI, Color, Style
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentAction, AgentFinish, LLMResult
-
-from ansi import ANSI, Color, Style
 from logger import logger
 
 
-class EVALCallbackHandler(BaseCallbackHandler):
-    def dim_multiline(self, message: str) -> str:
-        return message.split("\n")[0] + ANSI(
-            "\n... ".join(["", *message.split("\n")[1:]])
-        ).to(Color.black().bright())
+def dim_multiline(message: str) -> str:
+    return message.split("\n")[0] + ANSI(
+        "\n... ".join(["", *message.split("\n")[1:]])
+    ).to(Color.black().bright())
+
 
+class EVALCallbackHandler(BaseCallbackHandler):
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
@@ -33,7 +33,7 @@ class EVALCallbackHandler(BaseCallbackHandler):
         self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
     ) -> None:
         logger.info(ANSI(f"Entering new chain.").to(Color.green(), Style.italic()))
-        logger.info(ANSI("Prompted Text").to(Color.yellow()) + f': {inputs["input"]}')
+        logger.info(ANSI("Prompted Text").to(Color.yellow()) + f': {inputs["input"]}\n')
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         logger.info(ANSI(f"Finished chain.").to(Color.green(), Style.italic()))
@@ -52,14 +52,7 @@ class EVALCallbackHandler(BaseCallbackHandler):
         pass
 
     def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
-        logger.info(
-            ANSI("Action").to(Color.cyan()) + ": " + ANSI(action.tool).to(Style.bold())
-        )
-        logger.info(
-            ANSI("Input").to(Color.cyan())
-            + ": "
-            + self.dim_multiline(action.tool_input)
-        )
+        pass
 
     def on_tool_end(
         self,
@@ -69,7 +62,7 @@ class EVALCallbackHandler(BaseCallbackHandler):
         **kwargs: Any,
     ) -> None:
         logger.info(
-            ANSI("Observation").to(Color.magenta()) + ": " + self.dim_multiline(output)
+            ANSI("Observation").to(Color.magenta()) + ": " + dim_multiline(output)
         )
         logger.info(ANSI("Thinking...").to(Color.green(), Style.italic()))
 
@@ -93,5 +86,5 @@ class EVALCallbackHandler(BaseCallbackHandler):
         logger.info(
             ANSI("Final Answer").to(Color.yellow())
             + ": "
-            + self.dim_multiline(finish.return_values.get("output", ""))
+            + dim_multiline(finish.return_values.get("output", ""))
         )
diff --git a/core/agents/parser.py b/core/agents/parser.py
index 08832f6..a54c5dd 100644
--- a/core/agents/parser.py
+++ b/core/agents/parser.py
@@ -1,9 +1,12 @@
 import re
+import time
 from typing import Dict
 
-from langchain.output_parsers.base import BaseOutputParser
-
+from ansi import ANSI, Color, Style
+from core.agents.callback import dim_multiline
 from core.prompts.input import EVAL_FORMAT_INSTRUCTIONS
+from langchain.output_parsers.base import BaseOutputParser
+from logger import logger
 
 
 class EvalOutputParser(BaseOutputParser):
@@ -15,8 +18,19 @@ class EvalOutputParser(BaseOutputParser):
         match = re.search(regex, text, re.DOTALL)
         if not match:
             raise Exception("parse error")
+
         action = match.group(1).strip()
+        plan = match.group(2)
         action_input = match.group(3)
+
+        logger.info(ANSI("Plan").to(Color.blue()) + ": " + plan)
+        time.sleep(1)
+        logger.info(
+            ANSI("Action").to(Color.cyan()) + ": " + ANSI(action).to(Style.bold())
+        )
+        time.sleep(1)
+        logger.info(ANSI("Input").to(Color.cyan()) + ": " + dim_multiline(action_input))
+        time.sleep(1)
         return {"action": action, "action_input": action_input.strip(" ").strip('"')}
 
     def __str__(self):
diff --git a/core/tools/terminal/syscall.py b/core/tools/terminal/syscall.py
index fd16399..4541b1a 100644
--- a/core/tools/terminal/syscall.py
+++ b/core/tools/terminal/syscall.py
@@ -61,7 +61,7 @@ class SyscallTracer:
                 break
 
             try:
-                self.wait_syscall_with_timeout(60)
+                self.wait_syscall_with_timeout(30)
             except ProcessExit as event:
                 if event.exitcode is not None:
                     exitcode = event.exitcode