DARE: Aligning LLM Agents with the R Statistical Ecosystem via Distribution-Aware Retrieval

Large Language Model (LLM)-based agents have shown strong capabilities in automating data science workflows. However, many rigorous statistical methods implemented in the R ecosystem remain underutilized due to limitations in model knowledge and tool retrieval. Existing Retrieval-Augmented Generation (RAG) approaches rely primarily on function-level semantic similarity and often ignore data distribution characteristics that determine statistical method applicability, leading to inaccurate retrieval. To address this limitation, we propose DARE (Distribution-Aware Retrieval Embedding) a lightweight and plug-and-play retrieval model that incorporates data distribution into function representations for R package retrieval. Our main contributions are threefold: (i) We construct RPKB, a curated knowledge base of 8,191 high-quality R packages with functions spanning diverse statistical domains. (ii) We introduce DARE, an embedding model that incorporates data distribution for accurate R package retrieval. (iii) We design RCodingAgent, an R-oriented LLM agent, together with a suite of R-based statistical analysis tasks to systematically evaluate LLM agents under realistic analytical scenarios. Experimental results show that DARE achieves an NDCG@10 of 93.47%, outperforming open-source state-of-the-art embedding models by up to 17% on R package retrieval while using only 23M parameters. Furthermore, integrating DARE into RCodingAgent significantly improves downstream statistical analysis performance across multiple analytical tasks. This work bridges modern LLM-based automation with established statistical computing ecosystems.