Fix

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -18,11 +18,12 @@
 from autogen_agentchat.base import Response
 import json
 import os
-import asyncio
 from datetime import datetime
 
+
 class EmptyResponseUserProxyAgent(UserProxyAgent):
     """UserProxyAgent that automatically responds with empty messages."""
+
     def __init__(self, name):
         super().__init__(name=name)
         self._has_responded = False
@@ -35,6 +36,7 @@ async def on_messages_stream(self, messages, sender=None, config=None):
             yield message
         yield Response(chat_message=message)
 
+
 class AutoGenText2Sql:
     def __init__(self, engine_specific_rules: str, **kwargs: dict):
         self.use_query_cache = False
@@ -65,32 +67,31 @@ def get_all_agents(self):
         """Get all agents for the complete flow."""
         # Get current datetime for the Query Rewrite Agent
         current_datetime = datetime.now()
-        
+
         QUERY_REWRITE_AGENT = LLMAgentCreator.create(
-            "query_rewrite_agent",
-            current_datetime=current_datetime
+            "query_rewrite_agent", current_datetime=current_datetime
         )
-        
+
         SQL_QUERY_GENERATION_AGENT = LLMAgentCreator.create(
             "sql_query_generation_agent",
             target_engine=self.target_engine,
             engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
-        
+
         SQL_SCHEMA_SELECTION_AGENT = SqlSchemaSelectionAgent(
             target_engine=self.target_engine,
             engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
-        
+
         SQL_QUERY_CORRECTION_AGENT = LLMAgentCreator.create(
             "sql_query_correction_agent",
             target_engine=self.target_engine,
             engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
-        
+
         SQL_DISAMBIGUATION_AGENT = LLMAgentCreator.create(
             "sql_disambiguation_agent",
             target_engine=self.target_engine,
@@ -101,11 +102,9 @@ def get_all_agents(self):
         QUESTION_DECOMPOSITION_AGENT = LLMAgentCreator.create(
             "question_decomposition_agent"
         )
-        
+
         # Auto-responding UserProxyAgent
-        USER_PROXY = EmptyResponseUserProxyAgent(
-            name="user_proxy"
-        )
+        USER_PROXY = EmptyResponseUserProxyAgent(name="user_proxy")
 
         agents = [
             USER_PROXY,
@@ -114,7 +113,7 @@ def get_all_agents(self):
             SQL_SCHEMA_SELECTION_AGENT,
             SQL_QUERY_CORRECTION_AGENT,
             QUESTION_DECOMPOSITION_AGENT,
-            SQL_DISAMBIGUATION_AGENT
+            SQL_DISAMBIGUATION_AGENT,
         ]
 
         if self.use_query_cache:
@@ -192,7 +191,6 @@ def agentic_flow(self):
             allow_repeated_speaker=False,
             model_client=LLMModelCreator.get_model("4o-mini"),
             termination_condition=self.termination_condition,
-            selector_func=self.selector,
             selector_func=self.unified_selector,
         )
         return flow

text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/sql_query_cache_agent.py

@@ -47,19 +47,21 @@ async def on_messages_stream(
             # Initialize results dictionary
             cached_results = {
                 "cached_questions_and_schemas": [],
-                "contains_pre_run_results": False
+                "contains_pre_run_results": False,
             }
 
             # Process each question sequentially
             for question in user_questions:
                 # Fetch the queries from the cache based on the question
                 logging.info(f"Fetching queries from cache for question: {question}")
-                cached_query = await self.sql_connector.fetch_queries_from_cache(question)
-                
+                cached_query = await self.sql_connector.fetch_queries_from_cache(
+                    question
+                )
+
                 # If any question has pre-run results, set the flag
                 if cached_query.get("contains_pre_run_results", False):
                     cached_results["contains_pre_run_results"] = True
-                
+
                 # Add the cached results for this question
                 if cached_query.get("cached_questions_and_schemas"):
                     cached_results["cached_questions_and_schemas"].extend(
@@ -75,7 +77,9 @@ async def on_messages_stream(
         except json.JSONDecodeError:
             # If not JSON array, process as single question
             logging.info(f"Processing single question: {last_response}")
-            cached_queries = await self.sql_connector.fetch_queries_from_cache(last_response)
+            cached_queries = await self.sql_connector.fetch_queries_from_cache(
+                last_response
+            )
             yield Response(
                 chat_message=TextMessage(
                     content=json.dumps(cached_queries), source=self.name

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/query_rewrite_agent.yaml

@@ -14,7 +14,7 @@ system_message:
   - Use the current date/time above as reference point
   - Replace relative dates like 'last month', 'this year', 'previous quarter' with absolute dates
   - Maintain consistency in date formats (YYYY-MM-DD)
-  
+
   Examples of date resolution (assuming current date is {{ current_datetime }}):
   - 'last month' -> specific month name and year
   - 'this year' -> {{ current_datetime.year }}