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