fix tools

examples/cai/simple_one_tool_test.py

@@ -11,7 +11,7 @@ import asyncio
 import json
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
-from cai.sdk.agents import Runner, set_default_openai_client
+from cai.sdk.agents import Runner, set_default_openai_client, set_tracing_disabled
 from cai.agents import get_agent_by_name
 from cai.util import fix_litellm_transcription_annotations, color, cli_print_agent_messages
 from cai.sdk.agents.models._openai_shared import set_use_responses_by_default
@@ -19,6 +19,7 @@ from cai.sdk.agents.models._openai_shared import set_use_responses_by_default
 
 # Load environment variables
 load_dotenv()
+#set_tracing_disabled(True)
 
 # NOTE: This is needed when using LiteLLM Proxy Server
 #

src/cai/agents/__init__.py

@@ -123,8 +123,8 @@ def get_available_agents() -> Dict[str, Agent]:  # pylint: disable=R0912  # noqa
                         agent_name = attr_name
                         if agent_name not in agents_to_display:
                             agents_to_display[agent_name] = attr
-            except (ImportError, AttributeError):
-                pass
+            except (ImportError, AttributeError) as e:
+                print(f"Error importing {agent_name}: {e}")
 
     return agents_to_display
 

src/cai/agents/basic.py

@@ -1,146 +0,0 @@
-"""
-Collection of basic agents
-"""
-
-import os
-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
-
-# First party imports
-from cai.tools.reconnaissance.filesystem import (
-    list_dir,
-    cat_file,
-    find_file,
-    pwd_command
-)
-from cai.tools.reconnaissance.generic_linux_command import (
-    generic_linux_command
-)
-from cai.tools.reconnaissance.crypto_tools import (
-    decode64,
-    strings_command,
-    decode_hex_bytes
-)
-model = os.getenv('CAI_MODEL', "qwen2.5:14b")
-
-
-ctf_agent = Agent(
-    model=model,
-    name="CTF agent",
-    description="""Agent focused on CTF challenges.
-                   Expert in cybersecurity and exploitation.
-                   Specializes in executing Linux commands efficiently.""",
-    instructions="""You are a Cybersecurity expert Leader facing a CTF
-                    challenge.
-                    INSTRUCTIONS:
-                    1. Execute the generic_linux_command tool without any
-                    explanation.
-                    2. Be efficient and strategic when executing commands.
-                    3. Never assume the flag format - it could be any string
-                    4. NEVER GIVE UP. KEEP EXECUTING TOOLS UNTIL YOU FIND
-                    THE FLAG.
-                    5. IF YOU ENCOUNTER A PROBLEM, IT IS HIGHLY LIKELY
-                    THAT YOU ARE ASSUMING WRONG FILENAMES, PORTS,
-                    IP ADDRESSES, ETC. PAY ATTENTION TO YOUR HALLUCINATIONS.
-                    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)""",
-    tools=[
-        generic_linux_command,
-    ],
-)
-
-
-def transfer_to_ctf_agent(**kwargs):  # pylint: disable=W0613
-    """Transfer to ctf agent.
-    Accepts any keyword arguments but ignores them."""
-    return ctf_agent
-
-
-network_agent = Agent(
-    model=model,
-    name="Network Specialist",
-    description="""Agent specialized in network reconnaissance.
-                   Expert in using nmap and netcat for service discovery.
-                   Focuses on identifying network vulnerabilities.""",
-    instructions=("""INSTRUCTIONS:
-        1. You specialize in using nmap and netcat.
-        INSTRUCTIONS:
-        1. You always use real information of the
-        CTF challenge as arguments for the tools.
-        If you lack of that information, you find another
-        way to get it.
-        2. Keep calling agents until you find the flag.
-        3. Do not be verbose, just call the tools.
-        4. YOU NEVER USE PLACEHOLDERS, AS THE TOOLS WILL
-        NOT WORK WITH THEM.
-        5. Whenever you find a likely candidate flag, call
-        the flag_discriminator agent.
-        6. DO NOT SEND THE SAME COMMAND OVER AND OVER"""),
-    tools=[netcat, nmap],
-)
-
-
-def transfer_to_network_agent(**kwargs):  # pylint: disable=W0613
-    """ Transfer to network agent.
-    Accepts any keyword arguments but ignores them."""
-    return network_agent
-
-
-crypto_agent = Agent(
-    model=model,
-    name="Cryptography Expert",
-    description="""Agent specialized in cryptographic analysis.
-                   Expert in decoding and analyzing encrypted data.
-                   Focuses on breaking cryptographic challenges in CTFs.""",
-    instructions=("""
-        INSTRUCTIONS:
-        1. Your expertise is in decryption and
-                  data analysis.
-        2. Use your tools to decode and analyze
-                  encrypted data.
-        3. DO NOT BE VERBOSE, JUST CALL THE TOOLS AS NEEDED
-        4. Use your tools to solve the CTF challenge.
-        5. Do not generate a plan or verbose output.
-        """
-                  ),
-    tools=[decode64, strings_command, decode_hex_bytes],
-)
-
-
-def transfer_to_crypto_agent(**kwargs):  # pylint: disable=W0613
-    """ Transfer to crypto agent.
-    Accepts any keyword arguments but ignores them."""
-    return crypto_agent
-
-
-listing_agent = Agent(
-    model=model,
-    name="File Explorer",
-    description="""Agent specialized in filesystem navigation and exploration.
-                   Expert in finding and analyzing files and directories.
-                   Focuses on discovering hidden files and sensitive data.""",
-    instructions=("""INSTRUCTIONS:
-        1. You specialize in navigating file systems.
-        2. Use your tools to list, read, and find files in
-        a directory.
-        3. DO NOT USE THE SAME TOOL FOR THE SAME FILE OR
-        DIRECTORY more than 3 times.
-        4. DO NOT BE VERBOSE, JUST CALL THE TOOLS AS NEEDED
-        TO SOLVE THE CTF
-        5. KEEP CALLING THE TOOLS OR THE CTF Leader AGENT
-        UNTIL YOU FIND THE FLAG"""),
-    tools=[
-        list_dir,
-        cat_file,
-        find_file,
-        pwd_command,
-        generic_linux_command],
-)
-
-
-def transfer_to_listing_agent(**kwargs):  # pylint: disable=W0613
-    """ Transfer to listing agent.
-    Accepts any keyword arguments but ignores them."""
-    return listing_agent

src/cai/agents/bug_bounter.py

@@ -1,7 +1,8 @@
 """Red Team Base Agent"""
 import os
 from dotenv import load_dotenv
-from cai.sdk.agents import Agent  # pylint: disable=import-error
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
+from openai import AsyncOpenAI
 from cai.util import load_prompt_template
 from cai.tools.command_and_control.sshpass import (  # pylint: disable=import-error # noqa: E501
     run_ssh_command_with_credentials
@@ -26,8 +27,8 @@ from cai.tools.reconnaissance.shodan import (  # pylint: disable=import-error #
 load_dotenv()
 # Prompts
 bug_bounter_system_prompt = load_prompt_template("prompts/system_bug_bounter.md")
-# Define functions list based on available API keys
-functions = [
+# Define tools list based on available API keys
+tools = [
     generic_linux_command,
     execute_code,
     shodan_search,
@@ -35,13 +36,17 @@ functions = [
 ]
 
 if os.getenv('GOOGLE_SEARCH_API_KEY') and os.getenv('GOOGLE_SEARCH_CX'):
-    functions.append(make_google_search)
+    tools.append(make_google_search)
 
 bug_bounter_agent = Agent(
     name="Bug Bounter",
     instructions=bug_bounter_system_prompt,
     description="""Agent that specializes in bug bounty hunting and vulnerability discovery.
                    Expert in web security, API testing, and responsible disclosure.""",
-    model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
-    tools=functions,
+    tools=tools,
+    model=OpenAIChatCompletionsModel(
+        model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
+        openai_client=AsyncOpenAI(),
+    )
+   
 )

src/cai/agents/flag_discriminator.py

@@ -1,8 +1,9 @@
 """
-CTF Flag Discriminator Agent
+CTF Flag Discriminator Agent with test
 """
 import os
 from cai.sdk.agents import Agent
+from openai import AsyncOpenAI
 
 model = os.getenv('CAI_MODEL', "qwen2.5:14b")
 flag_discriminator = Agent(
@@ -10,15 +11,14 @@ flag_discriminator = Agent(
     name="Flag discriminator",
     description="Agent focused on extracting the flag from the output",
     instructions="""INSTRUCTIONS:
-                1. You are an agent tailored to extract the flag
-                from the output.
+                1. You are an agent tailored to extract the flag from a given output.
                 2. ONLY return the flag, nothing else.
-                3. IF THE FLAG IS NOT FOUND, CALL THE ctf_agent
-                again to CONTINUE the search""",
+                3. Flags can be in any format (not always CTF{...}).
+                4. If you do not find a flag, call `ctf_agent` to continue investigating.
+                """,
 )
 
-
-# TRANSFER FUNCTIONS
+# Transfer Function
 def transfer_to_flag_discriminator(**kwargs):  # pylint: disable=W0613
     """Transfer flag discriminator.
     Accepts any keyword arguments but ignores them."""

src/cai/agents/red_teamer.py

@@ -1,10 +1,11 @@
 """Red Team Base Agent"""
 import os
 from dotenv import load_dotenv
-from cai.sdk.agents import Agent  # pylint: disable=import-error
-from cai.tools.command_and_control.sshpass import (  # pylint: disable=import-error # noqa: E501
-    run_ssh_command_with_credentials
-)
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
+from openai import AsyncOpenAI
+# from cai.tools.command_and_control.sshpass import (  # pylint: disable=import-error # noqa: E501
+#     run_ssh_command_with_credentials
+# )
 
 from cai.tools.reconnaissance.generic_linux_command import (  # pylint: disable=import-error # noqa: E501
     generic_linux_command
@@ -18,25 +19,36 @@ from cai.tools.reconnaissance.exec_code import (  # pylint: disable=import-error
 )
 from cai.util import load_prompt_template
 
+load_dotenv()
+model_name = os.getenv("CAI_MODEL", "qwen2.5:14b")
 # Prompts
 redteam_agent_system_prompt = load_prompt_template("prompts/system_red_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,
+    #run_ssh_command_with_credentials,
     execute_code,
 ]
 
 # Add make_web_search_with_explanation function if PERPLEXITY_API_KEY environment variable is set
 if os.getenv('PERPLEXITY_API_KEY'):
-    functions.append(make_web_search_with_explanation)
+    tools.append(make_web_search_with_explanation)
     
 
 redteam_agent = Agent(
     name="Red Team Agent",
+    description="""Agent that mimics a red teamer in a security assessment.
+                   Expert in cybersecurity, recon, and exploitation.""",
     instructions=redteam_agent_system_prompt,
-    description="""Agent that mimic pentester/red teamer in a security assessment.
-                   Expert in cybersecurity and exploitation.""",
-    model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
-    tools=functions,
-)
+    tools=tools,
+    model=OpenAIChatCompletionsModel(
+        model=model_name,
+        openai_client=AsyncOpenAI(),
+    ),
+)
+
+# Transfer function
+def transfer_to_redteam_agent(**kwargs):  # pylint: disable=W0613
+    """Transfer to red team agent.
+    Accepts any keyword arguments but ignores them."""
+    return redteam_agent

src/cai/agents/thought.py

@@ -6,7 +6,8 @@ using reasoner as a tool call
 support meta agent may better @cai.sdk.agents.meta.reasoner_support
 """
 from cai.tools.misc.reasoning import thought
-from cai.sdk.agents import Agent  # pylint: disable=import-error
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel  # pylint: disable=import-error
+from openai import AsyncOpenAI
 from cai.util import load_prompt_template
 import os
 
@@ -15,7 +16,10 @@ thought_agent_system_prompt = load_prompt_template("prompts/system_thought_route
 # Thought Process Agent for analysis and planning
 thought_agent = Agent(
     name="ThoughAgent",
-    model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
+    model=OpenAIChatCompletionsModel(
+        model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
+        openai_client=AsyncOpenAI(),
+    ),
     description="""Agent focused on analyzing and planning the next steps
                    in a security assessment or CTF challenge.""",
     instructions=thought_agent_system_prompt,

src/cai/tools/command_and_control/sshpass.py

@@ -12,8 +12,11 @@ It handles privilege escalation very well and is autonomous regarding SSH passwo
 something that hasn't been seen in other cybersecurity frameworks yet (Feb 2025)
 """  # noqa: E501
 
-from cai.tools.misc.cli_utils import execute_cli_command  # pylint: disable=E0401 # noqa: E501
+from cai.tools.common import run_command  # pylint: disable=E0401 # noqa: E501
+from cai.sdk.agents import function_tool
 
+
+@function_tool
 def run_ssh_command_with_credentials(
         host: str,
         username: str,
@@ -43,4 +46,4 @@ def run_ssh_command_with_credentials(
         f"{username}@{host} -p {port} "
         f"'{escaped_command}'"
     )
-    return execute_cli_command(ssh_command)
+    return run_command(ssh_command)

src/cai/tools/misc/code_interpreter.py

@@ -5,8 +5,10 @@ Module for executing Python code and capturing its output.
 import io
 import sys
 from typing import Dict
+from cai.sdk.agents import function_tool
 
 
+@function_tool
 def execute_python_code(code: str, context: Dict = None) -> str:
     """
     Execute Python code and return the output.

src/cai/tools/misc/reasoning.py

@@ -3,8 +3,10 @@ Reasoning tools module for tracking thoughts, findings and analysis
 Provides utilities for recording and retrieving key information discovered
 during CTF progression.
 """
+from cai.sdk.agents import function_tool
 
 
+@function_tool
 def thought(breakdowns: str = "", reflection: str = "",  # pylint: disable=too-many-arguments  # noqa: E501
             action: str = "", next_step: str = "", key_clues: str = "",
             ctf=None) -> str:  # pylint: disable=unused-argument  # noqa: E501
@@ -34,7 +36,7 @@ def thought(breakdowns: str = "", reflection: str = "",  # pylint: disable=too-m
         output.append(f"Key Clues: {key_clues}")
     return "\n".join(output)
 
-
+@function_tool
 def write_key_findings(findings: str) -> str:
     """
     Write key findings to a state.txt file to track important CTF details.
@@ -58,7 +60,7 @@ def write_key_findings(findings: str) -> str:
     except OSError as e:
         return f"Error writing to state.txt: {str(e)}"
 
-
+@function_tool
 def read_key_findings() -> str:
     """
     Read key findings from the state.txt file to retrieve important data

src/cai/tools/others/scripting.py

@@ -5,8 +5,10 @@ This is used to create and execute a script in python
 # run_command is used in other parts of the codebase that import this module
 
 # pylint: disable=too-many-locals,too-many-branches
+from cai.sdk.agents import function_tool
 
 
+@function_tool
 def scripting_tool(
         command: str = "",
         args: str = "",

src/cai/tools/reconnaissance/crypto_tools.py

@@ -2,6 +2,9 @@
  Here are crypto tools
 """
 from cai.tools.common import run_command
+from cai.sdk.agents import function_tool
+
+
 
 # # URLDecodeTool
 # # HexDumpTool
@@ -9,7 +12,7 @@ from cai.tools.common import run_command
 # # ROT13DecodeTool
 # # BinaryAnalysisTool
 
-
+@function_tool
 def strings_command(file_path: str, ctf=None) -> str:
     """
     Extract printable strings from a binary file.
@@ -24,7 +27,7 @@ def strings_command(file_path: str, ctf=None) -> str:
     command = f'strings {file_path}'
     return run_command(command, ctf=ctf)
 
-
+@function_tool
 def decode64(input_data: str, ctf=None) -> str:
     """
     Decode a base64-encoded string.
@@ -39,7 +42,7 @@ def decode64(input_data: str, ctf=None) -> str:
     command = f"base64 --decode {input_data}"
     return run_command(command, ctf=ctf)
 
-
+@function_tool
 def decode_hex_bytes(input_data: str) -> str:
     """
     Decode a string of hex bytes into ASCII text.

src/cai/tools/reconnaissance/curl.py

@@ -2,8 +2,9 @@
 Here are the curl tools.
 """
 from cai.tools.common import run_command  # pylint: disable=import-error
+from cai.sdk.agents import function_tool
 
-
+@function_tool
 def curl(args: str = "", target: str = "", ctf=None) -> str:
     """
     A simple curl tool to make HTTP requests to a specified target.

src/cai/tools/reconnaissance/exec_code.py

@@ -2,7 +2,10 @@
 Tool for executing code via LLM tool calls.
 """
 from cai.tools.common import run_command  # pylint: disable=import-error
+from cai.sdk.agents import function_tool
 
+
+@function_tool
 def execute_code(code: str = "", language: str = "python",
                 filename: str = "exploit", timeout: int = 100, ctf=None) -> str:
     """

src/cai/tools/reconnaissance/filesystem.py

@@ -3,8 +3,9 @@ Here are the CLI tools for executing commands.
 """
 
 from cai.tools.common import run_command  # pylint: disable=E0401
+from cai.sdk.agents import function_tool
 
-
+@function_tool
 def list_dir(path: str, args: str = "", ctf=None) -> str:
     """
     List the contents of a directory.
@@ -19,7 +20,7 @@ def list_dir(path: str, args: str = "", ctf=None) -> str:
     command = f'ls {path} {args}'
     return run_command(command, ctf=ctf)
 
-
+@function_tool
 def cat_file(file_path: str, args: str = "", ctf=None) -> str:
     """
     Display the contents of a file.
@@ -44,7 +45,7 @@ def cat_file(file_path: str, args: str = "", ctf=None) -> str:
 # FilePermissionsTool
 # FileCompressionTool
 
-
+@function_tool
 def pwd_command(ctf=None) -> str:
     """
     Retrieve the current working directory.
@@ -55,7 +56,7 @@ def pwd_command(ctf=None) -> str:
     command = 'pwd'
     return run_command(command, ctf=ctf)
 
-
+@function_tool
 def find_file(file_path: str, args: str = "", ctf=None) -> str:
     """
     Find a file in the filesystem.

src/cai/tools/reconnaissance/netcat.py

@@ -2,8 +2,9 @@
  Here are the tools for netcat command
 """
 from cai.tools.common import run_command   # pylint: disable=import-error
+from cai.sdk.agents import function_tool
 
-
+@function_tool
 def netcat(host: str, port: int, data: str = '',
            args: str = '', ctf=None) -> str:
     """

src/cai/tools/reconnaissance/netstat.py

@@ -3,8 +3,9 @@
 Netstat tool
 """
 from cai.tools.common import run_command   # pylint: disable=import-error
+from cai.sdk.agents import function_tool
 
-
+@function_tool
 def netstat(args: str = '', ctf=None) -> str:
     """
     netstat tool to list all listening ports and their associated programs.

src/cai/tools/reconnaissance/nmap.py

@@ -3,8 +3,9 @@ Here are the nmap tools.
 """
 
 from cai.tools.common import run_command  # pylint: disable=E0401
+from cai.sdk.agents import function_tool
 
-
+@function_tool
 def nmap(args: str, target: str, ctf=None) -> str:
     """
     A simple nmap tool to scan a specified target.

src/cai/tools/reconnaissance/shodan.py

@@ -8,8 +8,10 @@ import os
 import requests
 from typing import Dict, List, Optional, Any
 from dotenv import load_dotenv
+from cai.sdk.agents import function_tool
 
 
+@function_tool
 def shodan_search(query: str, limit: int = 10) -> str:
     """
     Search Shodan for information based on the provided query.
@@ -41,7 +43,7 @@ def shodan_search(query: str, limit: int = 10) -> str:
     
     return formatted_results
 
-
+@function_tool
 def shodan_host_info(ip: str) -> str:
     """
     Get detailed information about a specific host from Shodan.

src/cai/tools/reconnaissance/wget.py

@@ -4,8 +4,9 @@
 Wget tool
 """
 from cai.tools.common import run_command   # pylint: disable=import-error
+from cai.sdk.agents import function_tool
 
-
+@function_tool
 def wget(url: str, args: str = '', ctf=None) -> str:
     """
     Wget tool to download files from the web.

src/cai/tools/web/search_web.py

@@ -6,7 +6,10 @@ from cai.tools.web.google_search import (
     google_dork_search, 
     google_search
 )
+from cai.sdk.agents import function_tool
 
+
+@function_tool
 def query_perplexity(query: str = "", context: str = "") -> str:
     """
     Query the Perplexity AI API with a user prompt.
@@ -55,7 +58,7 @@ def query_perplexity(query: str = "", context: str = "") -> str:
     )
     return response.choices[0].message.content
 
-
+@function_tool
 def make_web_search_with_explanation(context: str = "", query: str = "") -> str:
     """
     Executes an intelligent web search via the AI service for relevant
@@ -74,6 +77,7 @@ def make_web_search_with_explanation(context: str = "", query: str = "") -> str:
     """
     return query_perplexity(query, context)
 
+@function_tool
 def make_google_search(query: str, dorks = False) -> str:
     """
     Search Google for information.