Dump of work

Author: BenConstable9

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

@@ -41,7 +41,7 @@ def get_tool(cls, sql_helper, ai_search_helper, tool_name: str):
             )
         elif tool_name == "sql_get_entity_schemas_tool":
             return FunctionTool(
-                ai_search_helper.get_entity_schemas,
+                sql_helper.get_entity_schemas,
                 description="Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Extract key terms from the user question and use these as the search term. Several entities may be returned. Only use when the provided schemas in the system prompt are not sufficient to answer the question.",
             )
         elif tool_name == "sql_get_column_values_tool":

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

@@ -25,6 +25,8 @@ def __init__(self, **kwargs):
 
         self.open_ai_connector = ConnectorFactory.get_open_ai_connector()
 
+        self.sql_connector = ConnectorFactory.get_database_connector()
+
         system_prompt = load("sql_schema_selection_agent")["system_message"]
 
         self.system_prompt = Template(system_prompt).render(kwargs)
@@ -75,7 +77,7 @@ async def on_messages_stream(
             logging.info(f"Loaded entity result: {loaded_entity_result}")
 
             entity_search_tasks.append(
-                self.ai_search_connector.get_entity_schemas(
+                self.sql_connector.get_entity_schemas(
                     " ".join(loaded_entity_result["entities"]), as_json=False
                 )
             )

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

@@ -90,16 +90,27 @@ async def run_ai_search_query(
             async for result in results.by_page():
                 async for item in result:
                     if (
-                        minimum_score is not None
-                        and item["@search.reranker_score"] < minimum_score
+                        "@search.reranker_score" in item
+                        and item["@search.reranker_score"] is not None
                     ):
+                        score = item["@search.reranker_score"]
+                    elif "@search.score" in item and item["@search.score"] is not None:
+                        score = item["@search.score"]
+                    else:
+                        raise Exception("No score found in the search results.")
+
+                    if minimum_score is not None and score < minimum_score:
                         continue
 
                     if include_scores is False:
-                        del item["@search.reranker_score"]
-                        del item["@search.score"]
-                        del item["@search.highlights"]
-                        del item["@search.captions"]
+                        if "@search.reranker_score" in item:
+                            del item["@search.reranker_score"]
+                        if "@search.score" in item:
+                            del item["@search.score"]
+                        if "@search.highlights" in item:
+                            del item["@search.highlights"]
+                        if "@search.captions" in item:
+                            del item["@search.captions"]
 
                     logging.info("Item: %s", item)
                     combined_results.append(item)
@@ -131,19 +142,31 @@ async def get_column_values(
         text = " ".join([f"{word}~" for word in text.split()])
         values = await self.run_ai_search_query(
             text,
-            [],
-            ["FQN", "Column", "Value"],
-            os.environ[
+            vector_fields=[],
+            retrieval_fields=["FQN", "Column", "Value"],
+            index_name=os.environ[
                 "AIService__AzureSearchOptions__Text2SqlColumnValueStore__Index"
             ],
-            None,
+            semantic_config=None,
             top=10,
+            include_scores=True,
+            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"])
+
         if as_json:
-            return json.dumps(values, default=str)
+            return json.dumps(column_values, default=str)
         else:
-            return values
+            return column_values
 
     async def get_entity_schemas(
         self,
@@ -155,7 +178,6 @@ async def get_entity_schemas(
             list[str],
             "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
         ] = [],
-        as_json: bool = True,
     ) -> str:
         """Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned.
 
@@ -187,19 +209,14 @@ async def get_entity_schemas(
         )
 
         for schema in schemas:
-            entity = schema["Entity"]
-
             filtered_schemas = []
             for excluded_entity in excluded_entities:
-                if excluded_entity.lower() == entity.lower():
+                if excluded_entity.lower() == schema["Entity"].lower():
                     logging.info("Excluded entity: %s", excluded_entity)
                 else:
                     filtered_schemas.append(schema)
 
-        if as_json:
-            return json.dumps(schemas, default=str)
-        else:
-            return schemas
+        return filtered_schemas
 
     async def add_entry_to_index(document: dict, vector_fields: dict, index_name: str):
         """Add an entry to the search index."""

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/databricks_sql.py

@@ -6,6 +6,7 @@
 import asyncio
 import os
 import logging
+import json
 
 
 class DatabricksSqlConnector(SqlConnector):
@@ -69,3 +70,39 @@ async def query_execution(
             connection.close()
 
         return results
+
+    async def get_entity_schemas(
+        self,
+        text: Annotated[
+            str,
+            "The text to run a semantic search against. Relevant entities will be returned.",
+        ],
+        excluded_entities: Annotated[
+            list[str],
+            "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
+        ] = [],
+        as_json: bool = True,
+    ) -> str:
+        """Gets the schema of a view or table 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 schema of the views or tables in JSON format.
+        """
+
+        schemas = await self.ai_search_connector.get_entity_schemas(
+            text, excluded_entities
+        )
+
+        for schema in schemas:
+            schema["SelectFromEntity"] = ".".join(
+                [schema["Catalog"], schema["Schema"], schema["Entity"]]
+            )
+
+        if as_json:
+            return json.dumps(schemas, default=str)
+        else:
+            return schemas

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/snowflake_sql.py

@@ -6,6 +6,7 @@
 import asyncio
 import os
 import logging
+import json
 
 
 class SnowflakeSqlConnector(SqlConnector):
@@ -68,3 +69,44 @@ async def query_execution(
             conn.close()
 
         return results
+
+    async def get_entity_schemas(
+        self,
+        text: Annotated[
+            str,
+            "The text to run a semantic search against. Relevant entities will be returned.",
+        ],
+        excluded_entities: Annotated[
+            list[str],
+            "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
+        ] = [],
+        as_json: bool = True,
+    ) -> str:
+        """Gets the schema of a view or table 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 schema of the views or tables in JSON format.
+        """
+
+        schemas = await self.ai_search_connector.get_entity_schemas(
+            text, excluded_entities
+        )
+
+        for schema in schemas:
+            schema["SelectFromEntity"] = ".".join(
+                [
+                    schema["Warehouse"],
+                    schema["Database"],
+                    schema["Schema"],
+                    schema["Entity"],
+                ]
+            )
+
+        if as_json:
+            return json.dumps(schemas, default=str)
+        else:
+            return schemas

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

@@ -51,6 +51,29 @@ async def query_execution(
             list[dict]: The results of the SQL query.
         """
 
+    @abstractmethod
+    async def get_entity_schemas(
+        self,
+        text: Annotated[
+            str,
+            "The text to run a semantic search against. Relevant entities will be returned.",
+        ],
+        excluded_entities: Annotated[
+            list[str],
+            "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
+        ] = [],
+        as_json: bool = True,
+    ) -> str:
+        """Gets the schema of a view or table 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 schema of the views or tables in JSON format.
+        """
+
     async def query_execution_with_limit(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/tsql_sql.py

@@ -5,6 +5,7 @@
 from typing import Annotated
 import os
 import logging
+import json
 
 
 class TSQLSqlConnector(SqlConnector):
@@ -48,3 +49,37 @@ async def query_execution(
 
         logging.debug("Results: %s", results)
         return results
+
+    async def get_entity_schemas(
+        self,
+        text: Annotated[
+            str,
+            "The text to run a semantic search against. Relevant entities will be returned.",
+        ],
+        excluded_entities: Annotated[
+            list[str],
+            "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
+        ] = [],
+        as_json: bool = True,
+    ) -> str:
+        """Gets the schema of a view or table 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 schema of the views or tables in JSON format.
+        """
+
+        schemas = await self.ai_search_connector.get_entity_schemas(
+            text, excluded_entities
+        )
+
+        for schema in schemas:
+            schema["SelectFromEntity"] = ".".join([schema["Schema"], schema["Entity"]])
+
+        if as_json:
+            return json.dumps(schemas, default=str)
+        else:
+            return schemas

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/question_decomposition_agent.yaml

@@ -5,7 +5,7 @@ description:
 system_message:
   "You are a helpful AI Assistant that specialises in decomposing complex user questions into smaller parts that can be used in SQL queries.
 
-  If a user's question is actually a combination of multiple questions, break down the user's question into smaller questions that can be used in SQL queries.
+  If a user's question is actually a combination of multiple complex questions, break down the user's question into smaller questions that can be used in SQL queries. If a question is simple and is directly answerable by a set of SQL queries, do not decompose the question.
 
   Output Info:
     Return the decomposed questions to the user in the following format:

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_query_correction_agent.yaml

@@ -20,4 +20,4 @@ system_message:
 tools:
   - sql_get_entity_schemas_tool
   - sql_query_execution_tool
-  # - current_datetime_tool
+  - current_datetime_tool