Agent 上下文压缩实战：Compaction vs Sub-Agents 架构路线之争

"""
Compaction 架构：单窗口上下文压缩
模拟 OpenAI 风格的 Agent 上下文管理
"""

from dataclasses import dataclass
from typing import List, Dict, Optional
import json


@dataclass
class Message:
    role: str  # "user" | "assistant" | "system" | "tool"
    content: str
    tokens: int
    metadata: Optional[Dict] = None


class CompactionAgent:
    """
    Compaction 架构 Agent
    
    核心机制：当上下文快满时，自动压缩历史对话
    """
    
    def __init__(self, max_context_tokens: int = 128000, compaction_threshold: float = 0.8):
        self.max_tokens = max_context_tokens
        self.threshold = compaction_threshold
        self.messages: List[Message] = []
        self.compaction_count = 0
        self.total_compaction_time_ms = 0
    
    def add_message(self, role: str, content: str, tokens: int) -> None:
        """添加消息到上下文"""
        self.messages.append(Message(role=role, content=content, tokens=tokens))
        
        # 检查是否需要压缩
        current_tokens = sum(m.tokens for m in self.messages)
        if current_tokens > self.max_tokens * self.threshold:
            self._compact()
    
    def _compact(self) -> None:
        """
        压缩上下文
        
        策略：
        1. 保留最近 N 条消息（不压缩）
        2. 保留所有未完成的待办事项
        3. 将旧消息压缩为摘要
        """
        # 保留最近 10 条消息
        recent = self.messages[-10:]
        old = self.messages[:-10]
        
        if not old:
            return
        
        # 提取关键信息
        decisions = []
        errors = []
        todos = []
        facts = []
        
        for msg in old:
            content = msg.content.lower()
            if any(kw in content for kw in ["决定", "选择", "decision"]):
                decisions.append(msg.content[:200])
            if any(kw in content for kw in ["错误", "error", "失败", "fail"]):
                errors.append(msg.content[:200])
            if any(kw in content for kw in ["待办", "todo", "需要"]):
                todos.append(msg.content[:200])
            if msg.role == "tool" and "result" in content:
                facts.append(msg.content[:200])
        
        # 生成压缩摘要
        summary_parts = []
        if decisions:
            summary_parts.append(f"## 关键决策\n" + "\n".join(f"- {d}" for d in decisions[:5]))
        if errors:
            summary_parts.append(f"## 遇到的错误\n" + "\n".join(f"- {e}" for e in errors[:5]))
        if todos:
            summary_parts.append(f"## 待办事项\n" + "\n".join(f"- {t}" for t in todos[:5]))
        if facts:
            summary_parts.append(f"## 工具执行结果\n" + "\n".join(f"- {f}" for f in facts[:5]))
        
        summary_content = "\n\n".join(summary_parts) if summary_parts else "[历史对话已压缩，无关键信息]"
        
        # 用摘要替换旧消息
        summary_tokens = len(summary_content.split()) * 1.3
        summary_msg = Message(
            role="system",
            content=f"[上下文压缩 #{self.compaction_count + 1}]\n\n{summary_content}",
            tokens=int(summary_tokens),
            metadata={"compaction": True, "compressed_messages": len(old)}
        )
        
        self.messages = [summary_msg] + recent
        self.compaction_count += 1
        
        # 输出压缩统计
        old_tokens = sum(m.tokens for m in old)
        new_tokens = int(summary_tokens)
        saved = old_tokens - new_tokens
        print(f"[Compaction #{self.compaction_count}] "
              f"压缩 {len(old)} 条消息 | "
              f"{old_tokens} → {new_tokens} tokens | "
              f"节约 {saved} tokens ({saved/old_tokens*100:.0f}%)")
    
    def get_context(self) -> List[Message]:
        """获取当前上下文"""
        return self.messages
    
    def stats(self) -> Dict:
        """统计信息"""
        total_tokens = sum(m.tokens for m in self.messages)
        return {
            "messages": len(self.messages),
            "tokens": total_tokens,
            "utilization": f"{total_tokens / self.max_tokens * 100:.1f}%",
            "compactions": self.compaction_count,
        }


# 演示
if __name__ == "__main__":
    agent = CompactionAgent(max_context_tokens=1000, compaction_threshold=0.7)
    
    # 模拟长对话
    conversations = [
        ("user", "帮我分析这个项目的架构", 50),
        ("assistant", "好的，让我先读取项目结构...", 80),
        ("tool", "result: 项目包含 15 个模块，主要依赖 React + Node.js", 60),
        ("assistant", "发现项目采用前后端分离架构。决定：先分析后端 API 层", 70),
        ("user", "重点看认证模块", 30),
        ("assistant", "正在分析认证模块... 发现使用 JWT + OAuth2.0", 60),
        ("tool", "result: 认证模块有 3 个漏洞需要修复", 50),
        ("assistant", "错误：读取数据库配置时权限不足", 40),
        ("assistant", "决定：使用管理员权限重试", 30),
        ("tool", "result: 成功获取配置，发现 2 个待修复的安全问题", 50),
        ("assistant", "待办：1. 修复 SQL 注入漏洞 2. 更新 JWT 过期策略", 60),
        ("user", "继续分析前端模块", 30),
        # ... 更多对话触发压缩
    ]
    
    for role, content, tokens in conversations:
        agent.add_message(role, content, tokens)
        print(f"  + {role}: {content[:40]}... ({tokens} tokens)")
    
    print(f"\n最终状态: {agent.stats()}")

"""
Sub-Agents 架构：分而治之的 Agent 编排
模拟 Anthropic Claude Code 风格的子代理系统
"""

import asyncio
from dataclasses import dataclass
from typing import List, Dict, Optional
from enum import Enum


class SubAgentType(Enum):
    TASK = "task"          # 一次性任务
    PERSISTENT = "persistent"  # 长期运行


@dataclass
class SubAgentResult:
    agent_id: str
    task: str
    result: str
    tokens_used: int
    latency_ms: float
    success: bool


@dataclass
class SubAgent:
    id: str
    type: SubAgentType
    task: str
    model: str
    status: str = "pending"  # pending | running | completed | failed
    result: Optional[str] = None
    tokens_used: int = 0


class SubAgentOrchestrator:
    """
    Sub-Agents 编排器
    
    核心机制：spawn 多个子代理并行执行，整合结果
    """
    
    def __init__(self, max_concurrent: int = 5):
        self.max_concurrent = max_concurrent
        self.agents: List[SubAgent] = []
        self.completed_results: List[SubAgentResult] = []
    
    async def spawn(self, task: str, model: str = "claude-3.5-sonnet",
                    agent_type: SubAgentType = SubAgentType.TASK) -> str:
        """Spawn 一个子代理"""
        agent_id = f"agent-{len(self.agents) + 1:03d}"
        agent = SubAgent(
            id=agent_id,
            type=agent_type,
            task=task,
            model=model,
        )
        self.agents.append(agent)
        print(f"[Spawn] {agent_id}: {task[:50]}... (model: {model})")
        return agent_id
    
    async def execute_parallel(self, tasks: List[Dict]) -> List[SubAgentResult]:
        """
        并行执行多个子任务
        
        tasks: [{"task": "...", "model": "..."}, ...]
        """
        # Spawn 所有子代理
        agent_ids = []
        for t in tasks:
            aid = await self.spawn(
                task=t["task"],
                model=t.get("model", "claude-3.5-sonnet"),
            )
            agent_ids.append(aid)
        
        # 并行执行（受并发限制）
        semaphore = asyncio.Semaphore(self.max_concurrent)
        
        async def run_agent(agent_id: str):
            async with semaphore:
                agent = next(a for a in self.agents if a.id == agent_id)
                agent.status = "running"
                
                # 模拟执行
                await asyncio.sleep(0.1)  # 模拟延迟
                agent.result = f"[{agent.model}] 完成: {agent.task[:30]}..."
                agent.tokens_used = 500 + len(agent.task) * 2
                agent.status = "completed"
                
                result = SubAgentResult(
                    agent_id=agent.id,
                    task=agent.task,
                    result=agent.result,
                    tokens_used=agent.tokens_used,
                    latency_ms=100 + len(agent.task) * 0.5,
                    success=True,
                )
                self.completed_results.append(result)
                return result
        
        # 等待所有完成
        results = await asyncio.gather(*[run_agent(aid) for aid in agent_ids])
        return list(results)
    
    def synthesize(self, results: List[SubAgentResult]) -> str:
        """整合所有子代理的结果"""
        summary_parts = ["# 综合结果\n"]
        
        for r in results:
            status = "✅" if r.success else "❌"
            summary_parts.append(f"## {status} {r.agent_id}: {r.task[:40]}...")
            summary_parts.append(f"   结果: {r.result}")
            summary_parts.append(f"   Token: {r.tokens_used} | 延迟: {r.latency_ms:.0f}ms\n")
        
        total_tokens = sum(r.tokens_used for r in results)
        total_time = max(r.latency_ms for r in results)  # 并行，取最长
        summary_parts.append(f"---\n总计: {total_tokens} tokens | 总时间: {total_time:.0f}ms (并行)")
        
        return "\n".join(summary_parts)
    
    def stats(self) -> Dict:
        total_tokens = sum(r.tokens_used for r in self.completed_results)
        return {
            "total_agents": len(self.agents),
            "completed": len(self.completed_results),
            "failed": len([r for r in self.completed_results if not r.success]),
            "total_tokens": total_tokens,
            "avg_latency_ms": sum(r.latency_ms for r in self.completed_results) / max(len(self.completed_results), 1),
        }


# 演示
async def main():
    orchestrator = SubAgentOrchestrator(max_concurrent=3)
    
    # 拆解任务
    tasks = [
        {"task": "分析后端 API 架构，找出所有 REST 端点", "model": "claude-3.5-sonnet"},
        {"task": "审查前端 React 组件，检查性能问题", "model": "claude-3.5-sonnet"},
        {"task": "扫描安全漏洞：SQL 注入、XSS、CSRF", "model": "claude-3-opus"},
        {"task": "生成 API 文档（OpenAPI 格式）", "model": "claude-3.5-haiku"},
        {"task": "优化数据库查询，找出慢查询", "model": "claude-3.5-sonnet"},
    ]
    
    print("=== Sub-Agents 并行执行 ===\n")
    results = await orchestrator.execute_parallel(tasks)
    
    print("\n=== 整合结果 ===\n")
    synthesis = orchestrator.synthesize(results)
    print(synthesis)
    
    print(f"\n=== 统计: {orchestrator.stats()} ===")


asyncio.run(main())

"""
混合架构：Compaction + Sub-Agents
外层 Sub-Agents 并行，内层 Compaction 管理上下文
"""

import asyncio
from dataclasses import dataclass
from typing import List, Dict


@dataclass
class SubAgentConfig:
    task: str
    model: str
    max_context_tokens: int = 32000
    compaction_threshold: float = 0.8


class HybridAgent:
    """
    混合架构 Agent
    
    - 主 Agent：任务拆解 + 结果整合（Sub-Agents 模式）
    - 子代理：独立执行 + 内部 Compaction
    """
    
    def __init__(self):
        self.results: List[Dict] = []
    
    async def execute(self, task: str) -> str:
        """执行复杂任务"""
        # Step 1: 拆解任务
        subtasks = self._decompose(task)
        
        # Step 2: 为每个子任务 spawn 子代理（Sub-Agents 模式）
        configs = [
            SubAgentConfig(task=t, model=self._select_model(t))
            for t in subtasks
        ]
        
        # Step 3: 并行执行（每个子代理内部用 Compaction）
        results = await asyncio.gather(*[
            self._run_subagent_with_compaction(cfg) for cfg in configs
        ])
        
        # Step 4: 整合结果
        return self._synthesize(task, results)
    
    def _decompose(self, task: str) -> List[str]:
        """将复杂任务拆解为子任务"""
        # 实际实现中，这里应该用 LLM 进行智能拆解
        return [
            f"子任务 1: 分析 {task} 的结构",
            f"子任务 2: 评估 {task} 的风险",
            f"子任务 3: 生成 {task} 的优化方案",
        ]
    
    def _select_model(self, subtask: str) -> str:
        """根据子任务特征选择模型"""
        if "代码" in subtask or "code" in subtask.lower():
            return "deepseek-coder-32b"
        elif "安全" in subtask or "风险" in subtask:
            return "claude-opus-4.7"
        else:
            return "qwen-32b"
    
    async def _run_subagent_with_compaction(self, config: SubAgentConfig) -> Dict:
        """
        运行一个带 Compaction 的子代理
        
        子代理有独立的上下文窗口，当上下文快满时自动压缩
        """
        context_tokens = 0
        messages = []
        compactions = 0
        
        # 模拟子代理执行多步骤任务
        for step in range(15):  # 15 步，会触发压缩
            # 模拟执行
            msg_tokens = 2000 + step * 500
            messages.append({"step": step, "tokens": msg_tokens})
            context_tokens += msg_tokens
            
            # 检查是否需要压缩
            if context_tokens > config.max_context_tokens * config.compaction_threshold:
                # 压缩：保留关键信息，丢弃细节
                context_tokens = int(context_tokens * 0.3)  # 压缩到 30%
                compactions += 1
        
        return {
            "task": config.task,
            "model": config.model,
            "total_tokens": context_tokens,
            "compactions": compactions,
            "success": True,
            "result": f"完成: {config.task[:30]}...",
        }
    
    def _synthesize(self, original_task: str, results: List[Dict]) -> str:
        """整合所有子代理的结果"""
        parts = [f"# 任务完成: {original_task[:50]}...\n"]
        for r in results:
            parts.append(f"- ✅ {r['task'][:40]}... (model: {r['model']}, compactions: {r['compactions']})")
        parts.append(f"\n总计: {len(results)} 个子代理完成")
        return "\n".join(parts)


# 演示
async def main():
    agent = HybridAgent()
    result = await agent.execute("全面审查项目代码质量、安全性和性能")
    print(result)


asyncio.run(main())