minor modification to move the top_k param to the environment variabl…

…es for convenience. (#226) * minor modification to move the top_k param to the environment variables for convenience. * renamed MAX_QUERY_RETURN_ROWS -> UPPER_LIMIT_QUERY_RETURN_ROWS to make it easier to understand * updated docs to match env var change * os.getenv returns a 'str' type so it's necessary to cast to int --------- Co-authored-by: Paul Markus <[email protected]>
Dataherald · Oct 30, 2023 · efdd3b0 · efdd3b0
1 parent d8439f3
commit efdd3b0
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 2 deletions.
diff --git a/.env.example b/.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 = 

diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/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

diff --git a/dataherald/sql_generator/generates_nl_answer.py b/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

diff --git a/docs/envars.rst b/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"