diff --git a/src/cai/agents/memory_analysis_agent.py b/src/cai/agents/memory_analysis_agent.py index 1448e979..19c4a74d 100644 --- a/src/cai/agents/memory_analysis_agent.py +++ b/src/cai/agents/memory_analysis_agent.py @@ -1,6 +1,8 @@ """Memory Analysis and Manipulation Agent""" import os -from cai.sdk.agents import Agent # pylint: disable=import-error +from dotenv import load_dotenv +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error +from openai import AsyncOpenAI from cai.util import load_prompt_template # Add this import from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501 run_ssh_command_with_credentials @@ -17,6 +19,7 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error execute_code ) +load_dotenv() # Prompts memory_analysis_agent_system_prompt = load_prompt_template("prompts/memory_analysis_agent.md") @@ -38,7 +41,9 @@ memory_analysis_agent = Agent( description="""Agent for runtime memory analysis and manipulation. Specializes in process memory examination, monitoring, and modification for security assessment, vulnerability discovery, and runtime behavior analysis.""", - model=os.getenv('CAI_MODEL', "qwen2.5:14b"), tools=functions, - + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "alias0"), + openai_client=AsyncOpenAI(), + ) ) diff --git a/src/cai/agents/reporter.py b/src/cai/agents/reporter.py index cad39dbc..914ec28a 100644 --- a/src/cai/agents/reporter.py +++ b/src/cai/agents/reporter.py @@ -1,6 +1,8 @@ """Reporter Agent - Creates professional security assessment reports""" import os -from cai.sdk.agents import Agent # pylint: disable=import-error +from dotenv import load_dotenv +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error +from openai import AsyncOpenAI from cai.util import load_prompt_template # Add this import from cai.tools.reconnaissance.generic_linux_command import ( # pylint: disable=import-error # noqa: E501 @@ -11,6 +13,7 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error execute_code ) +load_dotenv() # Prompts reporting_agent_system_prompt = load_prompt_template("prompts/system_reporting_agent.md") @@ -26,7 +29,9 @@ reporting_agent = Agent( name="reporting agent", instructions=reporting_agent_system_prompt, description="""Agent that generates reports in html.""", - model=os.getenv('CAI_MODEL', "qwen2.5:14b"), tools=functions, - + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "alias0"), + openai_client=AsyncOpenAI(), + ) ) diff --git a/src/cai/agents/reverse_engineering_agent.py b/src/cai/agents/reverse_engineering_agent.py index 4ecf9051..6e490428 100644 --- a/src/cai/agents/reverse_engineering_agent.py +++ b/src/cai/agents/reverse_engineering_agent.py @@ -1,6 +1,8 @@ """Reverse Engineering and Binary Analysis Agent""" import os -from cai.sdk.agents import Agent # pylint: disable=import-error +from dotenv import load_dotenv +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error +from openai import AsyncOpenAI from cai.util import load_prompt_template # Add this import from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501 run_ssh_command_with_credentials @@ -17,6 +19,7 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error execute_code ) +load_dotenv() # Prompts reverse_engineering_agent_system_prompt = load_prompt_template("prompts/reverse_engineering_agent.md") @@ -39,7 +42,9 @@ reverse_engineering_agent = Agent( Specializes in firmware analysis, binary disassembly, decompilation, and vulnerability discovery using tools like Ghidra, Binwalk, and various binary analysis utilities.""", - model=os.getenv('CAI_MODEL', "qwen2.5:14b"), tools=functions, - + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "alias0"), + openai_client=AsyncOpenAI(), + ) ) diff --git a/src/cai/agents/subghz_sdr_agent.py b/src/cai/agents/subghz_sdr_agent.py index 96b89404..534374b9 100644 --- a/src/cai/agents/subghz_sdr_agent.py +++ b/src/cai/agents/subghz_sdr_agent.py @@ -1,6 +1,8 @@ """Sub-GHz Radio Frequency Analysis Agent using HackRF One""" import os -from cai.sdk.agents import Agent # pylint: disable=import-error +from dotenv import load_dotenv +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error +from openai import AsyncOpenAI from cai.util import load_prompt_template # Add this import from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501 run_ssh_command_with_credentials @@ -17,6 +19,7 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error execute_code ) +load_dotenv() # Prompts subghz_agent_system_prompt = load_prompt_template("prompts/subghz_agent.md") @@ -38,7 +41,9 @@ subghz_sdr_agent = Agent( description="""Agent for sub-GHz radio frequency analysis using HackRF One. Specializes in signal capture, replay, and protocol analysis for IoT, automotive, industrial, and wireless security applications.""", - model=os.getenv('CAI_MODEL', "qwen2.5:14b"), tools=functions, - + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "alias0"), + openai_client=AsyncOpenAI(), + ) ) diff --git a/src/cai/agents/wifi_security_tester.py b/src/cai/agents/wifi_security_tester.py index 535ba5f6..089a44b2 100644 --- a/src/cai/agents/wifi_security_tester.py +++ b/src/cai/agents/wifi_security_tester.py @@ -1,6 +1,8 @@ """Wi-Fi Security Testing Agent""" import os -from cai.sdk.agents import Agent # pylint: disable=import-error +from dotenv import load_dotenv +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel # pylint: disable=import-error +from openai import AsyncOpenAI from cai.util import load_prompt_template # Add this import from cai.tools.command_and_control.sshpass import ( # pylint: disable=import-error # noqa: E501 run_ssh_command_with_credentials @@ -17,6 +19,7 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error execute_code ) +load_dotenv() # Prompts wifi_security_agent_system_prompt = load_prompt_template("prompts/wifi_security_agent.md") @@ -37,7 +40,9 @@ wifi_security_agent = Agent( instructions=wifi_security_agent_system_prompt, description="""Agent for Wi-Fi network security testing and penetration. Specializes in wireless attacks, password recovery, and communication disruption.""", - model=os.getenv('CAI_MODEL', "qwen2.5:14b"), tools=functions, - + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "alias0"), + openai_client=AsyncOpenAI(), + ) )