agentclinic.py

import argparse
import anthropic
from transformers import pipeline
import openai, re, random, time, json, replicate, os
import datetime
from pathlib import Path

# ---------------------------------------------------------------------------
# Trajectory helpers
# ---------------------------------------------------------------------------

def make_trajectory(scenario_id, dataset, doctor_llm, patient_llm,
                    measurement_llm, moderator_llm, doctor_bias,
                    patient_bias, correct_diagnosis):
    """Create an empty trajectory record for one scenario."""
    return {
        "scenario_id": scenario_id,
        "dataset": dataset,
        "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
        "models": {
            "doctor": doctor_llm,
            "patient": patient_llm,
            "measurement": measurement_llm,
            "moderator": moderator_llm,
        },
        "biases": {
            "doctor": doctor_bias,
            "patient": patient_bias,
        },
        "correct_diagnosis": correct_diagnosis,
        "turns": [],                  # list of {turn_id, role, content, turn_type}
        "final_diagnosis": None,      # what the doctor ultimately said
        "is_correct": None,           # True / False / None (gave up)
        "total_turns": 0,
        "diagnosis_ready_issued": False,
        "tests_requested": [],        # names of tests requested
    }


def add_turn(trajectory, role, content, turn_type="dialogue"):
    """Append one labelled turn to the trajectory dict.

    Args:
        trajectory: dict created by make_trajectory()
        role:       'doctor' | 'patient' | 'measurement' | 'system'
        content:    raw text of the turn
        turn_type:  'dialogue' | 'test_request' | 'test_result' | 'diagnosis'
    """
    trajectory["turns"].append({
        "turn_id": len(trajectory["turns"]),
        "role": role,
        "content": content,
        "turn_type": turn_type,
    })
    trajectory["total_turns"] = len(trajectory["turns"])


def save_trajectory(trajectory, output_dir):
    """Write trajectory to output_dir/trajectory_{scenario_id:04d}.json."""
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    out_path = Path(output_dir) / "trajectory_{:04d}.json".format(trajectory["scenario_id"])
    with open(out_path, "w") as f:
        json.dump(trajectory, f, indent=2)
    return str(out_path)


llama2_url = "meta/llama-2-70b-chat"
llama3_url = "meta/meta-llama-3-70b-instruct"
mixtral_url = "mistralai/mixtral-8x7b-instruct-v0.1"

def load_huggingface_model(model_name):
    pipe = pipeline("text-generation", model=model_name, device_map="auto")
    return pipe

def inference_huggingface(prompt, pipe):
    response = pipe(prompt, max_new_tokens=100)[0]["generated_text"]
    response = response.replace(prompt, "")
    return response


# Global override for custom OpenAI-compatible endpoints (Voyager / local vLLM).
# Set by main() when --openai_api_base is provided.
_OPENAI_API_BASE_OVERRIDE = None
_LOCAL_MODEL_NAME = None  # actual model name sent to vLLM / Voyager

# ── Dual-endpoint globals ─────────────────────────────────────────────────────
# Local Gaudi vLLM  → used by doctor + patient when --doctor_llm local
_LOCAL_API_BASE = None
_LOCAL_API_KEY  = "EMPTY"
# Voyager API       → used by agents when --*_llm voyager / voyager_lite
_VOYAGER_API_BASE       = "https://openai.rc.asu.edu/v1"
_VOYAGER_API_KEY        = None
_VOYAGER_MODEL_NAME     = None  # strong model  e.g. "qwen3-235b-a22b-instruct-2507"
_VOYAGER_LITE_MODEL_NAME = None # fast model    e.g. "qwen3-30b-a3b-instruct-2507"
# ─────────────────────────────────────────────────────────────────────────────


def _strip_thinking(text: str) -> str:
    """Remove <think>…</think> blocks produced by Qwen3 / DeepSeek-R1 style models."""
    return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()


def query_model(model_str, prompt, system_prompt, tries=5, timeout=5.0, image_requested=False, scene=None, max_prompt_len=2**14, clip_prompt=False):
    _known = ["gpt4", "gpt3.5", "gpt4o", 'llama-2-70b-chat', "mixtral-8x7b",
              "gpt-4o-mini", "gpt-4.1-mini", "llama-3-70b-instruct", "gpt4v",
              "claude3.5sonnet", "o1-preview", "local", "voyager", "voyager_lite"]
    if model_str not in _known and "HF_" not in model_str:
        raise Exception("No model by the name {}".format(model_str))
    for _attempt in range(tries):
        if clip_prompt: prompt = prompt[:max_prompt_len]
        try:
            if image_requested:
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", 
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url",
                            "image_url": {
                                "url": "{}".format(scene.image_url),
                            },
                        },
                    ]},]
                if model_str == "gpt4v":
                    response = openai.ChatCompletion.create(
                            model="gpt-4-vision-preview",
                            messages=messages,
                            temperature=0.05,
                            max_tokens=200,
                        )
                elif model_str == "gpt-4o-mini":
                    response = openai.ChatCompletion.create(
                            model="gpt-4o-mini",
                            messages=messages,
                            temperature=0.05,
                            max_tokens=200,
                        )
                elif model_str == "gpt4":
                    response = openai.ChatCompletion.create(
                            model="gpt-4-turbo",
                            messages=messages,
                            temperature=0.05,
                            max_tokens=200,
                        )
                elif model_str == "gpt4o":
                    response = openai.ChatCompletion.create(
                            model="gpt-4o",
                            messages=messages,
                            temperature=0.05,
                            max_tokens=200,
                        )
                elif model_str == "local":
                    # Route multimodal request to local vLLM VLM endpoint.
                    _model_name = _LOCAL_MODEL_NAME or "local-model"
                    _saved_base, _saved_key = openai.api_base, openai.api_key
                    try:
                        if _LOCAL_API_BASE:
                            openai.api_base = _LOCAL_API_BASE
                        openai.api_key = _LOCAL_API_KEY
                        response = openai.ChatCompletion.create(
                                model=_model_name,
                                messages=messages,
                                temperature=0.05,
                                max_tokens=200,
                            )
                    finally:
                        openai.api_base = _saved_base
                        openai.api_key = _saved_key
                elif model_str in ("voyager", "voyager_lite"):
                    # Route multimodal request to Voyager API (e.g. Llama-4-Scout).
                    _voy_model = (_VOYAGER_MODEL_NAME if model_str == "voyager"
                                  else _VOYAGER_LITE_MODEL_NAME)
                    _saved_base, _saved_key = openai.api_base, openai.api_key
                    try:
                        openai.api_base = _VOYAGER_API_BASE
                        openai.api_key = _VOYAGER_API_KEY
                        response = openai.ChatCompletion.create(
                                model=_voy_model,
                                messages=messages,
                                temperature=0.05,
                                max_tokens=200,
                            )
                    finally:
                        openai.api_base = _saved_base
                        openai.api_key = _saved_key
                else:
                    raise Exception(f"Model '{model_str}' does not support image_requested=True.")
                answer = _strip_thinking(response["choices"][0]["message"]["content"])
            if model_str == "gpt4":
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}]
                response = openai.ChatCompletion.create(
                        model="gpt-4-turbo-preview",
                        messages=messages,
                        temperature=0.05,
                        max_tokens=200,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "gpt4v":
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}]
                response = openai.ChatCompletion.create(
                        model="gpt-4-vision-preview",
                        messages=messages,
                        temperature=0.05,
                        max_tokens=200,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "gpt-4o-mini":
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}]
                response = openai.ChatCompletion.create(
                        model="gpt-4o-mini",
                        messages=messages,
                        temperature=0.05,
                        max_tokens=200,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "gpt-4.1-mini":
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}]
                response = openai.ChatCompletion.create(
                        model="gpt-4.1-mini",
                        messages=messages,
                        temperature=0.05,
                        max_tokens=200,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "o1-preview":
                messages = [
                    {"role": "user", "content": system_prompt + prompt}]
                response = openai.ChatCompletion.create(
                        model="o1-preview-2024-09-12",
                        messages=messages,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "gpt3.5":
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}]
                response = openai.ChatCompletion.create(
                        model="gpt-3.5-turbo",
                        messages=messages,
                        temperature=0.05,
                        max_tokens=200,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "claude3.5sonnet":
                client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
                message = client.messages.create(
                    model="claude-3-5-sonnet-20240620",
                    system=system_prompt,
                    max_tokens=256,
                    messages=[{"role": "user", "content": prompt}])
                answer = json.loads(message.to_json())["content"][0]["text"]
            elif model_str == "gpt4o":
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}]
                response = openai.ChatCompletion.create(
                        model="gpt-4o",
                        messages=messages,
                        temperature=0.05,
                        max_tokens=200,
                    )
                answer = response["choices"][0]["message"]["content"]
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == 'llama-2-70b-chat':
                output = replicate.run(
                    llama2_url, input={
                        "prompt":  prompt, 
                        "system_prompt": system_prompt,
                        "max_new_tokens": 200})
                answer = ''.join(output)
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == 'mixtral-8x7b':
                output = replicate.run(
                    mixtral_url, 
                    input={"prompt": prompt, 
                            "system_prompt": system_prompt,
                            "max_new_tokens": 75})
                answer = ''.join(output)
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == 'llama-3-70b-instruct':
                output = replicate.run(
                    llama3_url, input={
                        "prompt":  prompt, 
                        "system_prompt": system_prompt,
                        "max_new_tokens": 200})
                answer = ''.join(output)
                answer = re.sub(r"\s+", " ", answer)
            elif model_str == "local":
                # Route to local Gaudi vLLM endpoint.
                _model_name = _LOCAL_MODEL_NAME or "local-model"
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": prompt}]
                _saved_base, _saved_key = openai.api_base, openai.api_key
                try:
                    if _LOCAL_API_BASE:
                        openai.api_base = _LOCAL_API_BASE
                    openai.api_key = _LOCAL_API_KEY
                    response = openai.ChatCompletion.create(
                            model=_model_name,
                            messages=messages,
                            temperature=0.05,
                            max_tokens=200,
                        )
                finally:
                    openai.api_base = _saved_base
                    openai.api_key = _saved_key
                answer = _strip_thinking(response["choices"][0]["message"]["content"])
                answer = re.sub(r"\s+", " ", answer)
            elif model_str in ("voyager", "voyager_lite"):
                # Route to Voyager API endpoint (ASU RC OpenAI-compatible).
                # "voyager"      → strong model (_VOYAGER_MODEL_NAME)      e.g. doctor, moderator
                # "voyager_lite" → fast model   (_VOYAGER_LITE_MODEL_NAME) e.g. patient, measurement
                if model_str == "voyager_lite" and _VOYAGER_LITE_MODEL_NAME:
                    _model_name = _VOYAGER_LITE_MODEL_NAME
                else:
                    _model_name = _VOYAGER_MODEL_NAME or "unknown-voyager-model"
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": prompt}]
                _saved_base, _saved_key = openai.api_base, openai.api_key
                try:
                    openai.api_base = _VOYAGER_API_BASE
                    openai.api_key = _VOYAGER_API_KEY
                    response = openai.ChatCompletion.create(
                            model=_model_name,
                            messages=messages,
                            temperature=0.05,
                            max_tokens=200,
                            request_timeout=180,  # 3 min — 235B model can be slow
                        )
                finally:
                    openai.api_base = _saved_base
                    openai.api_key = _saved_key
                answer = _strip_thinking(response["choices"][0]["message"]["content"])
                answer = re.sub(r"\s+", " ", answer)
            elif "HF_" in model_str:
                input_text = system_prompt + prompt
                raise Exception("HuggingFace local inference not yet implemented. Use 'local' + --openai_api_base with vLLM instead.")
            return answer

        except Exception as e:
            print(f"[query_model:{model_str}] attempt {_attempt+1}/{tries} failed — {type(e).__name__}: {e}", flush=True)
            time.sleep(timeout)
            continue
    raise Exception(f"Max retries exhausted for model '{model_str}' after {tries} attempts")



class ScenarioMedQA:
    def __init__(self, scenario_dict) -> None:
        self.scenario_dict = scenario_dict
        self.tests = scenario_dict["OSCE_Examination"]["Test_Results"]
        self.diagnosis = scenario_dict["OSCE_Examination"]["Correct_Diagnosis"]
        self.patient_info  = scenario_dict["OSCE_Examination"]["Patient_Actor"]
        self.examiner_info  = scenario_dict["OSCE_Examination"]["Objective_for_Doctor"]
        self.physical_exams = scenario_dict["OSCE_Examination"]["Physical_Examination_Findings"]
    
    def patient_information(self) -> dict:
        return self.patient_info

    def examiner_information(self) -> dict:
        return self.examiner_info
    
    def exam_information(self) -> dict:
        exams = self.physical_exams
        exams["tests"] = self.tests
        return exams
    
    def diagnosis_information(self) -> dict:
        return self.diagnosis


class ScenarioLoaderMedQA:
    def __init__(self) -> None:
        with open("agentclinic_medqa.jsonl", "r") as f:
            self.scenario_strs = [json.loads(line) for line in f]
        self.scenarios = [ScenarioMedQA(_str) for _str in self.scenario_strs]
        self.num_scenarios = len(self.scenarios)
    
    def sample_scenario(self):
        return self.scenarios[random.randint(0, len(self.scenarios)-1)]
    
    def get_scenario(self, id):
        if id is None: return self.sample_scenario()
        return self.scenarios[id]
        


class ScenarioMedQAExtended:
    def __init__(self, scenario_dict) -> None:
        self.scenario_dict = scenario_dict
        self.tests = scenario_dict["OSCE_Examination"]["Test_Results"]
        self.diagnosis = scenario_dict["OSCE_Examination"]["Correct_Diagnosis"]
        self.patient_info  = scenario_dict["OSCE_Examination"]["Patient_Actor"]
        self.examiner_info  = scenario_dict["OSCE_Examination"]["Objective_for_Doctor"]
        self.physical_exams = scenario_dict["OSCE_Examination"]["Physical_Examination_Findings"]
    
    def patient_information(self) -> dict:
        return self.patient_info

    def examiner_information(self) -> dict:
        return self.examiner_info
    
    def exam_information(self) -> dict:
        exams = self.physical_exams
        exams["tests"] = self.tests
        return exams
    
    def diagnosis_information(self) -> dict:
        return self.diagnosis


class ScenarioLoaderMedQAExtended:
    def __init__(self) -> None:
        with open("agentclinic_medqa_extended.jsonl", "r") as f:
            self.scenario_strs = [json.loads(line) for line in f]
        self.scenarios = [ScenarioMedQAExtended(_str) for _str in self.scenario_strs]
        self.num_scenarios = len(self.scenarios)
    
    def sample_scenario(self):
        return self.scenarios[random.randint(0, len(self.scenarios)-1)]
    
    def get_scenario(self, id):
        if id is None: return self.sample_scenario()
        return self.scenarios[id]
        


class ScenarioMIMICIVQA:
    def __init__(self, scenario_dict) -> None:
        self.scenario_dict = scenario_dict
        self.tests = scenario_dict["OSCE_Examination"]["Test_Results"]
        self.diagnosis = scenario_dict["OSCE_Examination"]["Correct_Diagnosis"]
        self.patient_info  = scenario_dict["OSCE_Examination"]["Patient_Actor"]
        self.examiner_info  = scenario_dict["OSCE_Examination"]["Objective_for_Doctor"]
        self.physical_exams = scenario_dict["OSCE_Examination"]["Physical_Examination_Findings"]
    
    def patient_information(self) -> dict:
        return self.patient_info

    def examiner_information(self) -> dict:
        return self.examiner_info
    
    def exam_information(self) -> dict:
        exams = self.physical_exams
        exams["tests"] = self.tests
        return exams
    
    def diagnosis_information(self) -> dict:
        return self.diagnosis


class ScenarioLoaderMIMICIV:
    def __init__(self) -> None:
        with open("agentclinic_mimiciv.jsonl", "r") as f:
            self.scenario_strs = [json.loads(line) for line in f]
        self.scenarios = [ScenarioMIMICIVQA(_str) for _str in self.scenario_strs]
        self.num_scenarios = len(self.scenarios)
    
    def sample_scenario(self):
        return self.scenarios[random.randint(0, len(self.scenarios)-1)]
    
    def get_scenario(self, id):
        if id is None: return self.sample_scenario()
        return self.scenarios[id]


class ScenarioNEJMExtended:
    def __init__(self, scenario_dict) -> None:
        self.scenario_dict = scenario_dict 
        self.question = scenario_dict["question"] 
        self.image_url = scenario_dict["image_url"] 
        self.diagnosis = [_sd["text"] 
            for _sd in scenario_dict["answers"] if _sd["correct"]][0]
        self.patient_info = scenario_dict["patient_info"]
        self.physical_exams = scenario_dict["physical_exams"]

    def patient_information(self) -> str:
        patient_info = self.patient_info
        return patient_info

    def examiner_information(self) -> str:
        return "What is the most likely diagnosis?"
    
    def exam_information(self) -> str:
        exams = self.physical_exams
        return exams
    
    def diagnosis_information(self) -> str:
        return self.diagnosis


class ScenarioLoaderNEJMExtended:
    def __init__(self) -> None:
        with open("agentclinic_nejm_extended.jsonl", "r") as f:
            self.scenario_strs = [json.loads(line) for line in f]
        self.scenarios = [ScenarioNEJMExtended(_str) for _str in self.scenario_strs]
        self.num_scenarios = len(self.scenarios)
    
    def sample_scenario(self):
        return self.scenarios[random.randint(0, len(self.scenarios)-1)]
    
    def get_scenario(self, id):
        if id is None: return self.sample_scenario()
        return self.scenarios[id]


class ScenarioNEJM:
    def __init__(self, scenario_dict) -> None:
        self.scenario_dict = scenario_dict 
        self.question = scenario_dict["question"] 
        self.image_url = scenario_dict["image_url"] 
        self.diagnosis = [_sd["text"] 
            for _sd in scenario_dict["answers"] if _sd["correct"]][0]
        self.patient_info = scenario_dict["patient_info"]
        self.physical_exams = scenario_dict["physical_exams"]

    def patient_information(self) -> str:
        patient_info = self.patient_info
        return patient_info

    def examiner_information(self) -> str:
        return "What is the most likely diagnosis?"
    
    def exam_information(self) -> str:
        exams = self.physical_exams
        return exams
    
    def diagnosis_information(self) -> str:
        return self.diagnosis


class ScenarioLoaderNEJM:
    def __init__(self) -> None:
        with open("agentclinic_nejm.jsonl", "r") as f:
            self.scenario_strs = [json.loads(line) for line in f]
        self.scenarios = [ScenarioNEJM(_str) for _str in self.scenario_strs]
        self.num_scenarios = len(self.scenarios)
    
    def sample_scenario(self):
        return self.scenarios[random.randint(0, len(self.scenarios)-1)]
    
    def get_scenario(self, id):
        if id is None: return self.sample_scenario()
        return self.scenarios[id]


class PatientAgent:
    def __init__(self, scenario, backend_str="gpt4", bias_present=None) -> None:
        # disease of patient, or "correct answer"
        self.disease = ""
        # symptoms that patient presents
        self.symptoms = ""
        # conversation history between doctor and patient
        self.agent_hist = ""
        # language model backend for patient agent
        self.backend = backend_str
        # presentation of any form of bias
        self.bias_present = (None if bias_present == "None" else bias_present)
        # sample initial question from dataset
        self.scenario = scenario
        self.reset()
        self.pipe = None

        self.biases = ["recency", "frequency", "false_consensus", "self_diagnosis", "gender", "race", "sexual_orientation", "cultural", "education", "religion", "socioeconomic"]

    def generate_bias(self) -> str:
        """ 
        ================
        Cognitive biases 
        ================
        """
        if self.bias_present == "recency":
            return "\nRecently, you had a friend with similar symptoms to yours that was diagnosed with cancer. This affects how you interact with the doctor.\n"
        elif self.bias_present == "frequency":
            return "\nYou recently learned that most frequently, doctors diagnose this as cancer based on past cases. This affects how you interact with the doctor.\n"
        elif self.bias_present == "false_consensus":
            return "\nMost of your friends who heard about your symptoms believe that you have cancer. This affects how you interact with the doctor.\n"
        elif self.bias_present == "self_diagnosis":
            return "\nYou recently performed a thorough investigation of your symptoms online and have come to the conclusion that you should be diagnosed with cancer. While the doctor may believe otherwise you remain convinced you have cancer. This affects how you interact with the doctor.\n"
        elif self.bias_present == "gender":
            return "\nYou are uncomfortable with your doctor because you find out that they are a particular gender and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "race":
            return "\nYou are uncomfortable with your doctor because you find out that they are a particular race and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "sexual_orientation":
            return "\nYou are uncomfortable with your doctor because you find out that they are a particular sexual orientation and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "cultural":
            return "\nYou are uncomfortable with your doctor because you find out that they are from a particular culture and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "education":
            return "\nYou are uncomfortable with your doctor because you find out that they went to a low ranked medical school and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "religion":
            return "\nYou are uncomfortable with your doctor because you find out that they believe in a particular religion and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "socioeconomic":
            return "\nYou are uncomfortable with your doctor because you find out that they are from a particular socioeconomic background and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present is None:
            pass
        else:
            print("BIAS TYPE {} NOT SUPPORTED, ignoring bias...".format(self.bias_present))
        return ""

    def inference_patient(self, question) -> str:
        answer = query_model(self.backend, "\nHere is a history of your dialogue: " + self.agent_hist + "\n Here was the doctor response: " + question + "Now please continue your dialogue\nPatient: ", self.system_prompt())
        self.agent_hist += question + "\n\n" + answer + "\n\n"
        return answer

    def system_prompt(self) -> str:
        bias_prompt = ""
        if self.bias_present is not None:
            bias_prompt = self.generate_bias()
        base = """You are a patient in a clinic who only responds in the form of dialogue. You are being inspected by a doctor who will ask you questions and will perform exams on you in order to understand your disease. Your answer will only be 1-3 sentences in length."""
        symptoms = "\n\nBelow is all of your information. {}. \n\n Remember, you must not reveal your disease explicitly but may only convey the symptoms you have in the form of dialogue if you are asked.".format(self.symptoms)
        return base + bias_prompt + symptoms
    
    def reset(self) -> None:
        self.agent_hist = ""
        self.symptoms = self.scenario.patient_information()

    def add_hist(self, hist_str) -> None:
        self.agent_hist += hist_str + "\n\n"


class DoctorAgent:
    def __init__(self, scenario, backend_str="gpt4", max_infs=20, bias_present=None, img_request=False) -> None:
        # number of inference calls to the doctor
        self.infs = 0
        # maximum number of inference calls to the doctor
        self.MAX_INFS = max_infs
        # conversation history between doctor and patient
        self.agent_hist = ""
        # presentation information for doctor
        self.presentation = ""
        # language model backend for doctor agent
        self.backend = backend_str
        # presentation of any form of bias
        self.bias_present = (None if bias_present == "None" else bias_present)
        # prepare initial conditions for LLM
        self.scenario = scenario
        self.reset()
        self.pipe = None
        self.img_request = img_request
        self.biases = ["recency", "frequency", "false_consensus", "confirmation", "status_quo", "gender", "race", "sexual_orientation", "cultural", "education", "religion", "socioeconomic"]

    def generate_bias(self) -> str:
        """ 
        ================
        Cognitive biases 
        ================
        """
        if self.bias_present == "recency":
            return "\nRecently, you learned that a doctor with similar symptoms to the patient that diagnosed their patient with cancer. This affects how you interact with the patient.\n"
        elif self.bias_present == "frequency":
            return "\nYou recently learned that most frequently, doctors diagnose patients with similar symptoms with cancer based on past cases. This affects how you interact with the patient.\n"
        elif self.bias_present == "false_consensus":
            return "\nMost of your colleagues who heard about the patient's symptoms believe that they have cancer. This affects how you interact with the patient.\n"
        elif self.bias_present == "status_quo":
            return "\nYour hosptial has been diagnosing these patient's symptoms with cancer for the past 10 years. This affects how you interact with the patient.\n"
        elif self.bias_present == "confirmation":
            return "\nYou are initially confident that the patient has cancer. This affects how you interact with the patient.\n"
        elif self.bias_present == "gender":
            return "\nYou are uncomfortable with your patient because you find out that they are a particular gender and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "race":
            return "\nYou are uncomfortable with your patient because you find out that they are a particular race and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "sexual_orientation":
            return "\nYou are uncomfortable with your patient because you find out that they are a particular sexual orientation and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "cultural":
            return "\nYou are uncomfortable with your patient because you find out that they are from a particular culture and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "education":
            return "\nYou are uncomfortable with your patient because you find out that they are uneducated and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "religion":
            return "\nYou are uncomfortable with your patient because you find out that they believe in a particular religion and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present == "socioeconomic":
            return "\nYou are uncomfortable with your patient because you find out that they are from a particular socioeconomic background and you do not trust their judgement. This affects how you interact with them.\n"
        elif self.bias_present is None:
            pass
        else:
            print("BIAS TYPE {} NOT SUPPORTED, ignoring bias...".format(self.bias_present))
        return ""

    def inference_doctor(self, question, image_requested=False) -> str:
        answer = str()
        if self.infs >= self.MAX_INFS: return "Maximum inferences reached"
        answer = query_model(self.backend, "\nHere is a history of your dialogue: " + self.agent_hist + "\n Here was the patient response: " + question + "Now please continue your dialogue\nDoctor: ", self.system_prompt(), image_requested=image_requested, scene=self.scenario)
        self.agent_hist += question + "\n\n" + answer + "\n\n"
        self.infs += 1
        return answer

    def system_prompt(self) -> str:
        bias_prompt = ""
        if self.bias_present is not None:
            bias_prompt = self.generate_bias()
        base = "You are a doctor named Dr. Agent who only responds in the form of dialogue. You are inspecting a patient who you will ask questions in order to understand their disease. You are only allowed to ask {} questions total before you must make a decision. You have asked {} questions so far. You can request test results using the format \"REQUEST TEST: [test]\". For example, \"REQUEST TEST: Chest_X-Ray\". Your dialogue will only be 1-3 sentences in length. Once you have decided to make a diagnosis please type \"DIAGNOSIS READY: [diagnosis here]\"".format(self.MAX_INFS, self.infs) + ("You may also request medical images related to the disease to be returned with \"REQUEST IMAGES\"." if self.img_request else "")
        presentation = "\n\nBelow is all of the information you have. {}. \n\n Remember, you must discover their disease by asking them questions. You are also able to provide exams.".format(self.presentation)
        return base + bias_prompt + presentation

    def reset(self) -> None:
        self.agent_hist = ""
        self.presentation = self.scenario.examiner_information()


class MeasurementAgent:
    def __init__(self, scenario, backend_str="gpt4") -> None:
        # conversation history between doctor and patient
        self.agent_hist = ""
        # presentation information for measurement 
        self.presentation = ""
        # language model backend for measurement agent
        self.backend = backend_str
        # prepare initial conditions for LLM
        self.scenario = scenario
        self.pipe = None
        self.reset()

    def inference_measurement(self, question) -> str:
        answer = str()
        answer = query_model(self.backend, "\nHere is a history of the dialogue: " + self.agent_hist + "\n Here was the doctor measurement request: " + question, self.system_prompt())
        self.agent_hist += question + "\n\n" + answer + "\n\n"
        return answer

    def system_prompt(self) -> str:
        base = "You are an measurement reader who responds with medical test results. Please respond in the format \"RESULTS: [results here]\""
        presentation = "\n\nBelow is all of the information you have. {}. \n\n If the requested results are not in your data then you can respond with NORMAL READINGS.".format(self.information)
        return base + presentation
    
    def add_hist(self, hist_str) -> None:
        self.agent_hist += hist_str + "\n\n"

    def reset(self) -> None:
        self.agent_hist = ""
        self.information = self.scenario.exam_information()


def compare_results(diagnosis, correct_diagnosis, moderator_llm, mod_pipe):
    answer = query_model(moderator_llm, "\nHere is the correct diagnosis: " + correct_diagnosis + "\n Here was the doctor dialogue: " + diagnosis + "\nAre these the same?", "You are responsible for determining if the corrent diagnosis and the doctor diagnosis are the same disease. Please respond only with Yes or No. Nothing else.")
    # Normalize: take first word, strip punctuation — handles "Yes.", "Yes!", "Yes, ..." etc.
    first_word = re.sub(r'[^a-z]', '', answer.lower().split()[0]) if answer.strip() else "no"
    return first_word


def main(api_key, replicate_api_key, inf_type, doctor_bias, patient_bias, doctor_llm, patient_llm,
         measurement_llm, moderator_llm, num_scenarios, dataset, img_request, total_inferences,
         anthropic_api_key=None, output_dir="./trajectories",
         openai_api_base=None, local_model_name=None, voyager_api_key=None, voyager_api_base=None,
         voyager_model_name=None, voyager_lite_model_name=None):
    global _OPENAI_API_BASE_OVERRIDE, _LOCAL_MODEL_NAME, \
           _LOCAL_API_BASE, _LOCAL_API_KEY, \
           _VOYAGER_API_BASE, _VOYAGER_API_KEY, _VOYAGER_MODEL_NAME, _VOYAGER_LITE_MODEL_NAME

    # ── Local vLLM endpoint (doctor + patient when --*_llm local) ─────────────
    if openai_api_base:
        _LOCAL_API_BASE = openai_api_base
        _OPENAI_API_BASE_OVERRIDE = openai_api_base
    if local_model_name:
        _LOCAL_MODEL_NAME = local_model_name

    # ── Voyager endpoint (measurement + moderator when --*_llm voyager) ────────
    if voyager_api_key:
        _VOYAGER_API_KEY = voyager_api_key
    if voyager_api_base:
        _VOYAGER_API_BASE = voyager_api_base
    if voyager_model_name:
        _VOYAGER_MODEL_NAME = voyager_model_name
    if voyager_lite_model_name:
        _VOYAGER_LITE_MODEL_NAME = voyager_lite_model_name

    # ── Default OpenAI config (for gpt4/gpt4o/etc. if directly selected) ──────
    openai.api_key = api_key or "EMPTY"
    # ──────────────────────────────────────────────────────────────────────────

    anthropic_llms = ["claude3.5sonnet"]
    replicate_llms = ["llama-3-70b-instruct", "llama-2-70b-chat", "mixtral-8x7b"]
    if patient_llm in replicate_llms or doctor_llm in replicate_llms:
        os.environ["REPLICATE_API_TOKEN"] = replicate_api_key
    if doctor_llm in anthropic_llms:
        os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key

    # Load MedQA, MIMICIV or NEJM agent case scenarios
    if dataset == "MedQA":
        scenario_loader = ScenarioLoaderMedQA()
    elif dataset == "MedQA_Ext":
        scenario_loader = ScenarioLoaderMedQAExtended()
    elif dataset == "NEJM":
        scenario_loader = ScenarioLoaderNEJM()
    elif dataset == "NEJM_Ext":
        scenario_loader = ScenarioLoaderNEJMExtended()
    elif dataset == "MIMICIV":
        scenario_loader = ScenarioLoaderMIMICIV()
    else:
        raise Exception("Dataset {} does not exist".format(str(dataset)))
    total_correct = 0
    total_presents = 0

    # Pipeline for huggingface models
    if "HF_" in moderator_llm:
        pipe = load_huggingface_model(moderator_llm.replace("HF_", ""))
    else:
        pipe = None
    if num_scenarios is None: num_scenarios = scenario_loader.num_scenarios
    for _scenario_id in range(0, min(num_scenarios, scenario_loader.num_scenarios)):
        total_presents += 1
        pi_dialogue = str()
        _scenario_start = time.time()
        # Initialize scenarios (MedQA/NEJM)
        scenario = scenario_loader.get_scenario(id=_scenario_id)

        # ── Trajectory setup ──────────────────────────────────────────────────
        traj = make_trajectory(
            scenario_id=_scenario_id,
            dataset=dataset,
            doctor_llm=doctor_llm,
            patient_llm=patient_llm,
            measurement_llm=measurement_llm,
            moderator_llm=moderator_llm,
            doctor_bias=doctor_bias,
            patient_bias=patient_bias,
            correct_diagnosis=str(scenario.diagnosis_information()),
        )
        # ─────────────────────────────────────────────────────────────────────

        # Initialize agents
        meas_agent = MeasurementAgent(
            scenario=scenario,
            backend_str=measurement_llm)
        patient_agent = PatientAgent(
            scenario=scenario, 
            bias_present=patient_bias,
            backend_str=patient_llm)
        doctor_agent = DoctorAgent(
            scenario=scenario, 
            bias_present=doctor_bias,
            backend_str=doctor_llm,
            max_infs=total_inferences, 
            img_request=img_request)

        doctor_dialogue = ""
        for _inf_id in range(total_inferences):
            # Check for medical image request
            if dataset == "NEJM":
                if img_request:
                    imgs = "REQUEST IMAGES" in doctor_dialogue
                else: imgs = True
            else: imgs = False
            # Check if final inference
            if _inf_id == total_inferences - 1:
                pi_dialogue += "This is the final question. Please provide a diagnosis.\n"
            # Obtain doctor dialogue (human or llm agent)
            if inf_type == "human_doctor":
                doctor_dialogue = input("\nQuestion for patient: ")
            else: 
                doctor_dialogue = doctor_agent.inference_doctor(pi_dialogue, image_requested=imgs)
            print("Doctor [{}%]:".format(int(((_inf_id+1)/total_inferences)*100)), doctor_dialogue)

            # ── Record doctor turn ────────────────────────────────────────────
            if "DIAGNOSIS READY" in doctor_dialogue:
                turn_type = "diagnosis"
            elif "REQUEST TEST" in doctor_dialogue:
                turn_type = "test_request"
                # Extract test name for summary
                try:
                    test_name = doctor_dialogue.split("REQUEST TEST:")[1].split("\n")[0].strip()
                    traj["tests_requested"].append(test_name)
                except Exception:
                    pass
            else:
                turn_type = "dialogue"
            add_turn(traj, role="doctor", content=doctor_dialogue, turn_type=turn_type)
            # ─────────────────────────────────────────────────────────────────

            # Doctor has arrived at a diagnosis, check correctness
            if "DIAGNOSIS READY" in doctor_dialogue:
                correctness = compare_results(doctor_dialogue, scenario.diagnosis_information(), moderator_llm, pipe) == "yes"
                if correctness: total_correct += 1
                print("\nCorrect answer:", scenario.diagnosis_information())
                print("Scene {}, The diagnosis was ".format(_scenario_id), "CORRECT" if correctness else "INCORRECT", int((total_correct/total_presents)*100))
                # ── Finalize trajectory ───────────────────────────────────────
                traj["diagnosis_ready_issued"] = True
                traj["final_diagnosis"] = doctor_dialogue
                traj["is_correct"] = bool(correctness)
                # ─────────────────────────────────────────────────────────────
                break
            # Obtain medical exam from measurement reader
            if "REQUEST TEST" in doctor_dialogue:
                pi_dialogue = meas_agent.inference_measurement(doctor_dialogue,)
                print("Measurement [{}%]:".format(int(((_inf_id+1)/total_inferences)*100)), pi_dialogue)
                patient_agent.add_hist(pi_dialogue)
                # ── Record measurement turn ───────────────────────────────────
                add_turn(traj, role="measurement", content=pi_dialogue, turn_type="test_result")
                # ─────────────────────────────────────────────────────────────
            # Obtain response from patient
            else:
                if inf_type == "human_patient":
                    pi_dialogue = input("\nResponse to doctor: ")
                else:
                    pi_dialogue = patient_agent.inference_patient(doctor_dialogue)
                print("Patient [{}%]:".format(int(((_inf_id+1)/total_inferences)*100)), pi_dialogue)
                meas_agent.add_hist(pi_dialogue)
                # ── Record patient turn ───────────────────────────────────────
                add_turn(traj, role="patient", content=pi_dialogue, turn_type="dialogue")
                # ─────────────────────────────────────────────────────────────
            # Prevent API timeouts
            time.sleep(1.0)

        # ── Save trajectory for this scenario ─────────────────────────────────
        if traj["is_correct"] is None:  # Doctor never issued DIAGNOSIS READY
            traj["final_diagnosis"] = "(Doctor did not issue DIAGNOSIS READY)"
            traj["is_correct"] = False
        _scenario_elapsed = time.time() - _scenario_start
        traj["scenario_duration_seconds"] = round(_scenario_elapsed, 1)
        saved_path = save_trajectory(traj, output_dir)
        print(f"[Trajectory saved → {saved_path}]  [{_scenario_elapsed:.1f}s for scenario {_scenario_id}]  [Running acc: {int((total_correct/total_presents)*100)}%]")
        # ─────────────────────────────────────────────────────────────────────


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Medical Diagnosis Simulation CLI')
    parser.add_argument('--openai_api_key', type=str, required=False, help='OpenAI API Key')
    parser.add_argument('--replicate_api_key', type=str, required=False, help='Replicate API Key')
    parser.add_argument('--inf_type', type=str, choices=['llm', 'human_doctor', 'human_patient'], default='llm')
    parser.add_argument('--doctor_bias', type=str, help='Doctor bias type', default='None', choices=["None", "recency", "frequency", "false_consensus", "confirmation", "status_quo", "gender", "race", "sexual_orientation", "cultural", "education", "religion", "socioeconomic"])
    parser.add_argument('--patient_bias', type=str, help='Patient bias type', default='None', choices=["None", "recency", "frequency", "false_consensus", "self_diagnosis", "gender", "race", "sexual_orientation", "cultural", "education", "religion", "socioeconomic"])
    parser.add_argument('--doctor_llm', type=str, default='gpt4')
    parser.add_argument('--patient_llm', type=str, default='gpt4')
    parser.add_argument('--measurement_llm', type=str, default='gpt4')
    parser.add_argument('--moderator_llm', type=str, default='gpt4')
    parser.add_argument('--agent_dataset', type=str, default='MedQA') # MedQA, MIMICIV or NEJM
    parser.add_argument('--doctor_image_request', type=bool, default=False) # whether images must be requested or are provided
    parser.add_argument('--num_scenarios', type=int, default=None, required=False, help='Number of scenarios to simulate')
    parser.add_argument('--total_inferences', type=int, default=20, required=False, help='Number of inferences between patient and doctor')
    parser.add_argument('--anthropic_api_key',  type=str, default=None, required=False, help='Anthropic API key for Claude 3.5 Sonnet')
    parser.add_argument('--output_dir',          type=str, default='./trajectories', required=False, help='Directory to save per-scenario JSON trajectory files')
    # ── Custom API endpoint arguments (Voyager / local vLLM on Gaudi) ─────────
    parser.add_argument('--openai_api_base',     type=str, default=None, required=False, help='Override OpenAI API base URL (e.g. http://127.0.0.1:8000/v1 for local vLLM)')
    parser.add_argument('--local_model_name',    type=str, default=None, required=False, help='Model name to pass to local vLLM server when using --doctor_llm local / --patient_llm local')
    parser.add_argument('--voyager_api_key',     type=str, default=None, required=False, help='Voyager API key (ASU RC OpenAI-compatible endpoint)')
    parser.add_argument('--voyager_api_base',    type=str, default='https://openai.rc.asu.edu/v1', required=False, help='Voyager API base URL')
    parser.add_argument('--voyager_model_name',       type=str, default=None, required=False, help='Strong Voyager model for doctor + moderator (e.g. qwen3-235b-a22b-instruct-2507)')
    parser.add_argument('--voyager_lite_model_name',  type=str, default=None, required=False, help='Fast Voyager model for patient + measurement (e.g. qwen3-30b-a3b-instruct-2507)')
    # ──────────────────────────────────────────────────────────────────────────

    args = parser.parse_args()

    main(args.openai_api_key, args.replicate_api_key, args.inf_type, args.doctor_bias, args.patient_bias,
         args.doctor_llm, args.patient_llm, args.measurement_llm, args.moderator_llm,
         args.num_scenarios, args.agent_dataset, args.doctor_image_request, args.total_inferences,
         args.anthropic_api_key, args.output_dir,
         args.openai_api_base, args.local_model_name, args.voyager_api_key, args.voyager_api_base,
         args.voyager_model_name, args.voyager_lite_model_name)