CLI output operational

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>

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

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

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

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,
+                    input_tokens=input_tokens,
-                    output_tokens=response.usage.completion_tokens,
+                    output_tokens=output_tokens,
-                    total_tokens=response.usage.total_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 = (
+            # Get final token counts using consistent method
-                ResponseUsage(
+            input_tokens = estimated_input_tokens
-                    input_tokens=usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0,
+            output_tokens = estimated_output_tokens
-                    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,
+            # Use API token counts if available and reasonable
-                    output_tokens_details=OutputTokensDetails(
+            if usage and hasattr(usage, 'prompt_tokens') and usage.prompt_tokens > 0:
-                        reasoning_tokens=usage.completion_tokens_details.reasoning_tokens
+                input_tokens = usage.prompt_tokens
-                        if usage and hasattr(usage, 'completion_tokens_details') 
+            if usage and hasattr(usage, 'completion_tokens') and usage.completion_tokens > 0:
-                        and usage.completion_tokens_details
+                output_tokens = usage.completion_tokens
-                        and hasattr(usage.completion_tokens_details, 'reasoning_tokens')
+                
-                        and usage.completion_tokens_details.reasoning_tokens
+            # # Debug information
-                        else 0
+            # print(f"\nDEBUG CONSISTENT TOKEN COUNTS - Streaming final tokens: input={input_tokens}, output={output_tokens}, total={input_tokens + output_tokens}")
-                    ),
+
-                    input_tokens_details={
+            # Create a proper usage object with our token counts
-                        "prompt_tokens": usage.prompt_tokens if usage and hasattr(usage, 'prompt_tokens') else 0,
+            final_response.usage = ResponseUsage(
-                        "cached_tokens": usage.prompt_tokens_details.cached_tokens
+                input_tokens=input_tokens,
-                        if usage and hasattr(usage, 'prompt_tokens_details')
+                output_tokens=output_tokens,
-                        and usage.prompt_tokens_details
+                total_tokens=input_tokens + output_tokens,
-                        and hasattr(usage.prompt_tokens_details, 'cached_tokens')
+                output_tokens_details=OutputTokensDetails(
-                        and usage.prompt_tokens_details.cached_tokens
+                    reasoning_tokens=usage.completion_tokens_details.reasoning_tokens
-                        else 0
+                    if usage and hasattr(usage, 'completion_tokens_details') 
-                    },
+                    and usage.completion_tokens_details
-                )
+                    and hasattr(usage.completion_tokens_details, 'reasoning_tokens')
-                if usage is not None
+                    and usage.completion_tokens_details.reasoning_tokens
-                else ResponseUsage(
+                    else 0
-                    input_tokens=0,
+                ),
-                    output_tokens=0,
+                input_tokens_details={
-                    total_tokens=0,
+                    "prompt_tokens": input_tokens,
-                    output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+                    "cached_tokens": usage.prompt_tokens_details.cached_tokens
-                    input_tokens_details={"prompt_tokens": 0, "cached_tokens": 0}
+                    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 = {
-                span_generation.span_data.usage = {
+                "input_tokens": input_tokens,
-                    "input_tokens": usage.prompt_tokens,
+                "output_tokens": output_tokens,
-                    "output_tokens": usage.completion_tokens,
+            }
-                }
-            else:
-                span_generation.span_data.usage = {
-                    "input_tokens": 0,
-                    "output_tokens": 0,
-                }
 
     @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

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(

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
+            model = agent.model
-
+        else:
-        return run_config.model_provider.get_model(agent.model)
+            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

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(
+                tc_id, []).append((i, "tool", None))
-                (i, "tool", None))
+
-    # Step 3: Mark invalid or extra occurrences for removal
+    # Step 3: Mark indices in the message list to remove.
-    removal_messages = set()  # Indices of messages (tool type) to remove
     # Maps message index (assistant) to set of indices (in tool_calls) to
-    # remove
+    # delete, or directly marks message indices (tool) to delete.
-    removal_assistant_entries = {}
+    to_remove = {}
     for tc_id, occurrences in tool_calls_occurrences.items():
-        # Only 2 occurrences allowed. Mark extras for removal.
+        if len(occurrences) > 2:
-        valid_occurrences = occurrences[:2]
+            # More than one assistant and tool message pair - trim down
-        extra_occurrences = occurrences[2:]
+            # by picking first pairing and removing the rest
-        for occ in extra_occurrences:
+            assistant_items = [
-            msg_idx, typ, j = occ
+                occ for occ in occurrences if occ[1] == "assistant"]
-            if typ == "assistant":
+            tool_items = [occ for occ in occurrences if occ[1] == "tool"]
-                removal_assistant_entries.setdefault(msg_idx, set()).add(j)
+
-            elif typ == "tool":
+            if assistant_items and tool_items:
-                removal_messages.add(msg_idx)
+                valid_assistant = assistant_items[0]
-        # If valid occurrences aren't exactly 2 (i.e., a lonely tool call),
+                valid_tool = tool_items[0]
-        # mark for removal
+                for item in occurrences:
-        if len(valid_occurrences) != 2:
+                    if item != valid_assistant and item != valid_tool:
-            for occ in valid_occurrences:
+                        if item[1] == "assistant":
-                msg_idx, typ, j = occ
+                            # If assistant message, mark specific tool_call index
-                if typ == "assistant":
+                            to_remove.setdefault(item[0], set()).add(item[2])
-                    removal_assistant_entries.setdefault(
+                        else:
-                        msg_idx, set()).add(j)
+                            # If tool message, mark whole message
-                elif typ == "tool":
+                            to_remove[item[0]] = None
-                    removal_messages.add(msg_idx)
+            else:
-        else:
+                # Only one type of message, no complete pairs - remove them all
-            # If exactly 2 occurrences, ensure both have content
+                for item in occurrences:
-            remove_pair = False
+                    if item[1] == "assistant":
-            for occ in valid_occurrences:
+                        to_remove.setdefault(item[0], set()).add(item[2])
-                msg_idx, typ, _ = occ
+                    else:
-                msg_content = messages[msg_idx].get("content")
+                        to_remove[item[0]] = None
-                if msg_content is None or not msg_content.strip():
+        elif len(occurrences) == 1:
-                    remove_pair = True
+            # Incomplete pair (only tool call without tool result or vice versa)
-                    break
+            item = occurrences[0]
-            if remove_pair:
+            if item[1] == "assistant":
-                for occ in valid_occurrences:
+                to_remove.setdefault(item[0], set()).add(item[2])
-                    msg_idx, typ, j = occ
+            else:
-                    if typ == "assistant":
+                to_remove[item[0]] = None
-                        removal_assistant_entries.setdefault(
+
-                            msg_idx, set()).add(j)
+    # Step 4: Apply the removals and reconstruct the message list
-                    elif typ == "tool":
+    sanitized_messages = []
-                        removal_messages.add(msg_idx)
-    # Step 4: Build new message list applying removals
-    new_messages = []
     for i, msg in enumerate(messages):
-        # Skip if message (tool type) is marked for removal
+        if i in to_remove and to_remove[i] is None:
-        if i in removal_messages:
+            # 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,
+            # If after modification message has no content and no tool_calls,
-        # discard it
+            # skip it
-        msg_content = msg.get("content")
+            if not (msg.get("content", "").strip() or
-        if ((msg_content is None or not msg_content.strip()) and
+                   not msg.get("tool_calls")):
-                not msg.get("tool_calls")):
+                continue
-            continue
+
-        new_messages.append(msg)
+        sanitized_messages.append(msg)
-    return new_messages
+
+    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)