""" CustomerSupportEnv — OpenEnv-compliant customer support inbox environment. Implements the standard OpenEnv interface: - reset(task_id, ticket_id) → Observation - step(action) → (Observation, Reward, done, info) - state() → TicketState Three tasks: Task 1 (easy): Ticket Triage Task 2 (medium): Guided Resolution Task 3 (hard): VIP Retention & Escalation """ from __future__ import annotations import random import uuid from copy import deepcopy from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple from .data import ( TICKET_TEMPLATES, CUSTOMER_FOLLOW_UPS, KNOWLEDGE_BASE, get_ticket_by_id, get_random_ticket, compute_sla_deadline, time_to_sla_human, ) from .graders import run_grader from .models import ( Action, ActionType, CustomerProfile, Message, Observation, Reward, RewardBreakdown, TicketCategory, TicketPriority, TicketState, TicketStatus, ) from .rewards import compute_step_reward from .tasks import TASK_DEFINITIONS, get_task, list_tasks VALID_ACTIONS_BY_STATUS = { TicketStatus.OPEN: [ ActionType.CLASSIFY, ActionType.RESPOND, ActionType.TAG, ActionType.REQUEST_INFO, ActionType.ESCALATE, ActionType.ASSIGN, ], TicketStatus.IN_PROGRESS: [ ActionType.RESPOND, ActionType.REQUEST_INFO, ActionType.ESCALATE, ActionType.RESOLVE, ActionType.TAG, ActionType.ASSIGN, ], TicketStatus.WAITING_CUSTOMER: [ ActionType.RESPOND, ActionType.ESCALATE, ActionType.RESOLVE, ActionType.TAG, ], TicketStatus.ESCALATED: [ ActionType.RESPOND, ActionType.RESOLVE, ActionType.TAG, ], TicketStatus.RESOLVED: [], TicketStatus.CLOSED: [], } class CustomerSupportEnv: """ OpenEnv-compliant Customer Support Inbox environment. Episode flow: 1. reset(task_id, ticket_id) → agent receives Observation 2. Agent calls step(action) repeatedly 3. Episode ends when: ticket is Resolved/Closed, max_turns reached, or agent takes a RESOLVE action 4. Final grader scores the episode (0.0–1.0) """ def __init__(self, seed: Optional[int] = None): self._seed = seed if seed is not None: random.seed(seed) self._state: Optional[TicketState] = None self._ticket_data: Optional[Dict[str, Any]] = None self._task_config: Optional[Any] = None self._episode_history: List[Dict[str, Any]] = [] self._cumulative_reward: float = 0.0 self._step_count: int = 0 self._session_id: str = str(uuid.uuid4()) # ─── OpenEnv interface ──────────────────────────────────────────────── def reset( self, task_id: str = "task1", ticket_id: Optional[str] = None, seed: Optional[int] = None, ) -> Observation: """ Reset the environment for a new episode. Args: task_id: Which task to run ("task1", "task2", "task3") ticket_id: Specific ticket to use (None = random from task's pool) seed: Optional random seed for reproducibility Returns: Initial Observation for the agent """ if seed is not None: random.seed(seed) self._seed = seed # Load task self._task_config = get_task(task_id) # Pick ticket if ticket_id: self._ticket_data = get_ticket_by_id(ticket_id) else: pool = self._task_config.ticket_pool chosen_id = random.choice(pool) self._ticket_data = get_ticket_by_id(chosen_id) # Build initial state now = datetime.now(timezone.utc) customer_data = self._ticket_data["customer"] customer = CustomerProfile(**customer_data) self._state = TicketState( ticket_id=self._ticket_data["id"], subject=self._ticket_data["subject"], body=self._ticket_data["body"], category=None, priority=None, status=TicketStatus.OPEN, tags=[], assigned_to=None, conversation=[ Message( role="customer", content=self._ticket_data["body"], timestamp=now.isoformat(), metadata={"initial": True}, ) ], customer=customer, created_at=now.isoformat(), updated_at=now.isoformat(), sla_deadline=compute_sla_deadline(self._ticket_data["sla_hours"], now), resolution_notes=None, escalation_reason=None, turn_count=0, max_turns=self._task_config.max_turns, task_id=task_id, task_metadata={ "ticket_pool": self._task_config.ticket_pool, "difficulty": self._task_config.difficulty, }, ) self._episode_history = [] self._cumulative_reward = 0.0 self._step_count = 0 self._session_id = str(uuid.uuid4()) return self._build_observation(last_action_result=None, last_reward=0.0) def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]: """ Execute one action in the environment. Args: action: The Action to perform Returns: (observation, reward, done, info) """ if self._state is None: raise RuntimeError("Call reset() before step()") if self._state.status in (TicketStatus.RESOLVED, TicketStatus.CLOSED): # Episode already done obs = self._build_observation("Episode already ended.", 0.0) reward = Reward(score=0.0, cumulative_score=self._cumulative_reward, done=True, breakdown=RewardBreakdown(), feedback="Episode already ended.") return obs, reward, True, {"episode_done": True} # Validate action valid_types = VALID_ACTIONS_BY_STATUS.get(self._state.status, []) if action.action_type not in valid_types: # Invalid action — small penalty, return same observation obs = self._build_observation( f"Invalid action '{action.action_type}' for status '{self._state.status}'.", -0.05, ) reward = Reward( score=0.0, cumulative_score=self._cumulative_reward, done=False, breakdown=RewardBreakdown(step_penalty=-0.05), feedback=f"Invalid action: '{action.action_type}' not allowed when status='{self._state.status}'.", ) return obs, reward, False, {"error": "invalid_action"} prev_state = deepcopy(self._state) now = datetime.now(timezone.utc) # ── Apply action to state ────────────────────────────────────────── action_result = self._apply_action(action, now) # Increment counters self._state.turn_count += 1 self._state.updated_at = now.isoformat() self._step_count += 1 # Record in history history_entry: Dict[str, Any] = { "step": self._step_count, "timestamp": now.isoformat(), "action_type": action.action_type.value, "action_detail": action.model_dump(exclude_none=True), "state_status": self._state.status.value, } self._episode_history.append(history_entry) # ── Check episode termination ───────────────────────────────────── done = self._check_done() # ── Compute reward ──────────────────────────────────────────────── final_grader_score = None if done: final_grader_score, breakdown, grader_feedback = run_grader( task_id=self._state.task_id, ticket_data=self._ticket_data, ticket_state=self._state, episode_history=self._episode_history, task_config={ "grader_config": self._task_config.grader_config, "task_id": self._state.task_id, }, ) action_result += f" | Grader: {grader_feedback}" reward = compute_step_reward( action=action, prev_state=prev_state, new_state=self._state, ticket_data=self._ticket_data, episode_history=self._episode_history, done=done, final_score=final_grader_score, ) self._cumulative_reward += reward.score # Update cumulative in reward object reward = reward.model_copy(update={ "cumulative_score": round(min(self._cumulative_reward / max(self._step_count, 1), 1.0), 3), "task_complete": done and (final_grader_score or 0.0) >= self._task_config.min_score_to_pass, }) obs = self._build_observation(action_result, reward.score) info: Dict[str, Any] = { "episode_id": self._session_id, "step": self._step_count, "task_id": self._state.task_id, "ticket_id": self._state.ticket_id, "status": self._state.status.value, "turns_used": self._state.turn_count, "turns_max": self._state.max_turns, } if done and final_grader_score is not None: info["final_grader_score"] = final_grader_score info["passed"] = final_grader_score >= self._task_config.min_score_to_pass return obs, reward, done, info def state(self) -> TicketState: """Return the current full ticket state.""" if self._state is None: raise RuntimeError("Call reset() before state()") return deepcopy(self._state) # ─── Action application ─────────────────────────────────────────────── def _apply_action(self, action: Action, now: datetime) -> str: """Mutate self._state based on the action. Returns human-readable result.""" state = self._state if action.action_type == ActionType.CLASSIFY: if action.category: state.category = action.category if action.priority: state.priority = action.priority if state.status == TicketStatus.OPEN: state.status = TicketStatus.IN_PROGRESS parts = [] if action.category: parts.append(f"category={action.category.value}") if action.priority: parts.append(f"priority={action.priority.value}") return f"Ticket classified: {', '.join(parts)}." elif action.action_type == ActionType.RESPOND: text = (action.response_text or "").strip() if not text: return "RESPOND action requires response_text." state.conversation.append(Message( role="agent", content=text, timestamp=now.isoformat(), metadata={"action_type": "respond"}, )) if state.status == TicketStatus.OPEN: state.status = TicketStatus.IN_PROGRESS # Simulate customer follow-up if available self._maybe_add_customer_followup(now) return f"Response sent to customer ({len(text)} chars)." elif action.action_type == ActionType.REQUEST_INFO: text = (action.response_text or "").strip() if not text: return "REQUEST_INFO action requires response_text." state.conversation.append(Message( role="agent", content=text, timestamp=now.isoformat(), metadata={"action_type": "request_info"}, )) state.status = TicketStatus.WAITING_CUSTOMER # Simulate customer response self._maybe_add_customer_followup(now, force=True) return f"Information requested from customer. Status → WAITING_CUSTOMER." elif action.action_type == ActionType.ESCALATE: reason = action.escalation_reason or "Escalated by agent." team = action.escalation_team or "tier2" state.escalation_reason = reason state.assigned_to = team state.status = TicketStatus.ESCALATED state.conversation.append(Message( role="agent", content=f"[Escalated to {team}]: {reason}", timestamp=now.isoformat(), metadata={"action_type": "escalate", "team": team}, )) return f"Ticket escalated to {team}: {reason[:60]}..." elif action.action_type == ActionType.RESOLVE: notes = (action.resolution_notes or "").strip() if not notes: return "RESOLVE action requires resolution_notes." state.resolution_notes = notes state.status = TicketStatus.RESOLVED state.conversation.append(Message( role="agent", content=f"[Resolved]: {notes}", timestamp=now.isoformat(), metadata={ "action_type": "resolve", "resolution_category": action.resolution_category or "unknown", }, )) return f"Ticket resolved. Notes: {notes[:80]}..." elif action.action_type == ActionType.TAG: new_tags = action.tags or [] # Deduplicate for t in new_tags: if t not in state.tags: state.tags.append(t.lower().replace(" ", "-")) return f"Tags added: {new_tags}. Total tags: {state.tags}." elif action.action_type == ActionType.ASSIGN: state.assigned_to = action.assigned_to if state.status == TicketStatus.OPEN: state.status = TicketStatus.IN_PROGRESS return f"Ticket assigned to: {action.assigned_to}." return f"Action {action.action_type} processed." def _maybe_add_customer_followup(self, now: datetime, force: bool = False): """Optionally inject a simulated customer follow-up message.""" tid = self._state.ticket_id followups = CUSTOMER_FOLLOW_UPS.get(tid, []) if not followups: return # Count existing customer follow-ups (beyond initial message) existing_followups = sum( 1 for m in self._state.conversation if m.role == "customer" and not m.metadata.get("initial") ) if existing_followups < len(followups): if force or (random.random() < 0.6): msg_text = followups[existing_followups] self._state.conversation.append(Message( role="customer", content=msg_text, timestamp=now.isoformat(), metadata={"followup_index": existing_followups}, )) def _check_done(self) -> bool: """Check if the episode should end.""" if self._state is None: return True if self._state.status in (TicketStatus.RESOLVED, TicketStatus.CLOSED): return True if self._state.turn_count >= self._state.max_turns: return True return False # ─── Observation builder ────────────────────────────────────────────── def _build_observation( self, last_action_result: Optional[str], last_reward: float, ) -> Observation: """Construct the Observation object from current state.""" s = self._state task = self._task_config available = [ a.value for a in VALID_ACTIONS_BY_STATUS.get(s.status, []) ] time_to_sla = time_to_sla_human(s.sla_deadline) return Observation( ticket_id=s.ticket_id, subject=s.subject, body=s.body, status=s.status, category=s.category, priority=s.priority, tags=s.tags, customer=s.customer, conversation=s.conversation, turn_count=s.turn_count, turns_remaining=max(0, s.max_turns - s.turn_count), time_to_sla=time_to_sla, task_id=s.task_id, task_description=task.description, task_objectives=task.objectives, available_actions=available, last_action_result=last_action_result, last_reward=last_reward, episode_done=self._check_done(), info={ "difficulty": task.difficulty, "session_id": self._session_id, "step": self._step_count, "customer_tier": s.customer.account_tier, "knowledge_base": list(KNOWLEDGE_BASE.keys()), }, ) # ─── Helper methods ─────────────────────────────────────────────────── def list_tasks(self) -> List[Dict[str, Any]]: """Return all available task definitions.""" return list_tasks() def get_knowledge_article(self, key: str) -> str: """Retrieve a knowledge base article.""" return KNOWLEDGE_BASE.get(key, "Article not found.") def get_episode_summary(self) -> Dict[str, Any]: """Return a summary of the current/last episode.""" if self._state is None: return {} return { "session_id": self._session_id, "task_id": self._state.task_id, "ticket_id": self._state.ticket_id, "status": self._state.status.value, "turns_used": self._state.turn_count, "cumulative_reward": round(self._cumulative_reward, 3), "history_length": len(self._episode_history), "conversation_length": len(self._state.conversation), }