Can we control number of documents to return in RetrievalQA Langchain

Question

Can we control the document query parameter in RetrievalQA() like we could do in vectorDBQA() in langchain before? Also, shall I use map_reduce chain type instead for my large documents?

I tried to look into the source code but could not find it.

Answer 1

You can customized the Retrieval and add your requirement

class URRetrieval(BaseRetriever):

    def __int__(self):
        pass

    def setparms(self, count):
        self.count = count

    def _get_relevant_documents(
            self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        
        # need to add your similarity search here and add set self.count as limit
    #results = your vectordb.similarity(limit_or_k = self.count)
        return results

    async def _aget_relevant_documents(
            self,
            query: str,
            *,
            run_manager: AsyncCallbackManagerForRetrieverRun,
            **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError()

In the above example, you need to add your vector db, it depends on you which you are using, and return the similarity results, the vector db supported for k/limit of relative docs.

_get_relevant_documents : method call when inside the chain, you need to pass this custom retrieval as chain retrieval.

Example:

retrieval = URRetrieval()
retrieval.setparms(5)
qa_with_sources_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retrieval,
    callbacks=[handler],