inference.py

"""
inference.py — Baseline inference script for CustomerSupportTriage-v0
=======================================================================

Runs an LLM agent (via OpenAI-compatible client) against all 3 tasks
and emits the mandatory [START] / [STEP] / [END] log format.

Environment variables
---------------------
API_BASE_URL   LLM endpoint  (default: https://router.huggingface.co/v1)
MODEL_NAME     Model to use  (default: Qwen/Qwen2.5-72B-Instruct)
HF_TOKEN       API key
ENV_BASE_URL   Triage env URL (default: http://localhost:7860)

Usage
-----
    python inference.py
    # or against a deployed HF Space:
    ENV_BASE_URL=https://your-space.hf.space python inference.py
"""

from __future__ import annotations

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

import requests
from openai import OpenAI

# ── Config ────────────────────────────────────────────────────────────────────
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") or os.getenv("API_KEY", "")
ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:7860").rstrip("/")

TASKS        = ["easy", "medium", "hard"]
TEMPERATURE  = 0.2      # low temp for reproducibility
MAX_TOKENS   = 1200
MAX_STEPS    = 5        # safety cap per episode
SUCCESS_THRESHOLD = 0.6

# ── Logging helpers ───────────────────────────────────────────────────────────

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()
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"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} "
        f"score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ── Env client ────────────────────────────────────────────────────────────────

def env_reset(task: str, seed: int = 42) -> Dict[str, Any]:
    r = requests.post(
        f"{ENV_BASE_URL}/reset",
        json={"task": task, "seed": seed},
        timeout=30,
    )
    r.raise_for_status()
    return r.json()


def env_step(actions: List[Dict]) -> Dict[str, Any]:
    r = requests.post(
        f"{ENV_BASE_URL}/step",
        json={"actions": actions},
        timeout=60,
    )
    r.raise_for_status()
    return r.json()


# ── Prompt builder ────────────────────────────────────────────────────────────

SYSTEM_PROMPT = textwrap.dedent("""
You are an expert customer support manager triaging a queue of support tickets.

For each ticket you must output a JSON object with EXACTLY these fields:
{
  "ticket_id":   "<string>",
  "priority":    "low" | "medium" | "high" | "urgent",
  "department":  "billing" | "technical" | "shipping" | "returns" | "general" | "escalation",
  "response":    "<1-3 sentence draft reply to the customer>",
  "needs_human": true | false,
  "reasoning":   "<brief chain of thought>"
}

Guidelines:
- "urgent": production down, security breach, legal threat, hours-away deadline
- "high": significant impact, billing errors, outages affecting work
- "medium": moderately impactful issues, can wait 1-2 days
- "low": feature requests, general questions, feedback
- Routing: billing=payment/invoice/refund, technical=bugs/API/performance,
  shipping=delivery/tracking, returns=damaged/return, 
  escalation=legal/security/compliance/account-suspension/press,
  general=questions/feature-requests/onboarding
- needs_human=true for: legal threats, security incidents, angry enterprise customers, 
  accessibility complaints, any ticket routed to escalation
- Response must directly address the customer's issue, be empathetic, and give next steps

Respond with a JSON array of action objects, one per ticket. Nothing else.
""").strip()


def build_user_prompt(queue: List[Dict]) -> str:
    tickets_text = []
    for t in queue:
        tickets_text.append(
            f"--- Ticket {t['ticket_id']} ---\n"
            f"Subject: {t['subject']}\n"
            f"Customer: {t['customer_name']} ({t['customer_tier']} tier)\n"
            f"Sentiment: {t['sentiment']}\n"
            f"Tags: {', '.join(t['tags']) if t['tags'] else 'none'}\n"
            f"Body:\n{t['body']}"
        )
    return "Triage all of the following tickets:\n\n" + "\n\n".join(tickets_text)


# ── LLM call ──────────────────────────────────────────────────────────────────

def get_agent_actions(
    client: OpenAI,
    queue: List[Dict],
) -> tuple[List[Dict], str]:
    """Call LLM and return (parsed_actions, raw_response_string)."""
    user_prompt = build_user_prompt(queue)
    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,
        )
        raw = (completion.choices[0].message.content or "").strip()
    except Exception as exc:
        print(f"[DEBUG] LLM call failed: {exc}", flush=True)
        # Fallback: medium priority / general for all tickets
        fallback = [
            {
                "ticket_id":   t["ticket_id"],
                "priority":    "medium",
                "department":  "general",
                "response":    f"Thank you for reaching out, {t['customer_name']}. We are reviewing your request and will respond shortly.",
                "needs_human": False,
                "reasoning":   "LLM unavailable, using fallback",
            }
            for t in queue
        ]
        return fallback, "FALLBACK"

    # Parse JSON — handle markdown fences
    json_str = raw
    if "```" in json_str:
        start = json_str.find("[")
        end   = json_str.rfind("]") + 1
        json_str = json_str[start:end] if start != -1 else json_str
    json_str = json_str.strip()
    if not json_str.startswith("["):
        start = json_str.find("[")
        if start != -1:
            json_str = json_str[start:]

    try:
        actions = json.loads(json_str)
        if isinstance(actions, dict):   # model returned single object
            actions = [actions]
    except json.JSONDecodeError as exc:
        print(f"[DEBUG] JSON parse error: {exc}\nRaw: {raw[:300]}", flush=True)
        actions = [
            {
                "ticket_id":   t["ticket_id"],
                "priority":    "medium",
                "department":  "general",
                "response":    f"Thank you for contacting us, {t['customer_name']}. Our team is looking into your request.",
                "needs_human": False,
                "reasoning":   "parse error fallback",
            }
            for t in queue
        ]
    return actions, raw


# ── Episode runner ────────────────────────────────────────────────────────────

def run_episode(client: OpenAI, task: str) -> float:
    """Run one episode on the given task. Returns final score (0–1)."""
    rewards:      List[float] = []
    steps_taken:  int         = 0
    score:        float       = 0.0
    success:      bool        = False

    log_start(task=task, env="CustomerSupportTriage-v0", model=MODEL_NAME)

    try:
        result    = env_reset(task=task, seed=42)
        obs       = result["observation"]
        queue     = obs["queue"]

        for step in range(1, MAX_STEPS + 1):
            if not queue or obs.get("time_remaining", 1) == 0:
                break

            actions, raw = get_agent_actions(client, queue)

            # Compact action string for log (list of ticket_id:priority)
            action_log = json.dumps(
                [{"id": a.get("ticket_id","?"), "p": a.get("priority","?"), "d": a.get("department","?")}
                 for a in actions]
            )

            try:
                step_result = env_step(actions)
            except Exception as exc:
                log_step(step=step, action=action_log, reward=0.0, done=True, error=str(exc))
                rewards.append(0.0)
                steps_taken = step
                break

            reward  = step_result.get("reward", 0.0)
            done    = step_result.get("done",   False)
            error   = step_result.get("info", {}).get("error")
            obs     = step_result["observation"]
            queue   = obs.get("queue", [])

            rewards.append(reward)
            steps_taken = step

            log_step(
                step=step,
                action=action_log,
                reward=reward,
                done=done,
                error=error,
            )

            if done:
                break

        score   = sum(rewards) / len(rewards) if rewards else 0.0
        score   = round(min(max(score, 0.0), 1.0), 3)
        success = score >= SUCCESS_THRESHOLD

    except Exception as exc:
        print(f"[DEBUG] Episode error: {exc}", flush=True)
        score   = 0.0
        success = False

    log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
    return score


# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "no-key")

    print(f"[INFO] Model:   {MODEL_NAME}",   flush=True)
    print(f"[INFO] API:     {API_BASE_URL}", flush=True)
    print(f"[INFO] Env:     {ENV_BASE_URL}", flush=True)
    print(f"[INFO] Tasks:   {TASKS}",        flush=True)
    print("", flush=True)

    all_scores: Dict[str, float] = {}
    for task in TASKS:
        score = run_episode(client, task)
        all_scores[task] = score
        print("", flush=True)

    print("=" * 50, flush=True)
    print("BASELINE SCORES", flush=True)
    for task, score in all_scores.items():
        bar = "#" * int(score * 20)
        print(f"  {task:<8} {score:.3f}  [{bar:<20}]", flush=True)
    overall = sum(all_scores.values()) / len(all_scores)
    print(f"  {'overall':<8} {overall:.3f}", flush=True)
    print("=" * 50, flush=True)


if __name__ == "__main__":
    main()