EHRAgent — an LLM that writes, runs, and debugs Python to answer complex EHR table queries with four-shot prompts

Overview

Production Readiness

0.6

Novelty Score

0.6

Cost Impact Score

0.5

Citation Count

2

Authors

Wenqi Shi, Ran Xu, Yuchen Zhuang, Yue Yu, Jieyu Zhang, Hang Wu, Yuanda Zhu, Joyce Ho, Carl Yang, May D. Wang

Links

Abstract / PDF

Why It Matters For Business

EHRAgent reduces dependence on data engineers by letting clinicians ask EHR questions in plain language and getting accurate answers; this can speed workflows but increases runtime/API calls and needs privacy safeguards.

Summary TLDR

EHRAgent is a code-centric LLM agent that answers complex, multi-table EHR queries by generating Python plans, executing them, and iteratively debugging using parsed error traces and a small memory of past successful examples. It injects query-specific medical metadata, dynamically selects four relevant few-shot examples from long-term memory, and uses a 'rubber duck' error-explainer to guide fixes. On three real EHR benchmarks (MIMIC-III, eICU, TREQS) it raises success rate vs prior best baselines (same tools and prompts) by up to 29.6 percentage points on evaluated datasets, while requiring only four demonstrations.

Problem Statement

Clinicians need ad hoc, multi-step data from large relational EHRs but rely on data engineers. Off-the-shelf LLMs struggle with multi-table schema linking, domain specifics, and iterative debugging. EHRAgent reframes EHR QA as code-based planning plus execution feedback to let an LLM autonomously generate, run, and refine queries.

Main Contribution

An LLM agent that produces executable Python plans to solve few-shot multi-table EHR questions.

Query-specific medical information injection to improve schema linking and table/column discovery.

Long-term memory that retrieves the most relevant past successful examples for adaptive few-shot prompting.

Interactive coding with an executor and a 'rubber duck' error tracer to iteratively debug and refine code.

Key Findings

EHRAgent substantially improves EHR multi-table QA success rates versus prior LLM agent baselines.

NumbersUp to +29.6 percentage points success rate (TREQS) vs strongest baseline

Interactive coding (code generation + execution feedback) is the largest single contributor to performance.

NumbersAblation: overall SR drops from 58.97% to 24.55% without interactive coding on MIMIC-III

EHRAgent can work with very few demonstrations.

NumbersUses only K=4 few-shot examples versus supervised text-to-SQL needing >10k labeled samples

Results

Success rate (MIMIC-III, overall)

Value58.97%

Success rate (eICU, overall)

Value53.10%

Success rate (TREQS, overall)

Value69.70%

Completion rate (MIMIC-III)

Value85.86%

Who Should Care

Ml EngineerData ScientistProduct ManagerCto

What To Try In 7 Days

Run a small prototype: wire a Python-based toolset (DB loader, SQLInterpreter) to an LLM and test 10 real queries.

Add EHR metadata injection (table/column descriptions) and 4 curated examples to your prompt; compare success vs plain prompts.

Implement a simple error parser that returns error type, location, and candidate causes to the LLM and measure iteration gains.

Agent Features

Memory

long-term memory of successful code examples
similarity-based retrieval for few-shot selection

Planning

code-based planning (Python programs as plans)
multi-step action decomposition

Tool Use

Python tool functions
SQLInterpreter
WolframAlpha calculator API

Frameworks

AutoGen
OpenAI API (GPT-4 / GPT-3.5)

Is Agentic

true

Architectures

LLM planner + code executor loop
Python-centered tool interface

Collaboration

single-agent interactive loop (planner-executor)

Optimization Features

Token Efficiency

selective few-shot example retrieval to stay within context window

Reproducibility

Code Urls

https://github.com/wshi83/EhrAgent

Data Urls

Code Available

Data Available

Open Source Status

partial

Risks & Boundaries

Limitations

Requires multiple LLM calls per query which increases cost and latency.
Relies on closed-box models (GPT-4) in experiments; behavior may differ for open LLMs.
Privacy constraints: MIMIC/eICU require credentialed access and authors avoided sending raw patient data to external services.
Context-length limits can still cause failures on very long dialogs or complex queries.

When Not To Use

For unsupervised, high-stakes clinical decisions without human review.
If strict cost or latency budgets forbid multiple LLM calls per query.
In settings that disallow any external API processing of patient data.

Failure Modes

Date/time miscalculations (uses system dates incorrectly)
Context-length truncation leading to missing info
Incorrect intermediate logic in multi-hop plans
Malformed SQL or empty SQL results from SQLInterpreter
Agent fails to follow prompt instructions
Agent cannot identify root cause within allowed iterations