Agent/agent/trace/compaction.py

"""
Context 压缩 — 两级压缩策略

Level 1: GoalTree 过滤（确定性，零成本）
  - 跳过 completed/abandoned goals 的消息（信息已在 GoalTree summary 中）
  - 始终保留：system prompt、第一条 user message、当前 focus goal 的消息

Level 2: LLM 总结（仅在 Level 1 后仍超限时触发）
  - 在消息列表末尾追加压缩 prompt → 主模型回复 → summary 存为新消息
  - summary 的 parent_sequence 跳过被压缩的范围

压缩不修改存储：原始消息永远保留在 messages/，通过 parent_sequence 树结构实现跳过。
"""

import logging
from dataclasses import dataclass
from typing import List, Dict, Any, Optional, Set

from .goal_models import GoalTree
from .models import Message
from agent.core.prompts import (
    COMPRESSION_EVAL_PROMPT_TEMPLATE,
    REFLECT_PROMPT,
    build_compression_eval_prompt,
)

logger = logging.getLogger(__name__)


# ===== 模型 Context Window（tokens）=====

MODEL_CONTEXT_WINDOWS: Dict[str, int] = {
    # Anthropic Claude
    "claude-sonnet-4": 200_000,
    "claude-opus-4": 200_000,
    "claude-3-5-sonnet": 200_000,
    "claude-3-5-haiku": 200_000,
    "claude-3-opus": 200_000,
    "claude-3-sonnet": 200_000,
    "claude-3-haiku": 200_000,
    # OpenAI
    "gpt-4o": 128_000,
    "gpt-4o-mini": 128_000,
    "gpt-4-turbo": 128_000,
    "gpt-4": 8_192,
    "o1": 200_000,
    "o3-mini": 200_000,
    # Google Gemini
    "gemini-2.5-pro": 1_000_000,
    "gemini-2.5-flash": 1_000_000,
    "gemini-2.0-flash": 1_000_000,
    "gemini-1.5-pro": 2_000_000,
    "gemini-1.5-flash": 1_000_000,
    # DeepSeek
    "deepseek-chat": 64_000,
    "deepseek-r1": 64_000,
}

DEFAULT_CONTEXT_WINDOW = 200_000


def get_context_window(model: str) -> int:
    """
    根据模型名称获取 context window 大小。

    支持带 provider 前缀的模型名（如 "anthropic/claude-sonnet-4.5"）和
    带版本后缀的名称（如 "claude-3-5-sonnet-20241022"）。
    """
    # 去掉 provider 前缀
    name = model.split("/")[-1].lower()

    # 精确匹配
    if name in MODEL_CONTEXT_WINDOWS:
        return MODEL_CONTEXT_WINDOWS[name]

    # 前缀匹配（处理版本后缀）
    for key, window in MODEL_CONTEXT_WINDOWS.items():
        if name.startswith(key):
            return window

    return DEFAULT_CONTEXT_WINDOW


# ===== 配置 =====

@dataclass
class CompressionConfig:
    """压缩配置"""
    max_tokens: int = 0                # 最大 token 数（0 = 自动：context_window * 0.5）
    threshold_ratio: float = 0.5       # 触发压缩的阈值 = context_window 的比例
    keep_recent_messages: int = 10     # Level 1 中始终保留最近 N 条消息
    max_messages: int = 50             # 最大消息数（超过此数量触发压缩，0 = 禁用）

    def get_max_tokens(self, model: str) -> int:
        """获取实际的 max_tokens（如果为 0 则自动计算）"""
        if self.max_tokens > 0:
            return self.max_tokens
        window = get_context_window(model)
        return int(window * self.threshold_ratio)


# ===== Level 1: GoalTree 过滤 =====

def filter_by_goal_status(
    messages: List[Message],
    goal_tree: Optional[GoalTree],
) -> List[Message]:
    """
    Level 1 过滤：跳过 completed/abandoned goals 的消息

    始终保留：
    - goal_id 为 None 的消息（system prompt、初始 user message）
    - 当前 focus goal 及其祖先链上的消息
    - in_progress 和 pending goals 的消息

    跳过：
    - completed 且不在焦点路径上的 goals 的消息
    - abandoned goals 的消息

    Args:
        messages: 主路径上的有序消息列表
        goal_tree: GoalTree 实例

    Returns:
        过滤后的消息列表
    """
    if not goal_tree or not goal_tree.goals:
        return messages

    # 构建焦点路径（当前焦点 + 父链 + 直接子节点）
    focus_path = _get_focus_path(goal_tree)

    # 构建需要跳过的 goal IDs
    skip_goal_ids: Set[str] = set()
    for goal in goal_tree.goals:
        if goal.id in focus_path:
            continue  # 焦点路径上的 goal 始终保留
        if goal.status in ("completed", "abandoned"):
            skip_goal_ids.add(goal.id)

    # 过滤消息
    result = []
    for msg in messages:
        if msg.goal_id is None:
            result.append(msg)  # 无 goal 的消息始终保留
        elif msg.goal_id not in skip_goal_ids:
            result.append(msg)  # 不在跳过列表中的消息保留

    return result


def _get_focus_path(goal_tree: GoalTree) -> Set[str]:
    """
    获取焦点路径上需要保留消息的 goal IDs

    保留：焦点自身 + 父链 + 未完成的直接子节点
    不保留：已完成/已放弃的直接子节点（信息已在 goal.summary 中）
    """
    focus_ids: Set[str] = set()

    if not goal_tree.current_id:
        return focus_ids

    # 焦点自身
    focus_ids.add(goal_tree.current_id)

    # 父链
    goal = goal_tree.find(goal_tree.current_id)
    while goal and goal.parent_id:
        focus_ids.add(goal.parent_id)
        goal = goal_tree.find(goal.parent_id)

    # 直接子节点：仅保留未完成的（completed/abandoned 的信息已在 summary 中）
    children = goal_tree.get_children(goal_tree.current_id)
    for child in children:
        if child.status not in ("completed", "abandoned"):
            focus_ids.add(child.id)

    return focus_ids


# ===== Token 估算 =====

def estimate_tokens(messages: List[Dict[str, Any]]) -> int:
    """
    估算消息列表的 token 数量

    对 CJK 字符和 ASCII 字符使用不同的估算系数：
    - ASCII/Latin 字符：~4 字符 ≈ 1 token
    - CJK 字符（中日韩）：~1 字符 ≈ 1.5 tokens（BPE tokenizer 特性）
    """
    total_tokens = 0
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total_tokens += _estimate_text_tokens(content)
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict):
                    if part.get("type") == "text":
                        total_tokens += _estimate_text_tokens(part.get("text", ""))
                    elif part.get("type") in ("image_url", "image"):
                        total_tokens += _estimate_image_tokens(part)
        # tool_calls
        tool_calls = msg.get("tool_calls")
        if tool_calls and isinstance(tool_calls, list):
            for tc in tool_calls:
                if isinstance(tc, dict):
                    func = tc.get("function", {})
                    total_tokens += len(func.get("name", "")) // 4
                    args = func.get("arguments", "")
                    if isinstance(args, str):
                        total_tokens += _estimate_text_tokens(args)

    return total_tokens


def _estimate_text_tokens(text: str) -> int:
    """
    估算文本的 token 数，区分 CJK 和 ASCII 字符。

    CJK 字符在 BPE tokenizer 中通常占 1.5-2 tokens，
    ASCII 字符约 4 个对应 1 token。
    """
    if not text:
        return 0
    cjk_chars = 0
    other_chars = 0
    for ch in text:
        if _is_cjk(ch):
            cjk_chars += 1
        else:
            other_chars += 1
    # CJK: 1 char ≈ 1.5 tokens; ASCII: 4 chars ≈ 1 token
    return int(cjk_chars * 1.5) + other_chars // 4


def _estimate_image_tokens(block: Dict[str, Any]) -> int:
    """
    估算图片块的 token 消耗。

    Anthropic 计算方式：tokens = (width * height) / 750
    优先从 _image_meta 读取真实尺寸，其次从 base64 数据量粗估，最小 1600 tokens。
    """
    MIN_IMAGE_TOKENS = 1600

    # 优先使用 _image_meta 中的真实尺寸
    meta = block.get("_image_meta")
    if meta and meta.get("width") and meta.get("height"):
        tokens = (meta["width"] * meta["height"]) // 750
        return max(MIN_IMAGE_TOKENS, tokens)

    # 回退：从 base64 数据长度粗估
    b64_data = ""
    if block.get("type") == "image":
        source = block.get("source", {})
        if source.get("type") == "base64":
            b64_data = source.get("data", "")
    elif block.get("type") == "image_url":
        url_obj = block.get("image_url", {})
        url = url_obj.get("url", "") if isinstance(url_obj, dict) else str(url_obj)
        if url.startswith("data:"):
            _, _, b64_data = url.partition(",")

    if b64_data:
        # base64 编码后大小约为原始字节的 4/3
        raw_bytes = len(b64_data) * 3 // 4
        # 粗估：假设 JPEG 压缩率 ~10:1，像素数 ≈ raw_bytes * 10 / 3 (RGB)
        estimated_pixels = raw_bytes * 10 // 3
        estimated_tokens = estimated_pixels // 750
        return max(MIN_IMAGE_TOKENS, estimated_tokens)

    return MIN_IMAGE_TOKENS


def _is_cjk(ch: str) -> bool:
    """判断字符是否为 CJK（中日韩）字符"""
    cp = ord(ch)
    return (
        0x2E80 <= cp <= 0x9FFF       # CJK 基本区 + 部首 + 笔画 + 兼容
        or 0xF900 <= cp <= 0xFAFF    # CJK 兼容表意文字
        or 0xFE30 <= cp <= 0xFE4F    # CJK 兼容形式
        or 0x20000 <= cp <= 0x2FA1F  # CJK 扩展 B-F + 兼容补充
        or 0x3000 <= cp <= 0x303F   # CJK 标点符号
        or 0xFF00 <= cp <= 0xFFEF   # 全角字符
    )


def estimate_tokens_from_messages(messages: List[Message]) -> int:
    """从 Message 对象列表估算 token 数"""
    return estimate_tokens([msg.to_llm_dict() for msg in messages])


def needs_level2_compression(
    token_count: int,
    config: CompressionConfig,
    model: str = "",
) -> bool:
    """判断是否需要触发 Level 2 压缩"""
    limit = config.get_max_tokens(model) if model else config.max_tokens
    return token_count > limit


# ===== Level 2: 压缩 Prompt =====
# 注意：这些 prompt 已迁移到 agent.core.prompts
# COMPRESSION_EVAL_PROMPT 和 REFLECT_PROMPT 现在从 prompts.py 导入


def build_compression_prompt(goal_tree: Optional[GoalTree]) -> str:
    """构建 Level 2 压缩 prompt"""
    goal_prompt = ""
    if goal_tree:
        goal_prompt = goal_tree.to_prompt(include_summary=True)

    return build_compression_eval_prompt(
        goal_tree_prompt=goal_prompt,
    )


def build_reflect_prompt() -> str:
    """构建反思 prompt"""
    return REFLECT_PROMPT