Add query rewriting option (#2437)

Upgrade Search SDK and add support for query rewriting by setting AZURE_SEARCH_QUERY_REWRITING

---------

Co-authored-by: Pamela Fox <[email protected]>

Author: mattgotteiner

.azdo/pipelines/azure-dev.yml

@@ -59,6 +59,7 @@ steps:
       AZURE_SEARCH_QUERY_LANGUAGE: $(AZURE_SEARCH_QUERY_LANGUAGE)
       AZURE_SEARCH_QUERY_SPELLER: $(AZURE_SEARCH_QUERY_SPELLER)
       AZURE_SEARCH_SEMANTIC_RANKER: $(AZURE_SEARCH_SEMANTIC_RANKER)
+      AZURE_SEARCH_QUERY_REWRITING: $(AZURE_SEARCH_QUERY_REWRITING)
       AZURE_STORAGE_ACCOUNT: $(AZURE_STORAGE_ACCOUNT)
       AZURE_STORAGE_RESOURCE_GROUP: $(AZURE_STORAGE_RESOURCE_GROUP)
       AZURE_STORAGE_SKU: $(AZURE_STORAGE_SKU)

.github/workflows/azure-dev.yml

@@ -49,6 +49,7 @@ jobs:
       AZURE_SEARCH_QUERY_LANGUAGE: ${{ vars.AZURE_SEARCH_QUERY_LANGUAGE }}
       AZURE_SEARCH_QUERY_SPELLER: ${{ vars.AZURE_SEARCH_QUERY_SPELLER }}
       AZURE_SEARCH_SEMANTIC_RANKER: ${{ vars.AZURE_SEARCH_SEMANTIC_RANKER }}
+      AZURE_SEARCH_QUERY_REWRITING: ${{ vars.AZURE_SEARCH_QUERY_REWRITING }}
       AZURE_STORAGE_ACCOUNT: ${{ vars.AZURE_STORAGE_ACCOUNT }}
       AZURE_STORAGE_RESOURCE_GROUP: ${{ vars.AZURE_STORAGE_RESOURCE_GROUP }}
       AZURE_STORAGE_SKU: ${{ vars.AZURE_STORAGE_SKU }}

app/backend/app.py

@@ -69,6 +69,7 @@
     CONFIG_INGESTER,
     CONFIG_LANGUAGE_PICKER_ENABLED,
     CONFIG_OPENAI_CLIENT,
+    CONFIG_QUERY_REWRITING_ENABLED,
     CONFIG_SEARCH_CLIENT,
     CONFIG_SEMANTIC_RANKER_DEPLOYED,
     CONFIG_SPEECH_INPUT_ENABLED,
@@ -291,6 +292,7 @@ def config():
         {
             "showGPT4VOptions": current_app.config[CONFIG_GPT4V_DEPLOYED],
             "showSemanticRankerOption": current_app.config[CONFIG_SEMANTIC_RANKER_DEPLOYED],
+            "showQueryRewritingOption": current_app.config[CONFIG_QUERY_REWRITING_ENABLED],
             "showVectorOption": current_app.config[CONFIG_VECTOR_SEARCH_ENABLED],
             "showUserUpload": current_app.config[CONFIG_USER_UPLOAD_ENABLED],
             "showLanguagePicker": current_app.config[CONFIG_LANGUAGE_PICKER_ENABLED],
@@ -453,6 +455,7 @@ async def setup_clients():
     AZURE_SEARCH_QUERY_LANGUAGE = os.getenv("AZURE_SEARCH_QUERY_LANGUAGE") or "en-us"
     AZURE_SEARCH_QUERY_SPELLER = os.getenv("AZURE_SEARCH_QUERY_SPELLER") or "lexicon"
     AZURE_SEARCH_SEMANTIC_RANKER = os.getenv("AZURE_SEARCH_SEMANTIC_RANKER", "free").lower()
+    AZURE_SEARCH_QUERY_REWRITING = os.getenv("AZURE_SEARCH_QUERY_REWRITING", "false").lower()
 
     AZURE_SPEECH_SERVICE_ID = os.getenv("AZURE_SPEECH_SERVICE_ID")
     AZURE_SPEECH_SERVICE_LOCATION = os.getenv("AZURE_SPEECH_SERVICE_LOCATION")
@@ -634,6 +637,9 @@ async def setup_clients():
 
     current_app.config[CONFIG_GPT4V_DEPLOYED] = bool(USE_GPT4V)
     current_app.config[CONFIG_SEMANTIC_RANKER_DEPLOYED] = AZURE_SEARCH_SEMANTIC_RANKER != "disabled"
+    current_app.config[CONFIG_QUERY_REWRITING_ENABLED] = (
+        AZURE_SEARCH_QUERY_REWRITING == "true" and AZURE_SEARCH_SEMANTIC_RANKER != "disabled"
+    )
     current_app.config[CONFIG_VECTOR_SEARCH_ENABLED] = os.getenv("USE_VECTORS", "").lower() != "false"
     current_app.config[CONFIG_USER_UPLOAD_ENABLED] = bool(USE_USER_UPLOAD)
     current_app.config[CONFIG_LANGUAGE_PICKER_ENABLED] = ENABLE_LANGUAGE_PICKER

app/backend/approaches/approach.py

@@ -149,6 +149,7 @@ async def search(
         use_semantic_captions: bool,
         minimum_search_score: Optional[float],
         minimum_reranker_score: Optional[float],
+        use_query_rewriting: Optional[bool] = None,
     ) -> List[Document]:
         search_text = query_text if use_text_search else ""
         search_vectors = vectors if use_vector_search else []
@@ -158,6 +159,7 @@ async def search(
                 filter=filter,
                 top=top,
                 query_caption="extractive|highlight-false" if use_semantic_captions else None,
+                query_rewrites="generative" if use_query_rewriting else None,
                 vector_queries=search_vectors,
                 query_type=QueryType.SEMANTIC,
                 query_language=self.query_language,

app/backend/approaches/chatreadretrieveread.py

@@ -89,6 +89,7 @@ async def run_until_final_call(
         use_vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         use_semantic_ranker = True if overrides.get("semantic_ranker") else False
         use_semantic_captions = True if overrides.get("semantic_captions") else False
+        use_query_rewriting = True if overrides.get("query_rewriting") else False
         top = overrides.get("top", 3)
         minimum_search_score = overrides.get("minimum_search_score", 0.0)
         minimum_reranker_score = overrides.get("minimum_reranker_score", 0.0)
@@ -147,6 +148,7 @@ async def run_until_final_call(
             use_semantic_captions,
             minimum_search_score,
             minimum_reranker_score,
+            use_query_rewriting,
         )
 
         # STEP 3: Generate a contextual and content specific answer using the search results and chat history
@@ -190,6 +192,7 @@ async def run_until_final_call(
                     {
                         "use_semantic_captions": use_semantic_captions,
                         "use_semantic_ranker": use_semantic_ranker,
+                        "use_query_rewriting": use_query_rewriting,
                         "top": top,
                         "filter": filter,
                         "use_vector_search": use_vector_search,

app/backend/approaches/chatreadretrievereadvision.py

@@ -81,6 +81,7 @@ async def run_until_final_call(
         use_text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         use_vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         use_semantic_ranker = True if overrides.get("semantic_ranker") else False
+        use_query_rewriting = True if overrides.get("query_rewriting") else False
         use_semantic_captions = True if overrides.get("semantic_captions") else False
         top = overrides.get("top", 3)
         minimum_search_score = overrides.get("minimum_search_score", 0.0)
@@ -151,6 +152,7 @@ async def run_until_final_call(
             use_semantic_captions,
             minimum_search_score,
             minimum_reranker_score,
+            use_query_rewriting,
         )
 
         # STEP 3: Generate a contextual and content specific answer using the search results and chat history
@@ -207,6 +209,7 @@ async def run_until_final_call(
                     {
                         "use_semantic_captions": use_semantic_captions,
                         "use_semantic_ranker": use_semantic_ranker,
+                        "use_query_rewriting": use_query_rewriting,
                         "top": top,
                         "filter": filter,
                         "vector_fields": vector_fields,

app/backend/approaches/retrievethenread.py

@@ -67,6 +67,7 @@ async def run(
         use_text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         use_vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         use_semantic_ranker = True if overrides.get("semantic_ranker") else False
+        use_query_rewriting = True if overrides.get("query_rewriting") else False
         use_semantic_captions = True if overrides.get("semantic_captions") else False
         top = overrides.get("top", 3)
         minimum_search_score = overrides.get("minimum_search_score", 0.0)
@@ -89,6 +90,7 @@ async def run(
             use_semantic_captions,
             minimum_search_score,
             minimum_reranker_score,
+            use_query_rewriting,
         )
 
         # Process results
@@ -118,6 +120,7 @@ async def run(
                     {
                         "use_semantic_captions": use_semantic_captions,
                         "use_semantic_ranker": use_semantic_ranker,
+                        "use_query_rewriting": use_query_rewriting,
                         "top": top,
                         "filter": filter,
                         "use_vector_search": use_vector_search,

app/backend/approaches/retrievethenreadvision.py

@@ -76,6 +76,7 @@ async def run(
         use_text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         use_vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         use_semantic_ranker = True if overrides.get("semantic_ranker") else False
+        use_query_rewriting = True if overrides.get("query_rewriting") else False
         use_semantic_captions = True if overrides.get("semantic_captions") else False
         top = overrides.get("top", 3)
         minimum_search_score = overrides.get("minimum_search_score", 0.0)
@@ -108,6 +109,7 @@ async def run(
             use_semantic_captions,
             minimum_search_score,
             minimum_reranker_score,
+            use_query_rewriting,
         )
 
         # Process results
@@ -145,6 +147,7 @@ async def run(
                     {
                         "use_semantic_captions": use_semantic_captions,
                         "use_semantic_ranker": use_semantic_ranker,
+                        "use_query_rewriting": use_query_rewriting,
                         "top": top,
                         "filter": filter,
                         "vector_fields": vector_fields,

app/backend/config.py

@@ -10,6 +10,7 @@
 CONFIG_AUTH_CLIENT = "auth_client"
 CONFIG_GPT4V_DEPLOYED = "gpt4v_deployed"
 CONFIG_SEMANTIC_RANKER_DEPLOYED = "semantic_ranker_deployed"
+CONFIG_QUERY_REWRITING_ENABLED = "query_rewriting_enabled"
 CONFIG_VECTOR_SEARCH_ENABLED = "vector_search_enabled"
 CONFIG_SEARCH_CLIENT = "search_client"
 CONFIG_OPENAI_CLIENT = "openai_client"

app/backend/requirements.in

@@ -7,7 +7,7 @@ tenacity
 azure-ai-documentintelligence==1.0.0b4
 azure-cognitiveservices-speech
 azure-cosmos
-azure-search-documents==11.6.0b6
+azure-search-documents==11.6.0b9
 azure-storage-blob
 azure-storage-file-datalake
 uvicorn