inference.py

"""
inference.py
====================================

Entry point for Meta PyTorch OpenEnv Hackathon.

REQUIRED LOG FORMAT:
[START]
[STEP]
[END]
"""

import os
import json
import traceback

from openai import OpenAI

from pydebug_optimizer.env import PyDebugOptimizerEnv
from pydebug_optimizer.models import Action
from pydebug_optimizer.utils import parse_agent_response


# ============================================================
# 🔧 ENV CONFIG
# ============================================================

API_BASE_URL = os.getenv("API_BASE_URL")
MODEL_NAME = os.getenv("MODEL_NAME")
HF_TOKEN = os.getenv("HF_TOKEN")


# ============================================================
# 🤖 LLM CALL
# ============================================================

def query_llm(prompt: str) -> dict:
    """
    Sends prompt to LLM and returns parsed JSON response.
    """

    try:
        client = OpenAI(
            base_url=API_BASE_URL,
            api_key=HF_TOKEN,
        )

        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a senior Python engineer. "
                        "Return ONLY valid JSON with fields:\n"
                        "error_type, error_justification, fixed_code, "
                        "fix_justification, optimized_code, complexity_justification"
                    ),
                },
                {"role": "user", "content": prompt},
            ],
            temperature=0.2,
        )

        raw_output = response.choices[0].message.content

        parsed = parse_agent_response(raw_output)

        return parsed

    except Exception as e:
        print("LLM ERROR:", str(e))
        return {}


# ============================================================
# 🚀 MAIN LOOP
# ============================================================

def main():
    print("[START]")

    try:
        # ----------------------------------------------------
        # INIT ENV
        # ----------------------------------------------------
        env = PyDebugOptimizerEnv()
        obs, _ = env.reset()

        prompt = f"""
You are given buggy Python code.

TASK:
{obs['task_description']}

CODE:
{obs['code_snippet']}

Respond ONLY in JSON with:
- error_type
- error_justification
- fixed_code
- fix_justification
- optimized_code
- complexity_justification
"""

        # ----------------------------------------------------
        # QUERY LLM
        # ----------------------------------------------------
        action_dict = query_llm(prompt)

        if not action_dict:
            raise ValueError("Empty or invalid LLM response")

        # ----------------------------------------------------
        # VALIDATE ACTION
        # ----------------------------------------------------
        action = Action(**action_dict)

        # ----------------------------------------------------
        # STEP ENV
        # ----------------------------------------------------
        obs, reward, terminated, truncated, info = env.step(action.model_dump())

        print("[STEP]")
        print("Reward:", reward)
        print("Info:", info)

    except Exception as e:
        print("[STEP]")
        print("ERROR:", str(e))
        print(traceback.format_exc())

    finally:
        print("[END]")


# ============================================================
# ENTRY
# ============================================================

if __name__ == "__main__":
    main()