feat: context grounding create index and ingest data

Author: radu-mocanu

src/uipath/_models/__init__.py

@@ -3,6 +3,7 @@
 from .assets import UserAsset
 from .connections import Connection, ConnectionToken
 from .context_grounding import ContextGroundingQueryResponse
+from .exceptions import IngestionInProgressException
 from .interrupt_models import CreateAction, InvokeProcess, WaitAction, WaitJob
 from .job import Job
 from .processes import Process
@@ -32,4 +33,5 @@
     "WaitJob",
     "WaitAction",
     "CreateAction",
+    "IngestionInProgressException",
 ]

src/uipath/_models/context_grounding_index.py

@@ -0,0 +1,58 @@
+from datetime import datetime
+from typing import Any, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ContextGroundingField(BaseModel):
+    id: Optional[str] = Field(default=None, alias="id")
+    name: Optional[str] = Field(default=None, alias="name")
+    description: Optional[str] = Field(default=None, alias="description")
+    type: Optional[str] = Field(default=None, alias="type")
+    is_filterable: Optional[bool] = Field(default=None, alias="isFilterable")
+    searchable_type: Optional[str] = Field(default=None, alias="searchableType")
+    is_user_defined: Optional[bool] = Field(default=None, alias="isUserDefined")
+
+
+class ContextGroundingDataSource(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 ContextGroundingIndex(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[ContextGroundingDataSource] = Field(default=None, alias="dataSource")
+    pre_processing: Any = Field(default=None, alias="preProcessing")
+    fields: Optional[List[ContextGroundingField]] = 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")
+
+    def in_progress_ingestion(self):
+        return (
+            self.last_ingestion_status == "Queued"
+            or self.last_ingestion_status == "In Progress"
+        )

src/uipath/_models/exceptions.py

@@ -0,0 +1,6 @@
+class IngestionInProgressException(Exception):
+    """An exception that is triggered when a search is attempted on an index that is currently undergoing ingestion."""
+
+    def __init__(self, index_name):
+        self.message = f"index {index_name} cannot be searched during ingestion"
+        super().__init__(self.message)

src/uipath/_services/__init__.py

@@ -4,6 +4,7 @@
 from .buckets_service import BucketsService
 from .connections_service import ConnectionsService
 from .context_grounding_service import ContextGroundingService
+from .folder_service import FolderService
 from .jobs_service import JobsService
 from .llm_gateway_service import UiPathLlmChatService, UiPathOpenAIService
 from .processes_service import ProcessesService
@@ -21,4 +22,5 @@
     "JobsService",
     "UiPathOpenAIService",
     "UiPathLlmChatService",
+    "FolderService",
 ]

src/uipath/_services/buckets_service.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Optional, Union
 
 from httpx import request
 
@@ -59,7 +59,9 @@ def download(
 
     def upload(
         self,
-        bucket_key: str,
+        *,
+        bucket_key: Optional[str] = None,
+        bucket_name: Optional[str] = None,
         blob_file_path: str,
         content_type: str,
         source_path: str,
@@ -68,11 +70,18 @@ def upload(
 
         Args:
             bucket_key: The key of the bucket
+            bucket_name: The name of the bucket
             blob_file_path: The path where the file will be stored in the bucket
             content_type: The MIME type of the file
             source_path: The local path of the file to upload
         """
-        bucket = self.retrieve_by_key(bucket_key)
+        if bucket_key:
+            bucket = self.retrieve_by_key(bucket_key)
+        elif bucket_name:
+            bucket = self.retrieve(bucket_name)
+        else:
+            raise ValueError("Must specify a bucket name or bucket key")
+
         bucket_id = bucket["Id"]
 
         endpoint = Endpoint(
@@ -99,6 +108,60 @@ def upload(
             else:
                 request("PUT", write_uri, headers=headers, files={"file": file})
 
+    def upload_from_memory(
+        self,
+        *,
+        bucket_key: Optional[str] = None,
+        bucket_name: Optional[str] = None,
+        blob_file_path: str,
+        content_type: str,
+        content: Union[str, bytes],
+    ) -> None:
+        """Upload content from memory to a bucket.
+
+        Args:
+            bucket_key: The key of the bucket
+            bucket_name: The name of the bucket
+            blob_file_path: The path where the content will be stored in the bucket
+            content_type: The MIME type of the content
+            content: The content to upload (string or bytes)
+        """
+        if bucket_key:
+            bucket = self.retrieve_by_key(bucket_key)
+        elif bucket_name:
+            bucket = self.retrieve(bucket_name)
+        else:
+            raise ValueError("Must specify a bucket name or bucket key")
+
+        bucket_id = bucket["Id"]
+
+        endpoint = Endpoint(
+            f"/orchestrator_/odata/Buckets({bucket_id})/UiPath.Server.Configuration.OData.GetWriteUri"
+        )
+
+        result = self.request(
+            "GET",
+            endpoint,
+            params={"path": blob_file_path, "contentType": content_type},
+        ).json()
+        write_uri = result["Uri"]
+
+        headers = {
+            key: value
+            for key, value in zip(
+                result["Headers"]["Keys"], result["Headers"]["Values"]
+            )
+        }
+
+        # Convert string to bytes if needed
+        if isinstance(content, str):
+            content = content.encode("utf-8")
+
+        if result["RequiresAuth"]:
+            self.request("PUT", write_uri, headers=headers, content=content)
+        else:
+            request("PUT", write_uri, headers=headers, content=content)
+
     @infer_bindings()
     def retrieve(self, name: str) -> Any:
         """Retrieve bucket information by its name.
@@ -192,14 +255,14 @@ def custom_headers(self) -> Dict[str, str]:
     def _retrieve_spec(self, name: str) -> RequestSpec:
         return RequestSpec(
             method="GET",
-            endpoint=Endpoint("/odata/Buckets"),
+            endpoint=Endpoint("/orchestrator_/odata/Buckets"),
             params={"$filter": f"Name eq '{name}'", "$top": 1},
         )
 
     def _retrieve_by_key_spec(self, key: str) -> RequestSpec:
         return RequestSpec(
             method="GET",
             endpoint=Endpoint(
-                f"/odata/Buckets/UiPath.Server.Configuration.OData.GetByKey(identifier={key})"
+                f"/orchestrator_/odata/Buckets/UiPath.Server.Configuration.OData.GetByKey(identifier={key})"
             ),
         )