Spaces:
Running
Running
| from __future__ import annotations | |
| import random | |
| from typing import Dict, List | |
| from models import DataRecord, EraseRequest, TaskConfig | |
| def _record( | |
| record_id: str, | |
| system: str, | |
| data_type: str, | |
| user_id: str, | |
| dependencies: List[str], | |
| legal_hold: bool, | |
| age_days: int, | |
| deletion_cost: float, | |
| ) -> DataRecord: | |
| return DataRecord( | |
| record_id=record_id, | |
| system=system, | |
| data_type=data_type, | |
| user_id=user_id, | |
| dependencies=dependencies, | |
| legal_hold=legal_hold, | |
| age_days=age_days, | |
| deletion_cost=deletion_cost, | |
| status="active", | |
| ) | |
| def _clean_sweep() -> Dict[str, object]: | |
| records = [ | |
| _record("rec_001", "postgres", "PII", "user_001", [], False, 120, 0.1), | |
| _record("rec_002", "s3", "PII", "user_001", [], False, 340, 0.1), | |
| _record("rec_003", "slack", "PII", "user_001", [], False, 89, 0.1), | |
| _record("rec_004", "email", "PII", "user_001", [], False, 200, 0.1), | |
| _record("rec_005", "postgres", "PII", "user_001", [], False, 450, 0.1), | |
| ] | |
| request = EraseRequest( | |
| request_id="gdpr_req_clean_001", | |
| user_id="user_001", | |
| deadline_steps=10, | |
| regulation="GDPR", | |
| submitted_at_step=0, | |
| ) | |
| return { | |
| "task_name": "clean_sweep", | |
| "step_budget": 10, | |
| "request": request, | |
| "records": records, | |
| } | |
| def _tangled_lake() -> Dict[str, object]: | |
| records = [ | |
| _record("rec_001", "postgres", "PII", "user_042", [], False, 120, 0.2), | |
| _record("rec_002", "s3", "PII", "user_042", [], False, 98, 0.2), | |
| _record("rec_003", "slack", "PII", "user_042", [], False, 76, 0.2), | |
| _record("rec_004", "email", "PII", "user_042", ["rec_001"], False, 145, 0.25), | |
| _record("rec_005", "postgres", "PII", "user_042", [], False, 311, 0.22), | |
| _record("rec_006", "s3", "PII", "user_042", [], False, 222, 0.18), | |
| _record("rec_007", "slack", "PII", "user_042", ["rec_004"], False, 84, 0.19), | |
| _record("rec_008", "email", "PII", "user_042", ["rec_002", "rec_003"], False, 51, 0.27), | |
| _record("rec_009", "postgres", "anonymized", "user_042", [], False, 400, 0.05), | |
| _record("rec_010", "s3", "anonymized", "user_042", [], False, 250, 0.05), | |
| _record("rec_011", "slack", "anonymized", "user_042", [], False, 160, 0.05), | |
| _record("rec_012", "email", "anonymized", "user_042", [], False, 365, 0.05), | |
| _record("rec_013", "postgres", "derived", "user_042", [], False, 45, 0.03), | |
| _record("rec_014", "s3", "derived", "user_042", [], False, 90, 0.03), | |
| _record("rec_015", "email", "derived", "user_042", [], False, 180, 0.03), | |
| ] | |
| request = EraseRequest( | |
| request_id="gdpr_req_tangled_042", | |
| user_id="user_042", | |
| deadline_steps=20, | |
| regulation="GDPR", | |
| submitted_at_step=0, | |
| ) | |
| return { | |
| "task_name": "tangled_lake", | |
| "step_budget": 20, | |
| "request": request, | |
| "records": records, | |
| } | |
| def _legal_minefield() -> Dict[str, object]: | |
| records = [ | |
| _record("rec_001", "postgres", "PII", "user_099", [], False, 280, 0.2), | |
| _record("rec_002", "s3", "PII", "user_099", [], False, 150, 0.2), | |
| _record("rec_003", "slack", "PII", "user_099", [], False, 65, 0.21), | |
| _record("rec_004", "email", "PII", "user_099", [], False, 100, 0.18), | |
| _record("rec_005", "postgres", "legal_hold", "user_099", [], True, 500, 0.4), | |
| _record("rec_006", "s3", "PII", "user_099", [], False, 208, 0.19), | |
| _record("rec_007", "email", "PII", "user_099", [], False, 72, 0.2), | |
| _record("rec_008", "postgres", "PII", "user_099", ["rec_002"], False, 117, 0.24), | |
| _record("rec_009", "s3", "anonymized", "user_099", [], False, 415, 0.05), | |
| _record("rec_010", "slack", "PII", "user_099", ["rec_005"], False, 132, 0.3), | |
| _record("rec_011", "email", "legal_hold", "user_099", [], True, 660, 0.4), | |
| _record("rec_012", "postgres", "anonymized", "user_099", [], False, 365, 0.05), | |
| _record("rec_013", "s3", "anonymized", "user_099", [], False, 322, 0.05), | |
| _record("rec_014", "slack", "anonymized", "user_099", [], False, 280, 0.05), | |
| _record("rec_015", "email", "derived", "user_099", ["rec_008"], False, 49, 0.04), | |
| _record("rec_016", "postgres", "derived", "user_099", [], False, 86, 0.04), | |
| _record("rec_017", "s3", "legal_hold", "user_099", [], True, 710, 0.4), | |
| _record("rec_018", "slack", "derived", "user_099", [], False, 140, 0.04), | |
| _record("rec_019", "email", "PII", "user_100", [], False, 35, 0.2), | |
| _record("rec_020", "postgres", "derived", "user_100", [], False, 80, 0.04), | |
| _record("rec_021", "s3", "legal_hold", "user_099", [], True, 740, 0.4), | |
| _record("rec_022", "slack", "PII", "user_100", [], False, 19, 0.2), | |
| _record("rec_023", "email", "derived", "user_100", [], False, 91, 0.04), | |
| _record("rec_024", "postgres", "PII", "user_100", [], False, 42, 0.2), | |
| _record("rec_025", "s3", "derived", "user_100", [], False, 110, 0.04), | |
| ] | |
| request = EraseRequest( | |
| request_id="gdpr_req_legal_099", | |
| user_id="user_099", | |
| deadline_steps=30, | |
| regulation="GDPR", | |
| submitted_at_step=0, | |
| ) | |
| return { | |
| "task_name": "legal_minefield", | |
| "step_budget": 30, | |
| "request": request, | |
| "records": records, | |
| } | |
| def _cascade_collapse() -> Dict[str, object]: | |
| records = [ | |
| _record("rec_001", "postgres", "PII", "user_777", ["rec_011"], False, 330, 0.23), | |
| _record("rec_002", "s3", "PII", "user_777", ["rec_011"], False, 290, 0.24), | |
| _record("rec_003", "slack", "PII", "user_777", ["rec_012"], False, 190, 0.2), | |
| _record("rec_004", "email", "PII", "user_777", ["rec_013"], False, 130, 0.2), | |
| _record("rec_005", "postgres", "PII", "user_777", ["rec_013"], False, 111, 0.21), | |
| _record("rec_006", "s3", "PII", "user_777", ["rec_014"], False, 88, 0.2), | |
| _record("rec_007", "slack", "PII", "user_777", ["rec_015"], False, 77, 0.2), | |
| _record("rec_008", "email", "PII", "user_777", ["rec_015"], False, 59, 0.21), | |
| _record("rec_009", "postgres", "PII", "user_777", ["rec_016"], False, 47, 0.19), | |
| _record("rec_010", "s3", "PII", "user_777", ["rec_016"], False, 35, 0.2), | |
| _record("rec_011", "postgres", "legal_hold", "user_777", [], True, 740, 0.4), | |
| _record("rec_012", "s3", "legal_hold", "user_777", [], True, 712, 0.4), | |
| _record("rec_013", "slack", "legal_hold", "user_777", [], True, 680, 0.4), | |
| _record("rec_014", "email", "legal_hold", "user_777", [], True, 650, 0.4), | |
| _record("rec_015", "postgres", "legal_hold", "user_777", [], True, 620, 0.4), | |
| _record("rec_016", "s3", "legal_hold", "user_777", [], True, 590, 0.4), | |
| _record("rec_017", "slack", "anonymized", "user_777", ["rec_003"], False, 400, 0.05), | |
| _record("rec_018", "email", "derived", "user_777", ["rec_006"], False, 220, 0.04), | |
| _record("rec_019", "postgres", "derived", "user_777", [], False, 180, 0.04), | |
| _record("rec_020", "s3", "anonymized", "user_777", [], False, 310, 0.05), | |
| _record("rec_021", "postgres", "PII", "user_778", [], False, 210, 0.2), | |
| _record("rec_022", "s3", "PII", "user_778", [], False, 201, 0.2), | |
| _record("rec_023", "slack", "PII", "user_778", [], False, 195, 0.2), | |
| _record("rec_024", "email", "PII", "user_778", [], False, 183, 0.2), | |
| _record("rec_025", "postgres", "derived", "user_778", [], False, 160, 0.04), | |
| _record("rec_026", "s3", "derived", "user_778", [], False, 142, 0.04), | |
| _record("rec_027", "slack", "anonymized", "user_778", [], False, 120, 0.05), | |
| _record("rec_028", "email", "derived", "user_778", [], False, 101, 0.04), | |
| ] | |
| request = EraseRequest( | |
| request_id="gdpr_req_cascade_777", | |
| user_id="user_777", | |
| deadline_steps=18, | |
| regulation="GDPR", | |
| submitted_at_step=0, | |
| ) | |
| return { | |
| "task_name": "cascade_collapse", | |
| "step_budget": 18, | |
| "request": request, | |
| "records": records, | |
| } | |
| TASK_BUILDERS = { | |
| "clean_sweep": _clean_sweep, | |
| "tangled_lake": _tangled_lake, | |
| "legal_minefield": _legal_minefield, | |
| "cascade_collapse": _cascade_collapse, | |
| } | |
| def get_task(task_name: str, seed: int) -> TaskConfig: | |
| if task_name not in TASK_BUILDERS: | |
| available = ", ".join(sorted(TASK_BUILDERS.keys())) | |
| raise ValueError(f"Unknown task '{task_name}'. Available tasks: {available}") | |
| base_task = TASK_BUILDERS[task_name]() | |
| records: List[DataRecord] = [r.model_copy(deep=True) for r in base_task["records"]] | |
| observation_order = [record.record_id for record in records] | |
| rng = random.Random(seed) | |
| rng.shuffle(observation_order) | |
| return TaskConfig( | |
| task_name=base_task["task_name"], | |
| step_budget=int(base_task["step_budget"]), | |
| request=base_task["request"], | |
| records=records, | |
| observation_order=observation_order, | |
| metadata={"seed": seed}, | |
| ) | |