#!/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()