mirror of https://github.com/aliasrobotics/cai.git
81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
"""
|
|
A simple example to test the one_tool agent with streaming.
|
|
|
|
This script demonstrates how to initialize and run the one_tool agent
|
|
using Runner.run_streamed() with a simple hello message to verify everything
|
|
is working correctly, with streaming output.
|
|
"""
|
|
|
|
import os
|
|
import asyncio
|
|
import time
|
|
import json
|
|
import sys
|
|
from dotenv import load_dotenv
|
|
from openai import AsyncOpenAI
|
|
from openai.types.responses import ResponseTextDeltaEvent
|
|
from cai.sdk.agents import Runner, set_default_openai_client
|
|
from cai.agents import get_agent_by_name
|
|
from cai.util import fix_litellm_transcription_annotations, color
|
|
from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
|
|
from cai.sdk.agents.models._openai_shared import set_use_responses_by_default
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# NOTE: This is needed when using LiteLLM Proxy Server
|
|
#
|
|
# external_client = AsyncOpenAI(
|
|
# base_url=os.getenv('LITELLM_BASE_URL', 'http://localhost:4000'),
|
|
# api_key=os.getenv('LITELLM_API_KEY', 'key')
|
|
# )
|
|
# set_default_openai_client(external_client)
|
|
|
|
async def main():
|
|
# Apply litellm patch to fix the __annotations__ error
|
|
patch_applied = fix_litellm_transcription_annotations()
|
|
if not patch_applied:
|
|
print(color("Something went wrong patching LiteLLM fix_litellm_transcription_annotations", color="red"))
|
|
|
|
# Force the use of OpenAIChatCompletionsModel instead of OpenAIResponsesModel
|
|
set_use_responses_by_default(False)
|
|
|
|
# Get the one_tool agent
|
|
agent = get_agent_by_name("one_tool_agent")
|
|
|
|
print(f"Using model: {os.getenv('CAI_MODEL', 'default')}")
|
|
|
|
# Stream indicator
|
|
print("\nAgent response (streaming):")
|
|
print("-" * 40)
|
|
print("Agent: ", end="", flush=True)
|
|
|
|
# Run the agent with a simple test message in streaming mode
|
|
result = Runner.run_streamed(agent, "Hello! Can you 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()
|
|
|
|
|
|
print("\n" + "-" * 40)
|
|
print("\nTest completed successfully!")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |