Update generated samples

Author: BenConstable9

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
     rev: v2.1.0
     hooks:
       - id: codespell
-        args: [--exclude-file, uv.lock]
+        args: [--skip, "uv.lock,*.json,*,jsonl"]
 
   - repo: https://github.com/psf/black-pre-commit-mirror
     rev: 23.12.1
@@ -50,4 +50,4 @@ repos:
     rev: 0.5.5
     hooks:
       # Update the uv lockfile
-      - id: uv-lock
+      - id: uv-lock

…ai_search/src/deploy_ai_search/README.md deploy_ai_search/README.mddeploy_ai_search/src/deploy_ai_search/README.md renamed to deploy_ai_search/README.md deploy_ai_search/src/deploy_ai_search/README.md renamed to deploy_ai_search/README.md

@@ -50,7 +50,7 @@
       "source": [
         "import dotenv\n",
         "import logging\n",
-        "from autogen_agentchat.task import Console\n",
+        "from autogen_agentchat.ui import Console\n",
         "from autogen_text_2_sql.autogen_text_2_sql import AutoGenText2Sql"
       ]
     },
@@ -85,7 +85,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "agentic_text_2_sql = AutoGenText2Sql(target_engine=\"TSQL\", engine_specific_rules=\"Use TOP X to limit the number of rows returned instead of LIMIT X. NEVER USE LIMIT X as it produces a syntax error.\").agentic_flow"
+        "agentic_text_2_sql = AutoGenText2Sql(engine_specific_rules=\"Use TOP X to limit the number of rows returned instead of LIMIT X. NEVER USE LIMIT X as it produces a syntax error.\", use_case=\"Analysing sales data across product categories.\").agentic_flow"
       ]
     },
     {
@@ -105,13 +105,6 @@
       "source": [
         "await Console(result)"
       ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": []
     }
   ],
   "metadata": {

text_2_sql/autogen/Iteration 5 - Agentic Vector Based Text2SQL.ipynb

@@ -5,9 +5,9 @@ description = "AutoGen Based Implementation"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "autogen-agentchat==0.4.0.dev7",
-    "autogen-core==0.4.0.dev7",
-    "autogen-ext[azure,openai]==0.4.0.dev7",
+    "autogen-agentchat==0.4.0.dev9",
+    "autogen-core==0.4.0.dev9",
+    "autogen-ext[azure,openai]==0.4.0.dev9",
     "grpcio>=1.68.1",
     "pyyaml>=6.0.2",
     "text_2_sql_core",

text_2_sql/autogen/pyproject.toml

@@ -76,7 +76,7 @@ def agents(self):
     @property
     def termination_condition(self):
         """Define the termination condition for the chat."""
-        termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(10)
+        termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(15)
         return termination
 
     @staticmethod

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -24,7 +24,7 @@ def load_agent_file(cls, name: str) -> dict:
         return load(name.lower())
 
     @classmethod
-    def get_tool(cls, sql_helper, tool_name: str):
+    def get_tool(cls, sql_helper, ai_search_helper, tool_name: str):
         """Gets the tool based on the tool name.
         Args:
         ----
@@ -41,9 +41,14 @@ def get_tool(cls, sql_helper, tool_name: str):
             )
         elif tool_name == "sql_get_entity_schemas_tool":
             return FunctionTool(
-                sql_helper.get_entity_schemas,
+                ai_search_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":
+            return FunctionTool(
+                ai_search_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.",
+            )
         elif tool_name == "sql_query_validation_tool":
             return FunctionTool(
                 sql_helper.query_validation,
@@ -93,11 +98,12 @@ 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, tool))
+                tools.append(cls.get_tool(sql_helper, ai_search_helper, tool))
 
         agent = AssistantAgent(
             name=name,

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

@@ -5,8 +5,8 @@
 from autogen_agentchat.agents import BaseChatAgent
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
-from autogen_core.base import CancellationToken
-from text_2_sql_core.connectors.sql import SqlConnector
+from autogen_core import CancellationToken
+from text_2_sql_core.connectors.factory import ConnectorFactory
 import json
 import logging
 
@@ -18,7 +18,7 @@ def __init__(self):
             "An agent that fetches the queries from the cache based on the user question.",
         )
 
-        self.sql_helper = SqlConnector()
+        self.sql_connector = ConnectorFactory.get_database_connector()
 
     @property
     def produced_message_types(self) -> List[type[ChatMessage]]:
@@ -43,7 +43,9 @@ async def on_messages_stream(
         # Fetch the queries from the cache based on the user question.
         logging.info("Fetching queries from cache based on the user question...")
 
-        cached_queries = await self.sql_helper.fetch_queries_from_cache(user_question)
+        cached_queries = await self.sql_connector.fetch_queries_from_cache(
+            user_question
+        )
 
         yield Response(
             chat_message=TextMessage(

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

@@ -5,20 +5,21 @@
 from autogen_agentchat.agents import BaseChatAgent
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 from text_2_sql_core.connectors.sql import SqlConnector
 import json
 import logging
 
 
 class SqlQueryCacheAgent(BaseChatAgent):
-    def __init__(self):
+    def __init__(self, **kwargs):
         super().__init__(
-            "sql_query_cache_agent",
-            "An agent that fetches the queries from the cache based on the user question.",
+            "sql_schema_selection_agent",
+            "An agent that fetches the schemas from the cache based on the user question.",
         )
 
-        self.sql_helper = SqlConnector()
+        self.kwargs = kwargs
+        self.sql_connector = SqlConnector()
 
     @property
     def produced_message_types(self) -> List[type[ChatMessage]]:
@@ -43,7 +44,9 @@ async def on_messages_stream(
         # Fetch the queries from the cache based on the user question.
         logging.info("Fetching queries from cache based on the user question...")
 
-        cached_queries = await self.sql_helper.fetch_queries_from_cache(user_question)
+        cached_queries = await self.sql_connector.fetch_queries_from_cache(
+            user_question
+        )
 
         yield Response(
             chat_message=TextMessage(

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

@@ -88,13 +88,13 @@ A full data dictionary must be built for all the views / tables you which to exp
 
 ## Indexing
 
-`./deploy_ai_search/text_2_sql.py` & `./deploy_ai_search/text_2_sql_query_cache.py` contains the scripts to deploy and index the data dictionary for use within the plugin. See instructions in `./deploy_ai_search/README.md`.
+`./deploy_ai_search/text_2_sql.py` & `./deploy_ai_search/text_2_sql_query_cache.py` contains the scripts to deploy and index the data dictionary for use within the plugin. See instructions in `./deploy_ai_search/README.md`. There is **no automatic mechanism** to upload these .json files currently to a storage account, once generated, you must automatically upload them to the appropriate storage account that the indexer is connected to.
 
 ## Automatic Generation
 
 > [!IMPORTANT]
 >
-> - The data dictioonary generation scripts have been moved to `text_2_sql_core`. Documentation will be updated shortly.
+> - The data dictionary generation scripts have been moved to `text_2_sql_core`. Documentation will be updated shortly.
 
 Manually creating the `entities.json` is a time consuming exercise. To speed up generation, a mixture of SQL Queries and an LLM can be used to generate a initial version. Existing comments and descriptions in the database, can be combined with sample values to generate the necessary descriptions. Manual input can then be used to tweak it for the use case and any improvements.
 
@@ -109,3 +109,20 @@ The following Databases have pre-built scripts for them:
 - **TSQL:** `./text_2_sql_core/data_dictionary/tsql_data_dictionary_creator.py`
 
 If there is no pre-built script for your database engine, take one of the above as a starting point and adjust it.
+
+## Running
+
+Fill out the `.env` template with connection details to your chosen database.
+
+Package and install the `text_2_sql_core` library. See [build](https://docs.astral.sh/uv/concepts/projects/build/) if you want to build as a wheel and install on an agent. Or you can run from within a `uv` environment.
+
+`data_dictionary <DATABASE ENGINE>`
+
+You can pass the following command line arguements:
+
+- `-- output_directory` or `-o`: Optional directory that the script will write the output files to.
+- `-- single_file` or `-s`: Optional flag that writes all schemas to a single file.
+
+> [!IMPORTANT]
+>
+> - The data dictioonary generation scripts will output column values for all possible filter clauses. This could lead to output of sensitive information. You should add exclusion criteria to exclude these for only columns that you may want to filter by.