PROMPT.md

You are an expert Python backend, ML, and infrastructure engineer.
Your task is to implement a complete, production-ready OpenEnv environment called **PolypharmacyEnv** for training and evaluating agentic RL policies that act as an "elderly polypharmacy safety agent" (clinical pharmacist assistant).

The deliverable MUST satisfy all of the following:
- Fully compliant with the OpenEnv spec (typed models, `step()` / `reset()` / `state()`, `openenv.yaml`, HTTP server, Dockerfile).
- Simulates a realistic healthcare workflow around elderly polypharmacy and dangerous drug combinations.
- Defines at least **3 tasks** (easy → medium → hard) with deterministic agent graders producing scores in (0.0, 1.0).
- Provides shaped rewards over the trajectory (not just sparse terminal rewards).
- Includes a baseline LLM-based inference script `inference.py` in the repo root, following the evaluation requirements:
  - Uses the OpenAI Python client.
  - Reads `OPENAI_API_KEY`, `API_BASE_URL`, `MODEL_NAME`, and `HF_TOKEN` from the environment.
  - Emits structured stdout logs in the exact `[START]`, `[STEP]`, `[END]` format from the OpenEnv sample inference script.
- Is containerized and deployable as a **Hugging Face Space** tagged with `openenv` that responds to OpenEnv-style `reset` / `step` / `state` HTTP calls.

Implement everything described below.

=================================================
1. Repository and folder structure
=================================================

Create a Python package repository with this structure (names are important unless clearly labeled as examples):

- `openenv-polypharmacy/`
  - `openenv.yaml`
  - `README.md`
  - `requirements.txt`
  - `Dockerfile`
  - `inference.py`                 # baseline LLM agent per spec
  - `pyproject.toml` or `setup.cfg` (optional but recommended)
  - `src/`
    - `polypharmacy_env/`
      - `__init__.py`
      - `config.py`
      - `models.py`                # Action, Observation, State, helper models
      - `env_core.py`              # PolypharmacyEnv implementation
      - `tasks.py`                 # task setup utilities
      - `graders.py`               # deterministic graders for each task
      - `rewards.py`               # reward shaping logic
      - `data_loader.py`           # load/preprocess patient and lookup data
      - `ddi_simulator.py`         # local DDI / guideline simulator
      - `api/`
        - `__init__.py`
        - `schemas.py`            # HTTP request/response schemas
        - `server.py`             # FastAPI app exposing OpenEnv endpoints
      - `baselines/`
        - `__init__.py`
        - `heuristic_agent.py`    # simple rule-based baseline agent
        - `random_agent.py`       # trivial random baseline (optional)
      - `tests/`
        - `__init__.py`
        - `test_env_core.py`
        - `test_api.py`
  - `data/`
    - `raw/`                      # placeholder for real/synthetic source data
    - `processed/`
    - `lookups/`
      - `ddi_rules.csv`
      - `beers_criteria.csv`
      - `drug_metadata.csv`
  - `scripts/`
    - `preprocess_data.py`
    - `run_validation.sh`         # optional; runs OpenEnv validator, tests, etc.

Use Python 3.10+ with full type hints, and keep the code black/isort-compatible.

=================================================
2. Domain, data, and clinical abstraction
=================================================

2.1. Core scenario

Model an elderly patient (age ≥ 65) with:
- Demographics: age, sex.
- Comorbidities: e.g., hypertension, diabetes, heart failure, CKD, dementia.
- Basic labs: kidney function (eGFR category), liver function category.
- A current medication list (polypharmacy, e.g., 3–15 drugs depending on task).

Each **episode** is one medication-review session where the agent:
- Observes patient info and current meds.
- Optionally **queries** a DDI/guideline tool for specific drug pairs.
- Proposes **interventions**:
  - `stop`: discontinue a drug.
  - `dose_reduce`: lower dose of a drug.
  - `substitute`: swap to a safer alternative.
  - `add_monitoring`: keep the drug but flag extra monitoring.
- Calls `finish_review` when it decides the regimen is acceptable or budgets are exhausted.

No external PHI, EHRs, or online APIs: all data is **synthetic** or de-identified and local to the container (CSV files).

2.2. Data files and CSV schemas

Implement local CSVs under `data/lookups/`:

**`drug_metadata.csv`**
- `drug_id` (string; unique key)
- `generic_name` (string)
- `atc_class` (string)
- `is_high_risk_elderly` (0/1)
- `default_dose_mg` (float)
- `min_dose_mg` (float)
- `max_dose_mg` (float)

**`beers_criteria.csv`**
- `drug_id` (string)
- `criterion_type` (enum string: `avoid`, `caution`, `dose_adjust`, `avoid_in_condition`)
- `condition` (nullable string; e.g., `CKD`, `dementia`)
- `rationale` (brief text)

**`ddi_rules.csv`**
- `drug_id_1` (string; normalized so `drug_id_1 < drug_id_2` lexicographically)
- `drug_id_2` (string)
- `severity` (enum string: `mild`, `moderate`, `severe`)
- `mechanism` (short text)
- `recommendation` (enum string: `avoid_combination`, `monitor_closely`, `dose_adjust`, `no_action`)
- `base_risk_score` (float in [0.0, 1.0])

Implement a synthetic patient-episode dataset under `data/processed/`:

**`patients_polypharmacy.csv`**
- `episode_id` (string)
- `age` (int)
- `sex` (enum: `M`, `F`, `O`)
- `conditions` (semicolon-separated; e.g., `HTN;DM;CKD`)
- `eGFR_category` (enum: `normal`, `mild`, `moderate`, `severe`)
- `liver_function_category` (enum: `normal`, `impaired`)
- `medication_ids` (semicolon-separated list of `drug_id`)
- `baseline_risk_score` (float in [0.0, 1.0])

2.3. Preprocessing script

In `scripts/preprocess_data.py`:
- If real data is not provided, procedurally generate synthetic but plausible data using:
  - Random combinations of conditions and drugs constrained by simple rules (e.g., CKD + renally-cleared drugs).
  - Controlled distribution of high-risk DDIs and Beers violations.
- Explicitly tag episodes as easy/medium/hard (e.g., via number of drugs, number/severity of DDIs, and number of Beers issues).
- Save `patients_polypharmacy.csv` ready for the environment to consume.

=================================================
3. OpenEnv models and environment implementation
=================================================

3.1. Models

In `models.py`, define dataclasses or Pydantic models that extend the appropriate OpenEnv base types (`Action`, `Observation`, `State`) and are JSON-compatible.

Auxiliary models:

**`MedicationEntry`**
- `drug_id: str`
- `generic_name: str`
- `atc_class: str`
- `dose_mg: float`
- `frequency: str`          # e.g., `qd`, `bid`
- `route: str`              # e.g., `po`
- `is_high_risk_elderly: bool`
- `beers_flags: list[str]`  # e.g., `["avoid", "dose_adjust_CKD"]`

**`InteractionQueryRecord`**
- `drug_id_1: str`
- `drug_id_2: str`
- `severity: str | None`
- `recommendation: str | None`
- `risk_score: float | None`
- `step_index: int`

**`InterventionRecord`**
- `target_drug_id: str`
- `action_type: Literal["stop", "dose_reduce", "substitute", "add_monitoring"]`
- `proposed_new_drug_id: str | None`
- `rationale: str`
- `step_index: int`

Core wire models:

**`PolypharmacyObservation`** (extends OpenEnv `Observation`)
- `episode_id: str`
- `task_id: Literal["easy_screening", "budgeted_screening", "complex_tradeoff"]`
- `age: int`
- `sex: str`
- `conditions: list[str]`
- `eGFR_category: str`
- `liver_function_category: str`
- `current_medications: list[MedicationEntry]`
- `interaction_queries: list[InteractionQueryRecord]`
- `interventions: list[InterventionRecord]`
- `step_index: int`
- `remaining_query_budget: int`
- `remaining_intervention_budget: int`
- `shaped_reward: float`  # reward from last step
- `done: bool`

**`PolypharmacyAction`** (extends OpenEnv `Action`)
- `action_type: Literal["query_ddi", "propose_intervention", "finish_review"]`
- `drug_id_1: str | None`        # for DDI queries or some interventions
- `drug_id_2: str | None`        # for DDI queries
- `target_drug_id: str | None`   # for interventions
- `intervention_type: Literal["stop", "dose_reduce", "substitute", "add_monitoring", "none"] | None`
- `proposed_new_drug_id: str | None`
- `rationale: str | None`

**`PolypharmacyState`** (extends OpenEnv `State`)
- `episode_id: str`
- `task_id: str`
- `step_count: int`
- `max_steps: int`
- `num_query_actions: int`
- `num_interventions: int`

3.2. Environment core

In `env_core.py`, implement `PolypharmacyEnv` extending the appropriate OpenEnv environment base class. It must implement:

**`reset(task_id: str | None = None) -> PolypharmacyObservation`**
- If `task_id` is `None`, default to medium (`budgeted_screening`).
- Sample an episode from `patients_polypharmacy.csv` filtered by difficulty.
- Initialize:
  - `episode_id`
  - `step_count = 0`
  - task-specific budgets (query, interventions, max_steps)
  - baseline regime and risk
  - empty `interaction_queries` and `interventions`
- Return the initial `PolypharmacyObservation` with:
  - `step_index = 0`
  - `shaped_reward = 0.0`
  - `done = False`

**`step(action: PolypharmacyAction) -> dict`**
- Validate the action; if invalid:
  - Apply a negative reward.
  - Do not modify regimen, but log error in `info`.
- If `action_type == "query_ddi"`:
  - If query budget exhausted, apply penalty and do not query.
  - Else:
    - Use `ddi_simulator.lookup_ddi(drug_id_1, drug_id_2)` to get severity, recommendation, base_risk_score.
    - Append an `InteractionQueryRecord`.
    - Apply a small negative reward for query cost.
- If `action_type == "propose_intervention"`:
  - If intervention budget exhausted, apply penalty and ignore change.
  - Else:
    - Update `current_medications` according to `intervention_type`:
      - `stop`: remove medication.
      - `dose_reduce`: adjust dose downward within [min_dose_mg, default_dose_mg].
      - `substitute`: replace with a safer alternative from same `atc_class`.
      - `add_monitoring`: keep drug but tag in internal state.
    - Append an `InterventionRecord`.
    - Recompute current regimen risk using the risk model (see 3.3).
    - Compute shaped reward = (previous_risk - new_risk) - small intervention cost.
- If `action_type == "finish_review"`:
  - Mark `done = True`.
  - Call the task’s grader to get episode-level score in [0.0, 1.0].
  - Add this as a terminal bonus to the current step reward.

- In all cases:
  - Increment `step_count`.
  - Check `max_steps`; if exceeded, auto-terminate:
    - `done = True`
    - apply time-out penalty
    - call grader with current trajectory for a final score if appropriate.
  - Construct next `PolypharmacyObservation` with updated fields.
  - Return a dict:
    - `observation`: `PolypharmacyObservation`
    - `reward`: float shaped reward for this step
    - `done`: bool
    - `info`: dict with fields like `current_risk`, `baseline_risk`, `grader_score_if_terminal`, and debug flags.

**`state` property**
- Returns `PolypharmacyState` reflecting the current internal state.

3.3. DDI simulator and risk model

In `ddi_simulator.py`:
- Load `ddi_rules.csv` once via `data_loader`.
- Implement `lookup_ddi(drug_id_1, drug_id_2) -> tuple[severity, recommendation, base_risk_score]`:
  - Normalize the pair ordering.
  - Look up row; if missing, return:
    - severity = `"none"`
    - recommendation = `"no_action"`
    - base_risk_score = 0.0

In `rewards.py` (or a dedicated module), implement:
- `compute_regimen_risk(current_drug_ids, patient_context, ddi_rules, beers_rules, drug_metadata) -> float`
  - Aggregate contributions from:
    - Beers violations (weighted by `criterion_type` and relevant conditions).
    - DDI base risk scores for all present drug pairs.
    - High-risk elderly drugs.
  - Normalize and clip to [0.0, 1.0].

Use this function to compute:
- `baseline_risk` at episode start.
- Risk after each intervention step.

Also implement:
- `compute_shaped_reward(previous_risk, new_risk, action, context, partial_metrics) -> float`
  - Positive component: `previous_risk - new_risk`.
  - Negative components: per-query cost, per-intervention cost, invalid-action penalty, time-out penalty.

=================================================
4. Tasks and graders (3 difficulty levels)
=================================================

Define three task IDs and semantics in `tasks.py` and `graders.py`:

Task IDs:
- `easy_screening`
- `budgeted_screening`
- `complex_tradeoff`

4.1. `easy_screening` (easy)

- Small regimen: 3–5 drugs.
- Exactly one **severe** DDI pair and possibly one simple Beers violation.
- Budgets:
  - query_budget ≈ 4
  - intervention_budget ≈ 2
  - max_steps ≈ 10

Grader:
- Input: full trajectory, baseline risk, final risk, list of interventions.
- Compute:
  - `risk_reduction = max(0.0, baseline_risk - final_risk) / max(baseline_risk, ε)` (normalized).
  - `targeted_intervention_flag = 1.0` if at least one intervention affects one of the drugs in the known severe DDI pair, else 0.0.
- Score:
  - `score = 0.5 * risk_reduction + 0.5 * targeted_intervention_flag`
  - Clip to [0.0, 1.0].

4.2. `budgeted_screening` (medium)

- Medium regimen: 6–10 drugs.
- Multiple DDIs (mild/moderate/severe) and multiple Beers issues.
- Budgets:
  - query_budget ≈ 8
  - intervention_budget ≈ 3
  - max_steps ≈ 20

Grader:
- Compute:
  - `risk_reduction_score` as normalized risk drop.
  - `intervention_precision_score` = fraction of interventions that actually reduce risk or fix guideline violations.
  - `query_efficiency_score` = (number of severe/moderate DDIs discovered) / (number of queries used), normalized.
- Weighted score, for example:
  - `score = 0.5 * risk_reduction_score + 0.3 * intervention_precision_score + 0.2 * query_efficiency_score`
  - Clip to [0.0, 1.0].

4.3. `complex_tradeoff` (hard)

- Larger regimen: 10–15 drugs.
- Some drugs are **clinically critical** (e.g., anticoagulants, insulin analogues) and encoded as such in `drug_metadata` or a small internal map.
- Episodes contain:
  - multiple DDIs and Beers issues, including ones involving critical drugs.
  - safer substitutes for some risky drugs.

Budgets:
- query_budget ≈ 12
- intervention_budget ≈ 5
- max_steps ≈ 30

Grader adds a **regimen disruption penalty** component:
- Metrics:
  - `risk_reduction_score` (as above).
  - `critical_drug_penalty` = penalty if a critical drug is stopped without substitution to another suitable agent.
  - `total_drug_changes` = number of drugs stopped or substituted.
  - `regimen_disruption_penalty` derived from `total_drug_changes` and `critical_drug_penalty`.

Example scoring:
- `base = risk_reduction_score`
- `penalty = α * regimen_disruption_penalty`
- `score = clamp(base - penalty, 0.0, 1.0)`

4.4. Reward shaping

In `rewards.py`, define a consistent shaping scheme:
- On each query:
  - Small negative reward (e.g., −0.01) plus any small bonus if it discovers a severe DDI, if desired.
- On each intervention:
  - Reward ≈ (previous_risk - new_risk) − small intervention cost.
- On invalid actions:
  - Larger negative reward (e.g., −0.1) and no state change.
- On `finish_review`:
  - Add the task-level `score` ∈ [0.0, 1.0] from the corresponding grader to that step’s shaped reward.

Ensure the sum of step rewards per episode remains in a reasonable numeric range (e.g., roughly -5 to +5) while still allowing meaningful differentiation by graders.

=================================================
5. HTTP API server and openenv.yaml
=================================================

5.1. HTTP server (FastAPI)

In `api/server.py`:
- Implement a FastAPI app that maintains a `PolypharmacyEnv` instance (or a multiplexing scheme if needed).
- Endpoints:
  - `POST /reset`:
    - Request body: may include `task_id` (string).
    - Response: serialized `PolypharmacyObservation`.
  - `POST /step`:
    - Request body: serialized `PolypharmacyAction`.
    - Response: dict with:
      - `observation`: `PolypharmacyObservation`
      - `reward`: float
      - `done`: bool
      - `info`: dict
  - `GET /state`:
    - Response: `PolypharmacyState`.

Provide a module-level `app = FastAPI(...)` object for use with uvicorn and Hugging Face Spaces. Ensure the JSON schema is consistent with OpenEnv clients (simple, flat JSON for observation/action/state).

5.2. `openenv.yaml`

At repo root, define `openenv.yaml` consistent with the latest OpenEnv spec. At minimum, include:
- `name`: `polypharmacy_env`
- `version`: e.g., `0.1.0`
- `description`: human-readable description.
- `author`: your details.
- `tags`: e.g., `["healthcare", "polypharmacy", "openenv"]`
- `tasks`:
  - One entry per task:
    - `id`: `"easy_screening"` / `"budgeted_screening"` / `"complex_tradeoff"`
    - `description`: one-line description
    - `difficulty`: `"easy"`, `"medium"`, `"hard"`

Ensure `openenv validate` (or equivalent validator) passes once implemented.

=================================================
6. Baseline heuristic (non-LLM) agent
=================================================

In `baselines/heuristic_agent.py`, implement a simple, deterministic baseline agent that:

For each episode:
- Iterates through all unordered medication pairs within query budget:
  - Calls `query_ddi` via the environment for each pair until the query budget is exhausted or all pairs are examined.
  - Records severe and moderate interactions.
- After querying:
  - For each severe DDI pair:
    - Try `substitute` one of the drugs using `drug_metadata`:
      - Prefer substitute within same `atc_class` that:
        - is not marked high-risk elderly.
        - does not participate in known severe DDIs with the rest of the regimen.
    - If no substitute exists, propose `stop` for the higher-risk drug.
  - Respect intervention budget limits.
- Finally, call `finish_review`.

This baseline should be callable as a simple Python function that interacts with `PolypharmacyEnv` directly (without HTTP).

=================================================
7. Baseline LLM inference script (inference.py)
=================================================

At repo root, create `inference.py` that:

7.1. Uses the OpenAI Python client

- Import and configure the official OpenAI Python client.
- Read environment variables:
  - `OPENAI_API_KEY` (required).
  - `API_BASE_URL` (base URL for LLM; default to OpenAI standard if not set).
  - `MODEL_NAME` (e.g., `gpt-4.1` or similar).
  - `HF_TOKEN` (if needed for HF auth; do not hardcode).
- Read `POLYPHARMACY_ENV_URL` (or similar) for the environment’s HTTP base URL.

7.2. Implements the required logging format

- For each **run** across all tasks:
  - Emit a `[START]` line with a JSON payload exactly matching the evaluation specification:
    - Fields such as `run_id`, `task_id`, `model`, etc., in the same order and naming as the sample OpenEnv inference script.
- For each **step** in an episode:
  - Emit a `[STEP]` line with JSON fields including:
    - `run_id`
    - `task_id`
    - `episode_id`
    - `step_index`
    - `observation_summary` (brief, machine-readable summary)
    - `action_payload` (the action sent to the env)
    - `reward`
    - `done`
- After finishing an episode for a task:
  - Emit an `[END]` line summarizing:
    - `run_id`
    - `task_id`
    - per-episode statistics (e.g., total reward, grader score from last step’s `info`).
- The stdout format MUST follow the sample exactly:
  - Same tags: `[START]`, `[STEP]`, `[END]`.
  - Same JSON field names and ordering as the provided reference.
  - No extra prints except these structured logs (and necessary error messages to stderr).

7.3. LLM agent loop

- For each task (`easy_screening`, `budgeted_screening`, `complex_tradeoff`):
  - Run a fixed small number of episodes (e.g., 5–10 per task) for baseline scoring.
  - For each episode:
    - Call `/reset` with the task id.
    - At each step:
      - Summarize the observation into a concise prompt for the LLM:
        - Include age, sex, conditions, high-risk flags, budgets, and a compressed view of meds and previous actions.
      - Ask the model to output a **strict JSON** representing `PolypharmacyAction` fields.
      - Parse and validate the JSON; if invalid, fall back to a safe default (e.g., `finish_review` or a no-op) and penalize in evaluation.
      - Send this action to `/step` and log `[STEP]`.
    - End when `done=True` or max_steps is reached.
- At the end, print aggregate scores per task and overall.

Make sure runtime < 20 minutes and that the script can run within 2 vCPUs and 8 GB RAM.

=================================================
8. Dockerfile and Hugging Face Space
=================================================

8.1. Dockerfile

Create a `Dockerfile` that:
- Starts from a slim Python image (e.g., `python:3.11-slim`).
- Installs system dependencies as needed (e.g., `build-essential`, `curl`).
- Copies the project into the container.
- Installs Python dependencies from `requirements.txt`.
- Sets appropriate environment variables for the app (e.g., `PORT=7860`).
- Exposes port 7860.
- Uses a `CMD` or `ENTRYPOINT` that runs the FastAPI server, for example:
  - `uvicorn polypharmacy_env.api.server:app --host 0.0.0.0 --port 7860`

8.2. Hugging Face Space

Ensure the repository is ready to be used as a Hugging Face Space:
- Space type: `docker`.
- Tag: `openenv`.
- On container start, the server must listen on the correct port and respond to:
  - `POST /reset`
  - `POST /step`
  - `GET /state`
- The environment must start cleanly with `docker build` + `docker run` locally.

=================================================
9. README and documentation
=================================================

In `README.md`, include:

- **Environment description & motivation**:
  - What PolypharmacyEnv simulates.
  - Why elderly polypharmacy safety matters.
- **Action and observation spaces**:
  - Describe `PolypharmacyAction`, `PolypharmacyObservation`, and `PolypharmacyState` fields and semantics.
- **Task descriptions**:
  - `easy_screening`, `budgeted_screening`, `complex_tradeoff`, their difficulty and goals.
- **Reward structure**:
  - Summarize shaping and terminal rewards.
- **Setup & usage**:
  - How to install dependencies.
  - How to run the API server locally (uvicorn command).
  - How to run the heuristic baseline.
  - How to run `inference.py` with environment variables.
- **Baseline scores**:
  - Document reproducible baseline scores for each task (heuristic agent, and LLM baseline if available).

=================================================
10. Validation and quality gates
=================================================

- Ensure:
  - `openenv.yaml` and the HTTP server pass the OpenEnv validation script.
  - `docker build` and `docker run` work without errors.
  - `inference.py` completes under 20 minutes, within 2 vCPUs / 8 GB RAM.
  - All graders:
    - Are deterministic.
    - Return scores strictly in [0.0, 1.0].
  - No grader returns a constant score irrespective of behavior.

Aim for clean, well-structured, well-documented code with clear separation of concerns between:
- Data loading,
- Environment state & dynamics,
- Reward/grade logic,
- HTTP serving,
- Baseline agents and inference.