Implement brief

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>

cai/core.py

@@ -47,6 +47,7 @@ class CAI:
             client = OpenAI(base_url=base_url, api_key=api_key)
         self.client = client
         self.ctf = ctf
+        self.brief = False
 
     def get_chat_completion(  # pylint: disable=too-many-arguments
         self,
@@ -68,7 +69,11 @@ class CAI:
             else agent.instructions
         )
         messages = [{"role": "system", "content": instructions}] + history
-        debug_print(debug, "Getting chat completion for...:", messages)
+        debug_print(
+            debug,
+            "Getting chat completion for...:",
+            messages,
+            brief=self.brief)
 
         tools = [function_to_json(f) for f in agent.functions]
         # hide context_variables from model
@@ -116,7 +121,7 @@ class CAI:
                     return Result(value=str(result))
                 except Exception as e:
                     error_message = f"Failed to cast response to string: {result}. Make sure agent functions return a string or Result object. Error: {str(e)}"  # noqa: E501 # pylint: disable=C0301
-                    debug_print(debug, error_message)
+                    debug_print(debug, error_message, brief=self.brief)
                     raise TypeError(error_message) from e
 
     def handle_tool_calls(
@@ -167,7 +172,10 @@ class CAI:
             name = tool_call.function.name
             # handle missing tool case, skip to next tool
             if name not in function_map:
-                debug_print(debug, f"Tool {name} not found in function map.")
+                debug_print(
+                    debug,
+                    f"Tool {name} not found in function map.",
+                    brief=self.brief)
                 partial_response.messages.append(
                     {
                         "role": "tool",
@@ -183,7 +191,8 @@ class CAI:
                 "Processing tool call",
                 name,
                 "with arguments",
-                args)
+                args,
+                brief=self.brief)
 
             func = function_map[name]
             # pass context_variables to agent functions
@@ -270,11 +279,15 @@ class CAI:
                 message.get("tool_calls", {}).values())
             if not message["tool_calls"]:
                 message["tool_calls"] = None
-            debug_print(debug, "Received completion:", message)
+            debug_print(
+                debug,
+                "Received completion:",
+                message,
+                brief=self.brief)
             history.append(message)
 
             if not message["tool_calls"] or not execute_tools:
-                debug_print(debug, "Ending turn.")
+                debug_print(debug, "Ending turn.", brief=self.brief)
                 break
 
             # convert tool_calls to objects
@@ -307,7 +320,7 @@ class CAI:
             )
         }
 
-    def run(  # pylint: disable=too-many-arguments,dangerous-default-value
+    def run(  # pylint: disable=too-many-arguments,dangerous-default-value, too-many-locals # noqa: E501
         self,
         agent: Agent,
         messages: List,
@@ -317,10 +330,12 @@ class CAI:
         debug: bool = False,
         max_turns: int = float("inf"),
         execute_tools: bool = True,
+        brief: bool = False,
     ) -> Response:
         """
         Run the cai and return the final response.
         """
+        self.brief = brief
         if stream:
             return self.run_and_stream(
                 agent=agent,
@@ -348,14 +363,18 @@ class CAI:
                 debug=debug,
             )
             message = completion.choices[0].message
-            debug_print(debug, "Received completion:", message)
+            debug_print(
+                debug,
+                "Received completion:",
+                message,
+                brief=self.brief)
             message.sender = active_agent.name
             history.append(
                 json.loads(message.model_dump_json())
             )  # to avoid OpenAI types (?)
 
             if not message.tool_calls or not execute_tools:
-                debug_print(debug, "Ending turn.")
+                debug_print(debug, "Ending turn.", brief=self.brief)
                 break
 
             # handle function calls, updating context_variables, and switching

cai/util.py

@@ -155,13 +155,21 @@ def format_chat_completion(msg, prev_msg=None) -> str:  # pylint: disable=unused
         COLORS['reset']}(\n    " + '\n    '.join(colored_lines) + "\n  )"
 
 
-def debug_print(debug: bool, intro: str, *args: Any) -> None:  # pylint: disable=too-many-locals # noqa: E501
+def debug_print(debug: bool, intro: str, *args: Any, brief: bool = False) -> None:  # pylint: disable=too-many-locals # noqa: E501
     """
     Print debug messages if debug mode is enabled with color-coded components.
+    If brief is True, prints a simplified timestamp and message format.
     """
     if not debug:
         return
 
+    if brief:
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        message = " ".join(map(str, [intro] + list(args)))
+        print(f"\033[97m[\033[90m{
+              timestamp}\033[97m]\033[90m {message}\033[0m")
+        return
+
     global _message_history  # pylint: disable=global-variable-not-assigned
 
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

examples/cybersecurity/1_arch_short_picoctf_static_flag.py

@@ -41,7 +41,11 @@ messages = [{
 }]
 
 # 5. Run the swarm
-response = client.run(agent=ctf_agent, messages=messages, debug=True)
+response = client.run(
+    agent=ctf_agent,
+    messages=messages,
+    debug=True,
+    brief=False)
 print(response.messages[-1]["content"])
 
 ctf.stop_ctf()