""" inference.py — LLM inference loop for AdaptiveWorld environment (v2). Updated from Round 1 api-debug-env inference.py to: - Use AdaptiveAction / AdaptiveObservation v2 types - Support all 5 action types: call_api, probe_schema, query_history, declare_belief, submit_result - System prompt guides the agent to detect drift from evidence (NOT from an explicit drift_occurred field — that's removed in v2) - Extract belief_state from LLM output for belief scoring - Track task_reward and belief_accuracy separately across episodes """ import os import json import re from typing import Optional from dotenv import load_dotenv from openai import OpenAI from client import AdaptiveWorldEnv from models import AdaptiveAction, AdaptiveObservation load_dotenv() # ── Config ───────────────────────────────────────────────────────────────────── ENV_URL = os.getenv("ENV_URL", "http://localhost:7860") OPENAI_MODEL = os.getenv("OPENAI_MODEL", "meta-llama/Llama-3.3-70B-Instruct") MAX_STEPS = int(os.getenv("MAX_STEPS", "12")) DIFFICULTY = os.getenv("DIFFICULTY", "easy") # Use Hugging Face serverless router using local HF token import huggingface_hub hf_token = os.getenv("HF_TOKEN") or huggingface_hub.get_token() or "dummy" client = OpenAI( base_url="https://api.groq.com/openai/v1", api_key=hf_token ) env = AdaptiveWorldEnv(base_url=ENV_URL).sync() # ── System prompt ────────────────────────────────────────────────────────────── SYSTEM_PROMPT = """You are an agent completing professional tasks via HTTP APIs. IMPORTANT: The API world can mutate mid-episode without warning. - Field names may change (e.g. "qty" → "quantity") - Endpoints may move (e.g. /book → /v2/book) - Policies may change (new required fields, new auth schemes) - Status values may change silently (200 OK but different meaning) You will NOT be told when a drift occurs. You must detect it from evidence. AVAILABLE ACTION TYPES: 1. call_api — Make an HTTP request to the mock API 2. probe_schema — GET /openapi.json to inspect current API contract 3. query_history — Review your last N API responses, compare for changes 4. declare_belief — State your current understanding of the world (saves for later scoring) 5. submit_result — End the episode. Include your final belief_state about what changed. STRATEGY: - Start with probe_schema to understand the initial API contract - After any unexpected error (422, 404, 401), call query_history first to compare responses - If responses look the same but downstream fails silently, the MEANING may have changed - Use declare_belief to record discoveries as you go - Always submit_result with your final belief_state — this is scored separately from task completion* *Belief accuracy is tracked independently. You can score high on belief even if the task fails. RESPONSE FORMAT (JSON, always): { "action_type": "call_api", // or probe_schema, query_history, declare_belief, submit_result "method": "POST", // for call_api "url": "/mock_api/orders", // for call_api "headers": {"Content-Type": "application/json"}, // for call_api "body": {"quantity": 2, "product_id": 5}, // for call_api "query_params": {}, // for call_api GET requests "belief_state": { // for declare_belief or submit_result "order_field": "quantity", "required_extra": "customer_id" }, "history_steps": 3 // for query_history } """ def build_user_message(obs: AdaptiveObservation, step: int) -> str: """Build the per-step user message from the current observation.""" parts = [ f"=== Step {step} ===", f"Task: {obs.task_description}", f"Domain: {obs.domain}", f"Progress: {obs.current_step}/{obs.max_steps} steps", f"Difficulty level: {obs.difficulty_level}/3", ] if obs.prior_world_model: parts.append( f"\\nYour prior world model (from previous episodes in this domain):\\n" f"{json.dumps(obs.prior_world_model, indent=2)}\\n" "NOTE: This may be stale if the world has changed since last episode." ) if obs.last_status_code > 0: parts.append(f"\\nLast API call:") parts.append(f" Status: {obs.last_status_code}") parts.append(f" Response: {obs.last_response_body[:800]}") if obs.episode_history: parts.append(f"\\nEpisode history (last {len(obs.episode_history)} steps):") parts.append(json.dumps(obs.episode_history, indent=2)[:1500]) parts.append(f"\\nFeedback: {obs.step_feedback}") parts.append( "\\nRespond with a JSON action object. " "If you detect drift, use declare_belief to record it before continuing. " "When ready to finish, use submit_result with your final belief_state." ) return "\\n".join(parts) def parse_action(llm_output: str) -> AdaptiveAction: """Parse LLM output to AdaptiveAction. Falls back to safe default on parse error.""" try: # Extract JSON from markdown code blocks or raw output match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', llm_output, re.DOTALL) if match: raw = match.group(1) else: raw_match = re.search(r'\{.*\}', llm_output, re.DOTALL) raw = raw_match.group() if raw_match else "{}" data = json.loads(raw) return AdaptiveAction( action_type=data.get("action_type", "call_api"), method=data.get("method", "GET"), url=data.get("url", "/mock_api/orders"), headers=data.get("headers", {}), body=data.get("body", {}), query_params=data.get("query_params", {}), belief_state=data.get("belief_state"), history_steps=int(data.get("history_steps", 3)), ) except Exception as e: print(f"[parse_action] Parse error: {e}. Defaulting to probe_schema.") return AdaptiveAction(action_type="probe_schema") def run_episode( scenario_id: str = "auto", difficulty: str = "easy", verbose: bool = True, ) -> dict: """ Run one complete episode. Returns: dict with task_reward, belief_accuracy, combined_reward, steps_taken, drift_detected """ step_result = env.reset(scenario_id=scenario_id, difficulty=difficulty) if verbose: print(f"\\n{'='*60}") print(f"EPISODE START") print(f"Task: {step_result.observation.task_description}") print(f"Domain: {step_result.observation.domain}") print(f"Difficulty: {step_result.observation.difficulty_level}/3") print(f"Prior beliefs: {step_result.observation.prior_world_model}") print(f"{'='*60}") conversation: list = [{"role": "system", "content": SYSTEM_PROMPT}] step = 0 final_reward = 0.0 task_reward = 0.0 belief_accuracy = 0.0 while not step_result.done and step < MAX_STEPS: step += 1 user_msg = build_user_message(step_result.observation, step) conversation.append({"role": "user", "content": user_msg}) # LLM call response = client.chat.completions.create( model=OPENAI_MODEL, messages=conversation, temperature=0.2, max_tokens=512, ) llm_output = response.choices[0].message.content conversation.append({"role": "assistant", "content": llm_output}) if verbose: print(f"\n--- Step {step} ---") print(f"LLM: {llm_output[:400]}") action = parse_action(llm_output) if verbose: print(f"Action: {action.action_type} | {action.method} {action.url}") step_result = env.step(action) final_reward = step_result.reward if verbose: print(f"Status: {step_result.observation.last_status_code} | Reward: {step_result.reward:.4f}") if step_result.observation.step_feedback: print(f"Feedback: {step_result.observation.step_feedback}") if step_result.done: task_reward = step_result.observation.task_reward belief_accuracy = step_result.observation.belief_accuracy break if verbose: print(f"\n{'='*60}") print(f"EPISODE COMPLETE") print(f"Steps taken: {step}") print(f"Task reward: {task_reward:.4f}") print(f"Belief accuracy: {belief_accuracy:.4f}") print(f"Combined reward: {final_reward:.4f}") print(f"{'='*60}\n") return { "task_reward": task_reward, "belief_accuracy": belief_accuracy, "combined_reward": final_reward, "steps_taken": step, } def run_evaluation( n_episodes: int = 10, scenario_id: str = "auto", difficulty: str = "easy", ) -> dict: """ Run N episodes and return aggregate statistics. Used for Day 3 reward signal verification. """ results = [] for i in range(n_episodes): print(f"\n=== Episode {i+1}/{n_episodes} ===") r = run_episode(scenario_id=scenario_id, difficulty=difficulty, verbose=True) results.append(r) task_rewards = [r["task_reward"] for r in results] belief_accs = [r["belief_accuracy"] for r in results] combined = [r["combined_reward"] for r in results] summary = { "n_episodes": n_episodes, "avg_task_reward": sum(task_rewards) / len(task_rewards), "avg_belief_accuracy": sum(belief_accs) / len(belief_accs), "avg_combined_reward": sum(combined) / len(combined), "max_task_reward": max(task_rewards), "max_belief_accuracy": max(belief_accs), } print("\n=== EVALUATION SUMMARY ===") for k, v in summary.items(): print(f" {k}: {v:.4f}" if isinstance(v, float) else f" {k}: {v}") return summary if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="AdaptiveWorld Inference") parser.add_argument("--scenario", default="auto", help="Scenario ID or 'auto'") parser.add_argument("--difficulty", default="easy", choices=["easy", "medium", "hard", "expert"]) parser.add_argument("--episodes", type=int, default=1) args = parser.parse_args() if args.episodes == 1: run_episode(scenario_id=args.scenario, difficulty=args.difficulty) else: run_evaluation( n_episodes=args.episodes, scenario_id=args.scenario, difficulty=args.difficulty, )