diff --git a/examples/cai/simple_one_tool_test.py b/examples/cai/simple_one_tool_test.py index 5ca90dca..669ddd0d 100644 --- a/examples/cai/simple_one_tool_test.py +++ b/examples/cai/simple_one_tool_test.py @@ -19,7 +19,7 @@ from cai.sdk.agents.models._openai_shared import set_use_responses_by_default # Load environment variables load_dotenv() -#set_tracing_disabled(True) +#set_tracing_disabled(True) #disable tracing or OPENAI_AGENTS_DISABLE_TRACING=1 # NOTE: This is needed when using LiteLLM Proxy Server # diff --git a/src/cai/agents/blue_teamer.py b/src/cai/agents/blue_teamer.py index d351521a..71583bfa 100644 --- a/src/cai/agents/blue_teamer.py +++ b/src/cai/agents/blue_teamer.py @@ -4,8 +4,10 @@ SSH_HOST SSH_USER """ import os -from cai.sdk.agents import Agent # pylint: disable=import-error +from openai import AsyncOpenAI +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error from cai.util import load_prompt_template +from dotenv import load_dotenv from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501 run_ssh_command_with_credentials ) @@ -14,7 +16,7 @@ from cai.tools.reconnaissance.generic_linux_command import ( # pylint: disable= generic_linux_command ) from cai.tools.web.search_web import ( # pylint: disable=import-error # noqa: E501 - search_web + make_web_search_with_explanation ) from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error # noqa: E501 @@ -23,23 +25,26 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error # Prompts blueteam_agent_system_prompt = load_prompt_template("prompts/system_blue_team_agent.md") -# Define functions list based on available API keys -functions = [ +# Define tools list based on available API keys +tools = [ generic_linux_command, run_ssh_command_with_credentials, execute_code, ] -# Add search_web function if PERPLEXITY_API_KEY environment variable is set +load_dotenv() +# Add search_web tools if PERPLEXITY_API_KEY environment variable is set if os.getenv('PERPLEXITY_API_KEY'): - functions.append(search_web) + tools.append(make_web_search_with_explanation) blueteam_agent = Agent( name="Blue Team Agent", instructions=blueteam_agent_system_prompt, description="""Agent that specializes in system defense and security monitoring. Expert in cybersecurity protection and incident response.""", - model=os.getenv('CAI_MODEL', "qwen2.5:14b"), - functions=functions, - parallel_tool_calls=False, + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ), + tools=tools, ) diff --git a/src/cai/agents/codeagent.py b/src/cai/agents/codeagent.py index a71bb78e..bd384744 100644 --- a/src/cai/agents/codeagent.py +++ b/src/cai/agents/codeagent.py @@ -31,8 +31,9 @@ from cai.agents.meta.local_python_executor import ( fix_final_answer_code, truncate_content, ) -from cai.sdk.agents import Agent, Result - +from cai.sdk.agents import Agent, Result, OpenAIChatCompletionsModel +from dotenv import load_dotenv +from openai import AsyncOpenAI class CodeAgentException(Exception): """Base exception class for CodeAgent-related errors.""" @@ -185,7 +186,7 @@ class CodeAgent(Agent): name: str = "CodeAgent", model: str = "qwen2.5:14b", instructions: Union[str, Callable[[], str]] = None, - functions: List[Callable] = None, + tools: List[Callable] = None, additional_authorized_imports: Optional[List[str]] = None, description: str = """Agent focused on writing and executing code. State-of-the-art in code production.""", @@ -201,7 +202,7 @@ class CodeAgent(Agent): name: Name of the agent model: Model to use for the agent instructions: Instructions for the agent - functions: List of functions available to the agent + tools: List of tools available to the agent additional_authorized_imports: List of additional imports to allow max_print_outputs_length: Maximum length of print outputs reasoning_effort: Level of reasoning effort (low, medium, high) diff --git a/src/cai/agents/mail.py b/src/cai/agents/mail.py index ba5435a0..88939c7a 100644 --- a/src/cai/agents/mail.py +++ b/src/cai/agents/mail.py @@ -2,9 +2,13 @@ Mail Agent module for checking email configuration security. """ +import os +from openai import AsyncOpenAI import dns.resolver # pylint: disable=import-error -from cai.sdk.agents import Agent +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel from cai.tools.misc.cli_utils import execute_cli_command +from cai.sdk.agents import function_tool + def get_txt_record(domain, record_type='TXT'): @@ -58,7 +62,7 @@ def check_dkim(domain: str, selector: str = "default"): return txt_records[0] return None - +@function_tool def check_mail_spoofing_vulnerability( domain: str, dkim_selector: str = "default") -> dict: @@ -100,7 +104,6 @@ def check_mail_spoofing_vulnerability( dns_smtp_agent = Agent( - model="gpt-4o", name="DNS_SMTP_Agent", description="Agent focused on assessing spoofing DMARC.", instructions=( @@ -110,5 +113,9 @@ dns_smtp_agent = Agent( "detailed reports. Use execute_cli_command for basic scans. " "USE ONLY TOOL CALLS, DONT RETURN REASON." ), - tools=[check_mail_spoofing_vulnerability, execute_cli_command] + tools=[check_mail_spoofing_vulnerability, execute_cli_command], + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) ) diff --git a/src/cai/agents/memory.py b/src/cai/agents/memory.py index 01bd5645..e1f4b6bd 100644 --- a/src/cai/agents/memory.py +++ b/src/cai/agents/memory.py @@ -82,7 +82,7 @@ Environment Variables enabling the episodic memory store """ import os -from cai.sdk.agents import Agent +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel from cai.tools.misc.rag import add_to_memory_semantic, add_to_memory_episodic from cai.rag.vector_db import get_previous_memory @@ -145,32 +145,35 @@ QUERY_PROMPT = """INSTRUCTIONS: """ semantic_builder = Agent( - model=model, name="Semantic_Builder", instructions=ADD_MEMORY_PROMPT, description="""Agent that stores semantic memories from security assessments and CTF exercises in semantic format.""", tool_choice="required", temperature=0, - functions=[add_to_memory_semantic], - parallel_tool_calls=True + tools=[add_to_memory_semantic], + model=OpenAIChatCompletionsModel( + model=model_name, + openai_client=AsyncOpenAI(), + ) ) episodic_builder = Agent( - model=model, name="Episodic_Builder", instructions=ADD_MEMORY_PROMPT, description="""Agent that stores episodic memories from security assessments and CTF exercises in episodic format.""", tool_choice="required", temperature=0, - functions=[add_to_memory_episodic], - parallel_tool_calls=True + tools=[add_to_memory_episodic], + model=OpenAIChatCompletionsModel( + model=model_name, + openai_client=AsyncOpenAI(), + ) ) query_agent = Agent( - model=model, name="Query_Agent", description="""Agent that queries the memory system to retrieve relevant historical information from previous security assessments @@ -178,4 +181,8 @@ query_agent = Agent( instructions=QUERY_PROMPT, tool_choice="required", temperature=0, + model=OpenAIChatCompletionsModel( + model=model_name, + openai_client=AsyncOpenAI(), + ) ) diff --git a/src/cai/tools/misc/cli_utils.py b/src/cai/tools/misc/cli_utils.py index e8206613..46dcd759 100644 --- a/src/cai/tools/misc/cli_utils.py +++ b/src/cai/tools/misc/cli_utils.py @@ -3,8 +3,9 @@ CLI utilities module for executing shell commands and processing their output. """ from cai.tools.common import run_command # pylint: disable=E0401 +from cai.sdk.agents import function_tool - +@function_tool def execute_cli_command(command: str) -> str: """ Execute a CLI command and return the output. diff --git a/src/cai/tools/misc/rag.py b/src/cai/tools/misc/rag.py index ab30c4af..7d4b16f2 100644 --- a/src/cai/tools/misc/rag.py +++ b/src/cai/tools/misc/rag.py @@ -5,12 +5,12 @@ querying and adding data to vector databases. import os import uuid from cai.rag.vector_db import QdrantConnector - +from cai.sdk.agents import function_tool # CTF BASED MEMORY collection_name = os.getenv('CAI_MEMORY_COLLECTION', "default") - +@function_tool def query_memory(query: str, top_k: int = 3, **kwargs) -> str: # pylint: disable=unused-argument,line-too-long # noqa: E501 """ Query memory to retrieve relevant context. From Previous CTFs executions. @@ -42,7 +42,7 @@ def query_memory(query: str, top_k: int = 3, **kwargs) -> str: # pylint: disabl except Exception: # pylint: disable=broad-exception-caught return results - +@function_tool def add_to_memory_episodic(texts: str, step: int = 0, **kwargs) -> str: # pylint: disable=unused-argument,line-too-long # noqa: E501 """ This is a persistent memory to add relevant context to our memory. @@ -76,7 +76,7 @@ def add_to_memory_episodic(texts: str, step: int = 0, **kwargs) -> str: # pylin except Exception as e: # pylint: disable=broad-exception-caught return f"Error adding documents to vector database: {str(e)}" - +@function_tool def add_to_memory_semantic(texts: str, step: int = 0, **kwargs) -> str: # pylint: disable=unused-argument,line-too-long # noqa: E501 """ This is a persistent memory to add relevant context to our memory. diff --git a/tests/agents/test_agent_inferece.py b/tests/agents/test_agent_inferece.py new file mode 100644 index 00000000..8b8ed122 --- /dev/null +++ b/tests/agents/test_agent_inferece.py @@ -0,0 +1,101 @@ +import os +import pytest +import pytest +from cai.sdk.agents import Runner +from cai.agents import get_agent_by_name +from cai.agents.bug_bounter import bug_bounter_agent +from cai.agents.flag_discriminator import flag_discriminator +from cai.agents.mail import dns_smtp_agent +from cai.agents.one_tool import one_tool_agent +from cai.agents.red_teamer import redteam_agent +from cai.agents.thought import thought_agent + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_blue_team_agent_inference(): + """ + Non-streaming inference test for the blueteam_agent. + Verifies that Runner.run returns the final output containing the expected text. + """ + prompt = "monitor login attempts for suspicious activity" + result = await Runner.run(get_agent_by_name("blueteam_agent"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert "login" in final_output.lower(), f"Expected 'login' in output, got: {final_output}" + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_bug_bounter_agent_inference(): + """ + Non-streaming inference test for the bug_bounter_agent. + """ + prompt = "find vulnerabilities in web application sample.com" + result = await Runner.run(get_agent_by_name("bug_bounter_agent"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert "sample.com" in final_output.lower(), f"Expected 'sample.com' in output, got: {final_output}" + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_flag_discriminator_inference(): + """ + Non-streaming inference test for the flag_discriminator. + """ + prompt = "is 'flag{test_flag}' a valid CTF flag?" + result = await Runner.run(get_agent_by_name("flag_discriminator"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert "valid" in final_output.lower(), f"Expected 'valid' in output, got: {final_output}" + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_dns_smtp_agent_inference(): + """ + Non-streaming inference test for the dns_smtp_agent. + """ + prompt = "check DKIM record for example.com" + result = await Runner.run(get_agent_by_name("dns_smtp_agent"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert "dkim" in final_output.lower(), f"Expected 'dkim' in output, got: {final_output}" + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_one_tool_agent_inference(): + """ + Non-streaming inference test for the one_tool_agent. + Verifies that Runner.run returns the final output containing the expected text. + """ + prompt = "use nmap to scan 192.168.2.0/24" + # Run the agent in non-streaming mode + result = await Runner.run(get_agent_by_name("one_tool_agent"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert "generic_linux_command" in final_output.lower(), f"Expected 'generic_linux_command' in output, got: {final_output}" + assert "nmap" in final_output.lower(), f"Expected 'nmap' in output, got: {final_output}" + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_red_team_agent_inference(): + """ + Non-streaming inference test for the redteam_agent. + """ + prompt = "perform penetration test on example.com ssh service" + result = await Runner.run(get_agent_by_name("redteam_agent"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert "ssh" in final_output.lower(), f"Expected 'ssh' in output, got: {final_output}" + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_thought_agent_inference(): + """ + Non-streaming inference test for the thought_agent. + """ + prompt = "what are the next steps for network security assessment?" + result = await Runner.run(get_agent_by_name("thought_agent"), prompt) + final_output = result.final_output or "" + assert final_output, "Expected non-empty final output" + assert ( + "plan" in final_output.lower() or "analysis" in final_output.lower() + ), f"Expected 'plan' or 'analysis' in output, got: {final_output}" \ No newline at end of file