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

CLI output operational 

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches 2025-04-03 08:19:42 +02:00
parent 5b29fbe4db
commit b2233b76d1
7 changed files with 657 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
examples/cai/simple_one_tool_test.py
View File

@ -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

6
examples/cai/simple_one_tool_test_streamed.py
View File

@ -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

2
src/cai/cli.py
View File

@ -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

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

@ -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

105
src/cai/sdk/agents/models/openai_responses.py
View File

@ -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(

17
src/cai/sdk/agents/run.py
View File

@ -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

320
src/cai/util.py
View File

@ -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)