aiagent/backend/app/agent_runtime/context.py

"""
Agent 会话上下文管理：维护消息历史、状态追踪。
"""
from __future__ import annotations

import uuid
from typing import Any, Dict, List, Optional


class AgentContext:
    """
    Agent 会话上下文：

    - 消息历史（messages 列表，OpenAI 格式）
    - 会话元信息（session_id, user_id 等）
    - 执行追踪（iteration 计数, 工具调用统计）
    """

    def __init__(
        self,
        system_prompt: str = "你是一个有用的AI助手。",
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
    ):
        self.session_id = session_id or str(uuid.uuid4())
        self.user_id = user_id
        self._messages: List[Dict[str, Any]] = []
        self._system_prompt = system_prompt
        # 执行状态
        self.iteration = 0
        self.tool_calls_made = 0

    @property
    def messages(self) -> List[Dict[str, Any]]:
        """获取完整消息列表（含 system prompt）。"""
        if self._system_prompt:
            # 确保 system prompt 始终在第一条
            has_system = (
                len(self._messages) > 0
                and self._messages[0].get("role") == "system"
            )
            if not has_system:
                return [
                    {"role": "system", "content": self._system_prompt},
                    *self._messages,
                ]
        return self._messages

    def add_user_message(self, content: str) -> None:
        """添加用户消息。"""
        self._messages.append({"role": "user", "content": content})

    def add_assistant_message(
        self,
        content: str,
        tool_calls: Optional[List[Dict[str, Any]]] = None,
        reasoning_content: Optional[str] = None,
    ) -> None:
        """添加助手回复。"""
        msg: Dict[str, Any] = {"role": "assistant", "content": content or ""}
        if tool_calls:
            msg["tool_calls"] = tool_calls
        if reasoning_content:
            msg["reasoning_content"] = reasoning_content
        self._messages.append(msg)

    def add_tool_result(
        self, tool_call_id: str, tool_name: str, result: str
    ) -> None:
        """添加工具执行结果。"""
        self._messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "content": result,
            "name": tool_name,
        })

    def set_system_prompt(self, prompt: str) -> None:
        """更新 system prompt（仅在未发送过消息时有效）。"""
        if not self._messages:
            self._system_prompt = prompt

    # ──────────────────── 消息操作（为 Compaction 提供） ────────────────────

    @property
    def raw_messages(self) -> List[Dict[str, Any]]:
        """获取原始消息列表（不含自动 prepend 的 system prompt）。

        用于 CompactionEngine 直接操作 _messages，避免 system prompt 重复。
        """
        return self._messages

    def replace_internal_messages(self, new_messages: List[Dict[str, Any]]) -> None:
        """替换全部内部消息（CompactionEngine 用）。"""
        self._messages = new_messages

    def remove_messages_before(self, index: int) -> None:
        """移除指定索引之前的所有消息（保留 system prompt 在首位时的位置）。

        注意：不修改 _system_prompt；调用方通过 messages 属性获取时会自动 prepend。
        """
        if index > 0:
            self._messages = self._messages[index:]

    def replace_message_range(
        self, start: int, end: int, new_messages: List[Dict[str, Any]]
    ) -> None:
        """替换消息列表中 [start, end) 区间的消息。"""
        self._messages[start:end] = new_messages

    def estimate_tokens(self, token_counter=None) -> int:
        """估算当前消息列表的总 token 数（含 system prompt）。"""
        from app.core.token_counter import TokenCounter
        if token_counter is None:
            token_counter = TokenCounter()
        return token_counter.count_messages(self.messages)

    # ──────────────────── 生命周期 ────────────────────

    def reset(self) -> None:
        """重置上下文（保留 system prompt 和 session_id）。"""
        self._messages = []
        self.iteration = 0
        self.tool_calls_made = 0