mirror of https://github.com/aliasrobotics/cai.git
844 lines
35 KiB
Python
844 lines
35 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Tool to convert JSONL files to a replay format that simulates the CLI output.
|
|
This allows reviewing conversations in a more readable format.
|
|
|
|
Usage:
|
|
JSONL_FILE_PATH="path/to/file.jsonl" REPLAY_DELAY="0.5" python3 tools/replay.py
|
|
|
|
# Or using positional arguments:
|
|
python3 tools/replay.py path/to/file.jsonl 0.5
|
|
cai-replay path/to/file.jsonl 0.5
|
|
|
|
# Or using command line arguments:
|
|
python3 tools/replay.py --jsonl-file-path path/to/file.jsonl --replay-delay 0.5
|
|
|
|
Usage with asciinema rec, generating a .cast file and then converting it to a gif:
|
|
asciinema rec --command="python3 tools/replay.py path/to/file.jsonl 0.5" --overwrite
|
|
|
|
Or alternatively:
|
|
asciinema rec --command="JSONL_FILE_PATH='caiextensions-memory/caiextensions/memory/it/pentestperf/hackableii/hackableII_autonomo.jsonl' REPLAY_DELAY='0.05' cai-replay"
|
|
|
|
Then convert the .cast file to a gif:
|
|
agg /tmp/tmp6c4dxoac-ascii.cast demo.gif
|
|
|
|
Environment Variables:
|
|
JSONL_FILE_PATH: Path to the JSONL file containing conversation history (required)
|
|
REPLAY_DELAY: Time in seconds to wait between actions (default: 0.5)
|
|
"""
|
|
|
|
import re
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import argparse
|
|
from typing import Dict, List, Tuple
|
|
|
|
# Disable session recording for replay tool
|
|
os.environ["CAI_DISABLE_SESSION_RECORDING"] = "true"
|
|
|
|
# Add the parent directory to the path to import cai modules
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
from rich.box import ROUNDED
|
|
from rich.text import Text
|
|
from rich.console import Group
|
|
from rich.columns import Columns
|
|
from rich.rule import Rule
|
|
|
|
from cai.util import cli_print_agent_messages, cli_print_tool_output, color, COST_TRACKER
|
|
from cai.sdk.agents.run_to_jsonl import get_token_stats, load_history_from_jsonl
|
|
from cai.repl.ui.banner import display_banner
|
|
from collections import defaultdict
|
|
|
|
# Initialize console object for rich printing
|
|
console = Console()
|
|
|
|
|
|
# Create our own display_execution_time function that uses our local console
|
|
def display_execution_time(metrics=None):
|
|
"""Display the total execution time with our local console."""
|
|
if metrics is None:
|
|
return
|
|
|
|
# Create a panel for the execution time
|
|
content = []
|
|
content.append(f"Session Time: {metrics['session_time']}")
|
|
content.append(f"Active Time: {metrics['active_time']}")
|
|
content.append(f"Idle Time: {metrics['idle_time']}")
|
|
|
|
if metrics.get("llm_time") and metrics["llm_time"] != "0.0s":
|
|
content.append(
|
|
f"LLM Processing Time: [bold yellow]{metrics['llm_time']}[/bold yellow] "
|
|
f"[dim]({metrics['llm_percentage']:.1f}% of session)[/dim]"
|
|
)
|
|
|
|
time_panel = Panel(
|
|
Group(*[Text(line) for line in content]),
|
|
border_style="blue",
|
|
box=ROUNDED,
|
|
padding=(0, 1),
|
|
title="[bold]Session Statistics[/bold]",
|
|
title_align="left",
|
|
)
|
|
console.print(time_panel)
|
|
|
|
|
|
def load_jsonl(file_path: str) -> List[Dict]:
|
|
"""Load a JSONL file and return its contents as a list of dictionaries."""
|
|
data = []
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
if line.strip():
|
|
try:
|
|
data.append(json.loads(line))
|
|
except json.JSONDecodeError:
|
|
print(f"Warning: Skipping invalid JSON line: {line[:50]}...")
|
|
return data
|
|
|
|
|
|
def normalize_content(content) -> str:
|
|
"""
|
|
Normalize message content from various formats to a simple string.
|
|
|
|
Handles:
|
|
- Simple strings: return as-is
|
|
- List of content blocks: extract text from each block
|
|
- None: return empty string
|
|
"""
|
|
if content is None:
|
|
return ""
|
|
|
|
if isinstance(content, str):
|
|
return content.strip()
|
|
|
|
if isinstance(content, list):
|
|
text_parts = []
|
|
for item in content:
|
|
if isinstance(item, str):
|
|
text_parts.append(item)
|
|
elif isinstance(item, dict):
|
|
# Handle various content block types
|
|
if "text" in item:
|
|
text_parts.append(item["text"])
|
|
elif "content" in item:
|
|
text_parts.append(str(item["content"]))
|
|
return "\n".join(text_parts).strip() if text_parts else str(content)
|
|
|
|
return str(content).strip()
|
|
|
|
|
|
def detect_parallel_agents(messages: List[Dict]) -> Dict[str, str]:
|
|
"""
|
|
Detect parallel agents from messages by analyzing sender field patterns.
|
|
Returns a mapping of agent_id to agent_name.
|
|
"""
|
|
agents = {}
|
|
|
|
# Look for messages with sender field that follows parallel pattern
|
|
for msg in messages:
|
|
sender = msg.get("sender", "")
|
|
# Match patterns like "Bug Bounter [P1]", "Red Team Agent [P2]" etc
|
|
match = re.match(r"(.+?)\s*\[(P\d+)\]$", sender)
|
|
if match:
|
|
agent_name = match.group(1).strip()
|
|
agent_id = match.group(2)
|
|
agents[agent_id] = agent_name
|
|
|
|
return agents
|
|
|
|
|
|
def replay_conversation(
|
|
messages: List[Dict],
|
|
replay_delay: float = 0.5,
|
|
usage: Tuple = None,
|
|
jsonl_file_path: str = None,
|
|
full_data: List[Dict] = None,
|
|
) -> None:
|
|
"""
|
|
Replay a conversation from a list of messages, printing in real-time.
|
|
|
|
Args:
|
|
messages: List of message dictionaries
|
|
replay_delay: Time in seconds to wait between actions
|
|
usage: Tuple containing (model_name, total_input_tokens, total_output_tokens,
|
|
total_cost, active_time, idle_time)
|
|
jsonl_file_path: Path to the original JSONL file for graph display
|
|
full_data: Full JSONL data for additional metadata lookup
|
|
"""
|
|
turn_counter = 0
|
|
interaction_counter = 0
|
|
debug = 0 # Always set debug to 2
|
|
|
|
# Detect parallel agents
|
|
parallel_agents = detect_parallel_agents(messages)
|
|
is_parallel = len(parallel_agents) > 0
|
|
|
|
# Store messages for graph display
|
|
agent_messages = defaultdict(list)
|
|
|
|
# Create a mapping of timestamps to agent names from full_data
|
|
timestamp_to_agent = {}
|
|
if full_data:
|
|
for entry in full_data:
|
|
if entry.get("agent_name") and entry.get("timestamp_iso"):
|
|
timestamp_to_agent[entry["timestamp_iso"]] = entry["agent_name"]
|
|
|
|
if not messages:
|
|
print(color("No valid messages found in the JSONL file", fg="yellow"))
|
|
return
|
|
|
|
print(color(f"Replaying conversation with {len(messages)} messages...", fg="green"))
|
|
|
|
if is_parallel:
|
|
print(color(f"Detected {len(parallel_agents)} parallel agents:", fg="cyan"))
|
|
for agent_id, agent_name in sorted(parallel_agents.items()):
|
|
print(color(f" • {agent_name} [{agent_id}]", fg="cyan"))
|
|
|
|
# Extract the usage stats from the usage tuple
|
|
# Handle both old format (4 elements) and new format (6 elements with timing)
|
|
file_model = usage[0]
|
|
total_input_tokens = usage[1]
|
|
total_output_tokens = usage[2]
|
|
total_cost = usage[3]
|
|
|
|
# Check if timing information is available
|
|
active_time = usage[4] if len(usage) > 4 else 0
|
|
idle_time = usage[5] if len(usage) > 5 else 0
|
|
|
|
# Display timing information if available
|
|
if active_time > 0 or idle_time > 0:
|
|
print(color(f"Active time: {active_time:.2f}s", fg="cyan"))
|
|
print(color(f"Idle time: {idle_time:.2f}s", fg="cyan"))
|
|
|
|
print(color(f"Total cost: ${total_cost:.6f}", fg="cyan"))
|
|
|
|
# Initialize COST_TRACKER with the total cost from the JSONL file
|
|
COST_TRACKER.session_total_cost = total_cost
|
|
|
|
# First pass: Process all tool outputs
|
|
tool_outputs = {}
|
|
for idx, message in enumerate(messages):
|
|
if message.get("role") == "tool" and message.get("tool_call_id"):
|
|
tool_id = message.get("tool_call_id")
|
|
content = message.get("content", "")
|
|
tool_outputs[tool_id] = content
|
|
|
|
# Process assistant messages to match tool calls with outputs
|
|
for message in messages:
|
|
if message.get("role") == "assistant" and message.get("tool_calls"):
|
|
for tool_call in message.get("tool_calls", []):
|
|
call_id = tool_call.get("id", "")
|
|
if call_id in tool_outputs:
|
|
# Add this output to the tool_outputs of the assistant message
|
|
if "tool_outputs" not in message:
|
|
message["tool_outputs"] = {}
|
|
message["tool_outputs"][call_id] = tool_outputs[call_id]
|
|
|
|
# Process all messages, including the last one
|
|
total_messages = len(messages)
|
|
cumulative_cost = 0.0 # Track cumulative cost for progressive updates
|
|
|
|
for i, message in enumerate(messages):
|
|
try:
|
|
# Add delay between actions
|
|
if i > 0:
|
|
time.sleep(replay_delay)
|
|
|
|
role = message.get("role", "")
|
|
content = normalize_content(message.get("content"))
|
|
sender = message.get("sender", role)
|
|
model = message.get("model", file_model)
|
|
|
|
# Update COST_TRACKER with cumulative cost up to this message
|
|
# Calculate cost from tokens if interaction_cost not available
|
|
message_cost = message.get("interaction_cost", 0.0)
|
|
if message_cost == 0 and role == "assistant":
|
|
# Estimate cost from tokens (rough estimate: $5/M input, $15/M output)
|
|
input_tokens = message.get("input_tokens", 0)
|
|
output_tokens = message.get("output_tokens", 0)
|
|
if input_tokens > 0 or output_tokens > 0:
|
|
message_cost = (input_tokens * 0.000005) + (output_tokens * 0.000015)
|
|
if message_cost > 0:
|
|
cumulative_cost += message_cost
|
|
COST_TRACKER.current_agent_total_cost = cumulative_cost
|
|
COST_TRACKER.session_total_cost = cumulative_cost
|
|
|
|
# Skip system messages
|
|
if role == "system":
|
|
continue
|
|
|
|
# Store message for graph if parallel agents detected
|
|
if is_parallel:
|
|
# Determine agent for this message
|
|
if role == "assistant":
|
|
# Extract agent ID from sender if present
|
|
agent_match = re.match(r"(.+?)\s*\[(P\d+)\]$", sender)
|
|
if agent_match:
|
|
agent_id = agent_match.group(2)
|
|
agent_messages[agent_id].append(message)
|
|
elif role == "user":
|
|
# User messages go to all agents
|
|
for agent_id in parallel_agents:
|
|
agent_messages[agent_id].append(message)
|
|
elif role == "tool":
|
|
# Tool messages go to the agent that called them
|
|
# Look back for the assistant message that made this tool call
|
|
tool_call_id = message.get("tool_call_id")
|
|
for j in range(i - 1, -1, -1):
|
|
prev_msg = messages[j]
|
|
if prev_msg.get("role") == "assistant":
|
|
prev_sender = prev_msg.get("sender", "")
|
|
agent_match = re.match(r"(.+?)\s*\[(P\d+)\]$", prev_sender)
|
|
if agent_match:
|
|
agent_id = agent_match.group(2)
|
|
agent_messages[agent_id].append(message)
|
|
break
|
|
|
|
# Handle user messages
|
|
if role == "user":
|
|
print(color(f"CAI> ", fg="cyan") + f"{content}")
|
|
turn_counter += 1
|
|
# Don't reset interaction_counter to maintain numbering across user prompts
|
|
|
|
# Handle assistant messages
|
|
elif role == "assistant":
|
|
# Check if there are tool calls
|
|
tool_calls = message.get("tool_calls", [])
|
|
tool_outputs = message.get("tool_outputs", {})
|
|
|
|
# Extract the actual agent name
|
|
display_sender = sender
|
|
|
|
# First, check if we have agent_name in the message metadata
|
|
agent_name = message.get("agent_name")
|
|
if agent_name:
|
|
display_sender = agent_name
|
|
else:
|
|
# If still not found, try to extract from content patterns
|
|
if display_sender in ["assistant", role] and content:
|
|
# Look for patterns like "Agent: Bug Bounter >>" or "[0] Agent: Bug Bounter"
|
|
agent_match = re.search(
|
|
r"(?:\[\d+\]\s*)?Agent:\s*([^>]+?)(?:\s*>>|\s*\[|$)", content
|
|
)
|
|
if agent_match:
|
|
display_sender = agent_match.group(1).strip()
|
|
|
|
# If still "assistant", default to a generic name
|
|
if display_sender == "assistant" or display_sender == role:
|
|
display_sender = "Assistant"
|
|
|
|
if tool_calls:
|
|
# Only print the assistant message if there's actual content
|
|
# Skip empty panels when only tool_calls are present
|
|
if content and content.strip():
|
|
cli_print_agent_messages(
|
|
display_sender,
|
|
content,
|
|
interaction_counter,
|
|
model,
|
|
debug,
|
|
interaction_input_tokens=message.get("input_tokens", 0),
|
|
interaction_output_tokens=message.get("output_tokens", 0),
|
|
interaction_reasoning_tokens=message.get("reasoning_tokens", 0),
|
|
total_input_tokens=total_input_tokens,
|
|
total_output_tokens=total_output_tokens,
|
|
total_reasoning_tokens=message.get("total_reasoning_tokens", 0),
|
|
interaction_cost=message.get("interaction_cost", 0.0),
|
|
total_cost=total_cost,
|
|
cache_read_tokens=message.get("cache_read_tokens", 0),
|
|
cache_creation_tokens=message.get("cache_creation_tokens", 0),
|
|
)
|
|
|
|
# Print each tool call with its output
|
|
for tool_call in tool_calls:
|
|
function = tool_call.get("function", {})
|
|
name = function.get("name", "")
|
|
arguments = function.get("arguments", "{}")
|
|
call_id = tool_call.get("id", "")
|
|
|
|
# Get the tool output if available
|
|
tool_output = ""
|
|
if call_id and call_id in tool_outputs:
|
|
tool_output = tool_outputs[call_id]
|
|
# Detect placeholder messages for empty outputs
|
|
if tool_output.startswith("Tool response for call_"):
|
|
tool_output = "(Tool returned no output)"
|
|
|
|
# Skip empty tool calls
|
|
if not name:
|
|
continue
|
|
|
|
try:
|
|
# Try to parse arguments as JSON
|
|
if (
|
|
arguments
|
|
and isinstance(arguments, str)
|
|
and arguments.strip().startswith("{")
|
|
):
|
|
args_obj = json.loads(arguments)
|
|
else:
|
|
args_obj = arguments
|
|
|
|
# Special handling for execute_code to show full code
|
|
# Don't modify args_obj for execute_code, we'll handle display separately
|
|
except json.JSONDecodeError:
|
|
args_obj = arguments
|
|
|
|
# Special handling for execute_code to show the code
|
|
if (
|
|
name == "execute_code"
|
|
and isinstance(args_obj, dict)
|
|
and args_obj.get("code")
|
|
):
|
|
# Show execute_code with full code content
|
|
from rich.panel import Panel
|
|
from rich.syntax import Syntax
|
|
|
|
code = args_obj.get("code", "")
|
|
language = args_obj.get("language", "python")
|
|
filename = args_obj.get("filename", "exploit")
|
|
|
|
# Create syntax highlighted code
|
|
syntax = Syntax(code, language, theme="monokai", line_numbers=True)
|
|
|
|
# Create the panel with code
|
|
code_panel = Panel(
|
|
syntax,
|
|
title=f"[bold yellow]execute_code({filename}.{language})[/bold yellow]",
|
|
border_style="yellow",
|
|
padding=(0, 1),
|
|
)
|
|
console.print(code_panel)
|
|
|
|
# If there's output, show it too
|
|
if tool_output:
|
|
output_panel = Panel(
|
|
tool_output,
|
|
title="[bold green]Output[/bold green]",
|
|
border_style="green",
|
|
padding=(0, 1),
|
|
)
|
|
console.print(output_panel)
|
|
|
|
console.print() # Add spacing
|
|
else:
|
|
# Print other tool calls normally
|
|
cli_print_tool_output(
|
|
tool_name=name,
|
|
args=args_obj,
|
|
output=tool_output, # Use the matched tool output
|
|
call_id=call_id,
|
|
token_info={
|
|
"interaction_input_tokens": message.get("input_tokens", 0),
|
|
"interaction_output_tokens": message.get("output_tokens", 0),
|
|
"interaction_reasoning_tokens": message.get(
|
|
"reasoning_tokens", 0
|
|
),
|
|
"total_input_tokens": total_input_tokens,
|
|
"total_output_tokens": total_output_tokens,
|
|
"total_reasoning_tokens": message.get(
|
|
"total_reasoning_tokens", 0
|
|
),
|
|
"model": model,
|
|
"interaction_cost": message.get("interaction_cost", 0.0),
|
|
"total_cost": total_cost,
|
|
"agent_name": f"{display_sender} [P1]",
|
|
"cache_read_tokens": message.get("cache_read_tokens", 0),
|
|
"cache_creation_tokens": message.get("cache_creation_tokens", 0),
|
|
},
|
|
)
|
|
else:
|
|
# Print regular assistant message
|
|
cli_print_agent_messages(
|
|
display_sender,
|
|
content or "",
|
|
interaction_counter,
|
|
model,
|
|
debug,
|
|
interaction_input_tokens=message.get("input_tokens", 0),
|
|
interaction_output_tokens=message.get("output_tokens", 0),
|
|
interaction_reasoning_tokens=message.get("reasoning_tokens", 0),
|
|
total_input_tokens=total_input_tokens,
|
|
total_output_tokens=total_output_tokens,
|
|
total_reasoning_tokens=message.get("total_reasoning_tokens", 0),
|
|
interaction_cost=message.get("interaction_cost", 0.0),
|
|
total_cost=total_cost,
|
|
cache_read_tokens=message.get("cache_read_tokens", 0),
|
|
cache_creation_tokens=message.get("cache_creation_tokens", 0),
|
|
)
|
|
interaction_counter += 1 # iterate the interaction counter
|
|
|
|
# Handle tool messages - only those not already displayed with assistant messages
|
|
elif role == "tool":
|
|
# Check if we've already displayed this tool output with an assistant message
|
|
tool_call_id = message.get("tool_call_id", "")
|
|
|
|
# Skip tool messages that have been displayed with an assistant message
|
|
is_already_displayed = False
|
|
for prev_msg in messages[:i]:
|
|
if prev_msg.get("role") == "assistant" and tool_call_id in prev_msg.get(
|
|
"tool_outputs", {}
|
|
):
|
|
is_already_displayed = True
|
|
break
|
|
|
|
if not is_already_displayed and content: # Only show if there's actual content
|
|
tool_name = message.get("name", message.get("tool_call_id", "unknown"))
|
|
cli_print_tool_output(
|
|
tool_name=tool_name,
|
|
args="",
|
|
output=content,
|
|
token_info={
|
|
"interaction_input_tokens": message.get("input_tokens", 0),
|
|
"interaction_output_tokens": message.get("output_tokens", 0),
|
|
"interaction_reasoning_tokens": message.get("reasoning_tokens", 0),
|
|
"total_input_tokens": total_input_tokens,
|
|
"total_output_tokens": total_output_tokens,
|
|
"total_reasoning_tokens": message.get("total_reasoning_tokens", 0),
|
|
"model": model,
|
|
"interaction_cost": message.get("interaction_cost", 0.0),
|
|
"total_cost": total_cost,
|
|
"cache_read_tokens": message.get("cache_read_tokens", 0),
|
|
"cache_creation_tokens": message.get("cache_creation_tokens", 0),
|
|
},
|
|
)
|
|
|
|
# Handle any other message types (including final messages)
|
|
else:
|
|
# Always show the last message even if it seems empty
|
|
if content or (i == total_messages - 1 and role not in ["system", "tool"]):
|
|
cli_print_agent_messages(
|
|
sender or role,
|
|
content or "[Session ended]",
|
|
interaction_counter,
|
|
model,
|
|
debug,
|
|
interaction_input_tokens=message.get("input_tokens", 0),
|
|
interaction_output_tokens=message.get("output_tokens", 0),
|
|
interaction_reasoning_tokens=message.get("reasoning_tokens", 0),
|
|
total_input_tokens=total_input_tokens,
|
|
total_output_tokens=total_output_tokens,
|
|
total_reasoning_tokens=message.get("total_reasoning_tokens", 0),
|
|
interaction_cost=message.get("interaction_cost", 0.0),
|
|
total_cost=total_cost,
|
|
)
|
|
|
|
# Force flush stdout to ensure immediate printing
|
|
sys.stdout.flush()
|
|
|
|
except Exception as e:
|
|
# Handle any errors during message processing
|
|
print(color(f"Warning: Error processing message {i+1}: {str(e)}", fg="yellow"))
|
|
print(color("Continuing with next message...", fg="yellow"))
|
|
continue
|
|
|
|
# Display graph at the end if parallel agents detected
|
|
if is_parallel and agent_messages:
|
|
display_parallel_graph(agent_messages, parallel_agents)
|
|
|
|
|
|
def display_parallel_graph(
|
|
agent_messages: Dict[str, List[Dict]], parallel_agents: Dict[str, str]
|
|
) -> None:
|
|
"""Display a graph showing the parallel agent interactions."""
|
|
print("\n" + "=" * 80)
|
|
print(color("\n🎯 Parallel Agent Interaction Graph", fg="cyan", style="bold"))
|
|
print("=" * 80 + "\n")
|
|
|
|
graphs = []
|
|
|
|
for agent_id in sorted(parallel_agents.keys()):
|
|
agent_name = parallel_agents[agent_id]
|
|
messages = agent_messages.get(agent_id, [])
|
|
|
|
if not messages:
|
|
continue
|
|
|
|
# Build graph for this agent
|
|
graph_lines = []
|
|
turn_counter = 0
|
|
|
|
for i, msg in enumerate(messages):
|
|
role = msg.get("role", "")
|
|
content = msg.get("content", "")
|
|
|
|
if role == "user":
|
|
# User messages don't get turn numbers
|
|
if len(content) > 50:
|
|
content = content[:47] + "..."
|
|
graph_lines.append(f"[cyan]● User[/cyan]")
|
|
graph_lines.append(f" {content}")
|
|
elif role == "assistant":
|
|
turn_counter += 1
|
|
tool_calls = msg.get("tool_calls", [])
|
|
if tool_calls:
|
|
tools_str = ", ".join(
|
|
[tc.get("function", {}).get("name", "?") for tc in tool_calls[:3]]
|
|
)
|
|
if len(tool_calls) > 3:
|
|
tools_str += f" (+{len(tool_calls)-3})"
|
|
graph_lines.append(
|
|
f"[bold red][{turn_counter}][/bold red] [yellow]▶ Agent[/yellow]"
|
|
)
|
|
graph_lines.append(f" [dim]Tools: {tools_str}[/dim]")
|
|
else:
|
|
graph_lines.append(
|
|
f"[bold red][{turn_counter}][/bold red] [yellow]▶ Agent[/yellow]"
|
|
)
|
|
if content and len(content.strip()) > 0:
|
|
preview = content[:50] + "..." if len(content) > 50 else content
|
|
graph_lines.append(f" [dim]{preview}[/dim]")
|
|
elif role == "tool":
|
|
# Tool responses get the same turn number as their assistant
|
|
graph_lines.append(
|
|
f"[bold red][{turn_counter}][/bold red] [magenta]◆ Tool[/magenta]"
|
|
)
|
|
if content:
|
|
preview = content[:50] + "..." if len(content) > 50 else content
|
|
graph_lines.append(f" [dim]{preview}[/dim]")
|
|
|
|
if i < len(messages) - 1:
|
|
graph_lines.append(" ↓")
|
|
|
|
# Create panel for this agent
|
|
agent_panel = Panel(
|
|
"\n".join(graph_lines),
|
|
title=f"[bold cyan]{agent_name} [{agent_id}][/bold cyan]",
|
|
border_style="blue",
|
|
padding=(0, 1),
|
|
expand=False,
|
|
)
|
|
graphs.append(agent_panel)
|
|
|
|
# Display graphs in columns
|
|
if len(graphs) > 1:
|
|
console.print(Columns(graphs, equal=False, expand=False, padding=(1, 2)))
|
|
elif graphs:
|
|
console.print(graphs[0])
|
|
|
|
# Print summary
|
|
console.print("\n[bold]Summary:[/bold]")
|
|
total_messages = sum(len(msgs) for msgs in agent_messages.values())
|
|
unique_user_messages = len(
|
|
set(
|
|
msg.get("content", "")
|
|
for msgs in agent_messages.values()
|
|
for msg in msgs
|
|
if msg.get("role") == "user"
|
|
)
|
|
)
|
|
|
|
console.print(f"• Total agents: {len(parallel_agents)}")
|
|
console.print(f"• Total messages: {total_messages}")
|
|
console.print(f"• User messages: {unique_user_messages}")
|
|
console.print(
|
|
f"• Average messages per agent: {total_messages / len(parallel_agents) if parallel_agents else 0:.1f}"
|
|
)
|
|
print("\n" + "=" * 80)
|
|
|
|
|
|
def parse_arguments():
|
|
"""Parse command line arguments."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Tool to convert JSONL files to a replay format that simulates the CLI output.",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Using environment variables:
|
|
JSONL_FILE_PATH="path/to/file.jsonl" REPLAY_DELAY="0.5" python3 tools/replay.py
|
|
|
|
# Using positional arguments:
|
|
python3 tools/replay.py path/to/file.jsonl 0.5
|
|
cai-replay path/to/file.jsonl 0.5
|
|
|
|
# Using command line arguments:
|
|
python3 tools/replay.py --jsonl-file-path path/to/file.jsonl --replay-delay 0.5
|
|
|
|
# Using positional argument for file only:
|
|
python3 tools/replay.py path/to/file.jsonl --replay-delay 0.5
|
|
|
|
# With asciinema:
|
|
asciinema rec --command="python3 tools/replay.py path/to/file.jsonl 0.5" --overwrite
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"jsonl_file",
|
|
nargs="?",
|
|
default=None,
|
|
help="Path to the JSONL file containing conversation history",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"replay_delay_pos",
|
|
nargs="?",
|
|
type=float,
|
|
default=None,
|
|
help="Time in seconds to wait between actions (positional argument)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--jsonl-file-path", type=str, help="Path to the JSONL file containing conversation history"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--replay-delay",
|
|
type=float,
|
|
default=0.5,
|
|
help="Time in seconds to wait between actions (default: 0.5)",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
"""Main function to process JSONL files and generate replay output."""
|
|
# Display banner
|
|
display_banner(console)
|
|
print("\n")
|
|
|
|
# Parse command line arguments
|
|
args = parse_arguments()
|
|
|
|
# Get environment variables or command line arguments
|
|
# First check for --jsonl-file-path, then positional argument, then environment variable
|
|
jsonl_file_path = args.jsonl_file_path or args.jsonl_file or os.environ.get("JSONL_FILE_PATH")
|
|
|
|
# For replay delay, prioritize: positional arg > --replay-delay > environment variable > default
|
|
if args.replay_delay_pos is not None:
|
|
replay_delay = args.replay_delay_pos
|
|
elif args.replay_delay != 0.5: # Check if --replay-delay was explicitly set
|
|
replay_delay = args.replay_delay
|
|
else:
|
|
replay_delay = float(os.environ.get("REPLAY_DELAY", "0.5"))
|
|
|
|
# Validate required parameters
|
|
if not jsonl_file_path:
|
|
print(
|
|
color(
|
|
"Error: JSONL file path is required. Use a positional argument, --jsonl-file-path option, or set JSONL_FILE_PATH environment variable.",
|
|
fg="red",
|
|
)
|
|
)
|
|
sys.exit(1)
|
|
|
|
print(color(f"Loading JSONL file: {jsonl_file_path}", fg="blue"))
|
|
|
|
try:
|
|
# Load the full JSONL file to extract tool outputs and agent names
|
|
full_data = load_jsonl(jsonl_file_path)
|
|
|
|
# Extract tool outputs from events and find last assistant message
|
|
tool_outputs = {}
|
|
agent_names = {} # Store agent names by timestamp or other identifier
|
|
|
|
# Extract agent names from full data
|
|
current_agent_name = None
|
|
for entry in full_data:
|
|
# Track the current agent name from various events
|
|
if entry.get("agent_name"):
|
|
current_agent_name = entry.get("agent_name")
|
|
# Store agent name with timestamp or other identifier
|
|
timestamp = entry.get("timestamp")
|
|
if timestamp:
|
|
agent_names[timestamp] = entry.get("agent_name")
|
|
|
|
# Also look for agent_run_start events which contain agent names
|
|
if entry.get("event") == "agent_run_start" and entry.get("agent_name"):
|
|
current_agent_name = entry.get("agent_name")
|
|
|
|
# Load the JSONL file for messages
|
|
messages = load_history_from_jsonl(jsonl_file_path)
|
|
|
|
# Attach tool outputs and agent names to messages
|
|
# Also track current agent for messages without timestamps
|
|
last_known_agent = current_agent_name
|
|
|
|
for i, message in enumerate(messages):
|
|
# Try to match agent names by timestamp
|
|
msg_timestamp = message.get("timestamp")
|
|
if msg_timestamp and msg_timestamp in agent_names:
|
|
message["agent_name"] = agent_names[msg_timestamp]
|
|
last_known_agent = agent_names[msg_timestamp]
|
|
elif (
|
|
message.get("role") == "assistant"
|
|
and not message.get("agent_name")
|
|
and last_known_agent
|
|
):
|
|
# If no timestamp match but we have a last known agent, use it
|
|
message["agent_name"] = last_known_agent
|
|
|
|
if message.get("role") == "assistant" and message.get("tool_calls"):
|
|
if "tool_outputs" not in message:
|
|
message["tool_outputs"] = {}
|
|
|
|
for tool_call in message.get("tool_calls", []):
|
|
call_id = tool_call.get("id", "")
|
|
if call_id in tool_outputs:
|
|
message["tool_outputs"][call_id] = tool_outputs[call_id]
|
|
|
|
print(color(f"Loaded {len(messages)} messages from JSONL file", fg="blue"))
|
|
|
|
# Get token stats and cost from the JSONL file
|
|
usage = get_token_stats(jsonl_file_path)
|
|
|
|
# Display timing information if available (new format)
|
|
if len(usage) > 4:
|
|
print(color(f"Active time: {usage[4]:.2f}s", fg="blue"))
|
|
print(color(f"Idle time: {usage[5]:.2f}s", fg="blue"))
|
|
|
|
# Pass full_data to replay_conversation for agent name lookup
|
|
replay_conversation(messages, replay_delay, usage, jsonl_file_path, full_data)
|
|
print(color("Replay completed successfully", fg="green"))
|
|
|
|
# Display the total cost
|
|
active_time = usage[4] if len(usage) > 4 else 0
|
|
idle_time = usage[5] if len(usage) > 5 else 0
|
|
total_time = active_time + idle_time
|
|
|
|
# Format time values as strings with units
|
|
def format_time(seconds):
|
|
"""Format time in seconds to a human-readable string."""
|
|
if seconds < 60:
|
|
return f"{seconds:.1f}s"
|
|
else:
|
|
# Convert seconds to hours, minutes, seconds
|
|
hours, remainder = divmod(seconds, 3600)
|
|
minutes, seconds = divmod(remainder, 60)
|
|
|
|
if hours > 0:
|
|
return f"{int(hours)}h {int(minutes)}m {int(seconds)}s"
|
|
else:
|
|
return f"{int(minutes)}m {int(seconds)}s"
|
|
|
|
metrics = {
|
|
"session_time": format_time(total_time),
|
|
"llm_time": "0.0s",
|
|
"llm_percentage": 0,
|
|
"active_time": format_time(active_time),
|
|
"idle_time": format_time(idle_time),
|
|
}
|
|
display_execution_time(metrics)
|
|
|
|
except FileNotFoundError:
|
|
print(color(f"Error: File {jsonl_file_path} not found", fg="red"))
|
|
sys.exit(1)
|
|
except json.JSONDecodeError:
|
|
print(color(f"Error: Invalid JSON in {jsonl_file_path}", fg="red"))
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(color(f"Error: {str(e)}", fg="red"))
|
|
sys.exit(1)
|
|
finally:
|
|
# Clean up the environment variable to avoid polluting other processes
|
|
if "CAI_DISABLE_SESSION_RECORDING" in os.environ:
|
|
del os.environ["CAI_DISABLE_SESSION_RECORDING"]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|