diff --git a/dataherald/services/sql_generations.py b/dataherald/services/sql_generations.py index c746802a..75e97f7d 100644 --- a/dataherald/services/sql_generations.py +++ b/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( diff --git a/dataherald/sql_generator/__init__.py b/dataherald/sql_generator/__init__.py index 87c66f3d..ba130da3 100644 --- a/dataherald/sql_generator/__init__.py +++ b/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 diff --git a/dataherald/sql_generator/dataherald_finetuning_agent.py b/dataherald/sql_generator/dataherald_finetuning_agent.py index 3cbdd04e..da936a2c 100644 --- a/dataherald/sql_generator/dataherald_finetuning_agent.py +++ b/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 diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py index d18273ed..a8ae5695 100644 --- a/dataherald/sql_generator/dataherald_sqlagent.py +++ b/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 diff --git a/dataherald/utils/agent_prompts.py b/dataherald/utils/agent_prompts.py index 23779def..0c486616 100644 --- a/dataherald/utils/agent_prompts.py +++ b/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.