Print with color in controller (All-Hands-AI#176)

* and color-coded printing action abd obs

* fix langchains agent's observation

* remove color printing from codeact agent

* fix mypy

* fix ruff

* resolve conflict

* fix ruff

agenthub/codeact_agent/codeact_agent.py

@@ -1,7 +1,6 @@
 import re
 from typing import List, Mapping
 
-from termcolor import colored
 
 from opendevin.agent import Agent
 from opendevin.state import State
@@ -75,19 +74,16 @@ def step(self, state: State) -> Action:
                 {"role": "system", "content": SYSTEM_MESSAGE},
                 {"role": "user", "content": state.plan.main_goal},
             ]
-            print(colored("===USER:===\n" + state.plan.main_goal, "green"))
         updated_info = state.updated_info
         if updated_info:
             for prev_action, obs in updated_info:
                 assert isinstance(prev_action, (CmdRunAction, AgentEchoAction)), "Expecting CmdRunAction or AgentEchoAction for Action"
                 if isinstance(obs, AgentMessageObservation):  # warning message from itself
                     self.messages.append({"role": "user", "content": obs.content})
-                    print(colored("===USER:===\n" + obs.content, "green"))
                 elif isinstance(obs, CmdOutputObservation):
                     content = "OBSERVATION:\n" + obs.content
                     content += f"\n[Command {obs.command_id} finished with exit code {obs.exit_code}]]"
                     self.messages.append({"role": "user", "content": content})
-                    print(colored("===ENV OBSERVATION:===\n" + content, "blue"))
                 else:
                     raise NotImplementedError(f"Unknown observation type: {obs.__class__}")
         response = self.llm.completion(
@@ -98,27 +94,23 @@ def step(self, state: State) -> Action:
         )
         action_str: str = parse_response(response)
         self.messages.append({"role": "assistant", "content": action_str})
-        print(colored("===ASSISTANT:===\n" + action_str, "yellow"))
 
         command = re.search(r"<execute>(.*)</execute>", action_str, re.DOTALL)
         if command is not None:
             # a command was found
             command_group = command.group(1)
             if command_group.strip() == "exit":
-                print(colored("Exit received. Exiting...", "red"))
                 return AgentFinishAction()
             return CmdRunAction(command = command_group)
             # # execute the code
             # # TODO: does exit_code get loaded into Message?
             # exit_code, observation = self.env.execute(command_group)
             # self._history.append(Message(Role.ASSISTANT, observation))
-            # print(colored("===ENV OBSERVATION:===\n" + observation, "blue"))
         else:
             # we could provide a error message for the model to continue similar to
             # https://github.com/xingyaoww/mint-bench/blob/main/mint/envs/general_env.py#L18-L23
             # observation = INVALID_INPUT_MESSAGE
             # self._history.append(Message(Role.ASSISTANT, observation))
-            # print(colored("===ENV OBSERVATION:===\n" + observation, "blue"))
             return AgentEchoAction(content=INVALID_INPUT_MESSAGE)  # warning message to itself
 
 

agenthub/planner_agent/prompt.py

@@ -1,7 +1,7 @@
 import json
 from typing import List, Tuple, Dict, Type
 
-from opendevin.controller.agent_controller import print_with_indent
+from opendevin.controller.agent_controller import print_with_color
 from opendevin.plan import Plan
 from opendevin.action import Action, action_from_dict
 from opendevin.observation import Observation
@@ -178,7 +178,7 @@ def get_prompt(plan: Plan, history: List[Tuple[Action, Observation]]):
         elif latest_action_id == "finish":
             hint = ""
 
-    print_with_indent("HINT:\n" + hint)
+    print_with_color("HINT:\n" + hint, "INFO")
     return prompt % {
         'task': plan.main_goal,
         'plan': plan_str,

opendevin/controller/agent_controller.py

@@ -1,7 +1,8 @@
 
 import asyncio
-from typing import List, Callable
 import traceback
+from typing import List, Callable, Literal, Mapping, Any
+from termcolor import colored
 
 from opendevin.plan import Plan
 from opendevin.state import State
@@ -21,8 +22,24 @@
 
 from .command_manager import CommandManager
 
-def print_with_indent(text: str):
-    print("\t"+text.replace("\n","\n\t"), flush=True)
+
+ColorType = Literal['red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'light_grey', 'dark_grey', 'light_red', 'light_green', 'light_yellow', 'light_blue', 'light_magenta', 'light_cyan', 'white']
+
+def print_with_color(text: Any, print_type: str = "INFO"):
+    TYPE_TO_COLOR: Mapping[str, ColorType] = {
+        "BACKGROUND LOG": "blue",
+        "ACTION": "green",
+        "OBSERVATION": "yellow",
+        "INFO": "cyan",
+        "ERROR": "red",
+        "PLAN": "light_magenta",
+    }
+    color = TYPE_TO_COLOR.get(print_type.upper(), TYPE_TO_COLOR["INFO"])
+    print(
+        colored(f"\n{print_type.upper()}:\n", color, attrs=["bold"])
+        + colored(str(text), color),
+        flush=True,
+    )
 
 class AgentController:
     def __init__(
@@ -72,14 +89,13 @@ async def start_loop(self, task: str):
     async def step(self, i: int):
         print("\n\n==============", flush=True)
         print("STEP", i, flush=True)
-        print_with_indent("\nPLAN:\n")
-        print_with_indent(self.state.plan.__str__())
+        print_with_color(self.state.plan.main_goal, "PLAN")
 
         log_obs = self.command_manager.get_background_obs()
         for obs in log_obs:
             self.add_history(NullAction(), obs)
             await self._run_callbacks(obs)
-            print_with_indent("\nBACKGROUND LOG:\n%s" % obs)
+            print_with_color(obs, "BACKGROUND LOG")
 
         self.update_state_for_step(i)
         action: Action = NullAction()
@@ -88,10 +104,10 @@ async def step(self, i: int):
             action = self.agent.step(self.state)
             if action is None:
                 raise ValueError("Agent must return an action")
-            print_with_indent("\nACTION:\n%s" % action)
+            print_with_color(action, "ACTION")
         except Exception as e:
             observation = AgentErrorObservation(str(e))
-            print_with_indent("\nAGENT ERROR:\n%s" % observation)
+            print_with_color(observation, "ERROR")
             traceback.print_exc()
             # TODO Change to more robust error handling
             if "The api_key client option must be set" in observation.content:
@@ -102,34 +118,34 @@ async def step(self, i: int):
 
         finished = isinstance(action, AgentFinishAction)
         if finished:
-            print_with_indent("\nFINISHED")
+            print_with_color(action, "INFO")
             return True
 
         if isinstance(action, AddTaskAction):
             try:
                 self.state.plan.add_subtask(action.parent, action.goal, action.subtasks)
             except Exception as e:
                 observation = AgentErrorObservation(str(e))
-                print_with_indent("\nADD TASK ERROR:\n%s" % observation)
+                print_with_color(observation, "ERROR")
                 traceback.print_exc()
         elif isinstance(action, ModifyTaskAction):
             try:
                 self.state.plan.set_subtask_state(action.id, action.state)
             except Exception as e:
                 observation = AgentErrorObservation(str(e))
-                print_with_indent("\nMODIFY TASK ERROR:\n%s" % observation)
+                print_with_color(observation, "ERROR")
                 traceback.print_exc()
 
         if action.executable:
             try:
                 observation = action.run(self)
             except Exception as e:
                 observation = AgentErrorObservation(str(e))
-                print_with_indent("\nACTION RUN ERROR:\n%s" % observation)
+                print_with_color(observation, "ERROR")
                 traceback.print_exc()
 
         if not isinstance(observation, NullObservation):
-            print_with_indent("\nOBSERVATION:\n%s" % observation)
+            print_with_color(observation, "OBSERVATION")
 
         self.add_history(action, observation)
         await self._run_callbacks(observation)