From 7b016e058ee72635b7b3d5b923aabf09462f0ed8 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 8 May 2025 12:38:57 +0200 Subject: [PATCH] Logging and timming from 0.3.0 --- .gitignore | 1 + src/cai/cli.py | 75 ++- .../agents/models/openai_chatcompletions.py | 118 ++++- src/cai/sdk/agents/run_to_jsonl.py | 461 ++++++++++++++++ src/cai/tools/common.py | 490 +++++++++++------- src/cai/util.py | 169 +++++- 6 files changed, 1103 insertions(+), 211 deletions(-) create mode 100644 src/cai/sdk/agents/run_to_jsonl.py diff --git a/.gitignore b/.gitignore index f0416924..7a723dbf 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,7 @@ lib/ lib64/ parts/ sdist/ +logs/ var/ wheels/ share/python-wheels/ diff --git a/src/cai/cli.py b/src/cai/cli.py index 19154e78..32e2f891 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -110,6 +110,8 @@ from rich.console import Console import asyncio from cai.util import fix_litellm_transcription_annotations, color, calculate_model_cost from cai.util import create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming +from cai.sdk.agents.run_to_jsonl import get_session_recorder +from cai.util import start_idle_timer, stop_idle_timer, start_active_timer, stop_active_timer # Import modules from cai.repl from cai.repl.commands import FuzzyCommandCompleter, handle_command as commands_handle_command @@ -184,6 +186,13 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Setup session logging history_file = setup_session_logging() + + # Initialize session logger and display the filename + session_logger = get_session_recorder() + # Use rich Text instead of wasabi color() for styling with underline + from rich.text import Text + log_text = Text(f"\nLog file: {session_logger.filename}", style="yellow underline") + console.print(log_text) # Display banner display_banner(console) @@ -209,6 +218,9 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= while turn_count < max_turns: try: + # Start measuring user idle time + start_idle_timer() + idle_start_time = time.time() # Check if model has changed and update if needed @@ -253,6 +265,10 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= current_text ) idle_time += time.time() - idle_start_time + + # Stop measuring user idle time and start measuring active time + stop_idle_timer() + start_active_timer() except KeyboardInterrupt: def format_time(seconds): @@ -263,21 +279,35 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= Total = time.time() - START_TIME idle_time += time.time() - idle_start_time try: - active_time = Total - idle_time + # Get more accurate active and idle time measurements from the timer functions + from cai.util import get_active_time_seconds, get_idle_time_seconds, COST_TRACKER + + # Use the precise measurements from our timers + active_time_seconds = get_active_time_seconds() + idle_time_seconds = get_idle_time_seconds() + + # Format for display + active_time_formatted = format_time(active_time_seconds) + idle_time_formatted = format_time(idle_time_seconds) + + # Get session cost from the global cost tracker + session_cost = COST_TRACKER.session_total_cost metrics = { "session_time": format_time(Total), - "active_time": format_time(active_time), - "idle_time": format_time(idle_time), - "llm_time": "0.0s", # Placeholder, update if available - "llm_percentage": 0.0, # Placeholder, update if available + "active_time": active_time_formatted, + "idle_time": idle_time_formatted, + "llm_time": format_time(active_time_seconds), # Using active time as LLM time + "llm_percentage": round((active_time_seconds / Total) * 100, 1) if Total > 0 else 0.0, + "session_cost": f"${session_cost:.6f}" # Add formatted session cost } - logging_path = None # Set this if you have a log file path + logging_path = session_logger.filename if hasattr(session_logger, 'filename') else None content = [] content.append(f"Session Time: {metrics['session_time']}") - content.append(f"Active Time: {metrics['active_time']}") + content.append(f"Active Time: {metrics['active_time']} ({metrics['llm_percentage']}%)") content.append(f"Idle Time: {metrics['idle_time']}") + content.append(f"Total Session Cost: {metrics['session_cost']}") # Add cost to display if logging_path: content.append(f"Log available at: {logging_path}") @@ -290,8 +320,21 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= from rich.box import ROUNDED from rich.console import Group + # Create Rich Text objects for each line + text_content = [] + for i, line in enumerate(content): + if "Total Session Cost" in line: + # Format cost line with special styling + cost_text = Text() + parts = line.split(":") + cost_text.append(parts[0] + ":", style="bold") + cost_text.append(parts[1], style="bold green") + text_content.append(cost_text) + else: + text_content.append(Text(line)) + time_panel = Panel( - Group(*[Text(line) for line in content]), + Group(*text_content), border_style="blue", box=ROUNDED, padding=(0, 1), @@ -301,6 +344,14 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= console.print(time_panel) print_session_summary(console, metrics, logging_path) + + # Log session end + if session_logger: + session_logger.log_session_end() + + # Prevent duplicate cost display from the COST_TRACKER exit handler + os.environ["CAI_COST_DISPLAYED"] = "true" + except Exception: pass break @@ -346,6 +397,10 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Use non-streamed response response = asyncio.run(Runner.run(agent, user_input)) turn_count += 1 + + # Stop measuring active time and start measuring idle time again + stop_active_timer() + start_idle_timer() except KeyboardInterrupt: # No need to clean up streaming context as model handles it @@ -358,6 +413,10 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= filename, line, func, text = tb_info[-1] console.print(f"[bold red]Error: {str(e)}[/bold red]") console.print(f"[bold red]Traceback: {tb_info}[/bold red]") + + # Make sure we switch back to idle mode even if there's an error + stop_active_timer() + start_idle_timer() def main(): # Apply litellm patch to fix the __annotations__ error diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 6b5c956f..279bd36c 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -15,7 +15,9 @@ from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming +from cai.util import start_idle_timer, stop_idle_timer, start_active_timer, stop_active_timer from wasabi import color +from cai.sdk.agents.run_to_jsonl import get_session_recorder from openai import NOT_GIVEN, AsyncOpenAI, AsyncStream, NotGiven from openai.types import ChatModel @@ -73,6 +75,22 @@ class InputTokensDetails(BaseModel): """The number of prompt tokens.""" cached_tokens: int = 0 """The number of cached tokens.""" + +# Custom ResponseUsage that makes prompt_tokens/input_tokens and completion_tokens/output_tokens compatible +class CustomResponseUsage(ResponseUsage): + """ + Custom ResponseUsage class that provides compatibility between different field naming conventions. + Works with both input_tokens/output_tokens and prompt_tokens/completion_tokens. + """ + @property + def prompt_tokens(self) -> int: + """Alias for input_tokens to maintain compatibility""" + return self.input_tokens + + @property + def completion_tokens(self) -> int: + """Alias for output_tokens to maintain compatibility""" + return self.output_tokens from .. import _debug from ..agent_output import AgentOutputSchema @@ -227,12 +245,16 @@ class OpenAIChatCompletionsModel(Model): self.total_input_tokens = 0 self.total_output_tokens = 0 self.total_reasoning_tokens = 0 + self.total_cost = 0.0 self.agent_name = "Agent" # Default name # Flags for CLI integration self.disable_rich_streaming = False # Prevents creating a rich panel in the model self.suppress_final_output = False # Prevents duplicate output at end of streaming + # Initialize the session logger + self.logger = get_session_recorder() + def set_agent_name(self, name: str) -> None: """Set the agent name for CLI display purposes.""" self.agent_name = name @@ -253,6 +275,10 @@ class OpenAIChatCompletionsModel(Model): # Increment the interaction counter for CLI display self.interaction_counter += 1 + # Stop idle timer and start active timer to track LLM processing time + stop_idle_timer() + start_active_timer() + with generation_span( model=str(self.model), model_config=dataclasses.asdict(model_settings) @@ -285,6 +311,8 @@ class OpenAIChatCompletionsModel(Model): "content": input } add_to_message_history(user_msg) + # Log the user message + self.logger.log_user_message(input) elif isinstance(input, list): for item in input: # Try to extract user messages @@ -295,6 +323,9 @@ class OpenAIChatCompletionsModel(Model): "content": item.get("content", "") } add_to_message_history(user_msg) + # Log the user message + if item.get("content"): + self.logger.log_user_message(item.get("content")) # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -416,6 +447,19 @@ class OpenAIChatCompletionsModel(Model): } add_to_message_history(tool_call_msg) + + # Log the assistant tool call message + tool_calls_list = [] + for tool_call in assistant_msg.tool_calls: + tool_calls_list.append({ + "id": tool_call.id, + "type": tool_call.type, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments + } + }) + self.logger.log_assistant_message(None, tool_calls_list) # If the assistant message is just text, add it as well elif hasattr(assistant_msg, "content") and assistant_msg.content: asst_msg = { @@ -423,6 +467,21 @@ class OpenAIChatCompletionsModel(Model): "content": assistant_msg.content } add_to_message_history(asst_msg) + # Log the assistant message + self.logger.log_assistant_message(assistant_msg.content) + + # Log the complete response for the session + self.logger.rec_training_data( + { + "model": str(self.model), + "messages": converted_messages, + "stream": False, + "tools": [t.params_json_schema for t in tools] if tools else [], + "tool_choice": model_settings.tool_choice + }, + response, + self.total_cost + ) usage = ( Usage( @@ -442,12 +501,25 @@ class OpenAIChatCompletionsModel(Model): } items = _Converter.message_to_output_items(response.choices[0].message) + + # For non-streaming responses, make sure we also log token usage with compatible field names + # This ensures both streaming and non-streaming use consistent naming + if not hasattr(response, 'usage'): + response.usage = {} + if hasattr(response.usage, 'prompt_tokens') and not hasattr(response.usage, 'input_tokens'): + response.usage.input_tokens = response.usage.prompt_tokens + if hasattr(response.usage, 'completion_tokens') and not hasattr(response.usage, 'output_tokens'): + response.usage.output_tokens = response.usage.completion_tokens return ModelResponse( output=items, usage=usage, referenceable_id=None, ) + + # Stop active timer and start idle timer when response is complete + stop_active_timer() + start_idle_timer() async def stream_response( self, @@ -465,6 +537,10 @@ class OpenAIChatCompletionsModel(Model): # Increment the interaction counter for CLI display self.interaction_counter += 1 + # Stop idle timer and start active timer to track LLM processing time + stop_idle_timer() + start_active_timer() + # Check if streaming should be shown in rich panel should_show_rich_stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' and not self.disable_rich_streaming @@ -508,6 +584,8 @@ class OpenAIChatCompletionsModel(Model): "content": input } add_to_message_history(user_msg) + # Log the user message + self.logger.log_user_message(input) elif isinstance(input, list): for item in input: if isinstance(item, dict): @@ -517,6 +595,9 @@ class OpenAIChatCompletionsModel(Model): "content": item.get("content", "") } add_to_message_history(user_msg) + # Log the user message + if item.get("content"): + self.logger.log_user_message(item.get("content")) # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -1017,7 +1098,7 @@ class OpenAIChatCompletionsModel(Model): output_tokens = usage.completion_tokens # Create a proper usage object with our token counts - final_response.usage = ResponseUsage( + final_response.usage = CustomResponseUsage( input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens, @@ -1069,6 +1150,8 @@ class OpenAIChatCompletionsModel(Model): interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0.00001) total_cost = max(float(total_cost if total_cost is not None else 0.0), 0.00001) + # Store the total cost for future recording + self.total_cost = total_cost # Create final stats with explicit type conversion for all values final_stats = { @@ -1110,6 +1193,14 @@ class OpenAIChatCompletionsModel(Model): # --- 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) + + # Log the assistant tool call message if any tool calls were collected + if streamed_tool_calls: + tool_calls_list = [] + for tool_call_msg in streamed_tool_calls: + for tool_call in tool_call_msg.get("tool_calls", []): + tool_calls_list.append(tool_call) + self.logger.log_assistant_message(None, tool_calls_list) # 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 = { @@ -1117,6 +1208,26 @@ class OpenAIChatCompletionsModel(Model): "content": state.text_content_index_and_output[1].text } add_to_message_history(asst_msg) + # Log the assistant message + self.logger.log_assistant_message(state.text_content_index_and_output[1].text) + + # Log the complete response + self.logger.rec_training_data( + { + "model": str(self.model), + "messages": converted_messages, + "stream": True, + "tools": [t.params_json_schema for t in tools] if tools else [], + "tool_choice": model_settings.tool_choice + }, + final_response, + self.total_cost + ) + + # Stop active timer and start idle timer when streaming is complete + stop_active_timer() + start_idle_timer() + except Exception as e: # Ensure streaming context is cleaned up in case of errors if streaming_context: @@ -1124,6 +1235,11 @@ class OpenAIChatCompletionsModel(Model): finish_agent_streaming(streaming_context, None) except Exception: pass + + # Stop active timer and start idle timer when streaming errors out + stop_active_timer() + start_idle_timer() + raise e @overload diff --git a/src/cai/sdk/agents/run_to_jsonl.py b/src/cai/sdk/agents/run_to_jsonl.py new file mode 100644 index 00000000..91d4b76c --- /dev/null +++ b/src/cai/sdk/agents/run_to_jsonl.py @@ -0,0 +1,461 @@ +""" +Data recorder +""" + +import os # pylint: disable=import-error +from datetime import datetime +import json +import socket +import urllib.request +import getpass +import platform +from urllib.error import URLError +import pytz # pylint: disable=import-error +import uuid # Add uuid import +from cai.util import get_active_time, get_idle_time +import time +import requests +import atexit + +# Global recorder instance for session-wide logging +_session_recorder = None + + +def get_session_recorder(workspace_name=None): + """ + Get the global session recorder instance. + If one doesn't exist, it will be created. + + Args: + workspace_name (str | None): Optional workspace name. + + Returns: + DataRecorder: The session recorder instance. + """ + global _session_recorder + if _session_recorder is None: + _session_recorder = DataRecorder(workspace_name) + return _session_recorder + + +class DataRecorder: # pylint: disable=too-few-public-methods + """ + Records training data from litellm.completion + calls in OpenAI-like JSON format. + + Stores both input messages and completion + responses during execution in a single JSONL file. + """ + + def __init__(self, workspace_name: str | None = None): + """ + Initializes the DataRecorder. + + Args: + workspace_name (str | None): The name of the current workspace. + """ + # Generate a session ID that will be used for the entire session + self.session_id = str(uuid.uuid4()) + + # Track the last message to ensure it's logged + self.last_assistant_message = None + self.last_assistant_tool_calls = None + self._last_message_logged = False + self._session_end_logged = False + + log_dir = 'logs' + os.makedirs(log_dir, exist_ok=True) + + # Get current username + try: + username = getpass.getuser() + except Exception: # pylint: disable=broad-except + username = "unknown" + + # Get operating system and version information + try: + os_name = platform.system().lower() + os_version = platform.release() + os_info = f"{os_name}_{os_version}" + except Exception: # pylint: disable=broad-except + os_info = "unknown_os" + + # Check internet connection and get public IP + public_ip = "127.0.0.1" + try: + # Quick connection check with minimal traffic + socket.create_connection(("1.1.1.1", 53), timeout=1) + + # If connected, try to get public IP + try: + # Using a simple and lightweight service + with urllib.request.urlopen( # nosec: B310 + "https://api.ipify.org", + timeout=2 + ) as response: + public_ip = response.read().decode('utf-8') + except (URLError, socket.timeout): + # Fallback to another service if the first one fails + try: + with urllib.request.urlopen( # nosec: B310 + "https://ifconfig.me", + timeout=2 + ) as response: + public_ip = response.read().decode('utf-8') + except (URLError, socket.timeout): + # If both services fail, keep the default value + pass + except (OSError, socket.timeout, socket.gaierror): + # No internet connection, keep the default value + pass + + # Create filename with username, OS info, and IP + timestamp = datetime.now().astimezone( + pytz.timezone("Europe/Madrid")).strftime("%Y%m%d_%H%M%S") + base_filename = f'cai_{self.session_id}_{timestamp}_{username}_{os_info}_{public_ip.replace(".", "_")}.jsonl' + + if workspace_name: + self.filename = os.path.join( + log_dir, f'{workspace_name}_{base_filename}' + ) + else: + self.filename = os.path.join(log_dir, base_filename) + + # Inicializar el coste total acumulado + self.total_cost = 0.0 + + # Log the session start + with open(self.filename, 'a', encoding='utf-8') as f: + session_start = { + "event": "session_start", + "timestamp": datetime.now().astimezone( + pytz.timezone("Europe/Madrid")).isoformat(), + "session_id": self.session_id + } + json.dump(session_start, f) + f.write('\n') + + def rec_training_data(self, create_params, msg, total_cost=None) -> None: + """ + Records a single training data entry to the JSONL file + + Args: + create_params: Parameters used for the LLM call + msg: Response from the LLM + total_cost: Optional total accumulated cost from CAI instance + """ + request_data = { + "model": create_params["model"], + "messages": create_params["messages"], + "stream": create_params["stream"] + } + if "tools" in create_params: + request_data.update({ + "tools": create_params["tools"], + "tool_choice": create_params["tool_choice"], + }) + + # Obtener el coste de la interacción + interaction_cost = 0.0 + if hasattr(msg, "cost"): + interaction_cost = float(msg.cost) + + # Usar el total_cost proporcionado o actualizar el interno + if total_cost is not None: + self.total_cost = float(total_cost) + else: + self.total_cost += interaction_cost + + # Get timing metrics (without units, just numeric values) + active_time_str = get_active_time() + idle_time_str = get_idle_time() + + # Convert string time to seconds for storage + def time_str_to_seconds(time_str): + if "h" in time_str: + parts = time_str.split() + hours = float(parts[0].replace("h", "")) + minutes = float(parts[1].replace("m", "")) + seconds = float(parts[2].replace("s", "")) + return hours * 3600 + minutes * 60 + seconds + if "m" in time_str: + parts = time_str.split() + minutes = float(parts[0].replace("m", "")) + seconds = float(parts[1].replace("s", "")) + return minutes * 60 + seconds + return float(time_str.replace("s", "")) + + active_time_seconds = time_str_to_seconds(active_time_str) + idle_time_seconds = time_str_to_seconds(idle_time_str) + + # Get token usage from the usage object - handle both field names + prompt_tokens = 0 + completion_tokens = 0 + total_tokens = 0 + + if hasattr(msg, "usage"): + # Try input_tokens first (ResponseUsage) + if hasattr(msg.usage, "input_tokens"): + prompt_tokens = msg.usage.input_tokens + # Fall back to prompt_tokens (ChatCompletion) + elif hasattr(msg.usage, "prompt_tokens"): + prompt_tokens = msg.usage.prompt_tokens + + # Try output_tokens first (ResponseUsage) + if hasattr(msg.usage, "output_tokens"): + completion_tokens = msg.usage.output_tokens + # Fall back to completion_tokens (ChatCompletion) + elif hasattr(msg.usage, "completion_tokens"): + completion_tokens = msg.usage.completion_tokens + + # Get total tokens - calculate if not available + if hasattr(msg.usage, "total_tokens"): + total_tokens = msg.usage.total_tokens + else: + total_tokens = prompt_tokens + completion_tokens + + completion_data = { + "id": msg.id, + "object": "chat.completion", + "created": int(datetime.now().timestamp()), + "model": msg.model, + "messages": [ + { + "role": m.role, + "content": m.content, + "tool_calls": [t.model_dump() for t in (m.tool_calls or [])] # pylint: disable=line-too-long # noqa: E501 + } + for m in msg.messages + ] if hasattr(msg, "messages") else [], + "choices": [{ + "index": 0, + "message": { + "role": msg.choices[0].message.role if hasattr(msg, "choices") and msg.choices else "assistant", + "content": msg.choices[0].message.content if hasattr(msg, "choices") and msg.choices else None, + "tool_calls": [t.model_dump() for t in (msg.choices[0].message.tool_calls or [])] if hasattr(msg, "choices") and msg.choices else [] # pylint: disable=line-too-long # noqa: E501 + }, + "finish_reason": msg.choices[0].finish_reason if hasattr(msg, "choices") and msg.choices else "stop" + }], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens + }, + "cost": { + "interaction_cost": interaction_cost, + "total_cost": self.total_cost + }, + "timing": { + "active_seconds": active_time_seconds, + "idle_seconds": idle_time_seconds + }, + "timestamp_iso": datetime.now().astimezone( + pytz.timezone("Europe/Madrid")).isoformat() + } + + # Append both request and completion to the instance's jsonl file + with open(self.filename, 'a', encoding='utf-8') as f: + json.dump(request_data, f) + f.write('\n') + json.dump(completion_data, f) + f.write('\n') + + def log_user_message(self, user_message): + """ + Logs a user message to the JSONL file. + + Args: + user_message: The message from the user to log + """ + with open(self.filename, 'a', encoding='utf-8') as f: + user_data = { + "event": "user_message", + "timestamp": datetime.now().astimezone( + pytz.timezone("Europe/Madrid")).isoformat(), + "content": user_message + } + json.dump(user_data, f) + f.write('\n') + + def log_assistant_message(self, assistant_message, tool_calls=None): + """ + Logs an assistant message to the JSONL file. + + Args: + assistant_message: The message from the assistant to log + tool_calls: Optional tool calls included in the assistant message + """ + # Store the last message in case we need to log it at exit + self.last_assistant_message = assistant_message + self.last_assistant_tool_calls = tool_calls + + with open(self.filename, 'a', encoding='utf-8') as f: + assistant_data = { + "event": "assistant_message", + "timestamp": datetime.now().astimezone( + pytz.timezone("Europe/Madrid")).isoformat(), + "content": assistant_message + } + if tool_calls: + assistant_data["tool_calls"] = tool_calls + json.dump(assistant_data, f) + f.write('\n') + + # Mark that the message has been logged + self._last_message_logged = True + + def log_session_end(self): + """ + Logs the end of the session to the JSONL file. + Includes timing metrics from active/idle time tracking. + """ + # Set a flag to indicate we've already logged the session end + self._session_end_logged = True + + try: + from cai.util import get_active_time_seconds, get_idle_time_seconds, COST_TRACKER + active_time = get_active_time_seconds() + idle_time = get_idle_time_seconds() + # Get the global session cost from COST_TRACKER + session_cost = COST_TRACKER.session_total_cost + except ImportError: + active_time = 0.0 + idle_time = 0.0 + session_cost = self.total_cost + + with open(self.filename, 'a', encoding='utf-8') as f: + session_end = { + "event": "session_end", + "timestamp": datetime.now().astimezone( + pytz.timezone("Europe/Madrid")).isoformat(), + "session_id": self.session_id, + "timing_metrics": { + "active_time_seconds": active_time, + "idle_time_seconds": idle_time, + "total_time_seconds": active_time + idle_time, + "active_percentage": round((active_time / (active_time + idle_time)) * 100, 2) if (active_time + idle_time) > 0 else 0.0 + }, + "cost": { + "total_cost": session_cost # Use the global session cost + } + } + json.dump(session_end, f) + f.write('\n') + + +def load_history_from_jsonl(file_path): + """ + Load conversation history from a JSONL file and + return it as a list of messages. + + Args: + file_path (str): The path to the JSONL file. + NOTE: file_path assumes it's either relative to the + current directory or absolute. + + 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 + + +def get_token_stats(file_path): + """ + Get token usage statistics from a JSONL file. + + Args: + file_path (str): Path to the JSONL file + + Returns: + tuple: (model_name, total_prompt_tokens, total_completion_tokens, + total_cost, active_time, idle_time) + """ + total_prompt_tokens = 0 + total_completion_tokens = 0 + total_cost = 0.0 + model_name = None + last_total_cost = 0.0 + last_active_time = 0.0 + last_idle_time = 0.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) + if "usage" in record: + total_prompt_tokens += record["usage"]["prompt_tokens"] + total_completion_tokens += ( + record["usage"]["completion_tokens"] + ) + if "cost" in record: + if isinstance(record["cost"], dict): + # Si cost es un diccionario, obtener total_cost + last_total_cost = record["cost"].get("total_cost", 0.0) + 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 "model" in record: + model_name = record["model"] + except Exception as e: # pylint: disable=broad-except + print(f"Error loading line: {line}: {e}") + continue + + # Usar el último total_cost encontrado como el total + total_cost = last_total_cost + + return (model_name, total_prompt_tokens, total_completion_tokens, + total_cost, last_active_time, last_idle_time) + +def atexit_handler(): + """ + Ensure session_end is logged when the program exits. + Only logs if a session recorder exists and session_end hasn't already been logged. + """ + global _session_recorder + if _session_recorder is None: + return + + # Check if we have an unlogged assistant message and log it + if hasattr(_session_recorder, 'last_assistant_message') and not getattr(_session_recorder, '_last_message_logged', False): + if _session_recorder.last_assistant_message or _session_recorder.last_assistant_tool_calls: + _session_recorder.log_assistant_message( + _session_recorder.last_assistant_message, + _session_recorder.last_assistant_tool_calls + ) + + # Check if we've already logged the session end (via KeyboardInterrupt) + if getattr(_session_recorder, '_session_end_logged', False): + return + + # Log the session end + _session_recorder.log_session_end() + +# Register the exit handler +atexit.register(atexit_handler) \ No newline at end of file diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index f687dded..3b2bc259 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -13,7 +13,7 @@ import uuid import sys import shlex from wasabi import color # pylint: disable=import-error -from cai.util import format_time +from cai.util import format_time, start_active_timer, stop_active_timer, start_idle_timer, stop_idle_timer # Instead of direct import @@ -440,44 +440,57 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t if stream and call_id: 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})" + # Make sure we're in active time mode for tool execution + stop_idle_timer() + start_active_timer() + try: - result = subprocess.run( - command, - shell=True, # nosec B602 - capture_output=True, - text=True, - check=False, - timeout=timeout, - cwd=target_dir - ) - 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 - - # Skip passing output to cli_print_tool_output when CAI_STREAM=true - # This prevents duplicate output in streaming mode - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if not is_streaming_enabled: - # Optional: Add cli_print_tool_output call here if needed for non-streaming - pass - - return output.strip() - except subprocess.TimeoutExpired as e: - error_output = e.stdout if e.stdout else str(e) - if stdout: - print("\033[32m" + error_output + "\033[0m") - 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 + target_dir = workspace_dir or _get_workspace_dir() + original_cmd_for_msg = command # For logging + context_msg = f"(local:{target_dir})" + try: + result = subprocess.run( + command, + shell=True, # nosec B602 + capture_output=True, + text=True, + check=False, + timeout=timeout, + cwd=target_dir + ) + 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 + + # Skip passing output to cli_print_tool_output when CAI_STREAM=true + # This prevents duplicate output in streaming mode + is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + if not is_streaming_enabled: + # Optional: Add cli_print_tool_output call here if needed for non-streaming + pass + + return output.strip() + except subprocess.TimeoutExpired as e: + error_output = e.stdout if e.stdout else str(e) + if stdout: + print("\033[32m" + error_output + "\033[0m") + 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 + finally: + # Always switch back to idle mode when function completes + stop_active_timer() + start_idle_timer() 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.""" + # Make sure we're in active time mode for tool execution + stop_idle_timer() + start_active_timer() + target_dir = workspace_dir or _get_workspace_dir() try: # Try to import Rich for nice display @@ -708,6 +721,10 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace cli_print_tool_output(tool_name, tool_args, error_msg, call_id=call_id) return error_msg + finally: + # Always switch back to idle mode when function completes + stop_active_timer() + start_idle_timer() def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arguments # noqa: E501 @@ -732,185 +749,256 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg Returns: 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" - 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) - 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] - - # 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 + # Use the active timer during tool execution + stop_idle_timer() + start_active_timer() + + try: + # If session_id is provided, send command to that session + if session_id: + if session_id not in ACTIVE_SESSIONS: + # Switch back to idle mode before returning error + stop_active_timer() + start_idle_timer() + return f"Session {session_id} not found" + session = ACTIVE_SESSIONS[session_id] + result = session.send_input(command) # Send the raw command string 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 + output = get_session_output(session_id, clear=False) + 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 + + # For async sessions, we don't switch back to idle mode here + # since the session continues to run in the background + if not async_mode: + # Switch back to idle mode after synchronous command completes + stop_active_timer() + start_idle_timer() + + return result # Return the result of sending input ("Input sent..." or error) - # Handle Streaming Container Execution - if stream: - # Ensure workspace directory exists inside the container first - mkdir_cmd = [ - "docker", "exec", container_id, - "mkdir", "-p", container_workspace - ] - subprocess.run( - mkdir_cmd, - capture_output=True, - text=True, - check=False, - timeout=10 - ) + # 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] - # Build docker exec command as a single shell string for streaming - docker_exec_cmd = ( - "docker exec -w " - f"{shlex.quote(container_workspace)} " - f"{shlex.quote(container_id)} sh -c " - f"{shlex.quote(command)}" - ) + # 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']) - # Re-use the local streaming helper to provide real-time output - return _run_local_streamed( - docker_exec_cmd, - call_id, - timeout, - tool_name, - workspace_dir=_get_workspace_dir() - ) + # --- 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 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 + # 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 + # Switch back to idle mode before returning error + stop_active_timer() + start_idle_timer() + 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 + + # For async sessions, switch back to idle mode after session creation + stop_active_timer() + start_idle_timer() + return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # 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 - ) + # Handle Streaming Container Execution + if stream: + # Ensure workspace directory exists inside the container first + mkdir_cmd = [ + "docker", "exec", container_id, + "mkdir", "-p", container_workspace + ] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False, + timeout=10 + ) - output = result.stdout if result.stdout else result.stderr - output = output.strip() # Clean trailing newline + # Build docker exec command as a single shell string for streaming + docker_exec_cmd = ( + "docker exec -w " + f"{shlex.quote(container_workspace)} " + f"{shlex.quote(container_id)} sh -c " + f"{shlex.quote(command)}" + ) - if stdout: - print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501 + # Re-use the local streaming helper to provide real-time output + result = _run_local_streamed( + docker_exec_cmd, + call_id, + timeout, + tool_name, + workspace_dir=_get_workspace_dir() + ) + + # Switch back to idle mode after streaming command completes + stop_active_timer() + start_idle_timer() + return result - # 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 + # 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 + # Switch back to idle mode before fallback execution + stop_active_timer() + start_idle_timer() + # Fallback to local execution, preserving workspace context + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + # Switch back to idle mode after command completes + stop_active_timer() + start_idle_timer() + 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")) + # Switch back to idle mode before fallback execution + stop_active_timer() + start_idle_timer() + # 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")) + # Switch back to idle mode before fallback execution + stop_active_timer() + start_idle_timer() + # Fallback to local execution on other errors return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 - return output # Return combined stdout/stderr + # --- CTF Execution --- + if ctf: + # 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 + result = _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir + + # Switch back to idle mode after CTF command completes + stop_active_timer() + start_idle_timer() + return result - 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 + # --- SSH Execution --- + if is_ssh_env: + # Async for SSH would require session management via SSH client features + if async_mode: + # Switch back to idle mode before returning message + stop_active_timer() + start_idle_timer() + 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 + result = _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + + # Switch back to idle mode after SSH command completes + stop_active_timer() + start_idle_timer() + return result - # --- CTF Execution --- - if ctf: - # 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 + # --- Local Execution (Default Fallback) --- + # Let _run_local handle determining the host workspace + # Handle Async Session Creation Locally 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 + # 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 + # Switch back to idle mode before returning error + stop_active_timer() + start_idle_timer() + 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") + + # For async, switch back to idle mode after session creation + stop_active_timer() + start_idle_timer() + return f"Started async session {new_session_id} locally. Use this ID to interact." - # --- 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) + # Handle Synchronous Execution Locally using _run_local default with streaming support + result = _run_local(command, stdout, timeout, stream, call_id, tool_name, None) + + # Switch back to idle mode after local command completes + stop_active_timer() + start_idle_timer() + return result + + except Exception as e: + # Ensure we switch back to idle mode if any unexpected error occurs + stop_active_timer() + start_idle_timer() + raise e diff --git a/src/cai/util.py b/src/cai/util.py index 52027fd5..2e66eafc 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -20,6 +20,171 @@ import atexit from dataclasses import dataclass, field from typing import Dict, Optional import time +import threading + +# Global timing variables for tracking active and idle time +_active_timer_start = None +_active_time_total = 0.0 +_idle_timer_start = None +_idle_time_total = 0.0 +_timing_lock = threading.Lock() + +def start_active_timer(): + """ + Start measuring active time (when LLM is processing or tool is executing). + Pauses the idle timer if it's running. + """ + global _active_timer_start, _idle_timer_start, _idle_time_total + + with _timing_lock: + # If idle timer is running, pause it and accumulate time + if _idle_timer_start is not None: + idle_duration = time.time() - _idle_timer_start + _idle_time_total += idle_duration + _idle_timer_start = None + + # Start active timer if not already running + if _active_timer_start is None: + _active_timer_start = time.time() + +def stop_active_timer(): + """ + Stop measuring active time and accumulate the total. + Restarts the idle timer. + """ + global _active_timer_start, _active_time_total, _idle_timer_start + + with _timing_lock: + # If active timer is running, pause it and accumulate time + if _active_timer_start is not None: + active_duration = time.time() - _active_timer_start + _active_time_total += active_duration + _active_timer_start = None + + # Start idle timer if not already running + if _idle_timer_start is None: + _idle_timer_start = time.time() + +def start_idle_timer(): + """ + Start measuring idle time (when waiting for user input). + Pauses the active timer if it's running. + """ + global _idle_timer_start, _active_timer_start, _active_time_total + + with _timing_lock: + # If active timer is running, pause it and accumulate time + if _active_timer_start is not None: + active_duration = time.time() - _active_timer_start + _active_time_total += active_duration + _active_timer_start = None + + # Start idle timer if not already running + if _idle_timer_start is None: + _idle_timer_start = time.time() + +def stop_idle_timer(): + """ + Stop measuring idle time and accumulate the total. + Restarts the active timer. + """ + global _idle_timer_start, _idle_time_total, _active_timer_start + + with _timing_lock: + # If idle timer is running, pause it and accumulate time + if _idle_timer_start is not None: + idle_duration = time.time() - _idle_timer_start + _idle_time_total += idle_duration + _idle_timer_start = None + + # Start active timer if not already running + if _active_timer_start is None: + _active_timer_start = time.time() + +def get_active_time(): + """ + Get the total active time (LLM processing, tool execution). + Returns a formatted string like "1h 30m 45s" or "45s" or "5m 30s". + """ + global _active_time_total, _active_timer_start + + with _timing_lock: + # Calculate total active time including current active period if running + total_active_seconds = _active_time_total + if _active_timer_start is not None: + current_active_duration = time.time() - _active_timer_start + total_active_seconds += current_active_duration + + # Format the time string + hours, remainder = divmod(int(total_active_seconds), 3600) + minutes, seconds = divmod(remainder, 60) + + if hours > 0: + return f"{hours}h {minutes}m {seconds}s" + elif minutes > 0: + return f"{minutes}m {seconds}s" + else: + return f"{seconds}s" + +def get_idle_time(): + """ + Get the total idle time (waiting for user input). + Returns a formatted string like "1h 30m 45s" or "45s" or "5m 30s". + """ + global _idle_time_total, _idle_timer_start + + with _timing_lock: + # Calculate total idle time including current idle period if running + total_idle_seconds = _idle_time_total + if _idle_timer_start is not None: + current_idle_duration = time.time() - _idle_timer_start + total_idle_seconds += current_idle_duration + + # Format the time string + hours, remainder = divmod(int(total_idle_seconds), 3600) + minutes, seconds = divmod(remainder, 60) + + if hours > 0: + return f"{hours}h {minutes}m {seconds}s" + elif minutes > 0: + return f"{minutes}m {seconds}s" + else: + return f"{seconds}s" + +def get_active_time_seconds(): + """ + Get the total active time in seconds for precise measurement. + Returns a float representing the total number of seconds. + """ + global _active_time_total, _active_timer_start + + with _timing_lock: + # Calculate total active time including current active period if running + total_active_seconds = _active_time_total + if _active_timer_start is not None: + current_active_duration = time.time() - _active_timer_start + total_active_seconds += current_active_duration + + return total_active_seconds + +def get_idle_time_seconds(): + """ + Get the total idle time in seconds for precise measurement. + Returns a float representing the total number of seconds. + """ + global _idle_time_total, _idle_timer_start + + with _timing_lock: + # Calculate total idle time including current idle period if running + total_idle_seconds = _idle_time_total + if _idle_timer_start is not None: + current_idle_duration = time.time() - _idle_timer_start + total_idle_seconds += current_idle_duration + + return total_idle_seconds + +# Initialize idle timer at module load - system starts in idle state +start_idle_timer() # Instead of direct import try: @@ -67,6 +232,9 @@ class CostTracker: def log_final_cost(self) -> None: """Display final cost information at exit""" + # Skip displaying cost if already shown in the session summary + if os.environ.get("CAI_COST_DISPLAYED", "").lower() == "true": + return print(f"\nTotal CAI Session Cost: ${self.session_total_cost:.6f}") def get_model_pricing(self, model_name: str) -> tuple: @@ -1124,7 +1292,6 @@ def finish_agent_streaming(context, final_stats=None): context["live"].stop() except Exception as e: context["error"] = str(e) - return False return True except Exception as e: