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.

[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"]))

[dosubot[bot]]

The similarity_top_k parameter is intended to limit the number of rows returned, but there might be other factors affecting the result. Here are a few possibilities to consider:

1. Internal Logic: The retrieve method might have internal logic that affects the number of rows returned, such as additional filtering or aggregation steps that are not directly controlled by similarity_top_k.

2. SQL Query Execution: If the SQL query generated by the retriever is not directly limiting the number of results (e.g., using a LIMIT clause), the database might return more rows than expected.

3. Configuration or Defaults: There might be other configurations or default settings within the NLSQLRetriever or related components that override or interact with similarity_top_k.

Unfortunately, without specific details from the codebase, it's challenging to pinpoint the exact cause. You might want to check the SQL query being executed and any additional configurations or logic in the retrieve method that could influence the result set size.

To continue talking to Dosu, mention @dosu.

[nrohan09-cloud]

@sambitlnt why don't you pass the top_k value in your prompt and tell the llm to add a 'limit top_k' clause at the end of the query?