from typing import List, Dict, Any
# 假设 llm_client.py 文件已存在，并从中导入 HelloAgentsLLM 类
from agent_client import HelloAgentsLLM

# --- 模块 1: 记忆模块 ---

class Memory:
    """
    一个简单的短期记忆模块，用于存储智能体的行动与反思轨迹。
    """
    def __init__(self):
        # 初始化一个空列表来存储所有记录
        self.records: List[Dict[str, Any]] = []

    def add_record(self, record_type: str, content: str):
        """
        向记忆中添加一条新记录。

        参数:
        - record_type (str): 记录的类型 ('execution' 或 'reflection')。
        - content (str): 记录的具体内容 (例如，生成的代码或反思的反馈)。
        """
        self.records.append({"type": record_type, "content": content})
        print(f"📝 记忆已更新，新增一条 '{record_type}' 记录。")

    def get_trajectory(self) -> str:
        """
        将所有记忆记录格式化为一个连贯的字符串文本，用于构建提示词。
        """
        trajectory = ""
        for record in self.records:
            if record['type'] == 'execution':
                trajectory += f"--- 上一轮尝试 (代码) ---\n{record['content']}\n\n"
            elif record['type'] == 'reflection':
                trajectory += f"--- 评审员反馈 ---\n{record['content']}\n\n"
        return trajectory.strip()

    def get_last_execution(self) -> str:
        """
        获取最近一次的执行结果 (例如，最新生成的代码)。
        """
        for record in reversed(self.records):
            if record['type'] == 'execution':
                return record['content']
        return None

# --- 模块 2: Reflection 智能体 ---

# 1. 初始执行提示词
INITIAL_PROMPT_TEMPLATE = """
你是一位资深的Python程序员。请根据以下要求，编写一个Python函数。
你的代码必须包含完整的函数签名、文档字符串，并遵循PEP 8编码规范。

要求: {task}

请直接输出代码，不要包含任何额外的解释。
"""

# 2. 反思提示词
REFLECT_PROMPT_TEMPLATE = """
你是一位极其严格的代码评审专家和资深算法工程师，对代码的性能有极致的要求。
你的任务是审查以下Python代码，并专注于找出其在**算法效率**上的主要瓶颈。

# 原始任务:
{task}

# 待审查的代码:
```python
{code}
```

请分析该代码的时间复杂度，并思考是否存在一种**算法上更优**的解决方案来显著提升性能。
如果存在，请清晰地指出当前算法的不足，并提出具体的、可行的改进算法建议（例如，使用筛法替代试除法）。
如果代码在算法层面已经达到最优，才能回答“无需改进”。

请直接输出你的反馈，不要包含任何额外的解释。
"""

# 3. 优化提示词
REFINE_PROMPT_TEMPLATE = """
你是一位资深的Python程序员。你正在根据一位代码评审专家的反馈来优化你的代码。

# 原始任务:
{task}

# 你上一轮尝试的代码:
{last_code_attempt}

# 评审员的反馈:
{feedback}

请根据评审员的反馈，生成一个优化后的新版本代码。
你的代码必须包含完整的函数签名、文档字符串，并遵循PEP 8编码规范。
请直接输出优化后的代码，不要包含任何额外的解释。
"""

class ReflectionAgent:
    def __init__(self, llm_client, max_iterations=3):
        self.llm_client = llm_client
        self.memory = Memory()
        self.max_iterations = max_iterations

    def run(self, task: str):
        print(f"\n--- 开始处理任务 ---\n任务: {task}")

        # --- 1. 初始执行 ---
        print("\n--- 正在进行初始尝试 ---")
        initial_prompt = INITIAL_PROMPT_TEMPLATE.format(task=task)
        initial_code = self._get_llm_response(initial_prompt)
        self.memory.add_record("execution", initial_code)

        # --- 2. 迭代循环：反思与优化 ---
        for i in range(self.max_iterations):
            print(f"\n--- 第 {i+1}/{self.max_iterations} 轮迭代 ---")

            # a. 反思
            print("\n-> 正在进行反思...")
            last_code = self.memory.get_last_execution()
            reflect_prompt = REFLECT_PROMPT_TEMPLATE.format(task=task, code=last_code)
            feedback = self._get_llm_response(reflect_prompt)
            self.memory.add_record("reflection", feedback)

            # b. 检查是否需要停止
            if "无需改进" in feedback or "no need for improvement" in feedback.lower():
                print("\n✅ 反思认为代码已无需改进，任务完成。")
                break

            # c. 优化
            print("\n-> 正在进行优化...")
            refine_prompt = REFINE_PROMPT_TEMPLATE.format(
                task=task,
                last_code_attempt=last_code,
                feedback=feedback
            )
            refined_code = self._get_llm_response(refine_prompt)
            self.memory.add_record("execution", refined_code)
        
        final_code = self.memory.get_last_execution()
        print(f"\n--- 任务完成 ---\n最终生成的代码:\n{final_code}")
        return final_code

    def _get_llm_response(self, prompt: str) -> str:
        """一个辅助方法，用于调用LLM并获取完整的流式响应。"""
        messages = [{"role": "user", "content": prompt}]
        # 确保能处理生成器可能返回None的情况
        response_text = self.llm_client.think(messages=messages) or ""
        return response_text

if __name__ == '__main__':
    # 1. 初始化LLM客户端 (请确保你的 .env 和 llm_client.py 文件配置正确)
    try:
        llm_client = HelloAgentsLLM()
    except Exception as e:
        print(f"初始化LLM客户端时出错: {e}")
        exit()

    # 2. 初始化 Reflection 智能体，设置最多迭代2轮
    agent = ReflectionAgent(llm_client, max_iterations=2)

    # 3. 定义任务并运行智能体
    task = "编写一个Python函数，找出1到n之间所有的素数 (prime numbers)。"
    agent.run(task)

输出：

(.venv) MacBook-Air section4 % python Reflection.py

--- 开始处理任务 ---
任务: 编写一个Python函数，找出1到n之间所有的素数 (prime numbers)。

--- 正在进行初始尝试 ---
🧠 正在调用 coding-glm-5-free 模型...
✅ 大语言模型响应成功:
```python
from typing import List

def find_primes(n: int) -> List[int]:
    """
    找出1到n之间所有的素数。

    使用埃拉托斯特尼筛法来高效地找出素数。

    Args:
        n (int): 查找素数的上限（包含）。

    Returns:
        List[int]: 包含1到n之间所有素数的列表。
                   如果n < 2，返回空列表。
    """
    if n < 2:
        return []

    # 初始化筛子，假设所有数字都是素数
    is_prime = [True] * (n + 1)
    is_prime[0] = is_prime[1] = False

    # 迭代到n的平方根
    for p in range(2, int(n**0.5) + 1):
        if is_prime[p]:
            # 标记p的倍数为非素数
            # 从p*p开始，因为较小的倍数已经被更小的素数标记过了
            for multiple in range(p * p, n + 1, p):
                is_prime[multiple] = False

    # 收集所有仍然是True的索引
    return [p for p, prime in enumerate(is_prime) if prime]
```
📝 记忆已更新，新增一条 'execution' 记录。

--- 第 1/2 轮迭代 ---

-> 正在进行反思...
🧠 正在调用 coding-glm-5-free 模型...
✅ 大语言模型响应成功:
### 代码审查反馈

**1. 算法复杂度与瓶颈分析**

*   **时间复杂度**: 当前代码使用的是标准的**埃拉托斯特尼筛法**，时间复杂度为 $O(N \log \log N)$。虽然这在一般场景下足够快，但在算法理论层面并非最优。
*   **主要瓶颈 (性能损耗点)**:
    1.  **重复标记 (核心算法缺陷)**: 埃氏筛法存在重复标记合数的问题。例如，合数 `12` 会被素数 `2` 和素数 `3` 分别标记一次。随着 $N$ 增大，这种冗余操作累加起来会显著增加时间开销。
    2.  **内存布局与访问效率 (工程实现缺陷)**: 代码使用了 `List[bool]`。在 Python 中，`list` 存储的是对象的引用（指针），且 `bool` 对象有内存开销。这导致两倍的性能下降：
        *   **内存浪费**: 相比于位数组或字节数组，内存占用高出 8 到 30 倍，导致 CPU 缓存命中率急剧下降。
        *   **指针解引用**: 每次访问 `is_prime[i]` 都需要一次指针解引用，这比直接的数组索引慢得多。

**2. 改进算法建议**

建议使用 **欧拉筛 (线性筛 / Euler's Sieve)** 结合 **`bytearray`** 进行优化。

*   **算法层面 (线性筛)**:
    *   **原理**: 欧拉筛通过保证每个合数**仅被其最小质因子标记一次**，成功将时间复杂度降低到严格的 **$O(N)$**。
    *   **优势**: 彻底消除了埃氏筛法中的重复标记问题。当 $N$ 很大时（例如 $10^8$ 或更高），性能提升非常明显。

*   **工程层面 (内存优化)**:
    *   **数据结构**: 使用 Python 内置的 `bytearray` 替代 `List[bool]`。
    *   **优势**: `bytearray` 在内存中是连续的字节序列，不仅内存占用极小（仅为 list 的 1/8 左右），而且遍历和访问速度极快，能极大地利用 CPU 缓存。

**3. 优化后的代码实现**

```python
from typing import List

def find_primes(n: int) -> List[int]:
    """
    使用欧拉筛（线性筛）找出1到n之间所有的素数。
    时间复杂度: O(N)
    空间复杂度: O(N) (使用 bytearray 优化内存布局)
    """
    if n < 2:
        return []

    # 使用 bytearray 替代 List[bool]，显著降低内存开销并提升访问速度
    # 0 代表非素数，1 代表素数
    is_prime = bytearray(b'\x01') * (n + 1)
    is_prime[0] = is_prime[1] = 0
    
    primes = []
    
    for i in range(2, n + 1):
        # 如果 i 是素数，加入列表
        if is_prime[i]:
            primes.append(i)
            
        # 遍历已找到的素数列表
        for p in primes:
            # 关键溢出检查：如果合数超过 n，停止标记
            if p * i > n:
                break
            
            # 标记合数
            is_prime[p * i] = 0
            
            # 核心逻辑：如果 i 能被 p 整除，说明 i = k * p
            # 那么 i * p = k * p * p。
            # 为了保证每个合数只被最小质因子标记，当 p 整除 i 时，
            # i * p 的最小质因子是 p，后续更大的素数 q 如果标记 i * q，
            # 则 q 不是最小质因子。因此此处必须 break。
            if i % p == 0:
                break
                
    return primes
```
📝 记忆已更新，新增一条 'reflection' 记录。

-> 正在进行优化...
🧠 正在调用 coding-glm-5-free 模型...
✅ 大语言模型响应成功:
```python
from typing import List

def find_primes(n: int) -> List[int]:
    """
    使用欧拉筛（线性筛）结合 bytearray 优化，找出 1 到 n 之间所有的素数。

    相比于埃拉托斯特尼筛法，欧拉筛保证了每个合数只被其最小质因子标记一次，
    从而将时间复杂度降低到严格的 O(N)。同时使用 bytearray 替代 List[bool]
    以减少内存开销并提高访问速度。

    Args:
        n (int): 查找素数的上限（包含）。

    Returns:
        List[int]: 包含 1 到 n 之间所有素数的列表。
                   如果 n < 2，返回空列表。
    """
    if n < 2:
        return []

    # 使用 bytearray 替代 List[bool]，优化内存布局和访问效率
    # 1 代表潜在素数，0 代表合数
    is_prime = bytearray(b'\x01') * (n + 1)
    is_prime[0] = is_prime[1] = 0

    primes = []

    for i in range(2, n + 1):
        if is_prime[i]:
            primes.append(i)

        # 遍历已找到的素数列表，标记合数
        for p in primes:
            composite = p * i
            
            # 如果合数超出范围 n，停止内层循环
            if composite > n:
                break
            
            is_prime[composite] = 0

            # 核心逻辑：保证每个合数只被最小质因子标记
            # 如果 i 是 p 的倍数，则 i 的最小质因子小于等于 p。
            # 若不 break，后续的素数 q 标记 i * q 时，q 就不是最小质因子了。
            if i % p == 0:
                break

    return primes
```
📝 记忆已更新，新增一条 'execution' 记录。

--- 第 2/2 轮迭代 ---

-> 正在进行反思...
🧠 正在调用 coding-glm-5-free 模型...
✅ 大语言模型响应成功:

解析

Reflection.py 是一个展示 Reflection（反思） 模式的 Agent 示例。这种模式通过让 LLM 自我审查和优化，来提升代码生成的质量。

🏗️ 架构解析

这个项目由两个核心模块组成：

1. 记忆模块 (Memory 类)

• 作用：充当 Agent 的短期记忆，记录整个优化过程的轨迹。
• 存储内容：

• • execution: 实际生成的代码（初始版本或优化版本）。
  • reflection: “评审员”给出的反馈意见。

• 关键方法：

• • get_last_execution(): 获取最新版本的代码，供下一轮评审使用。

2. Reflection 智能体 (ReflectionAgent 类)

这是核心控制逻辑，它通过一个循环来不断打磨代码。

工作流程 (Workflow):

1. 初始尝试 (Initial Attempt):

• • 使用 INITIAL_PROMPT_TEMPLATE 让 LLM 生成第一版代码。
  • 将代码存入记忆。

1. 迭代优化循环 (Iterative Loop):

• • a. 反思 (Reflect):

• • • 将当前最新代码发给 LLM（扮演严格的代码评审专家）。
    • 使用 REFLECT_PROMPT_TEMPLATE 要求 LLM 分析时间复杂度并找出瓶颈。
    • 检查反馈中是否包含“无需改进”，如果是则提前结束。

• • b. 优化 (Refine):

• • • 将原始任务、上一版代码、评审反馈一起发给 LLM（扮演资深程序员）。
    • 使用 REFINE_PROMPT_TEMPLATE 要求 LLM 根据反馈生成新版代码。
    • 将新代码存入记忆。

🧩 关键提示词 (Prompts)

代码中定义了三个精心设计的 Prompt 模板，分别对应三个角色/阶段：

• Coder (初级): 负责写出能跑的代码。
• Reviewer (专家): 负责挑刺，专注于算法效率（如建议用筛法代替试除法）。
• Coder (高级): 负责根据反馈修改代码。

🚀 运行逻辑

在 if __name__ == '__main__': 块中：

1. 初始化 HelloAgentsLLM 客户端。
2. 创建 ReflectionAgent，设定最大迭代次数为 2。
3. 执行任务：“找出1到n之间所有的素数”。

预期效果：

• 第一轮可能生成一个简单的双重循环（O(n^2)）代码。
• 评审员会指出效率低，建议用埃氏筛。
• 第二轮生成的代码就会变成优化后的筛法（O(n log log n)）