Move to new azure-search-documents package

Author: pamelafox

app/backend/app.py

@@ -188,7 +188,9 @@ async def ask(auth_claims: dict[str, Any]):
     try:
         approach: Approach = cast(Approach, current_app.config[CONFIG_ASK_APPROACH])
         r = await approach.run(
-            request_json["messages"], context=context, session_state=request_json.get("session_state")
+            request_json["messages"],
+            context=context,
+            session_state=request_json.get("session_state"),
         )
         return jsonify(r)
     except Exception as error:
@@ -329,7 +331,10 @@ async def speech():
             + "#"
             + current_app.config[CONFIG_SPEECH_SERVICE_TOKEN].token
         )
-        speech_config = SpeechConfig(auth_token=auth_token, region=current_app.config[CONFIG_SPEECH_SERVICE_LOCATION])
+        speech_config = SpeechConfig(
+            auth_token=auth_token,
+            region=current_app.config[CONFIG_SPEECH_SERVICE_LOCATION],
+        )
         speech_config.speech_synthesis_voice_name = current_app.config[CONFIG_SPEECH_SERVICE_VOICE]
         speech_config.speech_synthesis_output_format = SpeechSynthesisOutputFormat.Audio16Khz32KBitRateMonoMp3
         synthesizer = SpeechSynthesizer(speech_config=speech_config, audio_config=None)
@@ -339,7 +344,9 @@ async def speech():
         elif result.reason == ResultReason.Canceled:
             cancellation_details = result.cancellation_details
             current_app.logger.error(
-                "Speech synthesis canceled: %s %s", cancellation_details.reason, cancellation_details.error_details
+                "Speech synthesis canceled: %s %s",
+                cancellation_details.reason,
+                cancellation_details.error_details,
             )
             raise Exception("Speech synthesis canceled. Check logs for details.")
         else:
@@ -363,11 +370,22 @@ async def upload(auth_claims: dict[str, Any]):
         adls_manager: AdlsBlobManager = current_app.config[CONFIG_USER_BLOB_MANAGER]
         file_url = await adls_manager.upload_blob(file, file.filename, user_oid)
         ingester: UploadUserFileStrategy = current_app.config[CONFIG_INGESTER]
-        await ingester.add_file(File(content=file, url=file_url, acls={"oids": [user_oid]}), user_oid=user_oid)
+        await ingester.add_file(
+            File(content=file, url=file_url, acls={"oids": [user_oid]}),
+            user_oid=user_oid,
+        )
         return jsonify({"message": "File uploaded successfully"}), 200
     except Exception as error:
         current_app.logger.error("Error uploading file: %s", error)
-        return jsonify({"message": "Error uploading file, check server logs for details.", "status": "failed"}), 500
+        return (
+            jsonify(
+                {
+                    "message": "Error uploading file, check server logs for details.",
+                    "status": "failed",
+                }
+            ),
+            500,
+        )
 
 
 @bp.post("/delete_uploaded")
@@ -471,7 +489,7 @@ async def setup_clients():
     USE_CHAT_HISTORY_BROWSER = os.getenv("USE_CHAT_HISTORY_BROWSER", "").lower() == "true"
     USE_CHAT_HISTORY_COSMOS = os.getenv("USE_CHAT_HISTORY_COSMOS", "").lower() == "true"
     USE_AGENTIC_RETRIEVAL = os.getenv("USE_AGENTIC_RETRIEVAL", "").lower() == "true"
-    ENABLE_AGENTIC_RETRIEVAL_SOURCE_DATA = os.getenv("ENABLE_AGENTIC_RETRIEVAL_SOURCE_DATA", "").lower() == "true"
+    # TODO: ENABLE_AGENTIC_RETRIEVAL_SOURCE_DATA = os.getenv("ENABLE_AGENTIC_RETRIEVAL_SOURCE_DATA", "").lower() == "true"
 
     # WEBSITE_HOSTNAME is always set by App Service, RUNNING_IN_PRODUCTION is set in main.bicep
     RUNNING_ON_AZURE = os.getenv("WEBSITE_HOSTNAME") is not None or os.getenv("RUNNING_IN_PRODUCTION") is not None
@@ -487,15 +505,17 @@ async def setup_clients():
             # ManagedIdentityCredential should use AZURE_CLIENT_ID if set in env, but its not working for some reason,
             # so we explicitly pass it in as the client ID here. This is necessary for user-assigned managed identities.
             current_app.logger.info(
-                "Setting up Azure credential using ManagedIdentityCredential with client_id %s", AZURE_CLIENT_ID
+                "Setting up Azure credential using ManagedIdentityCredential with client_id %s",
+                AZURE_CLIENT_ID,
             )
             azure_credential = ManagedIdentityCredential(client_id=AZURE_CLIENT_ID)
         else:
             current_app.logger.info("Setting up Azure credential using ManagedIdentityCredential")
             azure_credential = ManagedIdentityCredential()
     elif AZURE_TENANT_ID:
         current_app.logger.info(
-            "Setting up Azure credential using AzureDeveloperCliCredential with tenant_id %s", AZURE_TENANT_ID
+            "Setting up Azure credential using AzureDeveloperCliCredential with tenant_id %s",
+            AZURE_TENANT_ID,
         )
         azure_credential = AzureDeveloperCliCredential(tenant_id=AZURE_TENANT_ID, process_timeout=60)
     else:
@@ -515,7 +535,9 @@ async def setup_clients():
         credential=azure_credential,
     )
     agent_client = KnowledgeAgentRetrievalClient(
-        endpoint=AZURE_SEARCH_ENDPOINT, agent_name=AZURE_SEARCH_AGENT, credential=azure_credential
+        endpoint=AZURE_SEARCH_ENDPOINT,
+        agent_name=AZURE_SEARCH_AGENT,
+        credential=azure_credential,
     )
 
     # Set up the global blob storage manager (used for global content/images, but not user uploads)
@@ -600,7 +622,9 @@ async def setup_clients():
             openai_deployment=AZURE_OPENAI_CHATGPT_DEPLOYMENT if OPENAI_HOST == OpenAIHost.AZURE else None,
         )
         search_info = await setup_search_info(
-            search_service=AZURE_SEARCH_SERVICE, index_name=AZURE_SEARCH_INDEX, azure_credential=azure_credential
+            search_service=AZURE_SEARCH_SERVICE,
+            index_name=AZURE_SEARCH_INDEX,
+            azure_credential=azure_credential,
         )
         text_embeddings_service = setup_embeddings_service(
             azure_credential=azure_credential,
@@ -690,7 +714,6 @@ async def setup_clients():
         query_speller=AZURE_SEARCH_QUERY_SPELLER,
         prompt_manager=prompt_manager,
         reasoning_effort=OPENAI_REASONING_EFFORT,
-        hydrate_references=ENABLE_AGENTIC_RETRIEVAL_SOURCE_DATA,
         multimodal_enabled=USE_MULTIMODAL,
         image_embeddings_client=image_embeddings_client,
         global_blob_manager=global_blob_manager,
@@ -718,7 +741,6 @@ async def setup_clients():
         query_speller=AZURE_SEARCH_QUERY_SPELLER,
         prompt_manager=prompt_manager,
         reasoning_effort=OPENAI_REASONING_EFFORT,
-        hydrate_references=ENABLE_AGENTIC_RETRIEVAL_SOURCE_DATA,
         multimodal_enabled=USE_MULTIMODAL,
         image_embeddings_client=image_embeddings_client,
         global_blob_manager=global_blob_manager,

app/backend/approaches/approach.py

@@ -6,13 +6,12 @@
 
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.agent.models import (
-    KnowledgeAgentAzureSearchDocReference,
-    KnowledgeAgentIndexParams,
     KnowledgeAgentMessage,
     KnowledgeAgentMessageTextContent,
     KnowledgeAgentRetrievalRequest,
     KnowledgeAgentRetrievalResponse,
-    KnowledgeAgentSearchActivityRecord,
+    KnowledgeAgentSearchIndexReference,
+    SearchIndexKnowledgeSourceParams,
 )
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.models import (
@@ -162,7 +161,6 @@ def __init__(
         openai_host: str,
         prompt_manager: PromptManager,
         reasoning_effort: Optional[str] = None,
-        hydrate_references: bool = False,
         multimodal_enabled: bool = False,
         image_embeddings_client: Optional[ImageEmbeddings] = None,
         global_blob_manager: Optional[BlobManager] = None,
@@ -180,7 +178,6 @@ def __init__(
         self.openai_host = openai_host
         self.prompt_manager = prompt_manager
         self.reasoning_effort = reasoning_effort
-        self.hydrate_references = hydrate_references
         self.include_token_usage = True
         self.multimodal_enabled = multimodal_enabled
         self.image_embeddings_client = image_embeddings_client
@@ -275,139 +272,58 @@ async def run_agentic_retrieval(
         search_index_name: str,
         top: Optional[int] = None,
         filter_add_on: Optional[str] = None,
-        minimum_reranker_score: Optional[float] = None,
-        max_docs_for_reranker: Optional[int] = None,
         results_merge_strategy: Optional[str] = None,
     ) -> tuple[KnowledgeAgentRetrievalResponse, list[Document]]:
         # STEP 1: Invoke agentic retrieval
         response = await agent_client.retrieve(
             retrieval_request=KnowledgeAgentRetrievalRequest(
                 messages=[
                     KnowledgeAgentMessage(
-                        role=str(msg["role"]), content=[KnowledgeAgentMessageTextContent(text=str(msg["content"]))]
+                        role=str(msg["role"]),
+                        content=[KnowledgeAgentMessageTextContent(text=str(msg["content"]))],
                     )
                     for msg in messages
                     if msg["role"] != "system"
                 ],
-                target_index_params=[
-                    KnowledgeAgentIndexParams(
-                        index_name=search_index_name,
-                        reranker_threshold=minimum_reranker_score,
-                        max_docs_for_reranker=max_docs_for_reranker,
+                knowledge_source_params=[
+                    SearchIndexKnowledgeSourceParams(
+                        knowledge_source_name="default-knowledge-source",
                         filter_add_on=filter_add_on,
-                        include_reference_source_data=True,
                     )
                 ],
             )
         )
 
-        # Map activity id -> agent's internal search query
-        activities = response.activity
-        activity_mapping: dict[int, str] = (
-            {
-                activity.id: activity.query.search
-                for activity in activities
-                if (
-                    isinstance(activity, KnowledgeAgentSearchActivityRecord)
-                    and activity.query
-                    and activity.query.search is not None
-                )
-            }
-            if activities
-            else {}
-        )
-
         # No refs? we're done
         if not (response and response.references):
             return response, []
 
         # Extract references
-        refs = [r for r in response.references if isinstance(r, KnowledgeAgentAzureSearchDocReference)]
+        refs = [r for r in response.references if isinstance(r, KnowledgeAgentSearchIndexReference)]
 
         documents: list[Document] = []
 
-        if self.hydrate_references:
-            # Hydrate references to get full documents
-            documents = await self.hydrate_agent_references(
-                references=refs,
-                top=top,
-            )
-        else:
-            # Create documents from reference source data
-            for ref in refs:
-                if ref.source_data:
-                    documents.append(
-                        Document(
-                            id=ref.doc_key,
-                            content=ref.source_data.get("content"),
-                            sourcepage=ref.source_data.get("sourcepage"),
-                        )
-                    )
-                    if top and len(documents) >= top:
-                        break
-
-        # Build mappings for agent queries and sorting
-        ref_to_activity: dict[str, int] = {}
-        doc_to_ref_id: dict[str, str] = {}
+        # Create documents from reference source data
         for ref in refs:
-            if ref.doc_key:
-                ref_to_activity[ref.doc_key] = ref.activity_source
-                doc_to_ref_id[ref.doc_key] = ref.id
-
-        # Inject agent search queries into all documents
-        for doc in documents:
-            if doc.id and doc.id in ref_to_activity:
-                activity_id = ref_to_activity[doc.id]
-                doc.search_agent_query = activity_mapping.get(activity_id, "")
-
-        # Apply sorting strategy to the documents
-        if results_merge_strategy == "interleaved":  # Use interleaved reference order
-            documents = sorted(
-                documents,
-                key=lambda d: int(doc_to_ref_id.get(d.id, 0)) if d.id and doc_to_ref_id.get(d.id) else 0,
-            )
-        # else: Default - preserve original order
+            if ref.source_data:
+                documents.append(
+                    Document(
+                        id=ref.source_data.get("id"),
+                        content=ref.source_data.get("content"),
+                        category=ref.source_data.get("category"),
+                        sourcepage=ref.source_data.get("sourcepage"),
+                        sourcefile=ref.source_data.get("sourcefile"),
+                        oids=ref.source_data.get("oids"),
+                        groups=ref.source_data.get("groups"),
+                        reranker_score=ref.reranker_score,
+                        images=ref.source_data.get("images"),
+                    )
+                )
+                if top and len(documents) >= top:
+                    break
 
         return response, documents
 
-    async def hydrate_agent_references(
-        self,
-        references: list[KnowledgeAgentAzureSearchDocReference],
-        top: Optional[int],
-    ) -> list[Document]:
-        doc_keys: set[str] = set()
-
-        for ref in references:
-            if not ref.doc_key:
-                continue
-            doc_keys.add(ref.doc_key)
-            if top and len(doc_keys) >= top:
-                break
-
-        if not doc_keys:
-            return []
-
-        # Build search filter only on unique doc IDs
-        id_csv = ",".join(doc_keys)
-        id_filter = f"search.in(id, '{id_csv}', ',')"
-
-        # Fetch full documents
-        hydrated_docs: list[Document] = await self.search(
-            top=len(doc_keys),
-            query_text=None,
-            filter=id_filter,
-            vectors=[],
-            use_text_search=False,
-            use_vector_search=False,
-            use_semantic_ranker=False,
-            use_semantic_captions=False,
-            minimum_search_score=None,
-            minimum_reranker_score=None,
-            use_query_rewriting=False,
-        )
-
-        return hydrated_docs
-
     async def get_sources_content(
         self,
         results: list[Document],
@@ -535,7 +451,11 @@ async def compute_multimodal_embedding(self, q: str):
         if not self.image_embeddings_client:
             raise ValueError("Approach is missing an image embeddings client for multimodal queries")
         multimodal_query_vector = await self.image_embeddings_client.create_embedding_for_text(q)
-        return VectorizedQuery(vector=multimodal_query_vector, k_nearest_neighbors=50, fields="images/embedding")
+        return VectorizedQuery(
+            vector=multimodal_query_vector,
+            k_nearest_neighbors=50,
+            fields="images/embedding",
+        )
 
     def get_system_prompt_variables(self, override_prompt: Optional[str]) -> dict[str, str]:
         # Allows client to replace the entire prompt, or to inject into the existing prompt using >>>

app/backend/approaches/chatreadretrieveread.py

@@ -57,7 +57,6 @@ def __init__(
         query_speller: str,
         prompt_manager: PromptManager,
         reasoning_effort: Optional[str] = None,
-        hydrate_references: bool = False,
         multimodal_enabled: bool = False,
         image_embeddings_client: Optional[ImageEmbeddings] = None,
         global_blob_manager: Optional[BlobManager] = None,
@@ -85,7 +84,6 @@ def __init__(
         self.query_rewrite_tools = self.prompt_manager.load_tools("chat_query_rewrite_tools.json")
         self.answer_prompt = self.prompt_manager.load_prompt("chat_answer_question.prompty")
         self.reasoning_effort = reasoning_effort
-        self.hydrate_references = hydrate_references
         self.include_token_usage = True
         self.multimodal_enabled = multimodal_enabled
         self.image_embeddings_client = image_embeddings_client
@@ -388,13 +386,9 @@ async def run_agentic_retrieval_approach(
         overrides: dict[str, Any],
         auth_claims: dict[str, Any],
     ):
-        minimum_reranker_score = overrides.get("minimum_reranker_score", 0)
         search_index_filter = self.build_filter(overrides, auth_claims)
         top = overrides.get("top", 3)
-        max_subqueries = overrides.get("max_subqueries", 10)
         results_merge_strategy = overrides.get("results_merge_strategy", "interleaved")
-        # 50 is the amount of documents that the reranker can process per query
-        max_docs_for_reranker = max_subqueries * 50
         send_text_sources = overrides.get("send_text_sources", True)
         send_image_sources = overrides.get("send_image_sources", True)
 
@@ -404,8 +398,6 @@ async def run_agentic_retrieval_approach(
             search_index_name=self.search_index_name,
             top=top,
             filter_add_on=search_index_filter,
-            minimum_reranker_score=minimum_reranker_score,
-            max_docs_for_reranker=max_docs_for_reranker,
             results_merge_strategy=results_merge_strategy,
         )
 
@@ -425,8 +417,6 @@ async def run_agentic_retrieval_approach(
                     "Use agentic retrieval",
                     messages,
                     {
-                        "reranker_threshold": minimum_reranker_score,
-                        "max_docs_for_reranker": max_docs_for_reranker,
                         "results_merge_strategy": results_merge_strategy,
                         "filter": search_index_filter,
                     },