environment/tasks.py

"""
Task definitions for the Customer Support Inbox environment.

Task 1 (Easy):   Ticket Triage — classify category + priority correctly
Task 2 (Medium): Guided Resolution — respond to customer, gather info, resolve within SLA
Task 3 (Hard):   Multi-Turn VIP Retention — handle complex angry VIP through full lifecycle
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class TaskDefinition:
    """Specification for a single task."""
    task_id: str
    name: str
    difficulty: str                    # "easy" | "medium" | "hard"
    description: str
    objectives: List[str]
    ticket_pool: List[str]            # Ticket IDs to use for this task
    max_turns: int
    min_score_to_pass: float          # Score threshold for "success"
    reward_weights: Dict[str, float]  # How to weight reward components
    grader_config: Dict[str, Any]     # Grader-specific settings


TASK_DEFINITIONS: Dict[str, TaskDefinition] = {

    # ────────────────────────────────────────────────────────────────────────
    # TASK 1 — EASY: Ticket Triage
    # Agent must classify tickets with correct category and priority.
    # Simple, deterministic grader. Tests basic understanding of support domain.
    # ────────────────────────────────────────────────────────────────────────
    "task1": TaskDefinition(
        task_id="task1",
        name="Ticket Triage",
        difficulty="easy",
        description=(
            "You are a customer support triage agent. Your job is to read incoming "
            "support tickets and accurately classify them by category and priority level. "
            "A correct classification helps route tickets to the right team and ensures "
            "urgent issues are handled first."
        ),
        objectives=[
            "Assign the correct ticket category (billing/technical/shipping/returns/account/complaint/general)",
            "Assign the correct priority level (low/medium/high/urgent)",
            "Apply at least 2 relevant tags to aid searchability",
        ],
        ticket_pool=["T001", "T004", "T006", "T007", "T010"],  # Easy/clear-cut tickets
        max_turns=5,
        min_score_to_pass=0.7,
        reward_weights={
            "classification_accuracy": 0.60,
            "tags": 0.20,
            "efficiency": 0.20,
        },
        grader_config={
            "category_weight": 0.50,
            "priority_weight": 0.30,
            "tag_weight": 0.20,
            "require_both": True,   # Must classify both category AND priority
            "partial_priority": True,  # Adjacent priority (e.g. medium vs high) = 0.5
        },
    ),

    # ────────────────────────────────────────────────────────────────────────
    # TASK 2 — MEDIUM: Guided Resolution
    # Agent must classify the ticket, send a helpful response, potentially ask
    # for info, and ultimately resolve the issue with appropriate notes.
    # Tests multi-step reasoning and response quality.
    # ────────────────────────────────────────────────────────────────────────
    "task2": TaskDefinition(
        task_id="task2",
        name="Guided Resolution",
        difficulty="medium",
        description=(
            "You are a customer support agent handling incoming tickets end-to-end. "
            "For each ticket you must: (1) classify it correctly, (2) send a helpful, "
            "empathetic response to the customer, (3) gather any needed information, "
            "and (4) resolve the ticket with clear resolution notes. "
            "You must resolve within the SLA deadline."
        ),
        objectives=[
            "Correctly classify the ticket (category + priority)",
            "Respond to the customer with empathy and clear next steps",
            "Request any necessary information to complete resolution",
            "Resolve the ticket within the SLA window with complete resolution notes",
            "Achieve resolution that matches the ticket's appropriate outcome",
        ],
        ticket_pool=["T001", "T003", "T005", "T007", "T008"],  # Medium complexity
        max_turns=8,
        min_score_to_pass=0.65,
        reward_weights={
            "classification_accuracy": 0.20,
            "response_quality": 0.30,
            "resolution_completeness": 0.30,
            "sla_compliance": 0.10,
            "efficiency": 0.10,
        },
        grader_config={
            "require_classification": True,
            "require_response": True,
            "require_resolution": True,
            "response_min_length": 50,
            "resolution_min_length": 30,
            "sla_breach_penalty": 0.3,
            "check_empathy_keywords": True,
            "check_resolution_type": True,
        },
    ),

    # ────────────────────────────────────────────────────────────────────────
    # TASK 3 — HARD: VIP Retention & Escalation
    # Agent handles an angry enterprise customer threatening to cancel.
    # Must: classify urgently, respond with executive empathy, properly
    # escalate to the right team, follow up, and retain the customer.
    # Tests strategic decision-making, escalation judgment, and retention.
    # ────────────────────────────────────────────────────────────────────────
    "task3": TaskDefinition(
        task_id="task3",
        name="VIP Retention & Escalation",
        difficulty="hard",
        description=(
            "You are a senior customer support specialist. An enterprise customer with a "
            "$4,800/year subscription is threatening to cancel due to repeated poor experiences. "
            "You must handle this delicate situation by: accurately triaging urgency, "
            "responding with executive-level empathy, escalating to the right specialized team, "
            "demonstrating ownership, and achieving a resolution that retains the customer. "
            "Every action matters — this is a high-stakes retention scenario."
        ),
        objectives=[
            "Identify ticket as 'complaint' with 'urgent' priority immediately",
            "Acknowledge the customer's frustration with genuine empathy",
            "Escalate to 'tier2' team with detailed escalation notes",
            "Commit to a specific callback time or follow-up action",
            "Resolve ticket with complete notes including retention outcome",
            "Achieve positive or neutral customer sentiment by resolution",
        ],
        ticket_pool=["T009"],  # Only the VIP complaint
        max_turns=10,
        min_score_to_pass=0.60,
        reward_weights={
            "classification_accuracy": 0.15,
            "response_quality": 0.25,
            "escalation_appropriateness": 0.25,
            "resolution_completeness": 0.20,
            "customer_sentiment_improvement": 0.15,
        },
        grader_config={
            "require_urgent_priority": True,
            "require_complaint_category": True,
            "require_escalation": True,
            "escalation_team": "tier2",
            "require_empathy_phrases": True,
            "require_callback_commitment": True,
            "require_resolution": True,
            "sentiment_improvement_bonus": 0.2,
            "check_tone_professional": True,
            "escalation_without_classify_penalty": 0.3,
        },
    ),
}


def get_task(task_id: str) -> TaskDefinition:
    """Retrieve task definition by ID."""
    if task_id not in TASK_DEFINITIONS:
        raise ValueError(f"Unknown task_id: {task_id}. Valid: {list(TASK_DEFINITIONS.keys())}")
    return TASK_DEFINITIONS[task_id]


def list_tasks() -> List[Dict[str, Any]]:
    """Return all task summaries."""
    return [
        {
            "task_id": t.task_id,
            "name": t.name,
            "difficulty": t.difficulty,
            "description": t.description,
            "objectives": t.objectives,
            "max_turns": t.max_turns,
            "min_score_to_pass": t.min_score_to_pass,
        }
        for t in TASK_DEFINITIONS.values()
    ]