MCP Client 的上下文裁剪策略

from typing import Dict, List, Any, Optional
from datetime import datetime
import hashlib

class ContextManager:
    def __init__(self, max_context_length: int = 8192, default_cut_strategy: str = "relevance"):
        self.contexts = {}  # 存储所有上下文，按会话 ID 组织
        self.max_context_length = max_context_length
        self.default_cut_strategy = default_cut_strategy
        self.cut_strategies = {
            "relevance": self._cut_by_relevance,
            "recent": self._cut_by_recency,
            "fixed": self._cut_by_fixed_length,
            "hybrid": self._cut_by_hybrid
        }
    
    def add_context(self, session_id: str, role: str, content: str, context_type: str = "general", metadata: Dict[str, Any] = None):
        """添加上下文"""
        if session_id not in self.contexts:
            self.contexts[session_id] = {
                "messages": [],
                "last_updated": datetime.now().isoformat(),
                "metadata": {}
            }
        
        # 生成上下文 ID
        context_id = hashlib.sha256((content + datetime.now().isoformat()).encode()).hexdigest()
        
        # 添加上下文消息
        self.contexts[session_id]["messages"].append({
            "id": context_id,
            "role": role,
            "content": content,
            "type": context_type,
            "timestamp": datetime.now().isoformat(),
            "metadata": metadata or {}
        })
        
        # 更新最后更新时间
        self.contexts[session_id]["last_updated"] = datetime.now().isoformat()
        
        # 自动裁剪上下文
        self._auto_cut_context(session_id)
    
    def get_context(self, session_id: str, cut_strategy: Optional[str] = None, max_length: Optional[int] = None) -> List[Dict[str, Any]]:
        """获取裁剪后的上下文"""
        if session_id not in self.contexts:
            return []
        
        # 使用指定的裁剪策略或默认策略
        strategy = cut_strategy or self.default_cut_strategy
        if strategy not in self.cut_strategies:
            strategy = self.default_cut_strategy
        
        # 使用指定的最大长度或默认长度
        max_len = max_length or self.max_context_length
        
        # 执行裁剪
        return self.cut_strategies[strategy](session_id, max_len)
    
    def _auto_cut_context(self, session_id: str):
        """自动裁剪上下文"""
        if session_id not in self.contexts:
            return
        
        # 检查上下文长度是否超过阈值
        current_length = self._calculate_context_length(session_id)
        if current_length > self.max_context_length * 1.2:  # 超过 20% 时触发自动裁剪
            self.contexts[session_id]["messages"] = self._cut_by_hybrid(session_id, self.max_context_length)
    
    def _calculate_context_length(self, session_id: str) -> int:
        """计算上下文长度（按 Token 估算）"""
        if session_id not in self.contexts:
            return 0
        
        total_length = 0
        for message in self.contexts[session_id]["messages"]:
            # 粗略估算：1 Token ≈ 4 字符
            total_length += len(message["content"]) // 4
        
        return total_length
    
    def _cut_by_relevance(self, session_id: str, max_length: int) -> List[Dict[str, Any]]:
        """基于相关性裁剪上下文"""
        if session_id not in self.contexts:
            return []
        
        # 获取所有上下文消息
        messages = self.contexts[session_id]["messages"].copy()
        
        # 这里简化实现，实际应使用相关性评分算法
        # 按时间倒序排列，保留最新的消息
        messages.sort(key=lambda x: x["timestamp"], reverse=True)
        
        # 裁剪到指定长度
        current_length = 0
        result = []
        
        for message in messages:
            message_length = len(message["content"]) // 4
            if current_length + message_length <= max_length:
                result.append(message)
                current_length += message_length
            else:
                break
        
        # 按时间正序返回
        result.sort(key=lambda x: x["timestamp"])
        return result
    
    def _cut_by_recency(self, session_id: str, max_length: int) -> List[Dict[str, Any]]:
        """基于时间裁剪上下文"""
        if session_id not in self.contexts:
            return []
        
        # 获取所有上下文消息
        messages = self.contexts[session_id]["messages"].copy()
        
        # 按时间倒序排列
        messages.sort(key=lambda x: x["timestamp"], reverse=True)
        
        # 裁剪到指定长度
        current_length = 0
        result = []
        
        for message in messages:
            message_length = len(message["content"]) // 4
            if current_length + message_length <= max_length:
                result.append(message)
                current_length += message_length
            else:
                break
        
        # 按时间正序返回
        result.sort(key=lambda x: x["timestamp"])
        return result
    
    def _cut_by_fixed_length(self, session_id: str, max_length: int) -> List[Dict[str, Any]]:
        """基于固定长度裁剪上下文"""
        if session_id not in self.contexts:
            return []
        
        # 获取所有上下文消息
        messages = self.contexts[session_id]["messages"].copy()
        
        # 简单的固定长度裁剪，保留最后 N 条消息
        # 这里简化实现，实际应考虑消息长度
        return messages[-max_length:]
    
    def _cut_by_hybrid(self, session_id: str, max_length: int) -> List[Dict[str, Any]]:
        """混合裁剪策略"""
        if session_id not in self.contexts:
            return []
        
        # 这里实现一个简单的混合策略：
        # 1. 保留最新的系统消息
        # 2. 按时间倒序保留用户和助手消息，直到达到长度限制
        
        messages = self.contexts[session_id]["messages"].copy()
        system_messages = [msg for msg in messages if msg["role"] == "system"]
        user_assistant_messages = [msg for msg in messages if msg["role"] in ["user", "assistant"]]
        
        # 按时间倒序排列用户和助手消息
        user_assistant_messages.sort(key=lambda x: x["timestamp"], reverse=True)
        
        # 计算系统消息的长度
        system_length = sum(len(msg["content"]) // 4 for msg in system_messages)
        remaining_length = max_length - system_length
        
        # 裁剪用户和助手消息
        current_length = 0
        selected_messages = []
        
        for message in user_assistant_messages:
            message_length = len(message["content"]) // 4
            if current_length + message_length <= remaining_length:
                selected_messages.append(message)
                current_length += message_length
            else:
                break
        
        # 合并并按时间正序返回
        selected_messages.sort(key=lambda x: x["timestamp"])
        result = system_messages + selected_messages
        result.sort(key=lambda x: x["timestamp"])
        
        return result
    
    def clear_context(self, session_id: str):
        """清空指定会话的上下文"""
        if session_id in self.contexts:
            del self.contexts[session_id]
    
    def clear_all_contexts(self):
        """清空所有上下文"""
        self.contexts.clear()
    
    def get_context_stats(self, session_id: str) -> Dict[str, Any]:
        """获取上下文统计信息"""
        if session_id not in self.contexts:
            return {}
        
        messages = self.contexts[session_id]["messages"]
        
        return {
            "total_messages": len(messages),
            "total_length": self._calculate_context_length(session_id),
            "last_updated": self.contexts[session_id]["last_updated"],
            "message_types": {
                "system": len([msg for msg in messages if msg["role"] == "system"]),
                "user": len([msg for msg in messages if msg["role"] == "user"]),
                "assistant": len([msg for msg in messages if msg["role"] == "assistant"]),
                "tool": len([msg for msg in messages if msg["role"] == "tool"])
            }
        }

from typing import Dict, List, Any, Tuple
import numpy as np
from sentence_transformers import SentenceTransformer

class RelevanceScorer:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.relevance_threshold = 0.5
        self.feedback_history = []
    
    def score_relevance(self, context_messages: List[Dict[str, Any]], current_query: str) -> List[Tuple[str, float]]:
        """评估上下文消息与当前查询的相关性"""
        if not context_messages:
            return []
        
        # 生成当前查询的嵌入
        query_embedding = self.model.encode([current_query])[0]
        
        # 生成所有上下文消息的嵌入
        context_texts = [msg["content"] for msg in context_messages]
        context_embeddings = self.model.encode(context_texts)
        
        # 计算余弦相似度
        similarities = np.dot(context_embeddings, query_embedding) / (np.linalg.norm(context_embeddings, axis=1) * np.linalg.norm(query_embedding))
        
        # 生成相关性评分列表
        relevance_scores = []
        for msg, similarity in zip(context_messages, similarities):
            relevance_scores.append((msg["id"], float(similarity)))
        
        # 按相关性降序排序
        relevance_scores.sort(key=lambda x: x[1], reverse=True)
        
        return relevance_scores
    
    def get_relevant_contexts(self, context_messages: List[Dict[str, Any]], current_query: str, top_k: int = 10) -> List[Dict[str, Any]]:
        """获取相关性最高的前 k 个上下文"""
        relevance_scores = self.score_relevance(context_messages, current_query)
        
        # 获取相关性最高的前 k 个上下文 ID
        relevant_ids = [msg_id for msg_id, score in relevance_scores[:top_k] if score >= self.relevance_threshold]
        
        # 过滤出相关的上下文消息
        relevant_contexts = [msg for msg in context_messages if msg["id"] in relevant_ids]
        
        # 按时间正序返回
        relevant_contexts.sort(key=lambda x: x["timestamp"])
        
        return relevant_contexts
    
    def update_threshold(self, new_threshold: float):
        """更新相关性阈值"""
        self.relevance_threshold = new_threshold
    
    def add_feedback(self, context_id: str, query: str, is_relevant: bool, score: float):
        """添加反馈数据，用于后续模型优化"""
        feedback_item = {
            "context_id": context_id,
            "query": query,
            "is_relevant": is_relevant,
            "score": score,
            "timestamp": datetime.now().isoformat()
        }
        self.feedback_history.append(feedback_item)
    
    def get_feedback_stats(self) -> Dict[str, Any]:
        """获取反馈统计信息"""
        total_feedback = len(self.feedback_history)
        positive_feedback = len([f for f in self.feedback_history if f["is_relevant"]])
        negative_feedback = total_feedback - positive_feedback
        
        avg_score = np.mean([f["score"] for f in self.feedback_history]) if total_feedback > 0 else 0.0
        
        return {
            "total_feedback": total_feedback,
            "positive_feedback": positive_feedback,
            "negative_feedback": negative_feedback,
            "avg_score": avg_score
        }

from typing import Dict, List, Any, Optional

class SmartCutter:
    def __init__(self, relevance_scorer, default_strategy: str = "relevance"):
        self.relevance_scorer = relevance_scorer
        self.default_strategy = default_strategy
        self.strategies = {
            "relevance": self._cut_by_relevance,
            "recency": self._cut_by_recency,
            "fixed": self._cut_by_fixed_length,
            "hybrid": self._cut_by_hybrid
        }
    
    def cut_context(self, context_messages: List[Dict[str, Any]], current_query: str, max_length: int, strategy: Optional[str] = None) -> List[Dict[str, Any]]:
        """执行上下文裁剪"""
        # 使用指定策略或默认策略
        cut_strategy = strategy or self.default_strategy
        if cut_strategy not in self.strategies:
            cut_strategy = self.default_strategy
        
        # 执行裁剪
        return self.strategies[cut_strategy](context_messages, current_query, max_length)
    
    def _cut_by_relevance(self, context_messages: List[Dict[str, Any]], current_query: str, max_length: int) -> List[Dict[str, Any]]:
        """基于相关性裁剪"""
        # 1. 获取相关性评分
        relevance_scores = self.relevance_scorer.score_relevance(context_messages, current_query)
        
        # 2. 按相关性排序上下文
        relevance_dict = {msg_id: score for msg_id, score in relevance_scores}
        context_messages.sort(key=lambda x: relevance_dict.get(x["id"], 0.0), reverse=True)
        
        # 3. 裁剪到指定长度
        current_length = 0
        result = []
        
        for message in context_messages:
            message_length = len(message["content"]) // 4
            if current_length + message_length <= max_length:
                result.append(message)
                current_length += message_length
            else:
                break
        
        # 4. 按时间正序返回
        result.sort(key=lambda x: x["timestamp"])
        return result
    
    def _cut_by_recency(self, context_messages: List[Dict[str, Any]], current_query: str, max_length: int) -> List[Dict[str, Any]]:
        """基于时间裁剪"""
        # 1. 按时间倒序排序
        context_messages.sort(key=lambda x: x["timestamp"], reverse=True)
        
        # 2. 裁剪到指定长度
        current_length = 0
        result = []
        
        for message in context_messages:
            message_length = len(message["content"]) // 4
            if current_length + message_length <= max_length:
                result.append(message)
                current_length += message_length
            else:
                break
        
        # 3. 按时间正序返回
        result.sort(key=lambda x: x["timestamp"])
        return result
    
    def _cut_by_fixed_length(self, context_messages: List[Dict[str, Any]], current_query: str, max_length: int) -> List[Dict[str, Any]]:
        """基于固定长度裁剪"""
        # 简单实现：返回最后 max_length 条消息
        return context_messages[-max_length:]
    
    def _cut_by_hybrid(self, context_messages: List[Dict[str, Any]], current_query: str, max_length: int) -> List[Dict[str, Any]]:
        """混合裁剪策略"""
        # 1. 分离不同类型的消息
        system_messages = [msg for msg in context_messages if msg["role"] == "system"]
        tool_messages = [msg for msg in context_messages if msg["role"] == "tool"]
        user_assistant_messages = [msg for msg in context_messages if msg["role"] in ["user", "assistant"]]
        
        # 2. 计算系统消息和工具消息的长度
        system_length = sum(len(msg["content"]) // 4 for msg in system_messages)
        tool_length = sum(len(msg["content"]) // 4 for msg in tool_messages)
        
        # 3. 计算剩余可用长度
        remaining_length = max_length - system_length - tool_length
        if remaining_length <= 0:
            # 如果剩余长度不足，只保留系统消息
            return sorted(system_messages, key=lambda x: x["timestamp"])
        
        # 4. 对用户和助手消息进行相关性裁剪
        relevant_user_assistant = self._cut_by_relevance(user_assistant_messages, current_query, remaining_length)
        
        # 5. 合并所有消息
        result = system_messages + relevant_user_assistant + tool_messages
        
        # 6. 按时间正序返回
        result.sort(key=lambda x: x["timestamp"])
        
        return result
    
    def get_cut_stats(self, original: List[Dict[str, Any]], cut: List[Dict[str, Any]]) -> Dict[str, Any]:
        """获取裁剪统计信息"""
        original_length = sum(len(msg["content"]) // 4 for msg in original)
        cut_length = sum(len(msg["content"]) // 4 for msg in cut)
        
        return {
            "original_messages": len(original),
            "cut_messages": len(cut),
            "original_length": original_length,
            "cut_length": cut_length,
            "reduction_ratio": (original_length - cut_length) / original_length if original_length > 0 else 0.0
        }

# 示例：上下文裁剪的完整示例
import asyncio
from mcp_client.context_manager import ContextManager
from mcp_client.relevance_scorer import RelevanceScorer
from mcp_client.smart_cutter import SmartCutter
from mcp_client.mcp_executor import MCPExecutor

async def main():
    # 1. 创建上下文管理器
    context_manager = ContextManager(max_context_length=4096)
    
    # 2. 创建相关性评分器
    relevance_scorer = RelevanceScorer()
    
    # 3. 创建智能裁剪器
    smart_cutter = SmartCutter(relevance_scorer)
    
    # 4. 创建 MCP 执行器
    mcp_executor = MCPExecutor(server_url="http://localhost:8000/mcp")
    
    # 5. 模拟对话过程
    session_id = "test_session_123"
    
    # 添加一些示例上下文
    sample_contexts = [
        {"role": "system", "content": "你是一个助手，可以使用工具来帮助用户解决问题。", "type": "system"},
        {"role": "user", "content": "我想了解今天的天气情况。", "type": "user"},
        {"role": "assistant", "content": "好的，我来帮你查询今天的天气。", "type": "assistant"},
        {"role": "tool", "content": "{\"tool_call\": {\"name\": \"weather_api\", \"arguments\": {\"city\": \"北京\"}}}", "type": "tool_call"},
        {"role": "tool", "content": "{\"temperature\": 22, \"humidity\": 45, \"condition\": \"晴\"}", "type": "tool_result"},
        {"role": "assistant", "content": "今天北京的天气是晴，温度22度，湿度45%。", "type": "assistant"},
        {"role": "user", "content": "那明天呢？", "type": "user"},
        {"role": "assistant", "content": "我来帮你查询明天的天气。", "type": "assistant"},
        {"role": "tool", "content": "{\"tool_call\": {\"name\": \"weather_api\", \"arguments\": {\"city\": \"北京\", \"days\": 1}}}", "type": "tool_call"},
        {"role": "tool", "content": "{\"temperature\": 24, \"humidity\": 50, \"condition\": \"多云\"}", "type": "tool_result"},
        {"role": "assistant", "content": "明天北京的天气是多云，温度24度，湿度50%。", "type": "assistant"}
    ]
    
    # 添加到上下文管理器
    for ctx in sample_contexts:
        context_manager.add_context(session_id, ctx["role"], ctx["content"], ctx["type"])
    
    # 6. 获取当前上下文
    current_query = "后天北京的天气怎么样？"
    
    # 7. 使用不同策略裁剪上下文
    print("=== 不同裁剪策略的效果对比 ===")
    
    # 7.1 原始上下文
    original_context = context_manager.get_context(session_id, strategy="fixed", max_length=10000)
    original_length = sum(len(msg["content"]) // 4 for msg in original_context)
    print(f"原始上下文：{len(original_context)}条消息，约{original_length}个Token")
    
    # 7.2 相关性裁剪
    relevance_cut = smart_cutter.cut_context(original_context, current_query, max_length=2048, strategy="relevance")
    relevance_length = sum(len(msg["content"]) // 4 for msg in relevance_cut)
    print(f"相关性裁剪：{len(relevance_cut)}条消息，约{relevance_length}个Token")
    
    # 7.3 时间裁剪
    recency_cut = smart_cutter.cut_context(original_context, current_query, max_length=2048, strategy="recency")
    recency_length = sum(len(msg["content"]) // 4 for msg in recency_cut)
    print(f"时间裁剪：{len(recency_cut)}条消息，约{recency_length}个Token")
    
    # 7.4 混合裁剪
    hybrid_cut = smart_cutter.cut_context(original_context, current_query, max_length=2048, strategy="hybrid")
    hybrid_length = sum(len(msg["content"]) // 4 for msg in hybrid_cut)
    print(f"混合裁剪：{len(hybrid_cut)}条消息，约{hybrid_length}个Token")
    
    # 8. 查看裁剪效果统计
    print("\n=== 裁剪效果统计 ===")
    stats = smart_cutter.get_cut_stats(original_context, hybrid_cut)
    print(f"原始消息数：{stats['original_messages']}")
    print(f"裁剪后消息数：{stats['cut_messages']}")
    print(f"原始长度：{stats['original_length']} Token")
    print(f"裁剪后长度：{stats['cut_length']} Token")
    print(f"缩减比例：{stats['reduction_ratio']:.2%}")
    
    # 9. 实际使用裁剪后的上下文调用 LLM
    print("\n=== 使用裁剪上下文调用 LLM ===")
    print(f"使用混合裁剪策略，裁剪后上下文长度：{hybrid_length} Token")
    
    # 这里可以继续实现使用裁剪上下文调用 LLM 的逻辑
    # ...
    
    print("上下文裁剪示例完成！")

if __name__ == "__main__":
    asyncio.run(main())

# MCP Client 上下文裁剪配置
context_management:
  # 上下文管理器配置
  manager:
    max_context_length: 4096
    default_cut_strategy: "hybrid"
    auto_cut_enabled: true
    auto_cut_threshold: 1.2  # 超过 max_length 的倍数时自动裁剪
  
  # 相关性评分器配置
  relevance_scorer:
    enabled: true
    model_name: "all-MiniLM-L6-v2"
    default_threshold: 0.5
    use_llm_assist: false  # 是否使用 LLM 辅助评估
    llm_model: "gpt-4o"
    llm_temperature: 0.1
  
  # 智能裁剪器配置
  smart_cutter:
    enabled: true
    strategies:
      relevance: true
      recency: true
      fixed: true
      hybrid: true
    hybrid_weights:
      system_messages: 0.4
      user_messages: 0.3
      assistant_messages: 0.2
      tool_messages: 0.1
  
  # 裁剪策略配置
  strategies:
    relevance:
      enabled: true
      top_k: 10
      min_score: 0.3
    recency:
      enabled: true
      lookback_window: 3600  # 秒
    fixed:
      enabled: true
      fixed_length: 2048
    hybrid:
      enabled: true
      default_max_length: 4096
  
  # 性能优化配置
  performance:
    cache_enabled: true
    cache_ttl: 3600  # 秒
    async_processing: true
    batch_size: 10
  
  # 监控配置
  monitoring:
    enabled: true
    metrics_port: 9090
    logging_level: "INFO"
    include_details: true
  
  # 安全与隐私配置
  security:
    sensitive_info_filtering: false
    sensitive_patterns: ["password", "api_key", "token"]
    encryption_enabled: false
  
  # 集成配置
  integrations:
    langchain:
      enabled: false
      langchain_url: "http://localhost:8080"
    llamaindex:
      enabled: false
      llamaindex_url: "http://localhost:8000"