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

merge with reset-v0.4.0

This commit is contained in:
Mery-Sanz 2025-05-07 10:36:30 +02:00
parent 17282f72ad 7b0c05ee82
commit 7fe22d7312
17 changed files with 2240 additions and 491 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
examples/cai/.gitkeep
View File

38
src/cai/cli.py
View File

@ -175,7 +175,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
ACTIVE_TIME = 0
idle_time = 0
console = Console()
last_model = os.getenv('CAI_MODEL', 'qwen2.5:14b')
last_agent_type = os.getenv('CAI_AGENT_TYPE', 'one_tool_agent')
# Initialize command completer and key bindings
command_completer = FuzzyCommandCompleter()
current_text = ['']
@ -208,6 +209,36 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
while turn_count < max_turns:
try:
idle_start_time = time.time()
# Check if model has changed and update if needed
current_model = os.getenv('CAI_MODEL', 'qwen2.5:14b')
if current_model != last_model and hasattr(agent, 'model'):
# Update the model in the agent
if hasattr(agent.model, 'model'):
agent.model.model = current_model
last_model = current_model
# Check if agent type has changed and recreate agent if needed
current_agent_type = os.getenv('CAI_AGENT_TYPE', 'one_tool_agent')
if current_agent_type != last_agent_type:
try:
# Import is already at the top level
agent = get_agent_by_name(current_agent_type)
last_agent_type = current_agent_type
# Configure the new agent's model flags
if hasattr(agent, 'model'):
if hasattr(agent.model, 'disable_rich_streaming'):
agent.model.disable_rich_streaming = True
if hasattr(agent.model, 'suppress_final_output'):
agent.model.suppress_final_output = True
# Apply current model to the new agent
if hasattr(agent.model, 'model'):
agent.model.model = current_model
except Exception as e:
console.print(f"[red]Error switching agent: {str(e)}[/red]")
# Get user input with command completion and history
user_input = get_user_input(
command_completer,
@ -280,8 +311,9 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
if commands_handle_command(command, args):
continue # Command was handled, continue to next iteration
# If command wasn't recognized, show error
console.print(f"[red]Unknown command: {command}[/red]")
# If command wasn't recognized, show error (skip for /shell or /s)
if command not in ("/shell", "/s"):
console.print(f"[red]Unknown command: {command}[/red]")
continue
# Process the conversation with the agent

6
src/cai/repl/commands/__init__.py
View File

@ -25,7 +25,6 @@ from cai.repl.commands.base import (
# Import all command modules
# These imports will register the commands with the registry
from cai.repl.commands import ( # pylint: disable=import-error,unused-import,line-too-long,redefined-builtin # noqa: E501,F401
memory,
help,
graph,
exit,
@ -34,10 +33,11 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li
platform,
kill,
model,
turns,
agent,
history,
config
config,
flush,
workspace,
)
# Define helper functions

202
src/cai/repl/commands/agent.py
View File

@ -155,45 +155,37 @@ class AgentCommand(Command):
table = Table(title="Available Agents")
table.add_column("#", style="dim")
table.add_column("Name", style="cyan")
table.add_column("Module", style="magenta")
table.add_column("Key", style="magenta")
table.add_column("Module", style="green")
table.add_column("Description", style="green")
table.add_column("Pattern", style="blue")
table.add_column("Model", style="yellow")
# Scan all agents from the agents folder
# Retrieve all registered agents
agents_to_display = get_available_agents()
# Display all agents
for i, (name, agent) in enumerate(agents_to_display.items(), 1):
description = agent.description
if not description and hasattr(agent, 'instructions'):
if callable(agent.instructions):
description = agent.instructions(context_variables={})
else:
description = agent.instructions
# Clean up description - remove newlines and strip spaces
for idx, (agent_key, agent) in enumerate(agents_to_display.items(), start=1):
# Human-friendly name (falls back to the dict key)
display_name = getattr(agent, "name", agent_key)
# Use provided description, otherwise derive from instructions
description = getattr(agent, "description", "") or ""
if not description and hasattr(agent, "instructions"):
instr = agent.instructions
description = instr(context_variables={}) if callable(instr) else instr
if isinstance(description, str):
description = " ".join(description.split())
if len(description) > 50:
description = description[:47] + "..."
# Get the module name for the agent
module_name = get_agent_module(name)
# Module where this agent lives
module_name = get_agent_module(agent_key)
# Get the pattern if it exists
pattern = getattr(agent, 'pattern', '')
if pattern:
pattern = pattern.capitalize()
# Handle model display based on agent type
model_display = self._get_model_display(name, agent)
# Add a row with all collected info
table.add_row(
str(i),
name,
str(idx),
display_name,
agent_key,
module_name,
description,
pattern,
model_display
description
)
console.print(table)
@ -210,76 +202,45 @@ class AgentCommand(Command):
"""
if not args:
console.print("[red]Error: No agent specified[/red]")
console.print("Usage: /agent select <name|number>")
console.print("Usage: /agent select <agent_key|number>")
return False
agent_id = args[0]
# Get the list of available agents
agents_to_display = get_available_agents()
agent_list = list(agents_to_display.items())
# Check if agent_id is a number
if agent_id.isdigit():
index = int(agent_id)
if 1 <= index <= len(agents_to_display):
agent_name = list(agents_to_display.keys())[index - 1]
if 1 <= index <= len(agent_list):
# Get the agent tuple from the list
selected_agent_key, selected_agent = agent_list[index - 1]
agent_name = getattr(selected_agent, "name", selected_agent_key)
agent = selected_agent
else:
console.print(
f"[red]Error: Invalid agent number: {agent_id}[/red]")
console.print(f"[red]Error: Invalid agent number: {agent_id}[/red]")
return False
else:
# Treat as agent name
agent_name = agent_id
if agent_name not in agents_to_display:
console.print(f"[red]Error: Unknown agent: {agent_name}[/red]")
# Treat as agent key
selected_agent_key = None
for key, agent_obj in agents_to_display.items():
if key == agent_id:
agent = agent_obj
selected_agent_key = key
agent_name = getattr(agent_obj, "name", key)
break
else:
console.print(f"[red]Error: Unknown agent key: {agent_id}[/red]")
return False
# Get the agent
agent = agents_to_display[agent_name]
# Set the agent as the current agent in the REPL
# We need to avoid circular imports, so we'll use a different approach
# to access the client and current_agent variables
# Import the module dynamically to avoid circular imports
if 'cai.repl.repl' in sys.modules:
repl_module = sys.modules['cai.repl.repl']
# Check if client is initialized
if hasattr(repl_module, 'client') and repl_module.client:
# Update the active_agent in the client
repl_module.client.active_agent = agent
# Update the global current_agent variable if it exists
if hasattr(repl_module, 'current_agent'):
repl_module.current_agent = agent
# Update the global agent variable if it exists
if hasattr(repl_module, 'agent'):
repl_module.agent = agent
# Also update the agent variable in the run_demo_loop
# function's frame if possible
try:
for frame_info in inspect.stack():
frame = frame_info.frame
if ('run_demo_loop' in frame.f_code.co_name and
'agent' in frame.f_locals):
frame.f_locals['agent'] = agent
break
except Exception: # pylint: disable=broad-except # nosec
# If this fails, we still have the global current_agent as
# a fallback
pass
console.print(
f"[green]Switched to agent: {agent_name}[/green]")
visualize_agent_graph(agent)
return True
console.print("[red]Error: CAI client not initialized[/red]")
return False
console.print("[red]Error: REPL module not initialized[/red]")
return False
# Set the agent key in environment variable (not the agent name)
os.environ["CAI_AGENT_TYPE"] = selected_agent_key
console.print(
f"[green]Switched to agent: {agent_name}[/green]")
visualize_agent_graph(agent)
return True
def handle_info(self, args: Optional[List[str]] = None) -> bool:
"""Handle /agent info command.
@ -292,57 +253,74 @@ class AgentCommand(Command):
"""
if not args:
console.print("[red]Error: No agent specified[/red]")
console.print("Usage: /agent info <name|number>")
console.print("Usage: /agent info <agent_key|number>")
return False
agent_id = args[0]
# Get the list of available agents
# Get available agents
agents_to_display = get_available_agents()
# Check if agent_id is a number
# Resolve agent_id to an agent key (by index or name)
if agent_id.isdigit():
index = int(agent_id)
if 1 <= index <= len(agents_to_display):
agent_name = list(agents_to_display.keys())[index - 1]
else:
console.print(
f"[red]Error: Invalid agent number: {agent_id}[/red]")
idx = int(agent_id)
if not (1 <= idx <= len(agents_to_display)):
console.print(f"[red]Error: Invalid agent number: {agent_id}[/red]")
return False
agent_key = list(agents_to_display.keys())[idx - 1]
else:
# Treat as agent name
agent_name = agent_id
if agent_name not in agents_to_display:
console.print(f"[red]Error: Unknown agent: {agent_name}[/red]")
agent_key = None
for key, ag in agents_to_display.items():
if key == agent_id or getattr(ag, "name", "").lower() == agent_id.lower():
agent_key = key
break
if agent_key is None:
console.print(f"[red]Error: Unknown agent key: {agent_id}[/red]")
return False
# Get the agent
agent = agents_to_display[agent_name]
agent = agents_to_display[agent_key]
# Display agent information
# Display agent information
instructions = agent.instructions
if callable(instructions):
instructions = instructions()
# Prepare agent properties
name = agent.name or agent_key
description = getattr(agent, "description", None) or "N/A"
clean_description = " ".join(line.strip() for line in description.splitlines())
functions = getattr(agent, "functions", [])
parallel = getattr(agent, "parallel_tool_calls", False)
handoff_desc = getattr(agent, "handoff_description", None) or "N/A"
handoffs = getattr(agent, "handoffs", [])
tools = getattr(agent, "tools", [])
guardrails_in = getattr(agent, "input_guardrails", [])
guardrails_out = getattr(agent, "output_guardrails", [])
output_type = getattr(agent, "output_type", None) or "N/A"
hooks = getattr(agent, "hooks", []) or []
# Handle model display based on agent type
model_display = self._get_model_display_for_info(agent_name, agent)
# Create a markdown table for agent details
# Build markdown content for agent info
markdown_content = f"""
# Agent: {agent_name}
# Agent Info: {name}
| Property | Value |
|----------|-------|
| Name | {agent.name} |
| Model | {model_display} |
| Functions | {len(agent.functions)} |
| Parallel Tool Calls | {'Yes' if agent.parallel_tool_calls else 'No'} |
| Property | Value |
|------------------------|-------------------------------|
| Key | {agent_key} |
| Name | {name} |
| Description | {clean_description} |
| Functions | {len(functions)} |
| Parallel Tool Calls | {"Yes" if parallel else "No"} |
| Handoff Description | {handoff_desc} |
| Handoffs | {len(handoffs)} |
| Tools | {len(tools)} |
| Input Guardrails | {len(guardrails_in)} |
| Output Guardrails | {len(guardrails_out)} |
| Output Type | {output_type} |
| Hooks | {len(hooks)} |
## Instructions
{instructions}
"""
"""
console.print(Markdown(markdown_content))
return True

10
src/cai/repl/commands/config.py
View File

@ -122,6 +122,16 @@ ENV_VARS = {
"name": "CAI_SUPPORT_INTERVAL",
"description": "Number of turns between support agent executions",
"default": "5"
},
22: {
"name": "CAI_STREAM",
"description": "Boolean to enable real-time, chunked responses instead of full messages.",
"default": "True"
},
23: {
"name": "CAI_WORKSPACE",
"description": "Name of the current workspace (affects log file naming)",
"default": None
},
}

120
src/cai/repl/commands/flush.py Normal file
View File

@ -0,0 +1,120 @@
"""
Flush command for CAI REPL.
This module provides commands for clear the context.
"""
import os
from typing import (
Dict,
List,
Optional
)
from rich.console import Console # pylint: disable=import-error
from rich.panel import Panel # pylint: disable=import-error
from cai.util import get_model_input_tokens
from cai.repl.commands.base import Command, register_command
from cai.sdk.agents.models.openai_chatcompletions import message_history
console = Console()
class FlushCommand(Command):
"""Command to flush the conversation history."""
def __init__(self):
"""Initialize the flush command."""
super().__init__(
name="/flush",
description="Clear the current conversation history.",
aliases=["/clear"]
)
def handle_no_args(self, messages: Optional[List[Dict]] = None) -> bool:
"""Handle the flush command when no args are provided.
Args:
messages: The conversation history messages
Returns:
True if the command was handled successfully
"""
# Use both the local messages parameter and the global message_history
local_messages = messages or []
global_history_length = len(message_history)
# Get token usage information before clearing
token_info = ""
context_usage = ""
# Access client through a function to avoid circular imports
# We can use globals() to get the client at runtime
client = self._get_client()
if client and hasattr(client, 'interaction_input_tokens') and hasattr(
client, 'total_input_tokens'):
model = os.getenv('CAI_MODEL', "qwen2.5:14b")
input_tokens = client.interaction_input_tokens if hasattr(
client, 'interaction_input_tokens') else 0
total_tokens = client.total_input_tokens if hasattr(
client, 'total_input_tokens') else 0
max_tokens = get_model_input_tokens(model)
context_pct = (input_tokens / max_tokens) * \
100 if max_tokens > 0 else 0
token_info = f"Current tokens: {input_tokens}, Total tokens: {total_tokens}"
context_usage = f"Context usage: {context_pct:.1f}% of {max_tokens} tokens"
# Clear both the local messages list and the global message_history
if local_messages:
local_messages.clear()
# Always clear the global message history
message_history.clear()
# Determine which length to report (use the greater of the two)
initial_length = max(len(local_messages) if messages else 0, global_history_length)
# Display information about the cleared messages
if initial_length > 0:
content = [
f"Conversation history cleared. Removed {initial_length} messages."
]
if token_info:
content.append(token_info)
if context_usage:
content.append(context_usage)
console.print(Panel(
"\n".join(content),
title="[bold cyan]Context Flushed[/bold cyan]",
border_style="blue",
padding=(1, 2)
))
else:
console.print(Panel(
"No conversation history to clear.",
title="[bold cyan]Context Flushed[/bold cyan]",
border_style="blue",
padding=(1, 2)
))
return True
def _get_client(self):
"""Get the CAI client from the global namespace.
This function avoids circular imports by accessing the client
at runtime instead of import time.
Returns:
The global CAI client instance or None if not available
"""
try:
# Import here to avoid circular import
from cai.repl.repl import client as global_client # pylint: disable=import-outside-toplevel # noqa: E501
return global_client
except (ImportError, AttributeError):
return None
# Register the /flush command
register_command(FlushCommand())

147
src/cai/repl/commands/graph.py
View File

@ -1,20 +1,64 @@
"""
Graph command for CAI REPL.
Graph command for CAI cli.
This module provides commands for visualizing the agent interaction graph.
It allows users to display a simple directed graph of the conversation history,
showing the sequence of user and agent interactions, including tool calls.
"""
from typing import (
List,
Optional
)
from typing import List, Optional
from rich.console import Console # pylint: disable=import-error
from rich.panel import Panel
from cai.repl.commands.base import Command, register_command
import os
import importlib.util
console = Console()
def find_agent_name_by_instructions(target_instructions: str, agents_dir: str) -> Optional[str]:
"""
Search all Python files in the agents directory for an agent whose 'instructions'
attribute matches the given target_instructions (ignoring leading/trailing whitespace).
Returns the agent's 'name' attribute if found, otherwise None.
Args:
target_instructions (str): The instructions string to match.
agents_dir (str): The directory containing agent files.
Returns:
Optional[str]: The agent name if found, else None.
"""
for filename in os.listdir(agents_dir):
if not filename.endswith(".py") or filename.startswith("__"):
continue
filepath = os.path.join(agents_dir, filename)
try:
spec = importlib.util.spec_from_file_location("agent_mod", filepath)
agent_mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(agent_mod)
for attr_name in dir(agent_mod):
attr = getattr(agent_mod, attr_name)
if hasattr(attr, "instructions"):
agent_instructions = getattr(attr, "instructions", None)
if agent_instructions and agent_instructions.strip() == target_instructions.strip():
agent_name = getattr(attr, "name", None)
if agent_name:
return agent_name
except Exception:
continue
return None
class GraphCommand(Command):
"""Command for visualizing the agent interaction graph."""
"""
Command for visualizing the agent interaction graph.
This command displays a directed graph of the conversation history,
showing the sequence of user and agent messages, and highlighting
tool calls made by the agent.
"""
def __init__(self):
"""Initialize the graph command."""
@ -25,31 +69,104 @@ class GraphCommand(Command):
)
def handle(self, args: Optional[List[str]] = None) -> bool:
"""Handle the graph command.
"""
Handle the /graph command.
Args:
args: Optional list of command arguments
Returns:
True if the command was handled successfully, False otherwise
bool: True if the command was handled successfully, False otherwise.
"""
return self.handle_graph_show()
def handle_graph_show(self) -> bool:
"""Handle /graph show command"""
from cai.repl.repl import client # pylint: disable=import-error
# Import here to avoid circular imports
if not client or not client._graph: # pylint: disable=protected-access
from cai.sdk.agents.models.openai_chatcompletions import message_history
if not message_history:
console.print("[yellow]No conversation graph available.[/yellow]")
return True
try:
agents_dir = os.path.join(os.path.dirname(__file__), "../../agents")
agents_dir = os.path.abspath(agents_dir)
import networkx as nx
G = nx.DiGraph()
last_agent_name = None
prev_node_idx = None # Track the last node actually added (not system)
for idx, msg in enumerate(message_history):
role = msg.get("role", "unknown")
# If the message is from the system, update last_agent_name but do not add a node
if role == "system":
system_msg = msg.get("content", "").strip()
agent_name = find_agent_name_by_instructions(system_msg, agents_dir)
if agent_name:
last_agent_name = agent_name
continue
label = role
extra_info = ""
if role == "assistant":
if last_agent_name:
label = last_agent_name
else:
label = "assistant"
if msg.get("tool_calls"):
tool_call = msg["tool_calls"][0]
if tool_call.get("function"):
func_name = tool_call["function"].get("name", "")
func_args = tool_call["function"].get("arguments", "")
extra_info = f"\n[cyan]Tool:[/cyan] [bold]{func_name}[/bold]\n[cyan]Args:[/cyan] {func_args}"
elif role == "user":
user_content = msg.get("content", "")
if user_content:
extra_info = f"\n{user_content}"
label = role
else:
label = role
G.add_node(idx, role=label, extra_info=extra_info)
if prev_node_idx is not None:
G.add_edge(prev_node_idx, idx)
prev_node_idx = idx
def ascii_graph(G):
"""
Render the conversation graph as a sequence of panels with arrows.
Args:
G (networkx.DiGraph): The conversation graph.
Returns:
List: List of rich Panel objects and arrow strings.
"""
lines = []
node_list = list(G.nodes(data=True))
for i, (idx, data) in enumerate(node_list):
role = data.get("role", "unknown")
extra_info = data.get("extra_info", "")
role_fmt = f"[bold][blue]{role[:1].upper()}{role[1:]}[/blue][/bold]"
panel_content = f"{role_fmt}"
if extra_info:
panel_content += f"{extra_info}"
panel = Panel(
panel_content,
expand=False,
border_style="cyan"
)
lines.append(panel)
if i < len(node_list) - 1:
lines.append("[cyan] │\n │\n ▼[/cyan]")
return lines
console.print("\n[bold]Conversation Graph:[/bold]")
console.print("------------------")
console.print(
client._graph.ascii()) # pylint: disable=protected-access
if len(G.nodes) == 0:
console.print("[yellow]No messages to display in graph.[/yellow]")
else:
for item in ascii_graph(G):
console.print(item)
console.print()
return True
except Exception as e: # pylint: disable=broad-except

119
src/cai/repl/commands/history.py
View File

@ -2,8 +2,12 @@
History command for CAI REPL.
This module provides commands for displaying conversation history.
"""
import json
from typing import Any, Dict, List, Optional
from rich.console import Console # pylint: disable=import-error
from rich.table import Table # pylint: disable=import-error
from rich.text import Text # pylint: disable=import-error
from cai.repl.commands.base import Command, register_command
@ -20,6 +24,20 @@ class HistoryCommand(Command):
description="Display the conversation history",
aliases=["/h"]
)
def handle(self, args: Optional[List[str]] = None,
messages: Optional[List[Dict]] = None) -> bool:
"""Handle the history command.
Args:
args: Optional list of command arguments
messages: Optional list of conversation messages
Returns:
True if the command was handled successfully, False otherwise
"""
# Currently, the history command doesn't take any arguments
return self.handle_no_args()
def handle_no_args(self) -> bool:
"""Handle the command when no arguments are provided.
@ -27,15 +45,15 @@ class HistoryCommand(Command):
Returns:
True if the command was handled successfully, False otherwise
"""
# Access messages directly from repl.py's global scope
# Access messages directly from openai_chatcompletions.py
try:
from cai.repl.repl import messages # pylint: disable=import-outside-toplevel # noqa: E501
from cai.sdk.agents.models.openai_chatcompletions import message_history # pylint: disable=import-outside-toplevel # noqa: E501
except ImportError:
console.print(
"[red]Error: Could not access conversation history[/red]")
return False
if not messages:
if not message_history:
console.print("[yellow]No conversation history available[/yellow]")
return True
@ -50,34 +68,87 @@ class HistoryCommand(Command):
table.add_column("Content", style="green")
# Add messages to the table
for idx, msg in enumerate(messages, 1):
role = msg.get("role", "unknown")
content = msg.get("content", "")
for idx, msg in enumerate(message_history, 1):
try:
role = msg.get("role", "unknown")
content = msg.get("content", "")
tool_calls = msg.get("tool_calls", None)
# Truncate long content for better display
if len(content) > 100:
content = content[:97] + "..."
# Create formatted content based on message type
formatted_content = self._format_message_content(
content, tool_calls)
# Color the role based on type
if role == "user":
role_style = "cyan"
elif role == "assistant":
role_style = "yellow"
else:
role_style = "red"
# Color the role based on type
if role == "user":
role_style = "cyan"
elif role == "assistant":
role_style = "yellow"
else:
role_style = "red"
# Add a newline between each role for better readability
if idx > 1:
table.add_row("", "", "")
# Add a newline between each role for better readability
if idx > 1:
table.add_row("", "", "")
table.add_row(
str(idx),
f"[{role_style}]{role}[/{role_style}]",
content
)
table.add_row(
str(idx),
f"[{role_style}]{role}[/{role_style}]",
formatted_content
)
except Exception as e:
# Log error but continue with next message
console.print(f"[red]Error displaying message {idx}: {e}[/red]")
continue
console.print(table)
return True
def _format_message_content(
self, content: Any, tool_calls: List[Dict[str, Any]]
) -> str:
"""Format message content for display, handling both text and tool calls.
Args:
content: Text content of the message
tool_calls: List of tool calls if present
Returns:
Formatted string representation of the message content
"""
if tool_calls:
# Format tool calls into a readable string
result = []
for tc in tool_calls:
func_details = tc.get("function", {})
func_name = func_details.get("name", "unknown_function")
# Format arguments (pretty-print JSON if possible)
args_str = func_details.get("arguments", "{}")
try:
# Parse and re-format JSON for better readability
args_dict = json.loads(args_str)
args_formatted = json.dumps(args_dict, indent=2)
# Limit to first 100 chars for display
if len(args_formatted) > 100:
args_formatted = args_formatted[:97] + "..."
except (json.JSONDecodeError, TypeError):
# If not valid JSON, use as is
args_formatted = args_str
if len(args_formatted) > 100:
args_formatted = args_formatted[:97] + "..."
result.append(f"Function: [bold blue]{func_name}[/bold blue]")
result.append(f"Args: {args_formatted}")
return "\n".join(result)
elif content:
# Regular text content (truncate if too long)
if len(content) > 100:
return content[:97] + "..."
return content
else:
# No content or tool calls (empty message)
return "[dim italic]Empty message[/dim italic]"
# Register the command

4
src/cai/repl/commands/memory.py
View File

@ -1,4 +0,0 @@
"""
Memory command for CAI REPL.
This module provides commands for managing memory collections.
"""

3
src/cai/repl/commands/model.py
View File

@ -145,6 +145,9 @@ class ModelCommand(Command):
{"name": "gpt-4-turbo",
"description": "Fast and powerful GPT-4 model"}
],
"OpenAI GPT-4o-mini": [
{"name": "gpt-4o-mini", "description": " GPT-4o mini model"}
],
"OpenAI GPT-4.5": [
{
"name": "gpt-4.5-preview",

96
src/cai/repl/commands/turns.py
View File

@ -1,96 +0,0 @@
"""
Turns command for CAI REPL.
This module provides commands for viewing and changing the maximum number
of turns.
"""
import os
from typing import (
List,
Optional
)
from rich.console import Console # pylint: disable=import-error
from rich.panel import Panel # pylint: disable=import-error
from cai.repl.commands.base import Command, register_command
console = Console()
class TurnsCommand(Command):
"""Command for viewing and changing the maximum number of turns."""
def __init__(self):
"""Initialize the turns command."""
super().__init__(
name="/turns",
description="View or change the maximum number of turns",
aliases=["/t"]
)
def handle(self, args: Optional[List[str]] = None) -> bool:
"""Handle the turns command.
Args:
args: Optional list of command arguments
Returns:
True if the command was handled successfully, False otherwise
"""
return self.handle_turns_command(args)
def handle_turns_command(self, args: List[str]) -> bool:
"""Change the maximum number of turns for CAI.
Args:
args: List containing the number of turns
Returns:
bool: True if the max turns was changed successfully
"""
if not args:
# Display current max turns
max_turns_info = os.getenv("CAI_MAX_TURNS", "inf")
console.print(Panel(
f"Current maximum turns: [bold green]{
max_turns_info}[/bold green]",
border_style="green",
title="Max Turns Setting"
))
# Usage instructions
console.print(
"\n[cyan]Usage:[/cyan] [bold]/turns <number_of_turns>[/bold]")
console.print("[cyan]Examples:[/cyan]")
console.print(" [bold]/turns 10[/bold] - Limit to 10 turns")
console.print(" [bold]/turns inf[/bold] - Unlimited turns")
return True
try:
turns = args[0]
# Check if it's a number or 'inf'
if turns.lower() == 'inf':
turns = 'inf'
else:
turns = int(turns)
# Set the max turns in environment variable
os.environ["CAI_MAX_TURNS"] = turns
console.print(Panel(
f"Maximum turns changed to: [bold green]{turns}[/bold green]\n"
"[yellow]Note: This will take effect on the next run[/yellow]",
border_style="green",
title="Max Turns Changed"
))
return True
except ValueError:
console.print(Panel(
"Error: Max turns must be a number or 'inf'",
border_style="red",
title="Invalid Input"
))
return False
# Register the command
register_command(TurnsCommand())

687
src/cai/repl/commands/workspace.py Normal file
View File

@ -0,0 +1,687 @@
"""
Virtualization command for CAI REPL.
This module provides commands for setting up and managing Docker virtualization
environments.
"""
# Standard library imports
import os
import json
import subprocess
import datetime
import time
from typing import List, Optional, Dict, Any, Tuple
# Third-party imports
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.markdown import Markdown
import rich.box
# Local imports
from cai.repl.commands.base import Command, register_command
console = Console()
class WorkspaceCommand(Command):
"""Command for workspace management within Docker containers or locally."""
def __init__(self):
"""Initialize the workspace command."""
super().__init__(
name="/workspace",
description=(
"Set or display the current workspace name and manage files."
" Affects log file naming and where files are stored."
),
aliases=["/ws"]
)
# Add subcommands
self.add_subcommand(
"set",
"Set the current workspace name",
self.handle_set
)
self.add_subcommand(
"get",
"Display the current workspace name",
self.handle_get
)
self.add_subcommand(
"ls",
"List files in the workspace",
self.handle_ls_subcommand
)
self.add_subcommand(
"exec",
"Execute a command in the workspace",
self.handle_exec_subcommand
)
self.add_subcommand(
"copy",
"Copy files between host and container",
self.handle_copy_subcommand
)
def handle(self, args: Optional[List[str]] = None) -> bool:
"""Handle the workspace command.
Args:
args: Optional list of command arguments
Returns:
True if the command was handled successfully, False otherwise
"""
# If there are subcommands, process them
if args and args[0] in self.subcommands:
return super().handle(args)
# No arguments means show workspace info (same as get)
return self.handle_get()
def handle_no_args(self) -> bool:
"""Handle the command when no arguments are provided."""
return self.handle_get()
def handle_get(self, _: Optional[List[str]] = None) -> bool:
"""Display the current workspace name and directory information."""
# Get workspace info
workspace_name = os.getenv("CAI_WORKSPACE", None)
# Check if a container is active
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
# Determine environment (container or host)
if active_container:
try:
# Get container details
result = subprocess.run(
["docker", "inspect", active_container],
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
container_info = json.loads(result.stdout)
if container_info:
image = container_info[0].get("Config", {}).get("Image", "unknown")
env_type = "container"
env_name = f"Container ({image})"
# For containers, if workspace is set, use container workspace path
# otherwise use root directory
if workspace_name:
# This will create the workspace in the container if it doesn't exist
workspace_dir = f"/workspace/workspaces/{workspace_name}"
# Ensure the directory exists in the container
subprocess.run(
["docker", "exec", active_container, "mkdir", "-p", workspace_dir],
capture_output=True,
check=False
)
else:
workspace_dir = "/"
else:
env_type = "host"
env_name = "Host System (container not running)"
# Use common._get_workspace_dir() for consistency
try:
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
workspace_dir = get_common_workspace_dir()
except ImportError:
workspace_dir = os.getcwd() # Basic fallback
except Exception:
env_type = "host"
env_name = "Host System (error inspecting container)"
# Use common._get_workspace_dir() for consistency
try:
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
workspace_dir = get_common_workspace_dir()
except ImportError:
workspace_dir = os.getcwd() # Basic fallback
else:
env_type = "host"
env_name = "Host System"
# Use common._get_workspace_dir() for consistency
try:
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
workspace_dir = get_common_workspace_dir()
except ImportError:
workspace_dir = os.getcwd() # Basic fallback
# Show workspace information
console.print(
Panel(
f"Current workspace: [bold green]{workspace_name or 'None'}[/bold green]\n"
f"Working in environment: [bold]{env_name}[/bold]\n"
f"Workspace directory: [bold]{workspace_dir}[/bold]",
title="Workspace Information",
border_style="green"
)
)
# Show available workspace commands
console.print("\n[cyan]Workspace Commands:[/cyan]")
console.print(
" [bold]/workspace set <name>[/bold] - "
"Set the current workspace name")
console.print(
" [bold]/workspace ls[/bold] - "
"List files in the workspace")
console.print(
" [bold]/workspace exec <cmd>[/bold] - "
"Execute a command in the workspace")
if active_container:
console.print(
" [bold]/workspace copy <src> <dst>[/bold] - "
"Copy files between host and container")
# List contents of the workspace
self._list_workspace_contents(env_type, workspace_dir)
return True
def handle_set(self, args: Optional[List[str]] = None) -> bool:
"""Set the current workspace name """
if not args or len(args) != 1:
console.print(
"[yellow]Usage: /workspace set <workspace_name>[/yellow]"
)
return False
workspace_name = args[0]
# Allow alphanumeric, underscores, hyphens
if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name):
console.print(
"[red]Invalid workspace name. "
"Use alphanumeric, underscores, or hyphens only.[/red]"
)
return False
# Import the necessary modules for setting environment variables
# And for getting workspace dir consistently
try:
from cai.repl.commands.config import set_env_var
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
from cai.tools.common import _get_container_workspace_path as get_common_container_path
# Set the environment variable
if not set_env_var("CAI_WORKSPACE", workspace_name):
console.print(
"[red]Failed to set workspace environment variable.[/red]"
)
return False
except ImportError:
# Fallback if import fails
os.environ["CAI_WORKSPACE"] = workspace_name
# Define basic fallbacks for path functions if import failed
def get_common_workspace_dir():
base = os.getenv("CAI_WORKSPACE_DIR", ".") # Default to current dir base
name = os.getenv("CAI_WORKSPACE")
if name:
return os.path.abspath(os.path.join(base, name))
return os.path.abspath(base) # Use base dir if no name
def get_common_container_path():
name = os.getenv("CAI_WORKSPACE")
if name:
return f"/workspace/workspaces/{name}"
return "/" # Default container path
# Get the new workspace directory using the common function
new_workspace_dir = get_common_workspace_dir()
# Create the directory if it doesn't exist on host
try: # Add try-except for robustness
os.makedirs(new_workspace_dir, exist_ok=True)
except OSError as e:
console.print(f"[red]Error creating host directory {new_workspace_dir}: {e}[/red]")
# Decide if this is fatal or just a warning
# If container is active, also create the directory in the container
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
if active_container:
# Check if container is running
check_process = subprocess.run(
["docker", "inspect", "--format", "{{.State.Running}}", active_container],
capture_output=True,
text=True,
check=False
)
if check_process.returncode == 0 and "true" in check_process.stdout.lower():
# Get container workspace path using the common function
container_workspace_path = get_common_container_path()
try:
mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path]
mkdir_result = subprocess.run(
mkdir_cmd,
capture_output=True,
text=True,
check=False
)
if mkdir_result.returncode == 0:
console.print(
f"[dim]Created workspace directory in container: {container_workspace_path}[/dim]"
)
else:
console.print(
f"[yellow]Warning: Could not create workspace directory in container: {mkdir_result.stderr}[/yellow]"
)
except Exception as e:
console.print(
f"[yellow]Warning: Failed to setup workspace in container: {str(e)}[/yellow]"
)
# Use a different panel style to indicate success
console.print(
Panel(
f"Workspace changed to: [bold green]{workspace_name}[/bold green]\n"
f"New workspace directory: [bold]{new_workspace_dir}[/bold]",
title="Workspace Updated",
border_style="green"
)
)
return True
def _get_workspace_dir(self) -> str:
"""Get the host workspace directory using the common utility.
Returns:
The host workspace directory path.
"""
try:
# Use the centralized function from common.py
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
return get_common_workspace_dir()
except ImportError:
# Provide a basic fallback if import fails, mirroring common.py logic
# without 'cai_default'
base_dir = os.getenv("CAI_WORKSPACE_DIR")
workspace_name = os.getenv("CAI_WORKSPACE")
if base_dir and workspace_name:
# Basic validation
if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name):
print(f"[yellow]Warning: Invalid CAI_WORKSPACE name '{workspace_name}' in fallback.[/yellow]")
# Fallback to base directory if name is invalid
return os.path.abspath(base_dir)
target_dir = os.path.join(base_dir, workspace_name)
return os.path.abspath(target_dir)
elif base_dir:
# If only base dir is set, use that
return os.path.abspath(base_dir)
else:
# Default to current working directory if nothing else is set
return os.getcwd()
def _list_workspace_contents(self, env_type: str, workspace_dir: str) -> None:
"""List the contents of the workspace.
Args:
env_type: The environment type (container or host)
workspace_dir: The workspace directory
"""
console.print("\n[bold]Workspace Contents:[/bold]")
if env_type == "container":
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
# For containers, use the workspace path provided
# This should already be the correct path from handle_get
# First ensure the workspace directory exists in the container
try:
mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", workspace_dir]
subprocess.run(
mkdir_cmd,
capture_output=True,
text=True,
check=False
)
# Now list the contents
result = subprocess.run(
["docker", "exec", active_container, "ls", "-la", workspace_dir],
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
console.print(result.stdout)
else:
console.print(f"[yellow]Error listing container files: {result.stderr}[/yellow]")
# Fallback to host
self._list_host_files(workspace_dir)
except Exception as e:
console.print(f"[yellow]Error accessing container: {str(e)}[/yellow]")
# Fallback to host
self._list_host_files(workspace_dir)
else:
# List files in host
self._list_host_files(workspace_dir)
def _list_host_files(self, workspace_dir: str) -> None:
"""List files in the host workspace.
Args:
workspace_dir: The workspace directory
"""
# Ensure the directory exists
os.makedirs(workspace_dir, exist_ok=True)
try:
result = subprocess.run(
["ls", "-la", workspace_dir],
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
console.print(result.stdout)
else:
console.print(f"[yellow]Error listing files: {result.stderr}[/yellow]")
except Exception as e:
console.print(f"[yellow]Error: {str(e)}[/yellow]")
def handle_ls_subcommand(self, args: Optional[List[str]] = None) -> bool:
"""Handle the ls subcommand.
Args:
args: Optional list of subcommand arguments
Returns:
True if the subcommand was handled successfully, False otherwise
"""
# Get workspace info using common functions
try:
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
from cai.tools.common import _get_container_workspace_path as get_common_container_path
except ImportError:
# Define basic fallbacks if import fails
def get_common_workspace_dir():
base = os.getenv("CAI_WORKSPACE_DIR", ".")
name = os.getenv("CAI_WORKSPACE")
if name: return os.path.abspath(os.path.join(base, name))
return os.path.abspath(base)
def get_common_container_path():
name = os.getenv("CAI_WORKSPACE")
if name: return f"/workspace/workspaces/{name}"
return "/"
host_workspace_dir = get_common_workspace_dir()
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
# Execute command in the appropriate environment
if active_container:
# Use the container workspace path from common function
container_workspace_path = get_common_container_path()
# Determine the target path within the container
target_path_in_container = container_workspace_path
if args:
# Ensure args[0] is treated as relative to the workspace
target_path_in_container = os.path.join(container_workspace_path, args[0])
# Ensure the base workspace directory exists in the container
mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path]
subprocess.run(
mkdir_cmd,
capture_output=True,
text=True,
check=False
)
# Try in container
result = subprocess.run(
["docker", "exec", active_container, "ls", "-la", target_path_in_container], # Use target path
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
console.print(result.stdout)
return True
# If failed, try on host
console.print(f"[yellow]Failed to list files in container: {result.stderr}[/yellow]")
console.print("[yellow]Falling back to host system...[/yellow]")
# List on host
# Determine target path on host relative to host workspace dir
target_path_on_host = host_workspace_dir
if args:
# Ensure args[0] is treated as relative to the workspace
target_path_on_host = os.path.join(host_workspace_dir, args[0])
# Ensure the target directory exists on host before listing
# Use os.path.dirname if target is potentially a file path
dir_to_ensure = os.path.dirname(target_path_on_host) if '.' in os.path.basename(target_path_on_host) else target_path_on_host
try:
os.makedirs(dir_to_ensure, exist_ok=True)
except OSError as e:
console.print(f"[red]Error creating directory {dir_to_ensure} on host: {e}[/red]")
# Potentially return False or handle error appropriately
try:
result = subprocess.run(
["ls", "-la", target_path_on_host], # Use target path
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
console.print(result.stdout)
return True
else:
console.print(f"[red]Error listing files: {result.stderr}[/red]")
return False
except Exception as e:
console.print(f"[red]Error: {str(e)}[/red]")
return False
return True
def handle_exec_subcommand(self, args: Optional[List[str]] = None) -> bool:
"""Handle the exec subcommand.
Args:
args: Optional list of subcommand arguments
Returns:
True if the subcommand was handled successfully, False otherwise
"""
if not args:
console.print("[yellow]Please specify a command to execute.[/yellow]")
return False
command = " ".join(args)
# Get workspace info using common functions
try:
from cai.tools.common import _get_workspace_dir as get_common_workspace_dir
from cai.tools.common import _get_container_workspace_path as get_common_container_path
except ImportError:
# Define basic fallbacks if import fails
def get_common_workspace_dir():
base = os.getenv("CAI_WORKSPACE_DIR", ".")
name = os.getenv("CAI_WORKSPACE")
if name: return os.path.abspath(os.path.join(base, name))
return os.path.abspath(base)
def get_common_container_path():
name = os.getenv("CAI_WORKSPACE")
if name: return f"/workspace/workspaces/{name}"
return "/"
host_workspace_dir = get_common_workspace_dir()
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
# Execute in container if active
if active_container:
try:
# Use the container workspace path from common function
container_workspace_path = get_common_container_path()
# First ensure the workspace directory exists in the container
mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path]
subprocess.run(
mkdir_cmd,
capture_output=True,
text=True,
check=False
)
# Execute the command in the container's workspace directory
result = subprocess.run(
["docker", "exec", "-w", container_workspace_path, active_container, "sh", "-c", command],
capture_output=True,
text=True,
check=False
)
console.print(f"[dim]$ {command}[/dim]")
if result.stdout:
console.print(result.stdout)
if result.stderr:
console.print(f"[yellow]{result.stderr}[/yellow]")
if result.returncode != 0:
console.print("[yellow]Command failed in container. Trying on host...[/yellow]")
return self._exec_on_host(command, host_workspace_dir) # Pass host_workspace_dir
return True
except Exception as e:
console.print(f"[yellow]Error executing in container: {str(e)}[/yellow]")
console.print("[yellow]Falling back to host execution...[/yellow]")
# Execute on host
return self._exec_on_host(command, host_workspace_dir) # Pass host_workspace_dir
def _exec_on_host(self, command: str, workspace_dir: str) -> bool:
"""Execute a command on the host.
Args:
command: The command to execute
workspace_dir: The workspace directory
Returns:
True if the command was executed successfully, False otherwise
"""
# Ensure the directory exists
os.makedirs(workspace_dir, exist_ok=True)
try:
result = subprocess.run(
command,
shell=True, # nosec B602
capture_output=True,
text=True,
check=False,
cwd=workspace_dir
)
console.print(f"[dim]$ {command}[/dim]")
if result.stdout:
console.print(result.stdout)
if result.stderr:
console.print(f"[yellow]{result.stderr}[/yellow]")
return result.returncode == 0
except Exception as e:
console.print(f"[red]Error executing command: {str(e)}[/red]")
return False
def handle_copy_subcommand(self, args: Optional[List[str]] = None) -> bool:
"""Handle the copy subcommand.
Args:
args: Optional list of subcommand arguments
Returns:
True if the subcommand was handled successfully, False otherwise
"""
if not args or len(args) < 2:
console.print("[yellow]Please specify source and destination for copy.[/yellow]")
console.print("Usage: /workspace copy <source> <destination>")
return False
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
if not active_container:
console.print("[yellow]No active container. Copy only works with containers.[/yellow]")
return False
source = args[0]
destination = args[1]
# Check if copying from container to host or vice versa
if source.startswith("container:"):
# Copy from container to host
container_path = source[10:] # Remove "container:" prefix
host_path = destination
if not container_path.startswith("/"):
container_path = f"/workspace/{container_path}"
try:
result = subprocess.run(
["docker", "cp", f"{active_container}:{container_path}", host_path],
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
console.print(f"[green]Copied from container:{container_path} to {host_path}[/green]")
return True
else:
console.print(f"[red]Error copying from container: {result.stderr}[/red]")
return False
except Exception as e:
console.print(f"[red]Error: {str(e)}[/red]")
return False
elif destination.startswith("container:"):
# Copy from host to container
host_path = source
container_path = destination[10:] # Remove "container:" prefix
if not container_path.startswith("/"):
container_path = f"/workspace/{container_path}"
try:
result = subprocess.run(
["docker", "cp", host_path, f"{active_container}:{container_path}"],
capture_output=True,
text=True,
check=False
)
if result.returncode == 0:
console.print(f"[green]Copied from {host_path} to container:{container_path}[/green]")
return True
else:
console.print(f"[red]Error copying to container: {result.stderr}[/red]")
return False
except Exception as e:
console.print(f"[red]Error: {str(e)}[/red]")
return False
else:
# Ambiguous copy - show help
console.print("[yellow]Ambiguous copy direction. Please specify container: prefix.[/yellow]")
console.print("Examples:")
console.print(" /workspace copy file.txt container:file.txt # Host to container")
console.print(" /workspace copy container:file.txt file.txt # Container to host")
return False
# Register the commands
register_command(WorkspaceCommand())

17
src/cai/repl/ui/toolbar.py
View File

@ -18,7 +18,7 @@ toolbar_last_refresh = [datetime.datetime.now()]
toolbar_cache = {
'html': "",
'last_update': datetime.datetime.now(),
'refresh_interval': 60 # Refresh every 60 seconds
'refresh_interval': 5 # Refresh every 60 seconds
}
# Cache for system information that rarely changes
@ -57,7 +57,22 @@ def update_toolbar_in_background():
ip_address = sys_info['ip_address']
os_name = sys_info['os_name']
os_version = sys_info['os_version']
# Get the current workspace and base directory
workspace_name = os.getenv("CAI_WORKSPACE")
base_dir = os.getenv("CAI_WORKSPACE_DIR", "workspaces")
# Construct the workspace path
standard_path = os.path.join(base_dir, workspace_name) if workspace_name else ""
workspace_path = ""
if workspace_name:
if os.path.isdir(standard_path):
workspace_path = standard_path
elif os.path.isdir(workspace_name):
workspace_path = os.path.abspath(workspace_name)
else:
workspace_path = standard_path
# Get Ollama information
ollama_status = "unavailable"
try:

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

@ -7,6 +7,9 @@ import os
import litellm
import tiktoken
import inspect
import hashlib
import re
import asyncio
from collections.abc import AsyncIterator, Iterable
from dataclasses import dataclass, field
@ -93,9 +96,40 @@ if TYPE_CHECKING:
# Suppress debug info from litellm
litellm.suppress_debug_info = True
if os.getenv('CAI_MODEL') == "o3-mini" or os.getenv('CAI_MODEL') == "gemini-1.5-pro":
litellm.drop_params = True
_USER_AGENT = f"Agents/Python {__version__}"
_HEADERS = {"User-Agent": _USER_AGENT}
message_history = []
# Function to add a message to history if it's not a duplicate
def add_to_message_history(msg):
"""Add a message to history if it's not a duplicate."""
if not message_history:
message_history.append(msg)
return
is_duplicate = False
if msg.get("role") in ["system", "user"]:
is_duplicate = any(
existing.get("role") == msg.get("role") and
existing.get("content") == msg.get("content")
for existing in message_history
)
elif msg.get("role") == "assistant" and msg.get("tool_calls"):
is_duplicate = any(
existing.get("role") == "assistant" and
existing.get("tool_calls") and
existing["tool_calls"][0].get("id") == msg["tool_calls"][0].get("id")
for existing in message_history
)
if not is_duplicate:
message_history.append(msg)
@dataclass
class _StreamingState:
@ -235,7 +269,32 @@ class OpenAIChatCompletionsModel(Model):
"role": "system",
},
)
# --- Add to message_history: user, system, and assistant tool call messages ---
# Add system prompt to message_history
if system_instructions:
sys_msg = {
"role": "system",
"content": system_instructions
}
add_to_message_history(sys_msg)
# Add user prompt(s) to message_history
if isinstance(input, str):
user_msg = {
"role": "user",
"content": input
}
add_to_message_history(user_msg)
elif isinstance(input, list):
for item in input:
# Try to extract user messages
if isinstance(item, dict):
if item.get("role") == "user":
user_msg = {
"role": "user",
"content": item.get("content", "")
}
add_to_message_history(user_msg)
# Get token count estimate before API call for consistent counting
estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages)
@ -335,6 +394,36 @@ class OpenAIChatCompletionsModel(Model):
tool_output=None, # Don't pass tool output here, we're using direct display
)
# --- Add assistant tool call to message_history if present ---
# If the response contains tool_calls, add them to message_history as assistant messages
assistant_msg = response.choices[0].message
if hasattr(assistant_msg, "tool_calls") and assistant_msg.tool_calls:
for tool_call in assistant_msg.tool_calls:
# Compose a message for the tool call
tool_call_msg = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": tool_call.id,
"type": tool_call.type,
"function": {
"name": tool_call.function.name,
"arguments": tool_call.function.arguments
}
}
]
}
add_to_message_history(tool_call_msg)
# If the assistant message is just text, add it as well
elif hasattr(assistant_msg, "content") and assistant_msg.content:
asst_msg = {
"role": "assistant",
"content": assistant_msg.content
}
add_to_message_history(asst_msg)
usage = (
Usage(
requests=1,
@ -404,7 +493,29 @@ class OpenAIChatCompletionsModel(Model):
"role": "system",
},
)
# --- Add to message_history: user, system prompts ---
if system_instructions:
sys_msg = {
"role": "system",
"content": system_instructions
}
add_to_message_history(sys_msg)
if isinstance(input, str):
user_msg = {
"role": "user",
"content": input
}
add_to_message_history(user_msg)
elif isinstance(input, list):
for item in input:
if isinstance(item, dict):
if item.get("role") == "user":
user_msg = {
"role": "user",
"content": item.get("content", "")
}
add_to_message_history(user_msg)
# Get token count estimate before API call for consistent counting
estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages)
@ -429,7 +540,25 @@ class OpenAIChatCompletionsModel(Model):
# Initialize a streaming text accumulator for rich display
streaming_text_buffer = ""
# For tool call streaming, accumulate tool_calls to add to message_history at the end
streamed_tool_calls = []
# Ollama specific: accumulate full content to check for function calls at the end
# Some Ollama models output the function call as JSON in the text content
ollama_full_content = ""
is_ollama = False
model_str = str(self.model).lower()
is_ollama = self.is_ollama or "ollama" in model_str or ":" in model_str or "qwen" in model_str
# Add visual separation before agent output
if streaming_context and should_show_rich_stream:
# If we're using rich context, we'll add separation through that
pass
else:
# Print clear visual separator
print("\n")
async for chunk in stream:
if not state.started:
state.started = True
@ -453,6 +582,10 @@ class OpenAIChatCompletionsModel(Model):
choices = [{"delta": chunk.delta}]
elif isinstance(chunk, dict) and 'choices' in chunk:
choices = chunk['choices']
# Special handling for Qwen/Ollama chunks
elif isinstance(chunk, dict) and ('content' in chunk or 'function_call' in chunk):
# Qwen direct delta format - convert to standard
choices = [{"delta": chunk}]
else:
# Skip chunks that don't contain choice data
continue
@ -478,6 +611,10 @@ class OpenAIChatCompletionsModel(Model):
content = delta['content']
if content:
# For Ollama, we need to accumulate the full content to check for function calls
if is_ollama:
ollama_full_content += content
# Add to the streaming text buffer
streaming_text_buffer += content
@ -589,11 +726,7 @@ class OpenAIChatCompletionsModel(Model):
# Handle tool calls
# Because we don't know the name of the function until the end of the stream, we'll
# save everything and yield events at the end
tool_calls = None
if hasattr(delta, 'tool_calls') and delta.tool_calls:
tool_calls = delta.tool_calls
elif isinstance(delta, dict) and 'tool_calls' in delta and delta['tool_calls']:
tool_calls = delta['tool_calls']
tool_calls = self._detect_and_format_function_calls(delta)
if tool_calls:
for tc_delta in tool_calls:
@ -636,9 +769,146 @@ class OpenAIChatCompletionsModel(Model):
call_id = tc_delta.id or ""
elif isinstance(tc_delta, dict) and 'id' in tc_delta:
call_id = tc_delta.get('id', "") or ""
else:
# For Qwen models, generate a predictable ID if none is provided
if state.function_calls[tc_index].name:
# Generate a stable ID from the function name and arguments
call_id = f"call_{hashlib.md5(state.function_calls[tc_index].name.encode()).hexdigest()[:8]}"
state.function_calls[tc_index].call_id += call_id
# --- Accumulate tool call for message_history ---
# Only add if not already present (avoid duplicates in streaming)
tool_call_msg = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": state.function_calls[tc_index].call_id,
"type": "function",
"function": {
"name": state.function_calls[tc_index].name,
"arguments": state.function_calls[tc_index].arguments
}
}
]
}
# Only add if not already in streamed_tool_calls
if tool_call_msg not in streamed_tool_calls:
streamed_tool_calls.append(tool_call_msg)
add_to_message_history(tool_call_msg)
# Special handling for Ollama - check if accumulated text contains a valid function call
if is_ollama and ollama_full_content and len(state.function_calls) == 0:
# Look for JSON object that might be a function call
try:
# Try to extract a JSON object from the content
json_start = ollama_full_content.find('{')
json_end = ollama_full_content.rfind('}') + 1
if json_start >= 0 and json_end > json_start:
json_str = ollama_full_content[json_start:json_end]
# Try to parse the JSON
parsed = json.loads(json_str)
# Check if it looks like a function call
if ('name' in parsed and 'arguments' in parsed):
logger.debug(f"Found valid function call in Ollama output: {json_str}")
# Create a tool call ID
tool_call_id = f"call_{hashlib.md5((parsed['name'] + str(time.time())).encode()).hexdigest()[:8]}"
# Ensure arguments is a valid JSON string
arguments_str = ""
if isinstance(parsed['arguments'], dict):
# Remove 'ctf' field if it exists
if 'ctf' in parsed['arguments']:
del parsed['arguments']['ctf']
arguments_str = json.dumps(parsed['arguments'])
elif isinstance(parsed['arguments'], str):
# If it's already a string, check if it's valid JSON
try:
# Try parsing to validate and remove 'ctf' if present
args_dict = json.loads(parsed['arguments'])
if isinstance(args_dict, dict) and 'ctf' in args_dict:
del args_dict['ctf']
arguments_str = json.dumps(args_dict)
except:
# If not valid JSON, encode it as a JSON string
arguments_str = json.dumps(parsed['arguments'])
else:
# For any other type, convert to string and then JSON
arguments_str = json.dumps(str(parsed['arguments']))
# Add it to our function_calls state
state.function_calls[0] = ResponseFunctionToolCall(
id=FAKE_RESPONSES_ID,
arguments=arguments_str,
name=parsed['name'],
type="function_call",
call_id=tool_call_id,
)
# Display the tool call in CLI
from cai.util import cli_print_agent_messages
try:
# Create a message-like object to display the function call
tool_msg = type('ToolCallWrapper', (), {
'content': None,
'tool_calls': [
type('ToolCallDetail', (), {
'function': type('FunctionDetail', (), {
'name': parsed['name'],
'arguments': arguments_str
}),
'id': tool_call_id,
'type': 'function'
})
]
})
# Print the tool call using the CLI utility
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 for Ollama
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 the tool is executed
)
except Exception as e:
logger.error(f"Error displaying tool call in CLI: {e}")
# Add to message history
tool_call_msg = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": tool_call_id,
"type": "function",
"function": {
"name": parsed['name'],
"arguments": arguments_str
}
}
]
}
streamed_tool_calls.append(tool_call_msg)
add_to_message_history(tool_call_msg)
logger.debug(f"Added function call: {parsed['name']} with args: {arguments_str}")
except Exception as e:
pass
function_call_starting_index = 0
if state.text_content_index_and_output:
function_call_starting_index += 1
@ -817,6 +1087,9 @@ 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)
# Add visual separation after agent output completes
print("\n")
# If we're not using rich streaming and not suppressing output, use old method
elif not self.suppress_final_output and final_response.output and any(isinstance(item, ResponseOutputMessage) for item in final_response.output):
# Find the assistant message to print
@ -837,8 +1110,22 @@ class OpenAIChatCompletionsModel(Model):
interaction_cost=interaction_cost,
total_cost=total_cost,
)
# Add visual separation after message
print("\n")
break
# --- Add assistant tool call(s) to message_history at the end of streaming ---
for tool_call_msg in streamed_tool_calls:
add_to_message_history(tool_call_msg)
# 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:
asst_msg = {
"role": "assistant",
"content": state.text_content_index_and_output[1].text
}
add_to_message_history(asst_msg)
if tracing.include_data():
span_generation.span_data.output = [final_response.model_dump()]
@ -964,36 +1251,69 @@ class OpenAIChatCompletionsModel(Model):
"extra_headers": _HEADERS,
}
# Error encountered: Error code: 400 - {'error': {'code': 'invalid_request_error',
# 'message': "'tool_choice' is only allowed when 'tools' are specified",
# 'type': 'invalid_request_error', 'param': None}}
#
# Only remove tool_choice if model starts with "gpt" and has no tools
if self.model.startswith("gpt") and not converted_tools:
kwargs.pop("tool_choice", None)
# Determine provider based on model string
model_str = str(self.model).lower()
# Provider-specific adjustments
if "/" in model_str:
# Handle provider/model format
provider = model_str.split("/")[0]
# TODO: review this. Remove tool_choice for Anthropic/Claude models when no tools are provided
if ("claude" in str(self.model).lower() or "anthropic" in str(self.model).lower()) and not converted_tools:
kwargs.pop("tool_choice", None)
# Model adjustments
if any(x in self.model for x in ["claude"]):
litellm.drop_params = True
# BadRequestError encountered: litellm.BadRequestError: AnthropicException -
# b'{"type":"error","error":
# {"type":"invalid_request_error","message":"store: Extra inputs are not permitted"}}'
#
kwargs.pop("store", None)
# Filter out NotGiven values to avoid JSON serialization issues
filtered_kwargs = {}
for key, value in kwargs.items():
if value is not NOT_GIVEN:
filtered_kwargs[key] = value
kwargs = filtered_kwargs
# Apply provider-specific configurations
if provider == "deepseek":
litellm.drop_params = True
kwargs.pop("parallel_tool_calls", None)
# Remove tool_choice if no tools are specified
if not converted_tools:
kwargs.pop("tool_choice", None)
elif provider == "claude":
litellm.drop_params = True
kwargs.pop("store", None)
# Remove tool_choice if no tools are specified
if not converted_tools:
kwargs.pop("tool_choice", None)
elif provider == "gemini":
kwargs.pop("parallel_tool_calls", None)
# Add any specific gemini settings if needed
else:
# Handle models without provider prefix
if "claude" in model_str:
litellm.drop_params = True
# Remove store parameter which isn't supported by Anthropic
kwargs.pop("store", None)
# Remove tool_choice if no tools are specified
if not converted_tools:
kwargs.pop("tool_choice", None)
elif "gemini" in model_str:
kwargs.pop("parallel_tool_calls", None)
elif "qwen" in model_str or ":" in model_str:
# Handle Ollama-served models with custom formats (e.g., qwen2.5:14b)
# These typically need the Ollama provider
litellm.drop_params = True
kwargs.pop("parallel_tool_calls", None)
# These models may not support certain parameters
if not converted_tools:
kwargs.pop("tool_choice", None)
# Don't add custom_llm_provider here to avoid duplication with Ollama provider
if self.is_ollama:
# Clean kwargs for ollama to avoid parameter conflicts
for param in ["custom_llm_provider"]:
kwargs.pop(param, None)
elif any(x in model_str for x in ["o1", "o3", "o4"]):
# Handle OpenAI reasoning models (o1, o3, o4)
kwargs.pop("parallel_tool_calls", None)
# Add reasoning effort if provided
if hasattr(model_settings, "reasoning_effort"):
kwargs["reasoning_effort"] = model_settings.reasoning_effort
# Filter out NotGiven values to avoid JSON serialization issues
filtered_kwargs = {}
for key, value in kwargs.items():
if value is not NOT_GIVEN:
filtered_kwargs[key] = value
kwargs = filtered_kwargs
try:
if self.is_ollama:
return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
@ -1003,21 +1323,71 @@ class OpenAIChatCompletionsModel(Model):
except litellm.exceptions.BadRequestError as e:
# 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
try:
return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
except litellm.exceptions.BadRequestError as e: # pylint: disable=W0621,C0301 # noqa: E501
#
# CTRL-C handler for ollama models
#
if "invalid message content type" in str(e):
kwargs["messages"] = fix_message_list(
kwargs["messages"])
model_str = str(self.model).lower()
provider = None
is_qwen = "qwen" in model_str or ":" in model_str
# Special handling for Qwen models
if is_qwen:
try:
# Use the specialized Qwen approach first
return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
except Exception as qwen_e:
print(qwen_e)
# If that fails, try our direct OpenAI approach
qwen_params = kwargs.copy()
qwen_params["api_base"] = get_ollama_api_base()
qwen_params["custom_llm_provider"] = "openai" # Use openai provider
# Make sure tools are passed
if "tools" in kwargs and kwargs["tools"]:
qwen_params["tools"] = kwargs["tools"]
if "tool_choice" in kwargs and kwargs["tool_choice"] is not NOT_GIVEN:
qwen_params["tool_choice"] = kwargs["tool_choice"]
try:
if stream:
# Streaming case
response = Response(
id=FAKE_RESPONSES_ID,
created_at=time.time(),
model=self.model,
object="response",
output=[],
tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else cast(Literal["auto", "required", "none"], tool_choice),
top_p=model_settings.top_p,
temperature=model_settings.temperature,
tools=[],
parallel_tool_calls=parallel_tool_calls or False,
)
stream_obj = await litellm.acompletion(**qwen_params)
return response, stream_obj
else:
# Non-streaming case
ret = litellm.completion(**qwen_params)
return ret
except Exception as direct_e:
# All approaches failed, log and raise the original error
print(f"All Qwen approaches failed. Original error: {str(e)}, Direct error: {str(direct_e)}")
raise e
# Try to detect provider from model string
if "/" in model_str:
provider = model_str.split("/")[0]
if provider:
# Add provider-specific settings based on detected provider
provider_kwargs = kwargs.copy()
if provider == "deepseek":
provider_kwargs["custom_llm_provider"] = "deepseek"
elif provider == "claude" or "claude" in model_str:
provider_kwargs["custom_llm_provider"] = "anthropic"
elif provider == "gemini":
provider_kwargs["custom_llm_provider"] = "gemini"
else:
raise e
# For unknown providers, try ollama as fallback
return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
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)): # noqa: E501 # pylint: disable=C0301
print(f"Error: {str(e)}")
@ -1121,77 +1491,165 @@ class OpenAIChatCompletionsModel(Model):
# Standard OpenAI handling for non-streaming
ret = litellm.completion(**kwargs)
return ret
async def _fetch_response_litellm_ollama(
self,
kwargs: dict,
model_settings: ModelSettings,
tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven,
stream: bool,
parallel_tool_calls: bool
parallel_tool_calls: bool,
provider="ollama"
) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
# Filter out parameters not supported by Ollama
# Extract only supported parameters for Ollama
ollama_supported_params = {
"model": kwargs["model"],
"messages": kwargs["messages"],
"temperature": kwargs["temperature"] if kwargs["temperature"] is not NOT_GIVEN else None,
"top_p": kwargs["top_p"] if kwargs["top_p"] is not NOT_GIVEN else None,
"max_tokens": kwargs["max_tokens"] if kwargs["max_tokens"] is not NOT_GIVEN else None,
"stream": kwargs["stream"],
"extra_headers": kwargs["extra_headers"]
"model": kwargs.get("model", ""),
"messages": kwargs.get("messages", []),
"stream": kwargs.get("stream", False)
}
# Modify the messages to remove system message for Ollama
if ollama_supported_params["messages"] and ollama_supported_params["messages"][0].get("role") == "system":
# Extract the system message
system_content = ollama_supported_params["messages"][0].get("content", "")
# Remove it from the messages
ollama_supported_params["messages"] = ollama_supported_params["messages"][1:]
# If there are user messages, prepend system to first user
if ollama_supported_params["messages"] and ollama_supported_params["messages"][0].get("role") == "user":
# Prepend the system instruction to the first user message, with a separator
user_content = ollama_supported_params["messages"][0].get("content", "")
if isinstance(user_content, str):
ollama_supported_params["messages"][0]["content"] = f"System: {system_content}\n\nUser: {user_content}"
# Add optional parameters if they exist and are not NOT_GIVEN
for param in ["temperature", "top_p", "max_tokens"]:
if param in kwargs and kwargs[param] is not NOT_GIVEN:
ollama_supported_params[param] = kwargs[param]
# Add extra headers if available
if "extra_headers" in kwargs:
ollama_supported_params["extra_headers"] = kwargs["extra_headers"]
# Add tools and tool_choice for compatibility with Qwen
if "tools" in kwargs and kwargs.get("tools") and kwargs.get("tools") is not NOT_GIVEN:
ollama_supported_params["tools"] = kwargs.get("tools")
if "tool_choice" in kwargs and kwargs.get("tool_choice") is not NOT_GIVEN:
ollama_supported_params["tool_choice"] = kwargs.get("tool_choice")
# Remove None values
ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None}
# Check if this is a Qwen model
model_str = str(self.model).lower()
is_qwen = "qwen" in model_str
api_base = get_ollama_api_base()
if "ollama" in provider:
api_base = api_base.rstrip('/v1')
# Create response object for streaming
if stream:
# For streaming with Ollama, we need to create a Response object first
response = Response(
id=FAKE_RESPONSES_ID,
created_at=time.time(),
model=self.model,
object="response",
output=[],
tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else cast(Literal["auto", "required", "none"], tool_choice),
tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else
cast(Literal["auto", "required", "none"], tool_choice),
top_p=model_settings.top_p,
temperature=model_settings.temperature,
tools=[],
parallel_tool_calls=parallel_tool_calls or False,
)
# Get the streaming object
# Get streaming response
stream_obj = await litellm.acompletion(
**ollama_kwargs,
api_base=get_ollama_api_base().rstrip('/v1'),
custom_llm_provider="ollama"
api_base=api_base,
custom_llm_provider=provider,
)
return response, stream_obj
else:
# Non-streaming mode
ret = litellm.completion(
# Get completion response
return litellm.completion(
**ollama_kwargs,
api_base=get_ollama_api_base().rstrip('/v1'),
custom_llm_provider="ollama"
api_base=api_base,
custom_llm_provider=provider,
)
return ret
def _get_client(self) -> AsyncOpenAI:
if self._client is None:
self._client = AsyncOpenAI()
return self._client
# Helper function to detect and format function calls from various models
def _detect_and_format_function_calls(self, delta):
"""
Helper to detect function calls in different formats and normalize them.
Handles Qwen specifics where function calls may be formatted differently.
Returns: List of normalized tool calls or None
"""
# Standard OpenAI-style tool_calls format
if hasattr(delta, 'tool_calls') and delta.tool_calls:
return delta.tool_calls
elif isinstance(delta, dict) and 'tool_calls' in delta and delta['tool_calls']:
return delta['tool_calls']
# Qwen/Ollama function_call format
if isinstance(delta, dict) and 'function_call' in delta:
function_call = delta['function_call']
return [{
'index': 0,
'id': f"call_{time.time_ns()}", # Generate a unique ID
'type': 'function',
'function': {
'name': function_call.get('name', ''),
'arguments': function_call.get('arguments', '')
}
}]
if isinstance(delta, dict) and 'content' in delta:
content = delta['content']
# Try to detect if the content is a JSON string with function call format
try:
if isinstance(content, str) and '{' in content and '}' in content:
# Try to extract JSON from the content (it might be embedded in text)
json_start = content.find('{')
json_end = content.rfind('}') + 1
if json_start >= 0 and json_end > json_start:
json_str = content[json_start:json_end]
parsed = json.loads(json_str)
if 'name' in parsed and 'arguments' in parsed:
# This looks like a function call in JSON format
return [{
'index': 0,
'id': f"call_{time.time_ns()}", # Generate a unique ID
'type': 'function',
'function': {
'name': parsed['name'],
'arguments': json.dumps(parsed['arguments']) if isinstance(parsed['arguments'], dict) else parsed['arguments']
}
}]
except Exception:
# If JSON parsing fails, just continue with normal processing
pass
# Anthropic-style tool_use format
if hasattr(delta, 'tool_use') and delta.tool_use:
tool_use = delta.tool_use
return [{
'index': 0,
'id': tool_use.get('id', f"tool_{time.time_ns()}"),
'type': 'function',
'function': {
'name': tool_use.get('name', ''),
'arguments': tool_use.get('input', '{}')
}
}]
elif isinstance(delta, dict) and 'tool_use' in delta and delta['tool_use']:
tool_use = delta['tool_use']
return [{
'index': 0,
'id': tool_use.get('id', f"tool_{time.time_ns()}"),
'type': 'function',
'function': {
'name': tool_use.get('name', ''),
'arguments': tool_use.get('input', '{}')
}
}]
return None
class _Converter:
@classmethod

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

@ -778,6 +778,7 @@ class Runner:
run_config: RunConfig,
tool_use_tracker: AgentToolUseTracker,
) -> SingleStepResult:
processed_response = RunImpl.process_model_response(
agent=agent,
all_tools=all_tools,
@ -785,6 +786,41 @@ class Runner:
output_schema=output_schema,
handoffs=handoffs,
)
# Log tools used with robust type checking
if hasattr(processed_response, 'tools_used') and processed_response.tools_used:
for i, tool_call in enumerate(processed_response.tools_used):
try:
# Safely extract tool name with multiple fallbacks
tool_name = "Unknown"
try:
if hasattr(tool_call, 'tool'):
if isinstance(tool_call.tool, str):
tool_name = tool_call.tool
elif hasattr(tool_call.tool, 'name'):
tool_name = tool_call.tool.name
else:
tool_name = str(tool_call.tool)
except Exception:
pass
# Safely extract call_id
call_id = "Unknown"
try:
if hasattr(tool_call, 'call_id'):
call_id = str(tool_call.call_id)
except Exception:
pass
# Safely extract parsed_args
parsed_args = "Unknown"
try:
if hasattr(tool_call, 'parsed_args'):
parsed_args = str(tool_call.parsed_args)
except Exception:
pass
except Exception:
pass
tool_use_tracker.add_tool_use(agent, processed_response.tools_used)

532
src/cai/tools/common.py
View File

@ -24,14 +24,70 @@ except ImportError:
# Global dictionary to store active sessions
ACTIVE_SESSIONS = {}
def _get_workspace_dir() -> str:
"""Determines the target workspace directory based on env vars for host."""
base_dir_env = os.getenv("CAI_WORKSPACE_DIR")
workspace_name = os.getenv("CAI_WORKSPACE")
# Determine the base directory
if base_dir_env:
base_dir = os.path.abspath(base_dir_env)
else: # Default base directory is 'workspaces'
if workspace_name:
base_dir = os.path.join(os.getcwd(), "workspaces")
else: # If no workspace name is set, the workspace IS the CWD.
return os.getcwd()
# If a workspace name is provided, append it to the base directory
if workspace_name:
if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name):
print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}'. "
f"Using directory '{base_dir}' instead.", fg="yellow"))
target_dir = base_dir
else:
target_dir = os.path.join(base_dir, workspace_name)
else:
target_dir = base_dir
# Ensure the final target directory exists on the host
try:
abs_target_dir = os.path.abspath(target_dir)
os.makedirs(abs_target_dir, exist_ok=True)
return abs_target_dir
except OSError as e:
print(color(f"Error creating/accessing host workspace directory '{abs_target_dir}': {e}",
fg="red"))
print(color(f"Falling back to current directory: {os.getcwd()}", fg="yellow"))
return os.getcwd()
def _get_container_workspace_path() -> str:
"""Determines the target workspace path inside the container."""
workspace_name = os.getenv("CAI_WORKSPACE")
if workspace_name:
if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name):
print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}' for container. "
f"Using '/workspace'.", fg="yellow"))
return "/"
# Standard path inside CAI containers
return f"/workspace/workspaces/{workspace_name}"
else:
return "/"
class ShellSession: # pylint: disable=too-many-instance-attributes
"""Class to manage interactive shell sessions"""
def __init__(self, command, session_id=None, ctf=None):
def __init__(self, command, session_id=None, ctf=None, workspace_dir=None, container_id=None): # noqa E501
self.session_id = session_id or str(uuid.uuid4())[:8]
self.command = command
self.command = command
self.ctf = ctf
self.container_id = container_id
# Determine workspace based on context (container, ctf or local host)
if self.container_id:
self.workspace_dir = _get_container_workspace_path()
elif self.ctf:
self.workspace_dir = workspace_dir or _get_workspace_dir()
else:
self.workspace_dir = _get_workspace_dir()
self.process = None
self.master = None
self.slave = None
@ -40,9 +96,40 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
self.last_activity = time.time()
def start(self):
"""Start the shell session"""
"""Start the shell session in the appropriate environment."""
start_message_cmd = self.command
# --- Start in Container ---
if self.container_id:
try:
self.master, self.slave = pty.openpty()
docker_cmd_list = [
"docker", "exec", "-i",
"-w", self.workspace_dir,
self.container_id,
"sh", "-c", # Use shell to handle complex commands if needed
self.command # The actual command to run
]
self.process = subprocess.Popen(
docker_cmd_list,
stdin=self.slave,
stdout=self.slave,
stderr=self.slave,
preexec_fn=os.setsid,
universal_newlines=True
)
self.is_running = True
self.output_buffer.append(
f"[Session {self.session_id}] Started in container {self.container_id[:12]}: "
f"{start_message_cmd} in {self.workspace_dir}")
threading.Thread(target=self._read_output, daemon=True).start()
except Exception as e:
self.output_buffer.append(f"Error starting container session: {str(e)}")
self.is_running = False
return
# --- Start in CTF ---
if self.ctf:
# For CTF environments
self.is_running = True
self.output_buffer.append(
f"[Session {
@ -52,81 +139,100 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
output = self.ctf.get_shell(self.command)
self.output_buffer.append(output)
except Exception as e: # pylint: disable=broad-except
self.output_buffer.append(f"Error: {str(e)}")
self.is_running = False
self.output_buffer.append(f"Error executing CTF command: {str(e)}")
self.is_running = False
return
# For local environment
# --- Start Locally (Host) ---
try:
# Create a pseudo-terminal
self.master, self.slave = pty.openpty()
# Start the process
self.process = subprocess.Popen( # pylint: disable=subprocess-popen-preexec-fn, consider-using-with # noqa: E501
self.command,
shell=True, # nosec B602
self.command,
shell=True, # nosec B602
stdin=self.slave,
stdout=self.slave,
stderr=self.slave,
preexec_fn=os.setsid, # Create a new process group
cwd=self.workspace_dir,
preexec_fn=os.setsid,
universal_newlines=True
)
self.is_running = True
self.output_buffer.append(
f"[Session {
self.session_id}] Started: {
self.command}")
# Start a thread to read output
threading.Thread(target=self._read_output, daemon=True).start()
except Exception as e: # pylint: disable=broad-except
self.output_buffer.append(f"Error starting session: {str(e)}")
self.output_buffer.append(f"Error starting local session: {str(e)}")
self.is_running = False
def _read_output(self):
"""Read output from the process"""
try:
while self.is_running:
while self.is_running and self.master is not None:
try:
# Check if process has exited before reading
if self.process and self.process.poll() is not None:
self.is_running = False
break
# Read the output
output = os.read(self.master, 1024).decode()
if output:
self.output_buffer.append(output)
self.last_activity = time.time()
except OSError:
# No data available or terminal closed
time.sleep(0.1)
if not self.is_process_running():
else:
self.is_running = False
break
except Exception as e: # pylint: disable=broad-except
self.output_buffer.append(f"Error reading output: {str(e)}")
except Exception as e:
self.output_buffer.append(f"Error reading output buffer: {str(read_err)}")
self.is_running = False
break
# Add a small sleep to prevent busy-waiting if no output
if is_process_running(self):
time.sleep(0.05)
except Exception as e:
self.output_buffer.append(f"Error in read_output loop: {str(e)}")
self.is_running = False
def is_process_running(self):
"""Check if the process is still running"""
# For CTF or container
if self.container_id or self.ctf:
return self.is_running
# For local host
if not self.process:
return False
return self.process.poll() is None
def send_input(self, input_data):
"""Send input to the process"""
if not self.is_running:
return "Session is not running"
"""Send input to the process (local or container)"""
if not self.is_running: # For CTF or container
if self.process and self.process.poll() is None:
self.is_running = True
else: # For local host
return "Session is not running"
try:
# --- Send to CTF ---
if self.ctf:
# For CTF environments
output = self.ctf.get_shell(input_data)
self.output_buffer.append(output)
return "Input sent to CTF session"
# For local environment
input_data = input_data.rstrip() + "\n"
os.write(self.master, input_data.encode())
self.last_activity = time.time()
return "Input sent to session"
# --- Send to Local or Container PTY ---
if self.master is not None:
input_data_bytes = (input_data.rstrip() + "\n").encode()
bytes_written = os.write(self.master, input_data_bytes)
if bytes_written != len(input_data_bytes):
self.output_buffer.append(f"[Session {self.session_id}] Warning: Partial input write.")
self.last_activity = time.time()
return "Input sent to session"
else:
return "Session PTY not available for input"
except Exception as e: # pylint: disable=broad-except
self.output_buffer.append(f"Error sending input: {str(e)}")
return f"Error sending input: {str(e)}"
def get_output(self, clear=True):
@ -138,8 +244,12 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
def terminate(self):
"""Terminate the session"""
session_id_short = self.session_id[:8]
if not self.is_running:
return "Session already terminated"
if self.process and self.process.poll() is None:
pass # Process is running, proceed with termination
else:
return f"Session {session_id_short} already terminated or finished."
try:
self.is_running = False
@ -147,28 +257,64 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
if self.process:
# Try to terminate the process group
try:
os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
except BaseException: # pylint: disable=bare-except,broad-except # noqa: E501
# If that fails, try to terminate just the process
self.process.terminate()
os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
except ProcessLookupError:
pass # Process already gone
except subprocess.TimeoutExpired:
print(color(f"Session {session_id_short} did not terminate gracefully, sending SIGKILL...", fg="yellow")) # noqa E501
try:
if pgid:
os.killpg(pgid, signal.SIGKILL) # Force kill
else:
self.process.kill()
except ProcessLookupError:
pass # Already gone
except Exception as kill_err:
termination_message = f" (Error during SIGKILL: {kill_err})"
except Exception as term_err: # Catch other errors during SIGTERM
termination_message = f" (Error during SIGTERM: {term_err})"
try:
self.process.kill()
except Exception: pass # Ignore nested errors
# Clean up resources
if self.master:
os.close(self.master)
if self.slave:
os.close(self.slave)
return f"Session {self.session_id} terminated"
# Final check
if self.process.poll() is None:
print(color(f"Session {session_id_short} process {self.process.pid} may still be running after termination attempts.", fg="red")) # noqa E501
termination_message += " (Warning: Process may still be running)"
# Clean up PTY resources if they exist
if self.master:
try: os.close(self.master)
except OSError: pass
self.master = None
if self.slave:
try: os.close(self.slave)
except OSError: pass
self.slave = None
return termination_message or f"Session {self.session_id} terminated"
except Exception as e: # pylint: disable=broad-except
return f"Error terminating session: {str(e)}"
return f"Error terminating session {session_id_short}: {str(e)}"
def create_shell_session(command, ctf=None):
"""Create a new shell session"""
session = ShellSession(command, ctf=ctf)
def create_shell_session(command, ctf=None, container_id=None, **kwargs):
"""Create a new shell session in the correct workspace/environment."""
if container_id:
session = ShellSession(command, ctf=ctf, container_id=container_id)
else:
workspace_dir = _get_workspace_dir()
session = ShellSession(command, ctf=ctf, workspace_dir=workspace_dir)
session.start()
ACTIVE_SESSIONS[session.session_id] = session
return session.session_id
if session.is_running or (ctf and not session.is_running):
ACTIVE_SESSIONS[session.session_id] = session
return session.session_id
else:
error_msg = session.get_output(clear=True)
print(color(f"Failed to start session: {error_msg}", fg="red"))
return f"Failed to start session: {error_msg}"
def list_shell_sessions():
@ -212,47 +358,100 @@ def get_session_output(session_id, clear=True):
def terminate_session(session_id):
"""Terminate a specific session"""
if session_id not in ACTIVE_SESSIONS:
return f"Session {session_id} not found"
return f"Session {session_id} not found or already terminated."
session = ACTIVE_SESSIONS[session_id]
result = session.terminate()
del ACTIVE_SESSIONS[session_id]
if session_id in ACTIVE_SESSIONS:
del ACTIVE_SESSIONS[session_id]
return result
def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None):
def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None):
"""Runs command in CTF env, changing to workspace_dir first."""
target_dir = workspace_dir or _get_workspace_dir()
full_command = f"cd '{target_dir}' && {command}"
original_cmd_for_msg = command # For logging
context_msg = f"(ctf:{target_dir})"
try:
# Ensure the command is executed in a shell that supports command
# chaining
output = ctf.get_shell(command, timeout=timeout)
# exploit_logger.log_ok()
output = ctf.get_shell(full_command, timeout=timeout)
if stdout:
print("\033[32m" + output + "\033[0m")
return output # output if output else result.stder
print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501
return output
except Exception as e: # pylint: disable=broad-except
print(color(f"Error executing CTF command: {e}", fg="red"))
# exploit_logger.log_error(str(e))
return f"Error executing CTF command: {str(e)}"
error_msg = f"Error executing CTF command '{original_cmd_for_msg}' in '{target_dir}': {e}" # noqa E501
print(color(error_msg, fg="red"))
return error_msg
def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None):
"""Runs command via SSH. Assumes SSH agent or passwordless setup unless sshpass is used externally.""" # noqa E501
ssh_user = os.environ.get('SSH_USER')
ssh_host = os.environ.get('SSH_HOST')
ssh_pass = os.environ.get('SSH_PASS')
remote_command = command
original_cmd_for_msg = command
context_msg = f"({ssh_user}@{ssh_host})"
# Construct base SSH command list
if ssh_pass:
ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", f"{ssh_user}@{ssh_host}"] # noqa E501
else:
ssh_cmd_list = ["ssh", f"{ssh_user}@{ssh_host}"]
ssh_cmd_list.append(remote_command)
try:
# Use subprocess.run with list of args for better security than shell=True
result = subprocess.run(
ssh_cmd_list,
capture_output=True,
text=True,
check=False, # Don't raise exception on non-zero exit code
timeout=timeout
)
output = result.stdout if result.stdout else result.stderr
if stdout:
print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501
# Return combined output, potentially including errors
return output.strip()
except subprocess.TimeoutExpired as e:
error_output = e.stdout if e.stdout else str(e)
timeout_msg = f"Timeout executing SSH command: {error_output}"
if stdout:
print(f"\033[33m{context_msg} $ {original_cmd_for_msg}\nTIMEOUT\n{error_output}\033[0m") # noqa E501
return timeout_msg
except FileNotFoundError:
# Handle case where ssh or sshpass isn't installed
error_msg = f"'sshpass' or 'ssh' command not found. Ensure they are installed and in PATH." # noqa E501
print(color(error_msg, fg="red"))
return error_msg
except Exception as e: # pylint: disable=broad-except
error_msg = f"Error executing SSH command '{original_cmd_for_msg}' on {ssh_host}: {e}" # noqa E501
print(color(error_msg, fg="red"))
return error_msg
def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None):
def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None):
"""Runs command locally in the specified workspace_dir."""
# If streaming is enabled and we have a call_id
if stream and call_id:
return _run_local_streamed(command, call_id, timeout, tool_name)
return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir)
target_dir = workspace_dir or _get_workspace_dir()
original_cmd_for_msg = command # For logging
context_msg = f"(local:{target_dir})"
try:
# nosec B602 - shell=True is required for command chaining
result = subprocess.run(
command,
shell=True, # nosec B602
capture_output=True,
text=True,
check=False,
timeout=timeout)
check=False,
timeout=timeout,
cwd=target_dir
)
output = result.stdout if result.stdout else result.stderr
if stdout:
print("\033[32m" + output + "\033[0m")
print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501
# Skip passing output to cli_print_tool_output when CAI_STREAM=true
# This prevents duplicate output in streaming mode
@ -261,20 +460,21 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
# Optional: Add cli_print_tool_output call here if needed for non-streaming
pass
return output
return output.strip()
except subprocess.TimeoutExpired as e:
error_output = e.stdout.decode() if e.stdout else str(e)
error_output = e.stdout if e.stdout else str(e)
if stdout:
print("\033[32m" + error_output + "\033[0m")
return error_output
return error_output
except Exception as e: # pylint: disable=broad-except
error_msg = f"Error executing local command: {e}"
print(color(error_msg, fg="red"))
return error_msg
error_msg = f"Error executing local command: {e}"
print(color(error_msg, fg="red"))
return error_msg
def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
"""Run a local command with streaming output to the Tool output panel"""
def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None):
"""Run a local command with streaming output to the Tool output panel."""
target_dir = workspace_dir or _get_workspace_dir()
try:
# Try to import Rich for nice display
try:
@ -298,7 +498,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
bufsize=1
bufsize=1,
cwd=target_dir # Set CWD for local process
)
# If tool_name is not provided, derive it from the command
@ -328,7 +529,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
header.append(")", style="yellow")
tool_time = 0
start_time = time.time()
total_time = time.time() - START_TIME
total_time = time.time() - START_TIME
timing_info = []
if total_time:
timing_info.append(f"Total: {format_time(total_time)}")
@ -505,7 +706,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
async_mode=False, session_id=None,
timeout=100, stream=False, call_id=None, tool_name=None):
"""
Run command either in CTF container or on the local attacker machine
Run command in the appropriate environment (Docker, CTF, SSH, Local)
and workspace.
Args:
command: The command to execute
@ -520,34 +722,168 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
If None, the tool name will be derived from the command.
Returns:
str: Command output, status message, or session ID
str: Command output, status message, or session ID.
"""
# If session_id is provided, send command to that session
if session_id:
if session_id not in ACTIVE_SESSIONS:
return f"Session {session_id} not found"
result = send_to_session(session_id, command)
session = ACTIVE_SESSIONS[session_id]
result = session.send_input(command) # Send the raw command string
if stdout:
output = get_session_output(session_id, clear=False)
print("\033[32m" + output + "\033[0m")
return result
# If async_mode, create a new session
if async_mode:
session_id = create_shell_session(command, ctf)
if stdout:
# Wait a moment for initial output
time.sleep(0.5)
output = get_session_output(session_id, clear=False)
print("\033[32m" + output + "\033[0m")
return f"Created session {session_id}. Use this ID to interact with the session."
env_type = "Local"
if session.container_id:
env_type = f"Container({session.container_id[:12]})"
elif session.ctf:
env_type = "CTF"
print(f"\033[32m(Session {session_id} in {env_type}:{session.workspace_dir}) >> {command}\n{output}\033[0m") # noqa E501
return result # Return the result of sending input ("Input sent..." or error)
# Generate a call_id if we're streaming and one wasn't provided
if stream and not call_id:
call_id = str(uuid.uuid4())[:8]
# Otherwise, run command normally
# 2. Determine Execution Environment (Container > CTF > SSH > Local)
active_container = os.getenv("CAI_ACTIVE_CONTAINER", "")
is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST'])
# --- Docker Container Execution ---
if active_container and not ctf and not is_ssh_env:
container_id = active_container
container_workspace = _get_container_workspace_path()
context_msg = f"(docker:{container_id[:12]}:{container_workspace})"
# Handle Async Session Creation in Container
if async_mode:
# Create a session specifically for the container environment
new_session_id = create_shell_session(command, container_id=container_id) # noqa E501
if "Failed" in new_session_id: # Check if session creation failed
return new_session_id
if stdout:
# Wait a moment for initial output
time.sleep(0.2)
output = get_session_output(new_session_id, clear=False)
print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501
return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # noqa E501
# Handle Streaming Container Execution - not yet implemented for containers
if stream:
# For now, display that streaming isn't supported for containers
from cai.util import cli_print_tool_output
if call_id and tool_name:
tool_args = {"command": command, "container": container_id[:12]}
cli_print_tool_output(
tool_name,
tool_args,
"Streaming not yet supported for container execution. Running normally...",
call_id=call_id
)
# Handle Synchronous Execution in Container
try:
# Ensure container workspace exists (best effort)
# Consider moving this to workspace set/container activation
mkdir_cmd = ["docker", "exec", container_id, "mkdir", "-p", container_workspace] # noqa E501
subprocess.run(mkdir_cmd, capture_output=True, text=True, check=False, timeout=10) # noqa E501
# Construct the docker exec command with workspace context
cmd_list = [
"docker", "exec",
"-w", container_workspace, # Set working directory
container_id,
"sh", "-c", command # Execute command via shell
]
result = subprocess.run(
cmd_list,
capture_output=True,
text=True,
check=False, # Don't raise exception on non-zero exit
timeout=timeout
)
output = result.stdout if result.stdout else result.stderr
output = output.strip() # Clean trailing newline
if stdout:
print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501
# Check if command failed specifically because container isn't running
if result.returncode != 0 and "is not running" in result.stderr:
print(color(f"{context_msg} Container is not running. Attempting execution on host instead.", fg="yellow")) # noqa E501
# Fallback to local execution, preserving workspace context
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501
return output # Return combined stdout/stderr
except subprocess.TimeoutExpired:
timeout_msg = "Timeout executing command in container."
if stdout:
print(f"\033[33m{context_msg} $ {command}\nTIMEOUT\033[0m") # noqa E501
print(color("Attempting execution on host instead.", fg="yellow"))
# Fallback to local execution on timeout
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501
except Exception as e: # pylint: disable=broad-except
error_msg = f"Error executing command in container: {str(e)}"
print(color(f"{context_msg} {error_msg}", fg="red"))
print(color("Attempting execution on host instead.", fg="yellow"))
# Fallback to local execution on other errors
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501
# --- CTF Execution ---
if ctf:
return _run_ctf(ctf, command, stdout, timeout, stream, call_id)
return _run_local(command, stdout, timeout, stream, call_id, tool_name)
# Handling streaming for CTF - not fully implemented yet
if stream:
from cai.util import cli_print_tool_output
if call_id and tool_name:
tool_args = {"command": command, "ctf": True}
cli_print_tool_output(
tool_name,
tool_args,
"Streaming not yet supported for CTF execution. Running normally...",
call_id=call_id
)
# _run_ctf handles workspace internally using _get_workspace_dir() default
return _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir
# --- SSH Execution ---
if is_ssh_env:
# Async for SSH would require session management via SSH client features
if async_mode:
return "Async mode not fully supported for SSH environment via this function yet."
# Handling streaming for SSH - not fully implemented yet
if stream:
from cai.util import cli_print_tool_output
if call_id and tool_name:
tool_args = {"command": command, "ssh": True}
cli_print_tool_output(
tool_name,
tool_args,
"Streaming not yet supported for SSH execution. Running normally...",
call_id=call_id
)
# _run_ssh handles command execution, workspace is relative to remote home
return _run_ssh(command, stdout, timeout) # Workspace dir less relevant here
# --- Local Execution (Default Fallback) ---
# Let _run_local handle determining the host workspace
# Handle Async Session Creation Locally
if async_mode:
# create_shell_session uses _get_workspace_dir() when container_id is None
new_session_id = create_shell_session(command)
if isinstance(new_session_id, str) and "Failed" in new_session_id: # Check failure
return new_session_id
# Retrieve the actual workspace dir the session is using
session = ACTIVE_SESSIONS.get(new_session_id)
actual_workspace = session.workspace_dir if session else "unknown"
if stdout:
time.sleep(0.2) # Allow session buffer to populate
output = get_session_output(new_session_id, clear=False)
print(f"\033[32m(Started Session {new_session_id} in local:{actual_workspace})\n{output}\033[0m")
return f"Started async session {new_session_id} locally. Use this ID to interact."
# Handle Synchronous Execution Locally using _run_local default with streaming support
return _run_local(command, stdout, timeout, stream, call_id, tool_name, None)

92
src/cai/util.py
View File

@ -268,87 +268,72 @@ def load_prompt_template(template_path):
except Exception as e:
raise ValueError(f"Failed to load template '{template_path}': {str(e)}")
# Start of Selection
def visualize_agent_graph(start_agent):
"""
Visualize agent graph showing all bidirectional connections between agents.
Uses Rich library for pretty printing.
"""
console = Console() # pylint: disable=redefined-outer-name
console = Console()
if start_agent is None:
console.print("[red]No agent provided to visualize.[/red]")
return
tree = Tree(
f"🤖 {
start_agent.name} (Current Agent)",
guide_style="bold blue")
tree = Tree(f"🤖 {start_agent.name} (Current Agent)", guide_style="bold blue")
# Track visited agents and their nodes to handle cross-connections
visited = {}
visited = set()
agent_nodes = {}
agent_positions = {} # Track positions in tree
position_counter = 0 # Counter for tracking positions
agent_positions = {}
position_counter = 0
def add_agent_node(agent, parent=None, is_transfer=False): # pylint: disable=too-many-branches # noqa: E501
"""Add agent node and track for cross-connections"""
def add_agent_node(agent, parent=None, is_transfer=False):
"""Add an agent node and track for cross-connections."""
nonlocal position_counter
if agent is None:
return None
aid = id(agent)
if aid in visited:
if is_transfer and parent:
original_pos = agent_positions.get(aid)
parent.add(f"[cyan]↩ Return to {agent.name} (Agent #{original_pos})[/cyan]")
return agent_nodes.get(aid)
# Create or get existing node for this agent
if id(agent) in visited:
if is_transfer:
# Add reference with position for repeated agents
original_pos = agent_positions[id(agent)]
parent.add(
f"[cyan]↩ Return to {
agent.name} (Top Level Agent #{original_pos})[/cyan]")
return agent_nodes[id(agent)]
visited[id(agent)] = True
visited.add(aid)
position_counter += 1
agent_positions[id(agent)] = position_counter
agent_positions[aid] = position_counter
# Create node for current agent
if is_transfer:
if is_transfer and parent:
node = parent
elif parent:
node = parent.add(f"[green]{agent.name} (#{position_counter})[/green]")
else:
node = parent.add(
f"[green]{agent.name} (#{position_counter})[/green]") if parent else tree # noqa: E501 pylint: disable=line-too-long
agent_nodes[id(agent)] = node
node = tree
agent_nodes[aid] = node
# Add tools as children
# Add tools
tools_node = node.add("[yellow]Tools[/yellow]")
for fn in getattr(agent, "functions", []):
if callable(fn):
fn_name = getattr(fn, "__name__", "")
if ("handoff" not in fn_name.lower() and
not fn_name.startswith("transfer_to")):
tools_node.add(f"[blue]{fn_name}[/blue]")
for tool in getattr(agent, "tools", []):
tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", "")
tools_node.add(f"[blue]{tool_name}[/blue]")
# Add Handoffs section
# Add handoffs
transfers_node = node.add("[magenta]Handoffs[/magenta]")
for handoff_fn in getattr(agent, "handoffs", []):
if callable(handoff_fn):
try:
next_agent = handoff_fn()
if next_agent:
transfer_node = transfers_node.add(f"🤖 {next_agent.name}")
add_agent_node(next_agent, transfer_node, True)
except Exception:
continue
# Process handoff functions
for fn in getattr(agent, "functions", []): # pylint: disable=too-many-nested-blocks # noqa: E501
if callable(fn):
fn_name = getattr(fn, "__name__", "")
if ("handoff" in fn_name.lower() or
fn_name.startswith("transfer_to")):
try:
next_agent = fn()
if next_agent:
# Show bidirectional connection
transfer = transfers_node.add(
f"🤖 {next_agent.name}") # noqa: E501
add_agent_node(next_agent, transfer, True)
except Exception: # nosec: B112 # pylint: disable=broad-exception-caught # noqa: E501
continue
return node
# Start recursive traversal from root agent
# Start traversal from the root agent
add_agent_node(start_agent)
console.print(tree)
# End of Selectio
def fix_litellm_transcription_annotations():
"""
@ -1005,6 +990,7 @@ def update_agent_streaming_content(context, text_delta):
# Force an update with the new panel
context["live"].update(updated_panel)
context["panel"] = updated_panel
context["live"].refresh()
def finish_agent_streaming(context, final_stats=None):
"""Finish the streaming session and display final stats if available."""