inference.py

import json
import os
import sys
import textwrap
from pathlib import Path
from typing import List, Optional

from openai import OpenAI


ROOT = Path(__file__).resolve().parent
PACKAGE_PARENT = ROOT.parent
if str(PACKAGE_PARENT) not in sys.path:
    sys.path.insert(0, str(PACKAGE_PARENT))

from helpdesk_env.environment import HelpdeskEnv
from helpdesk_env.models import Action


LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME", "helpdesk-openenv")
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.groq.com/openai/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "llama-3.3-70b-versatile")
API_KEY = os.getenv("GROQ_API_KEY") or os.getenv("HF_TOKEN") or os.getenv("API_KEY")
TASK_NAME = os.getenv("TASK_NAME", "easy")
BENCHMARK = os.getenv("BENCHMARK", "helpdesk_env")
TEMPERATURE = float(os.getenv("TEMPERATURE", "0"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "180"))
SUCCESS_SCORE_THRESHOLD = float(os.getenv("SUCCESS_SCORE_THRESHOLD", "0.50"))

MAX_STEPS_BY_TASK = {
    "easy": 1,
    "medium": 3,
    "hard": 8,
}

SYSTEM_PROMPT_BASE = (
    "You are a banking customer support agent for a UPI payments app. "
    "Never ask for PIN, OTP, CVV, or full card details. "
    "You must return exactly one JSON object with keys from: "
    "action_type, category, faq_id, message. "
    "Valid action_type values are exactly: classify, lookup_faq, ask_clarification, "
    "reply, escalate, resolve_ticket."
)


def system_prompt_for_task(task_id: str) -> str:
    if task_id == "easy":
        return (
            SYSTEM_PROMPT_BASE
            + " For easy tasks, classify the issue into exactly one category from "
            "observation.available_categories."
        )
    if task_id == "medium":
        return (
            SYSTEM_PROMPT_BASE
            + " For medium tasks, choose lookup_faq with the best faq_id from "
            "observation.knowledge_base, or use escalate when fraud or overdue review requires manual handling."
        )
    return (
        SYSTEM_PROMPT_BASE
        + " For hard tasks, ask for clarification first, then retrieve the right FAQ, "
        "then reply with safe guidance, and only resolve after the customer confirms the issue is fixed."
    )


def build_user_prompt(task_id: str, observation_json: str, history: List[str]) -> str:
    history_block = "\n".join(history[-4:]) if history else "None"
    return textwrap.dedent(
        f"""
        Task: {task_id}
        Observation JSON:
        {observation_json}

        Recent action history:
        {history_block}

        Return the next action as one JSON object only.
        """
    ).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"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"done={str(done).lower()} error={error_val}",
        flush=True,
    )


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


def _extract_json_object(text: str) -> str:
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        if len(lines) >= 2 and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    return text


_VALID_ACTIONS = frozenset(
    {
        "classify",
        "lookup_faq",
        "ask_clarification",
        "reply",
        "escalate",
        "resolve_ticket",
    }
)


def _normalize_action_type(raw: object) -> str:
    if raw is None:
        return ""
    value = str(raw).strip().lower().replace("-", "_")
    return value if value in _VALID_ACTIONS else ""


def _fallback_action(task_id: str, turn_number: int) -> Action:
    if task_id == "easy":
        return Action(action_type="classify", category="payment_failure")
    if task_id == "medium":
        return Action(action_type="escalate", message="Escalating for manual review.")
    if turn_number == 0:
        return Action(
            action_type="ask_clarification",
            message="Please share the UTR, amount, and exact issue.",
        )
    if turn_number == 1:
        return Action(action_type="lookup_faq", faq_id="faq_001")
    if turn_number in (2, 3):
        return Action(
            action_type="reply",
            message="Please follow the safe steps in the app and confirm the result.",
        )
    return Action(action_type="resolve_ticket")


def parse_action(response_text: str, task_id: str, turn_number: int) -> Action:
    text = _extract_json_object(response_text)
    try:
        payload = json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        end = text.rfind("}")
        if start != -1 and end != -1 and end > start:
            try:
                payload = json.loads(text[start : end + 1])
            except json.JSONDecodeError:
                payload = {}
        else:
            payload = {}

    action_type = _normalize_action_type(payload.get("action_type"))
    if not action_type:
        return _fallback_action(task_id, turn_number)

    try:
        return Action(
            action_type=action_type,
            category=payload.get("category"),
            faq_id=payload.get("faq_id"),
            message=payload.get("message"),
        )
    except Exception:
        return _fallback_action(task_id, turn_number)


def get_model_action(
    client: OpenAI,
    task_id: str,
    observation_json: str,
    history: List[str],
    turn_number: int,
) -> Action:
    user_prompt = build_user_prompt(task_id, observation_json, history)
    completion = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": system_prompt_for_task(task_id)},
            {"role": "user", "content": user_prompt},
        ],
        temperature=TEMPERATURE,
        max_tokens=MAX_TOKENS,
        response_format={"type": "json_object"},
    )
    text = completion.choices[0].message.content or ""
    return parse_action(text, task_id, turn_number)


def main() -> None:
    if not API_KEY:
        raise RuntimeError(
            "Set GROQ_API_KEY, HF_TOKEN, or API_KEY before running inference.py"
        )

    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
    env = HelpdeskEnv()

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

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

    try:
        observation = env.reset(TASK_NAME)
        done = False

        for step in range(1, MAX_STEPS_BY_TASK.get(TASK_NAME, 3) + 1):
            if done:
                break

            error: Optional[str] = None
            try:
                action = get_model_action(
                    client=client,
                    task_id=TASK_NAME,
                    observation_json=observation.model_dump_json(),
                    history=history,
                    turn_number=observation.turn_number,
                )
                observation, reward, done, _info = env.step(action)
                reward_value = reward.value
            except Exception as exc:
                action = _fallback_action(TASK_NAME, observation.turn_number)
                reward_value = 0.0
                done = True
                error = str(exc)

            action_str = json.dumps(action.model_dump(exclude_none=True), separators=(",", ":"))
            log_step(
                step=step,
                action=action_str,
                reward=reward_value,
                done=done,
                error=error,
            )

            rewards.append(reward_value)
            steps_taken = step
            history.append(f"step={step} action={action_str} reward={reward_value:.2f}")

        final_score = rewards[-1] if rewards else 0.0
        success = final_score >= SUCCESS_SCORE_THRESHOLD

    finally:
        log_end(success=success, steps=steps_taken, rewards=rewards)


if __name__ == "__main__":
    main()