Update agents

Author: BenConstable9

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

@@ -219,26 +219,26 @@ async def consume_inner_messages_from_agentic_flow(
                 # Create an instance of the InnerAutoGenText2Sql class
                 inner_autogen_text_2_sql = InnerAutoGenText2Sql(**self.kwargs)
 
-                identifier = ", ".join(parallel_message)
-
                 # Add database connection info to injected parameters
                 query_params = injected_parameters.copy() if injected_parameters else {}
                 if "Text2Sql__Tsql__ConnectionString" in os.environ:
                     query_params["database_connection_string"] = os.environ[
                         "Text2Sql__Tsql__ConnectionString"
                     ]
                 if "Text2Sql__Tsql__Database" in os.environ:
-                    query_params["database_name"] = os.environ["Text2Sql__Tsql__Database"]
+                    query_params["database_name"] = os.environ[
+                        "Text2Sql__Tsql__Database"
+                    ]
 
                 # Launch tasks for each sub-query
                 inner_solving_generators.append(
                     consume_inner_messages_from_agentic_flow(
                         inner_autogen_text_2_sql.process_user_message(
                             user_message=parallel_message,
                             injected_parameters=query_params,
-                            database_results=filtered_parallel_messages.database_results
+                            database_results=filtered_parallel_messages.database_results,
                         ),
-                        identifier,
+                        parallel_message,
                         filtered_parallel_messages,
                     )
                 )
@@ -267,7 +267,11 @@ async def consume_inner_messages_from_agentic_flow(
             # Check for disambiguation requests before processing the next round
 
             if (
-                max(map(len, filtered_parallel_messages.disambiguation_requests.values()))
+                max(
+                    map(
+                        len, filtered_parallel_messages.disambiguation_requests.values()
+                    )
+                )
                 > 0
             ):
                 # Final response

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

@@ -44,15 +44,15 @@ async def on_messages_stream(
         try:
             request_details = json.loads(messages[0].content)
             injected_parameters = request_details["injected_parameters"]
-            user_messages = request_details["user_message"]
-            logging.info(f"Processing messages: {user_messages}")
+            user_message = request_details["user_message"]
+            logging.info(f"Processing messages: {user_message}")
             logging.info(f"Input Parameters: {injected_parameters}")
         except json.JSONDecodeError:
             # If not JSON array, process as single message
             raise ValueError("Could not load message")
 
         cached_results = await self.agent.process_message(
-            user_messages, injected_parameters
+            user_message, injected_parameters
         )
         yield Response(
             chat_message=TextMessage(

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

@@ -43,19 +43,14 @@ async def on_messages_stream(
         # Try to parse as JSON first
         try:
             request_details = json.loads(messages[0].content)
-            messages = request_details["question"]
+            message = request_details["user_message"]
         except (json.JSONDecodeError, KeyError):
             # If not JSON or missing question key, use content directly
-            messages = messages[0].content
+            message = messages[0].content
 
-        if isinstance(messages, str):
-            messages = [messages]
-        elif not isinstance(messages, list):
-            messages = [str(messages)]
+        logging.info(f"Processing message: {message}")
 
-        logging.info(f"Processing questions: {messages}")
-
-        final_results = await self.agent.process_message(messages)
+        final_results = await self.agent.process_message(message)
 
         yield Response(
             chat_message=TextMessage(

text_2_sql/text_2_sql_core/src/text_2_sql_core/custom_agents/sql_query_cache_agent.py

@@ -8,39 +8,35 @@ class SqlQueryCacheAgentCustomAgent:
     def __init__(self):
         self.sql_connector = ConnectorFactory.get_database_connector()
 
-    async def process_message(
-        self, messages: list[str], injected_parameters: dict
-    ) -> dict:
+    async def process_message(self, message: str, injected_parameters: dict) -> dict:
         # Initialize results dictionary
         cached_results = {
             "cached_sql_queries_with_schemas_from_cache": [],
             "contains_cached_sql_queries_with_schemas_from_cache_database_results": False,
         }
 
-        # Process each question sequentially
-        for message in messages:
-            # Fetch the queries from the cache based on the question
-            logging.info(f"Fetching queries from cache for question: {message}")
-            cached_query = (
-                await self.sql_connector.fetch_sql_queries_with_schemas_from_cache(
-                    message, injected_parameters=injected_parameters
-                )
+        # Fetch the queries from the cache based on the question
+        logging.info(f"Fetching queries from cache for question: {message}")
+        cached_query = (
+            await self.sql_connector.fetch_sql_queries_with_schemas_from_cache(
+                message, injected_parameters=injected_parameters
             )
+        )
 
-            # If any question has pre-run results, set the flag
-            if cached_query.get(
-                "contains_cached_sql_queries_with_schemas_from_cache_database_results",
-                False,
-            ):
-                cached_results[
-                    "contains_cached_sql_queries_with_schemas_from_cache_database_results"
-                ] = True
+        # If any question has pre-run results, set the flag
+        if cached_query.get(
+            "contains_cached_sql_queries_with_schemas_from_cache_database_results",
+            False,
+        ):
+            cached_results[
+                "contains_cached_sql_queries_with_schemas_from_cache_database_results"
+            ] = True
 
-            # Add the cached results for this question
-            if cached_query.get("cached_sql_queries_with_schemas_from_cache"):
-                cached_results["cached_sql_queries_with_schemas_from_cache"].extend(
-                    cached_query["cached_sql_queries_with_schemas_from_cache"]
-                )
+        # Add the cached results for this question
+        if cached_query.get("cached_sql_queries_with_schemas_from_cache"):
+            cached_results["cached_sql_queries_with_schemas_from_cache"].extend(
+                cached_query["cached_sql_queries_with_schemas_from_cache"]
+            )
 
         logging.info(f"Final cached results: {cached_results}")
         return cached_results

text_2_sql/text_2_sql_core/src/text_2_sql_core/custom_agents/sql_schema_selection_agent.py

@@ -22,47 +22,35 @@ def __init__(self, **kwargs):
 
         self.system_prompt = Template(system_prompt).render(kwargs)
 
-    async def process_message(self, messages: list[str]) -> dict:
-        logging.info(f"user inputs: {messages}")
-
-        entity_tasks = []
-
-        for message in messages:
-            messages = [
-                {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": message},
-            ]
-            entity_tasks.append(
-                self.open_ai_connector.run_completion_request(
-                    messages, response_format=SQLSchemaSelectionAgentOutput
-                )
-            )
+    async def process_message(self, message: str) -> dict:
+        logging.info(f"Processing message: {message}")
 
-        entity_results = await asyncio.gather(*entity_tasks)
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": message},
+        ]
+        entity_result = await self.open_ai_connector.run_completion_request(
+            messages, response_format=SQLSchemaSelectionAgentOutput
+        )
 
         entity_search_tasks = []
         column_search_tasks = []
 
-        for entity_result in entity_results:
-            logging.info(f"Entity result: {entity_result}")
+        logging.info(f"Entity result: {entity_result}")
 
-            for entity_group in entity_result.entities:
-                logging.info("Searching for schemas for entity group: %s", entity_group)
-                entity_search_tasks.append(
-                    self.sql_connector.get_entity_schemas(
-                        " ".join(entity_group), as_json=False
-                    )
+        for entity_group in entity_result.entities:
+            logging.info("Searching for schemas for entity group: %s", entity_group)
+            entity_search_tasks.append(
+                self.sql_connector.get_entity_schemas(
+                    " ".join(entity_group), as_json=False
                 )
+            )
 
-            for filter_condition in entity_result.filter_conditions:
-                logging.info(
-                    "Searching for column values for filter: %s", filter_condition
-                )
-                column_search_tasks.append(
-                    self.sql_connector.get_column_values(
-                        filter_condition, as_json=False
-                    )
-                )
+        for filter_condition in entity_result.filter_conditions:
+            logging.info("Searching for column values for filter: %s", filter_condition)
+            column_search_tasks.append(
+                self.sql_connector.get_column_values(filter_condition, as_json=False)
+            )
 
         schemas_results = await asyncio.gather(*entity_search_tasks)
         column_value_results = await asyncio.gather(*column_search_tasks)

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/disambiguation_and_sql_query_generation_agent.yaml

@@ -274,3 +274,5 @@ system_message:
       TERMINATE
     </output_format>
   "
+tools:
+  - sql_get_entity_schemas_tool