242 lines
9.2 KiB
Python
242 lines
9.2 KiB
Python
# utils/context_manager.py - 上下文管理器（集成对话持久化和Token统计）
|
||
|
||
import os
|
||
import json
|
||
import base64
|
||
import mimetypes
|
||
import io
|
||
import uuid
|
||
import platform
|
||
import shutil
|
||
import subprocess
|
||
from copy import deepcopy
|
||
from typing import Dict, List, Optional, Any
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
try:
|
||
from config import (
|
||
MAX_CONTEXT_SIZE,
|
||
DATA_DIR,
|
||
PROMPTS_DIR,
|
||
TERMINAL_SANDBOX_MOUNT_PATH,
|
||
TERMINAL_SANDBOX_CPUS,
|
||
TERMINAL_SANDBOX_MEMORY,
|
||
PROJECT_MAX_STORAGE_MB,
|
||
TERMINAL_SANDBOX_MODE,
|
||
LINUX_SAFETY,
|
||
)
|
||
from config.model_profiles import (
|
||
get_model_prompt_replacements,
|
||
get_registered_model_keys,
|
||
model_supports_image,
|
||
model_supports_video,
|
||
)
|
||
except ImportError:
|
||
import sys
|
||
from pathlib import Path
|
||
project_root = Path(__file__).resolve().parents[1]
|
||
if str(project_root) not in sys.path:
|
||
sys.path.insert(0, str(project_root))
|
||
from config import (
|
||
MAX_CONTEXT_SIZE,
|
||
DATA_DIR,
|
||
PROMPTS_DIR,
|
||
TERMINAL_SANDBOX_MOUNT_PATH,
|
||
TERMINAL_SANDBOX_CPUS,
|
||
TERMINAL_SANDBOX_MEMORY,
|
||
PROJECT_MAX_STORAGE_MB,
|
||
TERMINAL_SANDBOX_MODE,
|
||
LINUX_SAFETY,
|
||
)
|
||
from config.model_profiles import (
|
||
get_model_prompt_replacements,
|
||
get_registered_model_keys,
|
||
model_supports_image,
|
||
model_supports_video,
|
||
)
|
||
from utils.conversation_manager import ConversationManager
|
||
from utils.host_workspace_debug import write_host_workspace_debug
|
||
from utils.media_store import MediaStore
|
||
from utils.token_usage import normalize_usage_payload
|
||
|
||
AUTO_SHALLOW_PLACEHOLDER = "过早的工具结果已经被自动压缩"
|
||
AUTO_SHALLOW_TOOL_WHITELIST = {
|
||
"write_file",
|
||
"read_file",
|
||
"edit_file",
|
||
"terminal_input",
|
||
"terminal_snapshot",
|
||
"web_search",
|
||
"extract_webpage",
|
||
"run_command",
|
||
"view_image",
|
||
"view_video",
|
||
}
|
||
|
||
|
||
class TokenMixin:
|
||
"""ContextManager token mixin 能力 mixin。"""
|
||
|
||
def _token_totals_path(self) -> Path:
|
||
return self.data_dir / "token_totals.json"
|
||
|
||
def _load_token_totals(self) -> Dict[str, Any]:
|
||
path = self._token_totals_path()
|
||
if not path.exists():
|
||
return {
|
||
"input_tokens": 0,
|
||
"output_tokens": 0,
|
||
"total_tokens": 0,
|
||
"updated_at": None,
|
||
}
|
||
try:
|
||
with open(path, 'r', encoding='utf-8') as fh:
|
||
payload = json.load(fh) or {}
|
||
return {
|
||
"input_tokens": int(payload.get("input_tokens") or payload.get("total_input_tokens") or 0),
|
||
"output_tokens": int(payload.get("output_tokens") or payload.get("total_output_tokens") or 0),
|
||
"total_tokens": int(payload.get("total_tokens") or 0),
|
||
"updated_at": payload.get("updated_at"),
|
||
}
|
||
except (OSError, json.JSONDecodeError, ValueError) as exc:
|
||
print(f"[TokenStats] 读取累计Token失败: {exc}")
|
||
return {
|
||
"input_tokens": 0,
|
||
"output_tokens": 0,
|
||
"total_tokens": 0,
|
||
"updated_at": None,
|
||
}
|
||
|
||
def _save_token_totals(self, data: Dict[str, Any]):
|
||
path = self._token_totals_path()
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(path, 'w', encoding='utf-8') as fh:
|
||
json.dump(data, fh, ensure_ascii=False, indent=2)
|
||
|
||
def _increment_workspace_token_totals(self, input_tokens: int, output_tokens: int, total_tokens: int):
|
||
if input_tokens <= 0 and output_tokens <= 0 and total_tokens <= 0:
|
||
return
|
||
snapshot = self._load_token_totals()
|
||
snapshot["input_tokens"] = snapshot.get("input_tokens", 0) + max(0, int(input_tokens))
|
||
snapshot["output_tokens"] = snapshot.get("output_tokens", 0) + max(0, int(output_tokens))
|
||
snapshot["total_tokens"] = snapshot.get("total_tokens", 0) + max(0, int(total_tokens))
|
||
snapshot["updated_at"] = datetime.now().isoformat()
|
||
self._save_token_totals(snapshot)
|
||
|
||
def apply_usage_statistics(self, usage: Dict[str, Any]) -> bool:
|
||
"""
|
||
根据模型返回的 usage 字段更新token统计
|
||
"""
|
||
normalized_usage = normalize_usage_payload(usage) or {}
|
||
prompt_tokens = int(normalized_usage.get("prompt_tokens") or 0)
|
||
completion_tokens = int(normalized_usage.get("completion_tokens") or 0)
|
||
total_tokens = int(normalized_usage.get("total_tokens") or (prompt_tokens + completion_tokens))
|
||
# 当前上下文长度优先取专用字段；缺失时回退到 prompt_tokens
|
||
current_context_tokens = int(normalized_usage.get("current_context_tokens") or prompt_tokens)
|
||
|
||
try:
|
||
self._increment_workspace_token_totals(prompt_tokens, completion_tokens, total_tokens)
|
||
except Exception as exc:
|
||
print(f"[TokenStats] 无法写入累计Token: {exc}")
|
||
|
||
if not self.current_conversation_id:
|
||
print("⚠️ 没有当前对话ID，跳过usage统计更新")
|
||
return False
|
||
|
||
try:
|
||
# 路由到正确的 conversation_manager（与 get_conversation_token_statistics 一致）
|
||
# 多智能体对话 ID 存储在 multi_agent_conversation_manager 中，
|
||
# 直接使用 self.conversation_manager 会找不到对话导致统计写入失败
|
||
target_manager = getattr(
|
||
self, "_get_conversation_manager_for_id", lambda _: self.conversation_manager
|
||
)(self.current_conversation_id)
|
||
success = target_manager.update_token_statistics(
|
||
self.current_conversation_id,
|
||
prompt_tokens,
|
||
completion_tokens,
|
||
total_tokens,
|
||
current_context_tokens=current_context_tokens,
|
||
)
|
||
|
||
if success:
|
||
self.safe_broadcast_token_update()
|
||
|
||
return success
|
||
except Exception as e:
|
||
print(f"更新usage统计失败: {e}")
|
||
return False
|
||
|
||
def get_conversation_token_statistics(self, conversation_id: str = None) -> Optional[Dict]:
|
||
"""
|
||
获取指定对话的token统计
|
||
|
||
Args:
|
||
conversation_id: 对话ID，默认为当前对话
|
||
|
||
Returns:
|
||
Dict: Token统计信息
|
||
"""
|
||
target_id = conversation_id or self.current_conversation_id
|
||
if not target_id:
|
||
return None
|
||
|
||
# 路由到正确的 conversation_manager（支持多智能体会话独立存储）
|
||
target_manager = getattr(
|
||
self, "_get_conversation_manager_for_id", lambda _: self.conversation_manager
|
||
)(target_id)
|
||
return target_manager.get_token_statistics(target_id)
|
||
|
||
def get_current_context_tokens(self, conversation_id: str = None) -> int:
|
||
"""
|
||
获取最近一次请求的上下文token数量
|
||
"""
|
||
stats = self.get_conversation_token_statistics(conversation_id)
|
||
if not stats:
|
||
return 0
|
||
return stats.get("current_context_tokens", 0)
|
||
|
||
def safe_broadcast_token_update(self):
|
||
"""安全的token更新广播（只广播累计统计，不重新计算）"""
|
||
try:
|
||
print(f"[Debug] 尝试广播token更新")
|
||
|
||
# 检查是否有回调函数
|
||
if not hasattr(self, '_web_terminal_callback'):
|
||
print(f"[Debug] 没有_web_terminal_callback属性")
|
||
return
|
||
|
||
if not self._web_terminal_callback:
|
||
print(f"[Debug] _web_terminal_callback为None")
|
||
return
|
||
|
||
if not self.current_conversation_id:
|
||
print(f"[Debug] 没有当前对话ID")
|
||
return
|
||
|
||
print(f"[Debug] 广播token统计，对话ID: {self.current_conversation_id}")
|
||
|
||
# 只获取已有的累计token统计，不重新计算
|
||
cumulative_stats = self.get_conversation_token_statistics()
|
||
|
||
# 准备广播数据
|
||
broadcast_data = {
|
||
'conversation_id': self.current_conversation_id,
|
||
'cumulative_input_tokens': cumulative_stats.get("total_input_tokens", 0) if cumulative_stats else 0,
|
||
'cumulative_output_tokens': cumulative_stats.get("total_output_tokens", 0) if cumulative_stats else 0,
|
||
'cumulative_total_tokens': cumulative_stats.get("total_tokens", 0) if cumulative_stats else 0,
|
||
'current_context_tokens': cumulative_stats.get("current_context_tokens", 0) if cumulative_stats else 0,
|
||
'updated_at': datetime.now().isoformat()
|
||
}
|
||
|
||
print(f"[Debug] Token统计: 累计输入={broadcast_data['cumulative_input_tokens']}, 累计输出={broadcast_data['cumulative_output_tokens']}")
|
||
|
||
# 广播到前端
|
||
self._web_terminal_callback('token_update', broadcast_data)
|
||
|
||
print(f"[Debug] token更新已广播")
|
||
|
||
except Exception as e:
|
||
print(f"[Debug] 广播token更新失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|