AI代码生成的机遇与挑战_02

import os
import openai
import json

def setup_openai():
    """设置OpenAI API密钥"""
    # 从环境变量获取API密钥，或直接在这里设置（不推荐硬编码）
    openai.api_key = os.getenv("OPENAI_API_KEY") or "your_api_key_here"
    
    # 检查API密钥是否设置
    if not openai.api_key or openai.api_key == "your_api_key_here":
        print("警告：请设置有效的OpenAI API密钥！")

def generate_code(prompt, model="gpt-4", max_tokens=1000, temperature=0.7):
    """
    使用OpenAI API生成代码
    
    参数:
        prompt: 自然语言描述或代码生成请求
        model: 使用的模型名称
        max_tokens: 生成的最大token数
        temperature: 控制生成内容的随机性，值越低越确定性
    
    返回:
        生成的代码字符串
    """
    try:
        # 发送请求到OpenAI API
        response = openai.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "system",
                    "content": "你是一个经验丰富的软件工程师，擅长多种编程语言和框架。请根据用户需求生成高质量、可运行的代码。"
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            max_tokens=max_tokens,
            temperature=temperature,
            n=1,
            stop=None
        )
        
        # 提取生成的代码
        code = response.choices[0].message.content.strip()
        
        # 处理可能的代码块标记
        if code.startswith("```"):
            # 寻找代码块的结束标记
            end_index = code.find("```", 3)
            if end_index != -1:
                code = code[3:end_index].strip()
                # 如果第一行是语言标识，移除它
                if '\n' in code and not code.split('\n')[0].strip().startswith("#") and not code.split('\n')[0].strip().startswith("//"):
                    code = '\n'.join(code.split('\n')[1:])
        
        return code
        
    except Exception as e:
        print(f"生成代码时出错: {e}")
        return None

def save_code_to_file(code, filename):
    """将生成的代码保存到文件"""
    try:
        # 确保目录存在
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        
        # 写入文件
        with open(filename, "w", encoding="utf-8") as f:
            f.write(code)
            
        print(f"代码已保存至: {filename}")
        return True
        
    except Exception as e:
        print(f"保存代码时出错: {e}")
        return False

def evaluate_code(code, task_description):
    """评估生成的代码是否满足需求"""
    try:
        response = openai.chat.completions.create(
            model="gpt-4",
            messages=[
                {
                    "role": "system",
                    "content": "你是一个代码评审专家，请评估以下代码是否满足给定的任务描述。如果不满足，请指出问题并提供改进建议。"
                },
                {
                    "role": "user",
                    "content": f"任务描述: {task_description}\n\n生成的代码:\n{code}"
                }
            ],
            max_tokens=500,
            temperature=0.5
        )
        
        evaluation = response.choices[0].message.content.strip()
        return evaluation
        
    except Exception as e:
        print(f"评估代码时出错: {e}")
        return "无法评估代码。"

# 示例使用
if __name__ == "__main__":
    # 设置OpenAI API
    setup_openai()
    
    # 定义代码生成任务
    tasks = [
        {
            "name": "数据可视化",
            "prompt": "使用Python的matplotlib库创建一个函数，该函数接收一个包含学生成绩的列表，然后生成一个柱状图，显示不同分数段的学生人数分布。确保图表有标题、坐标轴标签和图例。",
            "filename": "output_code/grade_visualization.py"
        },
        {
            "name": "Web API",
            "prompt": "使用Python的FastAPI框架创建一个简单的待办事项API，包含创建、获取、更新和删除待办事项的功能。数据可以存储在内存中。",
            "filename": "output_code/todo_api.py"
        },
        {
            "name": "文本处理",
            "prompt": "编写一个Python函数，该函数接收一个文本字符串，然后返回一个字典，其中包含文本中每个单词的出现次数。忽略大小写，并排除常见的停用词如'a', 'an', 'the', 'is', 'are'等。",
            "filename": "output_code/text_counter.py"
        }
    ]
    
    for task in tasks:
        print(f"\n" + "="*50 + f"\n生成 {task['name']} 代码...")
        
        # 生成代码
        code = generate_code(task['prompt'])
        
        if code:
            print(f"\n生成的代码:\n{code}")
            
            # 保存代码到文件
            save_code_to_file(code, task['filename'])
            
            # 评估代码
            print(f"\n代码评估:\n{evaluate_code(code, task['prompt'])}")
        else:
            print("代码生成失败！")
    
    print(f"\n" + "="*50 + "\n所有代码生成任务完成！")