From f149ddcb0a91da7b260f892729e0a3ff9c06dd32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 10 Jan 2025 11:39:38 +0000 Subject: [PATCH] Implement brief MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- cai/core.py | 37 ++++++++++++++----- cai/util.py | 10 ++++- .../1_arch_short_picoctf_static_flag.py | 6 ++- 3 files changed, 42 insertions(+), 11 deletions(-) diff --git a/cai/core.py b/cai/core.py index f041bf8b..948530fe 100644 --- a/cai/core.py +++ b/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 diff --git a/cai/util.py b/cai/util.py index f5d3836f..fa8f318b 100644 --- a/cai/util.py +++ b/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") diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py index 18c67027..6379d5d0 100644 --- a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py +++ b/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()