上下文与记忆：Agent 为什么会在长任务中失去方向，以及如何设计状态管理

from dataclasses import dataclass, field
from typing import Optional
from enum import Enum

class TaskStatus(Enum):
    IN_PROGRESS = "in_progress"
    BLOCKED = "blocked"
    COMPLETED = "completed"
    FAILED = "failed"

@dataclass
class SubTask:
    id: str
    description: str
    status: TaskStatus
    result: Optional[str] = None
    tool_used: Optional[str] = None

@dataclass
class TaskState:
    goal: str                              # 原始目标，永远不变
    current_focus: str                     # 当前正在做什么
    completed_subtasks: list[SubTask] = field(default_factory=list)
    pending_subtasks: list[SubTask] = field(default_factory=list)
    key_findings: list[str] = field(default_factory=list)   # 关键发现，蒸馏后保留
    blockers: list[str] = field(default_factory=list)       # 当前阻碍
    decisions_made: list[str] = field(default_factory=list) # 已经做出的决策

    def to_context_string(self) -> str:
        """把任务状态转成注入 context 的字符串"""
        lines = [
            f"## 当前任务状态",
            f"**目标**：{self.goal}",
            f"**当前焦点**：{self.current_focus}",
            "",
        ]
        if self.completed_subtasks:
            lines.append("**已完成步骤**：")
            for st in self.completed_subtasks:
                result_summary = f" → {st.result[:100]}..." if st.result else ""
                lines.append(f"  - ✅ {st.description}{result_summary}")
        if self.pending_subtasks:
            lines.append("**待执行步骤**：")
            for st in self.pending_subtasks:
                lines.append(f"  - ⬜ {st.description}")
        if self.key_findings:
            lines.append("**关键发现**：")
            for f in self.key_findings:
                lines.append(f"  - {f}")
        if self.decisions_made:
            lines.append("**已确定决策**：")
            for d in self.decisions_made:
                lines.append(f"  - {d}")
        if self.blockers:
            lines.append("**当前阻碍**：")
            for b in self.blockers:
                lines.append(f"  - ⚠️ {b}")
        return "\n".join(lines)

def build_messages(task_state: TaskState, user_input: str, history: list) -> list:
    system = f"""你是一个任务执行 Agent。

{task_state.to_context_string()}

---
根据以上任务状态，继续执行任务。如果当前焦点已完成，更新任务状态并推进到下一步。
"""
    messages = [{"role": "system", "content": system}]

    # 只保留最近 N 轮历史，不是全部
    recent_history = history[-10:]  # 后面会详细说这个策略
    messages.extend(recent_history)
    messages.append({"role": "user", "content": user_input})
    return messages

async def compress_tool_result(
    tool_name: str,
    raw_result: str,
    task_context: str,
    max_tokens: int = 500
) -> str:
    """
    在工具结果进入 context 之前，用一次轻量级 LLM 调用把它压缩成
    与当前任务相关的关键信息摘要。
    """
    if len(raw_result.split()) < 200:
        # 短结果不需要压缩
        return raw_result

    response = await client.chat.completions.create(
        model="gpt-4o-mini",   # 压缩用便宜模型就够了
        max_tokens=max_tokens,
        messages=[
            {
                "role": "system",
                "content": (
                    "你的任务是提取工具返回结果中与当前任务直接相关的关键信息，"
                    "去除无关内容，输出简洁的摘要。"
                    "保留具体的数据、结论、错误信息。不要添加分析，只做提取和压缩。"
                )
            },
            {
                "role": "user",
                "content": (
                    f"当前任务上下文：{task_context}\n\n"
                    f"工具名称：{tool_name}\n"
                    f"工具返回结果：\n{raw_result}"
                )
            }
        ]
    )
    return response.choices[0].message.content

class ContextWindow:
    def __init__(self, max_messages: int = 20, max_tokens: int = 8000):
        self.messages: list[dict] = []
        self.max_messages = max_messages
        self.max_tokens = max_tokens
        self.archived: list[dict] = []  # 滑出窗口的消息存档

    def add(self, message: dict):
        self.messages.append(message)
        self._trim()

    def _trim(self):
        """
        裁剪策略：
        1. 如果消息数量超过上限，把最老的消息移入 archived
        2. 工具调用消息（assistant + tool）要成对保留或成对移除，
           避免出现孤立的 tool_result 没有对应的 tool_call
        """
        while len(self.messages) > self.max_messages:
            # 找到第一对完整的 tool_call + tool_result 或普通消息
            if self.messages[0]["role"] == "assistant" and \
               hasattr(self.messages[0].get("content"), "__iter__") and \
               self.messages[1]["role"] == "tool":
                # 成对移除 tool_call + tool_result
                self.archived.append(self.messages.pop(0))
                self.archived.append(self.messages.pop(0))
            else:
                self.archived.append(self.messages.pop(0))

    def get_messages(self) -> list[dict]:
        return self.messages

    def search_archived(self, keyword: str) -> list[dict]:
        """在归档历史中检索相关消息，用于情节记忆检索"""
        return [m for m in self.archived
                if keyword.lower() in str(m.get("content", "")).lower()]

async def create_execution_plan(goal: str, context: str) -> dict:
    """任务开始前，先生成结构化执行计划"""
    response = await client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": "你是一个任务规划专家。根据给定目标和上下文，生成结构化执行计划。"
            },
            {
                "role": "user",
                "content": f"""
目标：{goal}
上下文：{context}

请生成执行计划，格式如下（JSON）：
{{
  "goal_summary": "一句话总结目标",
  "steps": [
    {{
      "id": "step_1",
      "description": "步骤描述",
      "depends_on": [],
      "success_criteria": "什么情况下这步算完成",
      "tools_likely_needed": ["tool_name"]
    }}
  ],
  "key_constraints": ["约束条件"],
  "definition_of_done": "整体任务完成的判断标准"
}}
"""
            }
        ]
    )
    return json.loads(response.choices[0].message.content)


async def checkpoint(plan: dict, completed_step_id: str,
                     result_summary: str, agent_notes: str) -> dict:
    """每完成一步，调用检查点更新计划状态"""
    # 更新步骤状态
    for step in plan["steps"]:
        if step["id"] == completed_step_id:
            step["status"] = "completed"
            step["result_summary"] = result_summary
            break

    # 让模型评估是否需要调整计划
    response = await client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": f"""
当前执行计划：
{json.dumps(plan, ensure_ascii=False, indent=2)}

刚完成步骤：{completed_step_id}
执行结果摘要：{result_summary}
执行者备注：{agent_notes}

请判断：
1. 是否需要调整后续步骤？（比如某个结果改变了后续的方向）
2. 如果需要调整，给出调整后的步骤列表
3. 整体任务是否已经完成？

以 JSON 格式返回：
{{
  "needs_replan": true/false,
  "updated_steps": [...] or null,
  "task_completed": true/false,
  "next_step_id": "step_x" or null,
  "notes": "调整原因或完成说明"
}}
"""
            }
        ]
    )
    checkpoint_result = json.loads(response.choices[0].message.content)

    if checkpoint_result["needs_replan"] and checkpoint_result["updated_steps"]:
        plan["steps"] = checkpoint_result["updated_steps"]

    return plan, checkpoint_result