mirror of https://github.com/aliasrobotics/cai.git
87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
"""
|
|
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()) |