diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index b7152da61..c0a4eab15 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -1,3 +1,4 @@ +import uuid from pathlib import Path from typing import Annotated, Any, Dict, List, Optional, Tuple, Union @@ -33,11 +34,13 @@ ) from .context_grounding_index import ContextGroundingIndex from .context_grounding_payloads import ( + AttachmentsDataSource, BucketDataSource, BucketSourceConfig, ConfluenceDataSource, ConfluenceSourceConfig, CreateIndexPayload, + CreateJitIndexPayload, DropboxDataSource, DropboxSourceConfig, GoogleDriveDataSource, @@ -398,6 +401,32 @@ def create_index( return ContextGroundingIndex.model_validate(response.json()) + @resource_override(resource_type="index") + @traced(name="contextgrounding_create_jit_index", run_type="uipath") + def create_jit_index( + self, + usage: str, + attachments: list[uuid.UUID], + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> ContextGroundingIndex: + """Create a new context jit grounding index.""" + spec = self._create_jit_spec( + usage, + attachments, + folder_path=folder_path, + folder_key=folder_key, + ) + + response = self.request( + spec.method, + spec.endpoint, + json=spec.json, + headers=spec.headers, + ) + + return ContextGroundingIndex.model_validate(response.json()) + @resource_override(resource_type="index") @traced(name="contextgrounding_create_index", run_type="uipath") async def create_index_async( @@ -1197,6 +1226,41 @@ def _create_spec( }, ) + def _create_jit_spec( + self, + usage: str, + attachments: list[uuid.UUID] = None, + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, + ) -> RequestSpec: + """Create request spec for index creation.""" + folder_key = self._resolve_folder_key(folder_key, folder_path) + + data_source_dict = self._build_jit_data_source(attachments) + + payload = CreateJitIndexPayload( + usage=usage, + data_source=data_source_dict, + ) + + return RequestSpec( + method="POST", + endpoint=Endpoint("/ecs_/v2/indexes/createephemeral"), + json=payload.model_dump(by_alias=True, exclude_none=True), + headers={ + **header_folder(folder_key, None), + }, + ) + + def _build_jit_data_source(self, attachments: list[uuid.UUID]) -> Dict[str, Any]: + data_source: AttachmentsDataSource + data_source = AttachmentsDataSource(attachments=attachments) + return data_source.model_dump( + by_alias=True, + exclude_none=True, + mode="json", + ) + def _build_data_source(self, source: SourceConfig) -> Dict[str, Any]: """Build data source configuration from typed source config. diff --git a/src/uipath/platform/context_grounding/context_grounding_payloads.py b/src/uipath/platform/context_grounding/context_grounding_payloads.py index 4b64778d7..1912610ab 100644 --- a/src/uipath/platform/context_grounding/context_grounding_payloads.py +++ b/src/uipath/platform/context_grounding/context_grounding_payloads.py @@ -1,6 +1,7 @@ """Payload models for context grounding index creation and configuration.""" import re +import uuid from typing import Any, Dict, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -82,6 +83,10 @@ class ConfluenceDataSource(DataSourceBase): space_id: str = Field(alias="spaceId", description="Space ID") +class AttachmentsDataSource(BaseModel): + attachments: list[uuid.UUID] = Field(description="List of attachment ids") + + class Indexer(BaseModel): """Configuration for periodic indexing of data sources.""" @@ -136,6 +141,17 @@ class CreateIndexPayload(BaseModel): model_config = ConfigDict(populate_by_name=True) +class CreateJitIndexPayload(BaseModel): + """ """ + + usage: str = Field(description="Index usage") + data_source: Dict[str, Any] = Field( + alias="dataSource", description="Data source configuration" + ) + + model_config = ConfigDict(populate_by_name=True) + + # user-facing source configuration models class BaseSourceConfig(BaseModel): """Base configuration for all source types."""