fix agent and add primary test

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
 #

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

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)

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(),
+    )
 )

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(),
+    )
 )

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.

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.

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