""" Typed models for the Customer Support RL Environment. Supports both single-agent (Round 1) and hierarchical multi-agent (Round 2) modes. Backward compatible: single-agent actions still work by defaulting role to 'support_agent'. """ from pydantic import BaseModel, Field, model_validator from typing import Optional, List, Literal, Dict, Any from enum import Enum # ── Agent Roles ──────────────────────────────────────────────────────────────── class AgentRole(str, Enum): SUPPORT_AGENT = "support_agent" SUPERVISOR = "supervisor" MANAGER = "manager" class SupervisorDecision(str, Enum): APPROVE = "approve" # L1 action is good, send to customer REJECT = "reject" # L1 action is bad, redo FEEDBACK = "feedback" # L1 action needs adjustment, provide guidance ESCALATE_TO_MANAGER = "escalate_to_manager" # Too complex for L2, send to L3 class ManagerDecision(str, Enum): OVERRIDE = "override" # Manager takes over and responds directly APPROVE_ESCALATION = "approve_escalation" # Approve L1's escalation request RESOLVE_DIRECTLY = "resolve_directly" # Manager resolves the issue SEND_BACK = "send_back" # Send back to L1 with directive # ── Action Types ─────────────────────────────────────────────────────────────── class ActionType(str, Enum): # L1 Support Agent actions RESPOND = "respond" ESCALATE = "escalate" CLOSE = "close" REQUEST_INFO = "request_info" # L1 DB query actions (internal, no customer reply) QUERY_USER_PROFILE = "query_user_profile" QUERY_ORDER_DETAILS = "query_order_details" # L2 Supervisor actions SUPERVISOR_APPROVE = "supervisor_approve" SUPERVISOR_REJECT = "supervisor_reject" SUPERVISOR_FEEDBACK = "supervisor_feedback" SUPERVISOR_ESCALATE = "supervisor_escalate" # L3 Manager actions MANAGER_OVERRIDE = "manager_override" MANAGER_RESOLVE = "manager_resolve" MANAGER_SEND_BACK = "manager_send_back" # Map action types to their originating role ACTION_ROLE_MAP: Dict[ActionType, AgentRole] = { ActionType.RESPOND: AgentRole.SUPPORT_AGENT, ActionType.ESCALATE: AgentRole.SUPPORT_AGENT, ActionType.CLOSE: AgentRole.SUPPORT_AGENT, ActionType.REQUEST_INFO: AgentRole.SUPPORT_AGENT, ActionType.QUERY_USER_PROFILE: AgentRole.SUPPORT_AGENT, ActionType.QUERY_ORDER_DETAILS: AgentRole.SUPPORT_AGENT, ActionType.SUPERVISOR_APPROVE: AgentRole.SUPERVISOR, ActionType.SUPERVISOR_REJECT: AgentRole.SUPERVISOR, ActionType.SUPERVISOR_FEEDBACK: AgentRole.SUPERVISOR, ActionType.SUPERVISOR_ESCALATE: AgentRole.SUPERVISOR, ActionType.MANAGER_OVERRIDE: AgentRole.MANAGER, ActionType.MANAGER_RESOLVE: AgentRole.MANAGER, ActionType.MANAGER_SEND_BACK: AgentRole.MANAGER, } # L1 action types (for backward compatibility checks) L1_ACTION_TYPES = { ActionType.RESPOND, ActionType.ESCALATE, ActionType.CLOSE, ActionType.REQUEST_INFO, ActionType.QUERY_USER_PROFILE, ActionType.QUERY_ORDER_DETAILS, } # DB query action types (internal lookups, no customer reply) DB_QUERY_ACTION_TYPES = { ActionType.QUERY_USER_PROFILE, ActionType.QUERY_ORDER_DETAILS, } L2_ACTION_TYPES = { ActionType.SUPERVISOR_APPROVE, ActionType.SUPERVISOR_REJECT, ActionType.SUPERVISOR_FEEDBACK, ActionType.SUPERVISOR_ESCALATE, } L3_ACTION_TYPES = { ActionType.MANAGER_OVERRIDE, ActionType.MANAGER_RESOLVE, ActionType.MANAGER_SEND_BACK, } # ── Core Models ──────────────────────────────────────────────────────────────── class Action(BaseModel): action_type: ActionType message: Optional[str] = Field(default=None, max_length=2000) reason: Optional[str] = Field(default=None, max_length=500) # DB query fields (L1 internal actions) email: Optional[str] = Field(default=None, max_length=254, description="Email for query_user_profile") order_id: Optional[str] = Field(default=None, max_length=100, description="Order ID for query_order_details") # Hierarchy fields (default for backward compat) role: AgentRole = Field(default=AgentRole.SUPPORT_AGENT) internal_note: Optional[str] = Field( default=None, max_length=1000, description="Internal reasoning, not shown to customer" ) feedback_to_agent: Optional[str] = Field( default=None, max_length=1000, description="Supervisor/Manager feedback to the lower-level agent" ) model_config = {"use_enum_values": True} @model_validator(mode="after") def validate_content(self): at = ActionType(self.action_type) # L1 validations if at == ActionType.RESPOND: if not self.message or not self.message.strip(): raise ValueError("message cannot be empty when action_type is 'respond'") if at == ActionType.ESCALATE: if not self.reason or not self.reason.strip(): raise ValueError("reason cannot be empty when action_type is 'escalate'") # DB query validations if at == ActionType.QUERY_USER_PROFILE: if not self.email or not self.email.strip(): raise ValueError("email required for query_user_profile") if at == ActionType.QUERY_ORDER_DETAILS: if not self.order_id or not self.order_id.strip(): raise ValueError("order_id required for query_order_details") # L2 validations if at == ActionType.SUPERVISOR_FEEDBACK: if not self.feedback_to_agent or not self.feedback_to_agent.strip(): raise ValueError("feedback_to_agent required for supervisor_feedback") if at == ActionType.SUPERVISOR_ESCALATE: if not self.reason or not self.reason.strip(): raise ValueError("reason required for supervisor_escalate") # L3 validations if at in (ActionType.MANAGER_OVERRIDE, ActionType.MANAGER_RESOLVE): if not self.message or not self.message.strip(): raise ValueError("message required for manager_override/manager_resolve") if at == ActionType.MANAGER_SEND_BACK: if not self.feedback_to_agent or not self.feedback_to_agent.strip(): raise ValueError("feedback_to_agent required for manager_send_back") # Auto-set role from action_type if not explicitly provided expected_role = ACTION_ROLE_MAP.get(at) if expected_role and self.role != expected_role: self.role = expected_role return self class Message(BaseModel): role: Literal["customer", "agent", "supervisor", "manager", "system"] content: str class HierarchyState(BaseModel): """Internal state of the agent hierarchy within an episode.""" support_agent_actions: int = 0 supervisor_reviews: int = 0 manager_interventions: int = 0 current_phase: str = Field( default="support_handling", description="Current phase: support_handling, supervisor_review, manager_override" ) escalation_reason: Optional[str] = None supervisor_feedback_history: List[str] = Field(default_factory=list) manager_directive_history: List[str] = Field(default_factory=list) pending_l1_action: Optional[Dict[str, Any]] = Field( default=None, description="The L1 action waiting for supervisor review" ) class Observation(BaseModel): session_id: str ticket_id: str category: str priority: Literal["low", "medium", "high", "critical"] subject: str conversation_history: List[Message] step: int max_steps: int customer_sentiment: float = Field(ge=-1.0, le=1.0) mood_trajectory: List[float] = Field(default_factory=list) unresolved_issues: List[str] is_done: bool task: str # DB query results (accumulated across episode) retrieved_data: Dict[str, Dict[str, Any]] = Field( default_factory=lambda: {"users": {}, "orders": {}}, description="Data fetched via query_user_profile / query_order_details" ) # ── Hierarchy fields (Round 2) ───────────────────────────────────────────── active_role: str = Field( default="support_agent", description="Which agent should act next: support_agent, supervisor, manager" ) supervisor_feedback: Optional[str] = Field( default=None, description="Latest feedback from supervisor to support agent" ) manager_directive: Optional[str] = Field( default=None, description="Latest directive from manager" ) hierarchy_state: Optional[HierarchyState] = None environment_event: Optional[str] = Field( default=None, description="Schema drift / policy change event" ) policy_context: str = Field( default="Standard operating procedure. Follow all default policies.", description="Current active policy text" ) escalation_chain: List[str] = Field( default_factory=list, description="History of escalations in this episode" ) class Reward(BaseModel): value: float = Field(ge=0.0, le=1.0) resolution_score: float tone_score: float efficiency_score: float accuracy_score: float breakdown: Dict[str, Any] # ── Hierarchy reward fields (Round 2) ────────────────────────────────────── empathy_score: float = Field(default=0.0, description="LLM-judged empathy") oversight_score: float = Field(default=0.0, description="Supervisor oversight quality") decision_quality_score: float = Field(default=0.0, description="Manager decision quality") policy_adherence_score: float = Field(default=0.0, description="LLM-judged policy adherence") role_rewards: Dict[str, float] = Field( default_factory=dict, description="Per-role reward breakdown: {support_agent: x, supervisor: y, manager: z}" ) @model_validator(mode="after") def _round_floats(self): self.value = round(self.value, 4) self.resolution_score = round(self.resolution_score, 4) self.tone_score = round(self.tone_score, 4) self.efficiency_score = round(self.efficiency_score, 4) self.accuracy_score = round(self.accuracy_score, 4) self.empathy_score = round(self.empathy_score, 4) self.oversight_score = round(self.oversight_score, 4) self.decision_quality_score = round(self.decision_quality_score, 4) self.policy_adherence_score = round(self.policy_adherence_score, 4) self.role_rewards = {k: round(v, 4) for k, v in self.role_rewards.items()} return self class Ticket(BaseModel): id: str category: str priority: Literal["low", "medium", "high", "critical"] subject: str opening_message: str follow_up_info: str required_info_before_close: List[str] expected_resolution_type: str ideal_max_steps: int customer_persona: Literal["impatient", "polite", "confused"] task: Literal["easy", "medium", "hard", "nightmare", "hierarchy_easy", "hierarchy_medium", "hierarchy_hard", "curriculum_basic", "curriculum_supervisor", "curriculum_full_hierarchy", "curriculum_nightmare", "multi_domain"] # DB linkage fields (optional; legacy tickets leave these empty) customer_email: Optional[str] = None related_order_ids: List[str] = Field(default_factory=list)