tasks/ medium_task.py

"""
MEDIUM TASK — Multi-Step Issue Resolution
Goal: Agent guides a customer through a multi-turn troubleshooting process.
"""

import random
from graders.medium_grader import MediumGrader


class MediumTask:
    """
    Task: Resolve a customer's technical/account issue over multiple turns.

    - 3–5 turn conversations.
    - Agent must ask clarifying questions, diagnose the problem, and resolve it.
    - Resolution requires following a correct sequence of steps.
    - Rewarded for each correct step, penalized for skipping or wrong order.
    """

    description = (
        "Help the customer resolve their issue step-by-step. "
        "Ask clarifying questions, diagnose the problem, and guide them to a resolution. "
        "Complete all required steps in a logical order."
    )
    max_steps = 6

    SCENARIOS = [
        {
            "id": "med_001",
            "topic": "billing_issue",
            "opening_message": "I was charged twice for my subscription this month. This is really frustrating!",
            "required_steps": [
                "acknowledge_frustration",   # "I understand, I'm sorry to hear that..."
                "collect_account_info",      # "Could I get your account email/ID?"
                "investigate",               # "Let me look into this for you..."
                "resolve_or_escalate",       # "I've initiated a refund..." or "escalating..."
            ],
            "resolution": "Issue a full refund for the duplicate charge and confirm via email.",
            "metadata": {
                "product": "Subscription Service",
                "account_id": "ACC-88321",
                "charge_amount": "$29.99",
                "charge_date": "March 28, 2026",
            },
            # Simulated customer replies for each agent step
            "customer_replies": {
                "acknowledge_frustration": "Thank you for understanding. Yes, I need this fixed.",
                "collect_account_info": "Sure, my email is user@example.com.",
                "investigate": "Okay, please check it. I can see two charges on my bank statement.",
                "resolve_or_escalate": "Great, thank you! When will the refund appear?",
            },
        },
        {
            "id": "med_002",
            "topic": "login_problem",
            "opening_message": "I can't log into my account. It keeps saying my password is wrong but I haven't changed it.",
            "required_steps": [
                "acknowledge_issue",
                "check_basic_steps",     # "Have you tried clearing your browser cache?"
                "offer_password_reset",  # "Let me send you a password reset link."
                "confirm_resolution",    # "Were you able to log in successfully?"
            ],
            "resolution": "Send password reset link and confirm the customer successfully logged in.",
            "metadata": {
                "product": "SaaS App",
                "account_status": "Active",
                "last_login": "3 days ago",
            },
            "customer_replies": {
                "acknowledge_issue": "Yes, I've been trying for an hour!",
                "check_basic_steps": "I tried that, it didn't work.",
                "offer_password_reset": "Okay, I got the email. Let me try.",
                "confirm_resolution": "Yes! I'm in now. Thank you so much!",
            },
        },
        {
            "id": "med_003",
            "topic": "wrong_item_shipped",
            "opening_message": "I ordered a blue jacket (size M) but received a red one in size L. What do I do?",
            "required_steps": [
                "apologize_sincerely",
                "collect_order_info",     # "Could I get your order number?"
                "confirm_item_details",   # "Let me confirm what you ordered..."
                "arrange_replacement",    # "I'll arrange a free return and send the correct item."
            ],
            "resolution": "Arrange free return and ship correct item with express delivery.",
            "metadata": {
                "product": "Online Clothing Store",
                "order_id": "ORD-44512",
                "ordered_item": "Blue Jacket Size M",
                "received_item": "Red Jacket Size L",
            },
            "customer_replies": {
                "apologize_sincerely": "Thank you, I hope you can fix this.",
                "collect_order_info": "My order number is ORD-44512.",
                "confirm_item_details": "Yes, that's correct, I ordered the blue one.",
                "arrange_replacement": "That's great, when will the correct one arrive?",
            },
        },
    ]

    def __init__(self):
        self._grader = MediumGrader()
        self._current_scenario = None
        self._steps_completed = []

    def reset(self) -> dict:
        self._steps_completed = []
        self._current_scenario = random.choice(self.SCENARIOS).copy()
        return self._current_scenario

    def step(self, action: str, scenario: dict, history: list, step_number: int) -> dict:
        score, grader_info, step_identified = self._grader.grade(
            action=action,
            scenario=scenario,
            history=history,
            steps_completed=self._steps_completed,
        )

        # Track completed steps
        if step_identified and step_identified not in self._steps_completed:
            self._steps_completed.append(step_identified)

        required = scenario["required_steps"]
        completed_count = len(self._steps_completed)
        total_required = len(required)

        # Reward formula:
        # Each correctly completed step earns proportional reward.
        # Final step earns a bonus. Missing steps or wrong answers penalize.
        if step_identified in required and step_identified not in grader_info.get("already_done", []):
            step_reward = (1.0 / total_required) * score    # partial step reward
        else:
            step_reward = -0.1 * (1 - score)                # small penalty for non-progress

        # Check if all required steps are done
        all_done = all(s in self._steps_completed for s in required)
        done = all_done or step_number >= self.max_steps

        if all_done:
            step_reward += 0.2   # bonus for completing all steps

        # Generate next customer message based on what step was just done
        next_msg = ""
        if not done:
            replies = scenario.get("customer_replies", {})
            next_msg = replies.get(step_identified, "Okay, what should I do next?")

        return {
            "reward": max(-1.0, min(1.0, step_reward)),
            "done": done,
            "next_customer_message": next_msg,
            "grader_info": {
                **grader_info,
                "score": score,
                "step_identified": step_identified,
                "steps_completed": list(self._steps_completed),
                "progress": f"{completed_count}/{total_required}",
            },
        }