From 89293adb9ed828465a2b2bff1389cf0e543d5e31 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 29 Apr 2025 10:30:34 +0200 Subject: [PATCH] test functionalities --- examples/cai/agent_patterns/LLM_as_judge.py | 87 +++++++++++++++ examples/cai/agent_patterns/agent_as_tool.py | 90 ++++++++++++++++ examples/cai/agent_patterns/deterministic.py | 60 +++++++++++ examples/cai/agent_patterns/guardrails.py | 100 ++++++++++++++++++ examples/cai/agent_patterns/handoffs.py | 92 ++++++++++++++++ examples/cai/agent_patterns/paralelization.py | 87 +++++++++++++++ examples/cai/basic_usage.py | 65 ++++++++++++ .../agents/models/openai_chatcompletions.py | 19 ++-- 8 files changed, 591 insertions(+), 9 deletions(-) create mode 100644 examples/cai/agent_patterns/LLM_as_judge.py create mode 100644 examples/cai/agent_patterns/agent_as_tool.py create mode 100644 examples/cai/agent_patterns/deterministic.py create mode 100644 examples/cai/agent_patterns/guardrails.py create mode 100644 examples/cai/agent_patterns/handoffs.py create mode 100644 examples/cai/agent_patterns/paralelization.py create mode 100644 examples/cai/basic_usage.py diff --git a/examples/cai/agent_patterns/LLM_as_judge.py b/examples/cai/agent_patterns/LLM_as_judge.py new file mode 100644 index 00000000..d5fb738d --- /dev/null +++ b/examples/cai/agent_patterns/LLM_as_judge.py @@ -0,0 +1,87 @@ +""" +LLM as a Judge Pattern + +Uses one LLM to perform a task and another to evaluate the result and provide feedback. +The loop continues until the evaluator is satisfied. This improves output quality +and allows cost optimization by combining smaller and larger models. +""" + +from __future__ import annotations + +import asyncio +import os +from dataclasses import dataclass +from typing import Literal +from cai.sdk.agents import Agent, ItemHelpers, Runner, TResponseInputItem, OpenAIChatCompletionsModel +from openai import AsyncOpenAI + + +# CTF task planner agent (performs planning) +ctf_task_planner = Agent( + name="CTF Task Planner", + description="Agent focused on creating a task plan to approach a CTF challenge.", + instructions=( + "You are a cybersecurity strategist. Given a CTF challenge description, " + "generate a clear and effective plan of tasks to solve it. " + "Use any feedback to improve your planning." + ), + model=OpenAIChatCompletionsModel( + model= os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ), + tools =[] +) + + +# Feedback structure for the judge +@dataclass +class EvaluationFeedback: + feedback: str + score: Literal["pass", "needs_improvement", "fail"] + + +# CTF task plan evaluator (judges planning quality) +ctf_plan_evaluator = Agent[None]( + name="CTF Plan Evaluator", + description="Agent that evaluates CTF task plans for effectiveness and completeness.", + instructions=( + "You evaluate a task plan created for solving a CTF challenge. " + "Ensure it covers all essential steps (recon, exploitation, post-exploitation, etc.). " + "Provide actionable feedback. Never approve on the first try." + ), + model=OpenAIChatCompletionsModel( + model= os.getenv(' z', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ), + output_type=EvaluationFeedback, +) + + +async def main() -> None: + challenge_desc = input("Describe the CTF challenge: ") + input_items: list[TResponseInputItem] = [{"content": challenge_desc, "role": "user"}] + + latest_plan: str | None = None + + while True: + planner_result = await Runner.run(ctf_task_planner, input_items) + input_items = planner_result.to_input_list() + latest_plan = ItemHelpers.text_message_outputs(planner_result.new_items) + print("CTF task plan generated.") + + eval_result = await Runner.run(ctf_plan_evaluator, input_items) + feedback: EvaluationFeedback = eval_result.final_output + print(f"Evaluator score: {feedback.score}") + + if feedback.score == "pass": + print("CTF task plan approved.") + break + + print("Refining plan based on evaluator feedback...") + input_items.append({"content": f"Feedback: {feedback.feedback}", "role": "user"}) + + print(f"Final CTF task plan:\n{latest_plan}") + + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/examples/cai/agent_patterns/agent_as_tool.py b/examples/cai/agent_patterns/agent_as_tool.py new file mode 100644 index 00000000..b70ec802 --- /dev/null +++ b/examples/cai/agent_patterns/agent_as_tool.py @@ -0,0 +1,90 @@ +""" +Agents as Tools Example. + +This module demonstrates how to embed one agent as a tool within another agent's workflow. +An agent-as-tool operates independently on its provided input —without access to prior conversation history or "taking over" +the conversation — completes its specific task, and returns the result to the calling (parent) agent. +""" +from __future__ import annotations +import os +import json +import random +import asyncio +from openai import AsyncOpenAI +from cai.sdk.agents import function_tool +from cai.tools.common import run_command +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, HandoffInputData, Runner, function_tool, handoff, trace, ItemHelpers, MessageOutputItem +from cai.sdk.agents.extensions import handoff_filters + + + +@function_tool +def execute_cli_command(command: str) -> str: + return run_command(command) + + +flag_discriminator = Agent( + name="Flag discriminator", + description="Agent focused on extracting the flag from the output", + instructions="You are an agent tailored to extract the flag from a given output.", + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ), + handoff_description = "Agent focused on extracting the flag from the output" +) + +# We'll use this variable to check if the flag_discriminator agent is used +flag_discriminator_used = False + +# Wrap the as_tool method to set the flag when used +def flag_discriminator_as_tool(*args, **kwargs): + global flag_discriminator_used + flag_discriminator_used = True + return flag_discriminator._as_tool(*args, **kwargs) + +# Patch the as_tool method +flag_discriminator._as_tool = flag_discriminator.as_tool +flag_discriminator.as_tool = flag_discriminator_as_tool + +ctf_agent = Agent( + name="CTF agent", + description="Agent focused on conquering security challenges", + instructions="You are a Cybersecurity expert Leader facing a CTF", + tools=[ + execute_cli_command, + flag_discriminator.as_tool( + tool_name="find_flag", + tool_description ="Find flag in output text" + ) + ], + model=OpenAIChatCompletionsModel( + model= os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + +# Main function to execute the workflow +async def main(): + + result = await Runner.run( + ctf_agent, + input= [ + {"content": "Here is some output from a task. Find the flag: nhwitm flag{1234} mlsk. And returns only the flag", "role": "user"} + ], + ) + + for item in result.new_items: + if isinstance(item, MessageOutputItem): + text = ItemHelpers.text_message_output(item) + if text: + print(f"Final step: {text}") + + # Print whether the flag_discriminator agent was used + if flag_discriminator_used: + print("Flag discriminator agent was used.") + else: + print("Flag discriminator agent was NOT used.") + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/examples/cai/agent_patterns/deterministic.py b/examples/cai/agent_patterns/deterministic.py new file mode 100644 index 00000000..8bb0a80f --- /dev/null +++ b/examples/cai/agent_patterns/deterministic.py @@ -0,0 +1,60 @@ +""" +This module demonstrates a deterministic agent pattern for solving CTF (Capture The Flag) challenges using composable AI agents. +The deterministic pattern here involves a fixed, sequential pipeline of agents, each with a well-defined role. +This pattern ensures reproducibility and traceability, as the flow of information and agent responsibilities are strictly defined and do not vary between runs. +""" + +import os +import asyncio +from pydantic import BaseModel +from openai import AsyncOpenAI +from cai.sdk.agents import Agent, Runner, OpenAIChatCompletionsModel, function_tool +from cai.tools.common import run_command +from cai.sdk.agents.extensions import handoff_filters + +# Function tool to execute CLI command +@function_tool +def execute_cli_command(command: str) -> str: + return run_command(command) + +# Define Flag Discriminator Agent (handles extracting flags from CTF output) +flag_discriminator = Agent( + name="Flag discriminator", + description="Agent focused on extracting the flag from the output", + instructions="You are an agent tailored to extract the flag from a given output.", + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + +# Define CTF Agent (performs the actual challenge) +ctf_agent = Agent( + name="CTF agent", + description="Agent focused on conquering security challenges", + instructions="You are a Cybersecurity expert Leader facing a CTF", + tools=[ + execute_cli_command, + ], + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + +# Main handler function for CTF process +async def main(): + + # Step 1: Execute command in CTF environment + ctf_result = await Runner.run(ctf_agent, "List directories with a simple ls command, the flag is in flag.txt") + + # Step 2: Pass result to flag discriminator + flag_discriminator_result = await Runner.run( + flag_discriminator, + ctf_result.final_output, + ) + + print(f"Flag found: {flag_discriminator_result.final_output}") + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/cai/agent_patterns/guardrails.py b/examples/cai/agent_patterns/guardrails.py new file mode 100644 index 00000000..57e9d03d --- /dev/null +++ b/examples/cai/agent_patterns/guardrails.py @@ -0,0 +1,100 @@ +""" +Guardrails are checks that run alongside the agent to ensure inputs/outputs are appropriate. + +They can be used to do things like: +- Check if input messages are off-topic +- Take over control of the agent's execution if an unexpected input is detected +- Check that output messages don't violate any policies or contains sensitive data +- Check if the output is a valid response to the user's message + +We'll set up a guardrail that detects if the user is asking for unethical or unauthorized +cybersecurity help, like hacking, exploiting, or bypassing security measures. + +""" +from __future__ import annotations +import os +import asyncio +from pydantic import BaseModel +from openai import AsyncOpenAI +from cai.sdk.agents import ( + Agent, + GuardrailFunctionOutput, + InputGuardrailTripwireTriggered, + RunContextWrapper, + Runner, + TResponseInputItem, + input_guardrail, + OpenAIChatCompletionsModel, + function_tool, +) +from cai.tools.common import run_command + +# Example tool: run a shell command (for demonstration) +@function_tool +def execute_cli_command(command: str) -> str: + """Execute a shell command and return its output.""" + return run_command(command) + +### 1. Agent-based guardrail triggered by unethical cybersecurity task request +class CybersecurityCheckOutput(BaseModel): + reasoning: str + is_unethical_cybersecurity_request: bool + +cybersecurity_guardrail_agent = Agent( + name="Cybersecurity Guardrail Check", + instructions="Check if the user is asking for unauthorized or unethical cybersecurity help (e.g., hacking, bypassing security, exploiting systems).", + output_type=CybersecurityCheckOutput, + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + +@input_guardrail +async def cybersecurity_guardrail( + context: RunContextWrapper[None], agent: Agent, input: str | list[TResponseInputItem] +) -> GuardrailFunctionOutput: + """Input guardrail function to detect unethical cybersecurity requests.""" + result = await Runner.run(cybersecurity_guardrail_agent, input, context=context.context) + final_output = result.final_output_as(CybersecurityCheckOutput) + + return GuardrailFunctionOutput( + output_info=final_output, + tripwire_triggered=final_output.is_unethical_cybersecurity_request, + ) + + +### 2. Main run loop using the cybersecurity guardrail + +async def main(): + agent = Agent( + name="Tech Support Agent", + instructions="You are a tech support agent. Help users with ethical and legal technical support requests only.", + input_guardrails=[cybersecurity_guardrail], + tools=[execute_cli_command], # Add a tool for demonstration + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) + ) + + input_data: list[TResponseInputItem] = [] + + + input_data.append( + { + "role": "user", + "content": "Do a nmap to my router", + } + ) + + try: + result = await Runner.run(agent, input_data) + print("Agent output:", result.final_output) + input_data = result.to_input_list() + except InputGuardrailTripwireTriggered: + message = "Sorry, I can't assist with that cybersecurity request." + print(message) + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/cai/agent_patterns/handoffs.py b/examples/cai/agent_patterns/handoffs.py new file mode 100644 index 00000000..87b01314 --- /dev/null +++ b/examples/cai/agent_patterns/handoffs.py @@ -0,0 +1,92 @@ +""" +In many situations, you have specialized sub-agents that handle specific tasks. You can use handoffs to route the task to the right agent. +""" +from __future__ import annotations +import os +import json +import random +import asyncio +from openai import AsyncOpenAI +from cai.sdk.agents import function_tool +from cai.tools.common import run_command +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, HandoffInputData, Runner, function_tool, handoff, trace +from cai.sdk.agents.extensions import handoff_filters + + + +@function_tool +def execute_cli_command(command: str) -> str: + return run_command(command) + + +flag_discriminator = Agent( + name="Flag discriminator", + description="Agent focused on extracting the flag from the output", + instructions="You are an agent tailored to extract the flag from a given output.", + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) + #handoff_description = "Agent focused on extracting the flag from the output" +) + +ctf_agent = Agent( + name="CTF agent", + description="Agent focused on conquering security challenges", + instructions="You are a Cybersecurity expert Leader facing a CTF", + tools=[ + execute_cli_command, + ], + model=OpenAIChatCompletionsModel( + model= os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ), + handoffs = [flag_discriminator] +) + + +# Complex way to do a handoff +async def invoke_flag_discriminator(context: RunContextWrapper[Any], args: str="") -> Agent: + """ + This function is called when we need to hand off the task to the flag_discriminator. + """ + # Check if args is empty + if not args: + print("No input provided.") + else: + print("Input provided, processing...") + print(f"Passing args to flag_discriminator: {args}") + + # Return the agent (flag_discriminator) that will handle extracting the flag + return flag_discriminator + +# input_filter: can be used for additional data filtering during the handoff +flag_discriminator_complex_handoff = handoff( + agent=flag_discriminator, + input_filter = invoke_flag_discriminator +) + + +ctf_agent.handoffs.append(flag_discriminator_complex_handoff) + +# Main function to execute the workflow +async def main(): + # Trace the entire run as a single workflow + with trace(workflow_name="CTF Workflow"): + # Step 1: Execute a command with the CTF agent + result = await Runner.run(ctf_agent, input="List all files in the current directory") + + # Step 2: Ask an additional question for calling the Flag Discriminator agent + result = await Runner.run( + ctf_agent, + input=result.to_input_list() + [ + {"content": "Here is some output from a task. The first file is the name of the flag", "role": "user"} + ], + ) + + for message in result.to_input_list(): + print(json.dumps(message, indent=2)) + +if __name__ == "__main__": + + asyncio.run(main()) \ No newline at end of file diff --git a/examples/cai/agent_patterns/paralelization.py b/examples/cai/agent_patterns/paralelization.py new file mode 100644 index 00000000..815bad19 --- /dev/null +++ b/examples/cai/agent_patterns/paralelization.py @@ -0,0 +1,87 @@ +""" +Parallelization Pattern: +This pattern runs multiple agents in parallel to perform a task, generating different responses. +Afterward, a separate agent is used to evaluate and pick the best result. +""" + +import asyncio +import os +import asyncio +from openai import AsyncOpenAI +from cai.sdk.agents import function_tool +from cai.tools.common import run_command +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, HandoffInputData, Runner, function_tool, handoff, trace, ItemHelpers +from cai.sdk.agents.extensions import handoff_filters + + + +@function_tool +def execute_cli_command(command: str) -> str: + return run_command(command) +# Create the CTF agent +ctf_agent = Agent( + name="CTF agent", + description="Agent focused on conquering security challenges", + instructions="You are a Cybersecurity expert Leader facing a CTF", + tools=[ + execute_cli_command, + ], + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + +# An agent to pick the best solution after multiple attempts +best_solution_picker = Agent( + name="best_solution_picker", + description="Agent focused on picking the best security solutio", + instructions="You pick the best security solution from the given attempts.", + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + +async def main(): + # Define your CTF challenge + challenge = input("Enter the CTF challenge you're facing:\n\n") + + # Ensure the entire workflow is a single trace + res_1, res_2, res_3 = await asyncio.gather( + Runner.run( + ctf_agent, + challenge, + ), + Runner.run( + ctf_agent, + challenge, + ), + Runner.run( + ctf_agent, + challenge, + ), + ) + + # Gather the results from the CTF attempts + outputs = [ + ItemHelpers.text_message_outputs(res_1.new_items), + ItemHelpers.text_message_outputs(res_2.new_items), + ItemHelpers.text_message_outputs(res_3.new_items), + ] + + # Show all the results + results = "\n\n".join(outputs) + print(f"\n\nCTF Results:\n\n{results}") + + # Run the best solution picker agent + best_solution = await Runner.run( + best_solution_picker, + f"Input: {challenge}\n\nResults:\n{results}", + ) + + print("\n\n-----") + print(f"Best solution: {best_solution.final_output}") + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/examples/cai/basic_usage.py b/examples/cai/basic_usage.py new file mode 100644 index 00000000..3e93069e --- /dev/null +++ b/examples/cai/basic_usage.py @@ -0,0 +1,65 @@ +""" +A common tactic is to break down a task into a series of smaller steps. +Each task can be performed by an agent, and the output of one agent is used as input to the next +try stream and normal +""" +import os +import sys +import time +import asyncio +from cai.sdk.agents import Runner, Agent, OpenAIChatCompletionsModel, set_tracing_disabled +from openai import AsyncOpenAI +from cai.sdk.agents import function_tool +from cai.tools.common import run_command + + +@function_tool +def execute_cli_command(command: str) -> str: + return run_command(command) + + +ctf_agent = Agent( + name="CTF agent", + description="Agent focused on conquering security challenges", + instructions="You are a Cybersecurity expert Leader facing a CTF", + tools=[ + execute_cli_command, + ], + model=OpenAIChatCompletionsModel( + model= os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) +async def main(): + result = await Runner.run(ctf_agent, "List the files in the current directory?") + print("\nAgent response:") + print(result.final_output) + +async def main_streamed(): + print("\nAgent response (streaming):") + result = Runner.run_streamed(ctf_agent, "List the files in the current directory?") + + # Process the streaming response events + event_count = 0 + start_time = time.time() + + # Process the streaming response + async for event in result.stream_events(): + event_count += 1 + # Add a small delay to allow the streaming panel to update properly + await asyncio.sleep(0.01) + + # # Print a progress indicator + # if event_count % 10 == 0: + # elapsed = time.time() - start_time + # sys.stdout.write(f"\rProcessed {event_count} events in {elapsed:.1f} seconds...") + # sys.stdout.flush() + + # Clear the progress line + sys.stdout.write("\r" + " " * 60 + "\r") + sys.stdout.flush() + +if __name__ == "__main__": + set_tracing_disabled(True) + asyncio.run(main()) + asyncio.run(main_streamed()) \ No newline at end of file diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index e294bee7..3ea39a94 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -964,18 +964,19 @@ class OpenAIChatCompletionsModel(Model): "extra_headers": _HEADERS, } + + + # Error encountered: Error code: 400 - {'error': {'code': 'invalid_request_error', + # 'message': "'tool_choice' is only allowed when 'tools' are specified", + # 'type': 'invalid_request_error', 'param': None}} + # + # Only remove tool_choice if model starts with "gpt" and has no tools + if self.model.startswith("gpt") and not converted_tools: + kwargs.pop("tool_choice", None) + # Model adjustments if any(x in self.model for x in ["claude"]): litellm.drop_params = True - - # Error encountered: Error code: 400 - {'error': {'code': 'invalid_request_error', - # 'message': "'tool_choice' is only allowed when 'tools' are specified", - # 'type': 'invalid_request_error', 'param': None}} - # - # if has no tools, remove tool_choice - if not converted_tools: - kwargs.pop("tool_choice", None) - # BadRequestError encountered: litellm.BadRequestError: AnthropicException - # b'{"type":"error","error": # {"type":"invalid_request_error","message":"store: Extra inputs are not permitted"}}'