inference.py

"""
PharmaAgent — inference.py
Hackathon evaluation script (OpenMV-compatible).

Runs an LLM agent through clinical decision episodes and reports
a grader score from 0.0 to 1.0.

Design:
  - The agent (Qwen 2.5 72B via HuggingFace router) makes clinical decisions.
  - All scoring is performed by environment.py using DrugBank data only.
  - Groq (Llama 3.3 70B) is called ONCE at the end of each episode,
    purely to format a human-readable clinical summary of the final regimen.
    It has zero influence on any reward value.

Usage:
    uv run inference.py

Environment variables (.env or export):
    HF_TOKEN          — HuggingFace token (required)
    GROQ_API_KEY      — Groq API key (required for end-of-episode summary)
    DOCKER_IMAGE_NAME — Docker image name (default: pharma-agent)
    MODEL_NAME        — HF model (default: Qwen/Qwen2.5-72B-Instruct)
"""

import os
import time
import json
import subprocess
import requests
from openai import OpenAI
from dotenv import load_dotenv

# Reward ceiling from environment — keeps normalisation in sync
from server.environment import MAX_EPISODE_REWARD

load_dotenv()

# ── Config ────────────────────────────────────────────────────────────────
HF_TOKEN      = os.environ.get("HF_TOKEN", "")
GROQ_API_KEY  = os.environ.get("GROQ_API_KEY", "")
DOCKER_IMAGE  = os.environ.get("DOCKER_IMAGE_NAME", "pharma-agent")
MODEL_NAME    = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
ENV_PORT      = 7860
ENV_BASE_URL  = f"http://localhost:{ENV_PORT}"
HF_ROUTER_URL = "https://router.huggingface.co/v1"
NUM_EPISODES  = 5

# ── Clients ───────────────────────────────────────────────────────────────

# Agent: Qwen via HuggingFace router (makes clinical decisions)
agent_client = OpenAI(base_url=HF_ROUTER_URL, api_key=HF_TOKEN)

# Formatter: Groq (formats the final regimen summary only — no reward influence)
try:
    from groq import Groq
    groq_client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None
except ImportError:
    groq_client = None


# ── Docker helpers ────────────────────────────────────────────────────────

def start_docker() -> bool:
    """Start the environment container. Returns True when server is ready."""
    print(f"Starting Docker container: {DOCKER_IMAGE}")
    proc = subprocess.Popen(
        ["docker", "run", "--rm", "-p", f"{ENV_PORT}:{ENV_PORT}",
         "-e", "ENABLE_WEB_INTERFACE=true", DOCKER_IMAGE],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
    )
    for _ in range(30):
        try:
            if requests.get(f"{ENV_BASE_URL}/health", timeout=2).status_code == 200:
                print("Environment server is ready.")
                return True
        except Exception:
            pass
        ret = proc.poll()
        if ret is not None:
            out = proc.stdout.read().decode(errors="replace")
            err = proc.stderr.read().decode(errors="replace")
            print(f"Docker container exited with code {ret}.")
            if out: print(f"  stdout: {out[:400]}")
            if err: print(f"  stderr: {err[:400]}")
            return False
        time.sleep(2)
    proc.terminate()
    print("Timed out waiting for environment server.")
    return False


def env_reset() -> dict:
    r = requests.post(f"{ENV_BASE_URL}/reset")
    r.raise_for_status()
    return r.json()


def env_step(session_id: str, action_type: str, value: str) -> dict:
    r = requests.post(
        f"{ENV_BASE_URL}/step?session_id={session_id}",
        json={"action_type": action_type, "value": value},
    )
    r.raise_for_status()
    return r.json()


# ── Groq formatter (runs after episode, no reward influence) ──────────────

def format_regimen_summary(episode_log: dict) -> str:
    """
    Ask Groq to write a plain-English clinical handover note for the
    final regimen. This is purely for human readability — the score
    has already been computed from DrugBank data before this is called.
    """
    if not groq_client:
        return ""
    try:
        prompt = (
            f"A clinical decision agent managed a patient case.\n\n"
            f"Condition: {episode_log.get('condition', 'unknown')}\n"
            f"Symptoms: {', '.join(episode_log.get('symptoms', []))}\n"
            f"Existing medications: {', '.join(episode_log.get('existing_meds', [])) or 'None'}\n"
            f"Agent diagnosis: {episode_log.get('diagnosis', 'not established')}\n"
            f"Final regimen: {', '.join(episode_log.get('final_drugs', [])) or 'None'}\n"
            f"DDI checks performed: {episode_log.get('ddi_checks', [])}\n"
            f"Score: {episode_log.get('score', 0)} / 1.0\n\n"
            f"Write a concise clinical handover note (3-5 sentences) summarising "
            f"this regimen for a human pharmacist to review. "
            f"Flag any concerns. Do not add any new clinical recommendations — "
            f"only summarise what the agent did."
        )
        response = groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=300,
            temperature=0.1,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"(Summary unavailable: {e})"


# ── Agent system prompt ───────────────────────────────────────────────────

SYSTEM_PROMPT = """You are PharmaAgent, a clinical pharmacist AI operating in a drug safety evaluation environment.

IMPORTANT: This environment scores your actions using a real pharmacological database (DrugBank, 19,842 drugs, 2.9M interaction pairs). Your decisions are evaluated against real drug data — not opinions.

Your task per episode:
1. DIAGNOSE — identify the condition from symptoms. Use standard clinical terminology.
2. SELECT_DRUG — add drugs to the regimen one at a time. Use exact approved drug names as they appear in drug databases (e.g. "Metformin", "Atorvastatin", "Amlodipine"). Misspelled or hallucinated names earn zero reward.
3. CHECK_DDI — check for interactions between the drugs you've selected and the patient's existing medications. Format: "Drug1,Drug2".
4. FINALIZE — submit the final regimen.

Scoring rules you must know:
- A drug earns full reward only if DrugBank confirms it is indicated for this condition.
- A drug that has a major/contraindicated interaction with the patient's existing medication earns a safety penalty.
- Performing DDI checks earns reward — skipping them when the patient has existing meds earns a penalty.
- Drug names not found in DrugBank earn zero, regardless of clinical reasoning.

Respond ONLY in this JSON format:
{
  "action_type": "diagnose|select_drug|check_ddi|finalize",
  "value": "your response here",
  "reasoning": "brief clinical reasoning (1-2 sentences)"
}"""


# ── Episode runner ────────────────────────────────────────────────────────

def run_agent_episode() -> float:
    """Run one full episode. Returns normalised reward (0.0–1.0)."""
    reset_data = env_reset()
    session_id = reset_data["session_id"]
    obs = reset_data["observation"]

    case_condition = obs["patient_case"].get("condition", "Unknown")
    symptoms = obs["patient_case"].get("symptoms", [])
    existing_meds = obs["patient_case"].get("existing_medications", [])

    print(f"\n{'─' * 60}")
    print(f"Condition:     {case_condition}")
    print(f"Symptoms:      {', '.join(symptoms)}")
    print(f"Existing meds: {', '.join(existing_meds) or 'None'}")

    conversation = [{"role": "system", "content": SYSTEM_PROMPT}]
    total_reward = 0.0
    episode_log = {
        "condition": case_condition,
        "symptoms": symptoms,
        "existing_meds": existing_meds,
        "diagnosis": None,
        "final_drugs": [],
        "ddi_checks": [],
        "score": 0.0,
    }

    for step_num in range(8):
        user_msg = (
            f"PATIENT CASE:\n"
            f"Condition (hidden — diagnose from symptoms): {case_condition}\n"
            f"Symptoms: {', '.join(obs['patient_case'].get('symptoms', []))}\n"
            f"Existing medications: {', '.join(obs['patient_case'].get('existing_medications', [])) or 'None'}\n"
            f"Current regimen: {', '.join(obs['patient_case'].get('current_regimen', [])) or 'None'}\n"
            f"Diagnosis so far: {obs['patient_case'].get('proposed_diagnosis') or 'Not yet established'}\n\n"
            f"Environment feedback:\n{obs['feedback']}\n\n"
            f"Valid actions: {obs['valid_options']}\n"
            f"Cumulative reward: {obs['reward_so_far']}\n\n"
            f"What is your next action? Respond in JSON."
        )
        conversation.append({"role": "user", "content": user_msg})

        try:
            response = agent_client.chat.completions.create(
                model=MODEL_NAME,
                messages=conversation,
                max_tokens=300,
                temperature=0.2,
            )
            agent_reply = response.choices[0].message.content.strip()
        except Exception as e:
            print(f"LLM error: {e}")
            break

        conversation.append({"role": "assistant", "content": agent_reply})

        try:
            clean = agent_reply.replace("```json", "").replace("```", "").strip()
            action_data = json.loads(clean)
            action_type = action_data.get("action_type", "finalize")
            value = action_data.get("value", "")
            reasoning = action_data.get("reasoning", "")
        except json.JSONDecodeError:
            print("Could not parse agent JSON — finalising.")
            action_type, value, reasoning = "finalize", "finalize", ""

        print(f"\nStep {step_num + 1} | {action_type.upper()} | {value}")
        if reasoning:
            print(f"  Reasoning: {reasoning[:100]}")

        try:
            step_result = env_step(session_id, action_type, value)
        except Exception as e:
            print(f"Environment step error: {e}")
            break

        obs = step_result["observation"]
        step_reward = step_result["reward"]
        total_reward = obs["reward_so_far"]
        done = step_result["done"]

        print(f"  Reward: +{step_reward:.3f} | Cumulative: {total_reward:.3f}")

        # Track episode log for Groq summary
        if action_type == "diagnose":
            episode_log["diagnosis"] = value
        elif action_type == "check_ddi":
            episode_log["ddi_checks"].append(value)

        if done:
            episode_log["final_drugs"] = obs["patient_case"].get("current_regimen", [])
            normalized = round(min(max(total_reward / MAX_EPISODE_REWARD, 0.0), 1.0), 4)
            episode_log["score"] = normalized

            # Groq formats a human-readable summary — does NOT affect the score
            summary = format_regimen_summary(episode_log)
            if summary:
                print(f"\n  Clinical Summary (Groq — informational only):\n  {summary}")

            print(f"\nEpisode complete. Reward: {total_reward:.3f} | Normalised: {normalized:.4f}")
            return normalized

    normalized = round(min(max(total_reward / MAX_EPISODE_REWARD, 0.0), 1.0), 4)
    return normalized


# ── Grader ────────────────────────────────────────────────────────────────

def grader(episode_rewards: list) -> float:
    """Average normalised reward across episodes, with consistency bonus."""
    if not episode_rewards:
        return 0.0
    avg = sum(episode_rewards) / len(episode_rewards)
    if len(episode_rewards) > 1:
        variance = sum((r - avg) ** 2 for r in episode_rewards) / len(episode_rewards)
        consistency_bonus = max(0.0, 0.1 - variance)
        avg = min(1.0, avg + consistency_bonus)
    return round(avg, 4)


# ── Main ──────────────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("  PharmaAgent — Clinical Decision RL Environment")
    print("  Scoring: DrugBank data only | LLM: formatting only")
    print("=" * 60)

    if not HF_TOKEN:
        print("HF_TOKEN not set. Add it to your .env file.")
        return

    try:
        r = requests.get(f"{ENV_BASE_URL}/health", timeout=3)
        if r.status_code == 200:
            print("Environment already running.")
    except Exception:
        if not start_docker():
            print("Could not start environment. Ensure Docker is running.")
            return

    episode_rewards = []
    for ep in range(1, NUM_EPISODES + 1):
        print(f"\n{'=' * 60}")
        print(f"  Episode {ep}/{NUM_EPISODES}")
        reward = run_agent_episode()
        episode_rewards.append(reward)
        print(f"  Episode normalised reward: {reward:.4f}")

    final_score = grader(episode_rewards)

    print(f"\n{'=' * 60}")
    print(f"  FINAL GRADER SCORE: {final_score:.4f} / 1.0000")
    print(f"  Episode rewards:    {episode_rewards}")
    print("=" * 60)

    return final_score


if __name__ == "__main__":
    main()