diff --git a/src/cai/agents/flag_discriminator.py b/src/cai/agents/flag_discriminator.py
index 1f2d39a0..6510a7ef 100644
--- a/src/cai/agents/flag_discriminator.py
+++ b/src/cai/agents/flag_discriminator.py
@@ -15,7 +15,7 @@ flag_discriminator = Agent(
                 2. ONLY return the flag, nothing else.
                 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 Function
diff --git a/src/cai/agents/patterns/red_team.py b/src/cai/agents/patterns/red_team.py
index 01d56b52..406934d2 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 tools to enable inter-agent communication pathways
+redteam_agent.tools.append(transfer_to_dns_agent)
+dns_smtp_agent.tools.append(redteam_agent_handoff)
+thought_agent.tools.append(redteam_agent_handoff)
 
 # Initialize the swarm pattern with the thought agent as the entry point
 redteam_swarm_pattern = thought_agent
diff --git a/tests/agents/test_agent_inferece.py b/tests/agents/test_agent_inferece.py
index 8b8ed122..22179f7d 100644
--- a/tests/agents/test_agent_inferece.py
+++ b/tests/agents/test_agent_inferece.py
@@ -15,9 +15,8 @@ from cai.agents.thought import thought_agent
 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"
+    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"
@@ -29,24 +28,12 @@ async def test_bug_bounter_agent_inference():
     """
     Non-streaming inference test for the bug_bounter_agent.
     """
-    prompt = "find vulnerabilities in web application sample.com"
+    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():
@@ -64,7 +51,6 @@ async def test_dns_smtp_agent_inference():
 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
@@ -84,18 +70,4 @@ async def test_red_team_agent_inference():
     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
+    assert "example.com" in final_output.lower(), f"Expected 'example.com' in output, got: {final_output}"