Spaces:
Sleeping
Sleeping
| """ | |
| ClinicalBench — Agentic Reasoning Baseline Inference | |
| ==================================================== | |
| Three architecturally distinct agent strategies: | |
| 1. NAIVE — single LLM call, small sample, no tools, no feedback loop | |
| 2. HEURISTIC — deterministic Python rules (honestly labeled, no LLM) | |
| 3. REACT — true multi-turn ReAct loop: LLM is the brain, Python is the hands | |
| The ReAct agent sends the raw protocol excerpt and raw patient data to the LLM. | |
| The LLM decides which actions to take. Python executes them via env.step() and | |
| feeds the observation back. The LLM then decides the next action. This continues | |
| until the episode ends or the step budget is exhausted. | |
| No deterministic detectors. No fake [THOUGHT] print statements. | |
| The LLM genuinely drives the agent loop. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import statistics | |
| import sys | |
| import time | |
| from collections import Counter | |
| from dataclasses import dataclass, field | |
| from datetime import datetime | |
| from types import SimpleNamespace | |
| from typing import Optional | |
| from openai import OpenAI | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from client import ClinicalTrialAuditorEnv | |
| from models import AuditAction | |
| try: | |
| from server.clinical_trial_auditor_environment import ClinicalTrialAuditorEnvironment | |
| except ImportError: | |
| ClinicalTrialAuditorEnvironment = None | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/hf-inference/v1") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Meta-Llama-3.3-70B-Instruct") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:8000") | |
| BASELINE_SEED = int(os.getenv("BASELINE_SEED", "20260402")) | |
| TASK_LIST = { | |
| "task_easy": "Dynamic Eligibility Screening (Easy)", | |
| "task_medium": "Protocol Timeline Audit (Medium)", | |
| "task_hard": "Equity + Protocol Audit (Hard)", | |
| } | |
| # ═══════════════════════════════════════════════════════════════ | |
| # SYSTEM PROMPT — teaches the LLM how to play the game | |
| # ═══════════════════════════════════════════════════════════════ | |
| REACT_SYSTEM_PROMPT = """You are a Senior Clinical Data Manager auditing a Phase III oncology trial. | |
| You will receive a protocol excerpt and a patient dataset in JSON format. | |
| Your job is to find protocol violations by using the available tools. | |
| ## AVAILABLE TOOLS (action_type values) | |
| 1. **investigate_pattern** — Investigate a variable to see its distribution. | |
| Required field: "variable" (e.g., "age", "death_date", "enrollment_date", "stage", "group", "ethnicity", "gender", "outcome") | |
| You MUST investigate the required variables before you can flag errors. | |
| 2. **compute_distribution** — Compute control-arm distribution for a variable. | |
| Required field: "variable" (e.g., "ethnicity", "gender", "outcome") | |
| Required before flagging selection_bias. | |
| 3. **flag_error** — Flag a protocol violation. | |
| Required fields: "patient_id" (e.g., "P0042"), "error_type", "reason", "confidence" (0.0-1.0) | |
| Valid error_types: "invalid_age", "temporal_inconsistency", "protocol_window_violation", "selection_bias" | |
| For selection_bias: patient_id is not required, but you must have computed ethnicity, gender, and outcome distributions first. | |
| 4. **submit_report** — Submit your final audit report. Do this LAST. | |
| Required field: "report" (text summarizing findings, root cause, risk, corrective actions, and fairness reasoning) | |
| ## CRITICAL WORKFLOW RULES | |
| - You MUST investigate required variables BEFORE flagging errors. The environment will BLOCK premature flags. | |
| - For easy tasks: investigate "age" first. | |
| - For medium tasks: investigate "age", "death_date", "enrollment_date", "stage" first. | |
| - For hard tasks: also compute distributions for "ethnicity", "gender", "outcome" before flagging selection_bias. | |
| - False positives are penalized HEAVILY (-0.26 reward). Do NOT guess. Only flag patients you are confident about. | |
| - Overconfident wrong flags (confidence > 0.8 on a wrong flag) get a 1.8x penalty multiplier. | |
| - BE CONSERVATIVE WITH CONFIDENCE. Set confidence > 0.8 ONLY for extreme violations (age=999, age=-1, death before treatment). For borderline cases (ages near protocol boundaries, delays near the window limit), use confidence 0.5-0.7. When in doubt, do NOT flag. | |
| - You have a LIMITED step budget. Do not waste steps investigating irrelevant variables like bmi, insurance_type, blood_pressure, etc. | |
| - The protocol excerpt contains EPISODE-SPECIFIC rules. Do NOT assume default ranges (e.g., 18-120 for age). READ the protocol. | |
| - Stage IV patients may have an EXTENDED treatment window. But check for comorbidity_index overrides — in some protocols, high comorbidity revokes the Stage IV exception. | |
| - "death_date must never precede treatment_start" — this is a temporal inconsistency. | |
| - For bias: compare mortality WITHIN stage strata, not overall. A raw mortality gap may be confounded by stage distribution. | |
| ## OUTPUT FORMAT | |
| Return a JSON array of actions. Each action is an object with the required fields. | |
| Return ONLY the JSON array, no other text. Example: | |
| [ | |
| {"action_type": "investigate_pattern", "variable": "age"}, | |
| {"action_type": "investigate_pattern", "variable": "death_date"}, | |
| {"action_type": "flag_error", "patient_id": "P0042", "error_type": "invalid_age", "reason": "Age 999 exceeds protocol max of 85", "confidence": 0.98}, | |
| {"action_type": "submit_report", "report": "Protocol-grounded audit found age violations..."} | |
| ] | |
| IMPORTANT: Return valid JSON only. No markdown, no backticks, no explanation text. | |
| """ | |
| NAIVE_SYSTEM_PROMPT = """You are auditing clinical trial patient records. | |
| Given a protocol excerpt and a SMALL SAMPLE of patient records, identify any protocol violations. | |
| Return one finding per line as: PATIENT_ID|ERROR_TYPE|REASON | |
| Valid ERROR_TYPE values: invalid_age, temporal_inconsistency, protocol_window_violation, selection_bias | |
| """ | |
| # ═══════════════════════════════════════════════════════════════ | |
| # Metrics Tracker | |
| # ═══════════════════════════════════════════════════════════════ | |
| class MetricsTracker: | |
| def __init__(self): | |
| self.true_pos = 0 | |
| self.false_pos = 0 | |
| self.total_flagged = 0 | |
| self.steps = 0 | |
| self.llm_calls = 0 | |
| def record(self, feedback: str) -> None: | |
| self.total_flagged += 1 | |
| if "✓" in feedback or "Correct" in feedback: | |
| self.true_pos += 1 | |
| elif "✗" in feedback or "REJECTED" in feedback: | |
| self.false_pos += 1 | |
| def precision(self) -> float: | |
| return self.true_pos / self.total_flagged if self.total_flagged else 0.0 | |
| def summary(self) -> str: | |
| return ( | |
| f" Metrics: {self.true_pos}/{self.total_flagged} correct " | |
| f"(precision={self.precision:.0%}) | {self.steps} steps | {self.llm_calls} LLM call(s)" | |
| ) | |
| # ═══════════════════════════════════════════════════════════════ | |
| # Environment Session Abstraction (with START/STEP/END wrapper) | |
| # ═══════════════════════════════════════════════════════════════ | |
| class EnvLoggerWrapper: | |
| """Wrapper that emits [START]/[STEP]/[END] structured output for Meta's validation bot.""" | |
| def __init__(self, env): | |
| self.env = env | |
| self._task_id = "" | |
| self._step_count = 0 | |
| self._score = 0.0 | |
| self._last_reward = 0.0 | |
| def __enter__(self): | |
| if hasattr(self.env, "__enter__"): | |
| self.env.__enter__() | |
| return self | |
| def __exit__(self, exc_type, exc, tb): | |
| # Emit END block with exact format Meta expects | |
| print(f"[END] task={self._task_id} score={self._score:.2f} steps={self._step_count}", flush=True) | |
| if hasattr(self.env, "__exit__"): | |
| return self.env.__exit__(exc_type, exc, tb) | |
| return False | |
| def reset(self, **kwargs): | |
| # Extract task_id and reset counters | |
| self._task_id = kwargs.get("task_id", "unknown") | |
| self._step_count = 0 | |
| self._score = 0.0 | |
| self._last_reward = 0.0 | |
| # Emit START block | |
| print(f"[START] task={self._task_id}", flush=True) | |
| return self.env.reset(**kwargs) | |
| def step(self, action): | |
| result = self.env.step(action) | |
| self._step_count += 1 | |
| # Safely extract reward and score from the observation | |
| try: | |
| if hasattr(result, 'reward'): | |
| self._last_reward = float(result.reward or 0.0) | |
| if hasattr(result, 'observation'): | |
| obs = result.observation | |
| if hasattr(obs, 'model_dump'): | |
| obs_dict = obs.model_dump() | |
| elif isinstance(obs, dict): | |
| obs_dict = obs | |
| else: | |
| obs_dict = {} | |
| self._score = float(obs_dict.get("score_so_far", self._score)) | |
| except Exception: | |
| pass # Never crash on logging | |
| # Emit STEP block | |
| print(f"[STEP] step={self._step_count} reward={self._last_reward:.2f}", flush=True) | |
| return result | |
| class InProcessEnvSession: | |
| def __init__(self): | |
| if ClinicalTrialAuditorEnvironment is None: | |
| raise RuntimeError("In-process environment is unavailable.") | |
| self._env = ClinicalTrialAuditorEnvironment() | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, exc_type, exc, tb): | |
| return False | |
| def reset(self, **kwargs): | |
| observation = self._env.reset(**kwargs) | |
| return SimpleNamespace(observation=observation, reward=observation.reward, done=observation.done) | |
| def step(self, action: AuditAction): | |
| observation = self._env.step(action) | |
| return SimpleNamespace(observation=observation, reward=observation.reward, done=observation.done) | |
| def open_env_session(): | |
| if ENV_BASE_URL.lower() == "inprocess": | |
| return EnvLoggerWrapper(InProcessEnvSession()) | |
| return EnvLoggerWrapper(ClinicalTrialAuditorEnv(base_url=ENV_BASE_URL).sync()) | |
| # ═══════════════════════════════════════════════════════════════ | |
| # Helper: strip dataset to essential columns to save tokens | |
| # (but do NOT compute anything — the LLM must do the reasoning) | |
| # ═══════════════════════════════════════════════════════════════ | |
| def prepare_dataset_for_llm(dataset: list[dict], task_id: str) -> list[dict]: | |
| """Pass raw data but strip None values and format dates cleanly. | |
| We do NOT calculate anything. We do NOT filter patients. | |
| The LLM must read ALL records and decide what matters. | |
| """ | |
| cleaned = [] | |
| for row in dataset: | |
| clean_row = {} | |
| for key, value in row.items(): | |
| if value is None: | |
| clean_row[key] = None # Keep None — it signals missing data | |
| elif isinstance(value, list) and len(value) == 0: | |
| continue # Skip empty lists to save tokens | |
| else: | |
| clean_row[key] = value | |
| cleaned.append(clean_row) | |
| return cleaned | |
| def truncate_dataset_for_display(dataset: list[dict], max_records: int = 40) -> str: | |
| """For console output only — not sent to LLM.""" | |
| if len(dataset) <= max_records: | |
| return f"{len(dataset)} records" | |
| return f"{len(dataset)} records (showing first {max_records})" | |
| # ═══════════════════════════════════════════════════════════════ | |
| # Parse LLM JSON output into actions | |
| # ═══════════════════════════════════════════════════════════════ | |
| def parse_llm_actions(raw_output: str) -> list[dict]: | |
| """Parse LLM output into action dicts. Handles various formats gracefully.""" | |
| # Strip markdown code fences if present | |
| text = raw_output.strip() | |
| text = re.sub(r'^```(?:json)?\s*', '', text) | |
| text = re.sub(r'\s*```$', '', text) | |
| text = text.strip() | |
| # Try parsing as JSON array | |
| try: | |
| parsed = json.loads(text) | |
| if isinstance(parsed, list): | |
| return parsed | |
| if isinstance(parsed, dict): | |
| # Maybe {"actions": [...]} | |
| if "actions" in parsed: | |
| return parsed["actions"] | |
| return [parsed] | |
| except json.JSONDecodeError: | |
| pass | |
| # Try to find JSON array in the text | |
| match = re.search(r'\[[\s\S]*\]', text) | |
| if match: | |
| try: | |
| return json.loads(match.group()) | |
| except json.JSONDecodeError: | |
| pass | |
| # Try individual JSON objects on separate lines | |
| actions = [] | |
| for line in text.splitlines(): | |
| line = line.strip() | |
| if line.startswith('{'): | |
| try: | |
| actions.append(json.loads(line)) | |
| except json.JSONDecodeError: | |
| pass | |
| if actions: | |
| return actions | |
| # Last resort: try pipe-delimited format (naive agent) | |
| for line in text.splitlines(): | |
| parts = [p.strip() for p in line.split("|")] | |
| if len(parts) >= 2: | |
| actions.append({ | |
| "action_type": "flag_error", | |
| "patient_id": parts[0] if parts[0] and parts[0] != "None" else None, | |
| "error_type": parts[1], | |
| "reason": parts[2] if len(parts) > 2 else "LLM finding", | |
| "confidence": 0.65, | |
| }) | |
| return actions | |
| def action_dict_to_audit_action(action_dict: dict) -> AuditAction: | |
| """Convert a raw dict to an AuditAction, handling missing fields.""" | |
| return AuditAction( | |
| action_type=action_dict.get("action_type", "flag_error"), | |
| patient_id=action_dict.get("patient_id"), | |
| error_type=action_dict.get("error_type"), | |
| reason=action_dict.get("reason"), | |
| proposed_value=action_dict.get("proposed_value"), | |
| variable=action_dict.get("variable"), | |
| report=action_dict.get("report"), | |
| confidence=action_dict.get("confidence"), | |
| ) | |
| # ═══════════════════════════════════════════════════════════════ | |
| # AGENT 1: NAIVE — Single LLM call, small sample, no feedback | |
| # ═══════════════════════════════════════════════════════════════ | |
| def run_naive_task(client: Optional[OpenAI], task_id: str, task_name: str, seed: int): | |
| """Naive agent: one LLM call on a 24-patient sample. No tools, no feedback loop.""" | |
| print(f"\n Task: {task_name}") | |
| print(" " + "-" * 54) | |
| metrics = MetricsTracker() | |
| final_score = 0.0 | |
| with open_env_session() as env: | |
| result = env.reset(task_id=task_id, seed=seed) | |
| obs = result.observation.model_dump() | |
| dataset = obs["dataset"] | |
| protocol_excerpt = obs["trial_protocol_excerpt"] | |
| max_steps = obs["attempts_remaining"] | |
| print(f" Protocol: {obs.get('protocol_title','')} | Patients: {len(dataset)} | Max steps: {max_steps}") | |
| # Naive: only look at 24 patients, single LLM call | |
| sample = dataset[:24] | |
| sample_json = json.dumps(prepare_dataset_for_llm(sample, task_id), default=str, indent=None) | |
| guessed_actions = [] | |
| if client is not None: | |
| try: | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[ | |
| {"role": "system", "content": NAIVE_SYSTEM_PROMPT}, | |
| { | |
| "role": "user", | |
| "content": ( | |
| f"Protocol excerpt:\n{protocol_excerpt}\n\n" | |
| f"Review only these {len(sample)} records:\n{sample_json}\n\n" | |
| "Return findings as PATIENT_ID|ERROR_TYPE|REASON, one per line." | |
| ), | |
| }, | |
| ], | |
| temperature=0, | |
| max_tokens=600, | |
| ) | |
| metrics.llm_calls += 1 | |
| raw = completion.choices[0].message.content or "" | |
| guessed_actions = parse_llm_actions(raw) | |
| print(f" LLM returned {len(guessed_actions)} findings from {len(sample)}-patient sample") | |
| except Exception as exc: | |
| print(f" LLM error: {exc}") | |
| # If no API key or LLM failed, use a trivial fallback | |
| if not guessed_actions: | |
| for row in sample: | |
| age = row.get("age") | |
| if age is None or age < 0 or age > 120: | |
| guessed_actions.append({ | |
| "action_type": "flag_error", | |
| "patient_id": row.get("patient_id"), | |
| "error_type": "invalid_age", | |
| "reason": f"Obvious age anomaly: {age}", | |
| "confidence": 0.55, | |
| }) | |
| # Execute: investigate required variables first, then flags, then report | |
| investigations = {"task_easy": ["age"], "task_medium": ["age", "death_date", "enrollment_date", "stage"], "task_hard": ["age", "death_date", "enrollment_date", "stage"]} | |
| for var in investigations.get(task_id, ["age"]): | |
| if result.done: | |
| break | |
| result = env.step(AuditAction(action_type="investigate_pattern", variable=var)) | |
| metrics.steps += 1 | |
| if task_id == "task_hard": | |
| for var in ["ethnicity", "gender", "outcome"]: | |
| if result.done: | |
| break | |
| result = env.step(AuditAction(action_type="compute_distribution", variable=var)) | |
| metrics.steps += 1 | |
| # Submit LLM's flags (up to step budget) | |
| for action_dict in guessed_actions: | |
| if result.done or metrics.steps >= max_steps - 1: | |
| break | |
| action = action_dict_to_audit_action(action_dict) | |
| if action.action_type != "flag_error": | |
| continue | |
| result = env.step(action) | |
| obs = result.observation.model_dump() | |
| final_score = obs["score_so_far"] | |
| metrics.steps += 1 | |
| metrics.record(obs["feedback"]) | |
| # Submit report | |
| if not result.done: | |
| result = env.step(AuditAction( | |
| action_type="submit_report", | |
| report="Sample review. Possible protocol violations found in limited sample. Recommend full manual audit." | |
| )) | |
| obs = result.observation.model_dump() | |
| final_score = obs["score_so_far"] | |
| metrics.steps += 1 | |
| print(metrics.summary()) | |
| print(f" Final score: {final_score:.2f}\n") | |
| return final_score, metrics | |
| # ═══════════════════════════════════════════════════════════════ | |
| # AGENT 2: HEURISTIC — Deterministic Python rules (honestly labeled) | |
| # No LLM involved. This is a transparent baseline. | |
| # ═══════════════════════════════════════════════════════════════ | |
| def parse_protocol_rules(excerpt: str) -> dict: | |
| """Parse protocol excerpt into rule dict. Used by heuristic agent only.""" | |
| age_match = re.search(r"age (\d+)-(\d+) inclusive", excerpt) | |
| window_match = re.search(r"Treatment must begin within (\d+) days", excerpt) | |
| stage_match = re.search(r"Stage IV exception: treatment may begin within (\d+) days", excerpt) | |
| comorbidity_match = re.search(r"comorbidity_index > (\d+)", excerpt) | |
| bias_match = re.search( | |
| r"dominance exceeds (\d+)%, male share exceeds (\d+)%, " | |
| r"and stage-adjusted mortality gap exceeds (\d+) percentage points", | |
| excerpt, | |
| ) | |
| return { | |
| "age_min": int(age_match.group(1)) if age_match else 18, | |
| "age_max": int(age_match.group(2)) if age_match else 120, | |
| "treatment_window": int(window_match.group(1)) if window_match else 21, | |
| "stage_iv_window": int(stage_match.group(1)) if stage_match else 35, | |
| "comorbidity_threshold": int(comorbidity_match.group(1)) if comorbidity_match else 99, | |
| "bias_dom": int(bias_match.group(1)) / 100 if bias_match else 1.0, | |
| "bias_male": int(bias_match.group(2)) / 100 if bias_match else 1.0, | |
| "bias_gap": int(bias_match.group(3)) / 100 if bias_match else 1.0, | |
| } | |
| def run_heuristic_task(client_unused: Optional[OpenAI], task_id: str, task_name: str, seed: int): | |
| """Heuristic agent: deterministic Python rules. Honestly labeled — no LLM involved. | |
| Deliberate weaknesses (for benchmark difficulty demonstration): | |
| - Uses off-by-3 age margins (misses boundary violations) | |
| - Ignores Stage IV comorbidity override (misses 2-hop window violations) | |
| - Uses overall mortality gap instead of stage-adjusted (falls for Simpson's paradox) | |
| """ | |
| print(f"\n Task: {task_name}") | |
| print(" " + "-" * 54) | |
| metrics = MetricsTracker() | |
| final_score = 0.0 | |
| with open_env_session() as env: | |
| result = env.reset(task_id=task_id, seed=seed) | |
| obs = result.observation.model_dump() | |
| dataset = obs["dataset"] | |
| rules = parse_protocol_rules(obs["trial_protocol_excerpt"]) | |
| max_steps = obs["attempts_remaining"] | |
| print(f" Protocol: {obs.get('protocol_title','')} | Patients: {len(dataset)} | Max steps: {max_steps}") | |
| # Investigations | |
| investigations = {"task_easy": ["age"], "task_medium": ["age", "death_date", "enrollment_date", "stage"], "task_hard": ["age", "death_date", "enrollment_date", "stage", "comorbidity_index"]} | |
| for var in investigations.get(task_id, ["age"]): | |
| if result.done: | |
| break | |
| result = env.step(AuditAction(action_type="investigate_pattern", variable=var)) | |
| metrics.steps += 1 | |
| if task_id == "task_hard": | |
| for var in ["ethnicity", "gender", "outcome"]: | |
| if result.done: | |
| break | |
| result = env.step(AuditAction(action_type="compute_distribution", variable=var)) | |
| metrics.steps += 1 | |
| # Age checks — uses off-by-3 margins (intentional weakness) | |
| findings = [] | |
| for row in dataset: | |
| age = row.get("age") | |
| if age is None or age < (rules["age_min"] - 3) or age > (rules["age_max"] + 3): | |
| findings.append(("invalid_age", row["patient_id"], f"Heuristic: age={age}", 0.82)) | |
| # Temporal checks | |
| for row in dataset: | |
| ts = row.get("treatment_start") | |
| dd = row.get("death_date") | |
| if ts and dd: | |
| try: | |
| t = datetime.strptime(ts, "%Y-%m-%d") | |
| d = datetime.strptime(dd, "%Y-%m-%d") | |
| if d < t: | |
| findings.append(("temporal_inconsistency", row["patient_id"], | |
| f"Death {(t-d).days}d before treatment", 0.90)) | |
| except ValueError: | |
| pass | |
| # Window checks — ignores Stage IV exception entirely (intentional weakness) | |
| if task_id in ("task_medium", "task_hard"): | |
| for row in dataset: | |
| try: | |
| e = datetime.strptime(row.get("enrollment_date", ""), "%Y-%m-%d") | |
| t = datetime.strptime(row.get("treatment_start", ""), "%Y-%m-%d") | |
| delay = (t - e).days | |
| if delay > rules["treatment_window"]: # Uses standard window for ALL stages | |
| findings.append(("protocol_window_violation", row["patient_id"], | |
| f"Delay {delay}d > {rules['treatment_window']}d", 0.80)) | |
| except (ValueError, TypeError): | |
| pass | |
| # Bias — uses overall gap, not stage-adjusted (intentional weakness) | |
| if task_id == "task_hard": | |
| control = [r for r in dataset if r.get("group") == "control"] | |
| if control: | |
| eth_counts = Counter(r.get("ethnicity", "?") for r in control) | |
| dom_eth, dom_count = eth_counts.most_common(1)[0] | |
| dom_ratio = dom_count / len(control) | |
| dom_group = [r for r in control if r.get("ethnicity") == dom_eth] | |
| min_group = [r for r in control if r.get("ethnicity") != dom_eth] | |
| dom_mort = sum(r.get("outcome") == "deceased" for r in dom_group) / max(1, len(dom_group)) | |
| min_mort = sum(r.get("outcome") == "deceased" for r in min_group) / max(1, len(min_group)) | |
| gap = min_mort - dom_mort | |
| if dom_ratio >= max(0.55, rules["bias_dom"] - 0.07) and gap >= 0.10: | |
| findings.append(("selection_bias", None, | |
| f"Overall gap={gap:.0%}, {dom_eth}={dom_ratio:.0%}", 0.74)) | |
| # Execute flags (within budget) | |
| flagged = set() | |
| for error_type, patient_id, reason, confidence in findings: | |
| if result.done or metrics.steps >= max_steps - 1: | |
| break | |
| if patient_id and patient_id in flagged: | |
| continue | |
| if patient_id: | |
| flagged.add(patient_id) | |
| result = env.step(AuditAction( | |
| action_type="flag_error", patient_id=patient_id, | |
| error_type=error_type, reason=reason, confidence=confidence | |
| )) | |
| obs = result.observation.model_dump() | |
| final_score = obs["score_so_far"] | |
| metrics.steps += 1 | |
| metrics.record(obs["feedback"]) | |
| # Report | |
| if not result.done: | |
| result = env.step(AuditAction( | |
| action_type="submit_report", | |
| report="Heuristic rule-based audit. Root cause: possible data-entry drift. Risk: moderate. Recommend validation checks." | |
| )) | |
| obs = result.observation.model_dump() | |
| final_score = obs["score_so_far"] | |
| metrics.steps += 1 | |
| print(metrics.summary()) | |
| print(f" Final score: {final_score:.2f}\n") | |
| return final_score, metrics | |
| def run_react_task(client: Optional[OpenAI], task_id: str, task_name: str, seed: int): | |
| """Genuine ReAct agent with batched patient review. | |
| Phase 1: LLM investigates variables (1 call, ~500 tokens) | |
| Phase 2: LLM reviews 25-patient batches (N calls, ~2K tokens each) | |
| MEMORY WIPE between batches — prevents token snowball | |
| Phase 3: LLM writes report (1 call, ~500 tokens) | |
| Python does NOT parse rules or run detectors. LLM does ALL reasoning. | |
| """ | |
| print(f"\n Task: {task_name}") | |
| print(" " + "-" * 54) | |
| metrics = MetricsTracker() | |
| final_score = 0.0 | |
| if client is None: | |
| print(" ERROR: ReAct agent requires LLM. Set HF_TOKEN or API_KEY.") | |
| return run_heuristic_task(None, task_id, task_name, seed) | |
| BATCH_SIZE = 25 | |
| with open_env_session() as env: | |
| result = env.reset(task_id=task_id, seed=seed) | |
| obs = result.observation.model_dump() | |
| dataset = obs["dataset"] | |
| protocol_excerpt = obs["trial_protocol_excerpt"] | |
| max_steps = obs["attempts_remaining"] | |
| task_start = time.time() | |
| print(f" Model: {MODEL_NAME}") | |
| print(f" Protocol: {obs.get('protocol_title', '')} | " | |
| f"Patients: {len(dataset)} | Max steps: {max_steps}") | |
| filtered = prepare_dataset_for_llm(dataset, task_id) | |
| batches = [filtered[i:i+BATCH_SIZE] for i in range(0, len(filtered), BATCH_SIZE)] | |
| print(f" Batches: {len(batches)} x ~{BATCH_SIZE} patients") | |
| # Task-specific config | |
| if task_id == "task_easy": | |
| inv_vars = ["age"] | |
| dist_vars = [] | |
| allowed = "invalid_age ONLY" | |
| hints = "Check each patient age against protocol range." | |
| elif task_id == "task_medium": | |
| inv_vars = ["age", "death_date", "enrollment_date", "stage"] | |
| dist_vars = [] | |
| allowed = "invalid_age, temporal_inconsistency, protocol_window_violation" | |
| hints = ("Check ages. Check death_date vs treatment_start ordering. " | |
| "Check enrollment-to-treatment delay. Stage IV may have extended window.") | |
| else: | |
| inv_vars = ["age", "death_date", "enrollment_date", "stage", "comorbidity_index"] | |
| dist_vars = ["ethnicity", "gender", "outcome"] | |
| allowed = "invalid_age, temporal_inconsistency, protocol_window_violation, selection_bias" | |
| hints = ("Check ages. Check death_date vs treatment_start. Check delays. " | |
| "CRITICAL: For Stage IV, check comorbidity_index — high comorbidity " | |
| "may REVOKE extended window. For bias: stage-adjusted mortality only.") | |
| # ── PHASE 1: Investigation ── | |
| inv_prompt = ( | |
| f"## PROTOCOL\n{protocol_excerpt}\n\n" | |
| f"## TASK\nAudit {len(dataset)} patients. ALLOWED ERRORS: {allowed}\n\n" | |
| f"## STEP 1: INVESTIGATE\nInvestigate: {', '.join(inv_vars)}\n" | |
| ) | |
| if dist_vars: | |
| inv_prompt += f"Compute distributions: {', '.join(dist_vars)}\n" | |
| inv_prompt += ("\nReturn JSON array of actions.\n" | |
| 'Example: [{"action_type":"investigate_pattern","variable":"age"}]') | |
| print(f"\n -- Phase 1: Investigation --") | |
| try: | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role":"system","content":REACT_SYSTEM_PROMPT}, | |
| {"role":"user","content":inv_prompt}], | |
| temperature=0, max_tokens=512, | |
| ) | |
| metrics.llm_calls += 1 | |
| inv_actions = parse_llm_actions(completion.choices[0].message.content or "") | |
| print(f" LLM: {len(inv_actions)} investigation actions") | |
| except Exception as e: | |
| print(f" [ERROR] {e}") | |
| inv_actions = [] | |
| inv_fb = [] | |
| for ad in inv_actions: | |
| if result.done or metrics.steps >= max_steps - 2: | |
| break | |
| action = action_dict_to_audit_action(ad) | |
| if action.action_type not in ("investigate_pattern", "compute_distribution"): | |
| continue | |
| result = env.step(action) | |
| obs = result.observation.model_dump() | |
| final_score = obs["score_so_far"] | |
| metrics.steps += 1 | |
| fb = obs.get("feedback", "") | |
| print(f" Step {metrics.steps}: [INV] {fb[:90]}") | |
| inv_fb.append(fb) | |
| inv_summary = "\n".join(f"- {f}" for f in inv_fb) | |
| # ── PHASE 2: Batched Patient Review ── | |
| print(f"\n -- Phase 2: Batched Review ({len(batches)} batches) --") | |
| flagged_ids = set() | |
| all_flag_fb = [] | |
| for bi, batch in enumerate(batches): | |
| if result.done or metrics.steps >= max_steps - 1: | |
| break | |
| remaining = max_steps - metrics.steps - 1 | |
| if remaining <= 0: | |
| break | |
| batch_json = json.dumps(batch, default=str, separators=(",",":")) | |
| # FRESH context — memory wipe each batch | |
| bp = ( | |
| f"## PROTOCOL\n{protocol_excerpt}\n\n" | |
| f"## INVESTIGATION SUMMARY\n{inv_summary}\n\n" | |
| f"## BATCH {bi+1}/{len(batches)}\n{batch_json}\n\n" | |
| f"## INSTRUCTIONS\nALLOWED ERRORS: {allowed}\n{hints}\n" | |
| f"Steps left: {remaining}. Already flagged: " | |
| f"{sorted(flagged_ids) if flagged_ids else 'none'}\n\n" | |
| f"Flag violations in THIS batch. Return JSON array.\n" | |
| f"Empty array [] if none. No investigate or submit actions." | |
| ) | |
| try: | |
| t0 = time.time() | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role":"system","content":REACT_SYSTEM_PROMPT}, | |
| {"role":"user","content":bp}], | |
| temperature=0, max_tokens=2048, | |
| ) | |
| el = time.time() - t0 | |
| metrics.llm_calls += 1 | |
| ba = parse_llm_actions(completion.choices[0].message.content or "") | |
| print(f" [Batch {bi+1}/{len(batches)}] {len(ba)} flags ({el:.1f}s)") | |
| except Exception as e: | |
| print(f" [Batch {bi+1}] ERROR: {e}") | |
| continue | |
| for ad in ba: | |
| if result.done or metrics.steps >= max_steps - 1: | |
| break | |
| action = action_dict_to_audit_action(ad) | |
| if action.action_type != "flag_error": | |
| continue | |
| pid = action.patient_id | |
| if pid and pid in flagged_ids: | |
| continue | |
| result = env.step(action) | |
| obs = result.observation.model_dump() | |
| final_score = obs["score_so_far"] | |
| metrics.steps += 1 | |
| fb = obs.get("feedback", "") | |
| metrics.record(fb) | |
| if pid: | |
| flagged_ids.add(pid) | |
| tag = "✓" if "✓" in fb else "✗" if "✗" in fb else "→" | |
| print(f" Step {metrics.steps}: score={final_score:.2f} [{tag}] {fb[:100]}") | |
| all_flag_fb.append(fb[:120]) | |
| # ── PHASE 3: Report ── | |
| if not result.done: | |
| print(f"\n -- Phase 3: Report --") | |
| rp = ( | |
| f"## PROTOCOL\n{protocol_excerpt[:500]}\n\n" | |
| f"## RESULTS\nFlags: {len(flagged_ids)}\n" | |
| + "\n".join(f"- {f}" for f in all_flag_fb[:15]) + "\n\n" | |
| f"Write audit report: protocol grounding, root cause, risk, corrections, fairness.\n" | |
| f'Return: [{{"action_type":"submit_report","report":"your text"}}]' | |
| ) | |
| try: | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role":"system","content":REACT_SYSTEM_PROMPT}, | |
| {"role":"user","content":rp}], | |
| temperature=0, max_tokens=512, | |
| ) | |
| metrics.llm_calls += 1 | |
| ra = parse_llm_actions(completion.choices[0].message.content or "") | |
| if ra and ra[0].get("report"): | |
| result = env.step(action_dict_to_audit_action(ra[0])) | |
| else: | |
| result = env.step(AuditAction(action_type="submit_report", | |
| report="Protocol audit complete. Corrective actions recommended.")) | |
| except Exception: | |
| result = env.step(AuditAction(action_type="submit_report", | |
| report="Protocol audit complete. Corrective actions recommended.")) | |
| obs = result.observation.model_dump() | |
| final_score = obs.get("score_so_far", 0.0) | |
| metrics.steps += 1 | |
| print(f" Step {metrics.steps}: [REPORT] submitted") | |
| task_elapsed = time.time() - task_start | |
| print(f" Time: {task_elapsed:.1f}s | LLM calls: {metrics.llm_calls}") | |
| print(metrics.summary()) | |
| print(f" Final score: {final_score:.2f}\n") | |
| return final_score, metrics | |
| # ═══════════════════════════════════════════════════════════════ | |
| # Agent Runner | |
| # ═══════════════════════════════════════════════════════════════ | |
| def run_agent(mode: str, client: Optional[OpenAI], seed: int): | |
| runner = { | |
| "naive": run_naive_task, | |
| "heuristic": run_heuristic_task, | |
| "full": run_react_task, | |
| }[mode] | |
| scores = [] | |
| metrics_list = [] | |
| start = time.time() | |
| for task_id, task_name in TASK_LIST.items(): | |
| score, metrics = runner(client, task_id, task_name, seed) | |
| scores.append(score) | |
| metrics_list.append(metrics) | |
| return { | |
| "mode": mode, | |
| "scores": dict(zip(TASK_LIST.keys(), scores)), | |
| "average": sum(scores) / len(scores), | |
| "elapsed": time.time() - start, | |
| "total_steps": sum(metric.steps for metric in metrics_list), | |
| "total_llm": sum(metric.llm_calls for metric in metrics_list), | |
| "avg_precision": statistics.mean(metric.precision for metric in metrics_list), | |
| } | |
| def main(): | |
| parser = argparse.ArgumentParser(description="ClinicalBench baseline inference") | |
| parser.add_argument("--mode", choices=["naive", "heuristic", "full", "all"], default="full") | |
| parser.add_argument("--seed", type=int, default=BASELINE_SEED) | |
| args = parser.parse_args() | |
| client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) if HF_TOKEN else None | |
| print("=" * 70) | |
| print(" ClinicalBench — Agentic Reasoning Baseline Inference") | |
| print(" LLM-Driven ReAct Loop | Protocol-Aware | Stage-Adjusted Fairness") | |
| print(f" Model: {MODEL_NAME}") | |
| print(f" Seed: {args.seed}") | |
| print("=" * 70) | |
| if client is None: | |
| print(" ⚠ No API key detected. ReAct agent will fall back to heuristic.") | |
| print(" Set HF_TOKEN or OPENAI_API_KEY to enable LLM-driven auditing.") | |
| modes = ["naive", "heuristic", "full"] if args.mode == "all" else [args.mode] | |
| results = [] | |
| for mode in modes: | |
| label = {"naive": "NAIVE LLM", "heuristic": "HEURISTIC (deterministic)", "full": "REACT (LLM-driven)"} | |
| print(f"\n{'═' * 70}") | |
| print(f" AGENT: {label.get(mode, mode.upper())}") | |
| print(f"{'═' * 70}") | |
| results.append(run_agent(mode, client, args.seed)) | |
| print("\n" + "=" * 70) | |
| print(" BENCHMARK RESULTS") | |
| print("=" * 70) | |
| if len(results) > 1: | |
| header = f" {'Agent':<15} {'Easy':>8} {'Medium':>8} {'Hard':>8} {'Avg':>8} {'Prec':>8} {'LLM':>6} {'Time':>8}" | |
| print(header) | |
| print(" " + "-" * 72) | |
| for result in results: | |
| scores = result["scores"] | |
| label = {"naive": "Naive LLM", "heuristic": "Heuristic", "full": "ReAct LLM"}.get(result['mode'], result['mode']) | |
| print( | |
| f" {label:<15} " | |
| f"{scores['task_easy']:.2f} {scores['task_medium']:.2f} " | |
| f"{scores['task_hard']:.2f} {result['average']:.2f} " | |
| f"{result['avg_precision']:.0%} {result['total_llm']:>3} " | |
| f"{result['elapsed']:.1f}s" | |
| ) | |
| else: | |
| result = results[0] | |
| for task_id, task_name in TASK_LIST.items(): | |
| print(f" {task_name:38s}: {result['scores'][task_id]:.2f}") | |
| print(f"\n Average score: {result['average']:.2f}") | |
| print(f" Total time: {result['elapsed']:.1f}s") | |
| print(f" LLM calls: {result['total_llm']}") | |
| print(f" Total steps: {result['total_steps']}") | |
| print(f" Average precision: {result['avg_precision']:.0%}") | |
| print("=" * 70) | |
| if __name__ == "__main__": | |
| main() | |