From 1a22013bda0e53a5a510f81ed7d8cc666f0909c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Thu, 27 Mar 2025 19:36:58 +0100 Subject: [PATCH] Advanced aligning agents but various missing bits: - agents can't translate directly it seems, see current error - in addition, it appears streaming mode leads to other problems MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- src/cai/agents/__init__.py | 81 +++++++++++++------ src/cai/agents/basic.py | 11 ++- src/cai/cli.py | 70 ++++++++++++---- src/cai/sdk/agents/run.py | 2 + .../reconnaissance/generic_linux_command.py | 2 + uv.lock | 14 ++++ 6 files changed, 132 insertions(+), 48 deletions(-) diff --git a/src/cai/agents/__init__.py b/src/cai/agents/__init__.py index f3975f1d..66dee8eb 100644 --- a/src/cai/agents/__init__.py +++ b/src/cai/agents/__init__.py @@ -49,6 +49,7 @@ import os import pkgutil import importlib from cai.sdk.agents import Agent +from cai.sdk.agents.handoffs import handoff from typing import Dict @@ -173,28 +174,58 @@ def get_agent_module(agent_name: str) -> str: return "unknown" - -load_dotenv() -cai_agent = os.getenv('CAI_AGENT_TYPE', "one_tool_agent").lower() -# Get all available agents -available_agents = get_available_agents() - -if cai_agent not in available_agents: - # stop and raise error - raise ValueError(f"Invalid CAI agent type: {cai_agent}") - -# Set the initial agent based on the environment variable -cai_initial_agent = available_agents[cai_agent] - -# Special handling for one_tool agent -# -# NOTE: maintained this for backwards compatibility -# -# NOTE 2: consider adding this if CTF_NAME defined -# -if cai_agent == "one_tool_agent": - from cai.agents.one_tool import transfer_to_one_tool_agent # noqa - cai_initial_agent.tools.append( - transfer_to_flag_discriminator - ) - flag_discriminator.tools.append(transfer_to_one_tool_agent) # noqa +def get_agent_by_name(agent_name: str) -> Agent: + """ + Get an agent instance by name. + + Args: + agent_name: Name of the agent to retrieve + + Returns: + Agent instance corresponding to the given name + + Raises: + ValueError: If the agent name is not found + """ + # Get all available agents from the agents module + available_agents = get_available_agents() + + # Convert agent_name to lowercase for case-insensitive comparison + agent_name = agent_name.lower() + + # Check if the agent exists in available_agents + if agent_name not in available_agents: + raise ValueError(f"Invalid agent type: {agent_name}. Available agents: {', '.join(available_agents.keys())}") + + # Get the agent instance + agent = available_agents[agent_name] + + # # Special handling for one_tool agent + # if agent_name == "one_tool_agent": + # from cai.agents.one_tool import one_tool_agent + + # # Create handoffs between agents + # # Add a handoff from one_tool_agent to flag_discriminator + # flag_discriminator_handoff = handoff( + # flag_discriminator, + # tool_name_override="transfer_to_flag_discriminator", + # tool_description_override="Transfer control to the flag discriminator agent" + # ) + + # # Add a handoff from flag_discriminator to one_tool_agent + # one_tool_agent_handoff = handoff( + # one_tool_agent, + # tool_name_override="transfer_to_one_tool_agent", + # tool_description_override="Transfer control back to the one tool agent" + # ) + + # # Add handoffs to agent.handoffs lists + # if not hasattr(agent, 'handoffs'): + # agent.handoffs = [] + # if not hasattr(flag_discriminator, 'handoffs'): + # flag_discriminator.handoffs = [] + + # agent.handoffs.append(flag_discriminator_handoff) + # flag_discriminator.handoffs.append(one_tool_agent_handoff) + + return agent \ No newline at end of file diff --git a/src/cai/agents/basic.py b/src/cai/agents/basic.py index 065414fa..3c075242 100644 --- a/src/cai/agents/basic.py +++ b/src/cai/agents/basic.py @@ -3,7 +3,7 @@ Collection of basic agents """ import os -from cai.util import Agent # pylint: disable=import-error +from cai.sdk.agents import Agent # pylint: disable=import-error from cai.tools.reconnaissance.netcat import netcat # pylint: disable=import-error # noqa: E501 from cai.tools.reconnaissance.nmap import nmap # pylint: disable=import-error # noqa: E501 @@ -46,10 +46,9 @@ ctf_agent = Agent( 6. YOU MUST USE THE flag_discriminator to check the flag when you have it in DECODED 64-bit format 7. Do a initial recon (nmap, service enumeration)""", - functions=[ + tools=[ generic_linux_command, ], - parallel_tool_calls=False ) @@ -79,7 +78,7 @@ network_agent = Agent( 5. Whenever you find a likely candidate flag, call the flag_discriminator agent. 6. DO NOT SEND THE SAME COMMAND OVER AND OVER"""), - functions=[netcat, nmap], + tools=[netcat, nmap], ) @@ -106,7 +105,7 @@ crypto_agent = Agent( 5. Do not generate a plan or verbose output. """ ), - functions=[decode64, strings_command, decode_hex_bytes], + tools=[decode64, strings_command, decode_hex_bytes], ) @@ -132,7 +131,7 @@ listing_agent = Agent( TO SOLVE THE CTF 5. KEEP CALLING THE TOOLS OR THE CTF Leader AGENT UNTIL YOU FIND THE FLAG"""), - functions=[ + tools=[ list_dir, cat_file, find_file, diff --git a/src/cai/cli.py b/src/cai/cli.py index d6d8ded7..d35affed 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -71,7 +71,7 @@ Usage Examples: CAI_AGENT_TYPE="one_tool_agent" CAI_MODEL="qwen2.5:14b" \ CAI_TRACING="false" python3 cai/cli.py - # Run a harder CTF + # Run a harder CTF CTF_NAME="hackableii" CAI_AGENT_TYPE="redteam_agent" \ CTF_INSIDE="False" CAI_MODEL="deepseek/deepseek-chat" \ CAI_TRACING="false" python3 cai/cli.py @@ -113,6 +113,9 @@ from cai.repl.ui.banner import display_banner from cai.repl.ui.prompt import get_user_input from cai.repl.ui.toolbar import get_toolbar_with_refresh +# Import agents-related functions +from cai.agents import get_agent_by_name + # Load environment variables from .env file load_dotenv() @@ -123,22 +126,22 @@ external_client = AsyncOpenAI( set_default_openai_client(external_client) set_tracing_disabled(True) -# llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini') -# llm_model=os.getenv('LLM_MODEL', 'claude-3-7') -llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b') +# # llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini') +# # llm_model=os.getenv('LLM_MODEL', 'claude-3-7') +# llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b') -# For Qwen models, we need to skip system instructions as they're not supported -instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant" +# # For Qwen models, we need to skip system instructions as they're not supported +# instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant" -agent = Agent( - name="Assistant", - instructions=instructions, - model=OpenAIChatCompletionsModel( - model=llm_model, - openai_client=external_client, - ) -) +# agent = Agent( +# name="Assistant", +# instructions=instructions, +# model=OpenAIChatCompletionsModel( +# model=llm_model, +# openai_client=external_client, +# ) +# ) def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=float('inf')): """ @@ -209,7 +212,9 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= return result except Exception as e: print() # Add a newline after any partial output - print(f"\n[Error occurred during streaming: {str(e)}]") + import traceback + tb = traceback.format_exc() + print(f"\n[Error occurred during streaming: {str(e)}]\nLocation: {tb}") return None asyncio.run(process_streamed_response()) @@ -222,11 +227,42 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= except KeyboardInterrupt: break except Exception as e: + import traceback + import sys + exc_type, exc_value, exc_traceback = sys.exc_info() + tb_info = traceback.extract_tb(exc_traceback) + filename, line, func, text = tb_info[-1] console.print(f"[bold red]Error: {str(e)}[/bold red]") + console.print(f"[bold red]Traceback: {tb_info}[/bold red]") -def main(): - run_cai_cli(agent, stream=True) +def main(): + # Get agent type from environment variables or use default + agent_type = os.getenv('CAI_AGENT_TYPE', "one_tool_agent") + + llm_model=os.getenv('LLM_MODEL', 'qwen2.5:14b') + # llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini') + + # For Qwen models, we need to skip system instructions as they're not supported + instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant" + + agent = Agent( + name="Assistant", + instructions=instructions, + model=OpenAIChatCompletionsModel( + model=llm_model, + openai_client=external_client, + ) + ) + + # Get the agent instance by name + agent = get_agent_by_name(agent_type) + + # Enable streaming by default, unless specifically disabled + stream = os.getenv('CAI_STREAM', 'false').lower() != 'false' + + # Run the CLI with the selected agent + run_cai_cli(agent, stream=stream) if __name__ == "__main__": main() \ No newline at end of file diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 934400fe..53791027 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -493,6 +493,8 @@ class Runner: # agent changes, or if the agent loop ends. if current_span is None: handoff_names = [h.agent_name for h in cls._get_handoffs(current_agent)] + # print("current_agent: ", current_agent) + # print("current_agent.tools: ", current_agent.tools) tool_names = [t.name for t in current_agent.tools] if output_schema := cls._get_output_schema(current_agent): output_type_name = output_schema.output_type_name() diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index 5356ad6e..b14ba4ad 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -5,8 +5,10 @@ from cai.tools.common import (run_command, list_shell_sessions, get_session_output, terminate_session) # pylint: disable=import-error # noqa E501 +from cai.sdk.agents import function_tool +@function_tool def generic_linux_command(command: str = "", args: str = "", ctf=None, async_mode: bool = False, diff --git a/uv.lock b/uv.lock index 063fd156..c0cf8098 100644 --- a/uv.lock +++ b/uv.lock @@ -206,6 +206,7 @@ dependencies = [ { name = "dotenv" }, { name = "griffe" }, { name = "litellm" }, + { name = "mako" }, { name = "openai" }, { name = "openinference-instrumentation-openai" }, { name = "prompt-toolkit" }, @@ -251,6 +252,7 @@ requires-dist = [ { name = "dotenv", specifier = ">=0.9.9" }, { name = "griffe", specifier = ">=1.5.6,<2" }, { name = "litellm", specifier = ">=1.63.7" }, + { name = "mako", specifier = ">=1.3.9" }, { name = "numpy", marker = "python_full_version >= '3.10' and extra == 'voice'", specifier = ">=2.2.0,<3" }, { name = "openai", specifier = ">=1.68.2" }, { name = "openinference-instrumentation-openai", specifier = ">=0.1.22" }, @@ -1018,6 +1020,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/0c/149ae51e16d0836d6ddc9aec03a110960227046860dfc68bd3746a23cdf5/litellm-1.63.14-py3-none-any.whl", hash = "sha256:d4c469f5990e142cc23dfa06c3fddd627001928e4df43682001f453af6a1fb51", size = 6964658 }, ] +[[package]] +name = "mako" +version = "1.3.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/62/4f/ddb1965901bc388958db9f0c991255b2c469349a741ae8c9cd8a562d70a6/mako-1.3.9.tar.gz", hash = "sha256:b5d65ff3462870feec922dbccf38f6efb44e5714d7b593a656be86663d8600ac", size = 392195 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/83/de0a49e7de540513f53ab5d2e105321dedeb08a8f5850f0208decf4390ec/Mako-1.3.9-py3-none-any.whl", hash = "sha256:95920acccb578427a9aa38e37a186b1e43156c87260d7ba18ca63aa4c7cbd3a1", size = 78456 }, +] + [[package]] name = "markdown" version = "3.7"