Spaces:
Running
Running
| from __future__ import annotations | |
| from typing import Any, Dict, List, Tuple | |
| from models import DataRecord, EraseOpsAction, EraseOpsObservation, EraseOpsReward, EraseRequest | |
| from server.tasks import get_task | |
| class EraseOpsEnvironment: | |
| """Stateful OpenEnv-compatible GDPR erasure environment.""" | |
| def __init__(self) -> None: | |
| self._records: Dict[str, DataRecord] = {} | |
| self._request: EraseRequest | None = None | |
| self._task_name: str | None = None | |
| self._observation_order: List[str] = [] | |
| self._history: List[Dict[str, Any]] = [] | |
| self._step: int = 0 | |
| self._step_budget: int = 0 | |
| self._cumulative_penalty: float = 0.0 | |
| self._last_action_result: str = "Environment not initialized" | |
| self._legal_hold_violations: int = 0 | |
| self._dependency_violations: int = 0 | |
| self._wrong_user_deletions: int = 0 | |
| self._contamination_violations: int = 0 | |
| self._requests_completed: int = 0 | |
| self._escalated_legal_holds: set[str] = set() | |
| self._blocked_since: Dict[str, int] = {} | |
| self._done: bool = False | |
| self._terminated_reason: str = "" | |
| def reset(self, task_name: str, seed: int = 42) -> EraseOpsObservation: | |
| task = get_task(task_name=task_name, seed=seed) | |
| self._task_name = task.task_name | |
| self._request = task.request.model_copy(deep=True) | |
| self._records = {record.record_id: record.model_copy(deep=True) for record in task.records} | |
| self._observation_order = list(task.observation_order) | |
| self._history = [] | |
| self._step = 0 | |
| self._step_budget = task.step_budget | |
| self._cumulative_penalty = 0.0 | |
| self._last_action_result = "Environment reset" | |
| self._legal_hold_violations = 0 | |
| self._dependency_violations = 0 | |
| self._wrong_user_deletions = 0 | |
| self._contamination_violations = 0 | |
| self._requests_completed = 0 | |
| self._escalated_legal_holds = set() | |
| self._blocked_since = {} | |
| self._done = False | |
| self._terminated_reason = "" | |
| return self._build_observation() | |
| def step(self, action: EraseOpsAction) -> Tuple[EraseOpsObservation, EraseOpsReward, bool, Dict[str, Any]]: | |
| if self._request is None: | |
| raise RuntimeError("Environment must be reset before calling step") | |
| if self._done: | |
| obs = self._build_observation() | |
| reward = EraseOpsReward( | |
| total=-1.0, | |
| deletion_reward=0.0, | |
| step_penalty=0.0, | |
| legal_hold_penalty=0.0, | |
| dependency_penalty=0.0, | |
| contamination_penalty=0.0, | |
| over_deletion_penalty=0.0, | |
| completion_bonus=0.0, | |
| ) | |
| info = { | |
| "error": "Episode already terminated. Call reset before taking new actions.", | |
| "terminated_reason": self._terminated_reason, | |
| } | |
| return obs, reward, True, info | |
| self._step += 1 | |
| deletion_reward = 0.0 | |
| step_penalty = -0.01 | |
| legal_hold_penalty = 0.0 | |
| dependency_penalty = 0.0 | |
| contamination_penalty = 0.0 | |
| over_deletion_penalty = 0.0 | |
| completion_bonus = 0.0 | |
| reward_total = step_penalty | |
| record = self._records.get(action.record_id) if action.record_id else None | |
| if action.action_type == "noop": | |
| self._last_action_result = "No action taken" | |
| elif action.action_type == "inspect": | |
| if record is None: | |
| reward_total -= 1.0 | |
| self._last_action_result = "ERROR: record not found" | |
| else: | |
| self._last_action_result = f"Inspected {action.record_id}: {record.model_dump()}" | |
| elif action.action_type == "delete": | |
| if record is None: | |
| reward_total -= 1.0 | |
| self._last_action_result = "ERROR: record not found" | |
| elif record.status != "active": | |
| reward_total -= 0.5 | |
| self._last_action_result = "ERROR: already deleted" | |
| elif record.legal_hold: | |
| legal_hold_penalty -= 5.0 | |
| reward_total += legal_hold_penalty | |
| self._legal_hold_violations += 1 | |
| self._cumulative_penalty += 5.0 | |
| self._last_action_result = "VIOLATION: legal hold record cannot be deleted" | |
| elif record.user_id != self._request.user_id: | |
| contamination_penalty -= 2.0 | |
| reward_total += contamination_penalty | |
| self._wrong_user_deletions += 1 | |
| self._contamination_violations += 1 | |
| self._cumulative_penalty += 2.0 | |
| self._last_action_result = "ERROR: wrong user - GDPR scope violation" | |
| else: | |
| unresolved_deps = self._unresolved_dependencies(record) | |
| if unresolved_deps: | |
| dependency_penalty -= 3.0 | |
| reward_total += dependency_penalty | |
| self._dependency_violations += 1 | |
| self._cumulative_penalty += 3.0 | |
| self._blocked_since.setdefault(record.record_id, self._step) | |
| self._last_action_result = f"VIOLATION: must delete dependencies first: {unresolved_deps}" | |
| elif record.data_type != "PII": | |
| over_deletion_penalty -= 2.0 | |
| reward_total += over_deletion_penalty | |
| record.status = "deleted" | |
| self._blocked_since.pop(record.record_id, None) | |
| self._last_action_result = "WARNING: over-deletion - this record is not PII" | |
| else: | |
| deletion_reward += 1.0 | |
| reward_total += deletion_reward | |
| record.status = "deleted" | |
| self._blocked_since.pop(record.record_id, None) | |
| self._requests_completed += 1 | |
| self._last_action_result = f"SUCCESS: deleted {record.record_id} from {record.system}" | |
| elif action.action_type == "anonymize": | |
| if record is None: | |
| reward_total -= 1.0 | |
| self._last_action_result = "ERROR: record not found" | |
| elif record.status != "active": | |
| reward_total -= 0.5 | |
| self._last_action_result = "ERROR: already deleted" | |
| elif record.legal_hold: | |
| legal_hold_penalty -= 5.0 | |
| reward_total += legal_hold_penalty | |
| self._legal_hold_violations += 1 | |
| self._cumulative_penalty += 5.0 | |
| self._last_action_result = "VIOLATION: legal hold record cannot be anonymized" | |
| elif record.user_id != self._request.user_id: | |
| contamination_penalty -= 2.0 | |
| reward_total += contamination_penalty | |
| self._wrong_user_deletions += 1 | |
| self._contamination_violations += 1 | |
| self._cumulative_penalty += 2.0 | |
| self._last_action_result = "ERROR: wrong user - GDPR scope violation" | |
| else: | |
| unresolved_deps = self._unresolved_dependencies(record) | |
| if unresolved_deps: | |
| dependency_penalty -= 3.0 | |
| reward_total += dependency_penalty | |
| self._dependency_violations += 1 | |
| self._cumulative_penalty += 3.0 | |
| self._blocked_since.setdefault(record.record_id, self._step) | |
| self._last_action_result = f"VIOLATION: must delete dependencies first: {unresolved_deps}" | |
| elif record.data_type == "PII": | |
| reward_total -= 1.0 | |
| self._last_action_result = "ERROR: anonymization is insufficient for PII erasure" | |
| elif record.data_type in {"derived", "anonymized"}: | |
| reward_total += 0.3 | |
| record.status = "anonymized" | |
| self._blocked_since.pop(record.record_id, None) | |
| self._last_action_result = f"SUCCESS: anonymized {record.record_id}" | |
| else: | |
| reward_total -= 1.0 | |
| self._last_action_result = "ERROR: invalid anonymization target" | |
| elif action.action_type == "escalate": | |
| if record is None: | |
| reward_total -= 1.0 | |
| self._last_action_result = "ERROR: record not found" | |
| elif record.legal_hold: | |
| if record.record_id in self._escalated_legal_holds: | |
| reward_total -= 0.1 | |
| self._last_action_result = f"ERROR: {record.record_id} already escalated" | |
| else: | |
| reward_total += 0.2 | |
| self._escalated_legal_holds.add(record.record_id) | |
| self._last_action_result = f"ESCALATED: {record.record_id} flagged for legal review" | |
| else: | |
| reward_total -= 0.5 | |
| self._last_action_result = "ERROR: escalation only valid for legal hold records" | |
| else: | |
| reward_total -= 1.0 | |
| self._last_action_result = f"ERROR: unsupported action type '{action.action_type}'" | |
| done = False | |
| pii_remaining = self._pii_remaining() | |
| if pii_remaining == 0: | |
| done = True | |
| self._terminated_reason = "all_pii_deleted" | |
| if self._step < self._step_budget: | |
| completion_bonus = 0.5 | |
| reward_total += completion_bonus | |
| elif self._step >= self._step_budget: | |
| done = True | |
| self._terminated_reason = "step_budget_exhausted" | |
| elif self._cumulative_penalty > 15.0: | |
| done = True | |
| self._terminated_reason = "catastrophic_penalty" | |
| self._done = done | |
| reward = EraseOpsReward( | |
| total=round(reward_total, 4), | |
| deletion_reward=round(deletion_reward, 4), | |
| step_penalty=round(step_penalty, 4), | |
| legal_hold_penalty=round(legal_hold_penalty, 4), | |
| dependency_penalty=round(dependency_penalty, 4), | |
| contamination_penalty=round(contamination_penalty, 4), | |
| over_deletion_penalty=round(over_deletion_penalty, 4), | |
| completion_bonus=round(completion_bonus, 4), | |
| ) | |
| obs = self._build_observation() | |
| info: Dict[str, Any] = { | |
| "task_name": self._task_name, | |
| "terminated_reason": self._terminated_reason if done else "", | |
| "wrong_user_deletions": self._wrong_user_deletions, | |
| "contamination_violations": self._contamination_violations, | |
| } | |
| self._history.append( | |
| { | |
| "step": self._step, | |
| "action": action.model_dump(), | |
| "reward": reward.model_dump(), | |
| "done": done, | |
| "info": info, | |
| "observation": { | |
| "step": obs.step, | |
| "pii_remaining": obs.pii_remaining, | |
| "legal_hold_violations": obs.legal_hold_violations, | |
| "dependency_violations": obs.dependency_violations, | |
| "contamination_violations": obs.contamination_violations, | |
| "requests_completed": obs.requests_completed, | |
| "cumulative_penalty": obs.cumulative_penalty, | |
| "last_action_result": obs.last_action_result, | |
| "deadline_pressure": obs.deadline_pressure, | |
| "escalation_queue_size": obs.escalation_queue_size, | |
| "blocked_dependency_count": obs.blocked_dependency_count, | |
| }, | |
| } | |
| ) | |
| return obs, reward, done, info | |
| def state(self) -> Dict[str, Any]: | |
| return { | |
| "task_name": self._task_name, | |
| "request": self._request.model_dump() if self._request else None, | |
| "step": self._step, | |
| "step_budget": self._step_budget, | |
| "done": self._done, | |
| "terminated_reason": self._terminated_reason, | |
| "counters": { | |
| "cumulative_penalty": self._cumulative_penalty, | |
| "legal_hold_violations": self._legal_hold_violations, | |
| "dependency_violations": self._dependency_violations, | |
| "wrong_user_deletions": self._wrong_user_deletions, | |
| "contamination_violations": self._contamination_violations, | |
| "requests_completed": self._requests_completed, | |
| "pii_remaining": self._pii_remaining(), | |
| }, | |
| "last_action_result": self._last_action_result, | |
| "history_length": len(self._history), | |
| "escalated_legal_holds": sorted(self._escalated_legal_holds), | |
| "records": [ | |
| self._records[record_id].model_dump() | |
| for record_id in sorted(self._records.keys()) | |
| ], | |
| "observation_order": list(self._observation_order), | |
| } | |
| def metrics(self, include_trace: bool = False) -> Dict[str, Any]: | |
| current_state = self.state() | |
| payload: Dict[str, Any] = { | |
| "task_name": self._task_name, | |
| "step": self._step, | |
| "step_budget": self._step_budget, | |
| "done": self._done, | |
| "terminated_reason": self._terminated_reason, | |
| "pii_remaining": self._pii_remaining(), | |
| "legal_hold_violations": self._legal_hold_violations, | |
| "dependency_violations": self._dependency_violations, | |
| "wrong_user_deletions": self._wrong_user_deletions, | |
| "contamination_violations": self._contamination_violations, | |
| "requests_completed": self._requests_completed, | |
| "cumulative_penalty": self._cumulative_penalty, | |
| "last_action_result": self._last_action_result, | |
| "deadline_pressure": self._deadline_pressure(), | |
| "escalation_queue_size": self._escalation_queue_stats()[0], | |
| "escalation_queue_max_age": self._escalation_queue_stats()[1], | |
| "blocked_dependency_count": self._blocked_dependency_stats()[0], | |
| "blocked_dependency_max_age": self._blocked_dependency_stats()[1], | |
| "state": current_state, | |
| } | |
| if include_trace: | |
| payload["trace"] = list(self._history) | |
| return payload | |
| def replay(self, limit: int | None = None) -> Dict[str, Any]: | |
| trace = self._history if limit is None else self._history[-max(limit, 0) :] | |
| return { | |
| "task_name": self._task_name, | |
| "step_budget": self._step_budget, | |
| "done": self._done, | |
| "terminated_reason": self._terminated_reason, | |
| "trace_length": len(self._history), | |
| "trace": list(trace), | |
| "final_state": self.state(), | |
| } | |
| def visualization(self) -> str: | |
| status_counts = {"active": 0, "deleted": 0, "anonymized": 0} | |
| for record in self._records.values(): | |
| status_counts[record.status] = status_counts.get(record.status, 0) + 1 | |
| lines = [ | |
| f"GDPR EraseOps | task={self._task_name} | step={self._step}/{self._step_budget}", | |
| f"pii_remaining={self._pii_remaining()} | legal_hold_violations={self._legal_hold_violations} | dependency_violations={self._dependency_violations}", | |
| f"requests_completed={self._requests_completed} | cumulative_penalty={self._cumulative_penalty:.2f} | done={self._done}", | |
| f"record_statuses={status_counts}", | |
| f"last_action_result={self._last_action_result}", | |
| ] | |
| return "\n".join(lines) | |
| def _build_observation(self) -> EraseOpsObservation: | |
| if self._request is None: | |
| raise RuntimeError("Environment is not initialized") | |
| ordered_records: List[DataRecord] = [] | |
| seen: set[str] = set() | |
| for record_id in self._observation_order: | |
| if record_id in self._records: | |
| ordered_records.append(self._records[record_id].model_copy(deep=True)) | |
| seen.add(record_id) | |
| for record_id, record in self._records.items(): | |
| if record_id not in seen: | |
| ordered_records.append(record.model_copy(deep=True)) | |
| escalation_queue_size, escalation_queue_max_age = self._escalation_queue_stats() | |
| blocked_dependency_count, blocked_dependency_max_age = self._blocked_dependency_stats() | |
| return EraseOpsObservation( | |
| records=ordered_records, | |
| request=self._request.model_copy(deep=True), | |
| step=self._step, | |
| step_budget=self._step_budget, | |
| cumulative_penalty=round(self._cumulative_penalty, 4), | |
| last_action_result=self._last_action_result, | |
| pii_remaining=self._pii_remaining(), | |
| legal_hold_violations=self._legal_hold_violations, | |
| dependency_violations=self._dependency_violations, | |
| contamination_violations=self._contamination_violations, | |
| requests_completed=self._requests_completed, | |
| deadline_pressure=self._deadline_pressure(), | |
| escalation_queue_size=escalation_queue_size, | |
| escalation_queue_max_age=escalation_queue_max_age, | |
| blocked_dependency_count=blocked_dependency_count, | |
| blocked_dependency_max_age=blocked_dependency_max_age, | |
| ) | |
| def _unresolved_dependencies(self, record: DataRecord) -> List[str]: | |
| unresolved: List[str] = [] | |
| for dep_id in record.dependencies: | |
| dep = self._records.get(dep_id) | |
| if dep is None: | |
| unresolved.append(dep_id) | |
| continue | |
| if dep.status == "deleted": | |
| continue | |
| # Escalated legal-hold dependencies are treated as policy-resolved. | |
| if dep.legal_hold and dep_id in self._escalated_legal_holds: | |
| continue | |
| unresolved.append(dep_id) | |
| return unresolved | |
| def _deadline_pressure(self) -> float: | |
| if self._request is None or self._request.deadline_steps <= 0: | |
| return 0.0 | |
| return round(min(1.0, self._step / float(self._request.deadline_steps)), 4) | |
| def _escalation_queue_stats(self) -> Tuple[int, int]: | |
| if self._request is None: | |
| return 0, 0 | |
| pending = [ | |
| record | |
| for record in self._records.values() | |
| if record.user_id == self._request.user_id | |
| and record.legal_hold | |
| and record.status == "active" | |
| and record.record_id not in self._escalated_legal_holds | |
| ] | |
| if not pending: | |
| return 0, 0 | |
| age = max(0, self._step - int(self._request.submitted_at_step or 0)) | |
| return len(pending), age | |
| def _blocked_dependency_stats(self) -> Tuple[int, int]: | |
| if self._request is None: | |
| return 0, 0 | |
| blocked_count = 0 | |
| max_age = 0 | |
| for record in self._records.values(): | |
| if record.user_id != self._request.user_id or record.data_type != "PII" or record.status != "active": | |
| continue | |
| unresolved = self._unresolved_dependencies(record) | |
| if not unresolved: | |
| self._blocked_since.pop(record.record_id, None) | |
| continue | |
| blocked_count += 1 | |
| first_seen = self._blocked_since.setdefault(record.record_id, self._step) | |
| max_age = max(max_age, max(0, self._step - first_seen)) | |
| return blocked_count, max_age | |
| def _pii_remaining(self) -> int: | |
| if self._request is None: | |
| return 0 | |
| return sum( | |
| 1 | |
| for record in self._records.values() | |
| if record.user_id == self._request.user_id | |
| and record.data_type == "PII" | |
| and record.status == "active" | |
| ) | |