From e12f8942880654847d4c230dd7937c0bc2b94a0d Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Tue, 13 Jan 2026 16:42:41 -0800 Subject: [PATCH 1/8] feat: add ecs jit sdk --- .../_context_grounding_service.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 31cadcc0a..4655538d7 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 @@ -398,6 +399,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( From 19856559a534c703ea2a2e90cc263b22b5155b32 Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Thu, 15 Jan 2026 17:13:49 -0800 Subject: [PATCH 2/8] add create ephemeral index --- .../_context_grounding_service.py | 32 ++++--------------- .../platform/resume_triggers/_protocol.py | 23 +++++++++++++ 2 files changed, 30 insertions(+), 25 deletions(-) diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 4655538d7..470682a08 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -292,17 +292,11 @@ def retrieve_by_id( Args: id (str): The unique identifier of the context index. - 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: Any: The index information, including its configuration and metadata. """ - spec = self._retrieve_by_id_spec( - id, - folder_key=folder_key, - folder_path=folder_path, - ) + spec = self._retrieve_by_id_spec(id) return self.request( spec.method, @@ -324,17 +318,11 @@ async def retrieve_by_id_async( Args: id (str): The unique identifier of the context index. - 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: Any: The index information, including its configuration and metadata. """ - spec = self._retrieve_by_id_spec( - id, - folder_key=folder_key, - folder_path=folder_path, - ) + spec = self._retrieve_by_id_spec(id) response = await self.request_async( spec.method, @@ -400,20 +388,16 @@ 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( + @traced(name="contextgrounding_create_ephemeral_index", run_type="uipath") + def create_ephemeral_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( + """Create a new context ephemeral grounding index.""" + spec = self._create_ephemeral_spec( usage, attachments, - folder_path=folder_path, - folder_key=folder_key, ) response = self.request( @@ -1407,9 +1391,7 @@ def _retrieve_by_id_spec( return RequestSpec( method="GET", endpoint=Endpoint(f"/ecs_/v2/indexes/{id}"), - headers={ - **header_folder(folder_key, None), - }, + headers={}, ) def _delete_by_id_spec( diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index f5a5b03c3..2aa2aaeac 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -268,7 +268,11 @@ async def read_trigger(self, trigger: UiPathResumeTrigger) -> Any | None: f"Index ingestion '{ephemeral_index.name}' did not finish successfully.", ) +<<<<<<< HEAD trigger_response = ephemeral_index.model_dump() +======= + trigger_response = ephemeral_index +>>>>>>> cf1f8f3 (add create ephemeral index) return trigger_response @@ -427,6 +431,10 @@ async def create_trigger(self, suspend_value: Any) -> UiPathResumeTrigger: await self._handle_ephemeral_index_job_trigger( suspend_value, resume_trigger ) + case UiPathResumeTriggerType.EPHEMERAL_INDEX: + await self._handle_ephemeral_index_job_trigger( + suspend_value, resume_trigger, uipath + ) case UiPathResumeTriggerType.BATCH_RAG: await self._handle_batch_rag_job_trigger( suspend_value, resume_trigger @@ -595,6 +603,7 @@ async def _handle_deep_rag_job_trigger( ) async def _handle_ephemeral_index_job_trigger( +<<<<<<< HEAD self, value: Any, resume_trigger: UiPathResumeTrigger ) -> None: """Handle ephemeral index. @@ -616,6 +625,20 @@ async def _handle_ephemeral_index_job_trigger( ) if not ephemeral_index: raise Exception("Failed to create ephemeral index") +======= + self, value: Any, resume_trigger: UiPathResumeTrigger, uipath: UiPath + ) -> None: + """Handle ephemeral index""" + if isinstance(value, WaitEphemeralIndex): + resume_trigger.item_key = value.ephemeral_index.id + elif isinstance(value, CreateEphemeralIndex): + ephemeral_index = uipath.context_grounding.create_ephemeral_index( + usage=value.usage, + attachments=value.attachments, + ) + if not ephemeral_index: + raise Exception("Failed to start ephemeral index") +>>>>>>> cf1f8f3 (add create ephemeral index) resume_trigger.item_key = ephemeral_index.id assert resume_trigger.item_key From 7ca69633630af92be48c6393554a454585dfa1de Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Fri, 16 Jan 2026 15:10:47 -0800 Subject: [PATCH 3/8] add tests --- src/uipath/platform/resume_triggers/_protocol.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index 2aa2aaeac..adbb4879d 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -268,11 +268,7 @@ async def read_trigger(self, trigger: UiPathResumeTrigger) -> Any | None: f"Index ingestion '{ephemeral_index.name}' did not finish successfully.", ) -<<<<<<< HEAD trigger_response = ephemeral_index.model_dump() -======= - trigger_response = ephemeral_index ->>>>>>> cf1f8f3 (add create ephemeral index) return trigger_response @@ -630,15 +626,19 @@ async def _handle_ephemeral_index_job_trigger( ) -> None: """Handle ephemeral index""" if isinstance(value, WaitEphemeralIndex): - resume_trigger.item_key = value.ephemeral_index.id + resume_trigger.item_key = value.index.id elif isinstance(value, CreateEphemeralIndex): ephemeral_index = uipath.context_grounding.create_ephemeral_index( usage=value.usage, attachments=value.attachments, ) if not ephemeral_index: +<<<<<<< HEAD raise Exception("Failed to start ephemeral index") >>>>>>> cf1f8f3 (add create ephemeral index) +======= + raise Exception("Failed to create ephemeral index") +>>>>>>> aca152a (add tests) resume_trigger.item_key = ephemeral_index.id assert resume_trigger.item_key From 5997229be20dd7c829ec66b53e4c909e1b606b96 Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Mon, 19 Jan 2026 13:54:22 -0800 Subject: [PATCH 4/8] fix tests --- .../_context_grounding_service.py | 23 ------------------- .../platform/resume_triggers/_protocol.py | 8 ++++--- .../test_context_grounding_service.py | 4 ++++ 3 files changed, 9 insertions(+), 26 deletions(-) diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 470682a08..9eb01140f 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -1,4 +1,3 @@ -import uuid from pathlib import Path from typing import Annotated, Any, Dict, List, Optional, Tuple, Union @@ -387,28 +386,6 @@ def create_index( return ContextGroundingIndex.model_validate(response.json()) - @resource_override(resource_type="index") - @traced(name="contextgrounding_create_ephemeral_index", run_type="uipath") - def create_ephemeral_index( - self, - usage: str, - attachments: list[uuid.UUID], - ) -> ContextGroundingIndex: - """Create a new context ephemeral grounding index.""" - spec = self._create_ephemeral_spec( - usage, - attachments, - ) - - 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( diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index adbb4879d..ad7ef57d1 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -628,9 +628,11 @@ async def _handle_ephemeral_index_job_trigger( if isinstance(value, WaitEphemeralIndex): resume_trigger.item_key = value.index.id elif isinstance(value, CreateEphemeralIndex): - ephemeral_index = uipath.context_grounding.create_ephemeral_index( - usage=value.usage, - attachments=value.attachments, + ephemeral_index = ( + await uipath.context_grounding.create_ephemeral_index_async( + usage=value.usage, + attachments=value.attachments, + ) ) if not ephemeral_index: <<<<<<< HEAD diff --git a/tests/sdk/services/test_context_grounding_service.py b/tests/sdk/services/test_context_grounding_service.py index a746d9e39..e25123acf 100644 --- a/tests/sdk/services/test_context_grounding_service.py +++ b/tests/sdk/services/test_context_grounding_service.py @@ -1837,7 +1837,11 @@ async def test_create_ephemeral_index_async( }, ) +<<<<<<< HEAD attachment_ids = [str(uuid.uuid4()), str(uuid.uuid4())] +======= + attachment_ids = [uuid.uuid4(), uuid.uuid4()] +>>>>>>> 7addb23 (fix tests) index = await service.create_ephemeral_index_async( usage="DeepRAG", attachments=attachment_ids, From 12c9be645c12e0795fb1344c61dee5b2560f1486 Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Thu, 22 Jan 2026 01:01:47 -0800 Subject: [PATCH 5/8] add new tool definitions --- src/uipath/agent/models/agent.py | 61 ++++++++++++++++++- .../platform/common/interrupt_models.py | 1 + .../_context_grounding_service.py | 29 ++++++--- .../platform/resume_triggers/_protocol.py | 1 + 4 files changed, 82 insertions(+), 10 deletions(-) diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index 6dbc85448..55af46691 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -62,6 +62,7 @@ class AgentInternalToolType(str, Enum): """Agent internal tool type enumeration.""" ANALYZE_FILES = "analyze-attachments" + DEEP_RAG = "deep-rag" class AgentEscalationRecipientType(str, Enum): @@ -119,6 +120,20 @@ class TextTokenType(str, Enum): EXPRESSION = "expression" +class CitationMode(str, Enum): + """Citation mode enumeration.""" + + INLINE = "Inline" + SKIP = "Skip" + + +class FileExtension(str, Enum): + """File extension enumeration.""" + + PDF = "pdf" + TXT = "txt" + + class BaseCfg(BaseModel): """Base configuration model with common settings.""" @@ -245,6 +260,18 @@ class AgentContextValueSetting(BaseCfg): value: Any = Field(...) +class DeepRagCitationModeSetting(BaseCfg): + """DeepRAG citation mode setting model.""" + + value: CitationMode = Field(...) + + +class DeepRagFileExtensionSetting(BaseCfg): + """DeepRAG file extension setting model.""" + + value: FileExtension = Field(...) + + class AgentContextOutputColumn(BaseCfg): """Agent context output column model.""" @@ -605,14 +632,42 @@ class AgentIntegrationToolProperties(BaseResourceProperties): ) -class AgentInternalToolProperties(BaseResourceProperties): - """Agent internal tool properties model.""" +class AgentInternalAnalyzeFilesToolProperties(BaseResourceProperties): + """Agent internal analyze files tool properties model.""" tool_type: Literal[AgentInternalToolType.ANALYZE_FILES] = Field( - ..., alias="toolType" + alias="toolType", default=AgentInternalToolType.ANALYZE_FILES, frozen=True ) +class AgentInternalDeepRagToolProperties(BaseResourceProperties): + """Agent internal DeepRAG tool properties model.""" + + tool_type: Literal[AgentInternalToolType.DEEP_RAG] = Field( + alias="toolType", default=AgentInternalToolType.DEEP_RAG, frozen=True + ) + settings: AgentInternalDeepRagSettings = Field(..., alias="settings") + + +AgentInternalToolProperties = Annotated[ + Union[ + AgentInternalAnalyzeFilesToolProperties, + AgentInternalDeepRagToolProperties, + ], + Field(discriminator="tool_type"), +] + + +class AgentInternalDeepRagSettings(BaseCfg): + """Agent internal DeepRAG tool settings model.""" + + context_type: str = Field(..., alias="contextType") + query: AgentContextQuerySetting = Field(None) + folder_path_prefix: AgentContextQuerySetting = Field(None, alias="folderPathPrefix") + citation_mode: DeepRagCitationModeSetting = Field(..., alias="citationMode") + file_extension: DeepRagFileExtensionSetting = Field(..., alias="fileExtension") + + class AgentIntegrationToolResourceConfig(BaseAgentToolResourceConfig): """Agent integration tool resource configuration model.""" diff --git a/src/uipath/platform/common/interrupt_models.py b/src/uipath/platform/common/interrupt_models.py index 7b6d084a1..b9d48eb8c 100644 --- a/src/uipath/platform/common/interrupt_models.py +++ b/src/uipath/platform/common/interrupt_models.py @@ -85,6 +85,7 @@ class CreateDeepRag(BaseModel): name: str index_name: Annotated[str, Field(max_length=512)] + index_id: Annotated[str, Field(max_length=512)] | None = None prompt: Annotated[str, Field(max_length=250000)] glob_pattern: Annotated[str, Field(max_length=512, default="*")] = "**" citation_mode: CitationMode = CitationMode.SKIP diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 9eb01140f..d425b74e2 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -447,6 +447,8 @@ def create_ephemeral_index( self, usage: EphemeralIndexUsage, attachments: list[str], + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, ) -> ContextGroundingIndex: """Create a new ephemeral context grounding index. @@ -460,6 +462,8 @@ def create_ephemeral_index( spec = self._create_ephemeral_spec( usage, attachments, + folder_path=folder_path, + folder_key=folder_key, ) response = self.request( @@ -477,6 +481,8 @@ async def create_ephemeral_index_async( self, usage: EphemeralIndexUsage, attachments: list[str], + folder_key: Optional[str] = None, + folder_path: Optional[str] = None, ) -> ContextGroundingIndex: """Create a new ephemeral context grounding index. @@ -490,6 +496,8 @@ async def create_ephemeral_index_async( spec = self._create_ephemeral_spec( usage, attachments, + folder_path=folder_path, + folder_key=folder_key, ) response = await self.request_async( @@ -883,6 +891,7 @@ async def start_deep_rag_async( self, name: str, index_name: Annotated[str, Field(max_length=512)], + index_id: Annotated[str, Field(max_length=512)], prompt: Annotated[str, Field(max_length=250000)], glob_pattern: Annotated[str, Field(max_length=512, default="*")] = "**", citation_mode: CitationMode = CitationMode.SKIP, @@ -904,14 +913,16 @@ async def start_deep_rag_async( Returns: DeepRagCreationResponse: The Deep RAG task creation response. """ - index = await self.retrieve_async( - index_name, folder_key=folder_key, folder_path=folder_path - ) - if index and index.in_progress_ingestion(): - raise IngestionInProgressException(index_name=index_name) + if not index_id: + index = await self.retrieve_async( + index_name, folder_key=folder_key, folder_path=folder_path + ) + if index and index.in_progress_ingestion(): + raise IngestionInProgressException(index_name=index_name) + index_id = index.id spec = self._deep_rag_creation_spec( - index_id=index.id, + index_id=index_id, name=name, glob_pattern=glob_pattern, prompt=prompt, @@ -1261,6 +1272,8 @@ def _create_ephemeral_spec( """ data_source_dict = self._build_ephemeral_data_source(attachments) + folder_key = self._resolve_folder_key(folder_key, folder_path) + payload = CreateEphemeralIndexPayload( usage=usage, data_source=data_source_dict, @@ -1270,7 +1283,9 @@ def _create_ephemeral_spec( method="POST", endpoint=Endpoint("/ecs_/v2/indexes/createephemeral"), json=payload.model_dump(by_alias=True, exclude_none=True), - headers={}, + headers={ + **header_folder(folder_key, None), + }, ) def _build_data_source(self, source: SourceConfig) -> Dict[str, Any]: diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index ad7ef57d1..f86fb2d4a 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -580,6 +580,7 @@ async def _handle_deep_rag_job_trigger( deep_rag = await uipath.context_grounding.start_deep_rag_async( name=value.name, index_name=value.index_name, + index_id=value.index_id, prompt=value.prompt, glob_pattern=value.glob_pattern, citation_mode=value.citation_mode, From 33f33e9e3813d0d1e99b17f66ae595b5018d652c Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Sun, 25 Jan 2026 11:44:01 -0800 Subject: [PATCH 6/8] cleanup --- .../_context_grounding_service.py | 42 +++++++++---------- .../platform/resume_triggers/_protocol.py | 25 ----------- .../test_context_grounding_service.py | 5 +-- 3 files changed, 21 insertions(+), 51 deletions(-) diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index d425b74e2..991a9f798 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -291,11 +291,17 @@ def retrieve_by_id( Args: id (str): The unique identifier of the context index. + 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: Any: The index information, including its configuration and metadata. """ - spec = self._retrieve_by_id_spec(id) + spec = self._retrieve_by_id_spec( + id, + folder_key=folder_key, + folder_path=folder_path, + ) return self.request( spec.method, @@ -317,11 +323,17 @@ async def retrieve_by_id_async( Args: id (str): The unique identifier of the context index. + 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: Any: The index information, including its configuration and metadata. """ - spec = self._retrieve_by_id_spec(id) + spec = self._retrieve_by_id_spec( + id, + folder_key=folder_key, + folder_path=folder_path, + ) response = await self.request_async( spec.method, @@ -444,11 +456,7 @@ async def create_index_async( @resource_override(resource_type="index") @traced(name="contextgrounding_create_ephemeral_index", run_type="uipath") def create_ephemeral_index( - self, - usage: EphemeralIndexUsage, - attachments: list[str], - folder_key: Optional[str] = None, - folder_path: Optional[str] = None, + self, usage: EphemeralIndexUsage, attachments: list[str] ) -> ContextGroundingIndex: """Create a new ephemeral context grounding index. @@ -462,8 +470,6 @@ def create_ephemeral_index( spec = self._create_ephemeral_spec( usage, attachments, - folder_path=folder_path, - folder_key=folder_key, ) response = self.request( @@ -478,11 +484,7 @@ def create_ephemeral_index( @resource_override(resource_type="index") @traced(name="contextgrounding_create_ephemeral_index", run_type="uipath") async def create_ephemeral_index_async( - self, - usage: EphemeralIndexUsage, - attachments: list[str], - folder_key: Optional[str] = None, - folder_path: Optional[str] = None, + self, usage: EphemeralIndexUsage, attachments: list[str] ) -> ContextGroundingIndex: """Create a new ephemeral context grounding index. @@ -496,8 +498,6 @@ async def create_ephemeral_index_async( spec = self._create_ephemeral_spec( usage, attachments, - folder_path=folder_path, - folder_key=folder_key, ) response = await self.request_async( @@ -1272,8 +1272,6 @@ def _create_ephemeral_spec( """ data_source_dict = self._build_ephemeral_data_source(attachments) - folder_key = self._resolve_folder_key(folder_key, folder_path) - payload = CreateEphemeralIndexPayload( usage=usage, data_source=data_source_dict, @@ -1283,9 +1281,7 @@ def _create_ephemeral_spec( method="POST", endpoint=Endpoint("/ecs_/v2/indexes/createephemeral"), json=payload.model_dump(by_alias=True, exclude_none=True), - headers={ - **header_folder(folder_key, None), - }, + headers={}, ) def _build_data_source(self, source: SourceConfig) -> Dict[str, Any]: @@ -1383,7 +1379,9 @@ def _retrieve_by_id_spec( return RequestSpec( method="GET", endpoint=Endpoint(f"/ecs_/v2/indexes/{id}"), - headers={}, + headers={ + **header_folder(folder_key, None), + }, ) def _delete_by_id_spec( diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index f86fb2d4a..920d5f4fe 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -427,10 +427,6 @@ async def create_trigger(self, suspend_value: Any) -> UiPathResumeTrigger: await self._handle_ephemeral_index_job_trigger( suspend_value, resume_trigger ) - case UiPathResumeTriggerType.EPHEMERAL_INDEX: - await self._handle_ephemeral_index_job_trigger( - suspend_value, resume_trigger, uipath - ) case UiPathResumeTriggerType.BATCH_RAG: await self._handle_batch_rag_job_trigger( suspend_value, resume_trigger @@ -600,7 +596,6 @@ async def _handle_deep_rag_job_trigger( ) async def _handle_ephemeral_index_job_trigger( -<<<<<<< HEAD self, value: Any, resume_trigger: UiPathResumeTrigger ) -> None: """Handle ephemeral index. @@ -622,26 +617,6 @@ async def _handle_ephemeral_index_job_trigger( ) if not ephemeral_index: raise Exception("Failed to create ephemeral index") -======= - self, value: Any, resume_trigger: UiPathResumeTrigger, uipath: UiPath - ) -> None: - """Handle ephemeral index""" - if isinstance(value, WaitEphemeralIndex): - resume_trigger.item_key = value.index.id - elif isinstance(value, CreateEphemeralIndex): - ephemeral_index = ( - await uipath.context_grounding.create_ephemeral_index_async( - usage=value.usage, - attachments=value.attachments, - ) - ) - if not ephemeral_index: -<<<<<<< HEAD - raise Exception("Failed to start ephemeral index") ->>>>>>> cf1f8f3 (add create ephemeral index) -======= - raise Exception("Failed to create ephemeral index") ->>>>>>> aca152a (add tests) resume_trigger.item_key = ephemeral_index.id assert resume_trigger.item_key diff --git a/tests/sdk/services/test_context_grounding_service.py b/tests/sdk/services/test_context_grounding_service.py index e25123acf..b85e6d192 100644 --- a/tests/sdk/services/test_context_grounding_service.py +++ b/tests/sdk/services/test_context_grounding_service.py @@ -1837,11 +1837,8 @@ async def test_create_ephemeral_index_async( }, ) -<<<<<<< HEAD attachment_ids = [str(uuid.uuid4()), str(uuid.uuid4())] -======= - attachment_ids = [uuid.uuid4(), uuid.uuid4()] ->>>>>>> 7addb23 (fix tests) + index = await service.create_ephemeral_index_async( usage="DeepRAG", attachments=attachment_ids, From 5bdbf83f46d73ceebbff4c789fc47fd7f842a43c Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Sun, 25 Jan 2026 17:36:40 -0800 Subject: [PATCH 7/8] add batch rag --- src/uipath/agent/models/agent.py | 59 +++++++++++++++++-- .../_context_grounding_service.py | 18 +++--- .../test_context_grounding_service.py | 1 - 3 files changed, 66 insertions(+), 12 deletions(-) diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index 55af46691..9a110387f 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -63,6 +63,7 @@ class AgentInternalToolType(str, Enum): ANALYZE_FILES = "analyze-attachments" DEEP_RAG = "deep-rag" + BATCH_TRANSFORM = "batch-transform" class AgentEscalationRecipientType(str, Enum): @@ -127,13 +128,26 @@ class CitationMode(str, Enum): SKIP = "Skip" -class FileExtension(str, Enum): - """File extension enumeration.""" +class DeepRagFileExtension(str, Enum): + """File extension enumeration for DeepRAG.""" PDF = "pdf" TXT = "txt" +class BatchTransformFileExtension(str, Enum): + """File extension enumeration for Batch Transform.""" + + CSV = "csv" + + +class BatchTransformWebSearchGrounding(str, Enum): + """Batch Transform web search grounding enumeration.""" + + ENABLED = "Enabled" + DISABLED = "Disabled" + + class BaseCfg(BaseModel): """Base configuration model with common settings.""" @@ -269,7 +283,19 @@ class DeepRagCitationModeSetting(BaseCfg): class DeepRagFileExtensionSetting(BaseCfg): """DeepRAG file extension setting model.""" - value: FileExtension = Field(...) + value: DeepRagFileExtension = Field(...) + + +class BatchTransformFileExtensionSetting(BaseCfg): + """Batch Transform file extension setting model.""" + + value: BatchTransformFileExtension = Field(...) + + +class BatchTransformWebSearchGroundingSetting(BaseCfg): + """DeepRAG file extension setting model.""" + + value: BatchTransformWebSearchGrounding = Field(...) class AgentContextOutputColumn(BaseCfg): @@ -649,10 +675,20 @@ class AgentInternalDeepRagToolProperties(BaseResourceProperties): settings: AgentInternalDeepRagSettings = Field(..., alias="settings") +class AgentInternalBatchTransformToolProperties(BaseResourceProperties): + """Agent internal Batch Tranform tool properties model.""" + + tool_type: Literal[AgentInternalToolType.BATCH_TRANSFORM] = Field( + alias="toolType", default=AgentInternalToolType.BATCH_TRANSFORM, frozen=True + ) + settings: AgentInternalBatchTransformSettings = Field(..., alias="settings") + + AgentInternalToolProperties = Annotated[ Union[ AgentInternalAnalyzeFilesToolProperties, AgentInternalDeepRagToolProperties, + AgentInternalBatchTransformToolProperties, ], Field(discriminator="tool_type"), ] @@ -662,12 +698,27 @@ class AgentInternalDeepRagSettings(BaseCfg): """Agent internal DeepRAG tool settings model.""" context_type: str = Field(..., alias="contextType") - query: AgentContextQuerySetting = Field(None) + query: AgentContextQuerySetting = Field(...) folder_path_prefix: AgentContextQuerySetting = Field(None, alias="folderPathPrefix") citation_mode: DeepRagCitationModeSetting = Field(..., alias="citationMode") file_extension: DeepRagFileExtensionSetting = Field(..., alias="fileExtension") +class AgentInternalBatchTransformSettings(BaseCfg): + """Agent internal DeepRAG tool settings model.""" + + context_type: str = Field(..., alias="contextType") + query: AgentContextQuerySetting = Field(...) + folder_path_prefix: AgentContextQuerySetting = Field(None, alias="folderPathPrefix") + file_extension: BatchTransformFileExtensionSetting = Field( + ..., alias="fileExtension" + ) + output_columns: List[AgentContextOutputColumn] = Field(..., alias="outputColumns") + web_search_grounding: BatchTransformWebSearchGroundingSetting = Field( + ..., alias="folderPathPrefix" + ) + + class AgentIntegrationToolResourceConfig(BaseAgentToolResourceConfig): """Agent integration tool resource configuration model.""" diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 991a9f798..283764a57 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -842,6 +842,7 @@ def start_deep_rag( prompt: Annotated[str, Field(max_length=250000)], glob_pattern: Annotated[str, Field(max_length=512, default="*")] = "**", citation_mode: CitationMode = CitationMode.SKIP, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> DeepRagCreationResponse: @@ -855,18 +856,20 @@ def start_deep_rag( citation_mode (CitationMode): The citation mode to use. Defaults to SKIP. folder_key (str, optional): The folder key where the index resides. Defaults to None. folder_path (str, optional): The folder path where the index resides. Defaults to None. + index_id (str): The id of the context index to search in, used in place of name if present Returns: DeepRagCreationResponse: The Deep RAG task creation response. """ - index = self.retrieve( - index_name, folder_key=folder_key, folder_path=folder_path - ) - if index and index.in_progress_ingestion(): - raise IngestionInProgressException(index_name=index_name) + if not index_id: + index = self.retrieve( + index_name, folder_key=folder_key, folder_path=folder_path + ) + if index and index.in_progress_ingestion(): + raise IngestionInProgressException(index_name=index_name) spec = self._deep_rag_creation_spec( - index_id=index.id, + index_id=index_id, name=name, glob_pattern=glob_pattern, prompt=prompt, @@ -891,10 +894,10 @@ async def start_deep_rag_async( self, name: str, index_name: Annotated[str, Field(max_length=512)], - index_id: Annotated[str, Field(max_length=512)], prompt: Annotated[str, Field(max_length=250000)], glob_pattern: Annotated[str, Field(max_length=512, default="*")] = "**", citation_mode: CitationMode = CitationMode.SKIP, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> DeepRagCreationResponse: @@ -909,6 +912,7 @@ async def start_deep_rag_async( citation_mode (CitationMode): The citation mode to use. Defaults to SKIP. folder_key (str, optional): The folder key where the index resides. Defaults to None. folder_path (str, optional): The folder path where the index resides. Defaults to None. + index_id (str): The id of the context index to search in, used in place of name if present Returns: DeepRagCreationResponse: The Deep RAG task creation response. diff --git a/tests/sdk/services/test_context_grounding_service.py b/tests/sdk/services/test_context_grounding_service.py index b85e6d192..a746d9e39 100644 --- a/tests/sdk/services/test_context_grounding_service.py +++ b/tests/sdk/services/test_context_grounding_service.py @@ -1838,7 +1838,6 @@ async def test_create_ephemeral_index_async( ) attachment_ids = [str(uuid.uuid4()), str(uuid.uuid4())] - index = await service.create_ephemeral_index_async( usage="DeepRAG", attachments=attachment_ids, From d0989ce505f3031d74ebfd6a9165c0242ca22c34 Mon Sep 17 00:00:00 2001 From: Caleb Martin Date: Wed, 28 Jan 2026 23:50:25 -0800 Subject: [PATCH 8/8] feat: add batch transform --- src/uipath/agent/models/agent.py | 12 +- .../platform/common/interrupt_models.py | 3 + .../_context_grounding_service.py | 169 ++++++++++++++++-- .../platform/resume_triggers/_protocol.py | 63 ++++--- 4 files changed, 209 insertions(+), 38 deletions(-) diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index 9a110387f..87a8352fe 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -263,9 +263,9 @@ class AgentUnknownResourceConfig(BaseAgentResourceConfig): class AgentContextQuerySetting(BaseCfg): """Agent context query setting model.""" - value: str | None = Field(None) - description: str | None = Field(None) - variant: str | None = Field(None) + value: str | None = Field(default=None) + description: str | None = Field(default=None) + variant: str | None = Field(default=None) class AgentContextValueSetting(BaseCfg): @@ -705,17 +705,17 @@ class AgentInternalDeepRagSettings(BaseCfg): class AgentInternalBatchTransformSettings(BaseCfg): - """Agent internal DeepRAG tool settings model.""" + """Agent internal Batch Transform tool settings model.""" context_type: str = Field(..., alias="contextType") query: AgentContextQuerySetting = Field(...) - folder_path_prefix: AgentContextQuerySetting = Field(None, alias="folderPathPrefix") + folder_path_prefix: AgentContextQuerySetting = Field(default=None, alias="folderPathPrefix") file_extension: BatchTransformFileExtensionSetting = Field( ..., alias="fileExtension" ) output_columns: List[AgentContextOutputColumn] = Field(..., alias="outputColumns") web_search_grounding: BatchTransformWebSearchGroundingSetting = Field( - ..., alias="folderPathPrefix" + ..., alias="webSearchGrounding" ) diff --git a/src/uipath/platform/common/interrupt_models.py b/src/uipath/platform/common/interrupt_models.py index b9d48eb8c..ac336335c 100644 --- a/src/uipath/platform/common/interrupt_models.py +++ b/src/uipath/platform/common/interrupt_models.py @@ -91,6 +91,7 @@ class CreateDeepRag(BaseModel): citation_mode: CitationMode = CitationMode.SKIP index_folder_key: str | None = None index_folder_path: str | None = None + is_ephemeral: bool | None = None class WaitDeepRag(BaseModel): @@ -119,6 +120,7 @@ class CreateBatchTransform(BaseModel): name: str index_name: str + index_id: Annotated[str, Field(max_length=512)] | None = None prompt: Annotated[str, Field(max_length=250000)] output_columns: list[BatchTransformOutputColumn] storage_bucket_folder_path_prefix: Annotated[str | None, Field(max_length=512)] = ( @@ -128,6 +130,7 @@ class CreateBatchTransform(BaseModel): destination_path: str index_folder_key: str | None = None index_folder_path: str | None = None + is_ephemeral: bool | None = None class WaitBatchTransform(BaseModel): diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 283764a57..57d624799 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -581,6 +581,7 @@ def start_batch_transform( str | None, Field(max_length=512) ] = None, enable_web_search_grounding: bool = False, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> BatchTransformCreationResponse: @@ -595,20 +596,23 @@ def start_batch_transform( output_columns (list[BatchTransformOutputColumn]): The output columns to add into the csv. storage_bucket_folder_path_prefix (str): The prefix pattern for filtering files in the storage bucket. Use "*" to include all files. Defaults to "*". enable_web_search_grounding (Optional[bool]): Whether to enable web search. Defaults to False. + index_id (str): The id of the context index to search in, used in place of name if present folder_key (str, optional): The folder key where the index resides. Defaults to None. folder_path (str, optional): The folder path where the index resides. Defaults to None. Returns: BatchTransformCreationResponse: The batch transform task creation response. """ - index = self.retrieve( - index_name, folder_key=folder_key, folder_path=folder_path - ) - if index and index.in_progress_ingestion(): - raise IngestionInProgressException(index_name=index_name) + if not index_id: + index = self.retrieve( + index_name, folder_key=folder_key, folder_path=folder_path + ) + if index and index.in_progress_ingestion(): + raise IngestionInProgressException(index_name=index_name) + index_id = index.id spec = self._batch_transform_creation_spec( - index_id=index.id, + index_id=index_id, name=name, storage_bucket_folder_path_prefix=storage_bucket_folder_path_prefix, prompt=prompt, @@ -639,6 +643,7 @@ async def start_batch_transform_async( str | None, Field(max_length=512) ] = None, enable_web_search_grounding: bool = False, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> BatchTransformCreationResponse: @@ -653,20 +658,23 @@ async def start_batch_transform_async( output_columns (list[BatchTransformOutputColumn]): The output columns to add into the csv. storage_bucket_folder_path_prefix (str): The prefix pattern for filtering files in the storage bucket. Use "*" to include all files. Defaults to "*". enable_web_search_grounding (Optional[bool]): Whether to enable web search. Defaults to False. + index_id (str): The id of the context index to search in, used in place of name if present folder_key (str, optional): The folder key where the index resides. Defaults to None. folder_path (str, optional): The folder path where the index resides. Defaults to None. Returns: BatchTransformCreationResponse: The batch transform task creation response. """ - index = await self.retrieve_async( - index_name, folder_key=folder_key, folder_path=folder_path - ) - if index and index.in_progress_ingestion(): - raise IngestionInProgressException(index_name=index_name) + if not index_id: + index = await self.retrieve_async( + index_name, folder_key=folder_key, folder_path=folder_path + ) + if index and index.in_progress_ingestion(): + raise IngestionInProgressException(index_name=index_name) + index_id = index.id spec = self._batch_transform_creation_spec( - index_id=index.id, + index_id=index_id, name=name, storage_bucket_folder_path_prefix=storage_bucket_folder_path_prefix, prompt=prompt, @@ -685,6 +693,52 @@ async def start_batch_transform_async( ) return BatchTransformCreationResponse.model_validate(response.json()) + @resource_override(resource_type="index", resource_identifier="index_name") + @traced(name="contextgrounding_start_batch_transform_async", run_type="uipath") + async def start_batch_transform_ephemeral_async( + self, + name: str, + prompt: Annotated[str, Field(max_length=250000)], + output_columns: list[BatchTransformOutputColumn], + storage_bucket_folder_path_prefix: Annotated[ + str | None, Field(max_length=512) + ] = None, + enable_web_search_grounding: bool = False, + index_id: Annotated[str, Field(max_length=512)] | None = None, + ) -> BatchTransformCreationResponse: + """Asynchronously starts a Batch Transform, task on the targeted index. + + Batch Transform tasks are processing and transforming csv files from the index. + + Args: + name (str): The name of the Deep RAG task. + prompt (str): Describe the task: what to research, what to synthesize. + output_columns (list[BatchTransformOutputColumn]): The output columns to add into the csv. + storage_bucket_folder_path_prefix (str): The prefix pattern for filtering files in the storage bucket. Use "*" to include all files. Defaults to "*". + enable_web_search_grounding (Optional[bool]): Whether to enable web search. Defaults to False. + index_id (str): The id of the context index to search in, used in place of name if present + + Returns: + BatchTransformCreationResponse: The batch transform task creation response. + """ + spec = self._batch_transform_ephemeral_creation_spec( + index_id=index_id, + name=name, + storage_bucket_folder_path_prefix=storage_bucket_folder_path_prefix, + prompt=prompt, + output_columns=output_columns, + enable_web_search_grounding=enable_web_search_grounding, + ) + + response = await self.request_async( + spec.method, + spec.endpoint, + json=spec.json, + params=spec.params, + headers=spec.headers, + ) + return BatchTransformCreationResponse.model_validate(response.json()) + @resource_override(resource_type="index", resource_identifier="index_name") @traced(name="contextgrounding_retrieve_batch_transform", run_type="uipath") def retrieve_batch_transform( @@ -867,6 +921,7 @@ def start_deep_rag( ) if index and index.in_progress_ingestion(): raise IngestionInProgressException(index_name=index_name) + index_id = index.id spec = self._deep_rag_creation_spec( index_id=index_id, @@ -945,6 +1000,47 @@ async def start_deep_rag_async( return DeepRagCreationResponse.model_validate(response.json()) + @resource_override(resource_type="index", resource_identifier="index_name") + @traced(name="contextgrounding_start_deep_rag_async", run_type="uipath") + async def start_deep_rag_ephemeral_async( + self, + name: str, + prompt: Annotated[str, Field(max_length=250000)], + glob_pattern: Annotated[str, Field(max_length=512, default="*")] = "**", + citation_mode: CitationMode = CitationMode.SKIP, + index_id: Annotated[str, Field(max_length=512)] | None = None, + ) -> DeepRagCreationResponse: + """Asynchronously starts a Deep RAG task on the targeted index. + + Args: + name (str): The name of the Deep RAG task. + name (str): The name of the Deep RAG task. + prompt (str): Describe the task: what to research across documents, what to synthesize and how to cite sources. + glob_pattern (str): The glob pattern to search in the index. Defaults to "**". + citation_mode (CitationMode): The citation mode to use. Defaults to SKIP. + index_id (str): The id of the context index to search in, used in place of name if present + + Returns: + DeepRagCreationResponse: The Deep RAG task creation response. + """ + spec = self._deep_rag_ephemeral_creation_spec( + index_id=index_id, + name=name, + glob_pattern=glob_pattern, + prompt=prompt, + citation_mode=citation_mode, + ) + + response = await self.request_async( + spec.method, + spec.endpoint, + params=spec.params, + json=spec.json, + headers=spec.headers, + ) + + return DeepRagCreationResponse.model_validate(response.json()) + @resource_override(resource_type="index") @traced(name="contextgrounding_search", run_type="uipath") def search( @@ -1455,6 +1551,29 @@ def _deep_rag_creation_spec( }, ) + def _deep_rag_ephemeral_creation_spec( + self, + index_id: str, + name: str, + glob_pattern: str, + prompt: str, + citation_mode: CitationMode, + ) -> RequestSpec: + return RequestSpec( + method="POST", + endpoint=Endpoint(f"/ecs_/v2/indexes/{index_id}/createDeepRag"), + json={ + "name": name, + "prompt": prompt, + "globPattern": glob_pattern, + "citationMode": citation_mode.value, + }, + params={ + "$select": "id,lastDeepRagStatus,createdDate", + }, + headers={}, + ) + def _batch_transform_creation_spec( self, index_id: str, @@ -1487,6 +1606,32 @@ def _batch_transform_creation_spec( }, ) + def _batch_transform_ephemeral_creation_spec( + self, + index_id: str, + name: str, + enable_web_search_grounding: bool, + output_columns: list[BatchTransformOutputColumn], + storage_bucket_folder_path_prefix: str | None, + prompt: str, + ) -> RequestSpec: + return RequestSpec( + method="POST", + endpoint=Endpoint(f"/ecs_/v2/indexes/{index_id}/createBatchRag"), + json={ + "name": name, + "prompt": prompt, + "targetFileGlobPattern": f"{storage_bucket_folder_path_prefix}/*" + if storage_bucket_folder_path_prefix + else "**", + "useWebSearchGrounding": enable_web_search_grounding, + "outputColumns": [ + column.model_dump(by_alias=True) for column in output_columns + ], + }, + headers={}, + ) + def _deep_rag_retrieve_spec( self, id: str, diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index 920d5f4fe..e71999b29 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -573,16 +573,28 @@ async def _handle_deep_rag_job_trigger( resume_trigger.item_key = value.deep_rag.id elif isinstance(value, CreateDeepRag): uipath = UiPath() - deep_rag = await uipath.context_grounding.start_deep_rag_async( - name=value.name, - index_name=value.index_name, - index_id=value.index_id, - prompt=value.prompt, - glob_pattern=value.glob_pattern, - citation_mode=value.citation_mode, - folder_path=value.index_folder_path, - folder_key=value.index_folder_key, - ) + if value.is_ephemeral: + deep_rag = ( + await uipath.context_grounding.start_deep_rag_ephemeral_async( + name=value.name, + index_id=value.index_id, + prompt=value.prompt, + glob_pattern=value.glob_pattern, + citation_mode=value.citation_mode, + ) + ) + + else: + deep_rag = await uipath.context_grounding.start_deep_rag_async( + name=value.name, + index_name=value.index_name, + index_id=value.index_id, + prompt=value.prompt, + glob_pattern=value.glob_pattern, + citation_mode=value.citation_mode, + folder_path=value.index_folder_path, + folder_key=value.index_folder_key, + ) if not deep_rag: raise Exception("Failed to start deep rag") @@ -642,16 +654,27 @@ async def _handle_batch_rag_job_trigger( resume_trigger.item_key = value.batch_transform.id elif isinstance(value, CreateBatchTransform): uipath = UiPath() - batch_transform = await uipath.context_grounding.start_batch_transform_async( - name=value.name, - index_name=value.index_name, - prompt=value.prompt, - output_columns=value.output_columns, - storage_bucket_folder_path_prefix=value.storage_bucket_folder_path_prefix, - enable_web_search_grounding=value.enable_web_search_grounding, - folder_path=value.index_folder_path, - folder_key=value.index_folder_key, - ) + if value.is_ephemeral: + batch_transform = await uipath.context_grounding.start_batch_transform_ephemeral_async( + name=value.name, + index_id=value.index_id, + prompt=value.prompt, + output_columns=value.output_columns, + storage_bucket_folder_path_prefix=value.storage_bucket_folder_path_prefix, + enable_web_search_grounding=value.enable_web_search_grounding, + ) + else: + batch_transform = await uipath.context_grounding.start_batch_transform_async( + name=value.name, + index_name=value.index_name, + index_id=value.index_id, + prompt=value.prompt, + output_columns=value.output_columns, + storage_bucket_folder_path_prefix=value.storage_bucket_folder_path_prefix, + enable_web_search_grounding=value.enable_web_search_grounding, + folder_path=value.index_folder_path, + folder_key=value.index_folder_key, + ) if not batch_transform: raise Exception("Failed to start batch transform")