mirror of https://github.com/aliasrobotics/cai.git
Logging and timming from 0.3.0
This commit is contained in:
luijait
parent
9fd1a2fa8d
commit
7b016e058e
|
|
@ -22,6 +22,7 @@ lib/
|
|||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
logs/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
169
src/cai/util.py
169
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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue