mvp.py

"""
MVP: 单题端到端闭环
流程: 加载HumanEval第1题 → DeepSeek生成代码 → subprocess执行测试 → 打印pass/fail
"""

import subprocess
import sys
import os
import textwrap
from datasets import load_dataset
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
)

def load_single_problem(index=0):
    print(f"[加载] 第{index}题...")
    ds = load_dataset("openai/openai_humaneval", trust_remote_code=True)
    problem = ds["test"][index]
    print(f"[加载] 题目ID: {problem['task_id']}")
    print(f"[加载] 函数名: {problem['entry_point']}")
    return problem

def generate_code(prompt, entry_point):
    print(f"\n[生成] 调用DeepSeek...")
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {
                "role": "system",
                "content": (
                    "你是Python代码专家。"
                    "用户给你函数签名和docstring，你直接输出完整的Python函数，"
                    "包含def行和函数体，不要任何解释，不要markdown代码块符号。"
                )
            },
            {
                "role": "user",
                "content": f"请实现以下Python函数：\n\n{prompt}"
            }
        ],
        temperature=0.0,
        max_tokens=512
    )
    code = response.choices[0].message.content.strip()
    # 清理markdown标记
    if "```" in code:
        lines = code.split('\n')
        lines = [l for l in lines if not l.strip().startswith('```')]
        code = '\n'.join(lines)
    print(f"[生成] 结果:\n{code.strip()}")
    return code.strip()

def execute_test(generated_code, test_code, entry_point, timeout=5):
    """
    直接用模型生成的完整函数 + 测试代码拼接执行
    不依赖prompt，避免缩进冲突
    """
    full_code = f"{generated_code}\n\n{test_code}\n\ncheck({entry_point})"
    print(f"\n[执行] 完整代码:\n{'-'*40}\n{full_code}\n{'-'*40}")

    try:
        result = subprocess.run(
            [sys.executable, "-c", full_code],
            capture_output=True, text=True, timeout=timeout
        )
        return {
            "passed": result.returncode == 0,
            "error": result.stderr.strip() or None
        }
    except subprocess.TimeoutExpired:
        return {"passed": False, "error": f"超时>{timeout}s"}
    except Exception as e:
        return {"passed": False, "error": str(e)}

def run_mvp():
    print("="*50)
    print("SAGE-Code MVP")
    print("="*50)

    problem = load_single_problem(index=0)
    generated_code = generate_code(problem["prompt"], problem["entry_point"])
    result = execute_test(
        generated_code=generated_code,
        test_code=problem["test"],
        entry_point=problem["entry_point"]
    )

    print("\n" + "="*50)
    if result["passed"]:
        print("✅ PASS")
    else:
        print("❌ FAIL")
        if result["error"]:
            print(f"错误:\n{result['error']}")
    print("="*50)
    return result

if __name__ == "__main__":
    run_mvp()