minor modification to move the top_k param to the environment variables for convenience.

.env.example

@@ -5,6 +5,7 @@ LLM_MODEL = 'gpt-4' #the openAI llm model that you want to use for evaluation an
 AGENT_LLM_MODEL = 'gpt-4-32k' # the llm model that you want to use for the agent, it should have a lrage context window. possible values: gpt-4-32k, gpt-3.5-turbo-16k
 
 DH_ENGINE_TIMEOUT = #timeout in seconds for the engine to return a response
+UPPER_LIMIT_QUERY_RETURN_ROWS = #The upper limit on number of rows returned from the query engine (equivalent to using LIMIT N in PostgreSQL/MySQL/SQlite). Defauls to 50
 
 #Encryption key for storing DB connection data in Mongo
 ENCRYPT_KEY = 

dataherald/sql_generator/dataherald_sqlagent.py

@@ -54,7 +54,7 @@
 logger = logging.getLogger(__name__)
 
 
-TOP_K = 50
+TOP_K = int(os.getenv("UPPER_LIMIT_QUERY_RETURN_ROWS", "50"))
 
 
 def catch_exceptions():  # noqa: C901

dataherald/sql_generator/generates_nl_answer.py

@@ -46,7 +46,7 @@ def execute(self, query_response: Response) -> Response:
         )
         database = SQLDatabase.get_sql_engine(database_connection)
         query_response = create_sql_query_status(
-            database, query_response.sql_query, query_response, top_k=50
+            database, query_response.sql_query, query_response, top_k=int(os.getenv("UPPER_LIMIT_QUERY_RETURN_ROWS", "50"))
         )
         system_message_prompt = SystemMessagePromptTemplate.from_template(
             SYSTEM_TEMPLATE

docs/envars.rst

@@ -36,6 +36,7 @@ provided in the .env.example file with the default values.
     S3_AWS_SECRET_ACCESS_KEY =
 
     DH_ENGINE_TIMEOUT =
+    UPPER_LIMIT_QUERY_RETURN_ROWS =
 
 
 .. csv-table::
@@ -63,3 +64,4 @@ provided in the .env.example file with the default values.
    "S3_AWS_ACCESS_KEY_ID", "The key used to access credential files if saved to S3", "None", "No"
    "S3_AWS_SECRET_ACCESS_KEY", "The key used to access credential files if saved to S3", "None", "No"
    "DH_ENGINE_TIMEOUT", "This is used to set the max seconds the process will wait for the response to be generate. If the specified time limit is exceeded, it will trigger an exception", "None", "No"
+   "UPPER_LIMIT_QUERY_RETURN_ROWS", "The upper limit on number of rows returned from the query engine (equivalent to using LIMIT N in PostgreSQL/MySQL/SQlite).", "None", "No"