feat: implement UnifiedSearch (search v1.2) method [ECS-1656] (#1427)

Author: gcuip

packages/uipath-platform/src/uipath/platform/context_grounding/__init__.py

@@ -9,12 +9,18 @@
     Citation,
     CitationMode,
     ContextGroundingQueryResponse,
+    ContextGroundingSearchResultItem,
     DeepRagContent,
     DeepRagCreationResponse,
     DeepRagResponse,
     DeepRagStatus,
     EphemeralIndexUsage,
     IndexStatus,
+    SearchMode,
+    SemanticSearchOptions,
+    SemanticSearchResult,
+    UnifiedQueryResult,
+    UnifiedSearchScope,
 )
 from .context_grounding_index import ContextGroundingIndex
 from .context_grounding_payloads import (
@@ -49,6 +55,7 @@
     "ConnectionSourceConfig",
     "ContextGroundingIndex",
     "ContextGroundingQueryResponse",
+    "ContextGroundingSearchResultItem",
     "ContextGroundingService",
     "CreateIndexPayload",
     "DeepRagCreationResponse",
@@ -64,6 +71,11 @@
     "Indexer",
     "OneDriveDataSource",
     "OneDriveSourceConfig",
+    "SearchMode",
+    "SemanticSearchOptions",
+    "SemanticSearchResult",
     "SourceConfig",
+    "UnifiedQueryResult",
+    "UnifiedSearchScope",
     "Citation",
 ]

packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py

@@ -3,6 +3,7 @@
 
 import httpx
 from pydantic import Field, TypeAdapter
+from typing_extensions import deprecated
 from uipath.core.tracing import traced
 
 from ..common._base_service import BaseService
@@ -33,6 +34,9 @@
     DeepRagCreationResponse,
     DeepRagResponse,
     EphemeralIndexUsage,
+    SearchMode,
+    UnifiedQueryResult,
+    UnifiedSearchScope,
 )
 from .context_grounding_index import ContextGroundingIndex
 from .context_grounding_payloads import (
@@ -1274,16 +1278,20 @@ async def start_deep_rag_ephemeral_async(
 
     @resource_override(resource_type="index")
     @traced(name="contextgrounding_search", run_type="uipath")
+    @deprecated("Use unified_search instead.")
     def search(
         self,
         name: str,
         query: str,
         number_of_results: int = 10,
+        threshold: Optional[float] = None,
         folder_key: Optional[str] = None,
         folder_path: Optional[str] = None,
     ) -> List[ContextGroundingQueryResponse]:
         """Search for contextual information within a specific index.
 
+        This method is deprecated. Use unified_search instead.
+
         This method performs a semantic search against the specified context index,
         helping to find relevant information that can be used in automation processes.
         The search is powered by AI and understands natural language queries.
@@ -1293,6 +1301,7 @@ def search(
             query (str): The search query in natural language.
             number_of_results (int, optional): Maximum number of results to return.
                 Defaults to 10.
+            threshold (float): Minimum similarity threshold. Defaults to 0.0.
 
         Returns:
             List[ContextGroundingQueryResponse]: A list of search results, each containing
@@ -1308,6 +1317,7 @@ def search(
             name,
             query,
             number_of_results,
+            threshold=threshold if threshold is not None else 0.0,
             folder_key=folder_key,
             folder_path=folder_path,
         )
@@ -1325,16 +1335,20 @@ def search(
 
     @resource_override(resource_type="index")
     @traced(name="contextgrounding_search", run_type="uipath")
+    @deprecated("Use unified_search_async instead.")
     async def search_async(
         self,
         name: str,
         query: str,
         number_of_results: int = 10,
+        threshold: Optional[float] = None,
         folder_key: Optional[str] = None,
         folder_path: Optional[str] = None,
     ) -> List[ContextGroundingQueryResponse]:
         """Search asynchronously for contextual information within a specific index.
 
+        This method is deprecated. Use unified_search_async instead.
+
         This method performs a semantic search against the specified context index,
         helping to find relevant information that can be used in automation processes.
         The search is powered by AI and understands natural language queries.
@@ -1344,6 +1358,7 @@ async def search_async(
             query (str): The search query in natural language.
             number_of_results (int, optional): Maximum number of results to return.
                 Defaults to 10.
+            threshold (float): Minimum similarity threshold. Defaults to 0.0.
 
         Returns:
             List[ContextGroundingQueryResponse]: A list of search results, each containing
@@ -1363,6 +1378,7 @@ async def search_async(
             name,
             query,
             number_of_results,
+            threshold=threshold if threshold is not None else 0.0,
             folder_key=folder_key,
             folder_path=folder_path,
         )
@@ -1378,6 +1394,120 @@ async def search_async(
             response.json()
         )
 
+    @resource_override(resource_type="index")
+    @traced(name="contextgrounding_unified_search", run_type="uipath")
+    def unified_search(
+        self,
+        name: str,
+        query: str,
+        search_mode: SearchMode = SearchMode.AUTO,
+        number_of_results: int = 10,
+        threshold: float = 0.0,
+        scope: Optional[UnifiedSearchScope] = None,
+        folder_key: Optional[str] = None,
+        folder_path: Optional[str] = None,
+    ) -> UnifiedQueryResult:
+        """Perform a unified search on a context grounding index.
+
+        This method performs a unified search (v1.2) against the specified context index,
+        supporting both semantic and tabular search modes.
+
+        Args:
+            name (str): The name of the context index to search in.
+            query (str): The search query in natural language.
+            search_mode (SearchMode): The search mode to use. Defaults to AUTO.
+            number_of_results (int): Maximum number of results to return. Defaults to 10.
+            threshold (float): Minimum similarity threshold. Defaults to 0.0.
+            scope (Optional[UnifiedSearchScope]): Optional search scope (folder, extension).
+            folder_key (Optional[str]): The key of the folder where the index resides.
+            folder_path (Optional[str]): The path of the folder where the index resides.
+
+        Returns:
+            UnifiedQueryResult: The unified search result containing semantic and/or tabular results.
+        """
+        index = self.retrieve(name, folder_key=folder_key, folder_path=folder_path)
+
+        folder_key = folder_key or index.folder_key
+
+        spec = self._unified_search_spec(
+            index_id=index.id,
+            query=query,
+            search_mode=search_mode,
+            number_of_results=number_of_results,
+            threshold=threshold,
+            scope=scope,
+            folder_key=folder_key,
+            folder_path=folder_path,
+        )
+
+        response = self.request(
+            spec.method,
+            spec.endpoint,
+            json=spec.json,
+            headers=spec.headers,
+        )
+
+        return UnifiedQueryResult.model_validate(response.json())
+
+    @resource_override(resource_type="index")
+    @traced(name="contextgrounding_unified_search", run_type="uipath")
+    async def unified_search_async(
+        self,
+        name: str,
+        query: str,
+        search_mode: SearchMode = SearchMode.AUTO,
+        number_of_results: int = 10,
+        threshold: float = 0.0,
+        scope: Optional[UnifiedSearchScope] = None,
+        folder_key: Optional[str] = None,
+        folder_path: Optional[str] = None,
+    ) -> UnifiedQueryResult:
+        """Asynchronously perform a unified search on a context grounding index.
+
+        This method performs a unified search (v1.2) against the specified context index,
+        supporting both semantic and tabular search modes.
+
+        Args:
+            name (str): The name of the context index to search in.
+            query (str): The search query in natural language.
+            search_mode (SearchMode): The search mode to use. Defaults to AUTO.
+            number_of_results (int): Maximum number of results to return. Defaults to 10.
+            threshold (float): Minimum similarity threshold. Defaults to 0.0.
+            scope (Optional[UnifiedSearchScope]): Optional search scope (folder, extension).
+            folder_key (Optional[str]): The key of the folder where the index resides.
+            folder_path (Optional[str]): The path of the folder where the index resides.
+
+        Returns:
+            UnifiedQueryResult: The unified search result containing semantic and/or tabular results.
+        """
+        index = await self.retrieve_async(
+            name, folder_key=folder_key, folder_path=folder_path
+        )
+        if index and index.in_progress_ingestion():
+            raise IngestionInProgressException(index_name=name)
+
+        folder_key = folder_key or index.folder_key
+
+        spec = self._unified_search_spec(
+            index_id=index.id,
+            query=query,
+            search_mode=search_mode,
+            number_of_results=number_of_results,
+            threshold=threshold,
+            scope=scope,
+            folder_key=folder_key,
+            folder_path=folder_path,
+        )
+
+        response = await self.request_async(
+            spec.method,
+            spec.endpoint,
+            json=spec.json,
+            headers=spec.headers,
+        )
+
+        return UnifiedQueryResult.model_validate(response.json())
+
     @traced(name="contextgrounding_ingest_data", run_type="uipath")
     def ingest_data(
         self,
@@ -1757,6 +1887,7 @@ def _search_spec(
         name: str,
         query: str,
         number_of_results: int = 10,
+        threshold: float = 0.0,
         folder_key: Optional[str] = None,
         folder_path: Optional[str] = None,
     ) -> RequestSpec:
@@ -1766,14 +1897,54 @@ def _search_spec(
             method="POST",
             endpoint=Endpoint("/ecs_/v1/search"),
             json={
-                "query": {"query": query, "numberOfResults": number_of_results},
+                "query": {
+                    "query": query,
+                    "numberOfResults": number_of_results,
+                    "threshold": threshold,
+                },
                 "schema": {"name": name},
             },
             headers={
                 **header_folder(folder_key, None),
             },
         )
 
+    def _unified_search_spec(
+        self,
+        index_id: str,
+        query: str,
+        search_mode: SearchMode = SearchMode.AUTO,
+        number_of_results: int = 10,
+        threshold: float = 0.0,
+        scope: Optional[UnifiedSearchScope] = None,
+        folder_key: Optional[str] = None,
+        folder_path: Optional[str] = None,
+    ) -> RequestSpec:
+        folder_key = self._resolve_folder_key(folder_key, folder_path)
+
+        json_body: Dict[str, Any] = {
+            "searchMode": search_mode.value
+            if isinstance(search_mode, SearchMode)
+            else search_mode,
+            "query": query,
+            "semanticSearchOptions": {
+                "numberOfResults": number_of_results,
+                "threshold": threshold,
+            },
+        }
+
+        if scope is not None:
+            json_body["scope"] = scope.model_dump(by_alias=True, exclude_none=True)
+
+        return RequestSpec(
+            method="POST",
+            endpoint=Endpoint(f"/ecs_/v1.2/search/{index_id}"),
+            json=json_body,
+            headers={
+                **header_folder(folder_key, None),
+            },
+        )
+
     def _deep_rag_creation_spec(
         self,
         index_id: str,