From b2233b76d178649ff1bfea3984d895ac76f13b73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Thu, 3 Apr 2025 08:19:42 +0200 Subject: [PATCH] CLI output operational MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- examples/cai/simple_one_tool_test.py | 8 +- examples/cai/simple_one_tool_test_streamed.py | 6 +- src/cai/cli.py | 2 +- .../agents/models/openai_chatcompletions.py | 332 ++++++++++++++---- src/cai/sdk/agents/models/openai_responses.py | 105 ++++++ src/cai/sdk/agents/run.py | 17 +- src/cai/util.py | 320 ++++++++++++++--- 7 files changed, 657 insertions(+), 133 deletions(-) diff --git a/examples/cai/simple_one_tool_test.py b/examples/cai/simple_one_tool_test.py index a4949612..31f791c8 100644 --- a/examples/cai/simple_one_tool_test.py +++ b/examples/cai/simple_one_tool_test.py @@ -8,11 +8,13 @@ is working correctly. import os import asyncio +import json from dotenv import load_dotenv from openai import AsyncOpenAI from cai.sdk.agents import Runner, set_default_openai_client from cai.agents import get_agent_by_name -from cai.util import fix_litellm_transcription_annotations, color +from cai.util import fix_litellm_transcription_annotations, color, cli_print_agent_messages +from cai.sdk.agents.models._openai_shared import set_use_responses_by_default # Load environment variables @@ -31,11 +33,13 @@ async def main(): patch_applied = fix_litellm_transcription_annotations() if not patch_applied: print(color("Something went wrong patching LiteLLM fix_litellm_transcription_annotations", color="red")) + + # Force the use of OpenAIChatCompletionsModel instead of OpenAIResponsesModel + set_use_responses_by_default(False) # Get the one_tool agent agent = get_agent_by_name("one_tool_agent") - print("Testing one_tool agent with a simple hello message...") print(f"Using model: {os.getenv('CAI_MODEL', 'default')}") # Run the agent with a simple test message diff --git a/examples/cai/simple_one_tool_test_streamed.py b/examples/cai/simple_one_tool_test_streamed.py index 9fc3bfa3..00632d49 100644 --- a/examples/cai/simple_one_tool_test_streamed.py +++ b/examples/cai/simple_one_tool_test_streamed.py @@ -8,6 +8,7 @@ is working correctly, with streaming output. import os import asyncio +import json from dotenv import load_dotenv from openai import AsyncOpenAI from openai.types.responses import ResponseTextDeltaEvent @@ -15,6 +16,7 @@ from cai.sdk.agents import Runner, set_default_openai_client from cai.agents import get_agent_by_name from cai.util import fix_litellm_transcription_annotations, color from cai.sdk.agents import Agent, OpenAIChatCompletionsModel +from cai.sdk.agents.models._openai_shared import set_use_responses_by_default # Load environment variables load_dotenv() @@ -33,10 +35,12 @@ async def main(): if not patch_applied: print(color("Something went wrong patching LiteLLM fix_litellm_transcription_annotations", color="red")) + # Force the use of OpenAIChatCompletionsModel instead of OpenAIResponsesModel + set_use_responses_by_default(False) + # Get the one_tool agent agent = get_agent_by_name("one_tool_agent") - print("Testing one_tool agent with a simple hello message (streaming mode)...") print(f"Using model: {os.getenv('CAI_MODEL', 'default')}") # Stream indicator diff --git a/src/cai/cli.py b/src/cai/cli.py index e19e4b89..f867131b 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -213,7 +213,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= 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() # Add a newline at the end + print("\n") # Add a newline at the end return result except Exception as e: print() # Add a newline after any partial output diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 8563c758..3807b15c 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -5,11 +5,12 @@ import json import time import os import litellm +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 +from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages from wasabi import color from openai import NOT_GIVEN, AsyncOpenAI, AsyncStream, NotGiven @@ -103,6 +104,77 @@ class _StreamingState: function_calls: dict[int, ResponseFunctionToolCall] = field(default_factory=dict) +# Add a new function for consistent token counting using tiktoken +def count_tokens_with_tiktoken(text_or_messages): + """ + Count tokens consistently using tiktoken library. + Works with both strings and message lists. + Returns a tuple of (input_tokens, reasoning_tokens). + """ + if not text_or_messages: + return 0, 0 + + try: + # Try to use cl100k_base encoding (used by GPT-4 and GPT-3.5-turbo) + encoding = tiktoken.get_encoding("cl100k_base") + except: + # Fall back to GPT-2 encoding if cl100k is not available + try: + encoding = tiktoken.get_encoding("gpt2") + except: + # If tiktoken fails, fall back to character estimate + if isinstance(text_or_messages, str): + return len(text_or_messages) // 4, 0 + elif isinstance(text_or_messages, list): + total_len = 0 + for msg in text_or_messages: + if isinstance(msg, dict) and 'content' in msg: + if isinstance(msg['content'], str): + total_len += len(msg['content']) + return total_len // 4, 0 + else: + return 0, 0 + + # Process different input types + if isinstance(text_or_messages, str): + token_count = len(encoding.encode(text_or_messages)) + return token_count, 0 + elif isinstance(text_or_messages, list): + total_tokens = 0 + reasoning_tokens = 0 + + # Add tokens for the messages format (ChatML format overhead) + # Each message has a base overhead (usually ~4 tokens) + total_tokens += len(text_or_messages) * 4 + + for msg in text_or_messages: + if isinstance(msg, dict): + # Add tokens for role + if 'role' in msg: + total_tokens += len(encoding.encode(msg['role'])) + + # Count content tokens + if 'content' in msg and msg['content']: + if isinstance(msg['content'], str): + content_tokens = len(encoding.encode(msg['content'])) + total_tokens += content_tokens + + # Count tokens in assistant messages as reasoning tokens + if msg.get('role') == 'assistant': + reasoning_tokens += content_tokens + elif isinstance(msg['content'], list): + for content_part in msg['content']: + if isinstance(content_part, dict) and 'text' in content_part: + part_tokens = len(encoding.encode(content_part['text'])) + total_tokens += part_tokens + if msg.get('role') == 'assistant': + reasoning_tokens += part_tokens + + return total_tokens, reasoning_tokens + else: + return 0, 0 + + class OpenAIChatCompletionsModel(Model): def __init__( self, @@ -115,6 +187,17 @@ class OpenAIChatCompletionsModel(Model): self.is_ollama = os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false' self.empty_content_error_shown = False + # Track interaction counter and token totals for cli display + self.interaction_counter = 0 + self.total_input_tokens = 0 + self.total_output_tokens = 0 + self.total_reasoning_tokens = 0 + self.agent_name = "Agent" # Default name + + def set_agent_name(self, name: str) -> None: + """Set the agent name for CLI display purposes.""" + self.agent_name = name + def _non_null_or_not_given(self, value: Any) -> Any: return value if value is not None else NOT_GIVEN @@ -128,12 +211,29 @@ class OpenAIChatCompletionsModel(Model): handoffs: list[Handoff], tracing: ModelTracing, ) -> ModelResponse: + # Increment the interaction counter for CLI display + self.interaction_counter += 1 + 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 the messages for consistent token counting + converted_messages = _Converter.items_to_messages(input) + if system_instructions: + converted_messages.insert( + 0, + { + "content": system_instructions, + "role": "system", + }, + ) + + # Get token count estimate before API call for consistent counting + estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) + response = await self._fetch_response( system_instructions, input, @@ -153,14 +253,68 @@ class OpenAIChatCompletionsModel(Model): f"LLM resp:\n{json.dumps(response.choices[0].message.model_dump(), indent=2)}\n" ) + # Ensure we have reasonable token counts + if response.usage: + input_tokens = response.usage.prompt_tokens + output_tokens = response.usage.completion_tokens + total_tokens = response.usage.total_tokens + + # Use estimated tokens if API returns zeroes or implausible values + if input_tokens == 0 or input_tokens < (len(str(input)) // 10): # Sanity check + input_tokens = estimated_input_tokens + total_tokens = input_tokens + output_tokens + + # # Debug information + # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - API tokens: input={input_tokens}, output={output_tokens}, total={total_tokens}") + # print(f"Estimated tokens were: input={estimated_input_tokens}") + else: + # If no usage info, use our estimates + input_tokens = estimated_input_tokens + output_tokens = 0 + total_tokens = input_tokens + # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - No API tokens, using estimates: input={input_tokens}, output={output_tokens}") + + + # Update token totals for CLI display + self.total_input_tokens += input_tokens + self.total_output_tokens += output_tokens + if (response.usage and + hasattr(response.usage, 'completion_tokens_details') and + response.usage.completion_tokens_details and + hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')): + self.total_reasoning_tokens += response.usage.completion_tokens_details.reasoning_tokens + + # Print the agent message for CLI display + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), # Default to 'Agent' if not available + message=response.choices[0].message, + counter=getattr(self, 'interaction_counter', 0), # Default to 0 if not available + model=str(self.model), + debug=False, + interaction_input_tokens=input_tokens, + interaction_output_tokens=output_tokens, + 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 hasattr(response.usage.completion_tokens_details, 'reasoning_tokens') + else 0 + ), + total_input_tokens=getattr(self, 'total_input_tokens', 0), # Will need to be tracked elsewhere + total_output_tokens=getattr(self, 'total_output_tokens', 0), # Will need to be tracked elsewhere + total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), # Will need to be tracked elsewhere + interaction_cost=None, # Would need cost calculation logic + total_cost=None, # Would need cost calculation logic + ) + usage = ( Usage( requests=1, - input_tokens=response.usage.prompt_tokens, - output_tokens=response.usage.completion_tokens, - total_tokens=response.usage.total_tokens, + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, ) - if response.usage + if response.usage or input_tokens > 0 else Usage() ) if tracing.include_data(): @@ -191,12 +345,29 @@ class OpenAIChatCompletionsModel(Model): """ Yields a partial message as it is generated, as well as the usage information. """ + # Increment the interaction counter for CLI display + self.interaction_counter += 1 + 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", + }, + ) + + # 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, @@ -211,7 +382,11 @@ class OpenAIChatCompletionsModel(Model): usage: CompletionUsage | None = None state = _StreamingState() - + + # Manual token counting (when API doesn't provide it) + output_text = "" + estimated_output_tokens = 0 + async for chunk in stream: if not state.started: state.started = True @@ -260,6 +435,11 @@ class OpenAIChatCompletionsModel(Model): content = delta['content'] if 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 = ( @@ -495,56 +675,91 @@ class OpenAIChatCompletionsModel(Model): final_response = response.model_copy() final_response.output = outputs - final_response.usage = ( - ResponseUsage( - input_tokens=usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0, - output_tokens=usage.completion_tokens if usage and hasattr(usage, 'completion_tokens') else 0, - total_tokens=usage.total_tokens if usage and hasattr(usage, 'total_tokens') else 0, - 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": usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0, - "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 usage is not None - else ResponseUsage( - input_tokens=0, - output_tokens=0, - total_tokens=0, - output_tokens_details=OutputTokensDetails(reasoning_tokens=0), - input_tokens_details={"prompt_tokens": 0, "cached_tokens": 0} - ) + # 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 + + # # Debug information + # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - Streaming final tokens: input={input_tokens}, output={output_tokens}, total={input_tokens + output_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 + }, ) 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 + + # 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): + # 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'), # Default to 'Agent' if not available + message=item, + counter=getattr(self, 'interaction_counter', 0), # Default to 0 if not available + model=str(self.model), + debug=False, + 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, + ) + break + if tracing.include_data(): span_generation.span_data.output = [final_response.model_dump()] - if usage: - span_generation.span_data.usage = { - "input_tokens": usage.prompt_tokens, - "output_tokens": usage.completion_tokens, - } - else: - span_generation.span_data.usage = { - "input_tokens": 0, - "output_tokens": 0, - } + span_generation.span_data.usage = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } @overload async def _fetch_response( @@ -677,30 +892,9 @@ class OpenAIChatCompletionsModel(Model): return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) else: return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - - # ret = await self._get_client().chat.completions.create(**kwargs) - - # if isinstance(ret, ChatCompletion): - # return ret - - # response = Response( - # id=FAKE_RESPONSES_ID, - # created_at=time.time(), - # model=self.model, - # object="response", - # output=[], - # tool_choice=cast(Literal["auto", "required", "none"], tool_choice) - # if tool_choice != NOT_GIVEN - # else "auto", - # top_p=model_settings.top_p, - # temperature=model_settings.temperature, - # tools=[], - # parallel_tool_calls=parallel_tool_calls or False, - # ) - # return response, ret except litellm.exceptions.BadRequestError as e: - print(color("BadRequestError encountered: " + str(e), fg="yellow")) + # print(color("BadRequestError encountered: " + str(e), fg="yellow")) if "LLM Provider NOT provided" in str(e): # Create a copy of params to avoid overwriting the original # ones diff --git a/src/cai/sdk/agents/models/openai_responses.py b/src/cai/sdk/agents/models/openai_responses.py index 17803fda..ce587045 100644 --- a/src/cai/sdk/agents/models/openai_responses.py +++ b/src/cai/sdk/agents/models/openai_responses.py @@ -54,8 +54,20 @@ class OpenAIResponsesModel(Model): model: str | ChatModel, openai_client: AsyncOpenAI, ) -> None: + print(f"\nDEBUG: OpenAIResponsesModel initialized with model: {model}\n") self.model = model self._client = openai_client + + # Track interaction counter and token totals for cli display + self.interaction_counter = 0 + self.total_input_tokens = 0 + self.total_output_tokens = 0 + self.total_reasoning_tokens = 0 + self.agent_name = "Agent" # Default name + + def set_agent_name(self, name: str) -> None: + """Set the agent name for CLI display purposes.""" + self.agent_name = name def _non_null_or_not_given(self, value: Any) -> Any: return value if value is not None else NOT_GIVEN @@ -70,6 +82,9 @@ class OpenAIResponsesModel(Model): handoffs: list[Handoff], tracing: ModelTracing, ) -> ModelResponse: + # Increment the interaction counter for CLI display + self.interaction_counter += 1 + with response_span(disabled=tracing.is_disabled()) as span_response: try: response = await self._fetch_response( @@ -104,6 +119,50 @@ class OpenAIResponsesModel(Model): if tracing.include_data(): span_response.span_data.response = response span_response.span_data.input = input + + # Print the agent message for CLI display + from cai.util import cli_print_agent_messages + try: + # Create a message-like object to display + message_obj = type('ResponseWrapper', (), { + 'content': '\n'.join([ + str(item.get('content', '')) if hasattr(item, 'get') + else str(getattr(item, 'text', '')) + for item in response.output + if hasattr(item, 'get') or hasattr(item, 'text') + ]), + 'tool_calls': [ + type('ToolCallWrapper', (), { + 'name': item.name, + 'arguments': item.arguments + }) + for item in response.output + if hasattr(item, 'name') and hasattr(item, 'arguments') + ] + }) + + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=message_obj, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=usage.input_tokens, + interaction_output_tokens=usage.output_tokens, + interaction_reasoning_tokens=0, # Not available in Responses API + 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, + ) + + # Update token totals + self.total_input_tokens += usage.input_tokens + self.total_output_tokens += usage.output_tokens + except Exception as e: + logger.error(f"Error printing agent message: {e}") + except Exception as e: span_response.set_error( SpanError( @@ -136,6 +195,9 @@ class OpenAIResponsesModel(Model): """ Yields a partial message as it is generated, as well as the usage information. """ + # Increment the interaction counter for CLI display + self.interaction_counter += 1 + with response_span(disabled=tracing.is_disabled()) as span_response: try: stream = await self._fetch_response( @@ -159,6 +221,49 @@ class OpenAIResponsesModel(Model): span_response.span_data.response = final_response span_response.span_data.input = input + # Print the agent message for CLI display + from cai.util import cli_print_agent_messages + try: + # Create a message-like object to display + message_obj = type('ResponseWrapper', (), { + 'content': '\n'.join([ + str(item.get('content', '')) if hasattr(item, 'get') + else str(getattr(item, 'text', '')) + for item in final_response.output + if hasattr(item, 'get') or hasattr(item, 'text') + ]), + 'tool_calls': [ + type('ToolCallWrapper', (), { + 'name': item.name, + 'arguments': item.arguments + }) + for item in final_response.output + if hasattr(item, 'name') and hasattr(item, 'arguments') + ] + }) + + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=message_obj, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=final_response.usage.input_tokens, + interaction_output_tokens=final_response.usage.output_tokens, + interaction_reasoning_tokens=0, # Not available in Responses API + 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, + ) + + # Update token totals + self.total_input_tokens += final_response.usage.input_tokens + self.total_output_tokens += final_response.usage.output_tokens + except Exception as e: + logger.error(f"Error printing agent message: {e}") + except Exception as e: span_response.set_error( SpanError( diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 0159822a..84bb664b 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -932,11 +932,18 @@ class Runner: @classmethod def _get_model(cls, agent: Agent[Any], run_config: RunConfig) -> Model: + model = None if isinstance(run_config.model, Model): - return run_config.model + model = run_config.model elif isinstance(run_config.model, str): - return run_config.model_provider.get_model(run_config.model) + model = run_config.model_provider.get_model(run_config.model) elif isinstance(agent.model, Model): - return agent.model - - return run_config.model_provider.get_model(agent.model) + model = agent.model + else: + model = run_config.model_provider.get_model(agent.model) + + # Set agent name if the model supports it (for CLI display) + if hasattr(model, 'set_agent_name'): + model.set_agent_name(agent.name) + + return model diff --git a/src/cai/util.py b/src/cai/util.py index a24217e4..a2f5cb21 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -9,6 +9,45 @@ from rich.console import Console from rich.tree import Tree from mako.template import Template # pylint: disable=import-error from wasabi import color +from rich.text import Text # pylint: disable=import-error +from rich.panel import Panel # pylint: disable=import-error +from rich.box import ROUNDED # pylint: disable=import-error +from rich.theme import Theme # pylint: disable=import-error +from rich.traceback import install # pylint: disable=import-error +from rich.pretty import install as install_pretty # pylint: disable=import-error # noqa: 501 +from datetime import datetime + +theme = Theme({ + # Primary colors - Material Design inspired + "timestamp": "#00BCD4", # Cyan 500 + "agent": "#4CAF50", # Green 500 + "arrow": "#FFFFFF", # White + "content": "#ECEFF1", # Blue Grey 50 + "tool": "#F44336", # Red 500 + + # Secondary colors + "cost": "#009688", # Teal 500 + "args_str": "#FFC107", # Amber 500 + + # UI elements + "border": "#2196F3", # Blue 500 + "border_state": "#FFD700", # Yellow (Gold), complementary to Blue 500 + "model": "#673AB7", # Deep Purple 500 + "dim": "#9E9E9E", # Grey 500 + "current_token_count": "#E0E0E0", # Grey 300 - Light grey + "total_token_count": "#757575", # Grey 600 - Medium grey + "context_tokens": "#0A0A0A", # Nearly black - Very high contrast + + # Status indicators + "success": "#4CAF50", # Green 500 + "warning": "#FF9800", # Orange 500 + "error": "#F44336" # Red 500 +}) + +console = Console(theme=theme) +install() +install_pretty() + def get_ollama_api_base(): """Get the Ollama API base URL from environment variable or default to localhost:8000.""" @@ -198,68 +237,239 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 elif msg.get("role") == "tool" and msg.get("tool_call_id"): tc_id = msg["tool_call_id"] tool_calls_occurrences.setdefault( - tc_id, []).append( - (i, "tool", None)) - # Step 3: Mark invalid or extra occurrences for removal - removal_messages = set() # Indices of messages (tool type) to remove + tc_id, []).append((i, "tool", None)) + + # Step 3: Mark indices in the message list to remove. # Maps message index (assistant) to set of indices (in tool_calls) to - # remove - removal_assistant_entries = {} + # delete, or directly marks message indices (tool) to delete. + to_remove = {} for tc_id, occurrences in tool_calls_occurrences.items(): - # Only 2 occurrences allowed. Mark extras for removal. - valid_occurrences = occurrences[:2] - extra_occurrences = occurrences[2:] - for occ in extra_occurrences: - msg_idx, typ, j = occ - if typ == "assistant": - removal_assistant_entries.setdefault(msg_idx, set()).add(j) - elif typ == "tool": - removal_messages.add(msg_idx) - # If valid occurrences aren't exactly 2 (i.e., a lonely tool call), - # mark for removal - if len(valid_occurrences) != 2: - for occ in valid_occurrences: - msg_idx, typ, j = occ - if typ == "assistant": - removal_assistant_entries.setdefault( - msg_idx, set()).add(j) - elif typ == "tool": - removal_messages.add(msg_idx) - else: - # If exactly 2 occurrences, ensure both have content - remove_pair = False - for occ in valid_occurrences: - msg_idx, typ, _ = occ - msg_content = messages[msg_idx].get("content") - if msg_content is None or not msg_content.strip(): - remove_pair = True - break - if remove_pair: - for occ in valid_occurrences: - msg_idx, typ, j = occ - if typ == "assistant": - removal_assistant_entries.setdefault( - msg_idx, set()).add(j) - elif typ == "tool": - removal_messages.add(msg_idx) - # Step 4: Build new message list applying removals - new_messages = [] + if len(occurrences) > 2: + # More than one assistant and tool message pair - trim down + # by picking first pairing and removing the rest + assistant_items = [ + occ for occ in occurrences if occ[1] == "assistant"] + tool_items = [occ for occ in occurrences if occ[1] == "tool"] + + if assistant_items and tool_items: + valid_assistant = assistant_items[0] + valid_tool = tool_items[0] + for item in occurrences: + if item != valid_assistant and item != valid_tool: + if item[1] == "assistant": + # If assistant message, mark specific tool_call index + to_remove.setdefault(item[0], set()).add(item[2]) + else: + # If tool message, mark whole message + to_remove[item[0]] = None + else: + # Only one type of message, no complete pairs - remove them all + for item in occurrences: + if item[1] == "assistant": + to_remove.setdefault(item[0], set()).add(item[2]) + else: + to_remove[item[0]] = None + elif len(occurrences) == 1: + # Incomplete pair (only tool call without tool result or vice versa) + item = occurrences[0] + if item[1] == "assistant": + to_remove.setdefault(item[0], set()).add(item[2]) + else: + to_remove[item[0]] = None + + # Step 4: Apply the removals and reconstruct the message list + sanitized_messages = [] for i, msg in enumerate(messages): - # Skip if message (tool type) is marked for removal - if i in removal_messages: + if i in to_remove and to_remove[i] is None: + # Skip entirely removed messages continue + # For assistant messages, remove marked tool_calls if msg.get("role") == "assistant" and "tool_calls" in msg: new_tool_calls = [] for j, tc in enumerate(msg["tool_calls"]): - if j not in removal_assistant_entries.get(i, set()): + if i not in to_remove or j not in to_remove[i]: new_tool_calls.append(tc) msg["tool_calls"] = new_tool_calls - # If after modification message has no content and no tool_calls, - # discard it - msg_content = msg.get("content") - if ((msg_content is None or not msg_content.strip()) and - not msg.get("tool_calls")): - continue - new_messages.append(msg) - return new_messages + # If after modification message has no content and no tool_calls, + # skip it + if not (msg.get("content", "").strip() or + not msg.get("tool_calls")): + continue + + sanitized_messages.append(msg) + + return sanitized_messages + +def cli_print_tool_call(tool_name="", args="", output="", prefix=" "): + """Print a tool call with pretty formatting""" + if not tool_name: + return + + print(f"\n{prefix}{color('Tool Call:', fg='cyan')}") + print(f"{prefix}{color('Name:', fg='cyan')} {tool_name}") + if args: + print(f"{prefix}{color('Args:', fg='cyan')} {args}") + if output: + print(f"{prefix}{color('Output:', fg='cyan')} {output}") + +def get_model_input_tokens(model): + """ + Get the number of input tokens for + max context window capacity for a given model. + """ + model_tokens = { + "gpt": 128000, + "o1": 200000, + "claude": 200000, + "qwen2.5": 32000, # https://ollama.com/library/qwen2.5, 128K input, 8K output # noqa: E501 # pylint: disable=C0301 + "llama3.1": 32000, # https://ollama.com/library/llama3.1, 128K input # noqa: E501 # pylint: disable=C0301 + "deepseek": 128000 # https://api-docs.deepseek.com/quick_start/pricing # noqa: E501 # pylint: disable=C0301 + } + for model_type, tokens in model_tokens.items(): + if model_type in model: + return tokens + return model_tokens["gpt"] + +def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals,too-many-statements,too-many-branches # noqa: E501 + interaction_input_tokens, + interaction_output_tokens, # noqa: E501, pylint: disable=R0913 + interaction_reasoning_tokens, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + model, + interaction_cost=0.0, + total_cost=None +) -> Text: # noqa: E501 + """ + Create a Text object displaying token usage information + with enhanced formatting. + """ + tokens_text = Text(justify="left") + + # Create a more compact, horizontal display + tokens_text.append(" ", style="bold") # Small padding + + # Current interaction tokens + tokens_text.append("Current: ", style="bold") + tokens_text.append(f"I:{interaction_input_tokens} ", style="green") + tokens_text.append(f"O:{interaction_output_tokens} ", style="red") + tokens_text.append(f"R:{interaction_reasoning_tokens} ", style="yellow") + + # Current cost + current_cost = float(interaction_cost) if interaction_cost is not None else 0.0 + tokens_text.append(f"(${current_cost:.4f}) ", style="bold") + + # Separator + tokens_text.append("| ", style="dim") + + # Total tokens + tokens_text.append("Total: ", style="bold") + tokens_text.append(f"I:{total_input_tokens} ", style="green") + tokens_text.append(f"O:{total_output_tokens} ", style="red") + tokens_text.append(f"R:{total_reasoning_tokens} ", style="yellow") + + # Total cost + total_cost_value = float(total_cost) if total_cost is not None else 0.0 + tokens_text.append(f"(${total_cost_value:.4f}) ", style="bold") + + # Separator + tokens_text.append("| ", style="dim") + + # Context usage + context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 + tokens_text.append("Context: ", style="bold") + tokens_text.append(f"{context_pct:.1f}% ", style="bold") + + # Context indicator + if context_pct < 50: + indicator = "🟩" + color_local = "green" + elif context_pct < 80: + indicator = "🟨" + color_local = "yellow" + else: + indicator = "🟥" + color_local = "red" + + tokens_text.append(f"{indicator}", style=color_local) + + return tokens_text + +def cli_print_agent_messages(agent_name, message, counter, model, debug, # pylint: disable=too-many-arguments,too-many-locals,unused-argument # noqa: E501 + interaction_input_tokens=None, + interaction_output_tokens=None, + interaction_reasoning_tokens=None, + total_input_tokens=None, + total_output_tokens=None, + total_reasoning_tokens=None, + interaction_cost=None, + total_cost=None): + """Print agent messages/thoughts with enhanced visual formatting.""" + # 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") + + # Create a more hacker-like header + text = Text() + + # Special handling for Reasoner Agent + if agent_name == "Reasoner Agent": + text.append(f"[{counter}] ", style="bold red") + text.append(f"Agent: {agent_name} ", style="bold yellow") + if message: + text.append(f">> {message} ", style="green") + text.append(f"[{timestamp}", style="dim") + if model: + text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})", + style="bold blue") + text.append("]", style="dim") + else: + text.append(f"[{counter}] ", style="bold cyan") + text.append(f"Agent: {agent_name} ", style="bold green") + if message: + text.append(f">> {message} ", style="yellow") + text.append(f"[{timestamp}", style="dim") + if model: + text.append(f" ({model})", style="bold magenta") + text.append("]", style="dim") + + # Add token information with enhanced formatting + tokens_text = None + if (interaction_input_tokens is not None and # pylint: disable=R0916 + 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, + model, + interaction_cost, + total_cost + ) + text.append(tokens_text) + + # Create a panel for better visual separation + panel = Panel( + text, + border_style="red" if agent_name == "Reasoner Agent" else "blue", + box=ROUNDED, + padding=(0, 1), + title=("[bold]Reasoning Analysis[/bold]" + if agent_name == "Reasoner Agent" + else "[bold]Agent Interaction[/bold]"), + title_align="left" + ) + console.print("\n") + console.print(panel) \ No newline at end of file