Merge branch 'main' into feature/semantic-chunking

Author: BenConstable9

deploy_ai_search/.env

@@ -20,4 +20,4 @@ OpenAI__Endpoint=<openAIEndpoint>
 OpenAI__EmbeddingModel=<openAIEmbeddingModelName>
 OpenAI__EmbeddingDeployment=<openAIEmbeddingDeploymentId>
 OpenAI__EmbeddingDimensions=1536
-Text2Sql__DatabaseName=<databaseName>
+Text2Sql__DatabaseEngine=<databaseEngine SQL Server / Snowflake / Databricks >

deploy_ai_search/text_2_sql_query_cache.py

@@ -7,6 +7,10 @@
     SearchableField,
     SimpleField,
     ComplexField,
+    SemanticField,
+    SemanticPrioritizedFields,
+    SemanticConfiguration,
+    SemanticSearch,
 )
 from ai_search import AISearch
 from environment import (
@@ -107,3 +111,22 @@ def get_index_fields(self) -> list[SearchableField]:
         ]
 
         return fields
+
+    def get_semantic_search(self) -> SemanticSearch:
+        """This function returns the semantic search configuration for sql index
+
+        Returns:
+            SemanticSearch: The semantic search configuration"""
+
+        semantic_config = SemanticConfiguration(
+            name=self.semantic_config_name,
+            prioritized_fields=SemanticPrioritizedFields(
+                content_fields=[
+                    SemanticField(field_name="Question"),
+                ],
+            ),
+        )
+
+        semantic_search = SemanticSearch(configurations=[semantic_config])
+
+        return semantic_search

deploy_ai_search/text_2_sql_schema_store.py

@@ -24,6 +24,16 @@
 from environment import (
     IndexerType,
 )
+import os
+from enum import StrEnum
+
+
+class DatabaseEngine(StrEnum):
+    """An enumeration to represent a database engine."""
+
+    SNOWFLAKE = "SNOWFLAKE"
+    SQL_SERVER = "SQL_SERVER"
+    DATABRICKS = "DATABRICKS"
 
 
 class Text2SqlSchemaStoreAISearch(AISearch):
@@ -42,13 +52,34 @@ def __init__(
             rebuild (bool, optional): Whether to rebuild the index. Defaults to False.
         """
         self.indexer_type = IndexerType.TEXT_2_SQL_SCHEMA_STORE
+        self.database_engine = DatabaseEngine[
+            os.environ["Text2Sql__DatabaseEngine"].upper()
+        ]
         super().__init__(suffix, rebuild)
 
         if single_data_dictionary:
             self.parsing_mode = BlobIndexerParsingMode.JSON_ARRAY
         else:
             self.parsing_mode = BlobIndexerParsingMode.JSON
 
+    @property
+    def excluded_fields_for_database_engine(self):
+        """A method to get the excluded fields for the database engine."""
+
+        all_engine_specific_fields = ["Warehouse", "Database", "Catalog"]
+        if self.database_engine == DatabaseEngine.SNOWFLAKE:
+            engine_specific_fields = ["Warehouse", "Database"]
+        elif self.database_engine == DatabaseEngine.SQL_SERVER:
+            engine_specific_fields = ["Database"]
+        elif self.database_engine == DatabaseEngine.DATABRICKS:
+            engine_specific_fields = ["Catalog"]
+
+        return [
+            field
+            for field in all_engine_specific_fields
+            if field not in engine_specific_fields
+        ]
+
     def get_index_fields(self) -> list[SearchableField]:
         """This function returns the index fields for sql index.
 
@@ -78,6 +109,10 @@ def get_index_fields(self) -> list[SearchableField]:
                 name="Warehouse",
                 type=SearchFieldDataType.String,
             ),
+            SearchableField(
+                name="Catalog",
+                type=SearchFieldDataType.String,
+            ),
             SearchableField(
                 name="Definition",
                 type=SearchFieldDataType.String,
@@ -161,6 +196,13 @@ def get_index_fields(self) -> list[SearchableField]:
             ),
         ]
 
+        # Remove fields that are not supported by the database engine
+        fields = [
+            field
+            for field in fields
+            if field.name not in self.excluded_fields_for_database_engine
+        ]
+
         return fields
 
     def get_semantic_search(self) -> SemanticSearch:
@@ -309,4 +351,12 @@ def get_indexer(self) -> SearchIndexer:
             parameters=indexer_parameters,
         )
 
+        # Remove fields that are not supported by the database engine
+        indexer.output_field_mappings = [
+            field_mapping
+            for field_mapping in indexer.output_field_mappings
+            if field_mapping.target_field_name
+            not in self.excluded_fields_for_database_engine
+        ]
+
         return indexer

text_2_sql/autogen/README.md

@@ -0,0 +1,3 @@
+# Multi-Shot Text2SQL Component - AutoGen
+
+Very much still work in progress, more documentation coming soon.

text_2_sql/autogen/agentic_text_2_sql.ipynb

@@ -0,0 +1,80 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import dotenv\n",
+        "import logging\n",
+        "from autogen_agentchat.task import Console\n",
+        "from agentic_text_2_sql import text_2_sql_generator"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "logging.basicConfig(level=logging.INFO)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "dotenv.load_dotenv()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "result = text_2_sql_generator.run_stream(task=\"What are the total number of sales within 2008?\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "await Console(result)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.12.6"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+}

text_2_sql/autogen/agentic_text_2_sql.py

@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from autogen_agentchat.task import TextMentionTermination, MaxMessageTermination
+from autogen_agentchat.teams import SelectorGroupChat
+from utils.models import MINI_MODEL
+from utils.llm_agent_creator import LLMAgentCreator
+import logging
+from custom_agents.sql_query_cache_agent import SqlQueryCacheAgent
+import json
+
+SQL_QUERY_GENERATION_AGENT = LLMAgentCreator.create(
+    "sql_query_generation_agent",
+    target_engine="Microsoft SQL Server",
+    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.",
+)
+SQL_SCHEMA_SELECTION_AGENT = LLMAgentCreator.create("sql_schema_selection_agent")
+SQL_QUERY_CORRECTION_AGENT = LLMAgentCreator.create(
+    "sql_query_correction_agent",
+    target_engine="Microsoft SQL Server",
+    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.",
+)
+SQL_QUERY_CACHE_AGENT = SqlQueryCacheAgent()
+ANSWER_AGENT = LLMAgentCreator.create("answer_agent")
+QUESTION_DECOMPOSITION_AGENT = LLMAgentCreator.create("question_decomposition_agent")
+
+
+def text_2_sql_generator_selector_func(messages):
+    logging.info("Messages: %s", messages)
+    decision = None  # Initialize decision variable
+
+    if len(messages) == 1:
+        decision = "sql_query_cache_agent"
+
+    elif (
+        messages[-1].source == "sql_query_cache_agent"
+        and messages[-1].content is not None
+    ):
+        cache_result = json.loads(messages[-1].content)
+        if cache_result.get("cached_questions_and_schemas") is not None:
+            decision = "sql_query_correction_agent"
+        else:
+            decision = "sql_schema_selection_agent"
+
+    elif messages[-1].source == "question_decomposition_agent":
+        decision = "sql_schema_selection_agent"
+
+    elif messages[-1].source == "sql_schema_selection_agent":
+        decision = "sql_query_generation_agent"
+
+    elif (
+        messages[-1].source == "sql_query_correction_agent"
+        and messages[-1].content == "VALIDATED"
+    ):
+        decision = "answer_agent"
+
+    elif messages[-1].source == "sql_query_correction_agent":
+        decision = "sql_query_correction_agent"
+
+    # Log the decision
+    logging.info("Decision: %s", decision)
+
+    return decision
+
+
+termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(10)
+text_2_sql_generator = SelectorGroupChat(
+    [
+        SQL_QUERY_GENERATION_AGENT,
+        SQL_SCHEMA_SELECTION_AGENT,
+        SQL_QUERY_CORRECTION_AGENT,
+        SQL_QUERY_CACHE_AGENT,
+        ANSWER_AGENT,
+        QUESTION_DECOMPOSITION_AGENT,
+    ],
+    allow_repeated_speaker=False,
+    model_client=MINI_MODEL,
+    termination_condition=termination,
+    selector_func=text_2_sql_generator_selector_func,
+)

text_2_sql/autogen/custom_agents/__init__.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from typing import AsyncGenerator, List, Sequence
+
+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 utils.sql_utils import fetch_queries_from_cache
+import json
+import logging
+
+
+class SqlQueryCacheAgent(BaseChatAgent):
+    def __init__(self):
+        super().__init__(
+            "sql_query_cache_agent",
+            "An agent that fetches the queries from the cache based on the user question.",
+        )
+
+    @property
+    def produced_message_types(self) -> List[type[ChatMessage]]:
+        return [TextMessage]
+
+    async def on_messages(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> Response:
+        # Calls the on_messages_stream.
+        response: Response | None = None
+        async for message in self.on_messages_stream(messages, cancellation_token):
+            if isinstance(message, Response):
+                response = message
+        assert response is not None
+        return response
+
+    async def on_messages_stream(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[AgentMessage | Response, None]:
+        user_question = messages[0].content
+
+        # 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 fetch_queries_from_cache(user_question)
+
+        yield Response(
+            chat_message=TextMessage(
+                content=json.dumps(cached_queries), source=self.name
+            )
+        )
+
+    async def on_reset(self, cancellation_token: CancellationToken) -> None:
+        pass

text_2_sql/autogen/custom_agents/sql_query_cache_agent.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import os
+from enum import Enum
+
+
+class IdentityType(Enum):
+    """The type of the indexer"""
+
+    USER_ASSIGNED = "user_assigned"
+    SYSTEM_ASSIGNED = "system_assigned"
+    KEY = "key"
+
+
+def get_identity_type() -> IdentityType:
+    """This function returns the identity type.
+
+    Returns:
+        IdentityType: The identity type
+    """
+    identity = os.environ.get("IdentityType")
+
+    if identity == "user_assigned":
+        return IdentityType.USER_ASSIGNED
+    elif identity == "system_assigned":
+        return IdentityType.SYSTEM_ASSIGNED
+    elif identity == "key":
+        return IdentityType.KEY
+    else:
+        raise ValueError("Invalid identity type")