RAG with over 4000 documents

[Tejaswgupta]

We have our database of over 4000 acts, each act has many sections which are stored in our Postgres database. The easiest method to answer queries seems to fetch -> re-rank -> answer . However, this can miss out essential documents if the retrieval doesn't go fine.

So the approach we're thinking of now is to build query retrievers for all the acts and pass them as tools using ReAct behavior. Though I'm not sure if this is the most efficient(or correct way) and how would the model figure out one or many right tools out of so many given.

Can someone share some resources or thoughts on how we should tackle this problem.

Thanks in advance!

[dosubot[bot]]

To effectively build and manage query retrievers for over 4000 acts in a Postgres database, you can leverage the PostgresMLRetriever and PostgresMLQueryEngine classes from the LlamaIndex framework. These classes provide a structured way to retrieve nodes based on queries, supporting both synchronous and asynchronous retrieval methods. Here is a basic outline of how you can implement and manage query retrievers:

1. Define a Custom Retriever: Use the PostgresMLRetriever class to handle the actual retrieval logic from your Postgres database.

2. Implement Retrieval Logic: Use SQL queries or an ORM like SQLAlchemy to fetch data from the Postgres database based on the query.

3. Handle Recursive Retrieval: Utilize the _handle_recursive_retrieval and _ahandle_recursive_retrieval methods to manage complex queries that may require recursive data fetching.

Here is a simplified example:

from llama_index.indices.managed.postgresml.retriever import PostgresMLRetriever
from llama_index.core.schema import QueryBundle, NodeWithScore, TextNode
from typing import List
import psycopg2

class PostgresRetriever(PostgresMLRetriever):
    def __init__(self, connection_params, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.connection_params = connection_params

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        query_str = query_bundle.query_str
        nodes = self._fetch_from_db(query_str)
        return [NodeWithScore(node=TextNode(text=node), score=1.0) for node in nodes]

    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        query_str = query_bundle.query_str
        nodes = await self._fetch_from_db_async(query_str)
        return [NodeWithScore(node=TextNode(text=node), score=1.0) for node in nodes]

    def _fetch_from_db(self, query_str: str) -> List[str]:
        conn = psycopg2.connect(**self.connection_params)
        cursor = conn.cursor()
        cursor.execute("SELECT content FROM acts WHERE content LIKE %s", (f"%{query_str}%",))
        rows = cursor.fetchall()
        cursor.close()
        conn.close()
        return [row[0] for row in rows]

    async def _fetch_from_db_async(self, query_str: str) -> List[str]:
        # Implement async database fetching logic here
        pass

# Usage
connection_params = {
    'dbname': 'your_db',
    'user': 'your_user',
    'password': 'your_password',
    'host': 'your_host',
    'port': 'your_port'
}

retriever = PostgresRetriever(connection_params)
query_bundle = QueryBundle("example query")
results = retriever.retrieve(query_bundle)

Additionally, you can explore the PostgresMLQueryEngine class to enhance your retrieval logic based on specific requirements. This class offers advanced capabilities for handling SQL queries and can be integrated into your custom retriever to improve performance and accuracy:

from llama_index.indices.managed.postgresml.query import PostgresMLQueryEngine
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.schema import QueryBundle

# Initialize the retriever
retriever = PostgresRetriever(connection_params)

# Initialize the query engine
query_engine = PostgresMLQueryEngine(
    retriever=retriever,
    streaming=False,
    callback_manager=CallbackManager([])
)

# Use the query engine to retrieve and synthesize responses
query_bundle = QueryBundle("example query")
response = query_engine._query(query_bundle)

This setup ensures that your query retrieval process is both efficient and accurate, leveraging the full capabilities of the LlamaIndex framework ^{[1][2][3][4][5]}.

To continue talking to Dosu, mention @dosu.