Add custom agent

Author: BenConstable9

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

@@ -94,7 +94,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "result = agentic_text_2_sql.run_stream(task=\"What are the total number of sales within 2008?\")"
+        "result = agentic_text_2_sql.run_stream(task=\"What are the total number of sales within 2008 for the mountain bike?\")"
       ]
     },
     {

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -1,11 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from autogen_agentchat.task import TextMentionTermination, MaxMessageTermination
+from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination
 from autogen_agentchat.teams import SelectorGroupChat
 from autogen_text_2_sql.creators.llm_model_creator import LLMModelCreator
 from autogen_text_2_sql.creators.llm_agent_creator import LLMAgentCreator
 import logging
 from autogen_text_2_sql.custom_agents.sql_query_cache_agent import SqlQueryCacheAgent
+from autogen_text_2_sql.custom_agents.sql_schema_selection_agent import (
+    SqlSchemaSelectionAgent,
+)
 import json
 import os
 
@@ -32,6 +35,10 @@ def set_mode(self):
             os.environ.get("Text2Sql__PreRunQueryCache", "False").lower() == "true"
         )
 
+        self.use_column_value_store = (
+            os.environ.get("Text2Sql__UseColumnValueStore", "False").lower() == "true"
+        )
+
     @property
     def agents(self):
         """Define the agents for the chat."""
@@ -41,8 +48,7 @@ def agents(self):
             engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,
         )
-        SQL_SCHEMA_SELECTION_AGENT = LLMAgentCreator.create(
-            "sql_schema_selection_agent",
+        SQL_SCHEMA_SELECTION_AGENT = SqlSchemaSelectionAgent(
             target_engine=self.target_engine,
             engine_specific_rules=self.engine_specific_rules,
             **self.kwargs,

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

@@ -1,9 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from autogen_ext.models import AzureOpenAIChatCompletionClient
-from text_2_sql_core.utils.environment import IdentityType, get_identity_type
+from text_2_sql_core.connectors.factory import ConnectorFactory
 
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import os
 import dotenv
 
@@ -28,32 +27,12 @@ def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
         else:
             raise ValueError(f"Model {model_name} not found")
 
-    @classmethod
-    def get_authentication_properties(cls) -> dict:
-        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
-            # Create the token provider
-            api_key = None
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
-            )
-        elif get_identity_type() == IdentityType.USER_ASSIGNED:
-            # Create the token provider
-            api_key = None
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["ClientId"]
-                ),
-                "https://cognitiveservices.azure.com/.default",
-            )
-        else:
-            token_provider = None
-            api_key = os.environ["OpenAI__ApiKey"]
-
-        return token_provider, api_key
-
     @classmethod
     def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
-        token_provider, api_key = cls.get_authentication_properties()
+        (
+            token_provider,
+            api_key,
+        ) = ConnectorFactory.get_open_ai_connector().get_authentication_properties()
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
             model=os.environ["OpenAI__MiniCompletionDeployment"],
@@ -70,7 +49,10 @@ def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
 
     @classmethod
     def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
-        token_provider, api_key = cls.get_authentication_properties()
+        (
+            token_provider,
+            api_key,
+        ) = ConnectorFactory.get_open_ai_connector().get_authentication_properties()
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__CompletionDeployment"],
             model=os.environ["OpenAI__CompletionDeployment"],

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

@@ -6,20 +6,28 @@
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
 from autogen_core import CancellationToken
-from text_2_sql_core.connectors.sql import SqlConnector
+from text_2_sql_core.connectors.factory import ConnectorFactory
 import json
 import logging
+from text_2_sql_core.prompts.load import load
+from jinja2 import Template
+import asyncio
 
 
-class SqlQueryCacheAgent(BaseChatAgent):
+class SqlSchemaSelectionAgent(BaseChatAgent):
     def __init__(self, **kwargs):
         super().__init__(
             "sql_schema_selection_agent",
             "An agent that fetches the schemas from the cache based on the user question.",
         )
 
-        self.kwargs = kwargs
-        self.sql_connector = SqlConnector()
+        self.ai_search_connector = ConnectorFactory.get_ai_search_connector()
+
+        self.open_ai_connector = ConnectorFactory.get_open_ai_connector()
+
+        system_prompt = load("sql_schema_selection_agent")["system_message"]
+
+        self.system_prompt = Template(system_prompt).render(kwargs)
 
     @property
     def produced_message_types(self) -> List[type[ChatMessage]]:
@@ -39,18 +47,63 @@ async def on_messages(
     async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentMessage | Response, None]:
-        user_question = messages[-1].content
+        last_response = messages[-1].content
 
-        # Fetch the queries from the cache based on the user question.
-        logging.info("Fetching queries from cache based on the user question...")
+        # load the json of the last message and get the user question's
 
-        cached_queries = await self.sql_connector.fetch_queries_from_cache(
-            user_question
-        )
+        user_questions = json.loads(last_response)
+
+        logging.info(f"User questions: {user_questions}")
+
+        entity_tasks = []
+
+        for user_question in user_questions:
+            messages = [
+                {"role": "system", "content": self.system_prompt},
+                {"role": "user", "content": user_question},
+            ]
+            entity_tasks.append(self.open_ai_connector.run_completion_request(messages))
+
+        entity_results = await asyncio.gather(*entity_tasks)
+
+        entity_search_tasks = []
+        column_search_tasks = []
+
+        for entity_result in entity_results:
+            loaded_entity_result = json.loads(entity_result)
+
+            logging.info(f"Loaded entity result: {loaded_entity_result}")
+
+            entity_search_tasks.append(
+                self.ai_search_connector.get_entity_schemas(
+                    " ".join(loaded_entity_result["entities"]), as_json=False
+                )
+            )
+
+            for filter_condition in loaded_entity_result["filter_conditions"]:
+                column_search_tasks.append(
+                    self.ai_search_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)
+
+        final_results = {
+            "schemas": [
+                schema for schema_result in schemas_results for schema in schema_result
+            ],
+            "column_values": [
+                column_values
+                for column_values_result in column_value_results
+                for column_values in column_values_result
+            ],
+        }
 
         yield Response(
             chat_message=TextMessage(
-                content=json.dumps(cached_queries), source=self.name
+                content=json.dumps(final_results), source=self.name
             )
         )
 

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

@@ -29,24 +29,29 @@ async def run_ai_search_query(
         """Run the AI search query."""
         identity_type = get_identity_type()
 
-        async with AsyncAzureOpenAI(
-            # This is the default and can be omitted
-            api_key=os.environ["OpenAI__ApiKey"],
-            azure_endpoint=os.environ["OpenAI__Endpoint"],
-            api_version=os.environ["OpenAI__ApiVersion"],
-        ) as open_ai_client:
-            embeddings = await open_ai_client.embeddings.create(
-                model=os.environ["OpenAI__EmbeddingModel"], input=query
-            )
+        if len(vector_fields) > 0:
+            async with AsyncAzureOpenAI(
+                # This is the default and can be omitted
+                api_key=os.environ["OpenAI__ApiKey"],
+                azure_endpoint=os.environ["OpenAI__Endpoint"],
+                api_version=os.environ["OpenAI__ApiVersion"],
+            ) as open_ai_client:
+                embeddings = await open_ai_client.embeddings.create(
+                    model=os.environ["OpenAI__EmbeddingModel"], input=query
+                )
 
-            # Extract the embedding vector
-            embedding_vector = embeddings.data[0].embedding
+                # Extract the embedding vector
+                embedding_vector = embeddings.data[0].embedding
 
-        vector_query = VectorizedQuery(
-            vector=embedding_vector,
-            k_nearest_neighbors=7,
-            fields=",".join(vector_fields),
-        )
+            vector_query = [
+                VectorizedQuery(
+                    vector=embedding_vector,
+                    k_nearest_neighbors=7,
+                    fields=",".join(vector_fields),
+                )
+            ]
+        else:
+            vector_query = None
 
         if identity_type == IdentityType.SYSTEM_ASSIGNED:
             credential = DefaultAzureCredential()
@@ -63,13 +68,20 @@ async def run_ai_search_query(
             index_name=index_name,
             credential=credential,
         ) as search_client:
+            if semantic_config is not None and vector_query is not None:
+                query_type = "semantic"
+            elif vector_query is not None:
+                query_type = "hybrid"
+            else:
+                query_type = "full"
+
             results = await search_client.search(
                 top=top,
                 semantic_configuration_name=semantic_config,
                 search_text=query,
                 select=",".join(retrieval_fields),
-                vector_queries=[vector_query],
-                query_type="semantic",
+                vector_queries=vector_query,
+                query_type=query_type,
                 query_language="en-GB",
             )
 
@@ -102,6 +114,7 @@ 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.
 
@@ -113,20 +126,24 @@ async def get_column_values(
         -------
             str: The values of the column in JSON format.
         """
+
+        # Adds tildes after each text word to do a fuzzy search
+        text = " ".join([f"{word}~" for word in text.split()])
         values = await self.run_ai_search_query(
             text,
             [],
             ["FQN", "Column", "Value"],
             os.environ[
                 "AIService__AzureSearchOptions__Text2SqlColumnValueStore__Index"
             ],
-            os.environ[
-                "AIService__AzureSearchOptions__Text2SqlColumnValueStore__SemanticConfig"
-            ],
+            None,
             top=10,
         )
 
-        return json.dumps(values, default=str)
+        if as_json:
+            return json.dumps(values, default=str)
+        else:
+            return values
 
     async def get_entity_schemas(
         self,
@@ -138,6 +155,7 @@ 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.
 
@@ -178,7 +196,10 @@ async def get_entity_schemas(
                 else:
                     filtered_schemas.append(schema)
 
-        return json.dumps(schemas, default=str)
+        if as_json:
+            return json.dumps(schemas, default=str)
+        else:
+            return 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/factory.py

@@ -1,5 +1,6 @@
 import os
 from text_2_sql_core.connectors.ai_search import AISearchConnector
+from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
 
 class ConnectorFactory:
@@ -36,3 +37,7 @@ def get_database_connector():
     @staticmethod
     def get_ai_search_connector():
         return AISearchConnector()
+
+    @staticmethod
+    def get_open_ai_connector():
+        return OpenAIConnector()

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/open_ai.py

@@ -0,0 +1,45 @@
+from openai import AsyncAzureOpenAI
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+import os
+import dotenv
+from text_2_sql_core.utils.environment import IdentityType, get_identity_type
+
+dotenv.load_dotenv()
+
+
+class OpenAIConnector:
+    @classmethod
+    def get_authentication_properties(cls) -> dict:
+        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
+            # Create the token provider
+            api_key = None
+            token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+        elif get_identity_type() == IdentityType.USER_ASSIGNED:
+            # Create the token provider
+            api_key = None
+            token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(
+                    managed_identity_client_id=os.environ["ClientId"]
+                ),
+                "https://cognitiveservices.azure.com/.default",
+            )
+        else:
+            token_provider = None
+            api_key = os.environ["OpenAI__ApiKey"]
+
+        return token_provider, api_key
+
+    async def run_completion_request(self, messages: list[dict], temperature=0):
+        async with AsyncAzureOpenAI(
+            api_key=os.environ["OpenAI__ApiKey"],
+            azure_endpoint=os.environ["OpenAI__Endpoint"],
+            api_version=os.environ["OpenAI__ApiVersion"],
+        ) as open_ai_client:
+            response = await open_ai_client.chat.completions.create(
+                model=os.environ["OpenAI__MiniCompletionDeployment"],
+                messages=messages,
+                temperature=temperature,
+            )
+        return response.choices[0].message.content