综合实战：给内部代码助手加一个自我修正的反思循环

"""reflexion.py — Reflexion 循环 + capstone 扩展(早停 + 反思护栏)"""
from dataclasses import dataclass, field
from typing import Protocol

class LLM(Protocol):
    def complete(self, prompt: str) -> str: ...

@dataclass
class Task:
    spec: str
    entry_point: str
    tests: list[str]

class Actor:
    def __init__(self, llm: LLM):
        self.llm = llm
    def act(self, task: Task, reflections: list[str]) -> str:
        memory_block = "\n".join(f"- {r}" for r in reflections)
        prompt = (
            f"实现函数 {task.entry_point}。\n任务:{task.spec}\n"
            + (f"\n过去尝试的教训(避免重犯):\n{memory_block}\n" if reflections else "")
            + "\n只输出函数定义代码。"
        )
        return self.llm.complete(prompt)

class Evaluator:
    def evaluate(self, code: str, task: Task) -> tuple[bool, str]:
        ns: dict = {}
        try:
            exec(code, ns)                  # 定义函数
            for assertion in task.tests:    # 逐条跑测试
                exec(assertion, ns)
            return True, ""
        except Exception as e:
            return False, f"{type(e).__name__}: {e}"

class SelfReflection:
    def __init__(self, llm: LLM):
        self.llm = llm
    def reflect(self, code: str, error: str, task: Task) -> str:
        prompt = (
            f"任务:{task.spec}\n我写的代码:\n{code}\n"
            f"测试失败:{error}\n"
            "用第一人称写一条简短教训:我哪一步错了、下次具体怎么改。"
        )
        return self.llm.complete(prompt)

@dataclass
class EpisodicMemory:
    capacity: int = 3
    reflections: list[str] = field(default_factory=list)
    def add(self, sr: str) -> None:
        self.reflections.append(sr)
        self.reflections = self.reflections[-self.capacity:]   # 只留最近 capacity 条

@dataclass
class TrialLog:
    """每个 trial 的可观测记录,供验收 checklist 检查行为。"""
    index: int
    passed: bool
    error: str
    reflected: bool          # 这一轮到底有没有触发反思
    skip_reason: str = ""    # 没反思时,护栏给出的原因

def reflexion(task: Task, actor: Actor, evaluator: Evaluator,
              reflector: SelfReflection, memory: EpisodicMemory,
              max_trials: int = 3,
              patience: int = 2) -> tuple[str, bool, list[TrialLog]]:
    """两道护栏:
       1) 反思护栏: Evaluator 判通过就立刻停、绝不在已通过的解上"再反思一版"
          —— 堵 04 章 reflection-hallucination 的 correct→incorrect 翻转。
       2) 早停: 连续 patience 次"同一个错误信号" → 反思没在推进, 提前停,
          不耗满 max_trials —— 堵 04 章 convergence 的 plateau / 白烧成本。
    """
    code = ""
    logs: list[TrialLog] = []
    last_error = None
    repeat = 0
    for t in range(max_trials):
        code = actor.act(task, memory.reflections)          # 生成 trajectory
        passed, error = evaluator.evaluate(code, task)       # 打分(reward)
        if passed:
            # 护栏 1: 通过即停, 不在正确解上反思(避免把对的改错)
            logs.append(TrialLog(t, True, "", reflected=False,
                                 skip_reason="已通过,无需反思"))
            return code, True, logs
        # 早停判定: 同一个错误信号重复 → 反思没在推进
        repeat = repeat + 1 if error == last_error else 0
        last_error = error
        if repeat >= patience - 1 and t < max_trials - 1:
            logs.append(TrialLog(t, False, error, reflected=False,
                                 skip_reason=f"连续 {repeat + 1} 次相同错误,反思未推进,提前停止"))
            return code, False, logs
        sr = reflector.reflect(code, error, task)            # 反思: 做 credit assignment
        memory.add(sr)                                       # 写入长期记忆
        logs.append(TrialLog(t, False, error, reflected=True))
    return code, False, logs                                 # 用尽 trials

class ScriptedLLM:
    """按调用顺序返回预设输出。真实使用时替换为 OpenAI / Anthropic 客户端封装。"""
    def __init__(self, outputs: list[str]):
        self._outputs = list(outputs)
    def complete(self, prompt: str) -> str:
        return self._outputs.pop(0) if self._outputs else ""

task = Task(
    spec="返回数字列表的中位数(偶数长度取中间两数平均)",
    entry_point="median",
    tests=[
        "assert median([1, 3, 2]) == 2",
        "assert median([1, 2, 3, 4]) == 2.5",     # 偶数长度
    ],
)

# 演示 1 — 失败→反思→通过(checklist 第 2 条)
llm = ScriptedLLM([
    "def median(nums):\n    s = sorted(nums)\n    return s[len(s)//2]",
    "我只取了中间单个元素,没处理偶数长度;下次对偶数长度取中间两数的平均。",
    "def median(nums):\n    s = sorted(nums)\n    n = len(s)\n    if n % 2:\n        return s[n//2]\n    return (s[n//2 - 1] + s[n//2]) / 2",
])
code, ok, logs = reflexion(task, Actor(llm), Evaluator(), SelfReflection(llm), EpisodicMemory())
print(f"[1] passed={ok} trials_run={len(logs)}")
assert ok and logs[0].reflected and not logs[-1].reflected   # 首轮反思过、末轮通过不再反思

# 演示 2 — 反思没推进就早停(checklist 第 3 条)
same_buggy = "def median(nums):\n    return nums[0]"          # 永远错且报同一个 AssertionError
llm2 = ScriptedLLM([same_buggy, "教训1", same_buggy, "教训2", same_buggy, "教训3"])
_, ok2, logs2 = reflexion(task, Actor(llm2), Evaluator(),
                          SelfReflection(llm2), EpisodicMemory(),
                          max_trials=5, patience=2)
print(f"[2] passed={ok2} trials_run={len(logs2)} (max_trials=5)")
assert not ok2 and len(logs2) < 5 and "提前停止" in logs2[-1].skip_reason   # 远未跑满就停

# 演示 3 — 首轮通过,护栏阻止反思(checklist 第 1 条)
llm3 = ScriptedLLM([
    "def median(nums):\n    s = sorted(nums)\n    n = len(s)\n    if n % 2:\n        return s[n//2]\n    return (s[n//2-1]+s[n//2])/2",
])
_, ok3, logs3 = reflexion(task, Actor(llm3), Evaluator(), SelfReflection(llm3), EpisodicMemory())
print(f"[3] passed={ok3} trials_run={len(logs3)} reflected={logs3[0].reflected}")
assert ok3 and len(logs3) == 1 and not logs3[0].reflected     # 首轮就过、绝不反思

print("ALL ASSERTIONS PASSED")

[1] passed=True trials_run=2
[2] passed=False trials_run=2 (max_trials=5)
[3] passed=True trials_run=1 reflected=False
ALL ASSERTIONS PASSED