diff --git a/examples/cai/simple_one_tool_test_streamed.py b/examples/cai/simple_one_tool_test_streamed.py index 00632d49..c1cf222d 100644 --- a/examples/cai/simple_one_tool_test_streamed.py +++ b/examples/cai/simple_one_tool_test_streamed.py @@ -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!") diff --git a/src/cai/cli.py b/src/cai/cli.py index f867131b..3eb1fafd 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -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) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 3807b15c..872f9069 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -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': diff --git a/src/cai/util.py b/src/cai/util.py index a2f5cb21..1cdb95ed 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -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) \ No newline at end of file + 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() \ No newline at end of file