Retrievers#

Retrievers are an important component of AI workflows utilizing Retrieval Augmented Generation (RAG) which allow LLMs to search a data store for content which is semantically similar to a query which can be used as context by the LLM when providing a response to the query. Within AgentIQ, retrievers are a configurable component which can be used within functions, similar to LLMs and Embedders, to provide a consistent read-only interface for connecting to different data store providers.

Features#

  • Standard Interface: Retrievers implement a standard search interface, allowing for compatibility across different retriever implementations.

  • Standard Output Format: Retrievers also implement a standard output format along with conversion functions to provide retriever output as a dictionary or string.

  • Extensible Via Plugins: Additional retrievers can be added as plugins by developers to support more data stores.

  • Additional Framework Implementations: Retrievers can be loaded using a framework implementation rather than the default AgentIQ retriever implementation.

Included Retrievers#

  • Milvus

  • NeMo Retriever

Usage#

Configuration#

Retrievers are configured similarly to other AgentIQ components, such as Functions and LLMs. Each Retriever provider (e.g., Milvus) has a Pydantic config object which defines its configurable parameters and type. These parameters can then be configured in the config file under the retrievers section.

Below is an example config object for the NeMo Retriever:

class NemoRetrieverConfig(RetrieverBaseConfig, name="nemo_retriever"):
"""
Configuration for a Retriever which pulls data from a Nemo Retriever service.
"""
uri: HttpUrl = Field(description="The uri of the Nemo Retriever service.")
collection_name: str | None = Field(description="The name of the collection to search", default=None)
top_k: int | None = Field(description="The number of results to return", gt=0, le=50, default=None)
output_fields: list[str] | None = Field(
    default=None,
    description="A list of fields to return from the datastore. If 'None', all fields but the vector are returned.")
timeout: int = Field(default=60, description="Maximum time to wait for results to be returned from the service.")
nvidia_api_key: str | None = Field(
    description="API key used to authenticate with the service. If 'None', will use ENV Variable 'NVIDIA_API_KEY'",
    default=None,
)

This retriever can be easily configured in the config file such as in the below example:

retrievers:
    my_retriever:
        _type: nemo_retriever
        uri: http://my-nemo-service-url
        collection_name: "test_collection"
        top_k: 10

In this example the uri, collection_name, and top_k are specified, while the default values for output_fields and timeout are used, and the nvidia_api_key will be pulled from the NVIDIA_API_KEY environment variable.

This configured retriever can then be used as an argument for a function which uses a retriever (such as the aiq_retriever function). The aiq_retriever function is a simple function to provide the configured retriever as an LLM tool. Its config is shown below

class AIQRetrieverConfig(FunctionBaseConfig, name="aiq_retriever"):
    """
    AIQRetriever tool which provides a common interface for different vectorstores. Its
    configuration uses clients, which are the vectorstore-specific implementaiton of the retriever interface.
    """
    retriever: RetrieverRef = Field(description="The retriever instance name from the workflow configuration object.")
    raise_errors: bool = Field(
        default=True,
        description="If true the tool will raise exceptions, otherwise it will log them as warnings and return []",
    )
    topic: str = Field(default=None, description="Used to provide a more detailed tool description to the agent")
    description: str = Field(default=None, description="If present it will be used as the tool description")

Here is an example configuration of an aiq_retriever function that uses a nemo_retriever:

retrievers:
    my_retriever:
        _type: nemo_retriever
        uri: http://my-nemo-service-url
        collection_name: "test_collection"
        top_k: 10

functions:
    aiq_retriever_tool:
        _type: aiq_retriever
        retriever: my_retriever
        topic: "AIQ documentation"

Developing with Retrievers#

Alternatively, you can use a retriever as a component in your own function, such as a custom built RAG workflow. When building a function that uses a retriever you can instantiate the retriever using the builder. Like other components, you can reference the retriever by name and specify the framework you want to use. Unlike other components, you can also omit the framework to get an instance of an AIQRetriever.

@register_function(config_type=MyFunctionConfig)
async def my_function(config: MyFunctionConfig, builder: Builder):

    # Build an AIQRetriever
    aiq_retriever = await builder.get_retriever(config.retriever)

    # Build a langchain Retriever
    langchain_retriever = await builder.get_retriever(config.retriever, wrapper_type=LLMFrameworkEnum.LANGCHAIN)

Retrievers expose a search method for retrieving data that takes a single required argument, “query”, and any number of optional keyword arguments. AgentIQ Retrievers support a bind method which can be used to set or override defaults for these optional keyword arguments. Any additional required, unbound, parameters can be inspected using the get_unbound_params method. This provides flexibility in how retrievers are used in functions, allowing for all search parameters to be specified in the config, or allowing some to be specified by the agent when the function is called.

Adding a Retriever Provider#

New retrievers can be added to AgentIQ by creating a plugin. The general process is the same as for most plugins, but the retriever-specific steps are outlined here.

First, create a retriever for the provider that implements the Retriever interface:

class AIQRetriever(ABC):
    """
    Abstract interface for interacting with data stores.

    A Retriever is resposible for retrieving data from a configured data store.

    Implemntations may integrate with vector stores or other indexing backends that allow for text-based search.
    """

    @abstractmethod
    async def search(self, query: str, **kwargs) -> RetrieverOutput:
        """
        Retireve max(top_k) items from the data store based on vector similarity search (implementation dependent).

        """
        raise NotImplementedError

Next, create the config for the provider and register it with AgentIQ:

class ExampleRetrieverConfig(RetrieverBaseConfig, name="example_retriever"):
    """
    Configuration for a Retriever provider. The parameters will depend on the particular provider. These are examples.
    """
    uri: HttpUrl = Field(description="The uri of the Nemo Retriever service.")
    collection_name: str = Field(description="The name of the collection to search")
    top_k: int = Field(description="The number of results to return", gt=0, le=50, default=5)
    output_fields: list[str] | None = Field(
        default=None,
        description="A list of fields to return from the datastore. If 'None', all fields but the vector are returned.")


@register_retriever_provider(config_type=ExampleRetrieverConfig)
async def example_retriever(retriever_config: ExampleRetrieverConfig, builder: Builder):
    yield RetrieverProviderInfo(config=retriever_config,
                                description="AIQ retriever provider for...")

Lastly, implement and register the retriever client:

@register_retriever_client(config_type=ExampleRetrieverConfig, wrapper_type=None)
async def nemo_retriever_client(config: ExampleRetrieverConfig, builder: Builder):
    from example_plugin.retriever import ExampleRetriever

    retriever = ExampleRetriever(**config.model_dump())

    yield retriever

You can then implement and register framework-specific clients for the retriever provider, our use the config to instantiate an existing framework implementation.