autool/DroidBot/guiagent_core/context_manager.py

"""
Context Manager - 管理对话历史和构建模型请求

Extracted from GuiAgent/core/context.py for modular use in DroidBot.
"""
import re
import logging
from typing import Optional, Dict, Any, List

logger = logging.getLogger(__name__)

# Constants
IMAGE_PLACEHOLDER = "<image>"
MAX_HISTORY_LENGTH = 20


class ContextManager:
    """
    上下文管理器，负责：
    1. 管理对话历史
    2. 构建发送给模型的请求消息
    3. 过滤和限制对话历史长度
    """
    
    def __init__(self, system_prompt: str, max_history: int = MAX_HISTORY_LENGTH):
        """
        初始化上下文管理器
        
        Args:
            system_prompt: 系统提示词
            max_history: 最大对话历史长度
        """
        self._system_prompt = system_prompt
        self.max_history = max_history
        
        # 对话历史：[{from: 'human'/'gpt', value: str, screenshot?: str}, ...]
        self.conversations: List[Dict[str, Any]] = []
        
        # 初始指令是否已添加
        self._has_instruction: bool = False

    @property
    def system_prompt(self) -> str:
        return self._system_prompt
    
    @system_prompt.setter
    def system_prompt(self, value: str) -> None:
        self._system_prompt = value

    def add_instruction(self, instruction: str) -> None:
        """添加用户初始指令（仅调用一次）"""
        if self._has_instruction:
            return
        
        self.conversations.append({
            'from': 'human',
            'value': instruction
        })
        self._has_instruction = True

    def add_user_message(self, message: str) -> None:
        """
        添加用户消息到对话历史
        
        Args:
            message: 用户消息内容
        """
        self.conversations.append({
            'from': 'human',
            'value': message
        })
        self._prune()

    def add_screenshot(self, screenshot_base64: str, width: int, height: int) -> None:
        """
        添加截图到对话历史
        
        Args:
            screenshot_base64: 截图的 Base64 编码（不含 data:image 前缀）
            width: 图像宽度
            height: 图像高度
        """
        self.conversations.append({
            'from': 'human',
            'value': IMAGE_PLACEHOLDER,
            'screenshot': screenshot_base64,
            'size': (width, height)
        })
        self._prune()

    def add_response(self, response: str) -> None:
        """
        添加模型响应到对话历史

        Args:
            response: 模型响应文本
        """
        # 提取摘要（移除 Reflection 部分）
        summary = re.sub(r'Reflection:[\s\S]*?(?=Action:|$)', '', response).strip()

        self.conversations.append({
            'from': 'gpt',
            'value': summary
        })
        self._prune()

    def add_execution_feedback(self, success: bool, detail: str = "") -> None:
        """
        添加动作执行结果反馈

        Args:
            success: 执行是否成功
            detail: 详细信息（失败原因或额外说明）
        """
        if success:
            feedback = "✓ 上次操作已成功执行"
            if detail:
                feedback += f": {detail}"
        else:
            feedback = f"✗ 上次操作失败: {detail}。请分析失败原因并调整策略（如调整坐标、使用替代方案等）"

        self.conversations.append({
            'from': 'system',
            'value': feedback
        })
        self._prune()

    def _prune(self) -> None:
        """限制对话历史长度，保留初始指令 + 最近对话"""
        if len(self.conversations) <= self.max_history:
            return
        
        # 保留第一条（初始指令）+ 最近的对话
        first = self.conversations[0] if self.conversations else None
        recent = self.conversations[-(self.max_history - 1):]
        self.conversations = [first] + recent if first else recent

    def _get_filtered(self) -> List[Dict[str, Any]]:
        """
        获取过滤后的对话历史：
        - 保留所有 AI 回复
        - 保留所有系统反馈（执行结果）
        - 保留初始指令 + 最近3轮的完整对话（包括截图）
        """
        # 找出所有截图索引
        screenshot_indices = [i for i, c in enumerate(self.conversations) if c.get('screenshot')]

        if not screenshot_indices:
            return self.conversations

        # 保留最近3轮的截图及其后续对话
        KEEP_RECENT_ROUNDS = 3
        keep_from_idx = screenshot_indices[-KEEP_RECENT_ROUNDS] if len(screenshot_indices) >= KEEP_RECENT_ROUNDS else 0

        filtered = []
        for i, conv in enumerate(self.conversations):
            # 始终保留初始指令
            if i == 0:
                filtered.append(conv)
            # 保留最近N轮的所有对话
            elif i >= keep_from_idx:
                filtered.append(conv)
            # 保留所有系统反馈（即使在旧轮次中）
            elif conv['from'] == 'system':
                filtered.append(conv)

        return filtered

    def build_messages(self) -> List[Dict[str, Any]]:
        """
        构建发送给模型的消息列表（OpenAI 格式）

        Returns:
            消息列表，每条消息包含 role 和 content
        """
        filtered = self._get_filtered()
        messages = []
        images = []

        # 收集图像
        for conv in filtered:
            if conv.get('screenshot'):
                images.append(conv['screenshot'])

        image_idx = 0
        for i, conv in enumerate(filtered):
            if conv.get('screenshot'):
                # 图像消息
                messages.append({
                    'role': 'user',
                    'content': [{
                        'type': 'image_url',
                        'image_url': {'url': f"data:image/png;base64,{images[image_idx]}"}
                    }]
                })
                image_idx += 1
            elif i == 0 and conv['from'] == 'human':
                # 第一条消息：嵌入系统提示词
                messages.append({
                    'role': 'system',
                    'content': f"{self._system_prompt}\n{conv['value']}"
                })
            else:
                # 普通消息和系统反馈
                role_map = {'human': 'user', 'gpt': 'assistant', 'system': 'user'}
                role = role_map.get(conv['from'], 'user')
                messages.append({
                    'role': role,
                    'content': conv['value']
                })

        return messages

    def clear(self) -> None:
        """清空对话历史"""
        self.conversations = []
        self._has_instruction = False