diff --git a/ci/test/.test.yml b/ci/test/.test.yml
index 7f4d74e6..1aac3be5 100644
--- a/ci/test/.test.yml
+++ b/ci/test/.test.yml
@@ -110,6 +110,12 @@
   variables:
     TEST_PATH: tests/agents/test_max_turns.py
 
+
+# 🤖 agents test_agent_inference:
+#   <<: *run_test
+#   variables:
+#     TEST_PATH: tests/agents/test_agent_inference.py
+
 # ⚙️ core test_openai_chatcompletions:
 #   <<: *run_test
 #   variables:
diff --git a/examples/cai/simple_one_tool_test.py b/examples/cai/simple_one_tool_test.py
index 31f791c8..669ddd0d 100644
--- a/examples/cai/simple_one_tool_test.py
+++ b/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) #disable tracing or OPENAI_AGENTS_DISABLE_TRACING=1
 
 # NOTE: This is needed when using LiteLLM Proxy Server
 #
diff --git a/src/cai/agents/__init__.py b/src/cai/agents/__init__.py
index 29ca71b6..356e7cfb 100644
--- a/src/cai/agents/__init__.py
+++ b/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
 
diff --git a/src/cai/agents/basic.py b/src/cai/agents/basic.py
deleted file mode 100644
index 3c075242..00000000
--- a/src/cai/agents/basic.py
+++ /dev/null
@@ -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
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/bug_bounter.py b/src/cai/agents/bug_bounter.py
index 0c9f451b..d01b42fc 100644
--- a/src/cai/agents/bug_bounter.py
+++ b/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(),
+    )
+   
 )
\ No newline at end of file
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/flag_discriminator.py b/src/cai/agents/flag_discriminator.py
index f75be316..b638448f 100644
--- a/src/cai/agents/flag_discriminator.py
+++ b/src/cai/agents/flag_discriminator.py
@@ -1,24 +1,27 @@
 """
-CTF Flag Discriminator Agent
+CTF Flag Discriminator Agent with test
 """
 import os
-from cai.sdk.agents import Agent
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
+from openai import AsyncOpenAI
 
 model = os.getenv('CAI_MODEL', "qwen2.5:14b")
 flag_discriminator = Agent(
-    model="gpt-4o" if os.getenv('CAI_MODEL') == "o3-mini" else model,
     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.
+                """,
+    model=OpenAIChatCompletionsModel(
+        model="gpt-4o" if os.getenv('CAI_MODEL') == "o3-mini" else model,
+        openai_client=AsyncOpenAI(),
+    )
 )
 
-
-# TRANSFER FUNCTIONS
+# Transfer Function
 def transfer_to_flag_discriminator(**kwargs):  # pylint: disable=W0613
     """Transfer flag discriminator.
     Accepts any keyword arguments but ignores them."""
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/agents/patterns/red_team.py b/src/cai/agents/patterns/red_team.py
index 01d56b52..00b408ad 100644
--- a/src/cai/agents/patterns/red_team.py
+++ b/src/cai/agents/patterns/red_team.py
@@ -37,10 +37,10 @@ def thought_agent_handoff(ctf=None):  # pylint: disable=unused-argument
     return thought_agent
 
 
-# Register handoff functions to enable inter-agent communication pathways
-redteam_agent.functions.append(transfer_to_dns_agent)
-dns_smtp_agent.functions.append(redteam_agent_handoff)
-thought_agent.functions.append(redteam_agent_handoff)
+# Register handoff to enable inter-agent communication pathways
+redteam_agent.handoffs.append(transfer_to_dns_agent)
+dns_smtp_agent.handoffs.append(redteam_agent_handoff)
+thought_agent.handoffs.append(redteam_agent_handoff)
 
 # Initialize the swarm pattern with the thought agent as the entry point
 redteam_swarm_pattern = thought_agent
diff --git a/src/cai/agents/red_teamer.py b/src/cai/agents/red_teamer.py
index a12ad652..69184e2e 100644
--- a/src/cai/agents/red_teamer.py
+++ b/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,
-)
\ No newline at end of file
+    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
\ No newline at end of file
diff --git a/src/cai/agents/thought.py b/src/cai/agents/thought.py
index 187daaef..032ac35a 100644
--- a/src/cai/agents/thought.py
+++ b/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,
diff --git a/src/cai/tools/command_and_control/sshpass.py b/src/cai/tools/command_and_control/sshpass.py
index 9bc577e7..6dc3d217 100644
--- a/src/cai/tools/command_and_control/sshpass.py
+++ b/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)
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/code_interpreter.py b/src/cai/tools/misc/code_interpreter.py
index 6fd4e5ce..b0a79f0a 100644
--- a/src/cai/tools/misc/code_interpreter.py
+++ b/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.
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/src/cai/tools/misc/reasoning.py b/src/cai/tools/misc/reasoning.py
index 651c6d1b..b793eca1 100644
--- a/src/cai/tools/misc/reasoning.py
+++ b/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
diff --git a/src/cai/tools/others/scripting.py b/src/cai/tools/others/scripting.py
index 4e4b7c74..b1b0744e 100644
--- a/src/cai/tools/others/scripting.py
+++ b/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 = "",
diff --git a/src/cai/tools/reconnaissance/crypto_tools.py b/src/cai/tools/reconnaissance/crypto_tools.py
index 8a5ebb4a..ee6509c7 100644
--- a/src/cai/tools/reconnaissance/crypto_tools.py
+++ b/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.
diff --git a/src/cai/tools/reconnaissance/curl.py b/src/cai/tools/reconnaissance/curl.py
index 08928e66..62f2102b 100644
--- a/src/cai/tools/reconnaissance/curl.py
+++ b/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.
diff --git a/src/cai/tools/reconnaissance/exec_code.py b/src/cai/tools/reconnaissance/exec_code.py
index ff4bf6cf..794ac990 100644
--- a/src/cai/tools/reconnaissance/exec_code.py
+++ b/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:
     """
diff --git a/src/cai/tools/reconnaissance/filesystem.py b/src/cai/tools/reconnaissance/filesystem.py
index c9df56de..8f9f4c98 100644
--- a/src/cai/tools/reconnaissance/filesystem.py
+++ b/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.
diff --git a/src/cai/tools/reconnaissance/netcat.py b/src/cai/tools/reconnaissance/netcat.py
index 64f0207b..2b03d09d 100644
--- a/src/cai/tools/reconnaissance/netcat.py
+++ b/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:
     """
diff --git a/src/cai/tools/reconnaissance/netstat.py b/src/cai/tools/reconnaissance/netstat.py
index 44d7ae09..9ed5ff02 100644
--- a/src/cai/tools/reconnaissance/netstat.py
+++ b/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.
diff --git a/src/cai/tools/reconnaissance/nmap.py b/src/cai/tools/reconnaissance/nmap.py
index 72233e4e..cafa82e8 100644
--- a/src/cai/tools/reconnaissance/nmap.py
+++ b/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.
diff --git a/src/cai/tools/reconnaissance/shodan.py b/src/cai/tools/reconnaissance/shodan.py
index 6a35c4c1..a771276c 100644
--- a/src/cai/tools/reconnaissance/shodan.py
+++ b/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.
diff --git a/src/cai/tools/reconnaissance/wget.py b/src/cai/tools/reconnaissance/wget.py
index dc68ca72..fecd5cac 100644
--- a/src/cai/tools/reconnaissance/wget.py
+++ b/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.
diff --git a/src/cai/tools/web/search_web.py b/src/cai/tools/web/search_web.py
index 32e0432d..0d8bf78e 100644
--- a/src/cai/tools/web/search_web.py
+++ b/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.
diff --git a/tests/agents/test_agent_inference.py b/tests/agents/test_agent_inference.py
new file mode 100644
index 00000000..43b690d2
--- /dev/null
+++ b/tests/agents/test_agent_inference.py
@@ -0,0 +1,93 @@
+import os
+import pytest
+from cai.sdk.agents import Runner
+from cai.agents import get_agent_by_name
+
+
+@pytest.mark.allow_call_model_methods
+@pytest.mark.asyncio
+async def test_blue_team_agent_inference():
+    """
+    Non-streaming inference test for the blueteam_agent.
+    """
+    prompt = "Monitor login attempts for suspicious activity what we can do?"
+    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_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.
+    """
+    prompt = "use nmap to scan 192.168.2.0/24"
+    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 "example.com" in final_output.lower(), f"Expected 'example.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 one_tool_agent.
+    """
+    prompt = "Hello! Can you find tell me which is the flag in this string: 'Hi there, your reward flag{1234}"
+    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 "flag{1234}" in final_output.lower(), f"Expected 'flag{1234}' 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 one_tool_agent.
+    """
+    prompt = """The phases of the cybersecurity kill chain are: 
+    Reconnaissance, Exploitation, Lateral Movement, Data Exfiltration, and Command and Control (C2). 
+    In which phase of the kill chain does this example fall: 
+    Gathering initial intelligence about the target using OSINT techniques?"""
+    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 "reconnaissance" in final_output.lower(), f"Expected 'reconnaissance' in output, got: {final_output}"    
\ No newline at end of file