diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index 6dbc85448..87a8352fe 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -62,6 +62,8 @@ class AgentInternalToolType(str, Enum): """Agent internal tool type enumeration.""" ANALYZE_FILES = "analyze-attachments" + DEEP_RAG = "deep-rag" + BATCH_TRANSFORM = "batch-transform" class AgentEscalationRecipientType(str, Enum): @@ -119,6 +121,33 @@ class TextTokenType(str, Enum): EXPRESSION = "expression" +class CitationMode(str, Enum): + """Citation mode enumeration.""" + + INLINE = "Inline" + SKIP = "Skip" + + +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.""" @@ -234,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): @@ -245,6 +274,30 @@ 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: 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): """Agent context output column model.""" @@ -605,11 +658,64 @@ 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") + + +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"), +] + + +class AgentInternalDeepRagSettings(BaseCfg): + """Agent internal DeepRAG tool settings model.""" + + context_type: str = Field(..., alias="contextType") + 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 Batch Transform tool settings model.""" + + context_type: str = Field(..., alias="contextType") + query: AgentContextQuerySetting = Field(...) + 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="webSearchGrounding" ) diff --git a/src/uipath/platform/common/interrupt_models.py b/src/uipath/platform/common/interrupt_models.py index 7b6d084a1..ac336335c 100644 --- a/src/uipath/platform/common/interrupt_models.py +++ b/src/uipath/platform/common/interrupt_models.py @@ -85,11 +85,13 @@ 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 index_folder_key: str | None = None index_folder_path: str | None = None + is_ephemeral: bool | None = None class WaitDeepRag(BaseModel): @@ -118,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)] = ( @@ -127,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 31cadcc0a..57d624799 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -456,9 +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], + self, usage: EphemeralIndexUsage, attachments: list[str] ) -> ContextGroundingIndex: """Create a new ephemeral context grounding index. @@ -486,9 +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], + self, usage: EphemeralIndexUsage, attachments: list[str] ) -> ContextGroundingIndex: """Create a new ephemeral context grounding index. @@ -585,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: @@ -599,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, @@ -643,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: @@ -657,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, @@ -689,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( @@ -846,6 +896,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: @@ -859,18 +910,21 @@ 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) + 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, @@ -898,6 +952,7 @@ async def start_deep_rag_async( 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: @@ -912,18 +967,21 @@ 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. """ - 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, @@ -942,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( @@ -1452,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, @@ -1484,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 f5a5b03c3..e71999b29 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -573,15 +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, - 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") @@ -641,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")