코드 생성 성능 비교: o4-mini, Claude 3.7, Gemini 2.5 Pro 평가


import subprocess
import sys
from litellm import completion


def clean_llm_code_block(text):
    return text.replace("```python", "").replace("```", "").strip()


def ask_llm_for_function_call(code: str, raw_input: str) -> str:
    prompt = (
        "You're given a Python function and a single input string. "
        "Format it into a valid Python function call using only standard types.\n\n"
        f"Function:\n{code}\n\n"
        f"Input:\n{raw_input.strip()}\n\n"
        "Return ONLY a valid function call (e.g., solve(3, 5))."
    )
    response = completion(
        model="openai/gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
        max_tokens=100,
    )
    raw = response["choices"][0]["message"]["content"]
    return clean_llm_code_block(raw)

def compare_output_with_llm(expected: str, actual: str) -> bool:
    prompt = (
        f"Expected output: {expected.strip()}\n"
        f"Actual output: {actual.strip()}\n\n"
        "Are these outputs semantically equivalent? Reply YES or NO."
    )
    response = completion(
        model="openai/gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
        max_tokens=10,
    )
    return response["choices"][0]["message"]["content"].strip().upper() == "YES"

def run_code_and_call_function(code: str, function_call: str, timeout=10):
    full_code = code + f"\n\nprint({function_call})"
    try:
        result = subprocess.run(
            [sys.executable, "-c", full_code],
            capture_output=True,
            text=True,
            timeout=timeout
        )
        return result.stdout.strip(), result.stderr.strip()
    except subprocess.TimeoutExpired:
        return "", "Execution timed out."
    except Exception as e:
        return "", str(e)

# Example LLM-generated code (you'd swap this in dynamically)
code_from_llm = """
def solve(a, b, mod):
    return pow(a, -1, mod)
"""

# Example test case block
test_case = {
    "input": [
        "3 2 998244353\n",
        "9 3 998244353\n",
        "3 1 998244353\n",
        "9 4 998244353\n"
    ],
    "output": [
        "665496236\n",
        "449209967\n",
        "499122178\n",
        "665496237\n"
    ]
}

# Run test cases
for i, input_line in enumerate(test_case["input"]):
    expected = test_case["output"][i]
    try:
        call = ask_llm_for_function_call(code_from_llm, input_line)
        result, error = run_code_and_call_function(code_from_llm, call)

        if error:
            print(f"[{i}] ERROR during execution: {error}")
            continue

        is_correct = compare_output_with_llm(expected, result)
        print(f"[{i}] input: {input_line.strip()} → output: {result} | expected: {expected.strip()} | PASS: {is_correct}")

    except Exception as e:
        print(f"[{i}] Failed: {repr(e)}")


import os
import sys
import time
import subprocess
from datasets import load_dataset
from litellm import completion as oai_completion
from anthropic import Anthropic
from google import genai
import weave

from weave.flow.eval_imperative import EvaluationLogger
from google import genai
from google.genai import types
from litellm import completion as oai_completion
import re
from litellm import completion as oai_completion


weave.init("codecontests_eval")

# API keys
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "sk-...")
CLAUDE_KEY = os.getenv("CLAUDE_API_KEY", "sk-...)
GENAI_KEY = os.getenv("GOOGLE_API_KEY", "sk-...")

# Clients
anthropic_client = Anthropic(api_key=CLAUDE_KEY)
gemini_client = genai.Client(api_key=GENAI_KEY)



def clean_llm_code_block(text):
    import re

    cleaned_text = text.replace("```python", "").replace("```", "").strip()
    code_blocks = re.findall(r"(def solve\(.*?)(?=^def |\Z)", cleaned_text, re.DOTALL | re.MULTILINE)
    source_text = code_blocks[-1] if code_blocks else cleaned_text

    prompt = (
        "Given the following response from a language model, extract ONLY the valid Python code for the function. "
        "Do not include any explanations, text, or formatting fences. Only the code.\n\n"
        f"Response:\n{source_text}\n\n"
        "Return ONLY the Python code, including any necessary imports:"
    )

    response = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )

    gpt4o_code = response["choices"][0]["message"]["content"]
    gpt4o_code = gpt4o_code.replace("```python", "").replace("```", "").strip()
    return gpt4o_code

@weave.op()
def generate_completion(model: str, prompt: str) -> str:
    if model.startswith("openai/"):
        response = oai_completion(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            reasoning_effort="low",
        )
        return response["choices"][0]["message"]["content"].strip()

    elif model.startswith("anthropic/"):
        response = anthropic_client.messages.create(
            model=model.replace("anthropic/", ""),
            max_tokens=8000,
            thinking={"type": "enabled", "budget_tokens": 4000},
            messages=[{"role": "user", "content": prompt}],
        )
        for block in response.content:
            if block.type == "text":
                return block.text.strip()
        return "[No Claude response]"

    elif model.startswith("gemini/"):
        result = gemini_client.models.generate_content(
            model=model.replace("gemini/", ""),
            config=types.GenerateContentConfig(
                thinking_config=types.ThinkingConfig(thinking_budget=4000)
            ),
            contents=[prompt]
        )
        return result.text.strip() if result.text else "[No Gemini response]"

    else:
        raise ValueError(f"Unsupported model: {model}")
    


def ask_llm_for_function_implementation(description: str, model: str) -> str:
    prompt = (
        f"Write a Python3 function named `solve` with typed input arguments for this problem -- eg the solve function should take arguments to handle different test cases:\n\n"
        f"{description.strip()}\n\n"
        "Return only a valid Python function -- no special packages that arent commonly used and NO MAIN function, no  if __name__ == __main__....., JUST write the function --  that returns the result. No comments, no explanations."
        f"HOWEVER, you still need to include necessary imports for libraries"
        f"IF you do not include the right imports, the code will not be executable, and your response will be judged as incorrect!"
    )
    return clean_llm_code_block(generate_completion(model, prompt))



@weave.op
def ask_llm_for_function_call(code: str, raw_input: str, model: str) -> str:

    prompt = (
        "You're given a Python function and a single input string. "
        "Format it into a valid Python function call using only standard types.\n\n"
        f"Function:\n{code}\n\n"
        f"Input:\n{raw_input.strip()}\n\n"
        "Return ONLY a valid function call (e.g., solve(3, 5)) WITH NO 'def' "
    )

    # Always use GPT-4o for this inference, regardless of the `model` argument.
    response = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    # The LLM may return markdown code blocks; strip them just in case.
    content = response["choices"][0]["message"]["content"]
    content = content.replace("```python", "").replace("```", "").strip()
    return content



def compare_output_with_llm(expected: str, actual: str, model: str) -> bool:
    prompt = (
        f"Expected output: {expected.strip()}\n"
        f"Actual output: {actual.strip()}\n\n"
        "Are these outputs equivalent? Eg ignore minor formatting errors etc, we are just looking for overall correctness in the output Reply YES or NO."
    )
    # response = generate_completion(model, prompt)

    response = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    # The LLM may return markdown code blocks; strip them just in case.
    res = 'YES' in str(response["choices"][0]["message"]["content"]).upper()
    return res 



def run_code_and_call_function(code: str, function_call: str, timeout=10):
    full_code = code + f"\n\nprint({function_call})"
    try:
        start = time.time()
        result = subprocess.run(
            [sys.executable, "-c", full_code],
            capture_output=True,
            text=True,
            timeout=timeout
        )
        latency = time.time() - start
        return result.stdout.strip(), result.stderr.strip(), latency
    except subprocess.TimeoutExpired:
        return "", "Execution timed out.", timeout
    except Exception as e:
        return "", str(e), 0.0


def ask_model_for_pip_command(error_msg):
    prompt = (
        "Given this Python error:\n\n"
        + error_msg +
        "\n\nWrite the pip install command needed to fix it. Only return the command, e.g.:\n"
        "pip install requests"
    )
    return generate_completion("openai/gpt-4o-2024-08-06", prompt)


def run_pip_install(pip_command):
    print(f"Running: {pip_command}")
    try:
        result = subprocess.run(
            pip_command.split(),
            capture_output=True,
            text=True,
            timeout=180
        )
        print(result.stdout.strip())
        if result.stderr:
            print(result.stderr.strip())
    except Exception as e:
        print(f"pip install failed: {e}")



def evaluate_model_on_code_contests(model_name: str):
    print(f"\n\nRunning evaluation for model: {model_name}\n")
    ds = load_dataset("deepmind/code_contests", split="test", streaming=True)
    ds = list(ds.take(31))

    eval_logger = EvaluationLogger(
        model=model_name.replace("-", "_").replace("/", "_").replace(".", "_"),
        dataset="code_contests_test"
    )
    all_latencies = []

    for i in range(30):
        row = ds[i]
        description = row["description"]
        raw_inputs = row["public_tests"]["input"]
        expected_outputs = row["public_tests"]["output"]

        try:
            code = ask_llm_for_function_implementation(description, model=model_name)
            print(f"\n=== Task {row['name']} ===", flush=True)
            # print("Generated code:\n", code)

            all_passed = True
            task_latencies = []
            results_lst, expected_lst = [], []

            for j, raw_input in enumerate(raw_inputs):
                expected = expected_outputs[j] if j < len(expected_outputs) else ""

                try:
                    function_call = ask_llm_for_function_call(code, raw_input, model=model_name)
                    result, error, latency = run_code_and_call_function(code, function_call)
                    if latency < 99:
                        task_latencies.append(latency)

                
                    if error:
                        print(f"[{j}] Runtime error: {error}")
                        if "ModuleNotFoundError" in error:
                            pip_cmd = ask_model_for_pip_command(error)
                            run_pip_install(pip_cmd)
                            # Re-run once after pip install
                            result, error, latency = run_code_and_call_function(code, function_call)
                            
                            task_latencies.append(latency)
                            if error:
                                print(f"[{j}] Retry failed: {error}")
                                all_passed = False
                                continue
                        else:
                            all_passed = False
                            continue

                    is_correct = compare_output_with_llm(expected, result, model="openai/gpt-4o-2024-08-06")###### 
                    results_lst.append(result)
                    expected_lst.append(expected)
                    if not is_correct:
                        all_passed = False
                    print(f"[{j}] input: {raw_input.strip()} → output: {result} | expected: {expected.strip()} | PASS: {is_correct} | latency: {latency:.2f}s")

                except Exception as inner:
                    print(f"[{j}] Inner error: {repr(inner)}")
                    all_passed = False
            
            task_avg_latency = sum(task_latencies) / len(task_latencies) if len(task_latencies) > 0 else 0.0
            all_latencies.extend(task_latencies)

            prediction_log = eval_logger.log_prediction(
                inputs={"description": description},
                output={'code': code, 'execution_result': results_lst, 'expected_execution_result': expected_lst}
            )
            prediction_log.log_score("correctness", all_passed)
            prediction_log.log_score("code_latency", task_avg_latency)
            prediction_log.finish()

        except Exception as e:
            print(f"[{i}] Top-level failure: {repr(e)}")
            prediction_log = eval_logger.log_prediction(
                inputs={"description": description},
                output=str(e)
            )
            prediction_log.log_score("correctness", False)
            prediction_log.finish()

    avg_latency = sum(all_latencies) / len(all_latencies) if all_latencies else 0.0
    eval_logger.log_summary({"avg_code_latency": avg_latency})
    print(f"Evaluation complete for {model_name}. View in Weave UI.")


# Run for all models

evaluate_model_on_code_contests("gemini/gemini-2.5-pro-preview-05-06")
evaluate_model_on_code_contests("anthropic/claude-3-7-sonnet-20250219")
evaluate_model_on_code_contests("openai/o4-mini")