Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from typing import Any, Dict, List | |
| import requests | |
| from openai import OpenAI | |
| from pydantic import ValidationError | |
| from models import EraseOpsAction | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| TASK_NAME = os.getenv("TASK_NAME", "clean_sweep") | |
| ENV_URL = os.getenv("ENV_URL", "http://localhost:7860") | |
| MAX_TOTAL_RUNTIME_SECONDS = int(os.getenv("MAX_TOTAL_RUNTIME_SECONDS", "1100")) | |
| MAX_TASK_RUNTIME_SECONDS = int(os.getenv("MAX_TASK_RUNTIME_SECONDS", "260")) | |
| SYSTEM_PROMPT = ( | |
| "You are a Data Privacy Officer agent managing GDPR Right-to-Erasure requests.\n" | |
| "At each step you receive a JSON observation showing all data records across enterprise systems.\n" | |
| "Your goal: delete all PII records belonging to the target user_id in the request.\n\n" | |
| "Rules you must follow:\n" | |
| "1. NEVER delete or anonymize records where legal_hold=True - escalate these instead\n" | |
| "2. NEVER delete records belonging to a different user_id than in the request\n" | |
| "3. Always delete dependency records BEFORE the records that depend on them\n" | |
| "4. Use 'inspect' action first if you are unsure about a record's dependencies\n" | |
| "5. Use 'anonymize' only for derived or anonymized type records, never for PII\n\n" | |
| "Respond with ONLY a valid JSON object matching this schema:\n" | |
| "{\n" | |
| " 'action_type': 'inspect' | 'delete' | 'anonymize' | 'escalate' | 'noop',\n" | |
| " 'record_id': 'rec_001',\n" | |
| " 'justification': 'brief reason'\n" | |
| "}" | |
| ) | |
| client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL) | |
| def _log_start(task: str, model: str) -> None: | |
| print(f"[START] task={task} env=gdpr-eraseops model={model}", flush=True) | |
| def _log_step(step: int, action_json: str, reward: float, done: bool, error: str | None) -> None: | |
| error_text = "null" if error in (None, "", "None") else error | |
| print( | |
| f"[STEP] step={step} action={action_json} reward={reward:.2f} done={str(done).lower()} error={error_text}", | |
| flush=True, | |
| ) | |
| def _log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: | |
| rewards_str = ",".join(f"{value:.2f}" for value in rewards) | |
| print(f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}", flush=True) | |
| def _safe_action() -> Dict[str, Any]: | |
| return {"action_type": "noop", "record_id": None, "justification": "safe fallback"} | |
| def _candidate_record(observation: Dict[str, Any]) -> str | None: | |
| request_user = observation.get("request", {}).get("user_id") | |
| records = observation.get("records", []) | |
| for record in records: | |
| if record.get("status") != "active": | |
| continue | |
| if record.get("legal_hold") and record.get("user_id") == request_user: | |
| return str(record.get("record_id")) | |
| for record in records: | |
| if record.get("status") != "active": | |
| continue | |
| if record.get("user_id") != request_user: | |
| continue | |
| if record.get("data_type") != "PII": | |
| continue | |
| dependencies = record.get("dependencies", []) | |
| deps_deleted = True | |
| for dep_id in dependencies: | |
| dep = next((r for r in records if r.get("record_id") == dep_id), None) | |
| if dep is None or dep.get("status") != "deleted": | |
| deps_deleted = False | |
| break | |
| if deps_deleted: | |
| return str(record.get("record_id")) | |
| return None | |
| def _heuristic_action(observation: Dict[str, Any]) -> Dict[str, Any]: | |
| request_user = observation.get("request", {}).get("user_id") | |
| records = observation.get("records", []) | |
| step = int(observation.get("step", 0) or 0) | |
| if step < min(3, len(records)): | |
| rec = records[step] | |
| return { | |
| "action_type": "inspect", | |
| "record_id": rec.get("record_id"), | |
| "justification": "progressive reconnaissance", | |
| } | |
| for record in records: | |
| if ( | |
| record.get("status") == "active" | |
| and record.get("user_id") == request_user | |
| and record.get("legal_hold") is True | |
| ): | |
| return { | |
| "action_type": "escalate", | |
| "record_id": record.get("record_id"), | |
| "justification": "legal hold discovered", | |
| } | |
| candidate = _candidate_record(observation) | |
| if candidate: | |
| return { | |
| "action_type": "delete", | |
| "record_id": candidate, | |
| "justification": "delete active in-scope PII with dependencies satisfied", | |
| } | |
| return {"action_type": "noop", "record_id": None, "justification": "no safe move"} | |
| def _llm_action(observation: Dict[str, Any]) -> Dict[str, Any]: | |
| if not HF_TOKEN: | |
| return _heuristic_action(observation) | |
| user_prompt = ( | |
| f"{json.dumps(observation, separators=(',', ':'), ensure_ascii=False)}\n" | |
| "Choose your next action. Respond with JSON only." | |
| ) | |
| try: | |
| response = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| temperature=0.0, | |
| max_tokens=220, | |
| ) | |
| raw = (response.choices[0].message.content or "").strip() | |
| if raw.startswith("```"): | |
| raw = raw.replace("```json", "").replace("```", "").strip() | |
| parsed = json.loads(raw) | |
| return { | |
| "action_type": parsed.get("action_type", "noop"), | |
| "record_id": parsed.get("record_id"), | |
| "justification": parsed.get("justification"), | |
| } | |
| except Exception: | |
| return _heuristic_action(observation) | |
| def _validated_action(raw_action: Dict[str, Any]) -> Dict[str, Any]: | |
| try: | |
| return EraseOpsAction(**raw_action).model_dump() | |
| except ValidationError: | |
| return _safe_action() | |
| def _run_task(task: str) -> float: | |
| _log_start(task, MODEL_NAME) | |
| task_started_at = time.monotonic() | |
| reset_resp = requests.post( | |
| f"{ENV_URL}/reset", | |
| params={"task_name": task, "seed": 42}, | |
| timeout=20, | |
| ) | |
| reset_resp.raise_for_status() | |
| observation = reset_resp.json() | |
| rewards: List[float] = [] | |
| done = False | |
| step = 0 | |
| max_steps = int(observation.get("step_budget", 30) or 30) | |
| while not done and step < max_steps: | |
| if time.monotonic() - task_started_at > MAX_TASK_RUNTIME_SECONDS: | |
| print(f"Runtime guard: task '{task}' timed out", file=sys.stderr, flush=True) | |
| done = True | |
| break | |
| step += 1 | |
| raw_action = _llm_action(observation) | |
| action = _validated_action(raw_action) | |
| action_json = json.dumps(action, separators=(",", ":"), ensure_ascii=False) | |
| error_text: str | None = None | |
| reward_total = 0.0 | |
| try: | |
| step_resp = requests.post(f"{ENV_URL}/step", json=action, timeout=20) | |
| step_resp.raise_for_status() | |
| payload = step_resp.json() | |
| observation = payload.get("observation", {}) | |
| reward_total = float(payload.get("reward", {}).get("total", 0.0) or 0.0) | |
| done = bool(payload.get("done", False)) | |
| info = payload.get("info", {}) | |
| error_text = info.get("error") if isinstance(info, dict) else None | |
| except Exception as exc: | |
| done = True | |
| reward_total = -1.0 | |
| error_text = str(exc) | |
| rewards.append(reward_total) | |
| _log_step(step=step, action_json=action_json, reward=reward_total, done=done, error=error_text) | |
| grade_resp = requests.post(f"{ENV_URL}/grade", params={"task_name": task}, timeout=20) | |
| grade_resp.raise_for_status() | |
| score = float(grade_resp.json().get("score", 0.0) or 0.0) | |
| success = score >= 0.5 | |
| _log_end(success=success, steps=step, score=score, rewards=rewards) | |
| return score | |
| def main() -> None: | |
| """Run the full baseline across all tasks.""" | |
| tasks = ["clean_sweep", "tangled_lake", "legal_minefield", "cascade_collapse"] | |
| run_started_at = time.monotonic() | |
| for task in tasks: | |
| if time.monotonic() - run_started_at > MAX_TOTAL_RUNTIME_SECONDS: | |
| print("Runtime guard: stopping before timeout", file=sys.stderr, flush=True) | |
| break | |
| _run_task(task) | |
| if __name__ == "__main__": | |
| main() | |