inference.py

"""
HARvestGym — Inference Script
==============================

Runs the RL agent (driven by an LLM via OpenAI client) through three tasks:
  1. har_classify_easy   — Template 1: list products in a category
  2. har_classify_medium — Template 3: add product to guest cart
  3. har_pipeline_hard   — Template 6: complete guest checkout

STDOUT FORMAT (strictly enforced by hackathon):
  [START] task=<task_name> env=<benchmark> model=<model_name>
  [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
  [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>

Usage:
  HF_TOKEN=hf_xxx uv run inference.py
  HF_TOKEN=hf_xxx MODEL_NAME=Qwen/Qwen2.5-72B-Instruct uv run inference.py
"""

import asyncio
import json
import os
import sys
import textwrap
from typing import Any, List, Optional

from openai import OpenAI

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    raise ValueError(
        "HF_TOKEN environment variable is required but not set. "
        "Export it with: export HF_TOKEN=hf_your_token_here"
    )

BENCHMARK = "harvgym"
MAX_STEPS = 20
TEMPERATURE = 0.2        # Lower temp → more deterministic tool calls
MAX_TOKENS = 1024        # More room for reasoning + JSON
SUCCESS_SCORE_THRESHOLD = 0.5

# Task definitions: use FIXED task descriptions so the model always knows
# exactly what to do (env.reset() may randomize, but we tell it the target)
TASKS = [
    {
        "task_name": "har_classify_easy",
        "template_id": 1,
        "description": "List products in the 'Gear' category",
        "app_base_url": "http://ec2-16-59-2-56.us-east-2.compute.amazonaws.com:7770/",
        "difficulty": "easy",
    },
    {
        "task_name": "har_classify_medium",
        "template_id": 3,
        "description": "Add 'Radiant Tee' (SKU: MH01-XS-Black) to a guest cart",
        "app_base_url": "http://ec2-16-59-2-56.us-east-2.compute.amazonaws.com:7770/",
        "difficulty": "medium",
    },
    {
        "task_name": "har_pipeline_hard",
        "template_id": 6,
        "description": "Complete a full guest checkout for 'Radiant Tee' (SKU: MH01-XS-Black)",
        "app_base_url": "http://ec2-16-59-2-56.us-east-2.compute.amazonaws.com:7770/",
        "difficulty": "hard",
    },
]

# ---------------------------------------------------------------------------
# Logging helpers (hackathon format)
# ---------------------------------------------------------------------------

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


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    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,
    )


# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = textwrap.dedent("""
You are an API agent. Your goal is to complete a real-world task by calling the correct
sequence of HTTP API endpoints on a live Magento shopping application.

You have exactly these tools available. Output ONE tool call per turn as a JSON object:

1. browser_agent(task, url)
   → Discovers API endpoints for this app from recorded traffic + catalog.
   → Call this FIRST and ONLY ONCE (at step 1).
   → Returns: {"app": "shopping", "total_endpoints": N, "endpoints": [...]}

2. search_endpoints(query)
   → Semantic/keyword search over all discovered endpoint schemas.
   → Returns full endpoint details: method, path, required parameters, auth.
   → Use this to find the exact path and params before making curl calls.
   → Example: search_endpoints("guest cart create") → POST /rest/V1/guest-carts

3. curl_exec(command)
   → Execute an HTTP request. Returns {"status_code": N, "body": {...}}.
   → Session cookies are auto-injected. Do NOT manually set Cookie headers.
   → Always use: curl -s -X METHOD 'URL' -H 'Content-Type: application/json' -d '{...}'

4. search_episode_data(query)
   → Search all prior API responses in this episode for a value.
   → Use when a prior response was long and you need to find a specific item/ID.

5. done(result?)
   → Call ONLY when the task is fully complete. Do not call early.

RULES:
- Output ONLY a JSON object: {"tool": "...", "args": {...}}. No explanation, no markdown.
- Step 1: ALWAYS call browser_agent to discover endpoints.
- Step 2+: Use search_endpoints to find the right endpoint before calling curl_exec.
- Read IDs and values (cart_id, sku, item_id) from LAST TOOL RESULT in the context.
- Magento REST API base: http://host/rest/V1/
- To add an item to guest cart: POST /rest/V1/guest-carts/{cartId}/items
  Body: {"cartItem": {"sku": "SKU", "qty": 1, "quote_id": "{cartId}"}}
- Do NOT call done() until the task is actually accomplished.

EXAMPLE SEQUENCE for "Add product to guest cart":
  Step 1: {"tool": "browser_agent", "args": {"task": "Add product to guest cart", "url": "http://..."}}
  Step 2: {"tool": "search_endpoints", "args": {"query": "create guest cart"}}
  Step 3: {"tool": "curl_exec", "args": {"command": "curl -s -X POST 'http://.../rest/V1/guest-carts' -H 'Content-Type: application/json'"}}
  Step 4: {"tool": "search_endpoints", "args": {"query": "add item to guest cart"}}
  Step 5: {"tool": "curl_exec", "args": {"command": "curl -s -X POST 'http://.../rest/V1/guest-carts/CART_ID/items' -H 'Content-Type: application/json' -d '{\"cartItem\":{\"sku\":\"MH01-XS-Black\",\"qty\":1,\"quote_id\":\"CART_ID\"}}'"}}
  Step 6: {"tool": "done", "args": {}}
""").strip()


# ---------------------------------------------------------------------------
# LLM agent loop
# ---------------------------------------------------------------------------

def _format_result_for_context(result: Any, max_chars: int = 3000) -> str:
    """Format tool result for the LLM context — more generous truncation."""
    if result is None:
        return "null"
    try:
        text = json.dumps(result, indent=2)
    except Exception:
        text = str(result)

    if len(text) <= max_chars:
        return text

    # Smart truncation: keep beginning (has structure/IDs) and hint at truncation
    kept = text[:max_chars]
    # Try to close the JSON gracefully at the last complete line
    last_newline = kept.rfind("\n")
    if last_newline > max_chars * 0.8:
        kept = kept[:last_newline]
    return kept + f"\n... [truncated, {len(text) - max_chars} chars omitted — use search_episode_data to find specific values]"


def build_user_prompt(task_desc: str, app_base_url: str, step: int,
                      last_result: Any, history: List[dict],
                      session_state: dict) -> str:
    """Build the user prompt for each step."""
    history_lines = []
    if history:
        # Show last 8 steps with meaningful result summaries
        for h in history[-8:]:
            result = h.get("result", {})
            # For curl results: show status_code + first 200 chars of body
            if isinstance(result, dict) and "status_code" in result:
                body_preview = str(result.get("body", ""))[:300]
                result_summary = f'status={result["status_code"]} body={body_preview}'
            else:
                result_summary = str(result)[:300]
            history_lines.append(
                f"  Step {h['step']}: {h['tool']}({json.dumps(h.get('args', {}))[:100]}) "
                f"→ {result_summary}"
            )

    session_str = json.dumps(session_state, indent=2)[:500] if session_state else "{}"
    last_result_str = _format_result_for_context(last_result)

    return textwrap.dedent(f"""
    TASK: {task_desc}
    APP URL: {app_base_url}
    STEP: {step}/{MAX_STEPS}

    SESSION STATE (cookies/tokens auto-managed):
    {session_str}

    LAST TOOL RESULT:
    {last_result_str}

    HISTORY (last {len(history_lines)} steps):
    {chr(10).join(history_lines) if history_lines else "  (none yet)"}

    What is your next tool call? Output ONLY the JSON object.
    """).strip()


def get_model_action(client: OpenAI, task_desc: str, app_base_url: str,
                     step: int, last_result: Any, history: List[dict],
                     session_state: dict) -> dict:
    """Ask the LLM for the next action. Returns parsed tool call dict."""
    user_prompt = build_user_prompt(task_desc, app_base_url, step,
                                    last_result, history, session_state)
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            stream=False,
        )
        text = (completion.choices[0].message.content or "").strip()

        # Strip markdown code fences
        if "```json" in text:
            text = text.split("```json")[1].split("```")[0].strip()
        elif "```" in text:
            text = text.split("```")[1].split("```")[0].strip()

        # Extract first {...} block
        start = text.find("{")
        end = text.rfind("}") + 1
        if start >= 0 and end > start:
            text = text[start:end]

        parsed = json.loads(text)
        if "tool" in parsed:
            return parsed

        return {"tool": "done", "args": {"result": "Model returned non-tool response"}}

    except json.JSONDecodeError:
        # Fallback heuristics
        if step == 1:
            return {"tool": "browser_agent", "args": {"task": task_desc, "url": app_base_url}}
        if "browser_agent" in text:
            return {"tool": "browser_agent", "args": {"task": task_desc, "url": app_base_url}}
        if "done" in text.lower():
            return {"tool": "done", "args": {}}
        return {"tool": "done", "args": {"result": f"Parse error: {text[:100]}"}}
    except Exception as exc:
        print(f"[DEBUG] LLM call failed: {exc}", flush=True)
        if step == 1:
            return {"tool": "browser_agent", "args": {"task": task_desc, "url": app_base_url}}
        return {"tool": "done", "args": {"result": f"LLM error: {exc}"}}


# ---------------------------------------------------------------------------
# Single task episode runner
# ---------------------------------------------------------------------------

async def run_episode(task_config: dict, client: OpenAI) -> dict:
    """
    Run a single episode for one task.
    Returns: {"task_name", "success", "steps", "score", "rewards"}
    """
    from server.models import HARvestGymEnvironment, HarvestGymAction

    task_name = task_config["task_name"]
    template_id = task_config["template_id"]
    task_description = task_config["description"]
    app_base_url = task_config["app_base_url"]

    # Pin the template via env var so reset() samples from the right pool
    os.environ["HARVGYM_TASK"] = task_name   # use name, not int

    env = HARvestGymEnvironment()

    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    last_result = None
    history: List[dict] = []
    session_state: dict = {}

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    try:
        obs = env.reset()
        # CRITICAL: use the env-sampled task description — the judge grades exactly
        # what env.reset() returned (random category/product), not our hardcoded string.
        task_desc = obs.task or task_description
        base_url = obs.app_base_url or app_base_url

        for step in range(1, MAX_STEPS + 1):
            if getattr(obs, "done", False):
                break

            action_dict = get_model_action(
                client=client,
                task_desc=task_desc,
                app_base_url=base_url,
                step=step,
                last_result=last_result,
                history=history,
                session_state=session_state,
            )

            tool = action_dict.get("tool", "done")
            args = action_dict.get("args", {})

            action_str = f"{tool}({json.dumps(args)[:150]})"
            error_str = None

            try:
                action = HarvestGymAction(tool=tool, args=args)
                obs = env.step(action)

                reward = float(obs.reward or 0.0)
                done = bool(obs.done)
                last_result = obs.last_tool_result
                session_state = dict(obs.session_state or {})

                history.append({
                    "step": step,
                    "tool": tool,
                    "args": args,
                    "result": last_result,
                })

            except Exception as exc:
                reward = -0.1
                done = False
                error_str = str(exc)[:200]

            rewards.append(reward)
            steps_taken = step
            log_step(step=step, action=action_str, reward=reward, done=done, error=error_str)

            if done:
                break

        # Score: terminal reward from judge dominates.
        # Reward range by design: terminal success = +2 to +5, terminal fail = -1.5
        # Use a generous baseline so partial credit shows up.
        total_reward = sum(rewards)
        # Normalise to [0,1]: shift by +1.5 (min), divide by max-possible per task
        # Template 1 max=2, Template 3 max=3.5, Template 6 max=5 → use 5.0 as ceiling
        score = max(0.0, min(1.0, (total_reward + 1.5) / (5.0 + 1.5)))
        success = total_reward >= 0.5   # any positive terminal reward = success

    except Exception as exc:
        error_str = str(exc)[:200]
        print(f"[DEBUG] Episode error: {error_str}", flush=True)
    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return {
        "task_name": task_name,
        "success": success,
        "steps": steps_taken,
        "score": score,
        "rewards": rewards,
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

async def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)

    results = []
    for task_config in TASKS:
        result = await run_episode(task_config, client)
        results.append(result)

    # Summary
    print("\n[SUMMARY]", flush=True)
    for r in results:
        status = "PASS" if r["success"] else "FAIL"
        print(
            f"  [{status}] {r['task_name']} — score={r['score']:.2f} steps={r['steps']}",
            flush=True,
        )

    overall_score = sum(r["score"] for r in results) / len(results) if results else 0.0
    print(f"\n  overall_score={overall_score:.2f}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())