Use Retriever as Filter to Another Retriever

[KavinKKuppusamy]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

def load_faiss_retriever(context_name, metadata_filter_retriever = None):

    if metadata_filter_retriever is not None:
        metadata_filter_retriever_dict = metadata_filter_retriever
    faiss_dir = faiss_dir_path
    faiss_k_value = 5

    vectorstore = FAISS.load_local(
        faiss_dir, embedding_model, allow_dangerous_deserialization=True
    )
    if metadata_filter_retriever is not None:
        faiss_retriever = vectorstore.as_retriever(
            search_type="similarity",
            search_kwargs={"k": faiss_k_value, "include_metadata": True, "filter": metadata_filter_retriever_dict},
        )
    else:
        faiss_retriever = vectorstore.as_retriever(
            search_type="similarity",
            search_kwargs={"k": faiss_k_value, "include_metadata": True},
            )

    return faiss_retriever

class TableRetriever(BaseRetriever):
    base_retriever: BaseRetriever

    def _get_relevant_documents(
        self, 
        query: str,
        * ,
        run_manager: CallbackManagerForRetrieverRun
    ) -> Dict[str, list]:
        """
        Retrieve documents and extract table names into a dictionary format.
        
        :param query: The search query
        :param run_manager: Callback manager for the retriever run
        :return: Dictionary with unique table names
        """
        # Retrieve documents from the base retriever
        documents: List[Document] = self.base_retriever.invoke(query)
        
        table_list = []
        for doc in documents:
            page_content = doc.page_content
            
            # Extract the table name
            start_index = page_content.find("table_name : ") + len("table_name : ")
            end_index = page_content.find(",", start_index)
            if start_index > -1 and end_index > -1:  # Ensure indices are valid
                table_name = page_content[start_index:end_index].strip()
                table_list.append(table_name)
        
        # Create metadata filter with unique table names
        metadata_filter = {"table_name": list(set(table_list))}        
        return metadata_filter
    
def get_retriever():
   table_retriever = load_faiss_retriever("TABLE_DDL")
   table_retriever_filter = TableRetriever(base_retriever=table_retriever)
   column_retriever = load_faiss_retriever("COLUMN_INFO", table_retriever_filter)
   return column_retriever


history_aware_retriever = create_history_aware_retriever(
    chat_model, get_retriever(), rephrase_prompt
)



Below is for testing

table_retriever = load_faiss_retriever("TABLE_DDL")
table_retriever_filter = TableRetriever(base_retriever=table_retriever)
column_retriever = load_faiss_retriever("COLUMN_INFO", table_retriever_filter)
result = column_retriever.get_relevant_documents(query="texas")

Description

I call FAISS retriever for TABLE_DDL and extract all the table names and use the table_names as filter for next Retriever. How to do this in lang chain. I tried many ways nothing seem to be working. Can someone please help?

System Info

%pip install langchain_community==0.0.32
%pip install langchain==0.1.16
%pip install faiss-cpu==1.8.0
%pip install langchain-openai==0.1.3