Delete trash

src/cai/agents/patterns/redteam_codeagent_as_tool.py

@@ -0,0 +1,63 @@
+"""
+Implementation of a Red Team Agent with Code Execution Agent as Tool Pattern
+
+This module establishes a specialized pattern where a Code Execution Agent
+serves as a tool for the Red Team Agent. This allows the Red Team Agent to
+delegate code execution tasks to a dedicated agent, enhancing security
+assessment capabilities through specialized code execution.
+"""
+import os
+from dotenv import load_dotenv
+from openai import AsyncOpenAI
+
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
+from cai.agents.red_teamer import redteam_agent
+from cai.tools.reconnaissance.exec_code import execute_code
+from cai.util import load_prompt_template
+
+
+# Load environment variables
+load_dotenv()
+model_name = os.getenv("CAI_MODEL", "qwen2.5:14b")
+
+# Create a specialized code execution agent with only the execute_code tool
+code_execution_agent = Agent(
+    name="Code Execution Agent",
+    description="Specialized agent for executing code in security assessments",
+    instructions=(
+        "You are a specialized code execution agent that helps with security "
+        "assessments. Execute the task and then transfer to red team"
+    ),
+    tools=[execute_code],
+    handoffs=[],
+    model=OpenAIChatCompletionsModel(
+        model=model_name,
+        openai_client=AsyncOpenAI(),
+    ),
+)
+
+# Create a handoff for the red team agent
+from cai.sdk.agents import handoff
+
+redteam_handoff = handoff(
+    agent=redteam_agent,
+    tool_description_override="Transfer to Red Team Agent for security assessment and exploitation tasks"
+)
+
+# Add the handoff to the code execution agent
+code_execution_agent.handoffs.append(redteam_handoff)
+
+# Register the code execution agent as a tool for the red team agent
+redteam_agent.tools.append(
+    code_execution_agent.as_tool(
+        tool_name="code_execution",
+        tool_description=(
+            "Use this tool when you need to execute code"
+        ),
+    )
+)
+if execute_code in redteam_agent.tools:
+    redteam_agent.tools.remove(execute_code)
+# Export the enhanced red team agent as the pattern
+redteam_with_code_agent_pattern = redteam_agent
+redteam_with_code_agent_pattern.pattern = "agent_as_tool"

src/cai/sdk/agents/models/openai_chatcompletions.py

@@ -1762,10 +1762,13 @@ class OpenAIChatCompletionsModel(Model):
         parallel_tool_calls: bool
     ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
         """Handle standard LiteLLM API calls for OpenAI and compatible models."""
+        # Make sure model is the first parameter
+        model = kwargs.pop("model", self.model)
+        
         if stream:
             # Standard LiteLLM handling for streaming
-            ret = litellm.completion(**kwargs)
-            stream_obj = await litellm.acompletion(**kwargs)
+            ret = litellm.completion(model=model, **kwargs)
+            stream_obj = await litellm.acompletion(model=model, **kwargs)
 
             response = Response(
                 id=FAKE_RESPONSES_ID,
@@ -1782,7 +1785,7 @@ class OpenAIChatCompletionsModel(Model):
             return response, stream_obj
         else:
             # Standard OpenAI handling for non-streaming
-            ret = litellm.completion(**kwargs)
+            ret = litellm.completion(model=model, **kwargs)
             return ret
             
     async def _fetch_response_litellm_ollama(
@@ -1794,9 +1797,11 @@ class OpenAIChatCompletionsModel(Model):
         parallel_tool_calls: bool,
         provider="ollama"
     ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
+        # Extract the model first to ensure it's the first parameter
+        model = kwargs.get("model", self.model)
+        
         # Extract only supported parameters for Ollama
         ollama_supported_params = {
-            "model": kwargs.get("model", ""),
             "messages": kwargs.get("messages", []),
             "stream": kwargs.get("stream", False)
         }
@@ -1821,7 +1826,7 @@ class OpenAIChatCompletionsModel(Model):
         ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None}
         
         # Check if this is a Qwen model
-        model_str = str(self.model).lower()
+        model_str = str(model).lower()
         is_qwen = "qwen" in model_str
                 
         api_base = get_ollama_api_base()
@@ -1842,21 +1847,21 @@ class OpenAIChatCompletionsModel(Model):
                 tools=[],
                 parallel_tool_calls=parallel_tool_calls or False,
             )
-            # Get streaming response
+            # Get streaming response - ensure model is first parameter
             stream_obj = await litellm.acompletion(
-                **ollama_kwargs,
+                model=model,
                 api_base=api_base,
                 custom_llm_provider=provider,
+                **ollama_kwargs
             )
             return response, stream_obj
         else:
-
-        
-            # Get completion response
+            # Get completion response - ensure model is first parameter
             return litellm.completion(
-                **ollama_kwargs,
+                model=model,
                 api_base=api_base,
                 custom_llm_provider=provider,
+                **ollama_kwargs
             )
 
     def _intermediate_logs(self):