[Paper Review] Machine Learning as a Tool (MLAT): A Framework for Integrating Statistical ML Models as Callable Tools within LLM Agent Workflows

Edwin Chen, Zulekha Bibi|arXiv (Cornell University)|Feb 15, 2026

Artificial Intelligence in Healthcare and Education0 citations

TL;DR

MLAT exposes pre-trained statistical ML models as callable tools in LLM agent workflows, enabling contextual, rationale-backed predictions integrated into structured outputs; demonstrated via PitchCraft pricing in a dual-agent Gemini setup with a small data regime.

ABSTRACT

We introduce Machine Learning as a Tool (MLAT), a design pattern in which pre-trained statistical machine learning models are exposed as callable tools within large language model (LLM) agent workflows. This allows an orchestrating agent to invoke quantitative predictions when needed and reason about their outputs in context. Unlike conventional pipelines that treat ML inference as a static preprocessing step, MLAT positions the model as a first-class tool alongside web search, database queries, and APIs, enabling the LLM to decide when and how to use it based on conversational context. To validate MLAT, we present PitchCraft, a pilot production system that converts discovery call recordings into professional proposals with ML-predicted pricing. The system uses two agents: a Research Agent that gathers prospect intelligence via parallel tool calls, and a Draft Agent that invokes an XGBoost pricing model as a tool call and generates a complete proposal through structured outputs. The pricing model, trained on 70 examples combining real and human-verified synthetic data, achieves R^2 = 0.807 on held-out data with a mean absolute error of 3688 USD. The system reduces proposal generation time from multiple hours to under 10 minutes. We describe the MLAT framework, structured output architecture, training methodology under extreme data scarcity, and sensitivity analysis demonstrating meaningful learned relationships. MLAT generalizes to domains requiring quantitative estimation combined with contextual reasoning.

Motivation & Objective

• Formalize the MLAT design pattern where a pre-trained ML model is exposed as a tool within an LLM agent registry.
• Demonstrate end-to-end MLAT implementation in a production-like system (PitchCraft) to generate proposals with ML-predicted pricing.
• Show how structured output parsing bridges LLM reasoning with ML feature vectors using JSON schemas.
• Evaluate MLAT under extreme data scarcity with small real+synthetic datasets and group-aware validation.

Proposed method

• Define MLAT as an agent-controlled tool-calling pattern where the LLM extracts a feature vector from structured context and calls a trained ML model to obtain a prediction.
• Register the ML model as a stateless REST endpoint tool and use schema-constrained extraction and output schemas to bridge LLM reasoning and ML inputs.
• Use Gemini's JSON schema constraints to enable reliable structured output parsing and inter-agent contracts between Research and Draft agents.
• Train an XGBoost regression model on a small dataset (N=70) with 40 real and 30 synthetic records, employing group-aware cross-validation and feature engineering (8 features, one-hot encoding of tech_stack).
• Perform a sensitivity analysis and a cross-validation-driven performance assessment to validate learned economic relationships.

Figure 2 : Target variable distribution across training and test sets. Left: Histogram showing the right-skewed price distribution. Training (blue) and test (orange) sets show comparable coverage. Right: Box plots confirming similar median and interquartile ranges ( $n_{\text{train}}=56$ , $n_{\text

Experimental results

Research questions

• RQ1 Can exposing a trained ML model as a tool within an LLM agent workflow improve contextual decision-making and interpretability of predictions?
• RQ2 How does the MLAT pattern perform in a low-data regime with synthetic augmentation, in terms of predictive accuracy and generalization?
• RQ3 Does structured output parsing enable reliable feature extraction and inter-agent communication between researchers and drafting agents?
• RQ4 What is the impact of MLAT on real-world tasks such as proposal generation and pricing in a production-like setting?

Key findings

• The XGBoost pricing model achieves R^2 = 0.807 on held-out test data with MAE = 3,688 and RMSE = 4,720, under a 70-sample regime.
• Cross-validation R^2 is 0.816 (±0.060), indicating reliable generalization despite small data and synthetic augmentation.
• The complete PitchCraft pipeline reduces proposal creation time from 3+ hours to under 10 minutes, and speeds up speed-to-lead by 12–18×.
• Sensitivity analysis shows economically coherent price increases with higher pain severity and integration complexity, suggesting learned meaningful relationships rather than memorization.
• Compared to Ridge regression, XGBoost shows substantially better CV R^2 (0.816±0.060 vs. 0.565±0.180), confirming nonlinear feature interactions matter.

Figure 3 : Predicted vs. actual price. Left: Training set ( $R^{2}=0.937$ ) shows tight clustering around the identity line. Right: Test set ( $R^{2}=0.807$ ) demonstrates generalization to unseen data, with slight overestimation in the mid-range and underestimation at the highest values—consistent