Agent-UniRAG: A Trainable Open-Source LLM Agent Framework for Unified Retrieval-Augmented Generation Systems

2025-05-28

AI Summary

This paper addresses key challenges in Retrieval-Augmented Generation (RAG) systems—namely, the lack of unified modeling for single-hop and multi-hop queries, poor interpretability, and difficulty in adapting to small language models (SLMs). To this end, we propose the first end-to-end trainable LLM agent framework capable of natively handling both query types within a single architecture. Our approach integrates dynamic reasoning path planning, synthetic data distillation, and lightweight fine-tuning, yielding SynAgent-RAG—a high-quality, SLM-optimized synthetic dataset. We demonstrate efficient adaptation on open-source small models such as Llama-3-8B. Key contributions include: (1) the first unified modeling framework for single- and multi-hop RAG; (2) substantial improvements in retrieval accuracy and reasoning transparency; and (3) competitive performance against larger closed- and open-source models across mainstream RAG benchmarks. All code and datasets are publicly released.

Technology Category

• Cognitive Modeling & Cognitive Systems: Agent Architectures
• Multiagent Systems: Modeling other Agents
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

This paper presents a novel approach for unified retrieval-augmented generation (RAG) systems using the recent emerging large language model (LLM) agent concept. Specifically, Agent LLM, which utilizes LLM as fundamental controllers, has become a promising approach to enable the interpretability of RAG tasks, especially for complex reasoning question-answering systems (e.g., multi-hop queries). Nonetheless, previous works mainly focus on solving RAG systems with either single-hop or multi-hop approaches separately, which limits the application of those approaches to real-world applications. In this study, we propose a trainable agent framework called Agent-UniRAG for unified retrieval-augmented LLM systems, which enhances the effectiveness and interpretability of RAG systems. The main idea is to design an LLM agent framework to solve RAG tasks step-by-step based on the complexity of the inputs, simultaneously including single-hop and multi-hop queries in an end-to-end manner. Furthermore, we introduce SynAgent-RAG, a synthetic dataset to enable the proposed agent framework for small open-source LLMs (e.g., Llama-3-8B). The results show comparable performances with closed-source and larger open-source LLMs across various RAG benchmarks. Our source code and dataset are publicly available for further exploitation.

Problem

• Unified RAG systems for single and multi-hop queries
• Enhancing interpretability in complex reasoning QA systems
• Trainable agent framework for small open-source LLMs

Innovation

• Trainable LLM agent for unified RAG tasks
• Handles single-hop and multi-hop queries end-to-end
• Uses synthetic dataset for small open-source LLMs