inference.py

import os
import random
from typing import List, Optional
from openai import OpenAI

from server.github_learner_env_environment import GitHubLearnerEnv
from models import Action

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

if "API_KEY" in os.environ and "API_BASE_URL" in os.environ:
    API_KEY = os.environ["API_KEY"]
    API_BASE_URL = os.environ["API_BASE_URL"]
else:
    API_KEY = None
    API_BASE_URL = None

MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")

TASK_NAME = "github_skill_path"
BENCHMARK = "github_learner_env"

MAX_STEPS = 5
random.seed(42)

# =========================
# LOGGING
# =========================

def log_start(task: str, env: str, model: str):
    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]):
    error_val = error if error else "null"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}",
        flush=True,
    )


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

# =========================
# ACTION SELECTION
# =========================
def choose_action(obs):
    state = obs.current_state
    repos = obs.available_repos

    mastered = set(getattr(state, "mastered_skills", []))

    valid = [
        r for r in repos
        if r.skill not in mastered
        and all(p in mastered for p in r.prereqs)
    ]

    if not valid:
        valid = repos

    def score(r):
        return r.difficulty * 2 + len(r.prereqs)

    best = max(valid, key=score)
    return best.id

# =========================
# MAIN EXECUTION
# =========================

def main():
    rewards = []
    steps_taken = 0
    success = False
    score = 0.0
    env = None  # ✅ prevent crash in finally

    # ✅ START LOG
    log_start(TASK_NAME, BENCHMARK, MODEL_NAME)

    # 🔥 LLM PING (ONLY FOR VALIDATOR)
    if API_KEY and API_BASE_URL:
        try:
            client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

            client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role": "user", "content": "1"}],
                max_tokens=1,
                temperature=0
            )

            print("[DEBUG] LLM ping success", flush=True)

        except Exception as e:
            print(f"[DEBUG] LLM ping failed: {e}", flush=True)

    try:
        # -------- ENV INIT --------
        try:
            env = GitHubLearnerEnv()
        except Exception as e:
            print(f"[FATAL] env init failed: {e}", flush=True)
            log_end(False, 0, 0.0, [])
            return

        # -------- RESET --------
        try:
            obs = env.reset()
            done = obs.done
        except Exception as e:
            print(f"[ERROR] reset failed: {e}", flush=True)
            log_end(False, 0, 0.0, [])
            return

        # -------- LOOP --------
        for step in range(1, MAX_STEPS + 1):
            if done:
                break

            action_id = choose_action(obs)

            if action_id is None:
                break

            action = Action(repo_id=action_id)

            try:
                obs = env.step(action)
                reward = obs.reward
                done = obs.done
                error = None
            except Exception as e:
                reward = 0.0
                done = True
                error = str(e)

            rewards.append(reward)
            steps_taken = step

            repo = next((r for r in env.repos if r.id == action_id), None)
            skill = repo.skill if repo else "unknown"

            log_step(step, f"{action_id}:{skill}", reward, done, error)

        # -------- SCORE --------
        total_reward = sum(rewards)
        max_possible = MAX_STEPS * 4.0

        score = total_reward / max_possible if max_possible > 0 else 0.0
        score = max(0.0, min(score, 1.0))

        success = score > 0.3

    except Exception as e:
        print(f"[FATAL] runtime crash: {e}", flush=True)

    finally:
        # ✅ SAFE CLOSE
        if env is not None:
            try:
                env.close()
            except:
                pass

        # ✅ ALWAYS END
        log_end(success, steps_taken, score, rewards)


if __name__ == "__main__":
    main()