RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix streaming responses in CAI and examples 

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches 2025-04-03 10:53:33 +02:00
parent b2233b76d1
commit e8e05e9e5a
4 changed files with 382 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
examples/cai/simple_one_tool_test_streamed.py
View File

@ -8,7 +8,9 @@ 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
@ -50,14 +52,27 @@ async def main():
# 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():
if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
# Print the delta with a visible marker for each token
# print(f"{event.data.delta}|", end="", flush=True)
print(f"{event.data.delta}", end="", flush=True)
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!")

183
src/cai/cli.py
View File

@ -55,6 +55,8 @@ Environment Variables
(default: "o3-mini")
CAI_SUPPORT_INTERVAL: Number of turns between support agent
executions (default: "5")
CAI_STREAM: Enable/disable streaming output in rich panel
(default: "true")
Extensions (only applicable if the right extension is installed):
@ -97,6 +99,7 @@ Usage Examples:
"""
import os
import sys
from dotenv import load_dotenv
from openai import AsyncOpenAI
from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner, AsyncOpenAI
@ -105,6 +108,7 @@ from openai.types.responses import ResponseTextDeltaEvent
from rich.console import Console
import asyncio
from cai.util import fix_litellm_transcription_annotations, color
from cai.util import create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming
# Import modules from cai.repl
from cai.repl.commands import FuzzyCommandCompleter, handle_command as commands_handle_command
@ -177,6 +181,24 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
# Display banner
display_banner(console)
# Function to get the short name of the agent for display
def get_agent_short_name(agent):
if hasattr(agent, 'name'):
# Split by spaces and take first word to get shortened name
return agent.name.split()[0]
return "Agent"
# Prevent the model from using its own rich streaming to avoid conflicts
# and suppress final output message to avoid duplicates
if hasattr(agent, 'model'):
if hasattr(agent.model, 'disable_rich_streaming'):
agent.model.disable_rich_streaming = True
if hasattr(agent.model, 'suppress_final_output'):
agent.model.suppress_final_output = True
# Track streaming context to ensure proper cleanup
current_streaming_context = None
while turn_count < max_turns:
try:
# Get user input with command completion and history
@ -204,19 +226,125 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
# Process the conversation with the agent
if stream:
# Use streamed response
print("Agent: ", end="", flush=True)
# For classic fallback when streaming fails
print_fallback = False
async def process_streamed_response():
nonlocal current_streaming_context, print_fallback
try:
# Get the model from the agent for display purposes
model_name = None
if hasattr(agent, 'model') and hasattr(agent.model, 'model'):
model_name = str(agent.model.model)
# Set the agent name in the model if available (for proper display in streaming panel)
if hasattr(agent, 'model'):
agent.model.agent_name = get_agent_short_name(agent)
# Make sure any previous streaming context is cleaned up
if current_streaming_context is not None:
try:
current_streaming_context["live"].stop()
except Exception:
pass # Ignore errors on cleanup
current_streaming_context = None
try:
# Create a new streaming context
current_streaming_context = create_agent_streaming_context(
agent_name=get_agent_short_name(agent),
counter=turn_count + 1, # 1-indexed for display
model=model_name
)
except Exception as e:
# If rich display fails, fall back to classic print mode
print(f"Agent: ", end="", flush=True)
print_fallback = True
import traceback
print(f"[Warning: Falling back to simple streaming: {str(e)}]", file=sys.stderr)
# Run the agent with streaming
result = Runner.run_streamed(agent, user_input)
# List to collect all deltas for computing final token counts
collected_text = []
# Process stream events
async for event in result.stream_events():
if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
print(event.data.delta, end="", flush=True)
print("\n") # Add a newline at the end
collected_text.append(event.data.delta)
# If using streaming context, update the panel
if current_streaming_context is not None:
update_agent_streaming_content(current_streaming_context, event.data.delta)
# Otherwise, print to console directly
elif print_fallback:
print(event.data.delta, end="", flush=True)
# Finish the streaming context if it exists
if current_streaming_context is not None:
# Get token stats for the final display
token_stats = None
# Try to get token stats from the model
if hasattr(agent, 'model'):
# Get the actual input/output token counts from the model when available
model = agent.model
# Calculate a more accurate output token estimate using tiktoken if available
output_text = "".join(collected_text)
output_tokens = len(output_text) // 4 # Fallback rough estimate
try:
import tiktoken
encoding = tiktoken.get_encoding("cl100k_base")
output_tokens = len(encoding.encode(output_text))
except Exception:
# Fallback to rough estimate if tiktoken fails
pass
# Store current input tokens to calculate difference next time
if not hasattr(model, 'previous_input_tokens'):
model.previous_input_tokens = 0
# Get the available token counts from the model, or use reasonable defaults
interaction_input = getattr(model, 'total_input_tokens', 0) - model.previous_input_tokens
if interaction_input <= 0:
interaction_input = output_tokens * 2 # Rough estimate based on output
# Update previous tokens for next calculation
model.previous_input_tokens = getattr(model, 'total_input_tokens', 0)
token_stats = {
"interaction_input_tokens": interaction_input,
"interaction_output_tokens": output_tokens,
"interaction_reasoning_tokens": 0,
"total_input_tokens": getattr(model, 'total_input_tokens', interaction_input),
"total_output_tokens": getattr(model, 'total_output_tokens', output_tokens),
"total_reasoning_tokens": getattr(model, 'total_reasoning_tokens', 0),
"interaction_cost": None,
"total_cost": None
}
finish_agent_streaming(current_streaming_context, token_stats)
current_streaming_context = None
elif print_fallback:
# Add a newline at the end of classic streaming
print("\n")
return result
except Exception as e:
print() # Add a newline after any partial output
# In case of errors, ensure streaming context is cleaned up
if current_streaming_context is not None:
try:
current_streaming_context["live"].stop()
except Exception:
pass
current_streaming_context = None
if print_fallback:
print() # Add a newline after any partial output
import traceback
tb = traceback.format_exc()
print(f"\n[Error occurred during streaming: {str(e)}]\nLocation: {tb}")
@ -230,16 +358,30 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
console.print(f"Agent: {response.final_output}")
turn_count += 1
except KeyboardInterrupt:
# Ensure streaming context is cleaned up on keyboard interrupt
if current_streaming_context is not None:
try:
current_streaming_context["live"].stop()
except Exception:
pass
current_streaming_context = None
break
# except Exception as e:
# import traceback
# import sys
# exc_type, exc_value, exc_traceback = sys.exc_info()
# tb_info = traceback.extract_tb(exc_traceback)
# filename, line, func, text = tb_info[-1]
# console.print(f"[bold red]Error: {str(e)}[/bold red]")
# console.print(f"[bold red]Traceback: {tb_info}[/bold red]")
except Exception as e:
# Ensure streaming context is cleaned up on any exception
if current_streaming_context is not None:
try:
current_streaming_context["live"].stop()
except Exception:
pass
current_streaming_context = None
import traceback
import sys
exc_type, exc_value, exc_traceback = sys.exc_info()
tb_info = traceback.extract_tb(exc_traceback)
filename, line, func, text = tb_info[-1]
console.print(f"[bold red]Error: {str(e)}[/bold red]")
console.print(f"[bold red]Traceback: {tb_info}[/bold red]")
def main():
# Apply litellm patch to fix the __annotations__ error
@ -252,9 +394,18 @@ def main():
# Get the agent instance by name
agent = get_agent_by_name(agent_type)
# Configure model flags to work well with CLI
if hasattr(agent, 'model'):
# Disable rich streaming in the model to avoid conflicts
if hasattr(agent.model, 'disable_rich_streaming'):
agent.model.disable_rich_streaming = True
# Suppress final output to avoid duplicates
if hasattr(agent.model, 'suppress_final_output'):
agent.model.suppress_final_output = True
# Enable streaming by default, unless specifically disabled
stream = os.getenv('CAI_STREAM', 'false').lower() != 'false'
stream = os.getenv('CAI_STREAM', 'true').lower() != 'false'
# Run the CLI with the selected agent
run_cai_cli(agent, stream=stream)

54
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -10,7 +10,7 @@ import tiktoken
from collections.abc import AsyncIterator, Iterable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal, cast, overload
from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages
from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming
from wasabi import color
from openai import NOT_GIVEN, AsyncOpenAI, AsyncStream, NotGiven
@ -194,6 +194,10 @@ class OpenAIChatCompletionsModel(Model):
self.total_reasoning_tokens = 0
self.agent_name = "Agent" # Default name
# Flags for CLI integration
self.disable_rich_streaming = False # Prevents creating a rich panel in the model
self.suppress_final_output = False # Prevents duplicate output at end of streaming
def set_agent_name(self, name: str) -> None:
"""Set the agent name for CLI display purposes."""
self.agent_name = name
@ -296,7 +300,7 @@ class OpenAIChatCompletionsModel(Model):
interaction_reasoning_tokens=(
response.usage.completion_tokens_details.reasoning_tokens
if response.usage and hasattr(response.usage, 'completion_tokens_details')
and response.usage.completion_tokens_details
and response.usage.completion_tokens_details
and hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')
else 0
),
@ -348,6 +352,18 @@ class OpenAIChatCompletionsModel(Model):
# Increment the interaction counter for CLI display
self.interaction_counter += 1
# Check if streaming should be shown in rich panel
should_show_rich_stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' and not self.disable_rich_streaming
# Create streaming context if needed
streaming_context = None
if should_show_rich_stream:
streaming_context = create_agent_streaming_context(
agent_name=self.agent_name,
counter=self.interaction_counter,
model=str(self.model)
)
with generation_span(
model=str(self.model),
model_config=dataclasses.asdict(model_settings)
@ -387,6 +403,9 @@ class OpenAIChatCompletionsModel(Model):
output_text = ""
estimated_output_tokens = 0
# Initialize a streaming text accumulator for rich display
streaming_text_buffer = ""
async for chunk in stream:
if not state.started:
state.started = True
@ -435,6 +454,13 @@ class OpenAIChatCompletionsModel(Model):
content = delta['content']
if content:
# Add to the streaming text buffer
streaming_text_buffer += content
# Update streaming display if enabled - always do this for text content
if streaming_context:
update_agent_streaming_content(streaming_context, content)
# More accurate token counting for text content
output_text += content
token_count, _ = count_tokens_with_tiktoken(output_text)
@ -726,8 +752,28 @@ class OpenAIChatCompletionsModel(Model):
hasattr(final_response.usage.output_tokens_details, 'reasoning_tokens')):
self.total_reasoning_tokens += final_response.usage.output_tokens_details.reasoning_tokens
# At the end of streaming, print the finalized agent message
if final_response.output and any(isinstance(item, ResponseOutputMessage) for item in final_response.output):
# Prepare final statistics for display
final_stats = {
"interaction_input_tokens": final_response.usage.input_tokens if final_response.usage else 0,
"interaction_output_tokens": final_response.usage.output_tokens if final_response.usage else 0,
"interaction_reasoning_tokens": (
final_response.usage.output_tokens_details.reasoning_tokens
if final_response.usage and final_response.usage.output_tokens_details
and hasattr(final_response.usage.output_tokens_details, 'reasoning_tokens')
else 0
),
"total_input_tokens": getattr(self, 'total_input_tokens', 0),
"total_output_tokens": getattr(self, 'total_output_tokens', 0),
"total_reasoning_tokens": getattr(self, 'total_reasoning_tokens', 0),
"interaction_cost": None,
"total_cost": None,
}
# At the end of streaming, finish the streaming context if we were using it
if streaming_context:
finish_agent_streaming(streaming_context, final_stats)
# If we're not using rich streaming and not suppressing output, use old method
elif not self.suppress_final_output and final_response.output and any(isinstance(item, ResponseOutputMessage) for item in final_response.output):
# Find the assistant message to print
for item in final_response.output:
if isinstance(item, ResponseOutputMessage) and item.role == 'assistant':

145
src/cai/util.py
View File

@ -472,4 +472,147 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli
title_align="left"
)
console.print("\n")
console.print(panel)
console.print(panel)
def create_agent_streaming_context(agent_name, counter, model):
"""Create a streaming context object that maintains state for streaming agent output."""
from rich.live import Live
import shutil
# Use the model from environment variable if available
model_override = os.getenv('CAI_MODEL')
if model_override:
model = model_override
timestamp = datetime.now().strftime("%H:%M:%S")
# Determine terminal size for best display
terminal_width, _ = shutil.get_terminal_size((100, 24))
panel_width = min(terminal_width - 4, 120) # Keep some margin
# Create base header for the panel
header = Text()
header.append(f"[{counter}] ", style="bold cyan")
header.append(f"Agent: {agent_name} ", style="bold green")
header.append(f">> ", style="yellow")
# Create the content area for streaming text
content = Text("")
# Add timestamp and model info
footer = Text()
footer.append(f"\n[{timestamp}", style="dim")
if model:
footer.append(f" ({model})", style="bold magenta")
footer.append("]", style="dim")
# Create the panel (initial state)
panel = Panel(
Text.assemble(header, content, footer),
border_style="blue",
box=ROUNDED,
padding=(1, 2), # Add more padding for better readability
title="[bold]Agent Streaming Response[/bold]",
title_align="left",
width=panel_width,
expand=True # Allow panel to expand to terminal width
)
# Start the live display with a higher refresh rate
live = Live(panel, refresh_per_second=20, console=console)
live.start()
# Return context object with all the elements needed for updating
return {
"live": live,
"panel": panel,
"header": header,
"content": content,
"footer": footer,
"timestamp": timestamp,
"model": model,
"agent_name": agent_name,
"panel_width": panel_width
}
def update_agent_streaming_content(context, text_delta):
"""Update the streaming content with new text."""
# Add the new text to the content
context["content"].append(text_delta)
# Update the live display with the latest content
updated_panel = Panel(
Text.assemble(context["header"], context["content"], context["footer"]),
border_style="blue",
box=ROUNDED,
padding=(1, 2), # Match padding from creation
title="[bold]Agent Streaming Response[/bold]",
title_align="left",
width=context.get("panel_width", 100),
expand=True # Allow panel to expand to terminal width
)
# Force an update with the new panel
context["live"].update(updated_panel)
context["panel"] = updated_panel
def finish_agent_streaming(context, final_stats=None):
"""Finish the streaming session and display final stats if available."""
# If we have token stats, add them
tokens_text = None
if final_stats:
interaction_input_tokens = final_stats.get("interaction_input_tokens")
interaction_output_tokens = final_stats.get("interaction_output_tokens")
interaction_reasoning_tokens = final_stats.get("interaction_reasoning_tokens")
total_input_tokens = final_stats.get("total_input_tokens")
total_output_tokens = final_stats.get("total_output_tokens")
total_reasoning_tokens = final_stats.get("total_reasoning_tokens")
interaction_cost = final_stats.get("interaction_cost")
total_cost = final_stats.get("total_cost")
if (interaction_input_tokens is not None and
interaction_output_tokens is not None and
interaction_reasoning_tokens is not None and
total_input_tokens is not None and
total_output_tokens is not None and
total_reasoning_tokens is not None):
tokens_text = _create_token_display(
interaction_input_tokens,
interaction_output_tokens,
interaction_reasoning_tokens,
total_input_tokens,
total_output_tokens,
total_reasoning_tokens,
context["model"],
interaction_cost,
total_cost
)
# Create the final panel with stats
final_panel = Panel(
Text.assemble(
context["header"],
context["content"],
Text("\n\n"),
tokens_text if tokens_text else Text(""),
context["footer"]
),
border_style="blue",
box=ROUNDED,
padding=(1, 2), # Match padding from creation
title="[bold]Agent Streaming Response[/bold]",
title_align="left",
width=context.get("panel_width", 100),
expand=True
)
# Update one last time
context["live"].update(final_panel)
# Ensure updates are displayed before stopping
import time
time.sleep(0.5)
# Stop the live display
context["live"].stop()