diff --git a/.env.example b/.env.example index c2132ebd..2efd2ba8 100644 --- a/.env.example +++ b/.env.example @@ -3,6 +3,9 @@ OPENAI_API_KEY = #This field is required for the engine to work. ORG_ID = LLM_MODEL = 'gpt-4-1106-preview' #the openAI llm model that you want to use. possible values: gpt-4-1106-preview. +# All of our SQL generation agents are using different tools to generate SQL queries, in order to limit the number of times that agents can +# use different tools you can set the "AGENT_MAX_ITERATIONS" env variable. By default it is set to 20 iterations. + FINETUNED_MODEL = '' # The id of the LLM model you fine-tuned and you want to use inside DataheraldFinetuningAgent agent. DH_ENGINE_TIMEOUT = #timeout in seconds for the engine to return a response diff --git a/README.md b/README.md index 1b4fa954..46a54c1c 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,10 @@ ORG_ID = #Encryption key for storing DB connection data in Mongo ENCRYPT_KEY = + +# All of our SQL generation agents are using different tools to generate SQL queries, in order to limit the number of times that agents can +# use different tools you can set the "AGENT_MAX_ITERATIONS" env variable. By default it is set to 20 iterations. + ``` While not strictly required, we also strongly suggest you change the MONGO username and password fields as well. diff --git a/dataherald/sql_generator/dataherald_finetuning_agent.py b/dataherald/sql_generator/dataherald_finetuning_agent.py index 55090f6a..1b331f5d 100644 --- a/dataherald/sql_generator/dataherald_finetuning_agent.py +++ b/dataherald/sql_generator/dataherald_finetuning_agent.py @@ -365,7 +365,7 @@ def create_sql_agent( suffix: str = FINETUNING_AGENT_SUFFIX, format_instructions: str = FORMAT_INSTRUCTIONS, input_variables: List[str] | None = None, - max_iterations: int | None = 12, + max_iterations: int | None = int(os.getenv("AGENT_MAX_ITERATIONS", "20")), max_execution_time: float | None = None, early_stopping_method: str = "force", verbose: bool = False, diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py index 05d124b5..f8818d9a 100644 --- a/dataherald/sql_generator/dataherald_sqlagent.py +++ b/dataherald/sql_generator/dataherald_sqlagent.py @@ -545,7 +545,7 @@ def create_sql_agent( input_variables: List[str] | None = None, max_examples: int = 20, number_of_instructions: int = 1, - max_iterations: int | None = 15, + max_iterations: int | None = int(os.getenv("AGENT_MAX_ITERATIONS", "20")), max_execution_time: float | None = None, early_stopping_method: str = "force", verbose: bool = False,