tasks.py

"""Task source for the HR environment — bundled samples + optional API upgrade."""

from __future__ import annotations

import logging
import os
import random
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

COLLINEAR_PLATFORM_URL = "https://platform.collinear.ai"
SCENARIO_MANAGER_API_URL = "https://rl-gym-api.collinear.ai"


@dataclass
class Task:
    """A single HR task for the agent to complete."""

    id: str
    instruction: str
    difficulty: str
    rubric: list[str] = field(default_factory=list)


BUNDLED_TASKS: list[Task] = [
    Task(
        id="hr-001",
        instruction=(
            "A new candidate, Priya Mehta, has applied for the Senior Software Engineer role. "
            "Create her employee record in the HRMS, schedule a phone screening for next "
            "Tuesday at 2 PM on the calendar, and send her a confirmation email at "
            "priya.mehta@gmail.com with the interview details."
        ),
        difficulty="medium",
        rubric=[
            "Employee record created in HRMS for Priya Mehta",
            "Phone screening event scheduled on calendar for next Tuesday at 2 PM",
            "Confirmation email sent to priya.mehta@gmail.com with interview details",
        ],
    ),
    Task(
        id="hr-002",
        instruction=(
            "Employee James Wilson (EMP-0042) has requested 5 days of annual leave starting "
            "next Monday. Check his remaining leave balance in the HRMS, approve the request "
            "if he has sufficient days, and notify his manager Sarah Chen via Rocket.Chat."
        ),
        difficulty="easy",
        rubric=[
            "Leave balance checked for employee EMP-0042",
            "Leave request approved or denied based on balance",
            "Manager Sarah Chen notified via RocketChat",
        ],
    ),
    Task(
        id="hr-003",
        instruction=(
            "Run a monthly attendance report: pull the attendance records for all employees "
            "from the HRMS for the current month, identify anyone with more than 2 absences, "
            "and send a summary email to hr-team@company.com with the findings."
        ),
        difficulty="medium",
        rubric=[
            "Attendance records retrieved from HRMS",
            "Employees with >2 absences identified",
            "Summary email sent to hr-team@company.com",
        ],
    ),
    Task(
        id="hr-004",
        instruction=(
            "The recruiting team needs to schedule a panel interview for candidate Alex Rivera "
            "for the Product Manager position. Check the availability of three interviewers "
            "(Sarah Chen, Mike Johnson, Lisa Park) on the calendar for this week, find a "
            "1-hour slot that works for all three, book the meeting, and send calendar "
            "invites via email to all participants including the candidate at alex.rivera@email.com."
        ),
        difficulty="hard",
        rubric=[
            "Availability checked for all three interviewers on the calendar",
            "Common 1-hour slot identified",
            "Meeting booked on the calendar",
            "Email invites sent to all participants including alex.rivera@email.com",
        ],
    ),
    Task(
        id="hr-005",
        instruction=(
            "Employee Maria Santos has been promoted from Junior Developer to Senior Developer. "
            "Update her designation and salary grade in the HRMS, send her a congratulatory "
            "email, and post an announcement in the #general channel on Rocket.Chat."
        ),
        difficulty="easy",
        rubric=[
            "Designation updated in HRMS to Senior Developer",
            "Congratulatory email sent to Maria Santos",
            "Announcement posted in #general on RocketChat",
        ],
    ),
    Task(
        id="hr-006",
        instruction=(
            "A new employee, David Kim, is starting next Monday. Complete the onboarding "
            "checklist: create his employee record in the HRMS with department 'Engineering', "
            "send him a welcome email at david.kim@company.com with first-day instructions, "
            "schedule a 30-minute orientation meeting on his start date, and add him to the "
            "#engineering channel on Rocket.Chat."
        ),
        difficulty="hard",
        rubric=[
            "Employee record created in HRMS with department Engineering",
            "Welcome email sent to david.kim@company.com",
            "Orientation meeting scheduled on calendar for start date",
            "Added to #engineering channel on RocketChat",
        ],
    ),
    Task(
        id="hr-007",
        instruction=(
            "Check how many open leave requests are pending approval in the HRMS. "
            "List them all and send a reminder email to the respective approving managers "
            "asking them to review the pending requests."
        ),
        difficulty="medium",
        rubric=[
            "Pending leave requests retrieved from HRMS",
            "Approving managers identified for each request",
            "Reminder emails sent to respective managers",
        ],
    ),
    Task(
        id="hr-008",
        instruction=(
            "The quarterly performance review cycle is starting. Look up all employees "
            "in the Engineering department from the HRMS, schedule individual 45-minute "
            "review meetings with their manager for next week on the calendar, and "
            "send each employee an email notification about their scheduled review time."
        ),
        difficulty="hard",
        rubric=[
            "Engineering department employees retrieved from HRMS",
            "Individual 45-minute review meetings scheduled on calendar",
            "Email notifications sent to each employee with their review time",
        ],
    ),
]


def get_task(task_index: int | None = None) -> Task:
    """Return a task — from the API if COLLINEAR_API_KEY is set, else from bundled set."""
    api_key = os.environ.get("COLLINEAR_API_KEY", "").strip()

    if api_key:
        try:
            return _fetch_api_task(api_key, task_index)
        except Exception:
            logger.warning(
                "Failed to fetch task from Scenario Manager API, falling back to bundled tasks.",
                exc_info=True,
            )

    if task_index is not None:
        return BUNDLED_TASKS[task_index % len(BUNDLED_TASKS)]
    return random.choice(BUNDLED_TASKS)


def _fetch_api_task(api_key: str, task_index: int | None) -> Task:
    """Fetch a task from the Collinear Scenario Manager API."""
    import requests

    base_url = os.environ.get("SIMLAB_SCENARIO_MANAGER_API_URL", SCENARIO_MANAGER_API_URL).rstrip(
        "/"
    )
    headers = {"Accept": "application/json", "API-Key": api_key}

    resp = requests.get(f"{base_url}/v1/scenarios", headers=headers, timeout=30)
    resp.raise_for_status()
    scenarios = resp.json()

    hr_scenario = None
    for s in scenarios:
        name = (s.get("name") or "").lower().replace(" ", "-").replace("_", "-")
        if "hr" in name or "human-resource" in name or "recruiting" in name:
            hr_scenario = s
            break

    if hr_scenario is None:
        raise ValueError("No HR scenario found in Scenario Manager")

    scenario_id = hr_scenario["scenario_id"]
    resp = requests.get(
        f"{base_url}/v1/scenarios/{scenario_id}/tasks", headers=headers, timeout=30
    )
    resp.raise_for_status()
    data = resp.json()
    tasks = data.get("tasks", [])

    if not tasks:
        raise ValueError(f"No tasks found for scenario {scenario_id}")

    if task_index is not None:
        api_task = tasks[task_index % len(tasks)]
    else:
        api_task = random.choice(tasks)

    return Task(
        id=api_task.get("task_id", "api-unknown"),
        instruction=api_task.get("description", ""),
        difficulty=api_task.get("difficulty", "unknown"),
        rubric=[],
    )