Update use of tools

Author: BenConstable9

text_2_sql/autogen/agents/llm_agents/sql_query_generation_agent.yaml

@@ -5,9 +5,16 @@ description:
 system_message:
   "You are a helpful AI Assistant that specialises in writing and executing SQL Queries to answer a given user's question.
 
-  If you need more information from the user to generate the SQL query, ask the user for the information you need with a question and end your answer with 'TERMINATE'.
+  You must:
+    1. Use the schema information provided to generate a SQL query that will answer the user's question.
+      If you need additional schema information, you can obtain it using the schema selection tool.
+    2. Validate the SQL query to ensure it is syntactically correct using the validation tool.
+    3. Run the SQL query to fetch the results.
 
-  If you are unsure how the question maps to the columns in the schema, return the possible columns that could be used in the SQL query and ask the user to provide more information to generate the SQL query. End your answer with 'TERMINATE'.
+  Important Info:
+    - If you need more information from the user to generate the SQL query, ask the user for the information you need with a question and end your answer with 'TERMINATE'.
+
+    - If you are unsure how the question maps to the columns in the schema, return the possible columns that could be used in the SQL query and ask the user to provide more information to generate the SQL query. End your answer with 'TERMINATE'.
 
   When generating the SQL query, you MUST follow these rules:
 
@@ -31,3 +38,4 @@ system_message:
 tools:
   - sql_query_execution_tool
   - sql_get_entity_schemas_tool
+  - sql_query_validation_tool

text_2_sql/autogen/utils/llm_agent_creator.py

@@ -3,10 +3,7 @@
 import yaml
 from autogen_core.components.tools import FunctionTool
 from autogen_agentchat.agents import AssistantAgent
-from utils.sql import (
-    query_execution,
-    get_entity_schemas,
-)
+from utils.sql import query_execution, get_entity_schemas, query_validation
 from utils.models import MINI_MODEL
 from jinja2 import Template
 
@@ -38,6 +35,11 @@ def get_tool(cls, tool_name):
                 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.",
             )
+        elif tool_name == "sql_query_validation_tool":
+            return FunctionTool(
+                query_validation,
+                description="Validates the SQL query to ensure that it is syntactically correct for the target database engine. Use this BEFORE executing any SQL statement.",
+            )
         else:
             raise ValueError(f"Tool {tool_name} not found")
 

text_2_sql/autogen/utils/models.py

@@ -14,6 +14,7 @@
 # )
 
 MINI_MODEL = AzureOpenAIChatCompletionClient(
+    azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
     model=os.environ["OpenAI__MiniCompletionDeployment"],
     api_version="2024-08-01-preview",
     azure_endpoint=os.environ["OpenAI__Endpoint"],

text_2_sql/autogen/utils/sql.py

@@ -97,7 +97,7 @@ async def query_execution(
     return results
 
 
-async def validate_sql_query(
+async def query_validation(
     sql_query: Annotated[
         str,
         "The SQL query to run against the database.",