openenv_bug_triage/env.py

"""
Main BugTriageEnv implementation following OpenEnv specification.
"""

from typing import Optional
from .models import (
    ObservationModel,
    ActionModel,
    RewardModel,
    StateModel,
    CurrentTicketModel,
    QueueStatsModel,
    TicketStateModel,
)
from .tasks import load_task, TaskDefinition
from .reward import RewardCalculator
from .grader import BugTriageGrader


DEFAULT_SEED = 42


class BugTriageEnv:
    """
    Bug Triage OpenEnv Environment.
    
    Simulates real-world software bug triage workflows.
    """
    
    def __init__(self):
        self.current_task: Optional[TaskDefinition] = None
        self.current_ticket_index: int = 0
        self.steps_used: int = 0
        self.episode_done: bool = False
        self.reward_calculator = RewardCalculator()
        self.cumulative_reward: float = 0.0
        self.ticket_states: list[dict] = []
        self.last_action_result: Optional[str] = None
        
        # Metrics for grading
        self.metrics = {}
        self._critical_ticket_ids: set[str] = set()
        self._duplicate_ticket_ids: set[str] = set()
        self._info_needed_ticket_ids: set[str] = set()
        self._severity_scored_tickets: set[str] = set()
        self._critical_severity_scored_tickets: set[str] = set()
        self._priority_scored_tickets: set[str] = set()
        self._component_scored_tickets: set[str] = set()
        self._team_scored_tickets: set[str] = set()
        self._duplicate_scored_tickets: set[str] = set()
        self._info_request_scored_tickets: set[str] = set()
        self._escalated_ticket_ids: set[str] = set()
        self._escalation_scored_tickets: set[str] = set()
        self._sla_met_ticket_ids: set[str] = set()
        self._missed_critical_ticket_ids: set[str] = set()
    
    def reset(
        self,
        task_id: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> ObservationModel:
        """
        Reset environment to start a new episode.
        
        Args:
            task_id: One of 'bug_triage_easy', 'bug_triage_medium', 'bug_triage_hard'
                     If None, defaults to 'bug_triage_easy'
            seed: Random seed for deterministic shuffling (default: 42)
        
        Returns:
            Initial observation
        """
        # Default values
        if task_id is None:
            task_id = "bug_triage_easy"
        if seed is None:
            seed = DEFAULT_SEED
        
        # Load task
        self.current_task = load_task(task_id, seed=seed)
        
        # Reset state
        self.current_ticket_index = 0
        self.steps_used = 0
        self.episode_done = False
        self.cumulative_reward = 0.0
        self.last_action_result = None
        
        # Initialize ticket states
        self.ticket_states = [
            {
                "ticket_id": ticket.ticket_id,
                "triaged": False,
                "actions_taken": [],
            }
            for ticket in self.current_task.tickets
        ]
        
        # Reset reward calculator
        self.reward_calculator.reset()
        
        self._critical_ticket_ids = {
            gt.ticket_id for gt in self.current_task.ground_truths
            if gt.true_severity in ["sev0", "sev1"]
        }
        self._duplicate_ticket_ids = {
            gt.ticket_id for gt in self.current_task.ground_truths
            if gt.duplicate_of is not None
        }
        self._info_needed_ticket_ids = {
            gt.ticket_id for gt in self.current_task.ground_truths
            if gt.needs_more_info
        }
        self._severity_scored_tickets = set()
        self._critical_severity_scored_tickets = set()
        self._priority_scored_tickets = set()
        self._component_scored_tickets = set()
        self._team_scored_tickets = set()
        self._duplicate_scored_tickets = set()
        self._info_request_scored_tickets = set()
        self._escalated_ticket_ids = set()
        self._escalation_scored_tickets = set()
        self._sla_met_ticket_ids = set()
        self._missed_critical_ticket_ids = set()

        label_total = sum(
            1 for gt in self.current_task.ground_truths
            if gt.duplicate_of is None
        )

        # Initialize metrics
        self.metrics = {
            "severity_correct": 0,
            "priority_correct": 0,
            "component_correct": 0,
            "team_correct": 0,
            "duplicate_correct": 0,
            "duplicate_total": 0,
            "duplicate_expected_total": len(self._duplicate_ticket_ids),
            "info_request_correct": 0,
            "info_needed_total": len(self._info_needed_ticket_ids),
            "major_mistakes": 0,
            "incorrect_close_count": 0,
            "critical_severity_correct": 0,
            "critical_severity_total": len(self._critical_ticket_ids),
            "sla_met": 0,
            "sla_total": len(self._critical_ticket_ids),
            "escalation_correct": 0,
            "escalation_total": len(self._critical_ticket_ids),
            "destructive_actions": 0,
            "missed_critical_escalation": 0,
            "steps_used": 0,
            "step_budget": self.current_task.step_budget,
            "label_total": label_total,
            "assignment_total": label_total,
        }
        
        return self._get_observation()
    
    def step(
        self,
        action: ActionModel,
    ) -> tuple[ObservationModel, RewardModel, bool, dict]:
        """
        Execute one step with the given action.
        
        Args:
            action: Action to take
        
        Returns:
            tuple of (observation, reward, done, info)
        """
        if self.episode_done:
            raise RuntimeError("Episode is done. Call reset() to start new episode.")
        
        if self.current_task is None:
            raise RuntimeError("No task loaded. Call reset() first.")
        
        # Validate action
        is_valid, validation_error = self._validate_action(action)
        
        # Get current ticket
        current_ticket = self.current_task.tickets[self.current_ticket_index]
        ground_truth = self.current_task.get_ground_truth(current_ticket.ticket_id)
        
        # Calculate reward
        is_critical = ground_truth.true_severity in ["sev0", "sev1"]
        step_reward, breakdown = self.reward_calculator.calculate_step_reward(
            action=action,
            ground_truth=ground_truth,
            ticket_id=current_ticket.ticket_id,
            is_valid=is_valid,
            is_critical_ticket=is_critical,
        )
        
        self.cumulative_reward += step_reward
        self.steps_used += 1
        self.metrics["steps_used"] = self.steps_used
        
        # Update metrics based on action
        if is_valid:
            self._update_metrics(action, ground_truth, current_ticket.ticket_id)
        else:
            self.metrics["major_mistakes"] += 1
        
        # Update ticket state
        self.ticket_states[self.current_ticket_index]["actions_taken"].append(
            action.action_type
        )
        
        # Handle action result
        if not is_valid:
            self.last_action_result = f"Invalid action: {validation_error}"
        else:
            self.last_action_result = self._execute_action(action, current_ticket)
        
        # Check if episode is done
        self._check_done()
        
        # Apply terminal bonus/penalty if done
        if self.episode_done:
            all_critical_triaged = self._all_critical_triaged()
            budget_exhausted = self.steps_used >= self.current_task.step_budget
            critical_remaining = self._has_critical_remaining()
            
            if all_critical_triaged:
                terminal_reward, _ = self.reward_calculator.calculate_step_reward(
                    action=action,
                    ground_truth=ground_truth,
                    ticket_id=current_ticket.ticket_id,
                    is_valid=True,
                    is_critical_ticket=False,
                    is_terminal=True,
                    all_critical_triaged=True,
                    budget_exhausted_with_critical=False,
                )
                self.cumulative_reward += terminal_reward
            elif budget_exhausted and critical_remaining:
                terminal_reward, _ = self.reward_calculator.calculate_step_reward(
                    action=action,
                    ground_truth=ground_truth,
                    ticket_id=current_ticket.ticket_id,
                    is_valid=True,
                    is_critical_ticket=False,
                    is_terminal=True,
                    all_critical_triaged=False,
                    budget_exhausted_with_critical=True,
                )
                self.cumulative_reward += terminal_reward
        
        # Create reward model
        reward_model = RewardModel(
            step_reward=step_reward,
            cumulative_reward=self.cumulative_reward,
            reward_breakdown=breakdown,
        )
        
        # Create observation
        observation = self._get_observation()
        
        # Info dict
        info = {
            "validation_error": validation_error if not is_valid else None,
            "metrics": self.metrics.copy(),
        }
        
        return observation, reward_model, self.episode_done, info
    
    def state(self) -> StateModel:
        """
        Get current environment state.
        
        Returns:
            Full state of the environment
        """
        if self.current_task is None:
            raise RuntimeError("No task loaded. Call reset() first.")
        
        ticket_states = [
            TicketStateModel(
                ticket_id=ts["ticket_id"],
                triaged=ts["triaged"],
                actions_taken=ts["actions_taken"],
            )
            for ts in self.ticket_states
        ]
        
        return StateModel(
            current_task_id=self.current_task.task_id,
            current_ticket_index=self.current_ticket_index,
            total_tickets=len(self.current_task.tickets),
            tickets_state=ticket_states,
            steps_used=self.steps_used,
            steps_remaining=max(0, self.current_task.step_budget - self.steps_used),
            cumulative_reward=self.cumulative_reward,
            episode_done=self.episode_done,
        )
    
    # Private helper methods
    
    def _get_observation(self) -> ObservationModel:
        """Create observation from current state."""
        if self.current_task is None or self.episode_done:
            # Return empty observation
            return ObservationModel(
                current_ticket=None,
                queue_stats=QueueStatsModel(
                    remaining_count=0,
                    urgent_count=0,
                    sla_at_risk_count=0,
                ),
                last_action_result=self.last_action_result,
                available_teams=[],
                available_components=[],
                steps_used=self.steps_used,
                steps_remaining=0,
                partial_score=self._calculate_partial_score(),
            )
        
        # Get current ticket
        current_ticket = self.current_task.tickets[self.current_ticket_index]
        
        # Calculate queue stats
        remaining_count = len(self.current_task.tickets) - self.current_ticket_index
        urgent_count = sum(
            1 for i in range(self.current_ticket_index, len(self.current_task.tickets))
            if self.current_task.get_ground_truth(
                self.current_task.tickets[i].ticket_id
            ).true_severity in ["sev0", "sev1"]
        )
        sla_at_risk_count = urgent_count  # Simplified: all urgent are at SLA risk
        
        return ObservationModel(
            current_ticket=CurrentTicketModel(**current_ticket.model_dump()),
            queue_stats=QueueStatsModel(
                remaining_count=remaining_count,
                urgent_count=urgent_count,
                sla_at_risk_count=sla_at_risk_count,
            ),
            last_action_result=self.last_action_result,
            available_teams=self.current_task.available_teams,
            available_components=self.current_task.available_components,
            steps_used=self.steps_used,
            steps_remaining=max(0, self.current_task.step_budget - self.steps_used),
            partial_score=self._calculate_partial_score(),
        )
    
    def _validate_action(self, action: ActionModel) -> tuple[bool, Optional[str]]:
        """Validate action is legal."""
        # Check if ticket exists for duplicate marking
        if action.action_type == "mark_duplicate" and action.mark_duplicate:
            canonical_id = action.mark_duplicate.canonical_ticket_id
            ticket_ids = [t.ticket_id for t in self.current_task.tickets]
            if canonical_id not in ticket_ids:
                return False, f"Unknown ticket ID: {canonical_id}"
        
        # Check if team is valid
        if action.action_type == "assign" and action.assign:
            if action.assign.team not in self.current_task.available_teams:
                return False, f"Unknown team: {action.assign.team}"
        
        # Check if component is valid
        if action.action_type == "classify" and action.classify:
            if action.classify.component not in self.current_task.available_components:
                return False, f"Unknown component: {action.classify.component}"
        
        return True, None
    
    def _execute_action(self, action: ActionModel, current_ticket) -> str:
        """Execute action and return result message."""
        if action.action_type == "next_ticket":
            if self.current_ticket_index < len(self.current_task.tickets) - 1:
                self.current_ticket_index += 1
                return "Moved to next ticket"
            else:
                # Move pointer past the end so _check_done can terminate the episode.
                self.current_ticket_index = len(self.current_task.tickets)
                return "No more tickets in queue"
        
        elif action.action_type == "classify":
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return f"Classified as {action.classify.severity}/{action.classify.priority}"
        
        elif action.action_type == "assign":
            return f"Assigned to team: {action.assign.team}"
        
        elif action.action_type == "mark_duplicate":
            return f"Marked as duplicate of {action.mark_duplicate.canonical_ticket_id}"
        
        elif action.action_type == "request_info":
            return f"Requested {action.request_info.info_type}"
        
        elif action.action_type == "defer":
            return f"Deferred: {action.defer.reason}"
        
        elif action.action_type == "close":
            return f"Closed: {action.close.reason}"
        
        elif action.action_type == "escalate_incident":
            return f"Escalated incident"
        
        return "Action executed"
    
    def _update_metrics(self, action: ActionModel, ground_truth, ticket_id: str):
        """Update metrics for grading."""
        is_duplicate_ticket = ground_truth.duplicate_of is not None
        is_critical_ticket = ground_truth.true_severity in ["sev0", "sev1"]

        if is_critical_ticket and action.action_type in ["classify", "mark_duplicate", "escalate_incident"]:
            if ticket_id not in self._sla_met_ticket_ids:
                self._sla_met_ticket_ids.add(ticket_id)
                self.metrics["sla_met"] += 1

        if action.action_type == "classify" and action.classify and not is_duplicate_ticket:
            if action.classify.severity == ground_truth.true_severity:
                if ticket_id not in self._severity_scored_tickets:
                    self._severity_scored_tickets.add(ticket_id)
                    self.metrics["severity_correct"] += 1

                if is_critical_ticket and ticket_id not in self._critical_severity_scored_tickets:
                    self._critical_severity_scored_tickets.add(ticket_id)
                    self.metrics["critical_severity_correct"] += 1

            if action.classify.priority == ground_truth.true_priority:
                if ticket_id not in self._priority_scored_tickets:
                    self._priority_scored_tickets.add(ticket_id)
                    self.metrics["priority_correct"] += 1

            if action.classify.component == ground_truth.true_component:
                if ticket_id not in self._component_scored_tickets:
                    self._component_scored_tickets.add(ticket_id)
                    self.metrics["component_correct"] += 1

        if action.action_type == "assign" and action.assign and not is_duplicate_ticket:
            if action.assign.team == ground_truth.true_assignee_team:
                if ticket_id not in self._team_scored_tickets:
                    self._team_scored_tickets.add(ticket_id)
                    self.metrics["team_correct"] += 1

        if action.action_type == "mark_duplicate" and action.mark_duplicate:
            self.metrics["duplicate_total"] += 1
            if (
                ground_truth.duplicate_of == action.mark_duplicate.canonical_ticket_id
                and ticket_id not in self._duplicate_scored_tickets
            ):
                self._duplicate_scored_tickets.add(ticket_id)
                self.metrics["duplicate_correct"] += 1

        if action.action_type == "request_info" and ground_truth.needs_more_info:
            if ticket_id not in self._info_request_scored_tickets:
                self._info_request_scored_tickets.add(ticket_id)
                self.metrics["info_request_correct"] += 1

        if action.action_type in ["close", "defer"]:
            if not ground_truth.duplicate_of and ground_truth.needs_more_info:
                self.metrics["incorrect_close_count"] += 1
                self.metrics["destructive_actions"] += 1

        if action.action_type == "escalate_incident":
            if is_critical_ticket:
                self._escalated_ticket_ids.add(ticket_id)
                if ticket_id not in self._escalation_scored_tickets:
                    self._escalation_scored_tickets.add(ticket_id)
                    self.metrics["escalation_correct"] += 1

        if (
            action.action_type == "next_ticket"
            and is_critical_ticket
            and ticket_id not in self._escalated_ticket_ids
            and ticket_id not in self._missed_critical_ticket_ids
        ):
            self._missed_critical_ticket_ids.add(ticket_id)
            self.metrics["missed_critical_escalation"] += 1
    
    def _check_done(self):
        """Check if episode should end."""
        # Episode ends if:
        # 1. Step budget exhausted
        # 2. All tickets processed (moved past last ticket)
        
        if self.steps_used >= self.current_task.step_budget:
            self.episode_done = True
        
        # If moved past last ticket
        if self.current_ticket_index >= len(self.current_task.tickets):
            self.episode_done = True
    
    def _all_critical_triaged(self) -> bool:
        """Check if all critical tickets have been triaged."""
        for i, ticket in enumerate(self.current_task.tickets):
            gt = self.current_task.get_ground_truth(ticket.ticket_id)
            if gt.true_severity in ["sev0", "sev1"]:
                if not self.ticket_states[i]["triaged"]:
                    return False
        return True
    
    def _has_critical_remaining(self) -> bool:
        """Check if there are untriaged critical tickets."""
        for i in range(self.current_ticket_index, len(self.current_task.tickets)):
            ticket = self.current_task.tickets[i]
            gt = self.current_task.get_ground_truth(ticket.ticket_id)
            if gt.true_severity in ["sev0", "sev1"]:
                if not self.ticket_states[i]["triaged"]:
                    return True
        return False
    
    def _calculate_partial_score(self) -> float:
        """Calculate partial score for intermediate feedback."""
        if self.steps_used == 0:
            return 0.0
        
        total_tickets = len(self.current_task.tickets) if self.current_task else 1
        
        # Simple weighted average of correctness
        score = 0.0
        score += self.metrics.get("severity_correct", 0) / total_tickets * 0.25
        score += self.metrics.get("priority_correct", 0) / total_tickets * 0.20
        score += self.metrics.get("component_correct", 0) / total_tickets * 0.25
        score += self.metrics.get("team_correct", 0) / total_tickets * 0.15
        
        return min(1.0, score)


if __name__ == "__main__":
    # Quick test
    env = BugTriageEnv()
    obs = env.reset(task_id="bug_triage_easy", seed=42)
    print(f"Reset complete. Current ticket: {obs.current_ticket.ticket_id if obs.current_ticket else 'None'}")
    print(f"Queue stats: {obs.queue_stats}")
    state = env.state()
    print(f"State: {state.current_task_id}, {state.total_tickets} tickets")