Adding the note to not found entity

dataherald/services/sql_generations.py

@@ -61,6 +61,7 @@ def update_the_initial_sql_generation(
         initial_sql_generation.status = sql_generation.status
         initial_sql_generation.error = sql_generation.error
         initial_sql_generation.intermediate_steps = sql_generation.intermediate_steps
+        initial_sql_generation.metadata.update(sql_generation.metadata)
         return self.sql_generation_repository.update(initial_sql_generation)
 
     def create(

dataherald/sql_generator/__init__.py

@@ -8,12 +8,12 @@
 from typing import Any, Dict, List, Tuple
 
 import sqlparse
-from sql_metadata import Parser
 from langchain.agents.agent import AgentExecutor
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentAction, LLMResult
 from langchain.schema.messages import BaseMessage
 from langchain_community.callbacks import get_openai_callback
+from sql_metadata import Parser
 
 from dataherald.config import Component, System
 from dataherald.db_scanner.models.types import TableDescription

dataherald/sql_generator/dataherald_finetuning_agent.py

@@ -253,7 +253,8 @@ class QuerySQLDataBaseTool(BaseSQLDatabaseTool, BaseTool):
     description = """
     Input: SQL query.
     Output: Result from the database or an error message if the query is incorrect.
-    Use this tool to execute the SQL query on the database, and return the results.
+    Use this tool to execute the SQL query on the database.
+    If the returned result is empty, mention as a comment in you final answer that the SQL query returned no results.
     """
     args_schema: Type[BaseModel] = SQLInput
 

dataherald/sql_generator/dataherald_sqlagent.py

@@ -153,6 +153,7 @@ class QuerySQLDataBaseTool(BaseSQLDatabaseTool, BaseTool):
     Output: Result from the database or an error message if the query is incorrect.
     If an error occurs, rewrite the query and retry.
     Use this tool to execute SQL queries.
+    If the returned result is empty, mention as a comment in you final answer that the SQL query returned no results.
     """
 
     @catch_exceptions()
@@ -315,10 +316,10 @@ class ColumnEntityChecker(BaseSQLDatabaseTool, BaseTool):
     Input: Column name and its corresponding table, and an entity.
     Output: cell-values found in the column similar to the given entity.
     Use this tool to get cell values similar to the given entity in the given column.
-    IF this tool returns no results, try to generate a SQL query regardless.
+    If this tool returns no results, try to generate a SQL query regardless. Also mention as a comment in the SQL query that the entity was not found in the column.
 
     Example Input: table1 -> column2, entity
-    """
+    """  # noqa: E501
 
     def find_similar_strings(
         self, input_list: List[tuple], target_string: str, threshold=0.4

dataherald/utils/agent_prompts.py

@@ -30,6 +30,7 @@
 tip5) If SQL results has None or NULL values, handle them by adding a WHERE clause to filter them out.
 tip6) The existance of the string values in the columns should always be checked using the DbColumnEntityChecker tool.
 tip7) You should always execute the SQL query by calling the SqlDbQuery tool to make sure the results are correct.
+tip8) If DBEntityChecker tool returns no entity values, you should still write a SQL query. But add a comment in the SQL query that the entity values were not found.
 """  # noqa: E501
 
 PLAN_WITH_INSTRUCTIONS = """1) Use the DbTablesWithRelevanceScores tool to find relevant tables.
@@ -46,6 +47,7 @@
 tip3) If SQL results has None or NULL values, handle them by adding a WHERE clause to filter them out.
 tip4) The existance of the string values in the columns should always be checked using the DbColumnEntityChecker tool.
 tip5) You should always execute the SQL query by calling the SqlDbQuery tool to make sure the results are correct.
+tip6) If DBEntityChecker tool returns no entity values, you should still write a SQL query. But add a comment in the SQL query that the entity values were not found.
 """  # noqa: E501
 
 PLAN_WITH_FEWSHOT_EXAMPLES = """1) Use the FewshotExamplesRetriever tool to retrieve samples of Question/SQL pairs that are similar to the given question, if there is a similar question among the examples, use the SQL query from the example and modify it to fit the given question.
@@ -63,6 +65,7 @@
 tip4) If SQL results has None or NULL values, handle them by adding a WHERE clause to filter them out.
 tip5) The existance of the string values in the columns should always be checked using the DbColumnEntityChecker tool.
 tip6) You should always execute the SQL query by calling the SqlDbQuery tool to make sure the results are correct.
+tip7) If DBEntityChecker tool returns no entity values, you should still write a SQL query. But add a comment in the SQL query that the entity values were not found.
 """  # noqa: E501
 
 PLAN_BASE = """1) Use the DbTablesWithRelevanceScores tool to find relevant tables.