# Simple Retry核心模式
MAX_RETRIES = 3

async def agent_with_retry(task, llm_call):
    error_history = []
    for attempt in range(MAX_RETRIES):
        result = await llm_call(task, previous_errors=error_history)
        if validate(result):  # 结构化校验（编译/格式/规则）
            return result
        error_history.append({
            "attempt": attempt + 1,
            "output": result,
            "error": get_validation_error(result)
        })
    return None  # 所有重试耗尽

# Self-Refine核心循环
async def self_refine(task, llm_call, max_rounds=3):
    draft = await llm_call(f"Generate: {task}")

    for round in range(max_rounds):
        # 自我批评阶段
        critique = await llm_call(
            f"Critique this output for the task '{task}':\n"
            f"Output: {draft}\n"
            f"Identify specific issues and suggest improvements."
        )

        # 基于批评的精炼阶段
        refined = await llm_call(
            f"Based on this critique:\n{critique}\n"
            f"Improve the original output:\n{draft}"
        )

        if refined == draft:  # 收敛检测
            break
        draft = refined

    return draft

# Reflexion核心架构
class ReflexionAgent:
    def __init__(self, llm, evaluator):
        self.llm = llm
        self.evaluator = evaluator
        self.reflection_memory = []  # 持久化的反思记忆

    async def execute_with_reflexion(self, task, max_trials=5):
        for trial in range(max_trials):
            # 构建包含历史反思的提示
            context = self._build_context(task)
            action = await self.llm(context)

            # 评估器判断是否成功
            eval_result = await self.evaluator(action, task)
            if eval_result.success:
                return action

            # 关键步骤：生成语言反思
            reflection = await self.llm(
                f"You failed this task. Here was your attempt:\n"
                f"{action}\n"
                f"The evaluation result was: {eval_result.feedback}\n"
                f"Reflect on what went wrong and how to improve. "
                f"Be specific about what to avoid next time."
            )

            # 将反思存入持久化记忆
            self.reflection_memory.append({
                "trial": trial,
                "task": task,
                "attempt": action,
                "feedback": eval_result.feedback,
                "reflection": reflection
            })

        return None  # 所有尝试失败

# LATS简化架构
class LATSAgent:
    def __init__(self, llm, evaluator, max_depth=5, n_candidates=3):
        self.llm = llm
        self.evaluator = evaluator
        self.max_depth = max_depth
        self.n_candidates = n_candidates

    async def search(self, task):
        root = Node(state=task, parent=None)

        for _ in range(self.max_depth):
            # 选择：UCB公式选择最有潜力的节点
            leaf = self._select(root)

            # 扩展：生成n个候选下一步
            children = []
            for _ in range(self.n_candidates):
                action = await self.llm(
                    f"Given state: {leaf.state}\n"
                    f"Generate next action. {self._get_reflections(leaf)}"
                )
                child = Node(state=action, parent=leaf)
                children.append(child)

            # 评估：对每个候选打分
            for child in children:
                child.value = await self.evaluator(child.state)

            # 反向传播：更新路径上所有节点的价值
            self._backpropagate(children)

            # 检查是否找到解
            best = max(children, key=lambda c: c.value)
            if best.value >= SUCCESS_THRESHOLD:
                return self._trace_path(best)

        return self._trace_path(self._best_leaf(root))

# Verifier-in-the-Loop生产架构
class VerifiedAgent:
    def __init__(self, actor_llm, verifier_llm, max_corrections=2):
        self.actor = actor_llm
        self.verifier = verifier_llm  # 可以是更小/更快的模型
        self.max_corrections = max_corrections

    async def execute(self, task):
        result = await self.actor(f"Execute: {task}")

        for attempt in range(self.max_corrections):
            # 独立验证（不信任Agent自我评估）
            verdict = await self.verifier(
                f"Evaluate this output for correctness:\n"
                f"Task: {task}\n"
                f"Output: {result}\n"
                f"Respond with PASS or FAIL and specific issues."
            )

            if verdict.status == "PASS":
                return result

            # 验证器提供具体纠错指导
            result = await self.actor(
                f"Your previous output had issues:\n"
                f"{verdict.feedback}\n"
                f"Original task: {task}\n"
                f"Please correct."
            )

        return result  # 降级：返回最后一次尝试

# Agent-R训练范式（简化）
# 不同于运行时纠错，这是训练数据生成流程

def generate_training_data(model, task_set):
    data = []
    for task in task_set:
        # 步骤1：模型执行任务
        trajectory = model.execute(task)

        # 步骤2：注入错误（随机干预某些步骤）
        corrupted = inject_errors(trajectory, error_rate=0.3)

        # 步骤3：让模型识别并修正自己的错误
        correction_trajectory = model.correct(corrupted)

        # 步骤4：构建 (错误轨迹, 修正轨迹) 训练对
        data.append({
            "input": format_with_error_marker(corrupted),
            "target": format_correction(correction_trajectory)
        })

    return data

# 推理时：模型仅需一次调用，内化的纠错能力自动生效

graph TD
    A[Agent执行任务] --> B{输出校验}
    B -->|通过| C[返回结果]
    B -->|失败| D{选择纠错策略}

    D -->|E1: Simple Retry| E[携带错误信息重新调用]
    E --> B

    D -->|E2: Self-Refine| F[自我批评 → 自我修改]
    F --> B

    D -->|E3: Reflexion| G[生成语言反思 → 存入记忆]
    G --> H[构建反思增强Prompt]
    H --> A

    D -->|E4: LATS| I[生成多候选 → MCTS搜索]
    I --> J[评估节点价值]
    J -->|高价值| K[选择最优路径]
    J -->|低价值| L[剪枝+反思回溯]
    L --> I
    K --> B

    D -->|E5: Verifier-in-Loop| M[独立验证器评估]
    M -->|PASS| C
    M -->|FAIL + 反馈| N[按指导修正]
    N --> M

    D -->|E6: Post-Training| O[模型内化纠错]
    O -->|一次调用自带纠错| C

    style A fill:#4A90D9,color:#fff
    style B fill:#F5A623,color:#fff
    style D fill:#7B68EE,color:#fff
    style C fill:#7ED321,color:#fff
    style G fill:#BD10E0,color:#fff
    style I fill:#BD10E0,color:#fff
    style M fill:#50E3C2,color:#fff
    style O fill:#50E3C2,color:#fff

graph LR
    subgraph Token开销倍率 vs 准确率提升
        A[Post-Training SC<br/>1.1x / +7%]
        B[Simple Retry<br/>1.6x / +3%]
        C[Verifier-in-Loop<br/>2.2x / +12%]
        D[Self-Refine<br/>4.5x / +10%]
        E[Reflexion<br/>8.0x / +15%]
        F[LATS<br/>25.0x / +20%]
    end

    style A fill:#7ED321,color:#000
    style C fill:#7ED321,color:#000
    style F fill:#D0021B,color:#fff

graph LR
    subgraph "METR任务时间窗口（50%成功率）"
        A["2023 GPT-4<br/>~10分钟"] --> B["2024 Claude 3.5<br/>~1小时"]
        B --> C["2025 Claude 3.5 Sonnet<br/>~8小时"]
        C --> D["2025末 o3<br/>~100小时"]
        D --> E["2026初<br/>~300小时"]
        E --> F["2026.04<br/>~700小时"]
    end

    style A fill:#95a5a6,color:#fff
    style F fill:#27ae60,color:#fff