azureaisearch: add semantic search mode support for async queries

llama-index-integrations/vector_stores/llama-index-vector-stores-azureaisearch/llama_index/vector_stores/azureaisearch/base.py

@@ -12,6 +12,7 @@
 from azure.search.documents.indexes.aio import (
     SearchIndexClient as AsyncSearchIndexClient,
 )
+
 from llama_index.core.bridge.pydantic import PrivateAttr
 from llama_index.core.schema import BaseNode, MetadataMode, TextNode
 from llama_index.core.vector_stores.types import (
@@ -30,9 +31,9 @@
 )
 from llama_index.vector_stores.azureaisearch.azureaisearch_utils import (
     create_node_from_result,
-    process_batch_results,
     create_search_request,
     handle_search_error,
+    process_batch_results,
 )
 
 logger = logging.getLogger(__name__)
@@ -398,10 +399,10 @@ async def _acreate_index(self, index_name: Optional[str]) -> None:
             ExhaustiveKnnParameters,
             HnswAlgorithmConfiguration,
             HnswParameters,
+            SearchableField,
             SearchField,
             SearchFieldDataType,
             SearchIndex,
-            SearchableField,
             SemanticConfiguration,
             SemanticField,
             SemanticPrioritizedFields,
@@ -1576,11 +1577,11 @@ def _create_query_vector(self) -> Optional[List[Any]]:
         return vector_queries
 
     def _create_query_result(
-        self, search_query: str, vector_queries: Optional[List[Any]]
+        self, search_query: str, vectors: Optional[List[Any]]
     ) -> VectorStoreQueryResult:
         results = self._search_client.search(
             search_text=search_query,
-            vector_queries=vector_queries,
+            vector_queries=vectors,
             top=self._query.similarity_top_k,
             select=self._select_fields,
             filter=self._odata_filter,
@@ -1631,5 +1632,61 @@ def _create_query_result(
             nodes=node_result, similarities=score_result, ids=id_result
         )
 
+    async def _acreate_query_result(
+        self, search_query: str, vectors: Optional[List[Any]]
+    ) -> VectorStoreQueryResult:
+        results = await self._search_client.search(
+            search_text=search_query,
+            vector_queries=vectors,
+            top=self._query.similarity_top_k,
+            select=self._select_fields,
+            filter=self._odata_filter,
+            query_type="semantic",
+            semantic_configuration_name="mySemanticConfig",
+        )
+
+        id_result = []
+        node_result = []
+        score_result = []
+        async for result in results:
+            node_id = result[self._field_mapping["id"]]
+            metadata_str = result[self._field_mapping["metadata"]]
+            metadata = json.loads(metadata_str) if metadata_str else {}
+            # use reranker_score instead of score
+            score = result["@search.reranker_score"]
+            chunk = result[self._field_mapping["chunk"]]
+
+            try:
+                node = metadata_dict_to_node(metadata)
+                node.set_content(chunk)
+            except Exception:
+                # NOTE: deprecated legacy logic for backward compatibility
+                metadata, node_info, relationships = legacy_metadata_dict_to_node(
+                    metadata
+                )
+
+                node = TextNode(
+                    text=chunk,
+                    id_=node_id,
+                    metadata=metadata,
+                    start_char_idx=node_info.get("start", None),
+                    end_char_idx=node_info.get("end", None),
+                    relationships=relationships,
+                )
+
+            logger.debug(f"Retrieved node id {node_id} with node data of {node}")
+
+            id_result.append(node_id)
+            node_result.append(node)
+            score_result.append(score)
+
+        logger.debug(
+            f"Search query '{search_query}' returned {len(id_result)} results."
+        )
+
+        return VectorStoreQueryResult(
+            nodes=node_result, similarities=score_result, ids=id_result
+        )
+
 
 CognitiveSearchVectorStore = AzureAISearchVectorStore