PaperlayerView exact revision on arXiv

Source-linked AI summary

EHRAgent: Code Empowers Large Language Models for Few-shot Complex Tabular Reasoning on Electronic Health Records

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

arXiv:2401.07128v3cs.CLcs.AI

TL;DR

EHRAgent addresses the difficulty of answering complex clinical questions across large, multi-table EHRs without extensive manual annotations. It generates and executes code as a tool-use plan, using medical information, retrieved successful examples, and iterative execution feedback. Across three real-world EHR datasets, it improves the strongest baseline by up to 29.6% in success rate while using four demonstrations.

  • Problem

    Complex EHR questions require reasoning across many records and tables, while existing supervised query-translation systems require extensive, costly fine-grained annotations.

    Source: §1 · p. 1

  • Method

    EHRAgent combines code-based tool-use planning, autonomous execution and debugging, query-specific medical information, and long-term-memory retrieval of relevant successful demonstrations.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    Up to 29.6% higher success rate than the strongest baseline on multi-hop reasoning across three EHR datasets.

    Source: §1 · p. 2

  • Takeaways & Limitations

    EHRAgent enables autonomous code generation and execution for complex clinical EHR tasks with only four demonstrations.

    Source: §1 · p. 2, §6 · pp. 8–9

    Experiments are restricted to the Microsoft Azure OpenAI API with gpt-3.5-turbo (0613) and gpt-4 (0613) because of API access, budget, and computation constraints.

    Source: Limitation and Future Work · p. 9

Abstract

from arXiv · show

Large language models (LLMs) have demonstrated exceptional capabilities in planning and tool utilization as autonomous agents, but few have been developed for medical problem-solving. We propose EHRAgent, an LLM agent empowered with a code interface, to autonomously generate and execute code for multi-tabular reasoning within electronic health records (EHRs). First, we formulate an EHR question-answering task into a tool-use planning process, efficiently decomposing a complicated task into a sequence of manageable actions. By integrating interactive coding and execution feedback, EHRAgent learns from error messages and improves the originally generated code through iterations. Furthermore, we enhance the LLM agent by incorporating long-term memory, which allows EHRAgent to effectively select and build upon the most relevant successful cases from past experiences. Experiments on three real-world multi-tabular EHR datasets show that EHRAgent outperforms the strongest baseline by up to 29.6% in success rate. EHRAgent leverages the emerging few-shot learning capabilities of LLMs, enabling autonomous code generation and execution to tackle complex clinical tasks with minimal demonstrations.

1 Introduction

EHRAgent addresses complex EHR questions by combining autonomous code-based planning, execution feedback, domain knowledge, and memory with few-shot demonstrations. Across three EHR datasets, it improves multi-tabular reasoning over the strongest baseline while requiring only four demonstrations.

1 Introduction

  • Clinicians can state tasks in natural language while the agent generates and executes code, reducing reliance on specialized data-engineering effort for EHR access.Existing supervised approaches require extensive fine-grained annotations that are expensive and difficult to obtain.Source: Figure 1 · p. 1, §1 · p. 1
  • EHRAgent targets EHR queries spanning many records and tables, where complex clinical tasks challenge LLMs’ ability to locate relevant information and perform multi-step operations.For example, MIMIC-III contains 26 tables covering 46K patients.Source: §1 · p. 2, Figure 2 · p. 2
  • EHRAgent combines query-specific medical information with long-term memory that retrieves relevant successful cases for few-shot reasoning.The approach uses only four demonstrations rather than the over 10K samples required by cited supervised settings.Source: §1 · p. 2
  • The agent formulates clinical problem solving as executable action sequences and iteratively refines them by examining execution errors and environmental feedback.This interactive coding mechanism links a code planner with a code executor.Source: §1 · p. 2, Figure 3 · p. 3
  • Up to 29.6% higher success rate than the strongest baseline on multi-hop reasoning across three EHR datasets.The evaluation targets challenging multi-tabular reasoning tasks on real-world EHRs.Source: §1 · p. 2

2 Preliminaries

The preliminaries formulate EHR question answering as an agent’s execution of a tool-based action plan. They motivate code as the planning interface so execution outcomes can drive iterative plan refinement.

2 Preliminaries

  • Given a natural-language query, the agent produces a T-step action plan from EHRs, column descriptions, and pre-built tools, then executes it to obtain the final output.The plan is represented as sequential actions selected from the available tool set.Source: §2 · p. 3
  • A code interface converts planning into executable code generation, with a code executor returning either successful results or error messages for refinement.The planner and executor interact through a multi-turn dialogue that optimizes plan execution.Source: §2 · p. 3
  • Algorithm 1 initializes medical information, retrieves relevant long-term-memory examples, generates code, executes it, and debugs the plan until termination or the step limit.The output is either a solved final answer or an error message.Source: §2 · p. 4

3 EHRAgent: LLMs as Medical Agents

EHRAgent addresses multi-hop EHR reasoning with query-specific medical information, long-term memory, interactive coding, and error-tracing feedback. These components help the agent locate relevant records, select useful demonstrations, refine executable plans, and debug errors iteratively.

3 EHRAgent: LLMs as Medical Agents

  • EHRAgent combines medical information integration, long-term memory, interactive coding, and rubber duck debugging into an LLM agent for multi-hop EHR reasoning.The workflow uses an LLM agent, code executor, retrieved demonstrations, and error explanations to refine solutions.Source: §3 · p. 4, §3.3 · p. 5, §3.4 · p. 5

3.1 Medical Information Integration

  • Query-specific medical information helps EHRAgent identify and locate tables and records relevant to a clinical question.The agent generates background knowledge from EHR metadata and the query to support table and record retrieval.Source: §3.1 · p. 4

3.2 Demonstration Optimization through Long-Term Memory

  • Long-term memory stores successful code snippets and replaces fixed few-shot examples with the most similar past cases for each query.Similarity is computed between questions using negative Levenshtein distance, and the retrieved examples reformulate the prompt.Source: §3.2 · p. 4

3.3 Interactive Coding with Execution

  • Interactive coding lets the agent generate code, execute it, inspect results, and either revise the plan or provide the final answer.The code executor returns execution outcomes to the LLM, which can iteratively modify code; conversations end with TERMINATE or a step threshold.Source: §3.3 · p. 5, §3.4 · p. 5

3.4 Rubber Duck Debugging via Error Tracing

  • Rubber duck debugging supplies parsed error types, messages, and locations to a separate LLM that proposes likely causes for subsequent refinement.The generated explanations are fed back into the agent–executor conversation to support deeper debugging rather than slight error-driven edits.Source: §3.4 · p. 5

4 Experiments

Experiments across three real-world EHR datasets evaluate EHRAgent against nine baselines, testing its effectiveness, component contributions, behavior across complexity and demonstrations, errors, and iterative debugging.

4.2 Main Results

  • 19.92%, 12.41%, and 29.60% performance gains show EHRAgent significantly outperforms all baselines across the three EHR datasets.Interactive coding with environment feedback and domain knowledge injection support iterative code refinement and planning.Source: §4.2 · p. 6

4.3 Ablation Studies

  • Interactive coding is the most significant contributor across MIMIC-III complexity levels, while knowledge integration helps more on challenging tasks.Knowledge integration supports schema linking and identification of relevant tables, columns, and condition values.Source: §4.3 · pp. 6–7, Table 2 · p. 7

4.4 Quantitative Analysis

  • Success and completion rates generally decrease as questions require more elements or columns, but EHRAgent consistently outperforms baselines across difficulty levels.More elements require calculations or domain knowledge, while more columns challenge accurate column location and association.Source: §4.4 · p. 7, Figure 4 · p. 7
  • EHRAgent remains stable as demonstrations increase, whereas AutoGen’s success and completion rates tend to decrease with more examples.The difference is associated with context limitations in AutoGen and EHRAgent’s debugging and adaptive demonstration selection.Source: §4.4 · p. 7, Figure 5 · p. 7

4.5 Error Analysis

  • Most EHRAgent errors arise when the agent fails to identify underlying causes within T-step trails, producing incomplete or inexcusable plans.The error categories are summarized for MIMIC-III through manual examination and analysis.Source: §4.5 · p. 7, Figure 6 · p. 8

4.6 Case Study

  • A case study shows EHRAgent resolving multiple initial coding errors one by one by analyzing execution messages across iterations.The initial LLM solution is unsatisfactory, but interactive coding with environment feedback enables successive correction.Source: §4.6 · p. 8, Figure 7 · p. 8

5 Related Work

Related work positions EHRAgent among tool-augmented LLM agents and distinguishes it through interactive coding, long-term memory, and complex few-shot EHR reasoning.

Augmenting LLMs with External Tools

  • Interactive coding addresses the disconnect between plan generation and execution by refining generated code from environmental error messages.EHRAgent also retrieves relevant successful experiences through long-term memory rather than relying only on short-term in-context learning.Source: §5 · p. 8

LLM Agents for Scientific Discovery

  • Unlike medical LLM agents focused on multiple-choice or general medical question answering, EHRAgent targets autonomous code generation and execution for few-shot multi-tabular reasoning on real-world EHRs.The comparison emphasizes its interactive code interface and structured EHR setting.Source: §5 · p. 8

6 Conclusion

EHRAgent applies LLMs to few-shot multi-tabular reasoning on real-world EHRs through autonomous code generation and execution. Interactive feedback and accumulated medical knowledge support plan optimization and improved medical reasoning over baseline agents.

6 Conclusion

  • EHRAgent is an LLM agent with external tools for few-shot multi-tabular reasoning on real-world EHRs.It uses autonomous code generation and execution for communication between clinicians and EHR systems.Source: §6 · pp. 8–9
  • Interactive coding with execution feedback and accumulative medical knowledge facilitate plan optimization for multi-step problem-solving.Experiments report advantages over baseline LLM agents in autonomous coding and medical reasoning.Source: §6 · p. 9

Limitation and Future Work

EHRAgent’s reported gains come with trade-offs in execution cost, clinical evaluation, privacy, and deployment scope. Future work targets more cost-effective, responsible, and broadly generalizable use.

Additional Execution Calls

  • Additional LLM calls trade execution cost for environmental feedback that adaptively refines plans and improves performance.Open-loop methods avoid multi-round adaptation, but their success rates remain below 40% on all three EHR datasets.Source: Limitation and Future Work · p. 9

Translational Clinical Research Considerations

  • Future development must address privacy, safety, bias, and ethical implications before broader clinical research and practice deployment.The authors frame responsible healthcare AI as an ongoing development goal.Source: Limitation and Future Work · p. 9
  • EHRAgent’s adaptation and generalization to low-resource languages are constrained by limited resources and training data.Current experiments are restricted to Azure OpenAI’s gpt-3.5-turbo (0613) and gpt-4 (0613) because of API, budget, and computation constraints.Source: Limitation and Future Work · p. 9
  • The authors plan to investigate fine-tuned white-box LLMs, including LLaMA-2, as a future direction.This is presented alongside efforts to expand model access beyond the currently used API-based systems.Source: Limitation and Future Work · p. 9

Completion Rate under Clinical Scenarios

  • Completion rate can reward executable plans even when answers are incorrect, so higher values may be unsafe in clinical settings.The authors plan stricter metrics addressing misinformation risk.Source: Limitation and Future Work · p. 9

Privacy and Ethical Statement

The study follows credentialed-data and privacy requirements for MIMIC-III and eICU while using Azure OpenAI without human review access to sensitive patient information.

Privacy and Ethical Statement

  • Confidential MIMIC-III and eICU patient data are prohibited from transfer to third parties, including through online APIs.The work states that this follows the PhysioNet Credentialed Health Data Use Agreement 1.5.04.Source: Privacy and Ethical Statement · p. 9
  • The researchers opted out of Azure OpenAI human review so third parties cannot access or process sensitive patient information.They report continuous monitoring for compliance with relevant guidelines and privacy laws.Source: Privacy and Ethical Statement · p. 9

A.1 Task Details

This section defines the multi-tabular EHR QA setting, datasets, tools, baselines, implementation choices, and four-shot experimental setup. It spans real-world clinical queries while standardizing shared resources across methods.

A.1 Task Details

  • The evaluation uses EHRSQL, TREQS, and related structured EHR benchmarks built from MIMIC-III and eICU for clinical text-to-SQL questions.The questions range from simple retrieval to complex calculations and were gathered from over 200 hospital staff.Source: §A.1 · p. 13

A.2 Question Complexity Level

  • Question complexity levels are based on the number of tables required, with additional analysis by question elements and solution columns.MIMIC-III uses Levels I–IV, while eICU and TREQS use Levels I–III.Source: §A.2 · p. 13

A.3 MIMIC-III

  • MIMIC-III contains 38,597 patients and 49,785 critical-care admissions with granular clinical and administrative information.Its data include demographics, vital signs, laboratory results, procedures, medications, notes, imaging reports, and mortality.Source: §A.3 · p. 13

A.4 eICU

  • eICU contains over 200,000 admissions from U.S. critical-care units in 2014–2015, with deidentified administrative and structured clinical data.The dataset includes vital signs, laboratory measurements, medications, treatment plans, diagnoses, and medical histories.Source: §A.4 · p. 13

A.5 TREQS

  • TREQS is a MIMIC-III-based benchmark restricted to five tables but containing many records per table.Its question types and SQL complexity are narrower than those in EHRSQL.Source: §A.5 · pp. 13–14

B Tool Set Details

  • The tool set covers table loading, filtering, value retrieval, calculation, date calculation, and SQL interpretation.These tools support both relational database operations and basic numerical or temporal computations.Source: B Tool Set Details · p. 14

C Baseline Details

  • Baselines and EHRAgent share tool definitions, table metadata, and few-shot demonstrations, while differing in prompting or technical design.The comparison includes natural-language planners and code-interface systems such as LLM2SQL, DIN-SQL, Self-Debugging, and AutoGen.Source: C Baseline Details · p. 14

D Selection of Primary Programming Language

E.5 Evaluation Metric Details

The evaluation measures both successful answers and executable plans, while the prompt design supplies dataset-specific schema knowledge, tools, examples, and debugging feedback. A GPT-3.5-turbo analysis tests whether the approach generalizes across backbones.

E.5 Evaluation Metric Details

  • Success rate measures correctly handled queries, while completion rate measures executable plans regardless of answer correctness.Completion rate is used to compare agents with and without environmental feedback.Source: §E.5 · p. 17

⋄MIMIC-III.

  • MIMIC-III metadata explains identifier-based table links, event tables, dictionary joins, patient-stay tables, and clinical data tables.This schema knowledge supports locating concepts and connecting records across the database.Source: ⋄MIMIC-III · p. 17

⋄eICU.

  • eICU metadata describes relational patient-care data while warning that missing interfaces can make measurements unavailable even when they occurred.The database includes vital signs, laboratory data, medications, diagnoses, treatments, and other care information.Source: ⋄eICU · p. 17

⋄TREQS.

  • TREQS organizes patient information into five category tables linked through patient ID and admission ID.The categories are demographics, laboratory tests, diagnosis, procedures, and prescriptions.Source: ⋄TREQS · p. 18

E.7 Prompt Details

  • EHRAgent prompts combine Python code generation, four-shot examples, tool definitions, medical knowledge integration, and execution-error debugging.The debugging prompt asks the agent to inspect generated code and execution results to identify likely errors.Source: §E.7 · p. 18

F.1 Effect of Base LLMs

  • 6.72% performance gain shows EHRAgent still outperforms all baselines with GPT-3.5-turbo, although both systems and baselines perform worse than with GPT-4.The authors attribute the decrease primarily to weaker instruction-following and reasoning in GPT-3.5-turbo.Source: §F.1 · p. 18, Table 6 · p. 19

F.2 Additional Ablation Studies

Ablations show that interactive coding contributes most to EHRAgent’s gains, while all four components improve performance. The analysis also reports query costs and sustained advantages on increasingly complex questions.

F.2 Additional Ablation Studies

  • All four EHRAgent components contribute significantly to performance gains, with medical information integration especially benefiting challenging tasks involving more tables and domain knowledge.The ablation studies on MIMIC-III and eICU report consistent results.Source: §F.2 · p. 19, Table 7 · p. 20
  • Long-term memory improves performance by selecting relevant demonstrations, but including unsuccessful completed cases raises completion while generally lowering success rates.The memory mechanism still outperforms the version without long-term memory.Source: §F.2 · p. 19, Table 8 · p. 20
  • Interactive coding is the most significant contributor to performance gains across complexity levels, increasing completion by avoiding overly complex contexts and refining plans with execution feedback.The ablation replaces code-based plans with natural-language plans while retaining other modules.Source: §F.2 · p. 19

F.3 Cost Estimation

  • $0.60, $0.17, and $0.52 are EHRAgent’s average per-query costs on MIMIC-III, eICU, and TREQS, respectively.Cost is mainly determined by question complexity and the difficulty of locating relevant information in tables.Source: §F.3 · p. 20

G.1 Additional Question Complexity Analysis

  • For questions involving more than 10 columns, open-loop baselines complete less than 20% of queries, whereas EHRAgent correctly answers around 50%.EHRAgent consistently outperforms baselines across difficulty levels, although performance generally declines as complexity increases.Source: §G.1 · p. 20

G.2 Additional Error Analysis

Error analysis identifies failures in date handling, context length, logic, SQL commands, instruction following, and debugging. The accompanying comparison discusses Python’s few-shot and feedback advantages alongside execution-time limitations on large databases.

G.2 Additional Error Analysis

  • Date queries can be inaccurate when the agent calculates using real-world time instead of the database system time used by the Calendar tool.The Calendar tool is described as typically reliable when used for date and time calculations.Source: §G.2 · p. 20
  • Excessively long queries or dialogue histories can exceed the context-length limit, while multi-hop plans can also contain incorrect intermediate logic.One example is filtering by patient_id instead of the correct admission_id when computing a hospital visit’s total cost.Source: §G.2 · p. 20
  • Incorrect SQL commands can return empty SQLInterpreter responses, causing subsequent parts of the Python plan to fail.This is identified as a recurring error type during solution generation.Source: §G.2 · p. 20
  • The agent often fails to follow initial or debugging instructions and may fail to identify error root causes despite using all T-step trials.These failures can leave plans incomplete or unusable.Source: §G.2 · p. 20

G.3 Additional Empirical Comparison of Primary Programming Languages

  • EHRAgent uses four demonstrations for few-shot multi-tabular reasoning, whereas SQL-based methods usually require at least tens of demonstrations to learn schemas and SQL grammar.The comparison attributes this difference partly to SQL’s smaller representation in training data.Source: §G.3 · p. 21
  • Python can be less efficient than SQL on large databases, but on MIMIC-III EHRAgent averages 52.63 seconds per question versus DIN-SQL’s 103.28 seconds.The paper identifies Python efficiency on large-scale databases as a future direction.Source: §G.3 · p. 21

H Additional Case Studies

  • The additional case studies cover medical information integration, long-term memory, the code interface, rubber duck debugging, and the complete EHRAgent workflow.Figures 13 and 14 illustrate the workflow with accumulated domain knowledge and coding capability.Source: H Additional Case Studies · p. 21

H.1 Case Study of Medical Information Integration

The case studies illustrate how EHRAgent integrates clinical knowledge, retrieves relevant demonstrations, plans with code, and debugs through environmental feedback. Together, these mechanisms support table identification, context-efficient action planning, and root-cause-based refinement.

H.1 Case Study of Medical Information Integration

  • EHRAgent links clinical concepts such as “aspirin ec” and “venous cath nec” to relevant tables, records, and identifiers including ICD9_CODE and HADM_ID.The case identifies potential locations such as prescriptions and d_icd_procedures.Source: §H.1 · p. 21, Figure 8 · p. 21

H.2 Case Study of Long-Term Memory

  • Long-term memory replaces fixed demonstrations with relevant successful cases, supplying matching logic for count and procedure questions absent from the original examples.This retrieval addresses limited context length and improves the relevance of few-shot guidance.Source: §H.2 · pp. 21–22, Figure 9 · p. 23

H.4 Case Study of Rubber Duck Debugging

  • EHRAgent’s debugging process analyzes error causes beyond surface messages, enabling accurate refinement such as correcting case-sensitive code issues.AutoGen instead makes slight modifications to the original error feedback without reasoning about the root cause.Source: §H.4 · p. 22, Figure 12 · p. 24

H.3 Case Study of Code Interface

  • Compared with ReAct’s repeated natural-language actions, EHRAgent uses code structures such as for loops to shorten trajectories and save context space.ReAct’s lengthy context can eventually lead to a low completion rate.Source: §H.3 · p. 22, Figure 10 · p. 23

H.4 Case Study of Rubber Duck Debugging

  • The complete workflow integrates medical information, retrieves memory examples, generates and executes code, and iteratively debugs until producing a final solution.The workflow is presented as a sequence from metadata and tool definitions through execution feedback.Source: Figure 14 · p. 25, Figure 13 · p. 25
Loading 2401.07128v3…