Reduce some complexity in sql query generation

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/creators/llm_agent_creator.py

@@ -42,18 +42,13 @@ def get_tool(cls, sql_helper, ai_search_helper, tool_name: str):
         elif tool_name == "sql_get_entity_schemas_tool":
             return FunctionToolAlias(
                 sql_helper.get_entity_schemas,
-                description="Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Extract key terms from the user question and use these as the search term. Several entities may be returned. Only use when the provided schemas in the system prompt are not sufficient to answer the question.",
+                description="Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Extract key terms from the user question and use these as the search term. Several entities may be returned. Only use when the provided schemas in the message history are not sufficient to answer the question.",
             )
         elif tool_name == "sql_get_column_values_tool":
             return FunctionToolAlias(
                 ai_search_helper.get_column_values,
                 description="Gets the values of a column in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned. Use this to get the correct value to apply against a filter for a user's question.",
             )
-        elif tool_name == "current_datetime_tool":
-            return FunctionToolAlias(
-                sql_helper.get_current_datetime,
-                description="Gets the current date and time.",
-            )
         else:
             raise ValueError(f"Tool {tool_name} not found")
 

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -146,26 +146,7 @@ def unified_selector(self, messages):
         elif current_agent == "sql_query_generation_agent":
             decision = "sql_query_correction_agent"
         elif current_agent == "sql_query_correction_agent":
-            try:
-                correction_result = json.loads(messages[-1].content)
-                if isinstance(correction_result, dict):
-                    if "answer" in correction_result and "sources" in correction_result:
-                        decision = "user_proxy"
-                    elif "corrected_query" in correction_result:
-                        if correction_result.get("executing", False):
-                            decision = "sql_query_correction_agent"
-                        else:
-                            decision = "sql_query_generation_agent"
-                    elif "error" in correction_result:
-                        decision = "sql_query_generation_agent"
-                elif isinstance(correction_result, list) and len(correction_result) > 0:
-                    if "requested_fix" in correction_result[0]:
-                        decision = "sql_query_generation_agent"
-
-                if decision is None:
-                    decision = "sql_query_generation_agent"
-            except json.JSONDecodeError:
-                decision = "sql_query_generation_agent"
+            decision = "sql_query_correction_agent"
 
         if decision:
             logging.info(f"Agent transition: {current_agent} -> {decision}")

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sql.py

@@ -7,7 +7,6 @@
 import asyncio
 import sqlglot
 from abc import ABC, abstractmethod
-from datetime import datetime
 from jinja2 import Template
 import json
 
@@ -30,22 +29,6 @@ def __init__(self):
 
         self.database_engine = None
 
-    def get_current_datetime(self) -> str:
-        """Get the current datetime."""
-        return datetime.now().strftime("%d/%m/%Y, %H:%M:%S")
-
-    def get_current_date(self) -> str:
-        """Get the current date."""
-        return datetime.now().strftime("%d/%m/%Y")
-
-    def get_current_time(self) -> str:
-        """Get the current time."""
-        return datetime.now().strftime("%H:%M:%S")
-
-    def get_current_unix_timestamp(self) -> int:
-        """Get the current unix timestamp."""
-        return int(datetime.now().timestamp())
-
     @abstractmethod
     async def query_execution(
         self,
@@ -169,19 +152,6 @@ async def fetch_queries_from_cache(
         if injected_parameters is None:
             injected_parameters = {}
 
-        # Populate the injected_parameters
-        if "date" not in injected_parameters:
-            injected_parameters["date"] = self.get_current_date()
-
-        if "time" not in injected_parameters:
-            injected_parameters["time"] = self.get_current_time()
-
-        if "datetime" not in injected_parameters:
-            injected_parameters["datetime"] = self.get_current_datetime()
-
-        if "unix_timestamp" not in injected_parameters:
-            injected_parameters["unix_timestamp"] = self.get_current_unix_timestamp()
-
         cached_schemas = await self.ai_search_connector.run_ai_search_query(
             question,
             ["QuestionEmbedding"],

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/agent_response.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from pydantic import BaseModel, RootModel, Field
+from pydantic import BaseModel, RootModel, Field, model_validator
 from enum import StrEnum
 
 from typing import Literal
@@ -59,6 +59,29 @@ class AgentRequestBody(BaseModel):
     question: str
     injected_parameters: dict = Field(default_factory=dict)
 
+    @model_validator(mode="before")
+    def add_defaults_to_injected_parameters(cls, values):
+        if "injected_parameters" not in values:
+            values["injected_parameters"] = {}
+
+        if "date" not in values["injected_parameters"]:
+            values["injected_parameters"]["date"] = datetime.now().strftime("%d/%m/%Y")
+
+        if "time" not in values["injected_parameters"]:
+            values["injected_parameters"]["time"] = datetime.now().strftime("%H:%M:%S")
+
+        if "datetime" not in values["injected_parameters"]:
+            values["injected_parameters"]["datetime"] = datetime.now().strftime(
+                "%d/%m/%Y, %H:%M:%S"
+            )
+
+        if "unix_timestamp" not in values["injected_parameters"]:
+            values["injected_parameters"]["unix_timestamp"] = int(
+                datetime.now().timestamp()
+            )
+
+        return values
+
 
 class AgentResponse(BaseModel):
     header: AgentResponseHeader | None = Field(default=None)

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_query_correction_agent.yaml

@@ -5,10 +5,11 @@ description:
 system_message:
   "<role_and_objective>
       You are a SQL syntax expert specializing in converting standard SQL to {{ target_engine }}-compliant SQL. Your job is to:
-      1. Take SQL queries with correct logic but potential syntax issues
-      2. Fix them according to {{ target_engine }} syntax rules
-      3. Execute the corrected queries
-      4. Return the results
+      1. Take SQL queries with correct logic but potential syntax issues.
+      2. Review the output from the SQL query being run and fix them according to {{ target_engine }} syntax rules if needed.
+      3. Execute the corrected queries if needed.
+      4. Verify that the results will answer all of the user's questions. If not, create additional queries and run them.
+      5. Return the results
   </role_and_objective>
 
   <engine_specific_rules>
@@ -85,18 +86,10 @@ system_message:
   </error_handling>
 
   <output_format>
-      - **When query executes successfully**:
+      - **When query executes successfully and answers all questions**:
         ```json
         {
-          \"answer\": \"<ANSWER BASED ON QUERY RESULTS>\",
-          \"sources\": [
-            {
-              \"sql_result_snippet\": \"<FORMATTED QUERY RESULTS>\",
-              \"sql_query_used\": \"<EXECUTED SQL QUERY>\",
-              \"original_query\": \"<QUERY BEFORE CONVERSION>\",
-              \"explanation\": \"<EXPLANATION OF CONVERSIONS AND RESULTS>\"
-            }
-          ]
+          \"validated\": \"<TRUE>\",
         }
         ```
         Followed by **TERMINATE**.
@@ -138,3 +131,5 @@ system_message:
   "
 tools:
   - sql_query_execution_tool
+  - sql_get_entity_schemas_tool
+  - sql_get_column_values_tool

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_query_generation_agent.yaml

@@ -54,5 +54,4 @@ system_message:
   Remember: Your job is to focus on the data relationships and logic while following basic {{ target_engine }} patterns. The correction agent will handle detailed syntax fixes and execution.
   "
 tools:
-  - sql_get_entity_schemas_tool
-  - current_datetime_tool
+  - sql_query_execution_tool