Milvus
Zilliz

How to add custom modules to UltraRag?

Adding custom modules to UltraRag, particularly in version 2.0 and later, is facilitated by its modular architecture, which is built upon the Model Context Protocol (MCP). This design allows developers to integrate new functionalities as independent components, orchestrated primarily through YAML configuration files. The framework is engineered for extensibility, enabling “hot-plugging” of custom modules similar to a plugin system, without requiring modifications to the core codebase. This approach significantly lowers the technical barrier for researchers and developers to experiment with new RAG algorithms and integrate specialized functionalities.

The process of adding a custom module to UltraRag generally involves encapsulating the custom logic within an MCP Server and defining its interface as Tools. UltraRag standardizes core RAG components, such as Retriever, Generator, and Evaluator, as independent MCP Servers, which can be flexibly extended. For a custom module, this means developing Python code that implements the desired functionality, registering it as a Server, and defining its specific functions as Tools that the UltraRag client can invoke. The framework supports defining the module’s parameters in a parameter.yaml file and its source code in a .py file, as seen in examples where a custom module might reference parameter.yaml and src/custom.py. Once the custom logic is developed and structured as an MCP Server with its Tools, the ultrarag build command can be used to automatically compile these servers and generate the necessary configuration files.

Finally, the integration and orchestration of these custom modules within a RAG pipeline are managed through UltraRag’s declarative YAML configuration. Complex workflows, including sequential, loop, and conditional branching logic, are defined entirely within these YAML files. This allows developers to specify how the custom module interacts with other components, such as a retriever utilizing a vector database like Milvus for efficient similarity search or a generator processing information. By defining the input and output expectations for each Tool within the YAML, UltraRag’s built-in MCP Client can coordinate tool calls and data transfer between Servers, ensuring precise execution of the custom pipeline. This clear separation of concerns, where functional logic resides in modular servers and workflow control in YAML, makes it straightforward to add, modify, and debug custom functionalities within UltraRag.

Like the article? Spread the word