From 3ebe4319c3f58c373a2d122cd49eeaa763f7a35d Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 13 May 2025 11:36:25 +0200 Subject: [PATCH] Delete trash --- .../patterns/redteam_codeagent_as_tool.py | 63 +++++++++++++++++++ .../agents/models/openai_chatcompletions.py | 27 ++++---- 2 files changed, 79 insertions(+), 11 deletions(-) create mode 100644 src/cai/agents/patterns/redteam_codeagent_as_tool.py diff --git a/src/cai/agents/patterns/redteam_codeagent_as_tool.py b/src/cai/agents/patterns/redteam_codeagent_as_tool.py new file mode 100644 index 00000000..fdbb195e --- /dev/null +++ b/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" diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index a61ab895..b2e72417 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/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):