feat: context grounding add create index and ingest data

Author: radu-mocanu

sdk/core/uipath_sdk/_models/ecs_index.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+from typing import List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class DataSource(BaseModel):
+    model_config = ConfigDict(
+        validate_by_name=True,
+        validate_by_alias=True,
+        use_enum_values=True,
+        arbitrary_types_allowed=True,
+        extra="allow",
+        json_encoders={datetime: lambda v: v.isoformat() if v else None},
+    )
+    id: Optional[str] = Field(default=None, alias="id")
+    folder: Optional[str] = Field(default=None, alias="folder")
+
+class EcsIndex(BaseModel):
+    model_config = ConfigDict(
+        validate_by_name=True,
+        validate_by_alias=True,
+        use_enum_values=True,
+        arbitrary_types_allowed=True,
+        extra="allow",
+        json_encoders={datetime: lambda v: v.isoformat() if v else None},
+    )
+    id: Optional[str] = Field(default=None, alias="id")
+    name: Optional[str] = Field(default=None, alias="name")
+    description: Optional[str] = Field(default=None, alias="description")
+    memory_usage: Optional[int] = Field(default=None, alias="memoryUsage")
+    disk_usage: Optional[int] = Field(default=None, alias="diskUsage")
+    data_source: Optional[DataSource] = Field(default=None, alias="dataSource")
+    pre_processing: Optional[dict] = Field(default=None, alias="preProcessing")
+    fields: Optional[List[Field]] = Field(default=None, alias="fields")
+    last_ingestion_status: Optional[str] = Field(default=None, alias="lastIngestionStatus")
+    last_ingested: Optional[datetime] = Field(default=None, alias="lastIngested")
+    last_queried: Optional[datetime] = Field(default=None, alias="lastQueried")
+    folder_key: Optional[str] = Field(default=None, alias="folderKey")

sdk/core/uipath_sdk/_services/__init__.py

@@ -8,6 +8,7 @@
 from .llm_gateway_service import UiPathLlmChatService, UiPathOpenAIService
 from .processes_service import ProcessesService
 from .queues_service import QueuesService
+from .folder_service import FolderService
 
 __all__ = [
     "ActionsService",
@@ -21,4 +22,5 @@
     "JobsService",
     "UiPathOpenAIService",
     "UiPathLlmChatService",
+    "FolderService",
 ]

sdk/core/uipath_sdk/_services/context_grounding_service.py

@@ -1,16 +1,18 @@
 import json
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
+from urllib.request import Request
 
 from pydantic import TypeAdapter
 
 from .._config import Config
 from .._execution_context import ExecutionContext
 from .._folder_context import FolderContext
 from .._models.context_grounding import ContextGroundingQueryResponse
+from .._models.ecs_index import EcsIndex
 from .._utils import Endpoint, RequestSpec
+from .._utils.constants import ORCHESTRATOR_STORAGE_BUCKET_DATA_SOURCE, ENV_FOLDER_PATH, HEADER_FOLDER_KEY
 from ._base_service import BaseService
-
-
+from .folder_service import FolderService
 class ContextGroundingService(FolderContext, BaseService):
     """Service for managing semantic automation contexts in UiPath.
 
@@ -24,10 +26,11 @@ class ContextGroundingService(FolderContext, BaseService):
     context.
     """
 
-    def __init__(self, config: Config, execution_context: ExecutionContext) -> None:
+    def __init__(self, config: Config, execution_context: ExecutionContext, folders_service: FolderService) -> None:
+        self._folders_service = folders_service
         super().__init__(config=config, execution_context=execution_context)
 
-    def retrieve(self, name: str) -> Any:
+    def retrieve(self, name: str) -> Optional[EcsIndex]:
         """Retrieve context grounding index information by its name.
 
         This method fetches details about a specific context index, which can be
@@ -38,17 +41,18 @@ def retrieve(self, name: str) -> Any:
             name (str): The name of the context index to retrieve.
 
         Returns:
-            Any: The index information, including its configuration and metadata.
+            Optional[EcsIndex]: The index information, including its configuration and metadata if found, otherwise None.
         """
         spec = self._retrieve_spec(name)
 
-        return self.request(
+        response = self.request(
             spec.method,
             spec.endpoint,
             params=spec.params,
         ).json()
+        return next((EcsIndex.model_validate(item) for item in response["value"] if item["name"] == name), None)
 
-    async def retrieve_async(self, name: str) -> Any:
+    async def retrieve_async(self, name: str) -> Optional[EcsIndex]:
         """Retrieve asynchronously context grounding index information by its name.
 
         This method fetches details about a specific context index, which can be
@@ -59,18 +63,17 @@ async def retrieve_async(self, name: str) -> Any:
             name (str): The name of the context index to retrieve.
 
         Returns:
-            Any: The index information, including its configuration and metadata.
+            Optional[EcsIndex]: The index information, including its configuration and metadata if found, otherwise None.
 
         """
         spec = self._retrieve_spec(name)
 
-        response = await self.request_async(
+        response = (await self.request_async(
             spec.method,
             spec.endpoint,
             params=spec.params,
-        )
-
-        return response.json()
+        )).json()
+        return next((EcsIndex.model_validate(item) for item in response["value"] if item["name"] == name), None)
 
     def retrieve_by_id(self, id: str) -> Any:
         """Retrieve context grounding index information by its ID.
@@ -183,19 +186,117 @@ async def search_async(
             response.json()
         )
 
+    def get_or_create_index(self,
+                            name: str,
+                            *,
+                            description: Optional[str],
+                            storage_bucket_name: Optional[str],
+                            file_name_glob: Optional[str],
+                            storage_bucket_folder_path: Optional[str],
+                            ) -> EcsIndex:
+        spec = self._create_or_get_spec(name, description, storage_bucket_name, file_name_glob, storage_bucket_folder_path)
+        index = self.retrieve(name=name)
+        if index:
+            return index
+
+        response = self.request(
+            spec.method,
+            spec.endpoint,
+            content=spec.content,
+            headers=spec.headers,
+        ).json()
+        return EcsIndex.model_validate(response)
+
+    async def get_or_create_index_async(self,
+                                        name: str,
+                                        *,
+                                        description: Optional[str],
+                                        storage_bucket_name: Optional[str],
+                                        file_name_glob: Optional[str],
+                                        storage_bucket_folder_path: Optional[str],
+                                        ) -> EcsIndex:
+        spec = self._create_or_get_spec(name, description, storage_bucket_name, file_name_glob, storage_bucket_folder_path)
+        index = await self.retrieve_async(name=name)
+        if index:
+            return index
+
+        response = (await self.request_async(
+            spec.method,
+            spec.endpoint,
+            content=spec.content,
+            headers=spec.headers,
+        )).json()
+        return EcsIndex.model_validate(response)
+
+    def ingest_data(self, index: EcsIndex) -> None:
+        spec = self._ingest_spec(index.id)
+        self.request(
+            spec.method,
+            spec.endpoint,
+            headers=spec.headers,
+        )
+
+    async def ingest_data(self, index: EcsIndex) -> None:
+        spec = self._ingest_spec(index.id)
+        await self.request_async(
+            spec.method,
+            spec.endpoint,
+            headers=spec.headers,
+        )
+
     @property
     def custom_headers(self) -> Dict[str, str]:
         if self.folder_headers["x-uipath-folderkey"] is None:
             raise ValueError("Folder key is not set (UIPATH_FOLDER_KEY)")
 
         return self.folder_headers
 
+    def _ingest_spec(self, key: str) -> RequestSpec:
+        if not self._folder_path:
+            raise ValueError(f"Folder path is not set ({ENV_FOLDER_PATH})")
+        folder_key = self._folder_key if self._folder_key else self._folders_service.retrieve_key_by_folder_path(self._folder_path)
+
+        return RequestSpec(
+            method="POST",
+            endpoint=Endpoint(f"/ecs_/v2/indexes/{key}/ingest"),
+            headers={HEADER_FOLDER_KEY: folder_key}
+        )
+
     def _retrieve_spec(self, name: str) -> RequestSpec:
         return RequestSpec(
             method="GET",
             endpoint=Endpoint("/ecs_/v2/indexes"),
             params={"$filter": f"Name eq '{name}'"},
         )
+    def _create_or_get_spec(self,
+                            name: str,
+                            description: Optional[str],
+                            storage_bucket_name: Optional[str],
+                            file_name_glob: Optional[str],
+                            storage_bucket_folder_path: Optional[str],
+                            ) -> RequestSpec:
+        if not self._folder_path:
+            raise ValueError(f"Folder path is not set ({ENV_FOLDER_PATH})")
+        storage_bucket_folder_path = storage_bucket_folder_path if storage_bucket_folder_path else self._folder_path
+        folder_key = self._folder_key if self._folder_key else self._folders_service.retrieve_key_by_folder_path(self._folder_path)
+        return RequestSpec(
+            method="POST",
+            endpoint=Endpoint("/ecs_/v2/indexes/create"),
+            content=json.dumps(
+                {
+                    "name": name,
+                    "description": description,
+                    "dataSource":
+                        {
+                            "@odata.type": ORCHESTRATOR_STORAGE_BUCKET_DATA_SOURCE,
+                            "folder": storage_bucket_folder_path,
+                            "bucketName": storage_bucket_name,
+                            "fileNameGlob": file_name_glob if file_name_glob is not None else "*",
+                        }
+                }
+            ),
+            headers={HEADER_FOLDER_KEY: folder_key},
+        )
 
     def _retrieve_by_id_spec(self, id: str) -> RequestSpec:
         return RequestSpec(

sdk/core/uipath_sdk/_services/folder_service.py

@@ -0,0 +1,36 @@
+from typing import Optional
+
+from .._config import Config
+from .._execution_context import ExecutionContext
+from .._utils import Endpoint, RequestSpec, header_folder
+from ._base_service import BaseService
+
+
+def _retrieve_spec(folder_path: str) -> RequestSpec:
+    folder_name = folder_path.split("/")[-1]
+    return RequestSpec(
+        method="GET",
+        endpoint=Endpoint(f"orchestrator_/api/FoldersNavigation/GetFoldersForCurrentUser?searchText={folder_name}")
+    )
+
+
+class FolderService(BaseService):
+    """Service for managing UiPath Folders.
+
+    A folder represents a single area for data organization
+    and access control - it is created when you need to categorize, manage, and enforce authorization rules for a group
+    of UiPath resources (i.e. processes, assets, connections, storage buckets etc.) or other folders
+    """
+
+    def __init__(self, config: Config, execution_context: ExecutionContext) -> None:
+        super().__init__(config=config, execution_context=execution_context)
+
+
+    def retrieve_key_by_folder_path(self, folder_path: str) -> Optional[str]:
+        spec = _retrieve_spec(folder_path)
+        response = self.request(
+            spec.method,
+            url=spec.endpoint
+        ).json()
+
+        return next((item['Key'] for item in response['PageItems'] if item['FullyQualifiedName'] == folder_path), None)

sdk/core/uipath_sdk/_uipath_sdk.py

@@ -15,6 +15,7 @@
     JobsService,
     ProcessesService,
     QueuesService,
+    FolderService
 )
 from ._utils import setup_logging
 from ._utils.constants import (
@@ -45,6 +46,7 @@ def __init__(
             base_url=base_url_value,  # type: ignore
             secret=secret_value,  # type: ignore
         )
+        self._folders_service = None
 
         setup_logging(debug)
         self._execution_context = ExecutionContext()
@@ -75,7 +77,9 @@ def connections(self) -> ConnectionsService:
 
     @property
     def context_grounding(self) -> ContextGroundingService:
-        return ContextGroundingService(self._config, self._execution_context)
+        if not self._folders_service:
+            self._folders_service = FolderService(self._config, self._execution_context)
+        return ContextGroundingService(self._config, self._execution_context, self._folders_service)
 
     @property
     def queues(self) -> QueuesService:
@@ -84,3 +88,9 @@ def queues(self) -> QueuesService:
     @property
     def jobs(self) -> JobsService:
         return JobsService(self._config, self._execution_context)
+
+    @property
+    def folders(self) -> FolderService:
+        if not self._folders_service:
+            self._folders_service = FolderService(self._config, self._execution_context)
+        return self._folders_service

sdk/core/uipath_sdk/_utils/constants.py

@@ -18,3 +18,6 @@
 
 # Entrypoint for plugins
 ENTRYPOINT = "uipath_sdk.connectors"
+
+# Data sources
+ORCHESTRATOR_STORAGE_BUCKET_DATA_SOURCE="#UiPath.Vdbs.Domain.Api.V20Models.StorageBucketDataSourceRequest"