mirror of https://github.com/aliasrobotics/cai.git
tool output streaming works
This commit is contained in:
lidia9
parent
734ef479f4
commit
7d6ddd3a4f
|
|
@ -1609,43 +1609,50 @@ class _Converter:
|
|||
# Display the tool output immediately with the matched tool call
|
||||
from cai.util import cli_print_tool_output
|
||||
|
||||
# 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
|
||||
|
||||
token_info = {}
|
||||
if model_instance:
|
||||
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', '')),
|
||||
}
|
||||
|
||||
# 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,
|
||||
execution_info=execution_info,
|
||||
token_info=token_info
|
||||
)
|
||||
# Check if we're in streaming mode - don't show tool output panel 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
|
||||
|
||||
token_info = {}
|
||||
if model_instance:
|
||||
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', '')),
|
||||
}
|
||||
|
||||
# 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
|
||||
)
|
||||
|
||||
# Continue with normal processing
|
||||
flush_assistant_message()
|
||||
|
|
|
|||
|
|
@ -245,6 +245,14 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None):
|
|||
output = result.stdout if result.stdout else result.stderr
|
||||
if stdout:
|
||||
print("\033[32m" + output + "\033[0m")
|
||||
|
||||
# Skip passing output to cli_print_tool_output when CAI_STREAM=true
|
||||
# This prevents duplicate output in streaming mode
|
||||
is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true'
|
||||
if not is_streaming_enabled:
|
||||
# Optional: Add cli_print_tool_output call here if needed for non-streaming
|
||||
pass
|
||||
|
||||
return output
|
||||
except subprocess.TimeoutExpired as e:
|
||||
error_output = e.stdout.decode() if e.stdout else str(e)
|
||||
|
|
@ -288,10 +296,16 @@ def _run_local_streamed(command, call_id, timeout=100):
|
|||
# Create panel content for Rich display
|
||||
if rich_available:
|
||||
tool_name = "generic_linux_command"
|
||||
# Parse command into command and args
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
header = Text()
|
||||
header.append(tool_name, style="#00BCD4")
|
||||
header.append("(", style="yellow")
|
||||
header.append(f"command='{command}'", style="yellow")
|
||||
# Format to match: generic_linux_command({"command":"ls","args":"-la","ctf":{},"async_mode":false,"session_id":""})
|
||||
header.append(f'{{"command":"{cmd}","args":"{args}","ctf":{{}},"async_mode":false,"session_id":""}}', style="yellow")
|
||||
header.append(")", style="yellow")
|
||||
|
||||
content = Text()
|
||||
|
|
@ -299,9 +313,9 @@ def _run_local_streamed(command, call_id, timeout=100):
|
|||
|
||||
panel = Panel(
|
||||
Text.assemble(header, "\n\n", content),
|
||||
title="[bold blue]Tool Execution[/bold blue]",
|
||||
title="[bold green]Tool Execution[/bold green]",
|
||||
subtitle="[bold green]Live Output[/bold green]",
|
||||
border_style="blue",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
)
|
||||
|
|
@ -320,9 +334,9 @@ def _run_local_streamed(command, call_id, timeout=100):
|
|||
content.append(line, style="bright_white")
|
||||
panel = Panel(
|
||||
Text.assemble(header, "\n\n", content),
|
||||
title="[bold blue]Tool Execution[/bold blue]",
|
||||
title="[bold green]Tool Execution[/bold green]",
|
||||
subtitle="[bold green]Live Output[/bold green]",
|
||||
border_style="blue",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
)
|
||||
|
|
@ -340,9 +354,9 @@ def _run_local_streamed(command, call_id, timeout=100):
|
|||
output_buffer.append("\nERROR OUTPUT:\n" + stderr_data)
|
||||
panel = Panel(
|
||||
Text.assemble(header, "\n\n", content),
|
||||
title="[bold blue]Tool Execution[/bold blue]",
|
||||
title="[bold green]Tool Execution[/bold green]",
|
||||
subtitle="[bold green]Live Output[/bold green]",
|
||||
border_style="blue",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
)
|
||||
|
|
@ -353,9 +367,9 @@ def _run_local_streamed(command, call_id, timeout=100):
|
|||
content.append(f"\nCommand {completion_status}", style="green")
|
||||
panel = Panel(
|
||||
Text.assemble(header, "\n\n", content),
|
||||
title="[bold blue]Tool Execution[/bold blue]",
|
||||
title="[bold green]Tool Execution[/bold green]",
|
||||
subtitle=f"[bold green]{completion_status}[/bold green]",
|
||||
border_style="blue",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
)
|
||||
|
|
@ -365,7 +379,11 @@ def _run_local_streamed(command, call_id, timeout=100):
|
|||
time.sleep(0.5)
|
||||
else:
|
||||
# Fallback to simpler streaming with cli_print_tool_output
|
||||
tool_args = {"command": command}
|
||||
# Parse command into command and args (same as rich mode)
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
tool_args = {"command": cmd, "args": args, "ctf": {}, "async_mode": False, "session_id": ""}
|
||||
|
||||
# Initial notification - just once
|
||||
cli_print_tool_output("generic_linux_command", tool_args, "Command started...", call_id=call_id)
|
||||
|
|
|
|||
|
|
@ -90,9 +90,20 @@ def generic_linux_command(command: str = "",
|
|||
stream = os.getenv('CAI_STREAM', 'false').lower() == 'true'
|
||||
|
||||
# Generate a call_id for streaming if needed
|
||||
# Only use call_id for streaming mode to prevent duplicate panels
|
||||
# When CAI_STREAM=true, we want the Tool Execution panel but not the Tool Output panel
|
||||
call_id = str(uuid.uuid4())[:8] if stream else None
|
||||
|
||||
# In streaming mode, prevent displaying both panels by forcing a call_id
|
||||
# This tricks cli_print_tool_output into thinking it's being called for a streaming update
|
||||
if stream and not call_id:
|
||||
call_id = str(uuid.uuid4())[:8]
|
||||
|
||||
# Run the command with the appropriate parameters
|
||||
return run_command(full_command, ctf=ctf,
|
||||
result = run_command(full_command, ctf=ctf,
|
||||
async_mode=async_mode, session_id=session_id,
|
||||
timeout=timeout, stream=stream, call_id=call_id)
|
||||
|
||||
# For better output formatting, if we're in streaming mode, we can modify the output
|
||||
# to include any additional information needed while still preventing the duplicate panel
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -877,8 +877,11 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli
|
|||
console.print(panel)
|
||||
|
||||
# If there are tool panels, print them after the main message panel
|
||||
for tool_panel in tool_panels:
|
||||
console.print(tool_panel)
|
||||
# But only in non-streaming mode to avoid duplicates
|
||||
is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true'
|
||||
if tool_panels and not is_streaming_enabled:
|
||||
for tool_panel in tool_panels:
|
||||
console.print(tool_panel)
|
||||
|
||||
def create_agent_streaming_context(agent_name, counter, model):
|
||||
"""Create a streaming context object that maintains state for streaming agent output."""
|
||||
|
|
@ -1072,10 +1075,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
if not output and not call_id:
|
||||
return
|
||||
|
||||
# CRITICAL CHECK: When in streaming mode (CAI_STREAM=true), ONLY show output panels
|
||||
# for streaming updates (those with call_id). This prevents duplicate output panels.
|
||||
is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true'
|
||||
if is_streaming_enabled and not call_id:
|
||||
# Skip all non-streaming tool output in streaming mode
|
||||
return
|
||||
|
||||
# Track seen call IDs to prevent duplicate panels
|
||||
if not hasattr(cli_print_tool_output, '_seen_calls'):
|
||||
cli_print_tool_output._seen_calls = {}
|
||||
|
||||
|
||||
# For streaming updates, only show updates for the same call_id
|
||||
# but allow the first appearance of each call_id
|
||||
if call_id:
|
||||
|
|
@ -1127,11 +1137,12 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
# Create the panel for display
|
||||
panel = Panel(
|
||||
Text.assemble(header, "\n\n", content),
|
||||
title=f"[bold blue]Tool Output (ID: {call_id})[/bold blue]",
|
||||
subtitle="[bold green]Live Update[/bold green]",
|
||||
title=f"[bold blue]Tool Output[/bold blue]", #(ID: {call_id})[/bold green]",
|
||||
#subtitle="[bold green]Live Update[/bold green]",
|
||||
border_style="blue",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
box=ROUNDED,
|
||||
title_align="left"
|
||||
)
|
||||
|
||||
# Display using Rich
|
||||
|
|
@ -1181,9 +1192,9 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
# Create the panel with all sections
|
||||
panel = Panel(
|
||||
Text.assemble(*sections),
|
||||
title=f"[bold blue]Tool Output: {tool_call}[/bold blue]",
|
||||
title=f"[bold green]Tool Output: {tool_call}[/bold green]",
|
||||
subtitle=subtitle,
|
||||
border_style="blue",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
box=ROUNDED
|
||||
)
|
||||
|
|
@ -1215,7 +1226,6 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
header = header[:term_width-4]
|
||||
|
||||
# Clear the screen for the tool output (alternative approach)
|
||||
# This works better with streamed content as it doesn't scroll the terminal
|
||||
print(f"\r{header}")
|
||||
|
||||
# Print the content
|
||||
|
|
|
|||
Loading…
Reference in New Issue