inference.py

"""
Baseline inference script — Customer Support Inbox OpenEnv.

STDOUT FORMAT (strict — deviations break evaluation scoring):
  [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>

Rules:
  - One [START] line at episode begin.
  - One [STEP] line per step, immediately after env.step() returns.
  - One [END] line after episode ends, always emitted (even on exception).
  - reward and rewards formatted to 2 decimal places.
  - done and success are lowercase booleans: true or false.
  - error is the raw error string, or null if none.
  - All fields on a single line with no newlines within a line.
  - Each task returns score in [0, 1].

Environment variables:
  API_BASE_URL   LLM API endpoint  (default: https://api.openai.com/v1)
  MODEL_NAME     Model identifier  (default: gpt-4o-mini)
  HF_TOKEN       API key (primary)
  OPENAI_API_KEY API key (fallback)
"""
from __future__ import annotations

import json
import os
import sys
import time
import textwrap
from typing import Any, Dict, List, Optional

# Ensure the package root is on sys.path so `environment` is importable
# regardless of the working directory the runner uses.
_here = os.path.dirname(os.path.abspath(__file__))
if _here not in sys.path:
    sys.path.insert(0, _here)

# ── OpenAI client ─────────────────────────────────────────────────────────────
try:
    from openai import OpenAI
except ImportError as _oi:
    OpenAI = None  # type: ignore[assignment,misc]
    print(f"WARNING: openai not installed ({_oi})", file=sys.stderr)

# ── Environment ───────────────────────────────────────────────────────────────
try:
    from environment import CustomerSupportEnv
    from environment.models import Action, ActionType, TicketCategory, TicketPriority
    from environment.tasks import TASK_DEFINITIONS
    _ENV_AVAILABLE = True
except Exception as _ie:
    _ENV_AVAILABLE = False
    CustomerSupportEnv = None  # type: ignore[assignment,misc]
    Action = ActionType = TicketCategory = TicketPriority = TASK_DEFINITIONS = None  # type: ignore[assignment]
    print(f"WARNING: environment import failed ({_ie})", file=sys.stderr)

# ── Configuration ─────────────────────────────────────────────────────────────
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME   = os.getenv("MODEL_NAME",   "gpt-4o-mini")
API_KEY      = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY") or "sk-placeholder"

BENCHMARK    = "customer-support-inbox"
SEED         = 42
MAX_RETRIES  = 3

# Fixed ticket per task for reproducible baseline
TASK_TICKET_MAP = {
    "task1": "T001",   # billing duplicate-charge (easy)
    "task2": "T005",   # missing package (medium)
    "task3": "T009",   # VIP retention (hard)
}

# ── OpenAI client setup (deferred to avoid module-level crash) ────────────────
client: Optional[OpenAI] = None

def _get_client() -> OpenAI:
    global client
    if client is None:
        client = OpenAI(api_key=API_KEY, base_url=API_BASE_URL)
    return client

# ── Prompts ───────────────────────────────────────────────────────────────────
SYSTEM_PROMPT = textwrap.dedent("""
    You are an expert customer support agent handling a support inbox.

    Available actions and their JSON schemas:
    - classify:     {"action_type":"classify","category":"<cat>","priority":"<pri>"}
    - respond:      {"action_type":"respond","response_text":"<message>"}
    - request_info: {"action_type":"request_info","response_text":"<question>"}
    - escalate:     {"action_type":"escalate","escalation_reason":"<why>","escalation_team":"<tier2|billing_specialist|engineering>"}
    - resolve:      {"action_type":"resolve","resolution_notes":"<summary>","resolution_category":"<fixed|refunded|explained|workaround>"}
    - tag:          {"action_type":"tag","tags":["tag1","tag2"]}

    Categories: billing | technical | shipping | returns | account | complaint | general
    Priorities: low | medium | high | urgent

    Strategy:
    1. ALWAYS classify first (category + priority) — do this as your very first action.
    2. Send an empathetic response acknowledging the issue.
    3. Escalate if the issue needs a specialist (VIP complaints → tier2, bugs → engineering, billing disputes → billing_specialist).
    4. Resolve with clear notes describing what was done.

    For angry/VIP customers: use empathy phrases like "I sincerely apologize", "I completely understand",
    "we value your loyalty", commit to a callback ("I will personally follow up within the hour").

    Respond with a SINGLE valid JSON object only. No explanation, no markdown, no extra text.
""").strip()


def build_user_prompt(obs: Dict[str, Any]) -> str:
    """Build the user-turn prompt from an observation dict."""
    customer = obs.get("customer", {})
    conversation = obs.get("conversation", [])
    # Show last 5 messages for context
    recent = conversation[-5:]
    conv_lines = "\n".join(
        f"  [{m['role'].upper()}]: {m['content']}"
        for m in recent
    )
    return (
        f"=== TICKET {obs.get('ticket_id')} ===\n"
        f"Subject   : {obs.get('subject')}\n"
        f"Status    : {obs.get('status')} | "
        f"Category  : {obs.get('category') or 'UNCLASSIFIED'} | "
        f"Priority  : {obs.get('priority') or 'UNSET'}\n"
        f"Customer  : {customer.get('name')} "
        f"({customer.get('account_tier')} tier, "
        f"sentiment={customer.get('sentiment', 'neutral')})\n"
        f"SLA       : {obs.get('time_to_sla')}\n"
        f"Turn      : {obs.get('turn_count')+1}/{obs.get('turn_count',0)+obs.get('turns_remaining',1)}\n"
        f"\n=== CONVERSATION (last {len(recent)}) ===\n"
        f"{conv_lines}\n"
        f"\n=== TASK: {obs.get('task_id')} ===\n"
        f"{obs.get('task_description','')}\n"
        f"\nObjectives:\n"
        + "\n".join(f"  - {o}" for o in obs.get("task_objectives", []))
        + f"\n\nAvailable actions: {obs.get('available_actions')}\n"
        f"Last result: {obs.get('last_action_result') or 'n/a'}\n"
        f"\nRespond with JSON only."
    )


def call_llm(messages: List[Dict[str, str]], retry: int = 0) -> str:
    """Call LLM with exponential back-off. Returns raw string."""
    try:
        resp = _get_client().chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            temperature=0.1,
            max_tokens=400,
            response_format={"type": "json_object"},
        )
        return resp.choices[0].message.content or ""
    except Exception:
        # Some APIs don't support response_format — retry without it
        try:
            resp = _get_client().chat.completions.create(
                model=MODEL_NAME,
                messages=messages,
                temperature=0.1,
                max_tokens=400,
            )
            return resp.choices[0].message.content or ""
        except Exception:
            pass
        if retry < MAX_RETRIES:
            time.sleep(2 ** retry)
            return call_llm(messages, retry + 1)
        # Final fallback — safe classify action
        return json.dumps({"action_type": "classify",
                           "category": "general", "priority": "medium"})


def parse_action(raw: str, available: List[str]) -> tuple[Action, Optional[str]]:
    """
    Parse LLM JSON into Action. Returns (action, error_str_or_None).
    On parse failure, returns a safe fallback + error string.
    """
    try:
        text = raw.strip()
        # Strip markdown fences if present
        if "```" in text:
            parts = text.split("```")
            text = parts[1].lstrip("json").strip() if len(parts) > 1 else text
        data = json.loads(text)

        atype_str = str(data.get("action_type", "classify")).lower()
        if atype_str not in available:
            atype_str = available[0] if available else "classify"

        atype = ActionType(atype_str)

        action = Action(
            action_type=atype,
            category=TicketCategory(data["category"])
                if data.get("category") and atype == ActionType.CLASSIFY else None,
            priority=TicketPriority(data["priority"])
                if data.get("priority") and atype == ActionType.CLASSIFY else None,
            response_text=data.get("response_text"),
            escalation_reason=data.get("escalation_reason"),
            escalation_team=data.get("escalation_team"),
            resolution_notes=data.get("resolution_notes"),
            resolution_category=data.get("resolution_category"),
            tags=data.get("tags"),
            assigned_to=data.get("assigned_to"),
        )
        return action, None

    except Exception as exc:
        fallback = Action(
            action_type=ActionType.CLASSIFY,
            category=TicketCategory.GENERAL,
            priority=TicketPriority.MEDIUM,
        )
        return fallback, str(exc)


def run_episode(task_id: str, ticket_id: str) -> Dict[str, Any]:
    """
    Run one full episode. Emits strict [START]/[STEP]/[END] lines.
    Returns result dict. Never raises — all exceptions are caught internally.
    """
    step_rewards: List[float] = []
    steps        = 0
    success      = False
    final_score  = 0.0
    last_error: Optional[str] = None
    task_name    = task_id  # fallback if TASK_DEFINITIONS lookup fails

    try:
        task_def  = TASK_DEFINITIONS[task_id]  # type: ignore[index]
        task_name = task_def.name.replace(" ", "-").replace("&", "and").lower()
    except Exception as e:
        print(f"[START] task={task_name} env={BENCHMARK} model={MODEL_NAME}", flush=True)
        print(f"[END]   success=false steps=0 score=0.00 rewards=0.00", flush=True)
        return {"task_id": task_id, "task_name": task_name, "ticket_id": ticket_id,
                "difficulty": "unknown", "score": 0.0, "success": False, "steps": 0, "rewards": []}

    if not _ENV_AVAILABLE:
        print(f"[START] task={task_name} env={BENCHMARK} model={MODEL_NAME}", flush=True)
        print(f"[END]   success=false steps=0 score=0.00 rewards=0.00", flush=True)
        return {"task_id": task_id, "task_name": task_name, "ticket_id": ticket_id,
                "difficulty": task_def.difficulty, "score": 0.0, "success": False, "steps": 0, "rewards": []}

    # ── [START] ───────────────────────────────────────────────────────────
    print(f"[START] task={task_name} env={BENCHMARK} model={MODEL_NAME}", flush=True)

    try:
        env  = CustomerSupportEnv(seed=SEED)
        obs  = env.reset(task_id=task_id, ticket_id=ticket_id, seed=SEED)
        obs_dict = obs.model_dump()

        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        done     = False

        while not done and steps < task_def.max_turns:
            steps += 1
            last_error = None

            # Build prompt + call LLM
            user_prompt = build_user_prompt(obs_dict)
            messages.append({"role": "user", "content": user_prompt})
            llm_raw = call_llm(messages)
            messages.append({"role": "assistant", "content": llm_raw})

            # Parse action
            available = obs_dict.get("available_actions", ["classify"])
            action, parse_err = parse_action(llm_raw, available)
            if parse_err:
                last_error = f"parse_error:{parse_err[:80]}"

            # Step environment
            try:
                obs, reward, done, info = env.step(action)
                obs_dict = obs.model_dump()
                step_rewards.append(reward.score)
                error_str = last_error or "null"
                done_str  = "true" if done else "false"

                # ── [STEP] ────────────────────────────────────────────────
                print(
                    f"[STEP]  step={steps} "
                    f"action={action.action_type.value} "
                    f"reward={reward.score:.2f} "
                    f"done={done_str} "
                    f"error={error_str}",
                    flush=True,
                )

            except Exception as step_exc:
                last_error = f"step_error:{str(step_exc)[:80]}"
                step_rewards.append(0.0)
                print(
                    f"[STEP]  step={steps} "
                    f"action={action.action_type.value} "
                    f"reward=0.00 "
                    f"done=false "
                    f"error={last_error}",
                    flush=True,
                )
                break

        # Compute final score from grader info
        final_info   = info if "info" in locals() else {}  # type: ignore[name-defined]
        final_score  = float(final_info.get("final_grader_score",
                        sum(step_rewards) / max(len(step_rewards), 1)))
        final_score  = round(min(max(final_score, 0.0), 1.0), 2)
        success      = final_score >= task_def.min_score_to_pass

    except Exception as ep_exc:
        last_error  = str(ep_exc)[:120]
        final_score = 0.0
        success     = False
        if not step_rewards:
            step_rewards = [0.0]

    # Rewards list, 2 dp
    rewards_str = ",".join(f"{r:.2f}" for r in step_rewards) if step_rewards else "0.00"
    success_str = "true" if success else "false"

    # ── [END] ─────────────────────────────────────────────────────────────
    print(
        f"[END]   success={success_str} "
        f"steps={steps} "
        f"score={final_score:.2f} "
        f"rewards={rewards_str}",
        flush=True,
    )

    return {
        "task_id":    task_id,
        "task_name":  task_name,
        "ticket_id":  ticket_id,
        "difficulty": task_def.difficulty,
        "score":      final_score,
        "success":    success,
        "steps":      steps,
        "rewards":    step_rewards,
    }


def main() -> int:
    print("=" * 62, flush=True)
    print(f"  Customer Support Inbox — OpenEnv Baseline Inference", flush=True)
    print(f"  Model : {MODEL_NAME}", flush=True)
    print(f"  API   : {API_BASE_URL}", flush=True)
    print(f"  Seed  : {SEED}", flush=True)
    print("=" * 62, flush=True)

    results = []

    for task_id in ["task1", "task2", "task3"]:
        ticket_id = TASK_TICKET_MAP.get(task_id, "T001")
        print(f"\n--- {task_id} | ticket {ticket_id} ---", flush=True)
        try:
            result = run_episode(task_id, ticket_id)
        except Exception as e:
            print(f"[END]   success=false steps=0 score=0.00 rewards=0.00", flush=True)
            result = {"task_id": task_id, "task_name": task_id, "ticket_id": ticket_id,
                      "difficulty": "unknown", "score": 0.0, "success": False, "steps": 0, "rewards": []}
        results.append(result)
        time.sleep(0.5)

    # ── Summary table ─────────────────────────────────────────────────────
    print("\n" + "=" * 62, flush=True)
    print("  RESULTS SUMMARY", flush=True)
    print("=" * 62, flush=True)
    for r in results:
        flag = "PASS" if r["success"] else "FAIL"
        print(
            f"  {r['task_id']} ({r['difficulty']:6s}) | "
            f"score={r['score']:.2f} | "
            f"steps={r['steps']:2d} | {flag}",
            flush=True,
        )

    avg   = round(sum(r["score"] for r in results) / len(results), 2)
    npassed = sum(1 for r in results if r["success"])
    print("-" * 62, flush=True)
    print(f"  Average score : {avg:.2f}", flush=True)
    print(f"  Tasks passed  : {npassed}/{len(results)}", flush=True)
    print("=" * 62, flush=True)

    # Machine-readable summary (stdout, for automated runners)
    summary = {
        "benchmark": BENCHMARK,
        "model":     MODEL_NAME,
        "seed":      SEED,
        "results":   results,
        "average_score": avg,
        "tasks_passed":  npassed,
    }
    print(f"\nJSON_SUMMARY: {json.dumps(summary)}", flush=True)

    # Always return 0 — the evaluator scores based on [END] lines, not exit code.
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except BaseException as _top_exc:
        print(f"[END]   success=false steps=0 score=0.00 rewards=0.00", flush=True)
        print(f"FATAL: {type(_top_exc).__name__}: {_top_exc}", file=sys.stderr, flush=True)
        sys.exit(0)