RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches 2025-03-27 19:36:58 +01:00
parent 326c3274fc
commit 1a22013bda
6 changed files with 132 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

81
src/cai/agents/__init__.py
View File

@ -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

11
src/cai/agents/basic.py
View File

@ -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,

70
src/cai/cli.py
View File

@ -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()

2
src/cai/sdk/agents/run.py
View File

@ -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()

2
src/cai/tools/reconnaissance/generic_linux_command.py
View File

@ -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,

14
uv.lock
View File

@ -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"