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}, )