diff --git a/src/cai/cli.py b/src/cai/cli.py index a6e50086..19154e78 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -199,12 +199,13 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # 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 + agent.model.disable_rich_streaming = False # Now True as the model handles streaming if hasattr(agent.model, 'suppress_final_output'): agent.model.suppress_final_output = True - - # Track streaming context to ensure proper cleanup - current_streaming_context = None + + # Set the agent name in the model for proper display in streaming panel + if hasattr(agent.model, 'set_agent_name'): + agent.model.set_agent_name(get_agent_short_name(agent)) while turn_count < max_turns: try: @@ -229,13 +230,17 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Configure the new agent's model flags if hasattr(agent, 'model'): if hasattr(agent.model, 'disable_rich_streaming'): - agent.model.disable_rich_streaming = True + agent.model.disable_rich_streaming = False # Now False to let model handle streaming if hasattr(agent.model, 'suppress_final_output'): agent.model.suppress_final_output = True # Apply current model to the new agent if hasattr(agent.model, 'model'): agent.model.model = current_model + + # Set agent name in the model for streaming display + if hasattr(agent.model, 'set_agent_name'): + agent.model.set_agent_name(get_agent_short_name(agent)) except Exception as e: console.print(f"[red]Error switching agent: {str(e)}[/red]") @@ -318,125 +323,19 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Process the conversation with the agent if stream: - # 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 + # Process stream events async for event in result.stream_events(): - if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent): - 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": calculate_model_cost(str(model), interaction_input, output_tokens), - "total_cost": calculate_model_cost(str(model), getattr(model, 'total_input_tokens', interaction_input), getattr(model, 'total_output_tokens', output_tokens)) - } - - 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") + # The openai_chatcompletions.py now handles all the streaming display + # We only need to pass through the events + pass return result except Exception as e: - # 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}") @@ -446,27 +345,12 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= else: # Use non-streamed response response = asyncio.run(Runner.run(agent, user_input)) - #console.print(f"Agent: {response.final_output}") # NOTE: this line is commented to avoid duplicate output turn_count += 1 except KeyboardInterrupt: - if stream: - # 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 + # No need to clean up streaming context as model handles it + pass 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() diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 18134f83..6b5c956f 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -477,678 +477,654 @@ class OpenAIChatCompletionsModel(Model): model=str(self.model) ) - with generation_span( - model=str(self.model), - model_config=dataclasses.asdict(model_settings) - | {"base_url": str(self._client.base_url)}, - disabled=tracing.is_disabled(), - ) as span_generation: - # Prepare messages for consistent token counting - converted_messages = _Converter.items_to_messages(input) - if system_instructions: - converted_messages.insert( - 0, - { - "content": system_instructions, + try: + with generation_span( + model=str(self.model), + model_config=dataclasses.asdict(model_settings) + | {"base_url": str(self._client.base_url)}, + disabled=tracing.is_disabled(), + ) as span_generation: + # Prepare messages for consistent token counting + converted_messages = _Converter.items_to_messages(input) + if system_instructions: + converted_messages.insert( + 0, + { + "content": system_instructions, + "role": "system", + }, + ) + # --- Add to message_history: user, system prompts --- + if system_instructions: + sys_msg = { "role": "system", + "content": system_instructions + } + add_to_message_history(sys_msg) + + if isinstance(input, str): + user_msg = { + "role": "user", + "content": input + } + add_to_message_history(user_msg) + elif isinstance(input, list): + for item in input: + if isinstance(item, dict): + if item.get("role") == "user": + user_msg = { + "role": "user", + "content": item.get("content", "") + } + add_to_message_history(user_msg) + # Get token count estimate before API call for consistent counting + estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) + + response, stream = await self._fetch_response( + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + span_generation, + tracing, + stream=True, + ) + + usage: CompletionUsage | None = None + state = _StreamingState() + + # Manual token counting (when API doesn't provide it) + output_text = "" + estimated_output_tokens = 0 + + # Initialize a streaming text accumulator for rich display + streaming_text_buffer = "" + # For tool call streaming, accumulate tool_calls to add to message_history at the end + streamed_tool_calls = [] + + # Ollama specific: accumulate full content to check for function calls at the end + # Some Ollama models output the function call as JSON in the text content + ollama_full_content = "" + is_ollama = False + + model_str = str(self.model).lower() + is_ollama = self.is_ollama or "ollama" in model_str or ":" in model_str or "qwen" in model_str + + # Add visual separation before agent output + if streaming_context and should_show_rich_stream: + # If we're using rich context, we'll add separation through that + pass + else: + # Print clear visual separator + print("\n") + + try: + async for chunk in stream: + if not state.started: + state.started = True + yield ResponseCreatedEvent( + response=response, + type="response.created", + ) + + # The usage is only available in the last chunk + if hasattr(chunk, 'usage'): + usage = chunk.usage + # For Ollama/LiteLLM streams that don't have usage attribute + else: + usage = None + + # Handle different stream chunk formats + if hasattr(chunk, 'choices') and chunk.choices: + choices = chunk.choices + elif hasattr(chunk, 'delta') and chunk.delta: + # Some providers might return delta directly + choices = [{"delta": chunk.delta}] + elif isinstance(chunk, dict) and 'choices' in chunk: + choices = chunk['choices'] + # Special handling for Qwen/Ollama chunks + elif isinstance(chunk, dict) and ('content' in chunk or 'function_call' in chunk): + # Qwen direct delta format - convert to standard + choices = [{"delta": chunk}] + else: + # Skip chunks that don't contain choice data + continue + + if not choices or len(choices) == 0: + continue + + # Get the delta content + delta = None + if hasattr(choices[0], 'delta'): + delta = choices[0].delta + elif isinstance(choices[0], dict) and 'delta' in choices[0]: + delta = choices[0]['delta'] + + if not delta: + continue + + # Handle text + content = None + if hasattr(delta, 'content') and delta.content is not None: + content = delta.content + elif isinstance(delta, dict) and 'content' in delta and delta['content'] is not None: + content = delta['content'] + + if content: + # For Ollama, we need to accumulate the full content to check for function calls + if is_ollama: + ollama_full_content += 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) + estimated_output_tokens = token_count + + if not state.text_content_index_and_output: + # Initialize a content tracker for streaming text + state.text_content_index_and_output = ( + 0 if not state.refusal_content_index_and_output else 1, + ResponseOutputText( + text="", + type="output_text", + annotations=[], + ), + ) + # Start a new assistant message stream + assistant_item = ResponseOutputMessage( + id=FAKE_RESPONSES_ID, + content=[], + role="assistant", + type="message", + status="in_progress", + ) + # Notify consumers of the start of a new output message + first content part + yield ResponseOutputItemAddedEvent( + item=assistant_item, + output_index=0, + type="response.output_item.added", + ) + yield ResponseContentPartAddedEvent( + content_index=state.text_content_index_and_output[0], + item_id=FAKE_RESPONSES_ID, + output_index=0, + part=ResponseOutputText( + text="", + type="output_text", + annotations=[], + ), + type="response.content_part.added", + ) + # Emit the delta for this segment of content + yield ResponseTextDeltaEvent( + content_index=state.text_content_index_and_output[0], + delta=content, + item_id=FAKE_RESPONSES_ID, + output_index=0, + type="response.output_text.delta", + ) + # Accumulate the text into the response part + state.text_content_index_and_output[1].text += content + + # Handle refusals (model declines to answer) + refusal_content = None + if hasattr(delta, 'refusal') and delta.refusal: + refusal_content = delta.refusal + elif isinstance(delta, dict) and 'refusal' in delta and delta['refusal']: + refusal_content = delta['refusal'] + + if refusal_content: + if not state.refusal_content_index_and_output: + # Initialize a content tracker for streaming refusal text + state.refusal_content_index_and_output = ( + 0 if not state.text_content_index_and_output else 1, + ResponseOutputRefusal(refusal="", type="refusal"), + ) + # Start a new assistant message if one doesn't exist yet (in-progress) + assistant_item = ResponseOutputMessage( + id=FAKE_RESPONSES_ID, + content=[], + role="assistant", + type="message", + status="in_progress", + ) + # Notify downstream that assistant message + first content part are starting + yield ResponseOutputItemAddedEvent( + item=assistant_item, + output_index=0, + type="response.output_item.added", + ) + yield ResponseContentPartAddedEvent( + content_index=state.refusal_content_index_and_output[0], + item_id=FAKE_RESPONSES_ID, + output_index=0, + part=ResponseOutputText( + text="", + type="output_text", + annotations=[], + ), + type="response.content_part.added", + ) + # Emit the delta for this segment of refusal + yield ResponseRefusalDeltaEvent( + content_index=state.refusal_content_index_and_output[0], + delta=refusal_content, + item_id=FAKE_RESPONSES_ID, + output_index=0, + type="response.refusal.delta", + ) + # Accumulate the refusal string in the output part + state.refusal_content_index_and_output[1].refusal += refusal_content + + # Handle tool calls + # Because we don't know the name of the function until the end of the stream, we'll + # save everything and yield events at the end + tool_calls = self._detect_and_format_function_calls(delta) + + if tool_calls: + for tc_delta in tool_calls: + tc_index = tc_delta.index if hasattr(tc_delta, 'index') else tc_delta.get('index', 0) + if tc_index not in state.function_calls: + state.function_calls[tc_index] = ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + arguments="", + name="", + type="function_call", + call_id="", + ) + + tc_function = None + if hasattr(tc_delta, 'function'): + tc_function = tc_delta.function + elif isinstance(tc_delta, dict) and 'function' in tc_delta: + tc_function = tc_delta['function'] + + if tc_function: + # Handle both object and dict formats + args = "" + if hasattr(tc_function, 'arguments'): + args = tc_function.arguments or "" + elif isinstance(tc_function, dict) and 'arguments' in tc_function: + args = tc_function.get('arguments', "") or "" + + name = "" + if hasattr(tc_function, 'name'): + name = tc_function.name or "" + elif isinstance(tc_function, dict) and 'name' in tc_function: + name = tc_function.get('name', "") or "" + + state.function_calls[tc_index].arguments += args + state.function_calls[tc_index].name += name + + # Handle call_id in both formats + call_id = "" + if hasattr(tc_delta, 'id'): + call_id = tc_delta.id or "" + elif isinstance(tc_delta, dict) and 'id' in tc_delta: + call_id = tc_delta.get('id', "") or "" + else: + # For Qwen models, generate a predictable ID if none is provided + if state.function_calls[tc_index].name: + # Generate a stable ID from the function name and arguments + call_id = f"call_{hashlib.md5(state.function_calls[tc_index].name.encode()).hexdigest()[:8]}" + + state.function_calls[tc_index].call_id += call_id + + # --- Accumulate tool call for message_history --- + # Only add if not already present (avoid duplicates in streaming) + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": state.function_calls[tc_index].call_id, + "type": "function", + "function": { + "name": state.function_calls[tc_index].name, + "arguments": state.function_calls[tc_index].arguments + } + } + ] + } + # Only add if not already in streamed_tool_calls + if tool_call_msg not in streamed_tool_calls: + streamed_tool_calls.append(tool_call_msg) + add_to_message_history(tool_call_msg) + + except Exception as e: + # Ensure streaming context is cleaned up in case of errors + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + except Exception: + pass + raise e + + # Special handling for Ollama - check if accumulated text contains a valid function call + if is_ollama and ollama_full_content and len(state.function_calls) == 0: + # Look for JSON object that might be a function call + try: + # Try to extract a JSON object from the content + json_start = ollama_full_content.find('{') + json_end = ollama_full_content.rfind('}') + 1 + + if json_start >= 0 and json_end > json_start: + json_str = ollama_full_content[json_start:json_end] + # Try to parse the JSON + parsed = json.loads(json_str) + + # Check if it looks like a function call + if ('name' in parsed and 'arguments' in parsed): + logger.debug(f"Found valid function call in Ollama output: {json_str}") + + # Create a tool call ID + tool_call_id = f"call_{hashlib.md5((parsed['name'] + str(time.time())).encode()).hexdigest()[:8]}" + + # Ensure arguments is a valid JSON string + arguments_str = "" + if isinstance(parsed['arguments'], dict): + # Remove 'ctf' field if it exists + if 'ctf' in parsed['arguments']: + del parsed['arguments']['ctf'] + arguments_str = json.dumps(parsed['arguments']) + elif isinstance(parsed['arguments'], str): + # If it's already a string, check if it's valid JSON + try: + # Try parsing to validate and remove 'ctf' if present + args_dict = json.loads(parsed['arguments']) + if isinstance(args_dict, dict) and 'ctf' in args_dict: + del args_dict['ctf'] + arguments_str = json.dumps(args_dict) + except: + # If not valid JSON, encode it as a JSON string + arguments_str = json.dumps(parsed['arguments']) + else: + # For any other type, convert to string and then JSON + arguments_str = json.dumps(str(parsed['arguments'])) + # Add it to our function_calls state + state.function_calls[0] = ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + arguments=arguments_str, + name=parsed['name'], + type="function_call", + call_id=tool_call_id, + ) + + # Display the tool call in CLI + from cai.util import cli_print_agent_messages + try: + # Create a message-like object to display the function call + tool_msg = type('ToolCallWrapper', (), { + 'content': None, + 'tool_calls': [ + type('ToolCallDetail', (), { + 'function': type('FunctionDetail', (), { + 'name': parsed['name'], + 'arguments': arguments_str + }), + 'id': tool_call_id, + 'type': 'function' + }) + ] + }) + + # Print the tool call using the CLI utility + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=tool_msg, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=estimated_input_tokens, + interaction_output_tokens=estimated_output_tokens, + interaction_reasoning_tokens=0, # Not available for Ollama + total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, + total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, + total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), + interaction_cost=None, + total_cost=None, + tool_output=None # Will be shown once the tool is executed + ) + except Exception as e: + logger.error(f"Error displaying tool call in CLI: {e}") + + # Add to message history + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": parsed['name'], + "arguments": arguments_str + } + } + ] + } + + streamed_tool_calls.append(tool_call_msg) + add_to_message_history(tool_call_msg) + + logger.debug(f"Added function call: {parsed['name']} with args: {arguments_str}") + except Exception as e: + pass + + function_call_starting_index = 0 + if state.text_content_index_and_output: + function_call_starting_index += 1 + # Send end event for this content part + yield ResponseContentPartDoneEvent( + content_index=state.text_content_index_and_output[0], + item_id=FAKE_RESPONSES_ID, + output_index=0, + part=state.text_content_index_and_output[1], + type="response.content_part.done", + ) + + if state.refusal_content_index_and_output: + function_call_starting_index += 1 + # Send end event for this content part + yield ResponseContentPartDoneEvent( + content_index=state.refusal_content_index_and_output[0], + item_id=FAKE_RESPONSES_ID, + output_index=0, + part=state.refusal_content_index_and_output[1], + type="response.content_part.done", + ) + + # Actually send events for the function calls + for function_call in state.function_calls.values(): + # First, a ResponseOutputItemAdded for the function call + yield ResponseOutputItemAddedEvent( + item=ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + call_id=function_call.call_id, + arguments=function_call.arguments, + name=function_call.name, + type="function_call", + ), + output_index=function_call_starting_index, + type="response.output_item.added", + ) + # Then, yield the args + yield ResponseFunctionCallArgumentsDeltaEvent( + delta=function_call.arguments, + item_id=FAKE_RESPONSES_ID, + output_index=function_call_starting_index, + type="response.function_call_arguments.delta", + ) + # Finally, the ResponseOutputItemDone + yield ResponseOutputItemDoneEvent( + item=ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + call_id=function_call.call_id, + arguments=function_call.arguments, + name=function_call.name, + type="function_call", + ), + output_index=function_call_starting_index, + type="response.output_item.done", + ) + + # Finally, send the Response completed event + outputs: list[ResponseOutputItem] = [] + if state.text_content_index_and_output or state.refusal_content_index_and_output: + assistant_msg = ResponseOutputMessage( + id=FAKE_RESPONSES_ID, + content=[], + role="assistant", + type="message", + status="completed", + ) + if state.text_content_index_and_output: + assistant_msg.content.append(state.text_content_index_and_output[1]) + if state.refusal_content_index_and_output: + assistant_msg.content.append(state.refusal_content_index_and_output[1]) + outputs.append(assistant_msg) + + # send a ResponseOutputItemDone for the assistant message + yield ResponseOutputItemDoneEvent( + item=assistant_msg, + output_index=0, + type="response.output_item.done", + ) + + for function_call in state.function_calls.values(): + outputs.append(function_call) + + final_response = response.model_copy() + final_response.output = outputs + + # Get final token counts using consistent method + input_tokens = estimated_input_tokens + output_tokens = estimated_output_tokens + + # Use API token counts if available and reasonable + if usage and hasattr(usage, 'prompt_tokens') and usage.prompt_tokens > 0: + input_tokens = usage.prompt_tokens + if usage and hasattr(usage, 'completion_tokens') and usage.completion_tokens > 0: + output_tokens = usage.completion_tokens + + # Create a proper usage object with our token counts + final_response.usage = ResponseUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + output_tokens_details=OutputTokensDetails( + reasoning_tokens=usage.completion_tokens_details.reasoning_tokens + if usage and hasattr(usage, 'completion_tokens_details') + and usage.completion_tokens_details + and hasattr(usage.completion_tokens_details, 'reasoning_tokens') + and usage.completion_tokens_details.reasoning_tokens + else 0 + ), + input_tokens_details={ + "prompt_tokens": input_tokens, + "cached_tokens": usage.prompt_tokens_details.cached_tokens + if usage and hasattr(usage, 'prompt_tokens_details') + and usage.prompt_tokens_details + and hasattr(usage.prompt_tokens_details, 'cached_tokens') + and usage.prompt_tokens_details.cached_tokens + else 0 }, ) - # --- Add to message_history: user, system prompts --- - if system_instructions: - sys_msg = { - "role": "system", - "content": system_instructions - } - add_to_message_history(sys_msg) - - if isinstance(input, str): - user_msg = { - "role": "user", - "content": input - } - add_to_message_history(user_msg) - elif isinstance(input, list): - for item in input: - if isinstance(item, dict): - if item.get("role") == "user": - user_msg = { - "role": "user", - "content": item.get("content", "") - } - add_to_message_history(user_msg) - # Get token count estimate before API call for consistent counting - estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) - - response, stream = await self._fetch_response( - system_instructions, - input, - model_settings, - tools, - output_schema, - handoffs, - span_generation, - tracing, - stream=True, - ) - usage: CompletionUsage | None = None - state = _StreamingState() - - # Manual token counting (when API doesn't provide it) - output_text = "" - estimated_output_tokens = 0 - - # Initialize a streaming text accumulator for rich display - streaming_text_buffer = "" - # For tool call streaming, accumulate tool_calls to add to message_history at the end - streamed_tool_calls = [] - - # Ollama specific: accumulate full content to check for function calls at the end - # Some Ollama models output the function call as JSON in the text content - ollama_full_content = "" - is_ollama = False - - model_str = str(self.model).lower() - is_ollama = self.is_ollama or "ollama" in model_str or ":" in model_str or "qwen" in model_str - - # Add visual separation before agent output - if streaming_context and should_show_rich_stream: - # If we're using rich context, we'll add separation through that - pass - else: - # Print clear visual separator - print("\n") - - async for chunk in stream: - if not state.started: - state.started = True - yield ResponseCreatedEvent( - response=response, - type="response.created", - ) - - # The usage is only available in the last chunk - if hasattr(chunk, 'usage'): - usage = chunk.usage - # For Ollama/LiteLLM streams that don't have usage attribute - else: - usage = None - - # Handle different stream chunk formats - if hasattr(chunk, 'choices') and chunk.choices: - choices = chunk.choices - elif hasattr(chunk, 'delta') and chunk.delta: - # Some providers might return delta directly - choices = [{"delta": chunk.delta}] - elif isinstance(chunk, dict) and 'choices' in chunk: - choices = chunk['choices'] - # Special handling for Qwen/Ollama chunks - elif isinstance(chunk, dict) and ('content' in chunk or 'function_call' in chunk): - # Qwen direct delta format - convert to standard - choices = [{"delta": chunk}] - else: - # Skip chunks that don't contain choice data - continue - - if not choices or len(choices) == 0: - continue - - # Get the delta content - delta = None - if hasattr(choices[0], 'delta'): - delta = choices[0].delta - elif isinstance(choices[0], dict) and 'delta' in choices[0]: - delta = choices[0]['delta'] - - if not delta: - continue - - # Handle text - content = None - if hasattr(delta, 'content') and delta.content is not None: - content = delta.content - elif isinstance(delta, dict) and 'content' in delta and delta['content'] is not None: - content = delta['content'] - - if content: - # For Ollama, we need to accumulate the full content to check for function calls - if is_ollama: - ollama_full_content += 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) - estimated_output_tokens = token_count - - if not state.text_content_index_and_output: - # Initialize a content tracker for streaming text - state.text_content_index_and_output = ( - 0 if not state.refusal_content_index_and_output else 1, - ResponseOutputText( - text="", - type="output_text", - annotations=[], - ), - ) - # Start a new assistant message stream - assistant_item = ResponseOutputMessage( - id=FAKE_RESPONSES_ID, - content=[], - role="assistant", - type="message", - status="in_progress", - ) - # Notify consumers of the start of a new output message + first content part - yield ResponseOutputItemAddedEvent( - item=assistant_item, - output_index=0, - type="response.output_item.added", - ) - yield ResponseContentPartAddedEvent( - content_index=state.text_content_index_and_output[0], - item_id=FAKE_RESPONSES_ID, - output_index=0, - part=ResponseOutputText( - text="", - type="output_text", - annotations=[], - ), - type="response.content_part.added", - ) - # Emit the delta for this segment of content - yield ResponseTextDeltaEvent( - content_index=state.text_content_index_and_output[0], - delta=content, - item_id=FAKE_RESPONSES_ID, - output_index=0, - type="response.output_text.delta", - ) - # Accumulate the text into the response part - state.text_content_index_and_output[1].text += content - - # Handle refusals (model declines to answer) - refusal_content = None - if hasattr(delta, 'refusal') and delta.refusal: - refusal_content = delta.refusal - elif isinstance(delta, dict) and 'refusal' in delta and delta['refusal']: - refusal_content = delta['refusal'] - - if refusal_content: - if not state.refusal_content_index_and_output: - # Initialize a content tracker for streaming refusal text - state.refusal_content_index_and_output = ( - 0 if not state.text_content_index_and_output else 1, - ResponseOutputRefusal(refusal="", type="refusal"), - ) - # Start a new assistant message if one doesn't exist yet (in-progress) - assistant_item = ResponseOutputMessage( - id=FAKE_RESPONSES_ID, - content=[], - role="assistant", - type="message", - status="in_progress", - ) - # Notify downstream that assistant message + first content part are starting - yield ResponseOutputItemAddedEvent( - item=assistant_item, - output_index=0, - type="response.output_item.added", - ) - yield ResponseContentPartAddedEvent( - content_index=state.refusal_content_index_and_output[0], - item_id=FAKE_RESPONSES_ID, - output_index=0, - part=ResponseOutputText( - text="", - type="output_text", - annotations=[], - ), - type="response.content_part.added", - ) - # Emit the delta for this segment of refusal - yield ResponseRefusalDeltaEvent( - content_index=state.refusal_content_index_and_output[0], - delta=refusal_content, - item_id=FAKE_RESPONSES_ID, - output_index=0, - type="response.refusal.delta", - ) - # Accumulate the refusal string in the output part - state.refusal_content_index_and_output[1].refusal += refusal_content - - # Handle tool calls - # Because we don't know the name of the function until the end of the stream, we'll - # save everything and yield events at the end - tool_calls = self._detect_and_format_function_calls(delta) - - if tool_calls: - for tc_delta in tool_calls: - tc_index = tc_delta.index if hasattr(tc_delta, 'index') else tc_delta.get('index', 0) - if tc_index not in state.function_calls: - state.function_calls[tc_index] = ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - arguments="", - name="", - type="function_call", - call_id="", - ) - - tc_function = None - if hasattr(tc_delta, 'function'): - tc_function = tc_delta.function - elif isinstance(tc_delta, dict) and 'function' in tc_delta: - tc_function = tc_delta['function'] - - if tc_function: - # Handle both object and dict formats - args = "" - if hasattr(tc_function, 'arguments'): - args = tc_function.arguments or "" - elif isinstance(tc_function, dict) and 'arguments' in tc_function: - args = tc_function.get('arguments', "") or "" - - name = "" - if hasattr(tc_function, 'name'): - name = tc_function.name or "" - elif isinstance(tc_function, dict) and 'name' in tc_function: - name = tc_function.get('name', "") or "" - - state.function_calls[tc_index].arguments += args - state.function_calls[tc_index].name += name - - # Handle call_id in both formats - call_id = "" - if hasattr(tc_delta, 'id'): - call_id = tc_delta.id or "" - elif isinstance(tc_delta, dict) and 'id' in tc_delta: - call_id = tc_delta.get('id', "") or "" - else: - # For Qwen models, generate a predictable ID if none is provided - if state.function_calls[tc_index].name: - # Generate a stable ID from the function name and arguments - call_id = f"call_{hashlib.md5(state.function_calls[tc_index].name.encode()).hexdigest()[:8]}" - - state.function_calls[tc_index].call_id += call_id - - # --- Accumulate tool call for message_history --- - # Only add if not already present (avoid duplicates in streaming) - tool_call_msg = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": state.function_calls[tc_index].call_id, - "type": "function", - "function": { - "name": state.function_calls[tc_index].name, - "arguments": state.function_calls[tc_index].arguments - } - } - ] - } - # Only add if not already in streamed_tool_calls - if tool_call_msg not in streamed_tool_calls: - streamed_tool_calls.append(tool_call_msg) - add_to_message_history(tool_call_msg) - - # Special handling for Ollama - check if accumulated text contains a valid function call - if is_ollama and ollama_full_content and len(state.function_calls) == 0: - # Look for JSON object that might be a function call - try: - # Try to extract a JSON object from the content - json_start = ollama_full_content.find('{') - json_end = ollama_full_content.rfind('}') + 1 - - if json_start >= 0 and json_end > json_start: - json_str = ollama_full_content[json_start:json_end] - # Try to parse the JSON - parsed = json.loads(json_str) - - # Check if it looks like a function call - if ('name' in parsed and 'arguments' in parsed): - logger.debug(f"Found valid function call in Ollama output: {json_str}") - - # Create a tool call ID - tool_call_id = f"call_{hashlib.md5((parsed['name'] + str(time.time())).encode()).hexdigest()[:8]}" - - # Ensure arguments is a valid JSON string - arguments_str = "" - if isinstance(parsed['arguments'], dict): - # Remove 'ctf' field if it exists - if 'ctf' in parsed['arguments']: - del parsed['arguments']['ctf'] - arguments_str = json.dumps(parsed['arguments']) - elif isinstance(parsed['arguments'], str): - # If it's already a string, check if it's valid JSON - try: - # Try parsing to validate and remove 'ctf' if present - args_dict = json.loads(parsed['arguments']) - if isinstance(args_dict, dict) and 'ctf' in args_dict: - del args_dict['ctf'] - arguments_str = json.dumps(args_dict) - except: - # If not valid JSON, encode it as a JSON string - arguments_str = json.dumps(parsed['arguments']) - else: - # For any other type, convert to string and then JSON - arguments_str = json.dumps(str(parsed['arguments'])) - # Add it to our function_calls state - state.function_calls[0] = ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - arguments=arguments_str, - name=parsed['name'], - type="function_call", - call_id=tool_call_id, - ) - - # Display the tool call in CLI - from cai.util import cli_print_agent_messages - try: - # Create a message-like object to display the function call - tool_msg = type('ToolCallWrapper', (), { - 'content': None, - 'tool_calls': [ - type('ToolCallDetail', (), { - 'function': type('FunctionDetail', (), { - 'name': parsed['name'], - 'arguments': arguments_str - }), - 'id': tool_call_id, - 'type': 'function' - }) - ] - }) - - # Print the tool call using the CLI utility - cli_print_agent_messages( - agent_name=getattr(self, 'agent_name', 'Agent'), - message=tool_msg, - counter=getattr(self, 'interaction_counter', 0), - model=str(self.model), - debug=False, - interaction_input_tokens=estimated_input_tokens, - interaction_output_tokens=estimated_output_tokens, - interaction_reasoning_tokens=0, # Not available for Ollama - total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, - total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, - total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), - interaction_cost=None, - total_cost=None, - tool_output=None # Will be shown once the tool is executed - ) - except Exception as e: - logger.error(f"Error displaying tool call in CLI: {e}") - - # Add to message history - tool_call_msg = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": tool_call_id, - "type": "function", - "function": { - "name": parsed['name'], - "arguments": arguments_str - } - } - ] - } - - streamed_tool_calls.append(tool_call_msg) - add_to_message_history(tool_call_msg) - - logger.debug(f"Added function call: {parsed['name']} with args: {arguments_str}") - except Exception as e: - pass - - function_call_starting_index = 0 - if state.text_content_index_and_output: - function_call_starting_index += 1 - # Send end event for this content part - yield ResponseContentPartDoneEvent( - content_index=state.text_content_index_and_output[0], - item_id=FAKE_RESPONSES_ID, - output_index=0, - part=state.text_content_index_and_output[1], - type="response.content_part.done", + yield ResponseCompletedEvent( + response=final_response, + type="response.completed", ) - - if state.refusal_content_index_and_output: - function_call_starting_index += 1 - # Send end event for this content part - yield ResponseContentPartDoneEvent( - content_index=state.refusal_content_index_and_output[0], - item_id=FAKE_RESPONSES_ID, - output_index=0, - part=state.refusal_content_index_and_output[1], - type="response.content_part.done", - ) - - # Actually send events for the function calls - for function_call in state.function_calls.values(): - # First, a ResponseOutputItemAdded for the function call - yield ResponseOutputItemAddedEvent( - item=ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - call_id=function_call.call_id, - arguments=function_call.arguments, - name=function_call.name, - type="function_call", + + # Update token totals for CLI display + if final_response.usage: + # Always update the total counters with the best available counts + self.total_input_tokens += final_response.usage.input_tokens + self.total_output_tokens += final_response.usage.output_tokens + if (final_response.usage.output_tokens_details and + hasattr(final_response.usage.output_tokens_details, 'reasoning_tokens')): + self.total_reasoning_tokens += final_response.usage.output_tokens_details.reasoning_tokens + + # Prepare final statistics for display + interaction_input = final_response.usage.input_tokens if final_response.usage else 0 + interaction_output = final_response.usage.output_tokens if final_response.usage else 0 + total_input = getattr(self, 'total_input_tokens', 0) + total_output = getattr(self, 'total_output_tokens', 0) + + # Calculate costs using the same token counts - ensure model is a string + model_name = str(self.model) + interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output) + total_cost = calculate_model_cost(model_name, total_input, total_output) + + # Explicit conversion to float with fallback to ensure they're never None or 0 + interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0.00001) + total_cost = max(float(total_cost if total_cost is not None else 0.0), 0.00001) + + + # Create final stats with explicit type conversion for all values + final_stats = { + "interaction_input_tokens": int(interaction_input), + "interaction_output_tokens": int(interaction_output), + "interaction_reasoning_tokens": int( + 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 ), - output_index=function_call_starting_index, - type="response.output_item.added", - ) - # Then, yield the args - yield ResponseFunctionCallArgumentsDeltaEvent( - delta=function_call.arguments, - item_id=FAKE_RESPONSES_ID, - output_index=function_call_starting_index, - type="response.function_call_arguments.delta", - ) - # Finally, the ResponseOutputItemDone - yield ResponseOutputItemDoneEvent( - item=ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - call_id=function_call.call_id, - arguments=function_call.arguments, - name=function_call.name, - type="function_call", - ), - output_index=function_call_starting_index, - type="response.output_item.done", - ) - - # Finally, send the Response completed event - outputs: list[ResponseOutputItem] = [] - if state.text_content_index_and_output or state.refusal_content_index_and_output: - assistant_msg = ResponseOutputMessage( - id=FAKE_RESPONSES_ID, - content=[], - role="assistant", - type="message", - status="completed", - ) - if state.text_content_index_and_output: - assistant_msg.content.append(state.text_content_index_and_output[1]) - if state.refusal_content_index_and_output: - assistant_msg.content.append(state.refusal_content_index_and_output[1]) - outputs.append(assistant_msg) - - # send a ResponseOutputItemDone for the assistant message - yield ResponseOutputItemDoneEvent( - item=assistant_msg, - output_index=0, - type="response.output_item.done", - ) - - for function_call in state.function_calls.values(): - outputs.append(function_call) - - final_response = response.model_copy() - final_response.output = outputs - - # Get final token counts using consistent method - input_tokens = estimated_input_tokens - output_tokens = estimated_output_tokens - - # Use API token counts if available and reasonable - if usage and hasattr(usage, 'prompt_tokens') and usage.prompt_tokens > 0: - input_tokens = usage.prompt_tokens - if usage and hasattr(usage, 'completion_tokens') and usage.completion_tokens > 0: - output_tokens = usage.completion_tokens + "total_input_tokens": int(total_input), + "total_output_tokens": int(total_output), + "total_reasoning_tokens": int(getattr(self, 'total_reasoning_tokens', 0)), + "interaction_cost": float(interaction_cost), + "total_cost": float(total_cost), + } - # # Debug information - # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - Streaming final tokens: input={input_tokens}, output={output_tokens}, total={input_tokens + output_tokens}") + # At the end of streaming, finish the streaming context if we were using it + if streaming_context: + # Create a direct copy of the costs to ensure they remain as floats + direct_stats = final_stats.copy() + direct_stats["interaction_cost"] = float(interaction_cost) + direct_stats["total_cost"] = float(total_cost) + # Use the direct copy with guaranteed float costs + finish_agent_streaming(streaming_context, direct_stats) + + # Add visual separation after agent output completes + print("\n") - # Create a proper usage object with our token counts - final_response.usage = ResponseUsage( - input_tokens=input_tokens, - output_tokens=output_tokens, - total_tokens=input_tokens + output_tokens, - output_tokens_details=OutputTokensDetails( - reasoning_tokens=usage.completion_tokens_details.reasoning_tokens - if usage and hasattr(usage, 'completion_tokens_details') - and usage.completion_tokens_details - and hasattr(usage.completion_tokens_details, 'reasoning_tokens') - and usage.completion_tokens_details.reasoning_tokens - else 0 - ), - input_tokens_details={ - "prompt_tokens": input_tokens, - "cached_tokens": usage.prompt_tokens_details.cached_tokens - if usage and hasattr(usage, 'prompt_tokens_details') - and usage.prompt_tokens_details - and hasattr(usage.prompt_tokens_details, 'cached_tokens') - and usage.prompt_tokens_details.cached_tokens - else 0 - }, - ) + if tracing.include_data(): + span_generation.span_data.output = [final_response.model_dump()] - yield ResponseCompletedEvent( - response=final_response, - type="response.completed", - ) - - # Update token totals for CLI display - if final_response.usage: - # Always update the total counters with the best available counts - self.total_input_tokens += final_response.usage.input_tokens - self.total_output_tokens += final_response.usage.output_tokens - if (final_response.usage.output_tokens_details and - hasattr(final_response.usage.output_tokens_details, 'reasoning_tokens')): - self.total_reasoning_tokens += final_response.usage.output_tokens_details.reasoning_tokens - - # Prepare final statistics for display - interaction_input = final_response.usage.input_tokens if final_response.usage else 0 - interaction_output = final_response.usage.output_tokens if final_response.usage else 0 - total_input = getattr(self, 'total_input_tokens', 0) - total_output = getattr(self, 'total_output_tokens', 0) - - # Calculate costs using the same token counts - ensure model is a string - model_name = str(self.model) - interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output) - total_cost = calculate_model_cost(model_name, total_input, total_output) - - # Explicit conversion to float with fallback to ensure they're never None or 0 - interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0.00001) - total_cost = max(float(total_cost if total_cost is not None else 0.0), 0.00001) - - - # Create final stats with explicit type conversion for all values - final_stats = { - "interaction_input_tokens": int(interaction_input), - "interaction_output_tokens": int(interaction_output), - "interaction_reasoning_tokens": int( - 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": int(total_input), - "total_output_tokens": int(total_output), - "total_reasoning_tokens": int(getattr(self, 'total_reasoning_tokens', 0)), - "interaction_cost": float(interaction_cost), - "total_cost": float(total_cost), - } - - # At the end of streaming, finish the streaming context if we were using it + span_generation.span_data.usage = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + + # --- Add assistant tool call(s) to message_history at the end of streaming --- + for tool_call_msg in streamed_tool_calls: + add_to_message_history(tool_call_msg) + # If there was only text output, add that as an assistant message + if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: + asst_msg = { + "role": "assistant", + "content": state.text_content_index_and_output[1].text + } + add_to_message_history(asst_msg) + except Exception as e: + # Ensure streaming context is cleaned up in case of errors if streaming_context: - # Create a direct copy of the costs to ensure they remain as floats - direct_stats = final_stats.copy() - direct_stats["interaction_cost"] = float(interaction_cost) - direct_stats["total_cost"] = float(total_cost) - # Use the direct copy with guaranteed float costs - finish_agent_streaming(streaming_context, direct_stats) - - # Add visual separation after agent output completes - print("\n") - # 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': - cli_print_agent_messages( - agent_name=getattr(self, 'agent_name', 'Agent'), - message=item, - counter=getattr(self, 'interaction_counter', 0), - model=str(self.model), - debug=False, - interaction_input_tokens=interaction_input, - interaction_output_tokens=interaction_output, - interaction_reasoning_tokens=final_stats["interaction_reasoning_tokens"], - total_input_tokens=total_input, - total_output_tokens=total_output, - total_reasoning_tokens=final_stats["total_reasoning_tokens"], - interaction_cost=interaction_cost, - total_cost=total_cost, - ) - - # Add visual separation after message - print("\n") - break - - # --- Add assistant tool call(s) to message_history at the end of streaming --- - for tool_call_msg in streamed_tool_calls: - add_to_message_history(tool_call_msg) - # If there was only text output, add that as an assistant message - if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: - asst_msg = { - "role": "assistant", - "content": state.text_content_index_and_output[1].text - } - add_to_message_history(asst_msg) - - if tracing.include_data(): - span_generation.span_data.output = [final_response.model_dump()] - - span_generation.span_data.usage = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - } - - # To avoid duplicate tool output display, we need to track tool calls - # Add this after the completion response is received - - if not stream and hasattr(response, 'choices') and len(response.choices) > 0: - # For non-streaming responses, make sure we capture tool call IDs - # to prevent duplicate printing - choice = response.choices[0] - if hasattr(choice, 'message') and hasattr(choice.message, 'tool_calls'): - for tool_call in choice.message.tool_calls: - if hasattr(tool_call, 'id'): - # Register this tool call ID as already seen - from cai.util import cli_print_tool_output - if not hasattr(cli_print_tool_output, '_seen_calls'): - cli_print_tool_output._seen_calls = {} - cli_print_tool_output._seen_calls[tool_call.id] = True + try: + finish_agent_streaming(streaming_context, None) + except Exception: + pass + raise e @overload async def _fetch_response( diff --git a/src/cai/util.py b/src/cai/util.py index 0eba2afe..52027fd5 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -899,179 +899,247 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli console.print(tool_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 + """ + Create a streaming context object that maintains state for streaming agent output. - # Use the model from env if available - model_override = os.getenv('CAI_MODEL') - if model_override: - model = model_override + Args: + agent_name: The name of the agent to display + counter: The interaction counter (turn number) + model: The model name - timestamp = datetime.now().strftime("%H:%M:%S") - - # Terminal size for better 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), - title="[bold]Agent Streaming Response[/bold]", - title_align="left", - width=panel_width, - expand=True - ) - - # Create Live display object but don't start it until we have content - live = Live(panel, refresh_per_second=20, console=console, auto_refresh=False) - - return { - "live": live, - "panel": panel, - "header": header, - "content": content, - "footer": footer, - "timestamp": timestamp, - "model": model, - "agent_name": agent_name, - "panel_width": panel_width, - "is_started": False # Track if we've started the display - } + Returns: + A dictionary with the streaming context + """ + try: + from rich.live import Live + import shutil + + # Use the model from env if available + model_override = os.getenv('CAI_MODEL') + if model_override: + model = model_override + + timestamp = datetime.now().strftime("%H:%M:%S") + + # Terminal size for better 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), + title="[bold]Agent Streaming Response[/bold]", + title_align="left", + width=panel_width, + expand=True + ) + + # Create Live display object but don't start it until we have content + live = Live(panel, refresh_per_second=20, console=console, auto_refresh=False) + + return { + "live": live, + "panel": panel, + "header": header, + "content": content, + "footer": footer, + "timestamp": timestamp, + "model": model, + "agent_name": agent_name, + "panel_width": panel_width, + "is_started": False, # Track if we've started the display + "error": None, # Track any errors + } + except Exception as e: + # If rich display fails, return None and log the error + import sys + print(f"Error creating streaming context: {e}", file=sys.stderr) + return None def update_agent_streaming_content(context, text_delta): - """Update the streaming content with new text.""" - # Parse the text_delta to get just the content if needed - parsed_delta = parse_message_content(text_delta) + """ + Update the streaming content with new text. - # Skip empty updates to avoid showing an empty panel - if not parsed_delta or parsed_delta.strip() == "": - return - - # Add the parsed text to the content - context["content"].append(parsed_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), - title="[bold]Agent Streaming Response[/bold]", - title_align="left", - width=context.get("panel_width", 100), - expand=True - ) - - # Check if we need to start the display - if not context.get("is_started", False): - context["live"].start() - context["is_started"] = True - - # Force an update with the new panel - context["live"].update(updated_panel) - context["panel"] = updated_panel - context["live"].refresh() + Args: + context: The streaming context created by create_agent_streaming_context + text_delta: The new text to add + """ + if not context: + return False + + try: + # Parse the text_delta to get just the content if needed + parsed_delta = parse_message_content(text_delta) + + # Skip empty updates to avoid showing an empty panel + if not parsed_delta or parsed_delta.strip() == "": + return True + + # Add the parsed text to the content + context["content"].append(parsed_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), + title="[bold]Agent Streaming Response[/bold]", + title_align="left", + width=context.get("panel_width", 100), + expand=True + ) + + # Check if we need to start the display + if not context.get("is_started", False): + try: + context["live"].start() + context["is_started"] = True + except Exception as e: + context["error"] = str(e) + return False + + # Force an update with the new panel + context["live"].update(updated_panel) + context["panel"] = updated_panel + context["live"].refresh() + return True + except Exception as e: + # If there's an error, set it in the context + context["error"] = str(e) + return False def finish_agent_streaming(context, final_stats=None): - """Finish the streaming session and display final stats if available.""" - # Check if there's actual content to display - don't show empty panels - if not context["content"] or context["content"].plain == "": - # If the display was never started, nothing to do - if not context.get("is_started", False): - return - # Otherwise, stop the display without showing final panel - context["live"].stop() - return + """ + 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") + Args: + context: The streaming context to finish + final_stats: Optional dictionary with token statistics and costs + """ + if not context: + return False - # Ensure costs are properly extracted and preserved as floats - interaction_cost = float(final_stats.get("interaction_cost", 0.0)) - total_cost = float(final_stats.get("total_cost", 0.0)) + try: + # Check if there's actual content to display - don't show empty panels + if not context["content"] or context["content"].plain == "": + # If the display was never started, nothing to do + if not context.get("is_started", False): + return True + # Otherwise, stop the display without showing final panel + try: + context["live"].stop() + except Exception: + pass + return True - model_name = context.get("model", "") - # If model is not a string, use env - if not isinstance(model_name, str): - model_name = os.environ.get('CAI_MODEL', 'gpt-4o-mini') - - 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): + # 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") - # Only calculate costs if they weren't provided or are zero - if interaction_cost is None or interaction_cost == 0.0: - interaction_cost = calculate_model_cost(model_name, interaction_input_tokens, interaction_output_tokens) - if total_cost is None or total_cost == 0.0: - total_cost = calculate_model_cost(model_name, total_input_tokens, total_output_tokens) + # Ensure costs are properly extracted and preserved as floats + interaction_cost = float(final_stats.get("interaction_cost", 0.0)) + total_cost = float(final_stats.get("total_cost", 0.0)) - tokens_text = _create_token_display( - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - model_name, # string model name! - interaction_cost, - total_cost - ) - - 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), - 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 - time.sleep(0.5) - - # Stop the live display - context["live"].stop() + model_name = context.get("model", "") + # If model is not a string, use env + if not isinstance(model_name, str): + model_name = os.environ.get('CAI_MODEL', 'gpt-4o-mini') + + 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): + + # Only calculate costs if they weren't provided or are zero + if interaction_cost is None or interaction_cost == 0.0: + interaction_cost = calculate_model_cost(model_name, interaction_input_tokens, interaction_output_tokens) + if total_cost is None or total_cost == 0.0: + total_cost = calculate_model_cost(model_name, total_input_tokens, total_output_tokens) + + tokens_text = _create_token_display( + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + model_name, # string model name! + interaction_cost, + total_cost + ) + + 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), + 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 + time.sleep(0.1) + + # Stop the live display + try: + context["live"].stop() + except Exception as e: + context["error"] = str(e) + return False + + return True + except Exception as e: + # If there's an error, print it if the context hasn't already tracked one + if not context.get("error"): + context["error"] = str(e) + + # Try to stop the live display even if there was an error + try: + if context.get("is_started", False) and context.get("live"): + context["live"].stop() + except Exception: + pass + + return False def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None): """