backend/solver.py

# solver.py
import sys
import re
import json
import os
import subprocess
from deep_convert import deep_convert
from openai import OpenAI
from dotenv import load_dotenv

def solve_puzzle(index, puzzle, expected_solution, sys_content):
    """
    使用 sys_content + puzzle 调用模型，生成 Z3 Python 代码，
    运行并比对输出是否与 expected_solution 相同。
    """
    load_dotenv()
    client = OpenAI(
        api_key=os.getenv("GEMINI_API_KEY"),
        base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
    )

    messages = [
        {"role": "user", "content": sys_content},  # 先把 sys_content 放进去
        {"role": "user", "content": puzzle},       # 再放 puzzle
    ]
    attempts = 0
    current_solution = None

    while attempts < 3:
        attempts += 1
        response = client.chat.completions.create(
            model="gemini-2.0-flash",
            messages=messages,
            temperature=0.0,
            stream=False
        )
        content = response.choices[0].message.content
        messages.append(response.choices[0].message)

        # 提取 code
        code_blocks = re.findall(r"```(?:python)?(.*?)```", content, re.DOTALL)
        if not code_blocks:
            messages.append({"role": "user", "content": "Please write a complete Python code in your response. Try again."})
            continue

        code_to_run = code_blocks[0].strip()
        result = subprocess.run(
            [sys.executable, "-c", code_to_run],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        output = result.stdout.strip()
        # print(output)

        try:
            current_solution = json.loads(output)
        except json.JSONDecodeError:
            messages.append({"role": "user", "content": "Your output is not valid JSON. Please ensure your code prints the solution as a JSON dictionary."})
            continue

        if 'rows' in current_solution.keys() and current_solution['rows'] == deep_convert(expected_solution)['rows']:
            return {
                "index": index,
                "success": True,
                "solution": current_solution,
                "attempts": attempts,
                "generatedCode": code_to_run,
                "modelResponse": content
            }
        else:
            messages.append({"role": "user", "content": "The solution does not match the expected answer. Please check your categories and constraints and provide the complete code again."})

    return {
        "index": index,
        "success": False,
        "solution": current_solution,
        "attempts": attempts,
        "generatedCode": code_to_run,
        "modelResponse": content
    }