inference.py

"""
Inference Script — Code Debug Environment
==========================================
Runs the baseline LLM agent against all three tasks (fix_syntax, fix_logic,
fix_algorithm) and emits the mandatory [START]/[STEP]/[END] log lines.

Environment variables:
    API_BASE_URL   API endpoint for the LLM  (default: HF router)
    MODEL_NAME     Model identifier           (default: Qwen2.5-72B-Instruct)
    HF_TOKEN       API key
    IMAGE_NAME     Docker image name (if using from_docker_image)
"""

import asyncio
import os
import textwrap
from typing import List, Optional

from openai import OpenAI

from code_debug_env import CodeDebugEnv, DebugAction

IMAGE_NAME = os.getenv("IMAGE_NAME")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")

BENCHMARK = "code_debug_env"
TASKS = ["fix_syntax", "fix_logic", "fix_algorithm"]
MAX_STEPS = 10
TEMPERATURE = 0.2
MAX_TOKENS = 1024


# ── logging helpers (match mandatory format exactly) ─────────────────────

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(
    step: int, action: str, reward: float, done: bool, error: Optional[str]
) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    safe_action = action.replace("\n", "\\n")
    print(
        f"[STEP] step={step} action={safe_action} "
        f"reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


def log_end(
    success: bool, steps: int, score: float, rewards: List[float]
) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} "
        f"score={score:.2f} rewards={rewards_str}",
        flush=True,
    )


# ── LLM interaction ─────────────────────────────────────────────────────

SYSTEM_PROMPT = textwrap.dedent("""\
    You are an expert Python debugger. You are given a buggy Python function,
    a description of what it should do, and test results showing which tests
    pass or fail.

    Your job is to output ONLY the corrected Python code — nothing else.
    Do not include explanations, markdown fences, or commentary.
    Output the complete function(s) so the code can be executed as-is.
""")


def build_user_prompt(obs_dict: dict) -> str:
    errors_section = ""
    if obs_dict.get("stderr"):
        errors_section = "## Errors\n" + obs_dict["stderr"]

    return textwrap.dedent(f"""\
        ## Task
        {obs_dict['task_description']}

        ## Buggy Code (original)
        ```
        {obs_dict['buggy_code']}
        ```

        ## Current Code (your last submission)
        ```
        {obs_dict['current_code']}
        ```

        ## Test Results ({obs_dict['tests_passed']}/{obs_dict['tests_total']} passed)
        {obs_dict['test_results']}

        {errors_section}

        Reply with ONLY the corrected Python code.
    """)


def call_llm(client: OpenAI, obs_dict: dict) -> str:
    user_prompt = build_user_prompt(obs_dict)
    try:
        resp = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            stream=False,
        )
        text = (resp.choices[0].message.content or "").strip()
        # Strip markdown fences if the model wraps the code
        if text.startswith("```"):
            lines = text.split("\n")
            lines = [l for l in lines if not l.startswith("```")]
            text = "\n".join(lines)
        return text
    except Exception as exc:
        print(f"[DEBUG] LLM call failed: {exc}", flush=True)
        return ""


# ── main loop ────────────────────────────────────────────────────────────

async def run_task(client: OpenAI, env: CodeDebugEnv, task: str) -> None:
    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False

    log_start(task=task, env=BENCHMARK, model=MODEL_NAME)

    try:
        result = await env.reset()
        obs = result.observation

        obs_dict = {
            "task_description": obs.task_description,
            "buggy_code": obs.buggy_code,
            "current_code": obs.current_code,
            "test_results": obs.test_results,
            "tests_passed": obs.tests_passed,
            "tests_total": obs.tests_total,
            "stderr": obs.stderr,
        }

        for step in range(1, MAX_STEPS + 1):
            fixed_code = call_llm(client, obs_dict)
            if not fixed_code:
                fixed_code = obs_dict["current_code"]

            result = await env.step(DebugAction(code=fixed_code))
            obs = result.observation
            reward = result.reward or 0.0
            done = result.done

            rewards.append(reward)
            steps_taken = step

            log_step(
                step=step,
                action=fixed_code[:120],
                reward=reward,
                done=done,
                error=obs.stderr if obs.stderr else None,
            )

            obs_dict = {
                "task_description": obs.task_description,
                "buggy_code": obs.buggy_code,
                "current_code": obs.current_code,
                "test_results": obs.test_results,
                "tests_passed": obs.tests_passed,
                "tests_total": obs.tests_total,
                "stderr": obs.stderr,
            }

            if done:
                break

        if obs_dict["tests_total"] > 0:
            score = obs_dict["tests_passed"] / obs_dict["tests_total"]
        score = min(max(score, 0.0), 1.0)
        success = score >= 0.5

    finally:
        try:
            await env.close()
        except Exception as e:
            print(f"[DEBUG] env.close() error: {e}", flush=True)
        log_end(
            success=success, steps=steps_taken, score=score, rewards=rewards
        )


async def main() -> None:
    openai_client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    for task in TASKS:
        # Each reset() cycles to the next task automatically on the server.
        # When using Docker, pass CODE_DEBUG_TASK if you want a specific task.
        if IMAGE_NAME:
            env = await CodeDebugEnv.from_docker_image(IMAGE_NAME)
        else:
            base = os.getenv("CODE_DEBUG_BASE_URL", "http://localhost:8000")
            env = CodeDebugEnv(base_url=base)

        await run_task(openai_client, env, task)


if __name__ == "__main__":
    asyncio.run(main())