How to limit the number of rows while using NLSQLRetriever?

[sambitlnt]

Is there a way to put a limit on the number of rows returned by NLSQLRetriever.retrieve? Say I always want only top 5 rows of the result set to be returned.

[dosubot[bot]]

To limit the number of rows returned by NLSQLRetriever.retrieve to always get only the top 5 rows, you can set the SIMILARITY_TOP_K parameter in the configuration:

SIMILARITY_TOP_K=5

This configuration ensures that only the top 5 rows are returned by the retriever ^[1][2].

To continue talking to Dosu, mention @dosu.

[sambitlnt]

@dosu NLSQLRetriever(sql_database, tables, llm, embed_model, return_raw) - This doesn't take SIMILARITY_TOP_K as an argument. How do I pass it in the parameter

[dosubot[bot]]

To pass the SIMILARITY_TOP_K parameter to the NLSQLRetriever constructor, you can set it using the similarity_top_k property after the object has been instantiated. Here is an example:

# Assuming you have an instance of NLSQLRetriever
nlsql_retriever = NLSQLRetriever(sql_database, tables, llm, embed_model, return_raw)

# Set the similarity_top_k parameter
nlsql_retriever.similarity_top_k = 5  # or any desired value

The similarity_top_k property has a setter method that allows you to set its value after the object has been created ^[1].

To continue talking to Dosu, mention @dosu.

[sambitlnt]

@dosu Setting nl_sql_retriever.similarity_top_k=3 still returns me more than 3 rows

[dosubot[bot]]

Setting nl_sql_retriever.similarity_top_k = 3 should ideally limit the number of results to 3. However, the _retrieve method in the MultiModalVectorIndexRetriever class combines results from both text and image vector stores. This means that if both stores are queried, the combined results can exceed the specified similarity_top_k.

Here is the relevant code snippet:

def _retrieve(
        self,
        query_bundle: QueryBundle,
    ) -> List[NodeWithScore]:
    res = []
    # If text vector store is not empty, retrieve text nodes
    if self._vector_store is not None:
        res.extend(self._text_retrieve(query_bundle))

    # If image vector store is not empty, retrieve text nodes
    if self._image_vector_store is not None:
        res.extend(self._text_to_image_retrieve(query_bundle))
    return res

In this method, results from _text_retrieve and _text_to_image_retrieve are combined, which can lead to more than 3 results being returned ^[1].

To continue talking to Dosu, mention @dosu.

[sambitlnt]

@dosu The following gives me 7 rows. I only have text data

nl_sql_retriever.similarity_top_k=3

query = "city with population higher than 100000"
answer = nl_sql_retriever.retrieve(query)

metadata = answer[0].metadata
print(metadata)
print(len(metadata["result"]))