inference.py

from __future__ import annotations

import argparse
import json
import os
from typing import Dict, List

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

from models import SupportAction
from server.customer_support_environment import CustomerSupportEnvironment
from tasks import grade_task


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: str | None = None) -> None:
    error_value = "null" if error is None else error.replace("\n", " ")
    print(
        f"[STEP] step={step} action={action!r} reward={reward:.4f} done={done} error={error_value}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    reward_text = ",".join(f"{r:.4f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards=[{reward_text}]",
        flush=True,
    )


def get_model_plan(client: OpenAI, model_name: str, task_prompt: str) -> Dict[str, str | bool]:
    completion = client.chat.completions.create(
        model=model_name,
        temperature=0,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a customer support policy agent. Reply with strict JSON only and no markdown. "
                    "Schema: {\"category\": string, \"search_kb\": bool, \"escalate\": bool, \"response\": string}."
                ),
            },
            {"role": "user", "content": task_prompt},
        ],
        response_format={"type": "json_object"},
    )
    raw = completion.choices[0].message.content or "{}"
    data = json.loads(raw)
    return {
        "category": str(data.get("category", "technical_issue")),
        "search_kb": bool(data.get("search_kb", True)),
        "escalate": bool(data.get("escalate", False)),
        "response": str(data.get("response", "Thank you for contacting support. We will assist shortly.")),
    }


def heuristic_plan(observation: object) -> Dict[str, str | bool]:
    query = str(getattr(observation, "query")).lower()
    requires_escalation = bool(getattr(observation, "requires_escalation"))

    category = "technical_issue"
    if any(t in query for t in ("refund", "charge", "invoice", "billing")):
        category = "billing"
    elif any(t in query for t in ("cancel", "termination", "close account", "cancellation")):
        category = "cancellation"
    elif any(t in query for t in ("warranty", "buy", "purchase", "availability", "product")):
        category = "product_inquiry"
    elif any(t in query for t in ("refund not received", "return", "reimburse")):
        category = "refund"

    return {
        "category": category,
        "search_kb": True,
        "escalate": requires_escalation,
        "response": "Thank you for contacting us. We will review your case and assist you shortly.",
    }


def build_actions(observation: object, plan: Dict[str, str | bool]) -> List[SupportAction]:
    actions = [
        SupportAction(action_type="classify", content=str(plan["category"])),
    ]

    if bool(plan.get("search_kb", True)):
        actions.append(SupportAction(action_type="search_kb", content=str(getattr(observation, "kb_id"))))

    if bool(plan.get("escalate", False)):
        actions.append(SupportAction(action_type="escalate", content="Escalating to specialist support for manual review."))
    else:
        actions.append(SupportAction(action_type="respond", content=str(plan.get("response", "Thank you for contacting support."))))
    return actions


def main() -> int:
    parser = argparse.ArgumentParser(description="OpenEnv baseline inference for customer support benchmark")
    parser.add_argument("--csv", default="dataset.csv")
    parser.add_argument("--limit-per-task", type=int, default=5)
    parser.add_argument("--max-steps", type=int, default=4)
    parser.add_argument("--offline", action="store_true", help="Use heuristic policy without API calls")
    args = parser.parse_args()

    api_base_url = os.getenv("API_BASE_URL")
    model_name = os.getenv("MODEL_NAME")
    hf_token = os.getenv("HF_TOKEN")

    if not args.offline:
        missing = [name for name, value in (("API_BASE_URL", api_base_url), ("MODEL_NAME", model_name), ("HF_TOKEN", hf_token)) if not value]
        if missing:
            raise RuntimeError(f"Missing required env vars: {', '.join(missing)}")

    client = OpenAI(base_url=api_base_url, api_key=hf_token) if not args.offline else None
    env = CustomerSupportEnvironment(csv_path=args.csv)

    all_rewards: List[float] = []
    task_scores: List[float] = []
    task_labels: List[str] = []
    total_steps = 0

    log_start(task="all", env="customer_support_benchmark", model=model_name or "offline-heuristic")
    difficulties = ["easy", "medium", "hard"]
    for difficulty in difficulties:
        for idx in range(args.limit_per_task):
            obs = env.reset(difficulty=difficulty, index=idx)
            rewards: List[float] = []
            done = False

            prompt = (
                "Return the best support plan for this ticket. "
                "Prefer policy-safe responses and escalate only when needed.\n"
                f"TicketID={obs.ticket_id}\n"
                f"Difficulty={obs.difficulty}\n"
                f"Task={obs.task_id}\n"
                f"Query={obs.query}\n"
                f"KB={obs.kb_id}\n"
                f"RequiresEscalation={obs.requires_escalation}\n"
            )

            error = None
            try:
                plan = heuristic_plan(obs)
                if client is not None and model_name is not None:
                    plan = get_model_plan(client, model_name, prompt)
            except Exception as exc:
                error = str(exc)
                plan = heuristic_plan(obs)

            planned_actions = build_actions(obs, plan)

            for step, action in enumerate(planned_actions[: args.max_steps], start=1):
                if done:
                    break
                obs = env.step(action)

                reward = float(obs.reward or 0.0)
                rewards.append(reward)
                all_rewards.append(reward)
                total_steps += 1
                done = bool(obs.done)
                step_error = error if step == 1 else None
                log_step(step=step, action=f"{action.action_type}|{action.content}", reward=reward, done=done, error=step_error)

            final_score = grade_task(obs.task_id, obs.metadata, obs.history)
            task_scores.append(final_score)
            task_labels.append(difficulty)
            log_step(
                step=total_steps,
                action=f"episode_score|difficulty={difficulty}|index={idx}",
                reward=final_score,
                done=True,
                error=None,
            )

    benchmark_score = sum(task_scores) / len(task_scores) if task_scores else 0.0
    by_task: dict[str, list[float]] = {"easy": [], "medium": [], "hard": []}
    for label, score in zip(task_labels, task_scores):
        by_task[label].append(score)
    for label in ("easy", "medium", "hard"):
        values = by_task[label]
        avg = sum(values) / len(values) if values else 0.0
        log_step(
            step=total_steps,
            action=f"task_average|difficulty={label}",
            reward=avg,
            done=False,
            error=None,
        )

    print(
        f"[END] success={str(benchmark_score >= 0.7).lower()} steps={total_steps} score={benchmark_score:.4f} rewards_count={len(all_rewards)}",
        flush=True,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())