mirror of https://github.com/aliasrobotics/cai.git
Merge branch 'stream_v2' into '0.4.0'
Stream v2 See merge request aliasrobotics/alias_research/cai!164
This commit is contained in:
Luis Javier Navarrete Lozano
commit
21bd5545f4
|
|
@ -8,7 +8,6 @@ license = {text = "Dual-licensed MIT and Proprietary"}
|
|||
authors = [{ name = "Alias Robotics", email = "research@aliasrobotics.com" }]
|
||||
dependencies = [
|
||||
# logging and visualization
|
||||
"flask>=2.0, <3",
|
||||
"folium>=0.15.0, <1",
|
||||
"matplotlib>=3.0, <4",
|
||||
"numpy>=1.21, <3",
|
||||
|
|
@ -26,7 +25,7 @@ dependencies = [
|
|||
"prompt_toolkit>=3.0.39",
|
||||
"dotenv>=0.9.9",
|
||||
"litellm>=1.63.7",
|
||||
"mako>=1.3.9",
|
||||
"mako>=1.3.8",
|
||||
"mcp; python_version >= '3.10'",
|
||||
"mkdocs>=1.6.0",
|
||||
"mkdocs-material>=9.6.0",
|
||||
|
|
|
|||
|
|
@ -355,7 +355,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
|
|||
title="[bold]Session Summary[/bold]",
|
||||
title_align="left"
|
||||
)
|
||||
console.print(time_panel)
|
||||
console.print(time_panel, end="")
|
||||
|
||||
print_session_summary(console, metrics, logging_path)
|
||||
|
||||
|
|
@ -442,6 +442,13 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
|
|||
}
|
||||
)
|
||||
|
||||
# Fix message list structure BEFORE sending to the model to prevent errors
|
||||
try:
|
||||
from cai.util import fix_message_list
|
||||
history_context = fix_message_list(history_context)
|
||||
except Exception as e:
|
||||
console.print(f"[yellow]Warning: Could not preprocess message history: {e}[/yellow]")
|
||||
|
||||
# Append the current user input as the last message in the list.
|
||||
conversation_input: list | str
|
||||
if history_context:
|
||||
|
|
@ -483,14 +490,85 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
|
|||
else:
|
||||
# Use non-streamed response
|
||||
response = asyncio.run(Runner.run(agent, conversation_input))
|
||||
|
||||
# Process the response items
|
||||
for item in response.new_items:
|
||||
# Handle tool call output items (tool results)
|
||||
if isinstance(item, ToolCallOutputItem):
|
||||
# First, ensure there's a corresponding assistant message with tool_calls
|
||||
# before adding the tool response to prevent the OpenAI error
|
||||
assistant_with_tool_call_exists = False
|
||||
tool_call_id = item.raw_item["call_id"]
|
||||
|
||||
for msg in message_history:
|
||||
if (msg.get("role") == "assistant" and
|
||||
msg.get("tool_calls") and
|
||||
any(tc.get("id") == tool_call_id for tc in msg.get("tool_calls", []))):
|
||||
assistant_with_tool_call_exists = True
|
||||
break
|
||||
|
||||
# If no matching assistant message exists, create one first
|
||||
if not assistant_with_tool_call_exists:
|
||||
tool_call_msg = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "unknown_function",
|
||||
"arguments": "{}"
|
||||
}
|
||||
}]
|
||||
}
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
# Now add the tool response
|
||||
tool_msg = {
|
||||
"role": "tool",
|
||||
"tool_call_id": item.raw_item["call_id"], # Use consistent format with streaming
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": item.output,
|
||||
}
|
||||
add_to_message_history(tool_msg)
|
||||
|
||||
# Make sure that assistant messages with tool calls are also added to message_history
|
||||
# This is especially important for non-streaming mode
|
||||
if hasattr(agent, 'model'):
|
||||
# Access the _Converter directly from the OpenAIChatCompletionsModel implementation
|
||||
from cai.sdk.agents.models.openai_chatcompletions import _Converter
|
||||
|
||||
# Check if recent_tool_calls exists and process them
|
||||
if hasattr(_Converter, 'recent_tool_calls'):
|
||||
for call_id, call_info in _Converter.recent_tool_calls.items():
|
||||
# Only process new tool calls that haven't been added to message history yet
|
||||
tool_call_found = False
|
||||
for msg in message_history:
|
||||
if (msg.get("role") == "assistant" and
|
||||
msg.get("tool_calls") and
|
||||
any(tc.get("id") == call_id for tc in msg.get("tool_calls", []))):
|
||||
tool_call_found = True
|
||||
break
|
||||
|
||||
if not tool_call_found:
|
||||
# Add the assistant message with the tool call
|
||||
tool_call_msg = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": call_info.get('name', ''),
|
||||
"arguments": call_info.get('arguments', '{}')
|
||||
}
|
||||
}]
|
||||
}
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
# Final validation to ensure message history follows OpenAI's requirements
|
||||
# Ensure every tool message has a preceding assistant message with matching tool_call_id
|
||||
from cai.util import fix_message_list
|
||||
message_history[:] = fix_message_list(message_history)
|
||||
turn_count += 1
|
||||
|
||||
# Stop measuring active time and start measuring idle time again
|
||||
|
|
|
|||
|
|
@ -238,7 +238,7 @@ class AgentCommand(Command):
|
|||
os.environ["CAI_AGENT_TYPE"] = selected_agent_key
|
||||
|
||||
console.print(
|
||||
f"[green]Switched to agent: {agent_name}[/green]")
|
||||
f"[green]Switched to agent: {agent_name}[/green]", end="")
|
||||
visualize_agent_graph(agent)
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -421,7 +421,7 @@ class ModelCommand(Command):
|
|||
change_message,
|
||||
border_style="green",
|
||||
title="Model Changed"
|
||||
)
|
||||
), end=""
|
||||
)
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -177,7 +177,7 @@ def display_banner(console: Console):
|
|||
[white] Bug bounty-ready AI[/white]
|
||||
"""
|
||||
|
||||
console.print(banner)
|
||||
console.print(banner, end="")
|
||||
|
||||
# # Create a table showcasing CAI framework capabilities
|
||||
# #
|
||||
|
|
@ -361,4 +361,4 @@ def display_quick_guide(console: Console):
|
|||
border_style="blue",
|
||||
padding=(1, 2),
|
||||
title_align="center"
|
||||
))
|
||||
), end="")
|
||||
|
|
|
|||
|
|
@ -96,7 +96,7 @@ def get_user_input(
|
|||
|
||||
# Get user input with all features
|
||||
return prompt(
|
||||
[('class:prompt', '\nCAI> ')],
|
||||
[('class:prompt', 'CAI> ')],
|
||||
completer=command_completer,
|
||||
style=create_prompt_style(),
|
||||
history=FileHistory(str(history_file)),
|
||||
|
|
|
|||
|
|
@ -336,6 +336,15 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# Log the user message
|
||||
if item.get("content"):
|
||||
self.logger.log_user_message(item.get("content"))
|
||||
|
||||
# IMPORTANT: Ensure the message list has valid tool call/result pairs
|
||||
# This needs to happen before the API call to prevent errors
|
||||
try:
|
||||
from cai.util import fix_message_list
|
||||
converted_messages = fix_message_list(converted_messages)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to fix message list: {e}")
|
||||
|
||||
# Get token count estimate before API call for consistent counting
|
||||
estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages)
|
||||
|
||||
|
|
@ -411,6 +420,10 @@ class OpenAIChatCompletionsModel(Model):
|
|||
|
||||
# Only display the agent message if we haven't already shown the tool output
|
||||
if should_display_message:
|
||||
# Ensure we're in non-streaming mode for proper markdown parsing
|
||||
previous_stream_setting = os.environ.get('CAI_STREAM', 'false')
|
||||
os.environ['CAI_STREAM'] = 'false' # Force non-streaming mode for markdown parsing
|
||||
|
||||
# Print the agent message for CLI display
|
||||
cli_print_agent_messages(
|
||||
agent_name=getattr(self, 'agent_name', 'Agent'),
|
||||
|
|
@ -433,7 +446,11 @@ class OpenAIChatCompletionsModel(Model):
|
|||
interaction_cost=None,
|
||||
total_cost=None,
|
||||
tool_output=None, # Don't pass tool output here, we're using direct display
|
||||
suppress_empty=True # Suppress empty panels
|
||||
)
|
||||
|
||||
# Restore previous streaming setting
|
||||
os.environ['CAI_STREAM'] = previous_stream_setting
|
||||
|
||||
# --- Add assistant tool call to message_history if present ---
|
||||
# If the response contains tool_calls, add them to message_history as assistant messages
|
||||
|
|
@ -457,6 +474,22 @@ class OpenAIChatCompletionsModel(Model):
|
|||
}
|
||||
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
# Save the tool call details for later matching with output
|
||||
# This is important for non-streaming mode to track tool calls properly
|
||||
if not hasattr(_Converter, 'recent_tool_calls'):
|
||||
_Converter.recent_tool_calls = {}
|
||||
|
||||
# Store the tool call by ID for later reference
|
||||
import time
|
||||
_Converter.recent_tool_calls[tool_call.id] = {
|
||||
'name': tool_call.function.name,
|
||||
'arguments': tool_call.function.arguments,
|
||||
'start_time': time.time(),
|
||||
'execution_info': {
|
||||
'start_time': time.time()
|
||||
}
|
||||
}
|
||||
|
||||
# Log the assistant tool call message
|
||||
tool_calls_list = []
|
||||
|
|
@ -583,7 +616,7 @@ class OpenAIChatCompletionsModel(Model):
|
|||
stop_idle_timer()
|
||||
start_active_timer()
|
||||
|
||||
# Check if streaming should be shown in rich panel
|
||||
# --- Check if streaming should be shown in rich panel ---
|
||||
should_show_rich_stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' and not self.disable_rich_streaming
|
||||
|
||||
# Create streaming context if needed
|
||||
|
|
@ -922,6 +955,55 @@ class OpenAIChatCompletionsModel(Model):
|
|||
if tool_call_msg not in streamed_tool_calls:
|
||||
streamed_tool_calls.append(tool_call_msg)
|
||||
add_to_message_history(tool_call_msg)
|
||||
|
||||
# NEW: Display tool call immediately when detected in streaming mode
|
||||
# But only if it has complete arguments and name
|
||||
if (state.function_calls[tc_index].name and
|
||||
state.function_calls[tc_index].arguments and
|
||||
state.function_calls[tc_index].call_id):
|
||||
# First, finish any existing streaming context if it exists
|
||||
if streaming_context:
|
||||
try:
|
||||
finish_agent_streaming(streaming_context, None)
|
||||
streaming_context = None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Create a message-like object for displaying the function call
|
||||
tool_msg = type('ToolCallStreamDisplay', (), {
|
||||
'content': None,
|
||||
'tool_calls': [
|
||||
type('ToolCallDetail', (), {
|
||||
'function': type('FunctionDetail', (), {
|
||||
'name': state.function_calls[tc_index].name,
|
||||
'arguments': state.function_calls[tc_index].arguments
|
||||
}),
|
||||
'id': state.function_calls[tc_index].call_id,
|
||||
'type': 'function'
|
||||
})
|
||||
]
|
||||
})
|
||||
|
||||
# Display the tool call during streaming
|
||||
cli_print_agent_messages(
|
||||
agent_name=getattr(self, 'agent_name', 'Agent'),
|
||||
message=tool_msg,
|
||||
counter=getattr(self, 'interaction_counter', 0),
|
||||
model=str(self.model),
|
||||
debug=False,
|
||||
interaction_input_tokens=estimated_input_tokens,
|
||||
interaction_output_tokens=estimated_output_tokens,
|
||||
interaction_reasoning_tokens=0, # Not available during streaming yet
|
||||
total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens,
|
||||
total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens,
|
||||
total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0),
|
||||
interaction_cost=None,
|
||||
total_cost=None,
|
||||
tool_output=None, # Will be shown once tool is executed
|
||||
suppress_empty=True # Prevent empty panels
|
||||
)
|
||||
# Set flag to suppress final output to avoid duplication
|
||||
self.suppress_final_output = True
|
||||
|
||||
except Exception as e:
|
||||
# Ensure streaming context is cleaned up in case of errors
|
||||
|
|
@ -983,8 +1065,15 @@ class OpenAIChatCompletionsModel(Model):
|
|||
)
|
||||
|
||||
# Display the tool call in CLI
|
||||
from cai.util import cli_print_agent_messages
|
||||
try:
|
||||
# First, finish any existing streaming context if it exists
|
||||
if streaming_context:
|
||||
try:
|
||||
finish_agent_streaming(streaming_context, None)
|
||||
streaming_context = None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Create a message-like object to display the function call
|
||||
tool_msg = type('ToolCallWrapper', (), {
|
||||
'content': None,
|
||||
|
|
@ -1015,8 +1104,12 @@ class OpenAIChatCompletionsModel(Model):
|
|||
total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0),
|
||||
interaction_cost=None,
|
||||
total_cost=None,
|
||||
tool_output=None # Will be shown once the tool is executed
|
||||
tool_output=None, # Will be shown once the tool is executed
|
||||
suppress_empty=True # Suppress empty panels during streaming
|
||||
)
|
||||
|
||||
# Set flag to suppress final output to avoid duplication
|
||||
self.suppress_final_output = True
|
||||
except Exception as e:
|
||||
logger.error(f"Error displaying tool call in CLI: {e}")
|
||||
|
||||
|
|
@ -1220,6 +1313,7 @@ class OpenAIChatCompletionsModel(Model):
|
|||
direct_stats["total_cost"] = float(total_cost)
|
||||
# Use the direct copy with guaranteed float costs
|
||||
finish_agent_streaming(streaming_context, direct_stats)
|
||||
streaming_context = None
|
||||
|
||||
# Removed extra newline after streaming completes to avoid blank lines
|
||||
pass
|
||||
|
|
@ -1243,8 +1337,10 @@ class OpenAIChatCompletionsModel(Model):
|
|||
for tool_call in tool_call_msg.get("tool_calls", []):
|
||||
tool_calls_list.append(tool_call)
|
||||
self.logger.log_assistant_message(None, tool_calls_list)
|
||||
# If there was only text output, add that as an assistant message
|
||||
if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text:
|
||||
|
||||
# If we've already shown the tool calls directly during streaming,
|
||||
# don't log the text message to avoid duplication
|
||||
elif (not self.suppress_final_output) and state.text_content_index_and_output and state.text_content_index_and_output[1].text:
|
||||
asst_msg = {
|
||||
"role": "assistant",
|
||||
"content": state.text_content_index_and_output[1].text
|
||||
|
|
@ -1253,6 +1349,9 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# Log the assistant message
|
||||
self.logger.log_assistant_message(state.text_content_index_and_output[1].text)
|
||||
|
||||
# Reset the suppress flag for future requests
|
||||
self.suppress_final_output = False
|
||||
|
||||
# Log the complete response
|
||||
self.logger.rec_training_data(
|
||||
{
|
||||
|
|
@ -1341,10 +1440,17 @@ class OpenAIChatCompletionsModel(Model):
|
|||
if tracing.include_data():
|
||||
span.span_data.input = converted_messages
|
||||
|
||||
# Ensure message list has correct structure regardless of error condition
|
||||
# IMPORTANT: Always sanitize the message list to prevent tool call errors
|
||||
# This is critical to fix common errors with tool/assistant sequences
|
||||
try:
|
||||
from cai.util import fix_message_list
|
||||
prev_length = len(converted_messages)
|
||||
converted_messages = fix_message_list(converted_messages)
|
||||
new_length = len(converted_messages)
|
||||
|
||||
# Log if the message list was changed significantly
|
||||
if new_length != prev_length:
|
||||
logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to fix message list: {e}")
|
||||
|
||||
|
|
@ -1531,15 +1637,49 @@ class OpenAIChatCompletionsModel(Model):
|
|||
|
||||
elif ("An assistant message with 'tool_calls'" in str(e) or
|
||||
"`tool_use` blocks must be followed by a user message with `tool_result`" in str(e) or # noqa: E501 # pylint: disable=C0301
|
||||
"An assistant message with 'tool_calls' must be followed by tool messages" in str(e)): # Añadir esta condición
|
||||
"An assistant message with 'tool_calls' must be followed by tool messages" in str(e) or
|
||||
"messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" in str(e)):
|
||||
print(f"Error: {str(e)}")
|
||||
|
||||
# Use the pretty message history printer instead of the simple loop
|
||||
try:
|
||||
from cai.util import print_message_history
|
||||
print("\nCurrent message sequence causing the error:")
|
||||
print_message_history(kwargs["messages"], title="Message Sequence Error")
|
||||
except ImportError:
|
||||
# Fall back to simple printing if the function isn't available
|
||||
print("\nCurrent message sequence causing the error:")
|
||||
for i, msg in enumerate(kwargs["messages"]):
|
||||
role = msg.get("role", "unknown")
|
||||
content_type = (
|
||||
"text" if isinstance(msg.get("content"), str) else
|
||||
"list" if isinstance(msg.get("content"), list) else
|
||||
"None" if msg.get("content") is None else
|
||||
type(msg.get("content")).__name__
|
||||
)
|
||||
tool_calls = "with tool_calls" if msg.get("tool_calls") else ""
|
||||
tool_call_id = f", tool_call_id: {msg.get('tool_call_id')}" if msg.get("tool_call_id") else ""
|
||||
|
||||
print(f" [{i}] {role}{tool_call_id} (content: {content_type}) {tool_calls}")
|
||||
|
||||
# NOTE: EDGE CASE: Report Agent CTRL C error
|
||||
#
|
||||
# This fix CTRL-C error when message list is incomplete
|
||||
# When a tool is not finished but the LLM generates a tool call
|
||||
try:
|
||||
from cai.util import fix_message_list
|
||||
kwargs["messages"] = fix_message_list(kwargs["messages"])
|
||||
print("Attempting to fix message sequence...")
|
||||
fixed_messages = fix_message_list(kwargs["messages"])
|
||||
|
||||
# Show the fixed messages if they're different
|
||||
if fixed_messages != kwargs["messages"]:
|
||||
try:
|
||||
from cai.util import print_message_history
|
||||
print_message_history(fixed_messages, title="Fixed Message Sequence")
|
||||
except ImportError:
|
||||
print("Messages fixed successfully.")
|
||||
|
||||
kwargs["messages"] = fixed_messages
|
||||
except Exception as fix_error:
|
||||
print(f"Failed to fix message sequence: {fix_error}")
|
||||
|
||||
|
|
@ -2042,7 +2182,11 @@ class _Converter:
|
|||
if current_assistant_msg is not None:
|
||||
# The API doesn't support empty arrays for tool_calls
|
||||
if not current_assistant_msg.get("tool_calls"):
|
||||
del current_assistant_msg["tool_calls"]
|
||||
# Ensure content is not None if tool_calls are absent and content is also None
|
||||
# Some models like Anthropic require some content, even if it's just a placeholder.
|
||||
if current_assistant_msg.get("content") is None:
|
||||
current_assistant_msg["content"] = "(No text content in this assistant message)" # Or just an empty string if preferred
|
||||
current_assistant_msg.pop("tool_calls", None) # Use pop with default to avoid KeyError
|
||||
result.append(current_assistant_msg)
|
||||
current_assistant_msg = None
|
||||
|
||||
|
|
@ -2079,16 +2223,28 @@ class _Converter:
|
|||
flush_assistant_message()
|
||||
tool_calls_param: list[ChatCompletionMessageToolCallParam] = []
|
||||
for tc in item["tool_calls"]:
|
||||
function_details = tc.get("function", {})
|
||||
arguments = function_details.get("arguments")
|
||||
# Ensure arguments is a valid JSON string, defaulting to "{}" if empty or None
|
||||
if arguments is None or (isinstance(arguments, str) and arguments.strip() == ""):
|
||||
arguments = "{}"
|
||||
elif isinstance(arguments, dict):
|
||||
# Ensure it's a string if it's a dict (should already be string per schema)
|
||||
arguments = json.dumps(arguments)
|
||||
|
||||
tool_calls_param.append(
|
||||
ChatCompletionMessageToolCallParam(
|
||||
id=tc.get("id", ""),
|
||||
type=tc.get("type", "function"),
|
||||
function=tc.get("function", {}),
|
||||
function={
|
||||
"name": function_details.get("name", "unknown_function"),
|
||||
"arguments": arguments, # Use sanitized arguments
|
||||
},
|
||||
)
|
||||
)
|
||||
msg_asst: ChatCompletionAssistantMessageParam = {
|
||||
"role": "assistant",
|
||||
"content": item.get("content"),
|
||||
"content": item.get("content"), # Content can be None here
|
||||
"tool_calls": tool_calls_param,
|
||||
}
|
||||
result.append(msg_asst)
|
||||
|
|
@ -2225,12 +2381,19 @@ class _Converter:
|
|||
}
|
||||
}
|
||||
|
||||
arguments = func_call.get("arguments") # func_call is a dict here
|
||||
# Ensure arguments is a valid JSON string, defaulting to "{}" if empty or None
|
||||
if arguments is None or (isinstance(arguments, str) and arguments.strip() == ""):
|
||||
arguments = "{}"
|
||||
elif isinstance(arguments, dict):
|
||||
arguments = json.dumps(arguments)
|
||||
|
||||
new_tool_call = ChatCompletionMessageToolCallParam(
|
||||
id=func_call["call_id"],
|
||||
type="function",
|
||||
function={
|
||||
"name": func_call["name"],
|
||||
"arguments": func_call["arguments"],
|
||||
"arguments": arguments, # Use sanitized arguments
|
||||
},
|
||||
)
|
||||
tool_calls.append(new_tool_call)
|
||||
|
|
@ -2244,22 +2407,22 @@ class _Converter:
|
|||
|
||||
# Update execution timing if we have the start time
|
||||
if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls:
|
||||
tool_call = cls.recent_tool_calls[call_id]
|
||||
if 'start_time' in tool_call:
|
||||
tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity
|
||||
if 'start_time' in tool_call_details:
|
||||
end_time = time.time()
|
||||
tool_execution_time = end_time - tool_call['start_time']
|
||||
tool_execution_time = end_time - tool_call_details['start_time']
|
||||
|
||||
# Update the execution info
|
||||
if 'execution_info' in tool_call:
|
||||
tool_call['execution_info']['end_time'] = end_time
|
||||
tool_call['execution_info']['tool_time'] = tool_execution_time
|
||||
if 'execution_info' in tool_call_details:
|
||||
tool_call_details['execution_info']['end_time'] = end_time
|
||||
tool_call_details['execution_info']['tool_time'] = tool_execution_time
|
||||
|
||||
# If this is the first tool being executed, record the total time from conversation start
|
||||
if not hasattr(cls, 'conversation_start_time'):
|
||||
cls.conversation_start_time = tool_call['start_time']
|
||||
cls.conversation_start_time = tool_call_details['start_time']
|
||||
|
||||
total_time = end_time - getattr(cls, 'conversation_start_time', tool_call['start_time'])
|
||||
tool_call['execution_info']['total_time'] = total_time
|
||||
total_time = end_time - getattr(cls, 'conversation_start_time', tool_call_details['start_time'])
|
||||
tool_call_details['execution_info']['total_time'] = total_time
|
||||
|
||||
# Store the output so it can be accessed later
|
||||
if not hasattr(cls, 'tool_outputs'):
|
||||
|
|
@ -2270,67 +2433,73 @@ class _Converter:
|
|||
# Display the tool output immediately with the matched tool call
|
||||
from cai.util import cli_print_tool_output
|
||||
|
||||
# Check if we're in streaming mode - don't show tool output panel in streaming mode
|
||||
# Look up the original tool call to get the name and arguments
|
||||
tool_name = "Unknown Tool"
|
||||
tool_args = {}
|
||||
execution_info = {}
|
||||
|
||||
if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls:
|
||||
tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity
|
||||
tool_name = tool_call_details.get('name', 'Unknown Tool')
|
||||
tool_args = tool_call_details.get('arguments', {})
|
||||
execution_info = tool_call_details.get('execution_info', {})
|
||||
|
||||
# Get token counts from the OpenAIChatCompletionsModel if available
|
||||
model_instance = None
|
||||
for frame in inspect.stack():
|
||||
if 'self' in frame.frame.f_locals:
|
||||
self_obj = frame.frame.f_locals['self']
|
||||
if isinstance(self_obj, OpenAIChatCompletionsModel):
|
||||
model_instance = self_obj
|
||||
break
|
||||
|
||||
# Always create a token_info dictionary, even if some values are zero
|
||||
token_info = {
|
||||
'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0),
|
||||
'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0),
|
||||
'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0),
|
||||
'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0),
|
||||
'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0),
|
||||
'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0),
|
||||
'model': str(getattr(model_instance, 'model', '')),
|
||||
}
|
||||
|
||||
# Calculate costs using standard cost model
|
||||
if model_instance and hasattr(model_instance, 'model'):
|
||||
from cai.util import calculate_model_cost
|
||||
model_name_str = str(model_instance.model) # Ensure model name is string
|
||||
token_info['interaction_cost'] = calculate_model_cost(
|
||||
model_name_str,
|
||||
token_info['interaction_input_tokens'],
|
||||
token_info['interaction_output_tokens']
|
||||
)
|
||||
token_info['total_cost'] = calculate_model_cost(
|
||||
model_name_str,
|
||||
token_info['total_input_tokens'],
|
||||
token_info['total_output_tokens']
|
||||
)
|
||||
|
||||
# Check if we're in streaming mode
|
||||
is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true'
|
||||
if is_streaming_enabled:
|
||||
# Don't display tool output in streaming mode - it will be handled elsewhere
|
||||
pass # Just skip the display, but preserve the tool output
|
||||
else:
|
||||
# For non-streaming mode, maintain the original behavior
|
||||
# Look up the original tool call to get the name and arguments
|
||||
if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls:
|
||||
tool_call = cls.recent_tool_calls[call_id]
|
||||
tool_name = tool_call.get('name', 'Unknown Tool')
|
||||
tool_args = tool_call.get('arguments', {})
|
||||
execution_info = tool_call.get('execution_info', {})
|
||||
|
||||
# Get token counts from the OpenAIChatCompletionsModel if available
|
||||
model_instance = None
|
||||
for frame in inspect.stack():
|
||||
if 'self' in frame.frame.f_locals:
|
||||
self_obj = frame.frame.f_locals['self']
|
||||
if isinstance(self_obj, OpenAIChatCompletionsModel):
|
||||
model_instance = self_obj
|
||||
break
|
||||
|
||||
# Always create a token_info dictionary, even if some values are zero
|
||||
token_info = {
|
||||
'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0),
|
||||
'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0),
|
||||
'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0),
|
||||
'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0),
|
||||
'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0),
|
||||
'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0),
|
||||
'model': str(getattr(model_instance, 'model', '')),
|
||||
}
|
||||
|
||||
# Calculate costs using standard cost model
|
||||
if model_instance and hasattr(model_instance, 'model'):
|
||||
from cai.util import calculate_model_cost
|
||||
model_name = str(model_instance.model)
|
||||
token_info['interaction_cost'] = calculate_model_cost(
|
||||
model_name,
|
||||
token_info['interaction_input_tokens'],
|
||||
token_info['interaction_output_tokens']
|
||||
)
|
||||
token_info['total_cost'] = calculate_model_cost(
|
||||
model_name,
|
||||
token_info['total_input_tokens'],
|
||||
token_info['total_output_tokens']
|
||||
)
|
||||
|
||||
# Use the cli_print_tool_output function with actual token values
|
||||
cli_print_tool_output(
|
||||
tool_name=tool_name,
|
||||
args=tool_args,
|
||||
output=output_content,
|
||||
call_id=call_id, # Keep call_id for non-streaming mode
|
||||
execution_info=execution_info,
|
||||
token_info=token_info
|
||||
)
|
||||
|
||||
# Always display tool output regardless of streaming mode
|
||||
cli_print_tool_output(
|
||||
tool_name=tool_name,
|
||||
args=tool_args,
|
||||
output=output_content,
|
||||
call_id=call_id,
|
||||
execution_info=execution_info,
|
||||
token_info=token_info
|
||||
)
|
||||
|
||||
# Continue with normal processing
|
||||
flush_assistant_message()
|
||||
|
||||
# REMOVED THE BLOCK THAT CREATED A SYNTHETIC ASSISTANT MESSAGE HERE
|
||||
# The responsibility for ensuring a preceding assistant message
|
||||
# is now fully deferred to fix_message_list, called later.
|
||||
|
||||
# Now add the tool message
|
||||
msg: ChatCompletionToolMessageParam = {
|
||||
"role": "tool",
|
||||
"tool_call_id": func_output["call_id"],
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -37,62 +37,101 @@ def execute_code(code: str = "", language: str = "python",
|
|||
"python": "py",
|
||||
"php": "php",
|
||||
"bash": "sh",
|
||||
"shell": "sh", # Add shell as alias for bash
|
||||
"ruby": "rb",
|
||||
"perl": "pl",
|
||||
"golang": "go",
|
||||
"go": "go", # Add go as alias for golang
|
||||
"javascript": "js",
|
||||
"js": "js", # Add js as alias for javascript
|
||||
"typescript": "ts",
|
||||
"ts": "ts", # Add ts as alias for typescript
|
||||
"rust": "rs",
|
||||
"csharp": "cs",
|
||||
"cs": "cs", # Add cs as alias for csharp
|
||||
"java": "java",
|
||||
"kotlin": "kt"
|
||||
"kotlin": "kt",
|
||||
"c": "c", # Add C language
|
||||
"cpp": "cpp", # Add C++ language
|
||||
"c++": "cpp" # Add C++ language alias
|
||||
}
|
||||
ext = extensions.get(language.lower(), "txt")
|
||||
# Normalize language to lowercase
|
||||
language = language.lower()
|
||||
ext = extensions.get(language, "txt")
|
||||
full_filename = f"{filename}.{ext}"
|
||||
|
||||
# Create code file with content
|
||||
create_cmd = f"cat << 'EOF' > {full_filename}\n{code}\nEOF"
|
||||
result = run_command(create_cmd, ctf=ctf)
|
||||
result = run_command(create_cmd, ctf=ctf, stream=True, tool_name="execute_code")
|
||||
if "error" in result.lower():
|
||||
return f"Failed to create code file: {result}"
|
||||
if language.lower() == "python":
|
||||
|
||||
# Prepare execution command based on language
|
||||
if language in ["python", "py"]:
|
||||
exec_cmd = f"python3 {full_filename}"
|
||||
elif language.lower() == "php":
|
||||
elif language in ["php"]:
|
||||
exec_cmd = f"php {full_filename}"
|
||||
elif language.lower() in ["bash", "sh"]:
|
||||
elif language in ["bash", "sh", "shell"]:
|
||||
exec_cmd = f"bash {full_filename}"
|
||||
elif language.lower() == "ruby":
|
||||
elif language in ["ruby", "rb"]:
|
||||
exec_cmd = f"ruby {full_filename}"
|
||||
elif language.lower() == "perl":
|
||||
elif language in ["perl", "pl"]:
|
||||
exec_cmd = f"perl {full_filename}"
|
||||
elif language.lower() == "golang" or language.lower() == "go":
|
||||
elif language in ["golang", "go"]:
|
||||
temp_dir = f"/tmp/go_exec_{filename}"
|
||||
run_command(f"mkdir -p {temp_dir}", ctf=ctf)
|
||||
run_command(f"cp {full_filename} {temp_dir}/main.go", ctf=ctf)
|
||||
run_command(f"cd {temp_dir} && go mod init temp", ctf=ctf)
|
||||
exec_cmd = f"cd {temp_dir} && go run main.go"
|
||||
elif language.lower() == "javascript":
|
||||
elif language in ["javascript", "js"]:
|
||||
exec_cmd = f"node {full_filename}"
|
||||
elif language.lower() == "typescript":
|
||||
elif language in ["typescript", "ts"]:
|
||||
exec_cmd = f"ts-node {full_filename}"
|
||||
elif language.lower() == "rust":
|
||||
elif language in ["rust", "rs"]:
|
||||
# For Rust, we need to compile first
|
||||
run_command(f"rustc {full_filename} -o {filename}", ctf=ctf)
|
||||
exec_cmd = f"./{filename}"
|
||||
elif language.lower() == "csharp":
|
||||
elif language in ["csharp", "cs"]:
|
||||
# For C#, compile with dotnet
|
||||
run_command(f"dotnet build {full_filename}", ctf=ctf)
|
||||
exec_cmd = f"dotnet run {full_filename}"
|
||||
elif language.lower() == "java":
|
||||
elif language in ["java"]:
|
||||
# For Java, compile first
|
||||
run_command(f"javac {full_filename}", ctf=ctf)
|
||||
exec_cmd = f"java {filename}"
|
||||
elif language.lower() == "kotlin":
|
||||
elif language in ["kotlin", "kt"]:
|
||||
# For Kotlin, compile first
|
||||
run_command(f"kotlinc {full_filename} -include-runtime -d {filename}.jar", ctf=ctf)
|
||||
exec_cmd = f"java -jar {filename}.jar"
|
||||
elif language in ["c"]:
|
||||
# For C, compile with gcc
|
||||
run_command(f"gcc {full_filename} -o {filename}", ctf=ctf)
|
||||
exec_cmd = f"./{filename}"
|
||||
elif language in ["cpp", "c++"]:
|
||||
# For C++, compile with g++
|
||||
run_command(f"g++ {full_filename} -o {filename}", ctf=ctf)
|
||||
exec_cmd = f"./{filename}"
|
||||
else:
|
||||
return f"Unsupported language: {language}"
|
||||
|
||||
output = run_command(exec_cmd, ctf=ctf, timeout=timeout)
|
||||
# Execute the code with syntax-highlighted output
|
||||
# Create a custom tool args dictionary to send language and code info to the tool output function
|
||||
tool_args = {
|
||||
"command": "execute",
|
||||
"language": language,
|
||||
"filename": filename,
|
||||
"code": code, # Include the code for syntax highlighting
|
||||
"timeout": timeout
|
||||
}
|
||||
|
||||
# Run the command with streaming to get syntax highlighting
|
||||
output = run_command(
|
||||
exec_cmd,
|
||||
ctf=ctf,
|
||||
timeout=timeout,
|
||||
stream=True,
|
||||
tool_name="execute_code",
|
||||
args=tool_args
|
||||
)
|
||||
|
||||
return output
|
||||
|
|
|
|||
1540
src/cai/util.py
1540
src/cai/util.py
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue