Fix agents

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/creators/llm_agent_creator.py

@@ -24,7 +24,7 @@ def load_agent_file(cls, name: str) -> dict:
         return load(name.lower())
 
     @classmethod
-    def get_tool(cls, sql_helper, ai_search_helper, tool_name: str):
+    def get_tool(cls, sql_helper, tool_name: str):
         """Gets the tool based on the tool name.
         Args:
         ----
@@ -46,7 +46,7 @@ def get_tool(cls, sql_helper, ai_search_helper, tool_name: str):
             )
         elif tool_name == "sql_get_column_values_tool":
             return FunctionToolAlias(
-                ai_search_helper.get_column_values,
+                sql_helper.get_column_values,
                 description="Gets the values of a column in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned. Use this to get the correct value to apply against a filter for a user's question.",
             )
         else:
@@ -88,12 +88,11 @@ def create(cls, name: str, **kwargs) -> AssistantAgent:
         agent_file = cls.load_agent_file(name)
 
         sql_helper = ConnectorFactory.get_database_connector()
-        ai_search_helper = ConnectorFactory.get_ai_search_connector()
 
         tools = []
         if "tools" in agent_file and len(agent_file["tools"]) > 0:
             for tool in agent_file["tools"]:
-                tools.append(cls.get_tool(sql_helper, ai_search_helper, tool))
+                tools.append(cls.get_tool(sql_helper, tool))
 
         agent = AssistantAgent(
             name=name,

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

@@ -108,25 +108,28 @@ async def consume_inner_messages_from_agentic_flow(
 
                 try:
                     if isinstance(inner_message, ToolCallResultMessage):
-                        # Check for SQL query results
-                        parsed_message = self.parse_inner_message(inner_message.content)
-
-                        logging.info(f"Inner Loaded: {parsed_message}")
-
-                        if isinstance(parsed_message, dict):
-                            if (
-                                "type" in parsed_message
-                                and parsed_message["type"]
-                                == "query_execution_with_limit"
-                            ):
-                                database_results[identifier].append(
-                                    {
-                                        "sql_query": parsed_message[
-                                            "sql_query"
-                                        ].replace("\n", " "),
-                                        "sql_rows": parsed_message["sql_rows"],
-                                    }
-                                )
+                        for call_result in inner_message.content:
+                            # Check for SQL query results
+                            parsed_message = self.parse_inner_message(
+                                call_result.content
+                            )
+                            logging.info(f"Inner Loaded: {parsed_message}")
+
+                            if isinstance(parsed_message, dict):
+                                if (
+                                    "type" in parsed_message
+                                    and parsed_message["type"]
+                                    == "query_execution_with_limit"
+                                ):
+                                    logging.info("Contains query results")
+                                    database_results[identifier].append(
+                                        {
+                                            "sql_query": parsed_message[
+                                                "sql_query"
+                                            ].replace("\n", " "),
+                                            "sql_rows": parsed_message["sql_rows"],
+                                        }
+                                    )
 
                     elif isinstance(inner_message, TextMessage):
                         parsed_message = self.parse_inner_message(inner_message.content)
@@ -138,6 +141,7 @@ async def consume_inner_messages_from_agentic_flow(
                             if ("contains_pre_run_results" in parsed_message) and (
                                 parsed_message["contains_pre_run_results"] is True
                             ):
+                                logging.info("Contains pre-run results")
                                 for pre_run_sql_query, pre_run_result in parsed_message[
                                     "cached_questions_and_schemas"
                                 ].items():

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/ai_search.py

@@ -9,7 +9,6 @@
 import logging
 import base64
 from datetime import datetime, timezone
-import json
 from typing import Annotated
 from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
@@ -111,7 +110,6 @@ async def get_column_values(
             str,
             "The text to run a semantic search against. Relevant entities will be returned.",
         ],
-        as_json: bool = True,
     ):
         """Gets the values of a column in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned.
 
@@ -139,24 +137,7 @@ async def get_column_values(
             minimum_score=5,
         )
 
-        # build into a common format
-        column_values = {}
-
-        for value in values:
-            trimmed_fqn = ".".join(value["FQN"].split(".")[:-1])
-            if trimmed_fqn not in column_values:
-                column_values[trimmed_fqn] = []
-
-            column_values[trimmed_fqn].append(value["Value"])
-
-        logging.info("Column Values: %s", column_values)
-
-        filter_to_column = {text: column_values}
-
-        if as_json:
-            return json.dumps(filter_to_column, default=str)
-        else:
-            return filter_to_column
+        return values
 
     async def get_entity_schemas(
         self,
@@ -185,6 +166,8 @@ async def get_entity_schemas(
 
         logging.info("Search Text: %s", text)
 
+        stringified_engine_specific_fields = list(map(str, engine_specific_fields))
+
         retrieval_fields = [
             "FQN",
             "Entity",
@@ -194,7 +177,7 @@ async def get_entity_schemas(
             "Columns",
             "EntityRelationships",
             "CompleteEntityRelationshipsGraph",
-        ] + list(map(str, engine_specific_fields))
+        ] + stringified_engine_specific_fields
 
         schemas = await self.run_ai_search_query(
             text,
@@ -207,6 +190,8 @@ async def get_entity_schemas(
             top=3,
         )
 
+        fqn_to_trim = ".".join(stringified_engine_specific_fields)
+
         if len(excluded_entities) == 0:
             return schemas
 
@@ -220,12 +205,16 @@ async def get_entity_schemas(
                 and len(schema["CompleteEntityRelationshipsGraph"]) == 0
             ):
                 del schema["CompleteEntityRelationshipsGraph"]
+            else:
+                schema["CompleteEntityRelationshipsGraph"] = list(
+                    map(
+                        lambda x: x.replace(fqn_to_trim, ""),
+                        schema["CompleteEntityRelationshipsGraph"],
+                    )
+                )
 
-            if (
-                schema["SammpleValues"] is not None
-                and len(schema["SammpleValues"]) == 0
-            ):
-                del schema["SammpleValues"]
+            if schema["SampleValues"] is not None and len(schema["SampleValues"]) == 0:
+                del schema["SampleValues"]
 
             if (
                 schema["EntityRelationships"] is not None

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sql.py

@@ -74,6 +74,48 @@ async def query_execution(
             list[dict]: The results of the SQL query.
         """
 
+    async def get_column_values(
+        self,
+        text: Annotated[
+            str,
+            "The text to run a semantic search against. Relevant entities will be returned.",
+        ],
+        as_json: bool = True,
+    ):
+        """Gets the values of a column in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned.
+
+        Args:
+        ----
+            text (str): The text to run the search against.
+
+        Returns:
+        -------
+            str: The values of the column in JSON format.
+        """
+
+        values = await self.ai_search_connector.get_column_values(text)
+
+        # build into a common format
+        column_values = {}
+
+        starting = len(self.engine_specific_fields)
+
+        for value in values:
+            trimmed_fqn = ".".join(value["FQN"].split(".")[starting:-1])
+            if trimmed_fqn not in column_values:
+                column_values[trimmed_fqn] = []
+
+            column_values[trimmed_fqn].append(value["Value"])
+
+        logging.info("Column Values: %s", column_values)
+
+        filter_to_column = {text: column_values}
+
+        if as_json:
+            return json.dumps(filter_to_column, default=str)
+        else:
+            return filter_to_column
+
     @abstractmethod
     async def get_entity_schemas(
         self,

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

@@ -55,7 +55,7 @@ async def process_message(self, user_questions: list[str]) -> dict:
 
             for filter_condition in entity_result.filter_conditions:
                 column_search_tasks.append(
-                    self.ai_search_connector.get_column_values(
+                    self.sql_connector.get_column_values(
                         filter_condition, as_json=False
                     )
                 )