Add a query cache POC into Text2SQL (#19)

Author: BenConstable9

README.md

@@ -6,8 +6,8 @@ It is intended that the plugins and skills provided in this repository, are adap
 
 ## Components
 
-- `./text_2_sql` contains an two Multi-Shot implementation for Text2SQL generation and querying which can be used to answer questions backed by a database as a knowledge base. A prompt based and vector based approach are shown, both of which exhibit great performance in answering sql queries. With these plugins, your RAG application can now access and pull data from any SQL table exposed to it to answer questions.
-- `./adi_function_app` contains code for linking Azure Document Intelligence with AI Search to process complex documents with charts and images, and uses multi-modal models (gpt4o) to interpret and understand these. With this custom skill, the RAG application can draw insights from complex charts and images during the vector search.
+- `./text_2_sql` contains an three Multi-Shot implementations for Text2SQL generation and querying which can be used to answer questions backed by a database as a knowledge base. A **prompt based** and **vector based** approach are shown, both of which exhibit great performance in answering sql queries. Additionally, a further iteration on the vector based approach is shown which uses a **query cache** to further speed up generation.  With these plugins, your RAG application can now access and pull data from any SQL table exposed to it to answer questions.
+- `./adi_function_app` contains code for linking **Azure Document Intelligence** with AI Search to process complex documents with charts and images, and uses **multi-modal models (gpt4o)** to interpret and understand these. With this custom skill, the RAG application can **draw insights from complex charts** and images during the vector search.
 - `./deploy_ai_search` provides an easy Python based utility for deploying an index, indexer and corresponding skillset for AI Search and for Text2SQL.
 
 The above components have been successfully used on production RAG projects to increase the quality of responses.

adi_function_app/README.md

@@ -206,6 +206,12 @@ If `chunk_by_page` header is `False`:
 
 **Page wise analysis in ADI is recommended to avoid splitting tables / figures across multiple chunks, when the chunking is performed.**
 
+## Other Provided Custom Skills
+
+Due to a AI Search product limitation that AI Search cannot connect to AI Services behind Private Endpoints, we provide a Custom Key Phrase Extraction Skill that will work within a Private Endpoint environment.
+
+Additionally, a custom cleaning skill is provided to clean the chunks before vectorisation takes place.
+
 ## Production Considerations
 
 Below are some of the considerations that should be made before using this custom skill in production:

deploy_ai_search/README.md

@@ -8,18 +8,30 @@ The associated scripts in this portion of the repository contains pre-built scri
 2. Adjust `rag_documents.py` with any changes to the index / indexer. The `get_skills()` method implements the skills pipeline. Make any adjustments here in the skills needed to enrich the data source.
 3. Run `deploy.py` with the following args:
 
-    - `indexer_type rag`. This selects the `RagDocumentsAISearch` sub class.
+    - `index_type rag`. This selects the `RagDocumentsAISearch` sub class.
     - `enable_page_chunking True`. This determines whether page wise chunking is applied in ADI, or whether the inbuilt skill is used for TextSplit. **Page wise analysis in ADI is recommended to avoid splitting tables / figures across multiple chunks, when the chunking is performed.**
     - `rebuild`. Whether to delete and rebuild the index.
     - `suffix`. Optional parameter that will apply a suffix onto the deployed index and indexer. This is useful if you want deploy a test version, before overwriting the main version.
 
 ## Steps for Text2SQL Index Deployment
 
+### Entity Schema Index
+
 1. Update `.env` file with the associated values. Not all values are required dependent on whether you are using System / User Assigned Identities or a Key based authentication.
 2. Adjust `text_2_sql.py` with any changes to the index / indexer. The `get_skills()` method implements the skills pipeline. Make any adjustments here in the skills needed to enrich the data source.
 3. Run `deploy.py` with the following args:
 
-    - `indexer_type text_2_sql`. This selects the `Text2SQLAISearch` sub class.
+    - `index_type text_2_sql`. This selects the `Text2SQLAISearch` sub class.
+    - `rebuild`. Whether to delete and rebuild the index.
+    - `suffix`. Optional parameter that will apply a suffix onto the deployed index and indexer. This is useful if you want deploy a test version, before overwriting the main version.
+
+### Query Cache Index
+
+1. Update `.env` file with the associated values. Not all values are required dependent on whether you are using System / User Assigned Identities or a Key based authentication.
+2. Adjust `text_2_sql_query_cache.py` with any changes to the index. **There is no provided indexer or skillset for this cache, it is expected that application code will write directly to it.**
+3. Run `deploy.py` with the following args:
+
+    - `index_type text_2_sql_query_cache`. This selects the `Text2SQLQueryCacheAISearch` sub class.
     - `rebuild`. Whether to delete and rebuild the index.
     - `suffix`. Optional parameter that will apply a suffix onto the deployed index and indexer. This is useful if you want deploy a test version, before overwriting the main version.
 

deploy_ai_search/deploy.py

@@ -3,23 +3,28 @@
 import argparse
 from rag_documents import RagDocumentsAISearch
 from text_2_sql import Text2SqlAISearch
+from text_2_sql_query_cache import Text2SqlQueryCacheAISearch
 
 
 def deploy_config(arguments: argparse.Namespace):
     """Deploy the indexer configuration based on the arguments passed.
 
     Args:
         arguments (argparse.Namespace): The arguments passed to the script"""
-    if arguments.indexer_type == "rag":
+    if arguments.index_type == "rag":
         index_config = RagDocumentsAISearch(
             suffix=arguments.suffix,
             rebuild=arguments.rebuild,
             enable_page_by_chunking=arguments.enable_page_chunking,
         )
-    elif arguments.indexer_type == "text_2_sql":
+    elif arguments.index_type == "text_2_sql":
         index_config = Text2SqlAISearch(
             suffix=arguments.suffix, rebuild=arguments.rebuild
         )
+    elif arguments.index_type == "text_2_sql_query_cache":
+        index_config = Text2SqlQueryCacheAISearch(
+            suffix=arguments.suffix, rebuild=arguments.rebuild
+        )
     else:
         raise ValueError("Invalid Indexer Type")
 
@@ -32,7 +37,7 @@ def deploy_config(arguments: argparse.Namespace):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Process some arguments.")
     parser.add_argument(
-        "--indexer_type",
+        "--index_type",
         type=str,
         required=True,
         help="Type of Indexer want to deploy.",

deploy_ai_search/environment.py

@@ -13,6 +13,7 @@ class IndexerType(Enum):
 
     RAG_DOCUMENTS = "rag-documents"
     TEXT_2_SQL = "text-2-sql"
+    TEXT_2_SQL_QUERY_CACHE = "text-2-sql-query-cache"
 
 
 class IdentityType(Enum):

deploy_ai_search/text_2_sql.py

@@ -40,8 +40,6 @@ def __init__(self, suffix: str | None = None, rebuild: bool | None = False):
 
         self.parsing_mode = BlobIndexerParsingMode.JSON_ARRAY
 
-        self.entities = []
-
     def get_index_fields(self) -> list[SearchableField]:
         """This function returns the index fields for sql index.
 

deploy_ai_search/text_2_sql_query_cache.py

@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from azure.search.documents.indexes.models import (
+    SearchFieldDataType,
+    SearchField,
+    SearchableField,
+    SemanticField,
+    SemanticPrioritizedFields,
+    SemanticConfiguration,
+    SemanticSearch,
+    SimpleField,
+    ComplexField,
+)
+from ai_search import AISearch
+from environment import (
+    IndexerType,
+)
+
+
+class Text2SqlQueryCacheAISearch(AISearch):
+    """This class is used to deploy the sql index."""
+
+    def __init__(self, suffix: str | None = None, rebuild: bool | None = False):
+        """Initialize the Text2SqlAISearch class. This class implements the deployment of the sql index.
+
+        Args:
+            suffix (str, optional): The suffix for the indexer. Defaults to None. If an suffix is provided, it is assumed to be a test indexer.
+            rebuild (bool, optional): Whether to rebuild the index. Defaults to False.
+        """
+        self.indexer_type = IndexerType.TEXT_2_SQL_QUERY_CACHE
+        super().__init__(suffix, rebuild)
+
+    def get_index_fields(self) -> list[SearchableField]:
+        """This function returns the index fields for sql index.
+
+        Returns:
+            list[SearchableField]: The index fields for sql index"""
+
+        fields = [
+            SimpleField(
+                name="Id", type=SearchFieldDataType.String, key=True, retrievable=False
+            ),
+            SearchableField(
+                name="Question",
+                type=SearchFieldDataType.String,
+                analyzer_name="keyword",
+            ),
+            SearchField(
+                name="QuestionEmbedding",
+                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+                vector_search_dimensions=self.environment.open_ai_embedding_dimensions,
+                vector_search_profile_name=self.vector_search_profile_name,
+            ),
+            SearchableField(
+                name="Query", type=SearchFieldDataType.String, filterable=True
+            ),
+            SearchField(
+                name="QueryEmbedding",
+                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+                vector_search_dimensions=self.environment.open_ai_embedding_dimensions,
+                vector_search_profile_name=self.vector_search_profile_name,
+            ),
+            ComplexField(
+                name="Schemas",
+                collection=True,
+                fields=[
+                    SearchableField(
+                        name="Entity",
+                        type=SearchFieldDataType.String,
+                        filterable=True,
+                    ),
+                    ComplexField(
+                        name="Columns",
+                        collection=True,
+                        fields=[
+                            SearchableField(
+                                name="Name", type=SearchFieldDataType.String
+                            ),
+                            SearchableField(
+                                name="Definition", type=SearchFieldDataType.String
+                            ),
+                            SearchableField(
+                                name="Type", type=SearchFieldDataType.String
+                            ),
+                            SimpleField(
+                                name="AllowedValues",
+                                type=SearchFieldDataType.String,
+                                collection=True,
+                            ),
+                            SimpleField(
+                                name="SampleValues",
+                                type=SearchFieldDataType.String,
+                                collection=True,
+                            ),
+                        ],
+                    ),
+                ],
+            ),
+        ]
+
+        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(
+                title_field=SemanticField(field_name="Question"),
+                keywords_fields=[
+                    SemanticField(field_name="Query"),
+                ],
+            ),
+        )
+
+        semantic_search = SemanticSearch(configurations=[semantic_config])
+
+        return semantic_search

images/Plugin Based RAG Flow.png

@@ -1,16 +1,19 @@
-OpenAI__CompletionDeployment=<openAICompletionDeploymentId>
-OpenAI__EmbeddingModel=<openAIEmbeddingModelName>
-OpenAI__Endpoint=<openAIEndpoint>
-OpenAI__ApiKey=<openAIKey if using non managed identity>
-OpenAI__ApiVersion=<openAIApiVersion>
-Text2Sql__DatabaseEngine=<databaseEngine>
-Text2Sql__DatabaseName=<databaseName>
-Text2Sql__DatabaseConnectionString=<databaseConnectionString>
-AIService__AzureSearchOptions__Endpoint=<searchServiceEndpoint>
-AIService__AzureSearchOptions__Key=<searchServiceKey if not using identity>
-AIService__AzureSearchOptions__RagDocuments__Index=<ragDocumentsIndexName>
-AIService__AzureSearchOptions__Text2Sql__Index=<text2SQLIndexName>
-AIService__AzureSearchOptions__RagDocuments__SemanticConfig=<ragDocumentsSemanticConfig>
-AIService__AzureSearchOptions__Text2Sql__SemanticConfig=<text2SQLSemanticConfig>
-IdentityType=<identityType> # system_assigned or user_assigned or key
-ClientId=<clientId if using user assigned identity>
+OpenAI__CompletionDeployment=<openAICompletionDeploymentId>
+OpenAI__EmbeddingModel=<openAIEmbeddingModelName>
+OpenAI__Endpoint=<openAIEndpoint>
+OpenAI__ApiKey=<openAIKey if using non managed identity>
+OpenAI__ApiVersion=<openAIApiVersion>
+Text2Sql__DatabaseEngine=<databaseEngine>
+Text2Sql__UseQueryCache=<whether to use the query cache first or not>
+Text2Sql__DatabaseName=<databaseName>
+Text2Sql__DatabaseConnectionString=<databaseConnectionString>
+AIService__AzureSearchOptions__Endpoint=<searchServiceEndpoint>
+AIService__AzureSearchOptions__Key=<searchServiceKey if not using identity>
+AIService__AzureSearchOptions__RagDocuments__Index=<ragDocumentsIndexName>
+AIService__AzureSearchOptions__Text2Sql__Index=<text2SQLIndexName>
+AIService__AzureSearchOptions__Text2SqlQueryCache__Index=<text2SQLIndexName>
+AIService__AzureSearchOptions__RagDocuments__SemanticConfig=<ragDocumentsSemanticConfig>
+AIService__AzureSearchOptions__Text2Sql__SemanticConfig=<text2SQLSemanticConfig>
+AIService__AzureSearchOptions__Text2SqlQueryCache__SemanticConfig=<text2SQLSemanticConfig>
+IdentityType=<identityType> # system_assigned or user_assigned or key
+ClientId=<clientId if using user assigned identity>