Blame for inference.py

#!/usr/bin/env python3
# inference.py
# 
# This is the main script evaluators run to test the environment.
# It spins up the server, uses an LLM to decide on actions, runs an episode,
# and prints logs in the required [START] [STEP] [END] format.
#
# Needed env vars:
#   API_BASE_URL - base URL for the LLM API (OpenAI-compatible)
#   MODEL_NAME   - which model to use
#   HF_TOKEN     - HuggingFace token (used if server is on HF Spaces)
#
# Usage:
#   python inference.py
#   python inference.py --task survive_8_hours
#   python inference.py --url http://localhost:8000

import os
import sys
import json
import time
import argparse
import threading
import subprocess
from typing import Optional

import httpx
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

# --- Config ---
API_BASE_URL = os.environ.get("API_BASE_URL", "https://openrouter.ai/api/v1")
MODEL_NAME = os.environ.get("MODEL_NAME", "google/gemma-2-9b-it:free")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
SERVER_URL = os.environ.get("SERVER_URL", "http://localhost:7860")

# how many steps before we force-stop (safety net)
MAX_STEPS = 100

# rate limit: space calls apart to avoid hitting limits
_last_llm_call = 0.0
MIN_CALL_INTERVAL = 3.0  # seconds between LLM calls


def wait_for_server(url: str, timeout: int = 30) -> bool:
    """Poll until the server is up. Returns True if it came up in time."""
    print(f"[START] Waiting for server at {url}...")
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            r = httpx.get(f"{url}/health", timeout=3)
            if r.status_code == 200:
                print(f"[START] Server is up!")
                return True
        except Exception:
            pass
        time.sleep(1)
    return False


def start_local_server() -> Optional[subprocess.Popen]:
    """Start the FastAPI server as a subprocess. Returns the process."""
    print("[START] Launching local server...")
    proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "server.app:app",
         "--host", "0.0.0.0", "--port", "7860", "--log-level", "warning"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    return proc


def _heuristic_action(obs: dict, task_id: str) -> int:
    """Smart fallback heuristic when LLM is unavailable."""
    energy = obs.get("energy", 50)
    focus = obs.get("focus", 50)
    consecutive = obs.get("consecutive_study_minutes", 0)

    if task_id == "smart_break_taker":
        # Only break when energy < 40
        if energy < 40:
            return 1 if energy > 20 else 2
        return 0  # study otherwise

    elif task_id == "survive_8_hours":
        # Conservative - preserve energy
        if energy < 25:
            return 2  # long break when very low
        elif energy < 45:
            return 1  # short break when low
        return 0  # study

    else:  # maximize_productivity
        # Study when focus is high, short break when focus drops
        if energy < 35:
            return 1  # short break to recover
        elif consecutive >= 40 or focus < 55:
            return 1  # break to restore focus
        return 0  # study to accumulate productive minutes


def get_llm_action(client: OpenAI, obs: dict, task_id: str, step_num: int) -> int:
    """
    Ask the LLM what action to take given the current observation.
    Retries up to 3 times on rate limit errors before falling back to heuristic.
    Returns 0, 1, or 2.
    """
    global _last_llm_call

    # Enforce minimum interval between calls to avoid rate limits
    now = time.time()
    wait = MIN_CALL_INTERVAL - (now - _last_llm_call)
    if wait > 0:
        time.sleep(wait)

    system_prompt = """You are controlling a virtual student in a study session simulation.
Your goal depends on the task:
- survive_8_hours: Keep energy above 0 for 8 hours (be conservative, rest often)
- maximize_productivity: Maximize focus>60 minutes in 4 hours (study when focused, short breaks when focus drops)
- smart_break_taker: Only take breaks when energy is below 40 (NEVER break when energy >= 40)

Available actions:
0 = study (5 minutes, drains energy and focus slightly)
1 = short break (10 minutes, recovers energy +15 and focus +20)
2 = long break (20 minutes, recovers energy +35 and focus +45)

Rules:
- For maximize_productivity: break when consecutive_study >= 40min OR focus < 55
- For smart_break_taker: ONLY break when energy < 40, never otherwise
- For survive_8_hours: break when energy < 40

Respond with ONLY a JSON object like: {"action": 0}
"""

    user_msg = f"""Step {step_num} | Task: {task_id}
Energy: {obs['energy']:.1f}/100 | Focus: {obs['focus']:.1f}/100
Consecutive study: {obs['consecutive_study_minutes']:.0f}min | Breaks: {obs['break_count']}
Productive minutes: {obs['productive_minutes']:.0f} | Session: {obs['session_minutes']:.0f}min

What action? Reply with only: {{"action": N}} where N is 0, 1, or 2."""

    # Try up to 3 times, with backoff on rate limit errors
    for attempt in range(3):
        try:
            _last_llm_call = time.time()
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_msg},
                ],
                max_tokens=30,
                temperature=0.1,
            )
            content = response.choices[0].message.content.strip()

            # parse JSON - try to extract action even from messy output
            try:
                parsed = json.loads(content)
            except json.JSONDecodeError:
                # try extracting number from response
                import re
                m = re.search(r'"action"\s*:\s*([012])', content)
                if m:
                    return int(m.group(1))
                # try just finding a digit
                m = re.search(r'\b([012])\b', content)
                if m:
                    return int(m.group(1))
                return _heuristic_action(obs, task_id)

            action = int(parsed.get("action", 0))
            return max(0, min(2, action))

        except Exception as ex:
            err_str = str(ex)
            is_rate_limit = "429" in err_str or "rate" in err_str.lower()

            if is_rate_limit and attempt < 2:
                wait_time = 8 * (attempt + 1)  # 8s, then 16s
                print(f"[STEP] Rate limited, waiting {wait_time}s (retry {attempt+2}/3)...")
                time.sleep(wait_time)
                continue
            else:
                print(f"[STEP] LLM failed (attempt {attempt+1}/3): using heuristic")
                return _heuristic_action(obs, task_id)

    return _heuristic_action(obs, task_id)


def run_episode(server_url: str, task_id: str, llm_client: OpenAI) -> dict:
    """
    Run a complete episode and return summary stats.
    This is where all the [STEP] logs happen.
    """
    http = httpx.Client(timeout=30.0)

    # reset the environment
    reset_resp = http.post(f"{server_url}/reset", json={"task_id": task_id})
    reset_resp.raise_for_status()
    reset_data = reset_resp.json()
    obs = reset_data["observation"]

    print(f"\n[START] Episode started | Task: {task_id}")
    print(f"[START] Initial state: energy={obs['energy']:.1f}, focus={obs['focus']:.1f}")
    print(f"[START] Using model: {MODEL_NAME} at {API_BASE_URL}")
    print("-" * 60)

    step_num = 0
    total_reward = 0.0
    done = False
    final_info = {}

    action_names = {0: "STUDY", 1: "SHORT_BREAK", 2: "LONG_BREAK"}

    while not done and step_num < MAX_STEPS:
        step_num += 1

        # get LLM decision
        action = get_llm_action(llm_client, obs, task_id, step_num)

        # execute action
        step_resp = http.post(f"{server_url}/step", json={"action": action})
        step_resp.raise_for_status()
        step_data = step_resp.json()

        new_obs = step_data["observation"]
        reward = step_data["reward"]
        done = step_data["done"]
        info = step_data.get("info", {})

        total_reward += reward
        final_info = info

        # required log format: [STEP]
        print(
            f"[STEP] {step_num:3d} | action={action_names[action]:<12} | "
            f"reward={reward:+.3f} | energy={new_obs['energy']:5.1f} | "
            f"focus={new_obs['focus']:5.1f} | productive={new_obs['productive_minutes']:5.0f}min | "
            f"done={done}"
        )

        obs = new_obs

        if done:
            break

    http.close()

    final_score = final_info.get("score", 0.5)
    print("-" * 60)
    print(f"[END] Episode finished after {step_num} steps")
    print(f"[END] Task: {task_id}")
    print(f"[END] Total reward: {total_reward:.4f}")
    print(f"[END] Final score: {final_score:.4f} (range: strictly 0 to 1 exclusive)")
    print(f"[END] Energy remaining: {obs['energy']:.1f}")
    print(f"[END] Productive minutes: {obs['productive_minutes']:.0f}")
    print(f"[END] All scores: {final_info.get('all_task_scores', {})}")

    return {
        "task_id": task_id,
        "steps": step_num,
        "total_reward": total_reward,
        "score": final_score,
        "info": final_info,
    }


def main():
    parser = argparse.ArgumentParser(description="Study Optimizer inference script")
    parser.add_argument("--task", default="maximize_productivity",
                        choices=["survive_8_hours", "maximize_productivity", "smart_break_taker"],
                        help="Task to run")
    parser.add_argument("--url", default=None,
                        help="Server URL (default: start local server)")
    parser.add_argument("--no-server", action="store_true",
                        help="Don't start local server, assume it's already running")
    args = parser.parse_args()

    server_url = args.url or SERVER_URL
    server_proc = None

    # start local server if not using remote
    if not args.no_server and args.url is None:
        server_proc = start_local_server()
        if not wait_for_server(server_url, timeout=30):
            print("[START] ERROR: Server failed to start in 30 seconds")
            if server_proc:
                server_proc.terminate()
            sys.exit(1)
    else:
        # just check it's up
        if not wait_for_server(server_url, timeout=10):
            print(f"[START] ERROR: Can't reach server at {server_url}")
            sys.exit(1)

    # set up OpenAI client
    llm_client = OpenAI(
    api_key=os.environ.get("OPENROUTER_API_KEY") or os.environ.get("OPENAI_API_KEY") or HF_TOKEN or "dummy-key",
    base_url=API_BASE_URL,
)

    try:
        results = run_episode(server_url, args.task, llm_client)
        print(f"\n[END] Done. Final score: {results['score']:.4f}")
    except KeyboardInterrupt:
        print("\n[END] Interrupted by user")
    except Exception as e:
        print(f"\n[END] ERROR: {e}")
        raise
    finally:
        if server_proc:
            print("[END] Shutting down local server...")
            server_proc.terminate()


if __name__ == "__main__":
    main()