PaperlayerView exact revision on arXiv

Source-linked AI summary

DrugPilot: LLM-based Parameterized Reasoning Agent for Drug Discovery

Kun Li, Zhennan Wu, Shoupeng Wang, Jia Wu, Shirui Pan, Wenbin Hu

arXiv:2505.13940v2cs.AIq-bio.BM

TL;DR

LLM agents for drug discovery face challenges handling multimodal data, automating multi-stage workflows, and supporting domain-specific tools. DrugPilot combines parameterized memory with tool-calling mechanisms, and its benchmark results show high task completion rates across three categories. The paper concludes that the framework has potential to reduce technical barriers and improve cross-domain collaboration in drug discovery.

Abstract

from arXiv · show

Large language models (LLMs) integrated with autonomous agents hold significant potential for advancing scientific discovery through automated reasoning and task execution. However, applying LLM agents to drug discovery is still constrained by challenges such as large-scale multimodal data processing, limited task automation, and poor support for domain-specific tools. To overcome these limitations, we introduce DrugPilot, a LLM-based agent system with a parameterized reasoning architecture designed for end-to-end scientific workflows in drug discovery. DrugPilot enables multi-stage research processes by integrating structured tool use with a novel parameterized memory pool. The memory pool converts heterogeneous data from both public sources and user-defined inputs into standardized representations. This design supports efficient multi-turn dialogue, reduces information loss during data exchange, and enhances complex scientific decision-making. To support training and benchmarking, we construct a drug instruction dataset covering eight core drug discovery tasks. Under the Berkeley function-calling benchmark, DrugPilot significantly outperforms state-of-the-art agents such as ReAct and LoT, achieving task completion rates of 98.0%, 93.5%, and 64.0% for simple, multi-tool, and multi-turn scenarios, respectively. These results highlight DrugPilot's potential as a versatile agent framework for computational science domains requiring automated, interactive, and data-integrated reasoning.

Results

DrugPilot combines multi-stage task execution with structured parameter handling, and its reported results cover tool-calling performance and high-volume molecule processing. Its strongest comparisons are higher accuracy across tool-calling categories and stable parameter handling where other methods encounter limits.

DrugPilot’s Framework

  • DrugPilot combines an LLM, parameterized memory pool, Fe-Fo error feedback, and an AI model zoo to plan and execute multi-stage tasks from user queries.Users can upload files, save dialogue parameters and results in the memory pool, and receive natural-language conclusions and visualizations.Source: DrugPilot’s Framework · pp. 5–6

Superior Performance over SOTA LLMs and Agents

  • DrugPilot outperforms all agents across simple, multi-function, and multi-turn categories, with accuracy exceeding 70% and 60% in the most complex category, respectively.DrugPilot also exceeds ReAct by 13.6% and 13.2% in simple function, 29.4% and 66.1% in multi-function, and 61.2% and 80.3% in multi-turn function on Acc.F and Acc.P.Source: Superior Performance over SOTA LLMs… · p. 6
  • DrugPilot reduces multi-turn latency to under 20 seconds, more than halving runtime for Qwen2, DeepSeek-R1, Llama3, and Llama3.1 relative to baseline methods.The tasks require sequential reasoning based on prior results.Source: Superior Performance over SOTA LLMs… · pp. 6–7, Figure 3 · p. 7

Parameterized Memory Pool enables High-volume Drug Data Processing

  • ChatGPT-4o stably selected the correct tool, passed parameters, and completed tasks for ≤51 molecules averaging ≤90 in length, but failed beyond either threshold; DrugPilot passed parameters for 91 molecules averaging 52.DrugPilot can theoretically handle parameters of any scale.Source: Parameterized Memory Pool enables… · p. 8

Case Studies on Two Classical Drug Discovery Tasks

DrugPilot case studies assess prediction capabilities on drug response and molecular property tasks, alongside ablations of its core components. The reported results include aligned predictions in both tasks and improvements in tool-calling performance and molecule handling.

Case Studies on Two Classical Drug Discovery Tasks

  • DrugPilot’s case studies evaluate model invocation on GDSCv2 drug response prediction and BACE molecular property prediction, following prior tests of task selection and parameter extraction.The study identifies reliable model invocation as a second capability needed alongside accurate module selection and parameter extraction.Source: Case Studies on Two Classical Drug… · p. 8

Model Capabilities for the DRP Task.

  • Across five held-out drug-cell-line tests totaling 26 samples, DrugPilot predictions follow the ascending order of actual measurements, with errors distributed around zero.The tests used GDSCv2 samples excluded from model training, and CLDR as the predictive engine.Source: Model Capabilities for the DRP Task · pp. 9–10, Figure 4 · p. 9

Model Capabilities for the MPP Task.

  • DrugPilot’s BACE inhibitor predictions largely match ground-truth labels, and t-SNE visualizations show alignment between true BACEi and high-probability predictions.The classification evaluation used the top and bottom five molecules; the t-SNE view covers the top and bottom 20 molecules.Source: Figure 4 · p. 9, Model Capabilities for the MPP Task · p. 10

Effect of SFT and Fe-Fo Mechanism.

  • With SFT and Fe-Fo, tool-selection accuracy reaches 95.0% and parameter-extraction accuracy reaches 93.7%, while latency falls from 30.91 to 15.18 seconds.SFT and Fe-Fo separately improve accuracy over the baseline, with larger parameter-extraction gains when paired with the other component.Source: Effect of SFT and Fe-Fo Mechanism · pp. 10–11

Effect of PMP

  • DrugPilot’s PMP maintains accuracy and response time as molecule inputs grow, while other methods degrade or become nonfunctional at larger inputs.Other methods’ accuracy drops significantly above 15 molecules; CoT essentially fails at 20, and DrugPilot without PMP, ReAct, and LoT become nonfunctional around 30 molecules.Source: Effect of PMP · p. 11

Study of key components

  • DrugPilot’s PMP remains stable as molecule counts rise, unlike methods that lose accuracy or functionality at larger inputs.The test increased molecules per query and measured parameter recognition, extraction accuracy, and latency under equivalent memory-pool inputs.Source: Effect of PMP · p. 11

Discussion

DrugPilot combines parameterized reasoning and error correction to address multimodal data alignment and fragmented drug-discovery workflows. On TCDD, it outperforms baseline methods across three tool-calling categories, while data construction and task extensibility remain future challenges.

Discussion

  • DrugPilot targets non-expert access, multimodal data alignment, and workflow fragmentation, which otherwise requires researchers to integrate tasks manually.The paper identifies alignment across drug-related data and fragmented workflows as persistent limitations of existing tools.Source: Discussion · p. 11
  • The parameterized memory pool converts drug-related entities such as SMILES and protein sequences into actionable key-value pairs for structured data transfer.The paper describes this representation as addressing multimodal data fragmentation and enabling lossless information transfer.Source: Discussion · p. 11
  • Above 95% task completion across inputs of 1–30 SMILES strings, and theoretically unlimited parameter handling versus GPT-4o's 51-parameter maximum.The parameter-handling capacity is described as theoretical.Source: Discussion · p. 13
  • 98.0% for simple function, 93.5% for multiple function, and 64.0% for multi-turn function: DrugPilot outperforms all baseline methods in each category.The evaluation uses TCDD, which covers eight core tasks and 2,800 annotated samples.Source: Discussion · p. 13
  • Data construction and task extensibility remain future challenges, despite the paper's conclusion that DrugPilot may reduce technical barriers and improve cross-domain collaboration.The authors also identify reinforcement learning for dynamic task planning and broader intelligent assistance as future directions.Source: Discussion · p. 13

Methods

DrugPilot combines an LLM, a parameterized memory pool, a feedback-focus mechanism, and an AI model zoo. It reasons over stored parameters, verifies tool calls, executes them, and uses feedback to correct detected errors.

Methods

  • DrugPilot combines TCDD-based fine-tuning, a parameterized memory pool for large-scale multimodal data and multi-turn conversations, and feedback-focus error correction.The design also includes LLM backbones and an AI model zoo.Source: Methods · p. 13
  • The LLM reads the PMP and decides which tools to call; Fe-Fo verifies calls before execution and returns feedback when it detects errors.Tool results return to the LLM, which continues reasoning until it produces a final answer.Source: Methods · pp. 13–14
  • Users can view or update PMP parameters during task execution.The system continues reasoning from the PMP and tool results until it produces the final answer.Source: Methods · pp. 13–14

Tool-Calling Dataset for Drug Discovery

TCDD is a drug-discovery tool-calling dataset designed to train and evaluate tool selection, parameter extraction, and memory-pool interaction. Its samples cover single- and multi-turn dialogues, error correction, and memory-pool updates.

Tool-Calling Dataset for Drug Discovery

  • The dataset tests tool selection, extraction of parameters for unique data forms, and retrieval of PMP data and saving of execution results.These capabilities correspond to selecting tools from natural-language queries, extracting required parameters, and interacting with the memory pool.Source: Tool-Calling Dataset for Drug Discovery · p. 14
  • TCDD contains 2,800 instruction samples—2,500 for training and 300 for testing—in ShareGPT format, spanning single-turn, multi-turn, error-correction, and memory-pool-update scenarios.The dataset is designed to fine-tune and test LLMs on drug-discovery tool calling.Source: Tool-Calling Dataset for Drug Discovery · p. 14
  • Each sample includes conversation, system instructions, and tool information; the conversation represents user input, JSON tool calls, returned observations, and the LLM response.The tools section specifies function names, descriptions, parameter types, and required parameters.Source: Tool-Calling Dataset for Drug Discovery · p. 14
  • About 50% of samples are single-turn dialogues, while multi-turn tasks account for 30% and parameter-error scenarios for 20%.The multi-turn scenarios model workflows that integrate intermediate results, retrieve parameters from the memory pool, and handle tool errors.Source: Tool-Calling Dataset for Drug Discovery · p. 15
  • Sample distributions vary by tool: optimization and generation have relatively more multi-turn examples, while target-affinity and target-interaction tools have more error-correction samples.Optimization also has a higher proportion of error-handling samples because it uses more parameters than other tools.Source: Tool-Calling Dataset for Drug Discovery · p. 15

Parameterized Memory Pool

DrugPilot’s parameterized memory pool (PMP) stores drug-related parameters as structured key-value pairs rather than conversation text, allowing the agent and users to access and update them during multi-stage workflows.

The Structure of PMP

  • PMP stores drug-related parameters as structured key-value pairs, with unique short keys and values formatted for direct use by drug-related tools.Parameters of the same type share a key, with their values stored in a list.Source: Parameterized Memory Pool · p. 15, The Structure of PMP · pp. 15–16
  • For large parameter lists, PMP replaces repeated interaction with complex data by short keys, while tools can use the stored values directly.The paper gives lists of tens of thousands of molecules as an example of large-scale parameters.Source: The Structure of PMP · pp. 15, 17

Parameter Reading and Updating

Feedback-focus Mechanism

DrugPilot’s feedback-focus (Fe-Fo) mechanism addresses tool-calling errors and loss of task focus in extended dialogues by feeding error information back to the LLM and reiterating the original task.

Feedback-focus Mechanism

Evaluation Metrics

The evaluation uses three tool-calling scenarios, compares DrugPilot with representative tool-calling agents, and assesses tool selection and parameter extraction through staged checks of each response.

Evaluation Metrics

  • The test set has 300 queries split evenly across simple-function, multiple-function, and multi-turn-function categories.The categories respectively involve one supplied function, choosing among eight tools, and multi-turn queries that may require different tool calls.Source: Evaluation Metrics · p. 19

Baselines

  • The comparison includes CoT, LoT, ReAct, and ChatGPT-4o Function Calling API as representative tool-calling baselines.The paper selects these agents because its focus is accurate tool calling across a drug-discovery workflow.Source: Baselines · p. 19

Fine-Tuning

  • DrugPilot fine-tunes small-scale LLMs on TCDD using LoRA and combines glaive_toolcall data with an 8:1:1 TCDD train-validation-test split.Source: Fine-Tuning · pp. 19–20

Evaluation Process

  • Responses are evaluated for tool selection and parameter extraction, with a 120-second per-query limit and staged checks of tool name, required and unexpected parameters, and parameter values.The checks correspond to function accuracy and parameter accuracy; exceeding the time limit counts as an error.Source: Evaluation Process · p. 20

Appendix A Memory Pool Prompt

The memory pool prompt tells the LLM how to use PMP keys to retrieve stored values and distinguishes retrieval from direct extraction of parameters already present in the input.

Appendix A Memory Pool Prompt

  • The prompt describes the LLM’s responsibility for transmitting parameters to tools and presents the user question or tool output alongside the current list of PMP keys.Source: Appendix A · p. 27
  • When the input already contains required parameters, the prompt directs the LLM to extract them directly; otherwise it directs the LLM to retrieve a matching PMP key.Correct and incorrect examples are included to discourage retrieval of nonexistent keys.Source: Appendix A · p. 27

Appendix B DrugPilot’s Reasoning Errors

The appendix describes JSON-formatted tool calls and common reasoning errors arising during interactions with the parameterized memory pool (PMP). Figure B2 presents these error types and the feedback provided by Fe-Fo.

Appendix B DrugPilot’s Reasoning Errors

  • Tool calls require an action input in JSON containing the tool name and required parameters.The appendix notes that interactions with PMP can introduce content- and format-related problems.Source: Appendix B · pp. 27–28
  • Frequent PMP interactions can lead to problems in the content and format of tool-calling outputs.These problems motivate classifying common reasoning errors.Source: Appendix B · p. 28
  • Figure B2 summarizes common LLM reasoning errors during drug-related tool calls and the feedback Fe-Fo provides.The figure concerns errors when LLMs call drug-related tools.Source: Figure B2 · p. 28

Appendix C Fine-Tuning Configuration

The appendix reports the LoRA fine-tuning setup for DrugPilot’s LLMs and points to Table C1 for hyperparameter settings.

Appendix C Fine-Tuning Configuration

  • LoRA fine-tuning was used to enhance the LLMs’ drug-discovery domain knowledge and tool-calling ability.The final inference-stage LLMs were deployed on Ollama.Source: Appendix C · p. 28
  • Fine-tuning used batch size 4 for smaller models and 8 for larger models.The appendix distinguishes the batch size by model size.Source: Appendix C · p. 28
  • Table C1 contains the hyperparameter settings used during fine-tuning.The final inference-stage LLMs were deployed on Ollama.Source: Table C1 · p. 28, Appendix C · p. 28

Appendix D Accuracy Calculation

The appendix defines Acc.F and Acc.P as accuracy measures for tool selection and parameter extraction, respectively, using per-category samples and correctness indicators.

Appendix D Accuracy Calculation

  • Acc.F measures tool-selection accuracy, while Acc.P measures parameter-extraction accuracy.Both metrics are defined for tool-calling evaluation.Source: Appendix D · p. 29
  • Nc denotes the number of samples in a category.The accuracy calculation uses the category sample count.Source: Appendix D · p. 29
  • F(·) and P(·) indicate whether a sample has the correct function selected and parameters extracted, respectively.These indicator functions are used in calculating the two accuracy measures.Source: Appendix D · p. 29
Loading 2505.13940v2…