DH-4902/adding_top_k_as_env_vars

.env.example

@@ -3,6 +3,7 @@ OPENAI_API_KEY = #This field is required for the engine to work.
 ORG_ID =
 LLM_MODEL = 'gpt-4' #the openAI llm model that you want to use for evaluation and generating the nl answer. possible values: gpt-4, gpt-3.5-turbo
 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
+TOP_K = #top k results to be returned to the agent. This values is used to limit the number of rows returned to the LLMs
 
 DH_ENGINE_TIMEOUT = #timeout in seconds for the engine to return a response
 

dataherald/sql_generator/dataherald_sqlagent.py

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

dataherald/sql_generator/generates_nl_answer.py

@@ -46,7 +46,10 @@ 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=os.environ.get("TOP_K", 50),
         )
         system_message_prompt = SystemMessagePromptTemplate.from_template(
             SYSTEM_TEMPLATE