Agents and tool use

import json

# Tools the agent may call, described in the JSON-schema style most model APIs use
TOOLS = {
    "egfr_ckd_epi_2021": {
        "description": "Estimate eGFR (mL/min/1.73m2) from serum creatinine (mg/dL), age and sex, "
                       "CKD-EPI 2021 equation.",
        "parameters": {"creatinine_mg_dl": "number", "age": "integer", "sex": "'female' or 'male'"},
    },
    "lookup_protocol": {
        "description": "Return the clinic protocol text for a topic.",
        "parameters": {"topic": "one of 'metformin', 'hypoglycaemia'"},
    },
}

def egfr_ckd_epi_2021(creatinine_mg_dl, age, sex):
    k, a = (0.7, -0.241) if sex == "female" else (0.9, -0.302)
    ratio = creatinine_mg_dl / k
    egfr = (142 * min(ratio, 1) ** a * max(ratio, 1) ** -1.200 * 0.9938 ** age
            * (1.012 if sex == "female" else 1))
    return round(egfr)

def lookup_protocol(topic):
    return {"metformin": "Do not start metformin if eGFR is below 45. Stop it if eGFR falls below 30.",
            "hypoglycaemia": "Treat glucose below 70 mg/dL with 15 g fast-acting carbohydrate."}[topic]

print(json.dumps(TOOLS["egfr_ckd_epi_2021"], indent=1))
print(egfr_ckd_epi_2021(1.4, 68, "female"))

# {
#  "description": "Estimate eGFR (mL/min/1.73m2) from serum creatinine (mg/dL), age and sex, CKD-EPI 2021 equation.",
#  "parameters": {
#   "creatinine_mg_dl": "number",
#   "age": "integer",
#   "sex": "'female' or 'male'"
#  }
# }
# 41

def scripted_model(messages):
    """Stand-in for a language model. A real model decides these steps itself;
    here they are fixed so the loop runs without an API."""
    tool_results = [m for m in messages if m["role"] == "tool"]
    if len(tool_results) == 0:
        return {"tool_call": {"name": "egfr_ckd_epi_2021",
                              "arguments": {"creatinine_mg_dl": 1.4, "age": 68, "sex": "female"}}}
    if len(tool_results) == 1:
        return {"tool_call": {"name": "lookup_protocol", "arguments": {"topic": "metformin"}}}
    egfr = tool_results[0]["content"]
    return {"answer": f"Estimated eGFR is {egfr} mL/min/1.73m2. The protocol says not to start metformin "
                      f"below 45, so it should not be started; please confirm with the prescriber. "
                      f"[tools: egfr_ckd_epi_2021, lookup_protocol]"}

IMPLEMENTATIONS = {"egfr_ckd_epi_2021": egfr_ckd_epi_2021, "lookup_protocol": lookup_protocol}

def run_agent(question, max_steps=5):
    messages = [{"role": "user", "content": question}]
    for step in range(1, max_steps + 1):
        reply = scripted_model(messages)
        if "answer" in reply:
            return reply["answer"], messages
        call = reply["tool_call"]
        if call["name"] not in IMPLEMENTATIONS:                      # allow-list
            raise ValueError(f"Unknown tool {call['name']}")
        result = IMPLEMENTATIONS[call["name"]](**call["arguments"])
        print(f"step {step}: {call['name']}({call['arguments']}) -> {result}")
        messages.append({"role": "tool", "name": call["name"], "content": result})
    raise RuntimeError("Step limit reached without an answer")

answer, log = run_agent("Synthetic case: 68-year-old woman, creatinine 1.4 mg/dL. Can metformin be started?")
print("\n" + answer)

# step 1: egfr_ckd_epi_2021({'creatinine_mg_dl': 1.4, 'age': 68, 'sex': 'female'}) -> 41
# step 2: lookup_protocol({'topic': 'metformin'}) -> Do not start metformin if eGFR is below 45. Stop it if eGFR falls below 30.
#
# Estimated eGFR is 41 mL/min/1.73m2. The protocol says not to start metformin below 45, so it should not be started; please confirm with the prescriber. [tools: egfr_ckd_epi_2021, lookup_protocol]

def validated_egfr(creatinine_mg_dl, age, sex):
    problems = []
    if not 0.2 <= creatinine_mg_dl <= 20:
        problems.append(f"creatinine {creatinine_mg_dl} mg/dL is implausible (was it µmol/L?)")
    if not 18 <= age <= 120:
        problems.append(f"age {age} outside the adult range of this equation")
    if sex not in ("female", "male"):
        problems.append(f"sex must be 'female' or 'male', got {sex!r}")
    if problems:
        return {"error": problems}            # returned to the model, never silently 'fixed'
    return {"egfr": egfr_ckd_epi_2021(creatinine_mg_dl, age, sex)}

print(validated_egfr(1.4, 68, "female"))
print(validated_egfr(124, 68, "female"))     # creatinine given in µmol/L by mistake

# {'egfr': 41}
# {'error': ['creatinine 124 mg/dL is implausible (was it µmol/L?)']}