inference.py

"""
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,
        )