New agent

Author: BenConstable9

text_2_sql/autogen/agentic_text_2_sql.ipynb

@@ -72,7 +72,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.12.6"
+      "version": "3.12.7"
     }
   },
   "nbformat": 4,

text_2_sql/autogen/agentic_text_2_sql.py

@@ -5,18 +5,15 @@
 from utils.models import MINI_MODEL
 from utils.llm_agent_creator import LLMAgentCreator
 import logging
-from custom_agents.sql_query_cache_agent import SqlQueryCacheAgent
-from custom_agents.sql_schema_extraction_agent import SqlSchemaExtractionAgent
+from agents.custom_agents.sql_query_cache_agent import SqlQueryCacheAgent
 import json
 
 SQL_QUERY_GENERATION_AGENT = LLMAgentCreator.create(
     "sql_query_generation_agent",
     target_engine="Microsoft SQL Server",
     engine_specific_rules="Use TOP X to limit the number of rows returned instead of LIMIT X. NEVER USE LIMIT X as it produces a syntax error.",
 )
-# SQL_SCHEMA_SELECTION_AGENT = LLMAgentCreator.create(
-#     "sql_schema_selection_agent")
-SQL_SCHEMA_SELECTION_AGENT = SqlSchemaExtractionAgent()
+SQL_SCHEMA_SELECTION_AGENT = LLMAgentCreator.create("sql_schema_selection_agent")
 SQL_QUERY_CORRECTION_AGENT = LLMAgentCreator.create(
     "sql_query_correction_agent",
     target_engine="Microsoft SQL Server",

text_2_sql/autogen/agents/llm_agents/sql_schema_selection_agent.yaml

@@ -21,15 +21,8 @@ system_message:
     You would then use the 'sql_get_entity_schemas_tool' tool to search for the correct schema based on these key terms.
 
   Output Info:
-    If there are multiple possible schemas, return the information to the use in the following form for clarrification:
-
-    [
-      {
-        'question': 'Which of the following columns does the key term <key_term> belong to?',
-        'columns': ['<column_1>', '<column_2>', '<column_3>']
-      }
-    ]
-
-    End your answer with 'TERMINATE' if you are unable to find the correct schema or if the user's question is too ambiguous to select a schema."
+    - Just return the output of the 'sql_get_entity_schemas_tool' tool.
+    - Concate the results of the 'sql_get_entity_schemas_tool' tool if you are calling this agent in parallel.
+    - Do not provide any analysis or additional information to the user."
 tools:
   - sql_get_entity_schemas_tool

text_2_sql/autogen/utils/llm_agent_creator.py

@@ -14,7 +14,7 @@
 class LLMAgentCreator:
     @classmethod
     def load_agent_file(cls, name):
-        with open(f"/agents/llm_agents/{name.lower()}.yaml", "r") as file:
+        with open(f"./agents/llm_agents/{name.lower()}.yaml", "r") as file:
             file = yaml.safe_load(file)
 
         return file

text_2_sql/autogen/utils/sql.py

@@ -4,7 +4,7 @@
 import os
 import aioodbc
 from typing import Annotated
-from text_2_sql.autogen.utils.ai_search import run_ai_search_query
+from utils.ai_search import run_ai_search_query
 import json
 import asyncio