DH-4879/decreasing the few-shot samples and returned results (#218)

Dataherald · Oct 18, 2023 · a40c248 · a40c248
1 parent e1d9d40
commit a40c248
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py
@@ -44,7 +44,7 @@
 logger = logging.getLogger(__name__)
 
 
-TOP_K = 100
+TOP_K = 50
 AGENT_PREFIX = """You are an agent designed to interact with a SQL database.
 Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
 You have access to tools for interacting with the database.
@@ -548,7 +548,7 @@ def get_tools(self) -> List[BaseTool]:
 class DataheraldSQLAgent(SQLGenerator):
     """Dataherald SQL agent"""
 
-    max_number_of_examples: int = 20  # maximum number of question/SQL pairs
+    max_number_of_examples: int = 5  # maximum number of question/SQL pairs
     llm: Any = None
 
     def remove_duplicate_examples(self, fewshot_exmaples: List[dict]) -> List[dict]:

diff --git a/dataherald/sql_generator/generates_nl_answer.py b/dataherald/sql_generator/generates_nl_answer.py
@@ -43,7 +43,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
+            database, query_response.sql_query, query_response, top_k=50
         )
         system_message_prompt = SystemMessagePromptTemplate.from_template(
             SYSTEM_TEMPLATE