Source-linked AI summary

TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools

Shanghua Gao, Richard Zhu, Zhenglun Kong, Ayush Noori, Xiaorui Su, Curtis Ginder, Theodoros Tsiligkaridis, Marinka Zitnik

arXiv:2503.10970v1cs.AIcs.LG

TL;DR

Personalized therapeutic reasoning must account for patient-specific factors while overcoming the limits of static, tool-free language models. TxAgent combines multi-step reasoning with real-time biomedical retrieval to support treatment decisions, identifying clinically relevant contraindications such as avoiding Prozac with Xolremdi.

  • Problem

    Therapeutic decisions require patient-specific factors, while language models relying solely on their internal knowledge have limited access to real-time biomedical information.

    Source: Discussion · p. 20
  • Method

    TxAgent uses goal-driven tool selection to call external databases and specialized machine-learning models for real-time biomedical knowledge retrieval.

    Source: Main · p. 2
  • Results

    TxAgent identifies that Xolremdi should not be used with Prozac, a CYP2D6 inhibitor, because the interaction alters treatment suitability.

    Source: Discussion · p. 20
  • Takeaways & Limitations

    TxAgent produces clinically relevant and personalized treatment recommendations for therapeutic reasoning.

    Source: Discussion · p. 20

    TxAgent depends on tool calls for external information, and gaps in ToolUniverse limit access to specific data types.

    Source: Discussion · p. 20

Abstract

from arXiv · show

Precision therapeutics require multimodal adaptive models that generate personalized treatment recommendations. We introduce TxAgent, an AI agent that leverages multi-step reasoning and real-time biomedical knowledge retrieval across a toolbox of 211 tools to analyze drug interactions, contraindications, and patient-specific treatment strategies. TxAgent evaluates how drugs interact at molecular, pharmacokinetic, and clinical levels, identifies contraindications based on patient comorbidities and concurrent medications, and tailors treatment strategies to individual patient characteristics. It retrieves and synthesizes evidence from multiple biomedical sources, assesses interactions between drugs and patient conditions, and refines treatment recommendations through iterative reasoning. It selects tools based on task objectives and executes structured function calls to solve therapeutic tasks that require clinical reasoning and cross-source validation. The ToolUniverse consolidates 211 tools from trusted sources, including all US FDA-approved drugs since 1939 and validated clinical insights from Open Targets. TxAgent outperforms leading LLMs, tool-use models, and reasoning agents across five new benchmarks: DrugPC, BrandPC, GenericPC, TreatmentPC, and DescriptionPC, covering 3,168 drug reasoning tasks and 456 personalized treatment scenarios. It achieves 92.1% accuracy in open-ended drug reasoning tasks, surpassing GPT-4o and outperforming DeepSeek-R1 (671B) in structured multi-step reasoning. TxAgent generalizes across drug name variants and descriptions. By integrating multi-step inference, real-time knowledge grounding, and tool-assisted decision-making, TxAgent ensures that treatment recommendations align with established clinical guidelines and real-world evidence, reducing the risk of adverse events and improving therapeutic decision-making.

Main

TXAGENT combines multi-step reasoning, goal-driven biomedical tool use, and real-time evidence retrieval to support accurate drug reasoning and personalized treatment selection. Across five benchmarks, it outperforms larger language models and tool-use systems while remaining robust to drug representation changes.

Method

  • TXAGENT delivers evidence-grounded treatment recommendations by combining multi-step reasoning with real-time biomedical tool integration, goal-driven tool selection, and TOOLUNIVERSE, which consolidates 211 expert-curated tools.Source: Main · p. 2

Evaluation and training

  • TXAGENT is trained and evaluated through a specialized multi-step reasoning system, five benchmarks, and 378,027 instruction-tuning samples derived from 85,340 reasoning traces.The benchmarks assess drug selection, treatment personalization, and reasoning robustness across structured and unstructured queries.Source: Main · pp. 2–3

Results

  • TXAGENT achieves 92.1% accuracy in open-ended drug reasoning, surpassing GPT-4o by 25.8% and Llama-3.1-70B-Instruct [2] by 39.3%.TXAGENT is based on a fine-tuned 8-billion-parameter Llama-3.1-8B-Instruct model and outperforms the nearly 9× larger Llama-3.1-70B-Instruct.Source: Main · p. 3
  • TXAGENT achieves 56.5% accuracy on DescriptionPC, outperforming GPT-4o by 8.3%, while maintaining accuracy variance below 0.01 across drug-name variants.GPT-4o shows variance of 9.96, indicating greater sensitivity to brand, generic, and descriptive drug representations.Source: Main · p. 3
  • On TreatmentPC’s 456 scenarios, TXAGENT outperforms GPT-4o by 13.6%, Llama-3.1-70B-Instruct by 25.4%, and DeepSeek-R1 by 7.5% in open-ended queries.These results establish stronger performance in personalized treatment recommendation and multi-step therapeutic reasoning.Source: Main · pp. 3–4

Ablation studies

  • Ablations show that larger toolboxes, real-world tool use, explicit reasoning before function calls, and longer multi-step traces each improve therapeutic reasoning.Tool-assisted decision-making outperforms LLM-only reasoning, while structured reasoning contributes more than multi-round function calls alone.Source: Main · p. 4

Results … TOOLUNIVERSE: A universe of tools and machine learning models

This section describes how TXAGENT combines multi-step therapeutic reasoning, tool selection, and real-time retrieval across TOOLUNIVERSE’s 211-tool biomedical toolbox, producing verified answers and transparent reasoning traces for complex patient and drug scenarios. It also covers TOOLUNIVERSE’s automated and expert tool validation and the use of specialized agents to support summarization, generation, checking, and robustness to unexpected inputs.

TXAGENT: Multi-step therapeutic reasoning with a universe of tools

  • TXAGENT performs multi-step therapeutic reasoning by generating thoughts, invoking tools through TOOLRAG, incorporating feedback, and returning a final answer with a reasoning trace.The iterative process ends with a FINISH tool and records function calls, tool feedback, and intermediate reasoning.Source: TXAGENT: Multi-step therapeutic… · p. 5

Capabilities of TXAGENT

  • TXAGENT retrieves verified biomedical information, adaptively selects tools, and integrates continuously updated sources to solve complex therapeutic tasks with transparent justifications.Examples include retrieving drug dosage beyond the base model’s training data, selecting adverse-reaction tools, and combining multiple retrieval steps when no single tool suffices.Source: Capabilities of TXAGENT · pp. 5–6

TOOLUNIVERSE: A universe of tools and machine learning models

TXAGENT-INSTRUCT dataset of therapeutic tasks and reasoning traces · TXAGENT outperforms larger LLMs in multi-step reasoning

TXAGENT-INSTRUCT provides large-scale therapeutic questions, tool-use data, and reasoning traces, while TXAGENT achieves superior and robust multi-step drug reasoning performance against larger LLMs and GPT-4o.

TXAGENT-INSTRUCT dataset of therapeutic tasks and reasoning traces

  • TXAGENT-INSTRUCT contains 378,027 instruction-tuning samples spanning tooling, therapeutic questions, and reasoning traces.Its reasoning-trace dataset includes 85,340 traces, 177,626 reasoning steps, and 281,695 function calls.Source: TXAGENT-INSTRUCT dataset of… · p. 8
  • The dataset augments 211 TOOLUNIVERSE tools and uses verified biomedical sources including FDA labeling documents, PrimeKG, and Open Targets.Tool descriptions are rephrased to promote tool-use learning rather than memorization.Source: TXAGENT-INSTRUCT dataset of… · p. 8
  • QUESTIONGEN generates therapeutic questions from verified knowledge bases through information extraction, question construction, answer explanation, and validation.It produces 85,340 questions and filters them for knowledge grounding, solvability, and reasonableness before trace generation.Source: TXAGENT-INSTRUCT dataset of… · p. 8
  • TRACEGEN uses HELPER, TOOL PROVIDER, SOLVER, and evaluation components to integrate tool outputs, manage unpredictable responses, and generate validated stepwise traces.The system addresses multi-factor questions, external-tool integration, noisy outputs, hallucinations, arbitrary outputs, and repetitive reasoning.Source: TXAGENT-INSTRUCT dataset of… · pp. 8–9

TXAGENT outperforms larger LLMs in multi-step reasoning

  • DrugPC benchmarks TXAGENT’s drug reasoning with 3,168 questions spanning 11 tasks and two settings, including multiple-choice, while using drugs approved in 2024 for evaluation after excluding post-2023 approvals from training.Source: TXAGENT outperforms larger LLMs in… · pp. 9–10
  • 93.8% multiple-choice and 92.1% open-ended accuracy surpass Llama3.1-70B-Instruct’s 75.1% and 52.8%, while exceeding GPT-4o by 17.4% and 25.8%.These results compare TXAGENT with larger models across both evaluation formats.Source: TXAGENT outperforms larger LLMs in… · p. 10
  • TXAGENT’s accuracy declines only 1.7% from multiple-choice to open-ended tasks, versus 10.1% for GPT-4o and 22.3% for Llama3.1-70B-Instruct.The open-ended setting removes answer choices and is therefore more challenging.Source: TXAGENT outperforms larger LLMs in… · p. 10
  • Across all 11 DrugPC tasks, TXAGENT surpasses every baseline and provides reasoning traces grounded in verified function-call results.The traces support response reliability assessment, unlike outputs requiring manual external verification.Source: TXAGENT outperforms larger LLMs in… · pp. 10–11

TXAGENT outperforms tool-use LLMs in multi-step reasoning · TXAGENT generalizes across drug name variants and descriptions · TXAGENT for precision treatment recommendation

TxAgent advances multi-step therapeutic reasoning by combining tool-assisted biomedical grounding with personalized treatment recommendation across specialized scenarios. It also generalizes across drug name variants and descriptions while outperforming tool-use and reasoning LLMs on these tasks.

TXAGENT outperforms tool-use LLMs in multi-step reasoning

  • TXAGENT’s gains reflect goal-oriented tool selection, iterative function calls, and adaptive multi-step reasoning rather than single-round tool use.Existing tool-use LLMs struggle with selecting among hundreds of tools, retrieving missing information, and refining calls from previous results.Source: TXAGENT outperforms tool-use LLMs in… · p. 11
  • TXAGENT outperforms ToolACE by 62.5% and WattTool by 59.1% in multiple-choice reasoning, and surpasses them by 59.4% and 55.0%, respectively, in open-ended reasoning.Source: TXAGENT outperforms tool-use LLMs in… · p. 11
  • Tool-use baselines produce invalid answers on 58.9%–63.1% of multiple-choice and 56.6%–60.7% of open-ended questions, unlike TXAGENT’s adaptive execution.WattTool-8B fails on 58.9% and 56.6%, while ToolACE-8B fails on 63.1% and 60.7%, respectively.Source: TXAGENT outperforms tool-use LLMs in… · p. 12

TXAGENT generalizes across drug name variants and descriptions

  • TXAGENT reaches 93.6% on BrandPC and 93.7% on GenericPC, exceeding the strongest pure-LLM references by 20.6% and 16.4%.The strongest pure-LLM references are Llama3.1-70B-Instruct at 73.0% on BrandPC and GPT-4o at 77.3% on GenericPC.Source: TXAGENT generalizes across drug name… · p. 12
  • TXAGENT’s variance is 0.00667 across original, BrandPC, and GenericPC, compared with GPT-4o’s 9.96 and WattTool-8B’s 13.07.These results indicate more consistent performance across drug name representations than the compared models.Source: TXAGENT generalizes across drug name… · p. 13
  • On DescriptionPC, TXAGENT achieves 90.4% answer correctness, 56.5% joint drug-identification-and-answer accuracy, and 60.1% drug-identification accuracy.The two-step evaluation separates drug identification from answer selection; joint accuracy requires both to be correct.Source: TXAGENT generalizes across drug name… · p. 13

TXAGENT for precision treatment recommendation

  • 86.8% multiple-choice and 75.0% open-ended accuracy show TXAGENT outperforming Llama-3.1-8B-Instruct’s 56.1% and 33.11% on TreatmentPC.TreatmentPC contains 456 specialized treatment scenarios requiring patient-condition analysis, drug selection, dosage considerations, contraindications, and interactions.Source: TXAGENT for precision treatment… · pp. 13–14
  • TXAGENT reaches 86.8% versus WattTool-8B’s 18.2% in multiple-choice tasks and 75.0% versus ToolACE-8B’s 13.4% in open-ended tasks.Tool-use baselines had full access to TOOLUNIVERSE, but performed substantially worse.Source: TXAGENT for precision treatment… · p. 14
  • TXAGENT’s advantage comes from integrating multiple sources, iterative function calls, query refinement, and dynamically adjusted reasoning for complex treatment recommendations.Unlike reasoning LLMs relying on internal knowledge, TxAgent uses trusted sources such as FDA drug labels, reducing hallucinations and misjudgments.Source: TXAGENT for precision treatment… · pp. 14–15
  • TXAGENT outperforms full DeepSeek-R1 by 10.3% in multiple-choice accuracy (86.8% vs. 76.5%) and by 7.5% in open-ended accuracy.The comparison evaluates DeepSeek-R1 on TreatmentPC’s multi-step reasoning over patient conditions and drug effects.Source: TXAGENT for precision treatment… · p. 14
  • Against DeepSeek-R1-Llama-8B, TXAGENT gains 36.1% in multiple-choice accuracy and 34.9% in open-ended tasks.The distilled model shares TXAGENT’s Llama-3.1-8B base, while TXAGENT additionally integrates verified external information from TOOLUNIVERSE.Source: TXAGENT for precision treatment… · p. 15

Examples of TXAGENT reasoning traces for specialized treatments

TXAGENT’s reasoning traces show how it combines drug mechanisms, interactions, age-specific guidance, comorbidities, and external evidence to produce personalized treatment recommendations. Across pediatric, interaction, geriatric, and cardiac cases, it iteratively retrieves and synthesizes relevant clinical information.

Treatment selection based on drug mechanism and pediatric use

  • TXAGENT recommends Duvyzat for a pediatric patient with Duchenne muscular dystrophy after identifying mechanism and pediatric-safety constraints.It distinguishes Duvyzat from steroid and exon-skipping options, then confirms safety for children over six years old.Source: Examples of TXAGENT reasoning traces… · p. 15

Treatment selection considering drug-drug interactions

  • TXAGENT concludes that Prozac and Xolremdi are unsuitable together because both affect CYP2D6, creating direct and compounded interaction risks.Xolremdi inhibits CYP2D6, while Prozac is both a CYP2D6 substrate and inhibitor, potentially increasing drug exposure and adverse effects.Source: Examples of TXAGENT reasoning traces… · p. 16

Treatment selection considering geriatric use

  • TXAGENT identifies 100 mg/20 mg twice daily as Cobenfy’s maximum recommended dose for elderly patients, below the 125 mg/30 mg twice-daily younger-patient dose.The geriatric adjustment reflects increased urinary-retention risk, and parallel tool calls support the patient-specific recommendation.Source: Examples of TXAGENT reasoning traces… · p. 16

Treatment selection considering comorbidities

  • TXAGENT filters hypertension treatments against second-degree AV-block contraindications and excludes two of the retrieved drugs.It searches FDA-approved labels, then summarizes mechanisms of the remaining non-contraindicated candidates.Source: Examples of TXAGENT reasoning traces… · p. 17

Impact of tools in TOOLUNIVERSE on TXAGENT’s performance · The critical role of reasoning in TXAGENT

TOOLUNIVERSE’s verified biomedical tools substantially outperform LLM-based tool substitutes, and expanding the tool set improves TXAGENT’s performance. Explicit, deeper reasoning during training and inference is likewise essential, especially for complex treatment tasks.

Impact of tools in TOOLUNIVERSE on TXAGENT’s performance

  • The tool comparison evaluates verified TOOLUNIVERSE retrieval against LLM-based substitutes using GPT-4o and Llama 3.1-Instruct-8B with other settings unchanged.The LLM-only condition prompts models to mimic tool functionality from structured descriptions of tool capabilities and arguments.Source: Impact of tools in TOOLUNIVERSE on… · p. 17
  • On DrugPC, replacing real TOOLUNIVERSE tools reduces accuracy from 93.8% to 68.7% with Llama3.1-8B-Instruct and 72.7% with GPT-4o, while TreatmentPC falls to 74.78% and 67.11% versus 86.84%.These results show that LLM-only tool emulation underperforms verified tools for precise biomedical information and treatment reasoning.Source: Impact of tools in TOOLUNIVERSE on… · p. 17
  • Expanding TOOLUNIVERSE raises DrugPC accuracy from 78.4% with 10% of tools to 93.8% with all tools, while TreatmentPC increases from 71.7% to 86.8%.The nested subsets contain all tools from smaller subsets, isolating the incremental effect of adding tools on both benchmarks.Source: Impact of tools in TOOLUNIVERSE on… · p. 18

The critical role of reasoning in TXAGENT

  • The reasoning experiments separately remove thought generation, restrict training traces to at most 1, 3, or 5 steps, and impose inference-time step limits.Inference limits force a final answer after the prescribed number of steps, while unrestricted inference is retained for models trained with different trace depths.Source: The critical role of reasoning in… · p. 18
  • Removing explicit thought generation lowers DrugPC accuracy from 93.8% to 71.5% and TreatmentPC accuracy from 86.4% to 64.9%.Unlike tool-use LLMs that produce only function calls, TXAGENT generates reasoning thoughts and function calls at each step.Source: The critical role of reasoning in… · p. 18
  • Training with only one reasoning step reduces TreatmentPC accuracy from 86.8% to 66.9% and DrugPC accuracy from 93.8% to 71.6%.The larger TreatmentPC decline indicates that complex treatment decisions require stronger multi-step reasoning traces during training.Source: The critical role of reasoning in… · p. 19
  • Limiting inference to one reasoning step yields 73.5% TreatmentPC accuracy, 13.3% below unrestricted reasoning; gains continue through five steps before plateauing.Maintaining full reasoning capacity remains optimal despite diminishing returns beyond five steps.Source: The critical role of reasoning in… · p. 19
  • TreatmentPC requires more reasoning steps and tool calls than DrugPC, particularly in open-ended settings compared with multiple-choice settings.DrugPC shows no significant reasoning-step or tool-call difference between multiple-choice and open-ended formats, whereas open-ended TreatmentPC requires significantly more of both.Source: The critical role of reasoning in… · p. 19

Discussion

TXAGENT combines adaptive multi-step reasoning with real-time biomedical tools to produce transparent, patient-specific therapeutic recommendations grounded in verified evidence. Its limitations include restricted tool coverage, unresolved uncertainty quantification, and support only for natural-language inputs, motivating multimodal and longer-memory extensions.

  • TXAGENT generates transparent reasoning traces while using real-time, multi-source biomedical evidence to support evidence-based therapeutic recommendations.It integrates tools such as FDA drug labels, Open Targets, and other trusted sources rather than relying solely on static training data.Source: Discussion · p. 20
  • TXAGENT adapts multi-step reasoning to patient age, comorbidities, pregnancy status, disease severity, immune function, phenotypes, and biological targets.It identifies diseases from phenotypes, retrieves candidate treatments, and evaluates drug suitability through iterative biomedical-tool calls grounded in verified sources.Source: Discussion · p. 20
  • TXAGENT identifies clinically relevant drug interactions, such as avoiding Xolremdi with Prozac because CYP2D6 inhibition alters Xolremdi metabolism.By incorporating patient-specific constraints, it produces personalized treatment recommendations.Source: Discussion · p. 20
  • Gaps in TOOLUNIVERSE restrict access to some data types, limiting TXAGENT’s ability to address a broader range of therapeutic questions.Although external tools improve verifiability, integrating internal knowledge with tool feedback could increase flexibility for exploratory tasks and address uncertainty quantification challenges.Source: Discussion · p. 20
  • TXAGENT currently processes only natural-language inputs and does not support pathology images, EHR data, or web-based laboratory results.Expanded multimodal support could enable more complex and specialized clinical cases.Source: Discussion · p. 20
  • Integrating clinical modalities and extended memory for patient histories could enable TXAGENT to analyze multimodal clinical data [33] and advance precision therapeutics.The framework is positioned to support personalized therapy selection and regulatory-compliant clinical decision-making.Source: Discussion · p. 21

Online Methods · 1 TXAGENT

This section describes TXAGENT, an agentic AI model for precision therapy that combines multi-step drug reasoning with verified biomedical tools and real-world knowledge. It outlines how TXAGENT dynamically selects and integrates tools, retrieves and summarizes biomedical information, and iteratively reasons over function calls to produce grounded therapeutic answers.

1.1 Overview

  • TXAGENT combines multi-step, white-box drug reasoning with verified real-world biomedical knowledge and a TOOLUNIVERSE containing 211 tools.The framework includes FDA drug information, specialized retrieval and termination tools, and supports fine-tuned open-source LLMs for local deployment and privacy protection.Source: §1.1 · p. 39
  • TXAGENT produces answers with detailed reasoning traces that expose thoughts, function calls, and tool-grounded information for verification.For a therapy question, each reasoning step contains a thought, function calls, and function responses, which are accumulated into a verified reasoning trace and final answer.Source: §1.1 · pp. 39–40
  • The framework is constructed with TOOLGEN, QUESTIONGEN, and TRACEGEN multi-agent systems that generate, verify, and filter tools, questions, and reasoning traces from biomedical sources.These systems use information and APIs from FDA documents, the OpenTarget database, and the PrimeKG graph to build the TXAGENT-INSTRUCT dataset.Source: §1.1 · p. 39

1.2 Skills of TXAGENT

  • TXAGENT’s instruction-finetuned skills support contextual thought generation, tool-argument generation, parallel tool execution, and iterative multi-step decision-making.At each step, it uses prior traces and available tools to generate thoughts and calls, evaluates whether retrieved outputs suffice, and either invokes additional tools or produces a final answer.Source: §1.2 · pp. 40–41
  • TXAGENT dynamically searches and selects tools through TOOLRAG when the available tools are insufficient, rather than relying only on tools memorized during training.TOOLRAG embeds tool descriptions and retrieves tools according to current requirements, enabling flexible and scalable tool use across the 211-tool TOOLUNIVERSE.Source: §1.2 · p. 42
  • TXAGENT embeds function-call requirements to retrieve similar tools, selects the most suitable options using prior reasoning and tool descriptions, and issues parallel calls.Retrieved tools are added to the available tool set, while selection filters candidates that are not immediately applicable.Source: §1.2 · p. 42
  • By compressing lengthy tool responses into concise summaries that preserve information relevant to the current thought, TXAGENT avoids context overflow and enables more reasoning steps.The mechanism addresses limited LLM context windows, which otherwise restrict the number of reasoning steps.Source: §1.2 · p. 42

1.3 Capabilities of TXAGENT

  • TXAGENT grounds answers in verified information retrieved through trusted-source tools, while reasoning traces let users review factual correctness and decision transparency.The system can retrieve medication dosage instructions from official FDA documents rather than generating answers directly from model memory.Source: §1.3 · p. 44
  • Goal-oriented tool selection lets TXAGENT plan actions, adapt to new tools, and expand the TOOLUNIVERSE without retraining when existing tools are insufficient.This process supports flexible handling of new scenarios by incorporating relevant tools into the tool universe.Source: §1.3 · p. 44
  • Multi-step therapeutic reasoning iteratively gathers information, generates new thoughts, and makes additional function calls until the therapeutic goal is achieved.The approach addresses complex problems requiring multiple perspectives and cases where an initial function call fails to retrieve necessary information.Source: §1.3 · p. 44
  • Direct function calls to continually updated Open Targets and FDA databases allow TXAGENT to answer questions about newly approved drugs even when training data lacks relevant information.The approach integrates complementary sources without requiring construction or maintenance of a vector database.Source: §1.3 · p. 45

2 TOOLUNIVERSE

TOOLUNIVERSE provides 211 biomedical tools for real-time information across diseases, drugs, targets, and clinical safety. TOOLGEN automates tool construction and validation, while the tool graph supports training-data construction without restricting inference-time tool integration.

2.1 Overview of TOOLUNIVERSE

  • TOOLUNIVERSE comprises 211 biomedical tools spanning diseases, drugs, targets, pharmacology, safety, clinical information, biological annotations, publications, and search.The tools use APIs from OpenFDA, OpenTargets, and the Monarch Initiative, with descriptions and backend mappings supporting TXAGENT function calls.Source: §2 · p. 45, §2.1 · p. 45

2.2 TOOLGEN: a multi-agent system for constructing tools

  • TOOLGEN converts heterogeneous API documentation into TXAGENT-compatible tools through summarization, specification generation, automated checking, and human verification.Its three agents create descriptions, arguments, and API mappings, test generated function calls, remove invalid tools, and admit validated tools to TOOLUNIVERSE.Source: §2.2 · p. 46

2.3 Tool Graph

  • The tool graph connects tools when one output can serve as another input, enabling sampled multi-call chains for constructing complex training questions.An LLM determines directed links from tool descriptions, but the graph is used only for training-data construction because precise inference-time graph construction is difficult.Source: §2.3 · p. 47
  • TXAGENT remains unrestricted by the tool graph during inference, allowing newly added tools to be integrated seamlessly.This separation preserves graph-assisted training while avoiding dependence on an exceptionally precise inference-time tool graph.Source: §2.3 · p. 47

3 Constructing TXAGENT-INSTRUCT dataset

TXAGENT-INSTRUCT is constructed to provide diverse, knowledge-grounded questions and complex, tool-integrated reasoning traces from FDA drug labels and biomedical knowledge bases. Its development combines multi-agent trace generation and verification with iterative training that jointly improves tool selection and reasoning-trace generation.

3 Constructing TXAGENT-INSTRUCT dataset

  • TXAGENT-INSTRUCT uses QUESTIONGEN to generate diverse treatment and drug questions, while TRACEGEN produces step-wise reasoning traces that integrate real-world tool feedback.The construction targets specialized treatment information and addresses the difficulty of generating valid traces for complex, tool-assisted reasoning.Source: §3 · p. 47

3.1 TXAGENT-INSTRUCT Data Sources

  • The dataset draws on OpenFDA drug labeling, Open Targets drug–disease and drug–target associations, HPO disease–phenotype links, and PrimeKG disease information.OpenFDA provides fields including indications and boxed warnings; the other resources supply complementary relationships for dataset construction.Source: §3.1 · p. 48

3.2 QUESTIONGEN Multi-agent System for Question Construction

  • QUESTIONGEN constructs drug-centered, disease-centered, and tool-chain-centered questions from extracted biomedical information, answers, explanations, and tool-retrievable references.Drug questions cover therapeutic properties and risks, disease questions incorporate patient profiles and candidate drugs, and tool-chain questions sample sequences from the TOOLUNIVERSE graph.Source: §3.2 · pp. 48–49
  • Generated questions are retained only after GPT-4o checks their knowledge grounding, answerability, and reasoning reasonableness.Questions that fail any evaluation are discarded before being sent to TRACEGEN for reasoning-trace construction.Source: §3.2 · p. 50

3.3 Reasoning Trace Generation

  • TRACEGEN addresses multi-step questions, external-tool integration, and unpredictable tool outputs when generating reasoning traces and final answers.Its design targets failure handling and continued progress when tool results deviate from expectations.Source: §3.3 · pp. 50–51
  • HELPER guides each next reasoning step using the question, ground-truth answer, explanation, and current trace, then checks whether the SOLVER’s answer matches the ground truth.When the answer is incorrect, HELPER prompts reflection and supplies another hint until the solution is consistent with the ground truth and explanation.Source: §3.3 · p. 51
  • TOOL PROVIDER first supplies tools identified from question references, then uses TOOLRAG to retrieve additional tools when the SOLVER lacks a suitable option.During trace generation, virtual calls to initially available tools are later replaced with actual TOOLRAG calls, enabling dynamic tool selection.Source: §3.3 · pp. 51–52
  • The SOLVER iteratively combines questions, prior reasoning, HELPER hints, and available tools to generate intermediate thoughts, function calls, and final answers.The process continues until the END tool proposes an answer; incorrect answers trigger removal and refinement of the corresponding reasoning step.Source: §3.3 · p. 52
  • Reasoning traces are evaluated for answer, trace, and function-call correctness, while behavior checks remove hallucinations, unverified results, and repeated thoughts or calls.Traces based on internal knowledge rather than tool feedback, including unseen drug, disease, or target IDs, are discarded; only traces passing all checks enter training data.Source: §3.3 · pp. 53–54

3.4 Iterative Training for TOOLRAG model

  • After initial training, TOOLRAG selects tools beyond the initially retrieved list, enabling more realistic traces and supplying new pairs for repeated retraining.The model is used in both TXAGENT inference and training-data collection, and the cycle continues through successive rounds.Source: §3.4 · pp. 54–55
  • The iterative process jointly refines TOOLRAG and the quality of generated reasoning traces.Generated reasoning traces provide training data that improves future trace generation and TOOLRAG performance.Source: §3.4 · pp. 54–55
  • Initially, reference information supplies the tool set used to generate reasoning traces and extract tool-requirement/tool-description training pairs.These pairs train TOOLRAG with a multiple negatives ranking loss using gte-Qwen2-1.5B-instruct [37] as the base model.Source: §3.4 · p. 54

4 Training TXAGENT model

This section describes the training of TXAGENT to support multi-step reasoning and function calling. It covers the training data and supervision, tool-focused augmentation, parameter-efficient adaptation, and distributed infrastructure used to handle large models and long contexts.

4.1 TXAGENT Training Dataset: TXAGENT-INSTRUCT dataset

  • TXAGENT-INSTRUCT contains 378,027 instruction-tuning samples built from therapeutic questions, reasoning steps, function calls, and augmented tools.The source data include 85,340 therapeutic questions, 177,626 reasoning steps, and 281,695 function calls, sampled from FDA drug labels and PrimeKG disease phenotypes while excluding drugs approved after 2023.Source: §4.1 · p. 55
  • Step-wise supervision decomposes each reasoning trace into samples that predict intermediate thoughts and function calls, then produce the final thought, FINISH call, and answer.Each intermediate input includes the system prompt, question, available tools, and prior trace; the final step additionally outputs the final answer.Source: §4.1 · p. 55

4.2 Training data augmentation

  • Tool augmentation rephrases tool fields, expands candidate sets, and shuffles tool order so TXAGENT selects functions from contextual descriptions rather than memorized names or positions.Each tool receives 20 rewritten versions of its name, descriptions, and arguments; training also includes TOOLRAG-retrieved and randomly sampled tools, supporting generalization to unseen tools and broader candidate selection.Source: §4.2 · pp. 56–57

4.3 Training design

  • The model uses instruction-tuned Llama3.1-8B-Instruct initialization with LoRA, while training minimizes autoregressive loss only on output tokens to learn thoughts and function calls.LoRA reduces computational cost and trained parameters while preserving pretrained knowledge; replacing long tool results with summaries keeps samples within the context window.Source: §4.3 · pp. 57–58

4.4 Training implementation

  • TXAGENT-8B training uses four Nvidia H100 GPUs with 320GB total memory and requires 9.93 GPU days.The infrastructure is based on TRL, Alignment Handbook, Transformers, and DeepSpeed libraries.Source: §4.4 · p. 58
  • Multi-node training uses PyTorch FSDP, whereas single-node training uses DeepSpeed FSDP, also called DeepSpeed Stage 3.These implementations provide the distributed training setup for TXAGENT.Source: §4.4 · p. 58
  • FSDP splits TXAGENT’s parameters across multiple GPUs for distributed training.This reduces memory usage while distributing model weights and computation.Source: §4.4 · p. 58
  • FSDP enables training with large backend LLMs and long context windows.The approach improves scalability while minimizing communication overhead.Source: §4.4 · p. 58

5 Benchmarking TXAGENT

TxAgent is benchmarked on five datasets spanning general drug reasoning, name-variant robustness, description-based identification, and personalized treatment selection. Evaluation uses multiple-choice and open-ended formats, with accuracy reported for both.

5.1 Benchmarks

  • The benchmark suite comprises DrugPC, BrandPC, GenericPC, DescriptionPC, and TreatmentPC, with DrugPC containing 3,168 questions across 11 therapeutic tasks.The datasets focus on FDA-approved drugs from 2024 to reduce potential exposure to evaluation information during pretraining.Source: §5.1 · p. 59
  • BrandPC and GenericPC replace DrugPC drug names with brand and generic forms to test robustness to naming variations while preserving unaffected questions.Questions requiring brand–generic conversion remain unchanged.Source: §5.1 · pp. 59–60
  • DescriptionPC replaces drug names with descriptions and evaluates drug identification before answer correctness, marking the second step incorrect when identification fails; it contains 626 questions.Descriptions include indications, mechanisms of action, contraindications, and drug interactions.Source: §5.1 · p. 60
  • TreatmentPC tests precision therapy by presenting disease treatments alongside patient-specific conditions, comorbidities, contraindications, and medication interactions requiring selection of the single suitable option.Its questions are generated from FDA information on indications, special populations, safety warnings, precautions, and contraindications.Source: §5.1 · pp. 60–61

5.2 Evaluation strategy

  • The evaluation combines multiple-choice selection and open-ended generation, converting open-ended answers into context for a subsequent multiple-choice judgment and reporting accuracy for both formats.This procedure enables quantitative evaluation despite the difficulty of directly assessing open-ended answers.Source: §5.2 · p. 61

6 Settings for analysis of TXAGENT

This section defines two analyses of TXAGENT: comparing TOOLUNIVERSE with LLM-simulated tools and removing explicit reasoning thoughts while retaining multi-step function-call inference.

6.1 TOOLUNIVERSE vs. LLM-as-tools

  • TOOLUNIVERSE is compared with LLM-as-tools by using TXAGENT-generated function-call arguments and TOOLUNIVERSE tool descriptions to prompt an LLM to simulate tool outputs.The LLM is instructed to answer from the supplied description and input, while abstaining from questions beyond its knowledge.Source: §6.1 · p. 62

6.2 Limit TXAGENT to function calls only, no thoughts

  • A thought-free TXAGENT variant directly generates function calls at each multi-step inference step instead of producing intermediate reasoning thoughts.The process initializes the reasoning trace and available tools, executes generated calls, updates tools through TOOLRAG when invoked, and finishes when textual content is produced.Source: §6.2 · p. 62

7 Prompt sketches

This section presents prompt sketches for constructing agents, generating tools and biomedical questions, and guiding TRACEGEN’s stepwise, tool-grounded reasoning. It covers schema-aligned function use, adaptive retrieval, self-reflection, and structured responses to patient- or drug-centered questions.

7 Prompt sketches

  • Prompt sketches summarize the full prompts used to build agents in the data-generation multi-agent systems while omitting unnecessary details.Source: §7 · p. 63

7.1 System prompt for TXAGENT

  • TXAGENT is instructed to solve problems through detailed, step-by-step reasoning and take actions using provided functions.The system prompt exposes the available functions through a {functions} placeholder.Source: §7.1 · p. 63

7.2 Prompts for TOOLGEN

  • TOOLGEN converts API schemas into specific natural-language functional commands and generates unique functions aligned with fields, descriptions, and capabilities.For openFDA tools, it generates paired functions for retrieving drug names from field information and retrieving field information from drug names.Source: §7.2 · p. 64
  • Generated tools represent inputs as search fields and outputs as return fields, with search-field keys serving as function parameters.Return fields are lists of field names whose information must be returned.Source: §7.2 · p. 64
  • TOOLGEN also creates diverse, specific natural-language test questions that require function use and function calls containing name and arguments.Source: §7.2 · p. 65

7.3 Prompt for QUESTIONGEN

  • QUESTIONGEN constructs disease-centered treatment cases by comparing drugs and encoding interactions, contraindications, warnings, age, population, pregnancy, and comorbidity considerations.The selected drug is made the definitive correct answer through a patient profile and comparison analysis with supporting evidence.Source: §7.3 · pp. 66–67
  • QUESTIONGEN generates expert biomedical multiple-choice questions that select the most suitable treatment for patient-specific diseases, phenotypes, genetic variation, age, comorbidities, or dosage.Additional variants generate tool-chain-centered questions from tool and PrimeKG information, or drug-centered questions from generic names, brand names, and specified fields without external information.Source: §7.3 · p. 67

7.4 Prompts for TRACEGEN

  • The Helper agent supplies suggestions for only one reasoning step and must not provide the final answer or information that directly leads to it.Source: §7.4 · p. 70
  • The prompts require self-reflection until the correct answer is found and specify final outputs containing the answer, its explanation, and the previous reasoning trace.Source: §7.4 · p. 70
  • TRACEGEN requires each reasoning step to include a thought and a correct function call, with multiple functions allowed when needed.Source: §7.4 · p. 71
  • The agent must reflect on whether it has answered the question, continue reasoning if necessary, and finish with a comprehensive answer explaining its derivation and relevance.Source: §7.4 · p. 71
  • If a function result is empty or unhelpful, TRACEGEN requires continued reasoning and ToolRAG retrieval, including virtual retrieval for listed tools and direct ToolRAG calls for unavailable tools.Source: §7.4 · p. 71
  • Answers must rely on information returned by tools rather than general knowledge, and failed attempts should prompt new solution approaches instead of repeated thoughts and calls.Source: §7.4 · p. 71
  • TRACEGEN supplies an initial Function List, structured JSON formats for intermediate and final steps, and placeholders for prior reasoning and the next-step hint.Source: §7.4 · p. 71
Loading 2503.10970v1…