CodeHeader = """from tools import tabtools, calculator
Calculate = calculator.WolframAlphaCalculator
LoadDB = tabtools.db_loader
FilterDB = tabtools.data_filter
GetValue = tabtools.get_value
SQLInterpreter = tabtools.sql_interpreter
Calendar = tabtools.date_calculator
"""

RetrKnowledge_example = """Read the following data descriptions, generate the background knowledge as the context information that could be helpful for answering the question.
(1) Data include vital signs, laboratory measurements, medications, APACHE components, care plan information, admission diagnosis, patient history, time-stamped diagnoses from a structured problem list, and similarly chosen treatments.
(2) Data from each patient is collected into a common warehouse only if certain “interfaces” are available. Each interface is used to transform and load a certain type of data: vital sign interfaces incorporate vital signs, laboratory interfaces provide measurements on blood samples, and so on. 
(3) It is important to be aware that different care units may have different interfaces in place, and that the lack of an interface will result in no data being available for a given patient, even if those measurements were made in reality. The data is provided as a relational database, comprising multiple tables joined by keys.
(4) All the databases are used to record information associated to patient care, such as allergy, cost, diagnosis, intakeoutput, lab, medication, microlab, patient, treatment, vitalperiodic.
For different tables, they contain the following information:
(1) allergy: allergyid, patientunitstayid, drugname, allergyname, allergytime
(2) cost: costid, uniquepid, patienthealthsystemstayid, eventtype, eventid, chargetime, cost
(3) diagnosis: diagnosisid, patientunitstayid, icd9code, diagnosisname, diagnosistime
(4) intakeoutput: intakeoutputid, patientunitstayid, cellpath, celllabel, cellvaluenumeric, intakeoutputtime
(5) lab: labid, patientunitstayid, labname, labresult, labresulttime
(6) medication: medicationid, patientunitstayid, drugname, dosage, routeadmin, drugstarttime, drugstoptime
(7) microlab: microlabid, patientunitstayid, culturesite, organism, culturetakentime
(8) patient: patientunitstayid, patienthealthsystemstayid, gender, age, ethnicity, hospitalid, wardid, admissionheight, hospitaladmitsource, hospitaldischargestatus, admissionweight, dischargeweight, uniquepid, hospitaladmittime, unitadmittime, unitdischargetime, hospitaldischargetime
(9) treatment: treatmentid, patientunitstayid, treatmentname, treatmenttime
(10) vitalperiodic: vitalperiodicid, patientunitstayid, temperature, sao2, heartrate, respiration, systemicsystolic, systemicdiastolic, systemicmean, observationtime

{examples}

Question: {question}
Knowledge:
"""

RetrKnowledge = """Read the following data descriptions, generate the background knowledge as the context information that could be helpful for answering the question.
(1) Data include vital signs, laboratory measurements, medications, APACHE components, care plan information, admission diagnosis, patient history, time-stamped diagnoses from a structured problem list, and similarly chosen treatments.
(2) Data from each patient is collected into a common warehouse only if certain “interfaces” are available. Each interface is used to transform and load a certain type of data: vital sign interfaces incorporate vital signs, laboratory interfaces provide measurements on blood samples, and so on. 
(3) It is important to be aware that different care units may have different interfaces in place, and that the lack of an interface will result in no data being available for a given patient, even if those measurements were made in reality. The data is provided as a relational database, comprising multiple tables joined by keys.
(4) All the databases are used to record information associated to patient care, such as allergy, cost, diagnosis, intakeoutput, lab, medication, microlab, patient, treatment, vitalperiodic.
For different tables, they contain the following information:
(1) allergy: allergyid, patientunitstayid, drugname, allergyname, allergytime
(2) cost: costid, uniquepid, patienthealthsystemstayid, eventtype, eventid, chargetime, cost
(3) diagnosis: diagnosisid, patientunitstayid, icd9code, diagnosisname, diagnosistime
(4) intakeoutput: intakeoutputid, patientunitstayid, cellpath, celllabel, cellvaluenumeric, intakeoutputtime
(5) lab: labid, patientunitstayid, labname, labresult, labresulttime
(6) medication: medicationid, patientunitstayid, drugname, dosage, routeadmin, drugstarttime, drugstoptime
(7) microlab: microlabid, patientunitstayid, culturesite, organism, culturetakentime
(8) patient: patientunitstayid, patienthealthsystemstayid, gender, age, ethnicity, hospitalid, wardid, admissionheight, hospitaladmitsource, hospitaldischargestatus, admissionweight, dischargeweight, uniquepid, hospitaladmittime, unitadmittime, unitdischargetime, hospitaldischargetime
(9) treatment: treatmentid, patientunitstayid, treatmentname, treatmenttime
(10) vitalperiodic: vitalperiodicid, patientunitstayid, temperature, sao2, heartrate, respiration, systemicsystolic, systemicdiastolic, systemicmean, observationtime

Question: was the fluticasone-salmeterol 250-50 mcg/dose in aepb prescribed to patient 035-2205 on their current hospital encounter?
Knowledge:
- We can find the patient 035-2205 information in the patient database.
- As fluticasone-salmeterol 250-50 mcg/dose in aepb is a drug, we can find the drug information in the medication database.
- We can find the patientunitstayid in the patient database and use it to find the drug precsription information in the medication database.

Question: in the last hospital encounter, when was patient 031-22988's first microbiology test time?
Knowledge:
- We can find the patient 031-22988 information in the patient database.
- We can find the microbiology test information in the microlab database.
- We can find the patientunitstayid in the patient database and use it to find the microbiology test information in the microlab database.

Question: what is the minimum hospital cost for a drug with a name called albumin 5% since 6 years ago?
Knowledge:
- As albumin 5% is a drug, we can find the drug information in the medication database.
- We can find the patientunitstayid in the medication database and use it to find the patienthealthsystemstayid information in the patient database.
- We can use the patienthealthsystemstayid information to find the cost information in the cost database.

Question: what are the number of patients who have had a magnesium test the previous year?
Knowledge:
- As magnesium is a lab test, we can find the lab test information in the lab database.
- We can find the patientunitstayid in the lab database and use it to find the patient information in the patient database.

Question: {question}
Knowledge:
"""

SYSTEM_PROMPT = """You are a helpful AI assistant. Solve tasks using your coding and language skills.
In the following cases, suggest python code (in a python coding block) or shell script (in a sh
coding block) for the user to execute.
1. When you need to collect info, use the code to output the info you need, for example, browse or
search the web, download/read a file, print the content of a webpage or a file, get the current
date/time. After sufficient info is printed and the task is ready to be solved based on your
language skill, you can solve the task by yourself.
2. When you need to perform some task with code, use the code to perform the task and output the
result. Finish the task smartly.
Solve the task step by step if you need to. If a plan is not provided, explain your plan first. Be
clear which step uses code, and which step uses your language skill.
When using code, you must indicate the script type in the code block. The user cannot provide any
other feedback or perform any other action beyond executing the code you suggest. The user can't
modify your code. So do not suggest incomplete code which requires users to modify. Don't use a
code block if it's not intended to be executed by the user.
If you want the user to save the code in a file before executing it, put # filename: <filename>
inside the code block as the first line. Don't include multiple code blocks in one response. Do not
ask users to copy and paste the result. Instead, use 'print' function for the output when relevant.
Check the execution result returned by the user.
If the result indicates there is an error, fix the error and output the code again. Suggest the
full code instead of partial code or code changes. If the error can't be fixed or if the task is
not solved even after the code is executed successfully, analyze the problem, revisit your
assumption, collect additional info you need, and think of a different approach to try.
When you find an answer, verify the answer carefully. Include verifiable evidence in your response
if possible.
Reply "TERMINATE" in the end when everything is done."""

EHRAgent_Message_Prompt = """Assume you have knowledge of several tables:
(1) Tables are linked by identifiers which usually have the suffix 'ID'. For example, SUBJECT_ID refers to a unique patient, HADM_ID refers to a unique admission to the hospital, and ICUSTAY_ID refers to a unique admission to an intensive care unit.
(2) Charted events such as notes, laboratory tests, and fluid balance are stored in a series of 'events' tables. For example the outputevents table contains all measurements related to output for a given patient, while the labevents table contains laboratory test results for a patient.
(3) Tables prefixed with 'd_' are dictionary tables and provide definitions for identifiers. For example, every row of chartevents is associated with a single ITEMID which represents the concept measured, but it does not contain the actual name of the measurement. By joining chartevents and d_items on ITEMID, it is possible to identify the concept represented by a given ITEMID.
(4) For the databases, four of them are used to define and track patient stays: admissions, patients, icustays, and transfers. Another four tables are dictionaries for cross-referencing codes against their respective definitions: d_icd_diagnoses, d_icd_procedures, d_items, and d_labitems. The remaining tables, including chartevents, cost, inputevents_cv, labevents, microbiologyevents, outputevents, prescriptions, procedures_icd, contain data associated with patient care, such as physiological measurements, caregiver observations, and billing information.
Write a python code to solve the given question. You can use the following functions:
(1) Calculate(FORMULA), which calculates the FORMULA and returns the result.
(2) LoadDB(DBNAME) which loads the database DBNAME and returns the database. The DBNAME can be one of the following: allergy, cost, diagnosis, intakeoutput, lab, medication, microlab, patient, treatment, vitalperiodic.
(3) FilterDB(DATABASE, CONDITIONS), which filters the DATABASE according to the CONDITIONS and returns the filtered database. The CONDITIONS is a string composed of multiple conditions, each of which consists of the column_name, the relation and the value (e.g., COST<10). The CONDITIONS is one single string (e.g., "admissions, SUBJECT_ID=24971"). Different conditions are separated by '||'.
(4) GetValue(DATABASE, ARGUMENT), which returns a string containing all the values of the column in the DATABASE (if multiple values, separated by ", "). When there is no additional operations on the values, the ARGUMENT is the column_name in demand. If the values need to be returned with certain operations, the ARGUMENT is composed of the column_name and the operation (like COST, sum). Please do not contain " or ' in the argument.
(5) SQLInterpreter(SQL), which interprets the query SQL and returns the result.
(6) Calendar(DURATION), which returns the date after the duration of time.
Use the variable 'answer' to store the answer of the code. Here are some examples:
{examples}
(END OF EXAMPLES)
Knowledge:
{knowledge}
Question: {question}
Solution: """

DEFAULT_USER_PROXY_AGENT_DESCRIPTIONS = {
    "ALWAYS": "An attentive HUMAN user who can answer questions about the task, and can perform tasks such as running Python code or inputting command line commands at a Linux terminal and reporting back the execution results.",
    "TERMINATE": "A user that can run Python code or input command line commands at a Linux terminal and report back the execution results.",
    "NEVER": "A user that can run Python code or input command line commands at a Linux terminal and report back the execution results.",
}

CodeDebugger = """Given a question:
{question}
The user have written code with the following functions:
(1) Calculate(FORMULA), which calculates the FORMULA and returns the result.
(2) LoadDB(DBNAME) which loads the database DBNAME and returns the database. The DBNAME can be one of the following: allergy, cost, diagnosis, intakeoutput, lab, medication, microlab, patient, treatment, vitalperiodic.
(3) FilterDB(DATABASE, CONDITIONS), which filters the DATABASE according to the CONDITIONS. The CONDITIONS is a string composed of multiple conditions, each of which consists of the column_name, the relation and the value (e.g., COST<10). The CONDITIONS is one single string (e.g., "admissions, SUBJECT_ID=24971"). Different conditions are separated by '||'.
(4) GetValue(DATABASE, ARGUMENT), which returns the values of the column in the DATABASE. When there is no additional operations on the values, the ARGUMENT is the column_name in demand. If the values need to be returned with certain operations, the ARGUMENT is composed of the column_name and the operation (like COST, sum). Please do not contain " or ' in the argument.
(5) SQLInterpreter(SQL), which interprets the query SQL and returns the result.
(6) Calendar(DURATION), which returns the date after the duration of time.

The code is as follows:
{code}

The execution result is:
{error_info}

Please check the code and point out the most possible reason to the error.
"""

EHRAgent_4Shots_Knowledge = """Question: was the fluticasone-salmeterol 250-50 mcg/dose in aepb prescribed to patient 035-2205 on their current hospital encounter?
Knowledge:
- We can find the patient 035-2205 information in the patient database.
- As fluticasone-salmeterol 250-50 mcg/dose in aepb is a drug, we can find the drug information in the medication database.
- We can find the patientunitstayid in the patient database and use it to find the drug precsription information in the medication database.
Solution: patient_db = LoadDB('patient')
filtered_patient_db = FilterDB(patient_db, 'uniquepid=035-2205||hospitaldischargetime=null')
patientunitstayid = GetValue(filtered_patient_db, 'patientunitstayid')
medication_db = LoadDB('medication')
filtered_medication_db = FilterDB(medication_db, 'patientunitstayid={}||drugname=fluticasone-salmeterol 250-50 mcg/dose in aepb'.format(patientunitstayid))
if len(filtered_medication_db) > 0:
	answer = 1
else:
	answer = 0

Question: in the last hospital encounter, when was patient 031-22988's first microbiology test time?
Knowledge:
- We can find the patient 031-22988 information in the patient database.
- We can find the microbiology test information in the microlab database.
- We can find the patientunitstayid in the patient database and use it to find the microbiology test information in the microlab database.
Solution: patient_db = LoadDB('patient')
filtered_patient_db = FilterDB(patient_db, 'uniquepid=031-22988||max(hospitaladmittime)')
patientunitstayid = GetValue(filtered_patient_db, 'patientunitstayid')
microlab_db = LoadDB('microlab')
filtered_microlab_db = FilterDB(microlab_db, 'patientunitstayid={}||min(culturetakentime)'.format(patientunitstayid))
culturetakentime = GetValue(filtered_microlab_db, 'culturetakentime')
answer = culturetakentime

Question: what is the minimum hospital cost for a drug with a name called albumin 5% since 6 years ago?
Knowledge:
- As albumin 5% is a drug, we can find the drug information in the medication database.
- We can find the patientunitstayid in the medication database and use it to find the patienthealthsystemstayid information in the patient database.
- We can use the patienthealthsystemstayid information to find the cost information in the cost database.
Solution: date = Calendar('-6 year')
medication_db = LoadDB('medication')
filtered_medication_db = FilterDB(medication_db, 'drugname=albumin 5%')
patientunitstayid_list = GetValue(filtered_medication_db, 'patientunitstayid, list')
patient_db = LoadDB('patient')
filtered_patient_db = FilterDB(patient_db, 'patientunitstayid in {}'.format(patientunitstayid_list))
patienthealthsystemstayid_list = GetValue(filtered_patient_db, 'patienthealthsystemstayid, list')
cost_db = LoadDB('cost')
min_cost = 1e9
for patienthealthsystemstayid in patienthealthsystemstayid_list:
	filtered_cost_db = FilterDB(cost_db, 'patienthealthsystemstayid={}||chargetime>{}'.format(patienthealthsystemstayid, date))
	cost = GetValue(filtered_cost_db, 'cost, sum')
	if cost < min_cost:
		min_cost = cost
answer = min_cost

Question: what are the number of patients who have had a magnesium test the previous year?
Knowledge:
- As magnesium is a lab test, we can find the lab test information in the lab database.
- We can find the patientunitstayid in the lab database and use it to find the patient information in the patient database.
Solution: answer = SQLInterpreter[select count( distinct patient.uniquepid ) from patient where patient.patientunitstayid in ( select lab.patientunitstayid from lab where lab.labname = 'magnesium' and datetime(lab.labresulttime,'start of year') = datetime(current_time,'start of year','-1 year') )]
"""