上手实操:跑通一个最小 Reflexion 循环

"""reflexion.py — 最小可运行的 Reflexion 循环(对应 Shinn et al. 2023, Algorithm 1)"""
from dataclasses import dataclass, field
from typing import Protocol

# ── LLM 接口:换成真实 API 即可接 GPT-4 / Claude ──────────────
class LLM(Protocol):
    def complete(self, prompt: str) -> str: ...

# ── 任务定义 ────────────────────────────────────────────────
@dataclass
class Task:
    spec: str            # 自然语言任务描述
    entry_point: str     # 待实现的函数名
    tests: list[str]     # 充当 Evaluator 判据的 assert 语句

# ── Actor:生成代码(产出 = trajectory)──────────────────────
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)

# ── Evaluator:跑测试,返回 (是否通过, 错误信息) ────────────────
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}"

# ── Self-Reflection:把失败转成第一人称语言教训 ──────────────────
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)

# ── 情景记忆:Ω 滑动窗口(论文 Ω=1-3) ─────────────────────────
@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 条

# ── 主循环:Algorithm 1 ──────────────────────────────────────
def reflexion(task: Task, actor: Actor, evaluator: Evaluator,
              reflector: SelfReflection, memory: EpisodicMemory,
              max_trials: int = 3) -> tuple[str, bool, int]:
    code = ""
    for t in range(max_trials):
        code = actor.act(task, memory.reflections)        # 生成 trajectory
        passed, error = evaluator.evaluate(code, task)     # 打分(reward）
        if passed:
            return code, True, t                            # 停止条件:通过
        sr = reflector.reflect(code, error, task)           # 反思
        memory.add(sr)                                      # 写入长期记忆
    return code, False, max_trials                          # 用尽 trials

class ScriptedLLM:
    """按调用顺序返回预设输出,模拟 Actor 从 buggy → fixed 的演化。
    真实使用时替换为 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",     # 偶数长度
    ],
)
# 三次 complete 调用:trial0 代码 → 反思文本 → trial1 代码
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, trials = reflexion(task, Actor(llm), Evaluator(), SelfReflection(llm), EpisodicMemory())
print(f"passed={ok}  trials_used={trials}")
print(code)

passed=True  trials_used=1
def median(nums):
    s = sorted(nums)
    n = len(s)
    if n % 2:
        return s[n//2]
    return (s[n//2 - 1] + s[n//2]) / 2

class Actor:
    def __init__(self, llm: LLM):
        self.llm = llm
    def act(self, task: Task, reflections: list[str]) -> str:
        # TODO-A1:把 reflections 列表拼成一段可读的 memory 文本
        memory_block = ...
        # TODO-A2:仅当有反思时,才把 memory 段落带标题地插进 prompt
        prompt = (
            f"实现函数 {task.entry_point}。\n任务:{task.spec}\n"
            + ...
            + "\n只输出函数定义代码。"
        )
        return self.llm.complete(prompt)

class SelfReflection:
    def __init__(self, llm: LLM):
        self.llm = llm
    def reflect(self, code: str, error: str, task: Task) -> str:
        # TODO-B:写一个 prompt,喂入 任务 / 代码 / 错误,
        #         引导模型用第一人称、定位哪一步错、给出下次的具体改法
        prompt = ...
        return self.llm.complete(prompt)