Merge agents into 1

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -59,14 +59,6 @@ def set_mode(self):
 
     def get_all_agents(self):
         """Get all agents for the complete flow."""
-        # Get current datetime for the Query Rewrite Agent
-        self.sql_query_generation_agent = LLMAgentCreator.create(
-            "sql_query_generation_agent",
-            target_engine=self.target_engine,
-            engine_specific_rules=self.engine_specific_rules,
-            **self.kwargs,
-        )
-
         # If relationship_paths not provided, use a generic template
         if "relationship_paths" not in self.kwargs:
             self.kwargs[
@@ -92,22 +84,16 @@ def get_all_agents(self):
             **self.kwargs,
         )
 
-        self.sql_disambiguation_agent = LLMAgentCreator.create(
-            "sql_disambiguation_agent",
+        self.disambiguation_and_sql_query_generation_agent = LLMAgentCreator.create(
+            "disambiguation_and_sql_query_generation_agent",
             target_engine=self.target_engine,
             engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
-
-        # Auto-responding UserProxyAgent
-        self.user_proxy = EmptyResponseUserProxyAgent(name="user_proxy")
-
         agents = [
-            self.user_proxy,
-            self.sql_query_generation_agent,
             self.sql_schema_selection_agent,
             self.sql_query_correction_agent,
-            self.sql_disambiguation_agent,
+            self.disambiguation_and_sql_query_generation_agent,
         ]
 
         if self.use_query_cache:
@@ -140,10 +126,8 @@ def unified_selector(self, messages):
             # Always go through schema selection after cache check
             decision = "sql_schema_selection_agent"
         elif current_agent == "sql_schema_selection_agent":
-            decision = "sql_disambiguation_agent"
-        elif current_agent == "sql_disambiguation_agent":
-            decision = "sql_query_generation_agent"
-        elif current_agent == "sql_query_generation_agent":
+            decision = "disambiguation_and_sql_query_generation_agent"
+        elif current_agent == "disambiguation_and_sql_query_generation_agent":
             decision = "sql_query_correction_agent"
         elif current_agent == "sql_query_correction_agent":
             decision = "sql_query_correction_agent"

…re/prompts/sql_disambiguation_agent.yaml …tion_and_sql_query_generation_agent.yamltext_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_disambiguation_agent.yaml renamed to text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/disambiguation_and_sql_query_generation_agent.yaml text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_disambiguation_agent.yaml renamed to text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/disambiguation_and_sql_query_generation_agent.yaml

@@ -6,6 +6,7 @@ system_message:
   "<role_and_objective>
     You are a helpful AI Assistant specializing in disambiguating questions about {{ use_case }} and mapping them to the relevant columns and schemas in the database.
     Your job is to create clear mappings between the user's intent and the available database schema.
+    If all mappings are clear, generate {{ target_engine }} compliant SQL query based on the mappings.
   </role_and_objective>
 
   <key_concepts>
@@ -106,6 +107,57 @@ system_message:
       }
     </examples>
 
+    <sql_query_generation_rules>
+
+      <engine_specific_rules>
+          {{ engine_specific_rules }}
+      </engine_specific_rules>
+
+      Your primary focus is on:
+        1. Understanding what data the user wants to retrieve
+        2. Identifying the necessary tables and their relationships
+        3. Determining any required calculations or aggregations
+        4. Specifying any filtering conditions based on the user's criteria
+
+      When generating SQL queries, focus on these key aspects:
+
+        - Data Selection:
+          * Identify the main pieces of information the user wants to see
+          * Include any calculated fields or aggregations needed
+          * Consider what grouping might be required
+          * Follow basic {{ target_engine }} syntax patterns
+
+        - Table Relationships:
+          * Use the schema information to identify required tables
+          * Join tables as needed to connect related information
+          * Request additional schema information if needed using the schema selection tool
+          * Use {{ target_engine }}-compatible join syntax
+
+        - Filtering Conditions:
+          * Translate user criteria into WHERE conditions
+          * Handle date ranges, categories, or numeric thresholds
+          * Consider both explicit and implicit filters in the user's question
+          * Use {{ target_engine }}-compatible date and string functions
+
+        - Result Organization:
+          * Determine if specific sorting is needed
+          * Consider if grouping is required
+          * Include any having conditions for filtered aggregates
+          * Follow {{ target_engine }} ordering syntax
+
+      Guidelines:
+
+        - Focus on getting the right tables and relationships
+        - Ensure all necessary data is included
+        - Follow basic {{ target_engine }} syntax patterns
+        - The correction agent will handle:
+          * Detailed syntax corrections
+          * Query execution
+          * Result formatting
+
+      Remember: Your job is to focus on the data relationships and logic while following basic {{ target_engine }} patterns.
+  </sql_query_generation_rules>
+
   <output_format>
     If all mappings are clear:
     {
@@ -128,6 +180,8 @@ system_message:
       }
     }
 
+    Then use the mapping to generate the SQL query following the engine-specific rules. Run this query to retrieve the data afterwards.
+
     If disambiguation needed:
     {
       \"disambiguation\": [{
@@ -140,3 +194,5 @@ system_message:
     TERMINATE
   </output_format>
   "
+tools:
+  - sql_query_execution_tool