How can I set the max number of docs (k) that I want a retriever to retrieve?

[fantauzzi]

I am using a retriever within a chain, and would like the retriever to limit the number of docs it returns to a given number k (e.g. k=1). How can I do that? I have code like:

    self_query_retriever = SelfQueryRetriever.from_llm(llm,
                                                       vectordb,
                                                       document_content_description,
                                                       metadata_field_info,
                                                       verbose=True)
    qa_chain_with_filter = RetrievalQA.from_chain_type(llm,
                                                       retriever=self_query_retriever,
                                                       return_source_documents=True,
                                                       chain_type_kwargs={"prompt": qa_chain_prompt}
                                                       )

I have tried passing a parameter k=1 to from_llm() but it got quietly ignored, langsmith still shows me that the retriever is returning 4 documents, not 1.

Otherwise in code like

qa_chain_with_context = RetrievalQA.from_chain_type(llm, retriever=vectordb.as_retriever())

again I tried passing a k=1 to as_retriever(), which got quietly ignored too.

[dosubot[bot]]

Answer generated by a 🤖

Answer

I understand that you're trying to limit the number of documents returned by a retriever in the LangChain framework. You've tried passing a k parameter to the from_llm() and as_retriever() methods, but it seems to be ignored.

To address this, you can modify the _get_docs() method in the ConversationalRetrievalChain class. This method is where the retriever fetches the relevant documents. By adding a k parameter to this method and slicing the returned documents list, you can limit the number of documents returned.

Here's how you can modify the _get_docs() method:

def _get_docs(
    self,
    question: str,
    inputs: Dict[str, Any],
    *,
    run_manager: CallbackManagerForChainRun,
    k: int = 1  # Add this line
) -> List[Document]:
    """Get docs."""
    docs = self.retriever.get_relevant_documents(
        question, callbacks=run_manager.get_child()
    )
    return self._reduce_tokens_below_limit(docs[:k])  # Modify this line

In this version of _get_docs(), I've added a k parameter to control the number of documents returned. The docs[:k] operation is used to limit the number of documents to k.

Please note that this is one potential solution and it will affect all retrievals using this class. If you want to set k for a specific retriever, you might need to extend the retriever class and override its method to include the k parameter.

If you find this solution correct and it's not covered in the existing documentation, I encourage you to make a pull request to update the documentation. This will help other users who might encounter the same issue.

Sources

• libs/langchain/langchain/chains/conversational_retrieval/base.py

---

This response is meant to be useful, save you time, and share context. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.