pip install langchain_openai openai python-dotenv
## .env 文件中需要有 OPENAI_API_KEY

## MIT License
## Copyright (c) 2025 Mahtab Syed
## https://www.linkedin.com/in/mahtabsyed/

"""
实践代码示例 - 迭代 2
为了说明目标设定与监控模式，我们有一个使用 LangChain 和 OpenAI API 的示例：

目标：构建一个 AI Agent，可以根据指定的目标为指定的用例编写代码：
- 接受编码问题（用例）作为代码输入或可以作为输入。
- 接受目标列表（例如，"简单"、"经过测试"、"处理边缘情况"）作为代码输入或可以作为输入。
- 使用 LLM（如 GPT-4o）生成和完善 Python 代码，直到满足目标。（我使用最多 5 次迭代，这也可以基于设定的目标）
- 要检查我们是否达到了目标，我要求 LLM 判断这一点并仅回答 True 或 False，这使得更容易停止迭代。
- 将最终代码保存在 .py 文件中，使用清晰的文件名和头部注释。
"""

import os
import random
import re
from pathlib import Path
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv, find_dotenv

## 🔐 加载环境变量
_ = load_dotenv(find_dotenv())
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
   raise EnvironmentError("❌ 请设置 OPENAI_API_KEY 环境变量。")

## ✅ 初始化 OpenAI 模型
print("📡 初始化 OpenAI LLM (gpt-4o)...")
llm = ChatOpenAI(
   model="gpt-4o",  # 如果你无法访问 gpt-4o，请使用其他 OpenAI LLM
   temperature=0.3,
   openai_api_key=OPENAI_API_KEY,
)

## --- 实用函数 ---

def generate_prompt(
   use_case: str, goals: list[str], previous_code: str = "", feedback: str = ""
) -> str:
   print("📝 构建代码生成提示词...")
   base_prompt = f"""
你是一个 AI 编码 Agent。你的工作是根据以下用例编写 Python 代码：

用例：{use_case}

你的目标是：
{chr(10).join(f"- {g.strip()}" for g in goals)}
"""
   if previous_code:
       print("🔄 将之前的代码添加到提示词中以进行完善。")
       base_prompt += f"\n之前生成的代码：\n{previous_code}"
   
   if feedback:
       print("📋 包含反馈以进行修订。")
       base_prompt += f"\n对之前版本的反馈：\n{feedback}\n"
   
   base_prompt += "\n请仅返回修订后的 Python 代码。不要在代码之外包含注释或解释。"
   return base_prompt

def get_code_feedback(code: str, goals: list[str]) -> str:
   print("🔍 根据目标评估代码...")
   feedback_prompt = f"""
你是一个 Python 代码审查员。下面显示了一个代码片段。

基于以下目标：
{chr(10).join(f"- {g.strip()}" for g in goals)}

请对此代码进行批评并确定是否满足目标。提及是否需要改进清晰度、简单性、正确性、边缘情况处理或测试覆盖率。

代码：
{code}
"""
   return llm.invoke(feedback_prompt)

def goals_met(feedback_text: str, goals: list[str]) -> bool:
   """
   使用 LLM 根据反馈文本评估目标是否已达成。
   返回 True 或 False（从 LLM 输出中解析）。
   """
   review_prompt = f"""
你是一个 AI 审查员。这些是目标：
{chr(10).join(f"- {g.strip()}" for g in goals)}

这是关于代码的反馈：
\"\"\"
{feedback_text}
\"\"\"

根据上述反馈，目标是否已达成？仅用一个词回答：True 或 False。
"""
   response = llm.invoke(review_prompt).content.strip().lower()
   return response == "true"

def clean_code_block(code: str) -> str:
   lines = code.strip().splitlines()
   if lines and lines[0].strip().startswith("```"):
       lines = lines[1:]
   if lines and lines[-1].strip() == "```":
       lines = lines[:-1]
   return "\n".join(lines).strip()

def add_comment_header(code: str, use_case: str) -> str:
   comment = f"# 此 Python 程序实现以下用例：\n# {use_case.strip()}\n"
   return comment + "\n" + code

def to_snake_case(text: str) -> str:
   text = re.sub(r"[^a-zA-Z0-9 ]", "", text)
   return re.sub(r"\s+", "_", text.strip().lower())

def save_code_to_file(code: str, use_case: str) -> str:
   print("💾 保存最终代码到文件...")
   summary_prompt = (
       f"将以下用例总结为一个小写单词或短语，"
       f"不超过 10 个字符，适合作为 Python 文件名：\n\n{use_case}"
   )
   raw_summary = llm.invoke(summary_prompt).content.strip()
   short_name = re.sub(r"[^a-zA-Z0-9_]", "", raw_summary.replace(" ", "_").lower())[:10]
   random_suffix = str(random.randint(1000, 9999))
   filename = f"{short_name}_{random_suffix}.py"
   filepath = Path.cwd() / filename
   
   with open(filepath, "w") as f:
       f.write(code)
   
   print(f"✅ 代码保存到：{filepath}")
   return str(filepath)

## --- 主智能体 ---

def run_code_agent(use_case: str, goals_input: str, max_iterations: int = 5) -> str:
   goals = [g.strip() for g in goals_input.split(",")]
   print(f"\n🎯 用例：{use_case}")
   print("🎯 目标：")
   for g in goals:
       print(f"  - {g}")
   
   previous_code = ""
   feedback = ""
   
   for i in range(max_iterations):
       print(f"\n=== 🔁 迭代 {i + 1} / {max_iterations} ===")
       
       prompt = generate_prompt(use_case, goals, previous_code, 
                               feedback if isinstance(feedback, str) else feedback.content)
       print("🚧 生成代码...")
       code_response = llm.invoke(prompt)
       raw_code = code_response.content.strip()
       code = clean_code_block(raw_code)
       
       print("\n🧾 生成的代码：\n" + "-" * 50 + f"\n{code}\n" + "-" * 50)
       
       print("\n📤 提交代码进行反馈审查...")
       feedback = get_code_feedback(code, goals)
       feedback_text = feedback.content.strip()
       
       print("\n📥 收到反馈：\n" + "-" * 50 + f"\n{feedback_text}\n" + "-" * 50)
       
       if goals_met(feedback_text, goals):
           print("✅ LLM 确认目标已达成。停止迭代。")
           break
       
       print("🛠️ 目标尚未完全达成。准备下一次迭代...")
       previous_code = code
   
   final_code = add_comment_header(code, use_case)
   return save_code_to_file(final_code, use_case)

## --- CLI 测试运行 ---

if __name__ == "__main__":
   print("\n🧠 欢迎使用 AI 代码生成 Agent")
   
   # 示例 1
   use_case_input = "编写代码查找给定正整数的 BinaryGap"
   goals_input = "代码简单易懂，功能正确，处理全面的边缘情况，仅接受正整数输入，打印结果并附带几个示例"
   run_code_agent(use_case_input, goals_input)
   
   # 示例 2
   # use_case_input = "编写代码计算当前目录及其所有嵌套子目录中的文件数量，并打印总数"
   # goals_input = (
   #     "代码简单易懂，功能正确，处理全面的边缘情况，忽略性能建议，忽略关于使用 unittest 或 pytest 等测试套件的建议"
   # )
   # run_code_agent(use_case_input, goals_input)
   
   # 示例 3
   # use_case_input = "编写代码，接受 word doc 或 docx 文件的命令行输入，打开它并计算其中的单词数和字符数并全部打印"
   # goals_input = "代码简单易懂，功能正确，处理边缘情况"
   # run_code_agent(use_case_input, goals_input)

充当一位对产生清晰、正确和简单代码有着深刻承诺的专家代码审查员。你的核心使命是通过确保每个建议都基于现实和最佳实践来消除代码"幻觉"。当我向你提供代码片段时，我希望你：

-- 识别和纠正错误：指出任何逻辑缺陷、错误或潜在的运行时错误。
-- 简化和重构：建议使代码更易读、高效和可维护的更改，而不牺牲正确性。
-- 提供清晰的解释：对于每个建议的更改，解释为什么它是改进，引用清晰代码、性能或安全性的原则。
-- 提供更正后的代码：显示建议更改的"之前"和"之后"，以便改进清晰可见。

你的反馈应该是直接的、建设性的，并始终旨在提高代码质量。