diff --git a/.env.example b/.env.example index aa3f8ba7..0f83b133 100644 --- 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 index 7d6a6cac..ccac6be0 100644 --- 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 index 8fd503a8..54720cb0 100644 --- 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 index f15c519f..80f8b2d5 100644 --- 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"