inference.py

import asyncio
import json
import os
import re
import textwrap
from typing import List, Optional

from dotenv import load_dotenv
from openai import AsyncOpenAI

load_dotenv()

from support_env import SupportEnvWrapper, SupportAction

IMAGE_NAME = os.getenv("IMAGE_NAME")
API_BASE_URL = os.getenv("API_BASE_URL")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")   # HF_TOKEN is primary per spec
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
TASK_NAME = os.getenv("SUPPORT_ENV_TASK", "easy")
BENCHMARK = os.getenv("SUPPORT_ENV_BENCHMARK", "support_env")

MAX_STEPS = 5
TEMPERATURE = 0.2
MAX_TOKENS = 500
SUCCESS_SCORE_THRESHOLD = 0.5

_MAX_REWARD_PER_STEP = 1.0
MAX_TOTAL_REWARD = MAX_STEPS * _MAX_REWARD_PER_STEP

SYSTEM_PROMPT = textwrap.dedent(
    """
    You are an autonomous customer support agent for a large e-commerce and SaaS company.

    EVERY episode runs through exactly 3 phases. Complete each phase in order:

    ── PHASE 1: TRIAGE ──────────────────────────────────────────────────────
    Read the ticket and classify the issue type.
    Action: {"action_type": "classify"}
    (No team, no response needed in this phase)

    ── PHASE 2: ROUTE ───────────────────────────────────────────────────────
    Assign the ticket to the correct specialist team.
    Action: {"action_type": "assign", "team": "<team_name>"}
    Team directory:
      logistics_team      → lost packages, shipping, delivery
      tech_support_team   → login, password, bugs, API errors, data loss
      safety_team         → defects, overheating, recalls, hazards
      finance_team        → invoices, billing, tax, payment issues
      orders_team         → bulk orders, corporate accounts, returns, subscriptions
      management_team     → escalated complaints, refund delays, manager requests

    ── PHASE 3: RESOLVE ─────────────────────────────────────────────────────
    Take the ONE correct resolution action:
      escalate — customer demands manager, extreme anger, OR safety/data emergency
                 Always include team + response. PENALTY -0.30 if unnecessary.
      refund   — product definitively broken/wrong and company is at fault
                 Always include response. PENALTY -0.50 if unnecessary.
      respond  — customer needs information (policy, pricing, features, technical)
                 Write a detailed, specific response addressing their exact question.

    OUTPUT FORMAT — MANDATORY:
    Respond with ONLY a raw JSON object. No markdown, no explanation.
    {"action_type": "...", "team": "...", "response": "..."}
    """
).strip()


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} 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:.3f} rewards={rewards_str}", flush=True)


PHASE_PROMPTS = {
    1: (
        "CURRENT PHASE: 1 — TRIAGE\n"
        "Read the customer message carefully and classify the ticket.\n"
        "Identify the issue type and include it in the response field.\n"
        "Issue types: shipping, billing, technical, returns, safety, cancellation, complaint, orders, sales\n"
        "Output: {\"action_type\": \"classify\", \"response\": \"<issue_type>\"}"
    ),
    2: (
        "CURRENT PHASE: 2 — ROUTE\n"
        "The issue type has been revealed in the history. Assign to the correct team.\n"
        "Output: {\"action_type\": \"assign\", \"team\": \"<team_name>\"}"
    ),
    3: (
        "CURRENT PHASE: 3 — RESOLVE\n"
        "Choose the correct final action: escalate / refund / respond.\n"
        "Include team if escalating. Include a detailed response for respond/refund/escalate."
    ),
}


def build_user_prompt(step: int, obs: any, history: List[str], phase: int = 1) -> str:
    obs_dict = {
        "ticket_id": obs.ticket_id,
        "issue_type": obs.issue_type,
        "sentiment": obs.sentiment,
        "priority": obs.priority,
        "message": obs.message,
    }
    history_block = "\n".join(history[-6:]) if history else "None"
    phase_instruction = PHASE_PROMPTS.get(phase, PHASE_PROMPTS[3])
    return textwrap.dedent(
        f"""
        {phase_instruction}

        Step: {step}
        Ticket: {json.dumps(obs_dict, indent=2)}
        History:
        {history_block}

        Respond with JSON only.
        """
    ).strip()


def fallback_policy(phase: int, state: dict) -> dict:
    if phase == 1:
        msg = state.get("message", "").lower()
        if any(k in msg for k in ["ship", "deliver", "package", "track", "address"]):
            issue_type = "shipping"
        elif any(k in msg for k in ["bill", "charge", "invoice", "payment", "refund", "price"]):
            issue_type = "billing"
        elif any(k in msg for k in ["login", "password", "bug", "api", "error", "sync", "data", "account"]):
            issue_type = "technical"
        elif any(k in msg for k in ["return", "wrong item", "broken", "damaged"]):
            issue_type = "returns"
        elif any(k in msg for k in ["overheat", "fire", "safety", "defect", "recall", "hazard"]):
            issue_type = "safety"
        elif any(k in msg for k in ["cancel", "subscription", "renewal"]):
            issue_type = "cancellation"
        elif any(k in msg for k in ["angry", "complaint", "manager", "terrible"]):
            issue_type = "complaint"
        elif any(k in msg for k in ["order", "bulk", "corporate", "gift", "stock"]):
            issue_type = "orders"
        else:
            issue_type = "billing"
        return {"action_type": "classify", "response": issue_type}
    elif phase == 2:
        msg = state["message"].lower()
        if any(k in msg for k in ["ship", "deliver", "package", "track"]):
            team = "logistics_team"
        elif any(k in msg for k in ["login", "password", "bug", "api", "data", "error"]):
            team = "tech_support_team"
        elif any(k in msg for k in ["overheat", "fire", "safety", "defect"]):
            team = "safety_team"
        elif any(k in msg for k in ["invoice", "bill", "charge", "payment"]):
            team = "finance_team"
        elif any(k in msg for k in ["bulk", "corporate", "order", "return", "cancel"]):
            team = "orders_team"
        else:
            team = "management_team"
        return {"action_type": "assign", "team": team}
    else:
        if state.get("sentiment") == "angry":
            return {
                "action_type": "escalate",
                "team": "management_team",
                "response": "We sincerely apologize. I am escalating your issue to our management team immediately.",
            }
        return {
            "action_type": "respond",
            "response": "Thank you for contacting us. Our team will look into this and get back to you shortly.",
        }


async def call_api_model(client: AsyncOpenAI, user_prompt: str) -> dict:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt},
    ]

    response = await client.chat.completions.create(
        model=MODEL_NAME,
        messages=messages,
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
    )

    text = response.choices[0].message.content.strip()

    json_match = re.search(r"(\{.*\})", text, re.DOTALL)
    if json_match:
        clean_json = json_match.group(1)
        try:
            return json.loads(clean_json)
        except json.JSONDecodeError as decode_exc:
            raise ValueError(f"JSON decode failed. Response: {text}") from decode_exc
    else:
        raise ValueError(f"No JSON object detected in response. Raw response: {text}")


async def get_action(phase: int, state: dict, client: AsyncOpenAI, user_prompt: str) -> tuple[SupportAction, Optional[str]]:
    """Returns (action, error_message). error_message is None on success."""
    try:
        data = await call_api_model(client, user_prompt)
        return SupportAction(**data), None
    except Exception as exc:
        err = f"{type(exc).__name__}: {exc}"
        print(f"[DEBUG] ⚠️ API/parse error: {err}", flush=True)
        return SupportAction(**fallback_policy(phase, state)), err


async def run_episode(task_name: str, client: AsyncOpenAI) -> None:
    """Run one full 3-phase episode for the given task and emit START/STEP/END logs."""
    env = await SupportEnvWrapper.from_docker_image(IMAGE_NAME, task=task_name)

    history: List[str] = []
    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False

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

    try:
        result = await env.reset()
        obs = result.observation

        for step in range(1, MAX_STEPS + 1):
            if result.done:
                break

            phase = env.env.phase  # read actual phase from env, not a local counter
            user_prompt = build_user_prompt(step, obs, history, phase=phase)
            state_dict = {
                "message": obs.message,
                "sentiment": obs.sentiment,
            }
            action, step_error = await get_action(phase, state_dict, client, user_prompt)
            action_str = json.dumps(action.model_dump())

            result = await env.step(action)
            obs = result.observation
            reward = result.reward or 0.0
            done = result.done

            rewards.append(reward)
            steps_taken = step

            log_step(step=step, action=action_str, reward=reward, done=done, error=step_error)
            history.append(f"Step {step} (phase {phase}): {action_str} -> reward {reward:+.2f}")

            if done:
                break

        score = sum(rewards) / float(len(rewards)) if rewards else 0.0
        score = min(max(score, 0.002), 0.998)
        success = score >= SUCCESS_SCORE_THRESHOLD

    finally:
        try:
            await env.close()
        except Exception as exc:
            print(f"[DEBUG] env.close() error: {exc}", flush=True)
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)


async def main() -> None:
    if not API_BASE_URL:
        raise RuntimeError("[FATAL] API_BASE_URL is not set. The LLM proxy URL must be provided via environment variable.")
    if not API_KEY:
        raise RuntimeError("[FATAL] HF_TOKEN (or API_KEY) is not set. A valid API key must be provided via environment variable.")

    print(f"[INFO] Using API_BASE_URL={API_BASE_URL} MODEL_NAME={MODEL_NAME}", flush=True)
    client = AsyncOpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    # Always run all three tasks so the validator can enumerate tasks and
    # verify each grader produces scores in [0.002, 0.998].
    for task in ["easy", "medium", "hard"]:
        try:
            await run_episode(task, client)
        except Exception as exc:
            print(f"[DEBUG] ⚠️ Episode failed for task={task}: {type(exc).__name__}: {exc}", flush=True)
            log_end(success=False, steps=0, score=0.002, rewards=[0.002])
            continue


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