""" environment.py — Strict RL-style Customer Support Environment Handles: step enforcement, repeat detection, fail conditions, reward calculation """ from __future__ import annotations import re from typing import List, Tuple, Optional from models import ( Episode, EpisodeStatus, StepName, StepResult, DifficultyLevel, Task ) from graders.base_grader import BaseGrader # ── Step order ──────────────────────────────────────────────────────────────── STEP_ORDER = [ StepName.EMPATHY, StepName.COLLECT_INFO, StepName.INVESTIGATE, StepName.RESOLUTION, ] # ── Reward constants ─────────────────────────────────────────────────────────── BASE_SCORE_CORRECT = 1.0 BASE_SCORE_INCORRECT = 0.2 STEP_BONUS = 0.2 # bonus when step is correct WRONG_STEP_PENALTY = 0.3 # wrong action in correct step position REPEAT_PENALTY = 0.2 # repeated question / response SKIP_STEP_PENALTY = 0.3 # jumped ahead EARLY_SOLUTION_PENALTY = 0.25 # gave resolution before investigation EMOTION_IGNORE_PENALTY = 0.25 # angry customer → neutral / cold reply GENERIC_RESPONSE_PENALTY = 0.15 # "ok", "done", vague one-liners WRONG_ASSUMPTION_PENALTY = 0.2 # stated wrong facts LOOP_PENALTY = 0.35 # agent stuck in loop (same step repeated 2+) # ── Overlap threshold for repeat detection ───────────────────────────────────── REPEAT_WORD_OVERLAP_MIN = 10 # words in common → flagged as repeat class CustomerSupportEnv: """ Strict step-based RL environment for customer support training. """ def __init__(self, task: Task, grader: BaseGrader): self.task = task self.grader = grader self.episode = Episode( task_id = task.task_id, difficulty = task.difficulty, ) self._response_history: List[str] = [] self._step_index = 0 # which step we expect next self._consecutive_wrong = 0 # wrong attempts at current step # ── Public API ───────────────────────────────────────────────────────────── def step(self, agent_response: str) -> Tuple[StepResult, bool]: """ Process one agent response. Returns (StepResult, done:bool). """ if self.episode.status != EpisodeStatus.RUNNING: raise RuntimeError("Episode is already finished.") expected_step = STEP_ORDER[self._step_index] result = self._evaluate(agent_response, expected_step) self.episode.add_step(result) self._response_history.append(agent_response.lower().strip()) done = False if result.correct: self._step_index += 1 self._consecutive_wrong = 0 else: self._consecutive_wrong += 1 # Loop fail: 3 consecutive wrong attempts at the same step if self._consecutive_wrong >= 3 and self.episode.status == EpisodeStatus.RUNNING: self.episode.status = EpisodeStatus.FAIL self.episode.fail_reason = ( f"Agent stuck in loop at step '{expected_step.value}' " f"({self._consecutive_wrong} consecutive failures)" ) if self.episode.status != EpisodeStatus.RUNNING: done = True elif self._step_index >= len(STEP_ORDER): done = True # episode.add_step() already set SUCCESS/FAIL return result, done def reset(self) -> None: self.episode = Episode( task_id = self.task.task_id, difficulty = self.task.difficulty, ) self._response_history = [] self._step_index = 0 self._consecutive_wrong = 0 def summary(self): return self.episode.summary() # ── Internal evaluation ──────────────────────────────────────────────────── def _evaluate(self, response: str, expected_step: StepName) -> StepResult: penalties: List[str] = [] total_penalty = 0.0 # 1. Grade the response against expected step grader_result = self.grader.grade( response = response, expected_step = expected_step, task = self.task, ) correct = grader_result["correct"] base_score = BASE_SCORE_CORRECT if correct else BASE_SCORE_INCORRECT detected_action = grader_result.get("detected_action", "unknown") # 2. Step bonus step_bonus = STEP_BONUS if correct else 0.0 # 3. Wrong-step penalty if not correct: total_penalty += WRONG_STEP_PENALTY penalties.append( f"Wrong action detected ('{detected_action}' " f"≠ '{expected_step.value}'): -{WRONG_STEP_PENALTY}" ) # 4. Repeat detection if self._is_repeated_response(response): total_penalty += REPEAT_PENALTY penalties.append(f"Repeated/duplicate response: -{REPEAT_PENALTY}") # 5. Early solution penalty if self._is_early_solution(response, expected_step): total_penalty += EARLY_SOLUTION_PENALTY penalties.append(f"Solution given too early: -{EARLY_SOLUTION_PENALTY}") # 6. Emotion mismatch penalty if self._is_emotion_mismatch(response): total_penalty += EMOTION_IGNORE_PENALTY penalties.append( f"Angry customer ignored (no empathy/de-escalation): " f"-{EMOTION_IGNORE_PENALTY}" ) # 7. Generic / too-short response if self._is_generic_response(response): total_penalty += GENERIC_RESPONSE_PENALTY penalties.append(f"Generic/too-short response: -{GENERIC_RESPONSE_PENALTY}") # 8. Wrong assumption detection wrong_assumption = grader_result.get("wrong_assumption", False) if wrong_assumption: total_penalty += WRONG_ASSUMPTION_PENALTY penalties.append(f"Incorrect assumption stated: -{WRONG_ASSUMPTION_PENALTY}") # 9. Skip-step penalty (grader signals this) if grader_result.get("skipped_step", False): total_penalty += SKIP_STEP_PENALTY penalties.append(f"Step skipped: -{SKIP_STEP_PENALTY}") # ── Reward formula ───────────────────────────────────────────────────── reward = max(0.0, base_score + step_bonus - total_penalty) # ── Fail trigger (individual step) ──────────────────────────────────── fail_triggered = False fail_reason = "" if total_penalty >= 0.8: fail_triggered = True fail_reason = f"Single step penalty exceeded threshold ({total_penalty:.2f})" return StepResult( step = expected_step, agent_response = response, detected_action = detected_action, expected_action = expected_step.value, correct = correct, base_score = base_score, step_bonus = step_bonus, penalty = total_penalty, penalty_reasons = penalties, reward = reward, fail_triggered = fail_triggered, fail_reason = fail_reason, ) # ── Helper detectors ─────────────────────────────────────────────────────── def _is_repeated_response(self, response: str) -> bool: if not self._response_history: return False words_new = set(response.lower().split()) for prev in self._response_history: words_prev = set(prev.split()) overlap = len(words_new & words_prev) if overlap >= REPEAT_WORD_OVERLAP_MIN: return True return False def _is_early_solution(self, response: str, step: StepName) -> bool: """Penalise giving resolution keywords before the resolution step.""" if step in (StepName.EMPATHY, StepName.COLLECT_INFO): resolution_signals = [ "refund", "replacement", "we will fix", "we will credit", "here is the solution", "the fix is", "escalate your", ] r = response.lower() return any(sig in r for sig in resolution_signals) return False def _is_emotion_mismatch(self, response: str) -> bool: """Flag cold/neutral replies when customer is angry/frustrated.""" if self.task.customer_emotion not in ("angry", "frustrated"): return False empathy_signals = [ "sorry", "apologize", "apology", "understand your frustration", "i hear you", "i completely understand", "that must be", "deeply sorry", "sincerely apologize", ] r = response.lower() return not any(sig in r for sig in empathy_signals) def _is_generic_response(self, response: str) -> bool: stripped = response.strip().lower() # Very short replies if len(stripped.split()) <= 4: return True # Generic filler phrases generic_phrases = [ "ok", "okay", "done", "sure", "got it", "no problem", "understood", "alright", ] return stripped in generic_phrases