mirror of https://github.com/aliasrobotics/cai.git
Add replay
This commit is contained in:
luijait
parent
36a85a3bdf
commit
f3c8d93ab7
|
|
@ -356,25 +356,70 @@ def load_history_from_jsonl(file_path):
|
|||
Returns:
|
||||
list: A list of messages extracted from the JSONL file.
|
||||
"""
|
||||
history = []
|
||||
max_length = 0
|
||||
with open(file_path, encoding='utf-8') as file:
|
||||
for line in file:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
record = json.loads(line)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
print(f"Error loading line: {line}")
|
||||
continue
|
||||
if isinstance(record, dict) and "messages" \
|
||||
in record and isinstance(
|
||||
record["messages"], list):
|
||||
if len(record["messages"]) > max_length:
|
||||
max_length = len(record["messages"])
|
||||
history = record["messages"]
|
||||
return history
|
||||
messages = []
|
||||
|
||||
try:
|
||||
with open(file_path, encoding='utf-8') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
record = json.loads(line)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
print(f"Error loading line: {line}")
|
||||
continue
|
||||
|
||||
# Extract messages from model record
|
||||
if "model" in record and "messages" in record and isinstance(record["messages"], list):
|
||||
# Store only complete conversation message objects
|
||||
for msg in record["messages"]:
|
||||
if "role" in msg:
|
||||
# Skip system messages
|
||||
if msg.get("role") == "system":
|
||||
continue
|
||||
|
||||
# Add this message if we haven't seen it already
|
||||
if not any(m.get("role") == msg.get("role") and
|
||||
m.get("content") == msg.get("content") for m in messages):
|
||||
messages.append(msg)
|
||||
|
||||
# Extract assistant messages and tool responses from model record choices
|
||||
elif "choices" in record and isinstance(record["choices"], list) and record["choices"]:
|
||||
choice = record["choices"][0]
|
||||
if "message" in choice and "role" in choice["message"]:
|
||||
msg = choice["message"]
|
||||
if not any(m.get("role") == msg.get("role") and
|
||||
m.get("content") == msg.get("content") for m in messages):
|
||||
messages.append(msg)
|
||||
|
||||
# Check for tool_calls in the message
|
||||
if msg.get("tool_calls"):
|
||||
for tool_call in msg.get("tool_calls", []):
|
||||
if tool_call.get("id") and "function" in tool_call:
|
||||
name = tool_call["function"].get("name", "")
|
||||
arguments = tool_call["function"].get("arguments", "")
|
||||
if name and arguments:
|
||||
# Add a placeholder tool message - will be filled later
|
||||
tool_message = {
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.get("id"),
|
||||
"content": ""
|
||||
}
|
||||
messages.append(tool_message)
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
print(f"Error loading history from {file_path}: {e}")
|
||||
|
||||
# Clean up duplicates and reorder
|
||||
unique_messages = []
|
||||
for msg in messages:
|
||||
if not any(m.get("role") == msg.get("role") and
|
||||
m.get("content") == msg.get("content") and
|
||||
m.get("tool_call_id", "") == msg.get("tool_call_id", "") for m in unique_messages):
|
||||
unique_messages.append(msg)
|
||||
|
||||
return unique_messages
|
||||
|
||||
|
||||
def get_token_stats(file_path):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,416 @@
|
|||
#!/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/jsonl_to_replay.py
|
||||
|
||||
Usage with asciinema rec, generating a .cast file and then converting it to a gif:
|
||||
asciinema rec --command="JSONL_FILE_PATH=\"/workspace/caiextensions-memory/caiextensions/memory/it/htb/challenges/insomnia/cai_20250307_114836.jsonl\" REPLAY_DELAY=\"0.5\" python3 tools/jsonl_to_replay.py" --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 json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
# 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 cai.util import (
|
||||
cli_print_agent_messages,
|
||||
cli_print_tool_output,
|
||||
color
|
||||
)
|
||||
from cai.sdk.agents.run_to_jsonl import get_token_stats, load_history_from_jsonl
|
||||
|
||||
# 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 replay_conversation(messages: List[Dict], replay_delay: float = 0.5, usage: Tuple = 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)
|
||||
"""
|
||||
turn_counter = 0
|
||||
interaction_counter = 0
|
||||
debug = 0 # Always set debug to 2
|
||||
|
||||
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"))
|
||||
|
||||
# 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"))
|
||||
|
||||
# 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]
|
||||
|
||||
for i, message in enumerate(messages):
|
||||
# Add delay between actions
|
||||
if i > 0:
|
||||
time.sleep(replay_delay)
|
||||
|
||||
role = message.get("role", "")
|
||||
content = message.get("content", "")
|
||||
sender = message.get("sender", role)
|
||||
model = message.get("model", file_model)
|
||||
|
||||
# Skip system messages
|
||||
if role == "system":
|
||||
continue
|
||||
|
||||
# Handle user messages
|
||||
if role == "user":
|
||||
# Use cli_print_agent_messages for user messages
|
||||
print(color(f"CAI> ", fg="cyan") + f"{content}")
|
||||
|
||||
turn_counter += 1
|
||||
interaction_counter = 0
|
||||
|
||||
# Handle assistant messages
|
||||
elif role == "assistant":
|
||||
# Check if there are tool calls
|
||||
tool_calls = message.get("tool_calls", [])
|
||||
tool_outputs = message.get("tool_outputs", {})
|
||||
|
||||
if tool_calls:
|
||||
# Print the assistant message with tool calls
|
||||
cli_print_agent_messages(
|
||||
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
|
||||
)
|
||||
|
||||
# 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]
|
||||
|
||||
# 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
|
||||
except json.JSONDecodeError:
|
||||
args_obj = arguments
|
||||
|
||||
# Print the tool call and output
|
||||
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
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Print regular assistant message
|
||||
cli_print_agent_messages(
|
||||
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
|
||||
)
|
||||
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
|
||||
}
|
||||
)
|
||||
|
||||
# Handle any other message types
|
||||
else:
|
||||
if content: # Only display if there's actual content
|
||||
cli_print_agent_messages(
|
||||
sender or role,
|
||||
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
|
||||
)
|
||||
|
||||
# Force flush stdout to ensure immediate printing
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to process JSONL files and generate replay output."""
|
||||
# Get environment variables
|
||||
jsonl_file_path = os.environ.get("JSONL_FILE_PATH")
|
||||
replay_delay = float(os.environ.get("REPLAY_DELAY", "0.5"))
|
||||
|
||||
# Validate environment variables
|
||||
if not jsonl_file_path:
|
||||
print(color("Error: JSONL_FILE_PATH environment variable is required",
|
||||
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
|
||||
full_data = load_jsonl(jsonl_file_path)
|
||||
|
||||
# Extract tool outputs from events and find last assistant message
|
||||
tool_outputs = {}
|
||||
last_assistant_message = None
|
||||
|
||||
for entry in full_data:
|
||||
if entry.get("event") == "tool_message":
|
||||
tool_call_id = entry.get("tool_call_id", "")
|
||||
content = entry.get("content", "")
|
||||
if tool_call_id and content:
|
||||
tool_outputs[tool_call_id] = content
|
||||
elif entry.get("event") == "assistant_message":
|
||||
last_assistant_message = entry
|
||||
|
||||
# Load the JSONL file for messages
|
||||
messages = load_history_from_jsonl(jsonl_file_path)
|
||||
|
||||
# Attach tool outputs to messages
|
||||
for message in messages:
|
||||
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"))
|
||||
|
||||
# Generate the replay with live printing
|
||||
replay_conversation(messages, replay_delay, usage)
|
||||
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"
|
||||
if seconds < 3600:
|
||||
minutes = seconds / 60
|
||||
return f"{minutes:.1f}m"
|
||||
hours = seconds / 3600
|
||||
return f"{hours:.1f}h"
|
||||
|
||||
if last_assistant_message:
|
||||
# Display the last assistant message in a panel
|
||||
console.print(Panel(
|
||||
last_assistant_message.get("content", "No content available"),
|
||||
title="[bold]Final Answer[/bold]",
|
||||
title_align="left",
|
||||
border_style="green",
|
||||
box=ROUNDED,
|
||||
padding=(1, 2)
|
||||
))
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue