CodeContests에서 Claude 4, Codex, Gemini 2.5 Pro 비교 테스트

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.genai import types
import re

weave.init("codecontests_evalv2")

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_api_key")
OAIKEY = os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_api_key")
CLAUDE_KEY = os.getenv("CLAUDE_API_KEY", "your_api_key")

GENAI_KEY = os.getenv("GOOGLE_API_KEY", "your_api_key")
from openai import OpenAI
client = OpenAI(api_key=OAIKEY)


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

def clean_llm_code_block(text):
    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:
    # Codex models (openai/codex-*, openai/codex-mini-latest, etc)
    if model.startswith("openai/codex-"):
        # Use the .responses.create API as per the latest OpenAI SDK
        codex_model = model.replace("openai/", "")
        # You can optionally add custom instructions (here, we use a neutral persona)
        response = client.responses.create(
            model=codex_model,
            instructions="You are a helpful, accurate Python coding assistant.",
            input=prompt
        )
        return response.output_text.strip()
    
    # General OpenAI chat-completions
    elif model.startswith("openai/"):
        from litellm import completion as oai_completion
        response = oai_completion(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            reasoning_effort="low",
        )
        return response["choices"][0]["message"]["content"].strip()

    # Anthropic Claude (API unchanged)
    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]"

    # Gemini (API unchanged)
    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' "
    )
    response = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    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 = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    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)
            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 TARGET MODELS ----

evaluate_model_on_code_contests("gemini/gemini-2.5-pro-preview-05-06")
evaluate_model_on_code_contests("anthropic/claude-sonnet-4-20250514")
evaluate_model_on_code_contests("openai/codex-mini-latest")
evaluate_model_on_code_contests("anthropic/claude-opus-4-20250514")

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.genai import types
import re

weave.init("codecontests_evalv2")

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_api_key")
OAIKEY = os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_api_key")
CLAUDE_KEY = os.getenv("CLAUDE_API_KEY", "your_api_key")

GENAI_KEY = os.getenv("GOOGLE_API_KEY", "your_api_key")


from openai import OpenAI
client = OpenAI(api_key=OAIKEY)

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

def clean_llm_code_block(text):
    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, thinking_budget=None, streaming=True) -> str:
    response_text = ""
    if model.startswith("anthropic/"):
        # Calculate correct total max_tokens
        if thinking_budget and int(thinking_budget) > 0:
            max_tokens = 8000 + int(thinking_budget)
        else:
            max_tokens = 8000

        create_kwargs = dict(
            model=model.replace("anthropic/", ""),
            max_tokens=max_tokens,
            messages=[{"role": "user", "content": prompt}],
        )

        # Only provide thinking if budget is positive and not None
        if thinking_budget and int(thinking_budget) > 0:
            create_kwargs['thinking'] = {"type": "enabled", "budget_tokens": int(thinking_budget)}

        client = anthropic_client

        if streaming:
            with client.messages.stream(**create_kwargs) as stream:
                for event in stream:
                    if event.type == "content_block_start":
                        block_type = event.content_block.type
                        if block_type == "thinking":
                            print("\n[THINKING]: ", end="", flush=True)
                        elif block_type == "text":
                            print("\n[RESPONSE]: ", end="", flush=True)
                    elif event.type == "content_block_delta":
                        d = event.delta
                        if getattr(d, "type", None) == "thinking_delta":
                            print(d.thinking, end="", flush=True)
                        elif getattr(d, "type", None) == "text_delta":
                            print(d.text, end="", flush=True)
                            response_text += d.text
                    elif event.type == "content_block_stop":
                        print()
        else:
            response = client.messages.create(**create_kwargs)
            for block in response.content:
                if block.type == "thinking" and getattr(block, "thinking", "").strip():
                    print("\n[THINKING]:", block.thinking.strip(), flush=True)
                elif block.type == "text" and getattr(block, "text", "").strip():
                    print("\n[RESPONSE]:", block.text.strip(), flush=True)
                    response_text += block.text.strip()
        return str(response_text)
    else:
        raise ValueError(f"Unsupported model: {model}")


def ask_llm_for_function_implementation(description: str, model: str, thinking_budget=None) -> 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, thinking_budget=thinking_budget))

@weave.op
def ask_llm_for_function_call(code: str, raw_input: str, model: str, thinking_budget=None) -> 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' "
    )
    response = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    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 = oai_completion(
        model="openai/gpt-4o-2024-08-06",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0,
    )
    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"
    )
    # In this context, no "thinking" is likely needed; you could propagate a thinking_budget if you want
    return generate_completion("anthropic/claude-sonnet-4-20250514", prompt, thinking_budget=None)

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, thinking_budget=None):
    tb_str = "nothinking" if not thinking_budget or int(thinking_budget) == 0 else f"tb{thinking_budget}"
    print(f"\n\nRunning evaluation for model: {model_name} | thinking_budget={tb_str}\n")
    ds = load_dataset("deepmind/code_contests", split="test", streaming=True)
    ds = list(ds.take(31))

    eval_logger = EvaluationLogger(
        model=f"{model_name.replace('-', '_').replace('/', '_').replace('.', '_')}_{tb_str}",
        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, thinking_budget=thinking_budget)
            print(f"\n=== Task {row['name']} ===", flush=True)
            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, thinking_budget=thinking_budget)
                    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} (thinking_budget={tb_str}). View in Weave UI.")

# ---- RUN FOR CLAUDE/SONNET4 FOR MULTIPLE THINKING BUDGETS INCLUDING "NO THINKING" ----
thinking_budgets = [None, 1024, 4000, 8000, 12000 ]
for tb in thinking_budgets:
    evaluate_model_on_code_contests("anthropic/claude-sonnet-4-20250514", thinking_budget=tb)