Fixes in replay and /load WIP

Signed-off-by: Víctor Mayoral Vilches <v.mayoralv@gmail.com>
This commit is contained in:
Víctor Mayoral Vilches 2025-05-09 16:02:58 +02:00
parent 65a015323c
commit c45685d78d
4 changed files with 72 additions and 51 deletions

View File

@ -38,7 +38,8 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li
config,
flush,
workspace,
virtualization
virtualization,
load
)
# Define helper functions

View File

@ -13,6 +13,7 @@ from typing import (
from rich.console import Console # pylint: disable=import-error
from cai.repl.commands.base import Command, register_command
from cai.sdk.agents.models.openai_chatcompletions import message_history
from cai.sdk.agents.run_to_jsonl import get_token_stats, load_history_from_jsonl
console = Console()
@ -56,17 +57,20 @@ class LoadCommand(Command):
jsonl_file = args[0]
# Try to load the jsonl file
try:
with open(jsonl_file, 'r') as f:
for line in f:
print(line)
messages = load_history_from_jsonl(jsonl_file)
console.print(f"[green]Jsonl file {jsonl_file} loaded[/green]")
except BaseException: # pylint: disable=broad-exception-caught
# If killing the process group fails, try killing just the
# process
console.print(f"[red]Error: Failed to load jsonl file {jsonl_file}[/red]")
# fetch longest message from jsonl file and send to message_history
# TODO @luijait
# fetch messages from JSONL file and add them to message_history
import pprint
pprint.pprint(messages)
# for message in messages:
# message_history.append(message)
except Exception as e: # pylint: disable=broad-exception-caught
console.print(f"[red]Error loading jsonl file: {str(e)}[/red]")

View File

@ -357,6 +357,7 @@ def load_history_from_jsonl(file_path):
list: A list of messages extracted from the JSONL file.
"""
messages = []
last_assistant_message = None
try:
with open(file_path, encoding='utf-8') as f:
@ -364,13 +365,26 @@ def load_history_from_jsonl(file_path):
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
# process assistant messages and keep the last one
# for additing it manually at the end
#
# NOTE: it might be the case that if the last message is of type "tool_message"
# we might be missing it. For that purpose, leaving here the corresponding code
# in case of need:
# 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
if record.get("event") == "assistant_message":
last_assistant_message = record.get("content")
# Extract messages from model record
if "model" in record and "messages" in record and isinstance(record["messages"], list):
# Store only complete conversation message objects
@ -379,12 +393,12 @@ def load_history_from_jsonl(file_path):
# 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]
@ -410,7 +424,7 @@ def load_history_from_jsonl(file_path):
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:
@ -418,7 +432,15 @@ def load_history_from_jsonl(file_path):
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)
# Add last message to the end of the list
if last_assistant_message:
unique_messages.append(
{
"role": "assistant",
"content": last_assistant_message
}
)
return unique_messages
@ -460,14 +482,23 @@ def get_token_stats(file_path):
else:
# Si cost es un valor directo
last_total_cost = float(record["cost"])
if "timing" in record:
if isinstance(record["timing"], dict):
last_active_time = record["timing"].get(
"active_seconds", 0.0)
last_idle_time = record["timing"].get(
"idle_seconds", 0.0)
if "timing_metrics" in record:
if isinstance(record["timing_metrics"], dict):
last_active_time = record["timing_metrics"].get(
"active_time_seconds", 0.0)
last_idle_time = record["timing_metrics"].get(
"idle_time_seconds", 0.0)
if "model" in record:
model_name = record["model"]
# Keep track of the last record for session_end event
if record.get("event") == "session_end":
if "timing_metrics" in record and isinstance(record["timing_metrics"], dict):
last_active_time = record["timing_metrics"].get(
"active_time_seconds", 0.0)
last_idle_time = record["timing_metrics"].get(
"idle_time_seconds", 0.0)
if "cost" in record and isinstance(record["cost"], dict):
last_total_cost = record["cost"].get("total_cost", 0.0)
except Exception as e: # pylint: disable=broad-except
print(f"Error loading line: {line}: {e}")
continue

View File

@ -132,7 +132,7 @@ def replay_conversation(messages: List[Dict], replay_delay: float = 0.5, usage:
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"):
@ -150,7 +150,7 @@ def replay_conversation(messages: List[Dict], replay_delay: float = 0.5, usage:
time.sleep(replay_delay)
role = message.get("role", "")
content = message.get("content", "")
content = message.get("content", "").strip()
sender = message.get("sender", role)
model = message.get("model", file_model)
@ -326,17 +326,7 @@ def main():
# 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)
@ -350,17 +340,19 @@ def main():
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)
print(usage)
# 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"))
@ -375,22 +367,16 @@ def main():
"""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)
))
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",
@ -400,7 +386,6 @@ def main():
}
display_execution_time(metrics)
except FileNotFoundError:
print(color(f"Error: File {jsonl_file_path} not found", fg="red"))
sys.exit(1)