From c45685d78d4b1ed99a9e3a4a08149f64a1103836 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20Mayoral=20Vilches?= Date: Fri, 9 May 2025 16:02:58 +0200 Subject: [PATCH] Fixes in replay and /load WIP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Víctor Mayoral Vilches --- src/cai/repl/commands/__init__.py | 3 +- src/cai/repl/commands/load.py | 14 +++++--- src/cai/sdk/agents/run_to_jsonl.py | 55 +++++++++++++++++++++++------- tools/replay.py | 51 ++++++++++----------------- 4 files changed, 72 insertions(+), 51 deletions(-) diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index bd2b471a..41bac88a 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -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 diff --git a/src/cai/repl/commands/load.py b/src/cai/repl/commands/load.py index 64ba0232..c78a8f9c 100644 --- a/src/cai/repl/commands/load.py +++ b/src/cai/repl/commands/load.py @@ -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]") diff --git a/src/cai/sdk/agents/run_to_jsonl.py b/src/cai/sdk/agents/run_to_jsonl.py index edee3efe..1be0520b 100644 --- a/src/cai/sdk/agents/run_to_jsonl.py +++ b/src/cai/sdk/agents/run_to_jsonl.py @@ -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 diff --git a/tools/replay.py b/tools/replay.py index 96e0f9a5..4ab42921 100644 --- a/tools/replay.py +++ b/tools/replay.py @@ -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)