diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index f1053b654..aca323f86 100644 --- a/src/uipath/agent/models/agent.py +++ b/src/uipath/agent/models/agent.py @@ -63,6 +63,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): @@ -120,6 +122,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.""" @@ -236,9 +265,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): @@ -247,6 +276,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.""" @@ -607,11 +660,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 | None = Field(default=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 | None = 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 2d76a1349..3e6f61fbf 100644 --- a/src/uipath/platform/common/interrupt_models.py +++ b/src/uipath/platform/common/interrupt_models.py @@ -86,12 +86,21 @@ class CreateDeepRag(BaseModel): """Model representing a Deep RAG task creation.""" name: str - index_name: Annotated[str, Field(max_length=512)] + index_name: Annotated[str, Field(max_length=512)] | None = None + 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_index: bool | None = None + + @model_validator(mode="after") + def validate_ephemeral_index_requires_index_id(self) -> "CreateBatchTransform": + """Validate that if it is an ephemeral index that it is using index id.""" + if self.is_ephemeral_index is True and self.index_id is None: + raise ValueError("Index id must be provided for an ephemeral index") + return self class WaitDeepRag(BaseModel): @@ -119,7 +128,8 @@ class CreateBatchTransform(BaseModel): """Model representing a Batch Transform task creation.""" name: str - index_name: str + index_name: str | None = None + 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)] = ( @@ -129,6 +139,14 @@ class CreateBatchTransform(BaseModel): destination_path: str index_folder_key: str | None = None index_folder_path: str | None = None + is_ephemeral_index: bool | None = None + + @model_validator(mode="after") + def validate_ephemeral_index_requires_index_id(self) -> "CreateBatchTransform": + """Validate that if it is an ephemeral index that it is using index id.""" + if self.is_ephemeral_index is True and self.index_id is None: + raise ValueError("Index id must be provided for an ephemeral index") + return self 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..8dff961b9 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. @@ -578,13 +574,14 @@ async def retrieve_deep_rag_async( def start_batch_transform( self, name: str, - index_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_name: str | None = None, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> BatchTransformCreationResponse: @@ -599,20 +596,28 @@ 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_name or index_id): + raise ValueError("Index name or id is required") + if index_name and index_id: + raise ValueError("Index name or id are mutually exclusive") + + 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, @@ -636,13 +641,14 @@ def start_batch_transform( async def start_batch_transform_async( self, name: str, - index_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_name: str | None = None, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> BatchTransformCreationResponse: @@ -657,20 +663,28 @@ 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_name or index_id): + raise ValueError("Index name or id is required") + if index_name and index_id: + raise ValueError("Index name or id are mutually exclusive") + + 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 +703,98 @@ 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", run_type="uipath") + async def start_batch_transform_ephemeral( + 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 = self.request( + 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_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( @@ -842,10 +948,11 @@ async def download_batch_transform_result_async( def start_deep_rag( self, name: str, - index_name: 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_name: Annotated[str, Field(max_length=512)] | None = None, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> DeepRagCreationResponse: @@ -859,18 +966,26 @@ 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_name or index_id): + raise ValueError("Index name or id is required") + if index_name and index_id: + raise ValueError("Index name or id are mutually exclusive") + + 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, @@ -894,10 +1009,11 @@ def start_deep_rag( async def start_deep_rag_async( self, name: str, - index_name: 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_name: Annotated[str, Field(max_length=512)] | None = None, + index_id: Annotated[str, Field(max_length=512)] | None = None, folder_key: str | None = None, folder_path: str | None = None, ) -> DeepRagCreationResponse: @@ -912,18 +1028,26 @@ 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_name or index_id): + raise ValueError("Index name or id is required") + if index_name and index_id: + raise ValueError("Index name or id are mutually exclusive") + + 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 +1066,88 @@ 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", run_type="uipath") + async def start_deep_rag_ephemeral( + 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 = self.request( + spec.method, + spec.endpoint, + params=spec.params, + json=spec.json, + headers=spec.headers, + ) + + 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 +1658,29 @@ def _deep_rag_creation_spec( }, ) + def _deep_rag_ephemeral_creation_spec( + self, + index_id: str | None, + 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 +1713,32 @@ def _batch_transform_creation_spec( }, ) + def _batch_transform_ephemeral_creation_spec( + self, + index_id: str | None, + 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 9d98460cd..ffd27a1f3 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_index: + 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_index: + 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") diff --git a/tests/cli/test_hitl.py b/tests/cli/test_hitl.py index e1e3506e9..abdbe955b 100644 --- a/tests/cli/test_hitl.py +++ b/tests/cli/test_hitl.py @@ -1026,6 +1026,7 @@ async def test_create_resume_trigger_create_deep_rag( mock_start_deep_rag_async.assert_called_once_with( name=create_deep_rag.name, index_name=create_deep_rag.index_name, + index_id=None, prompt=create_deep_rag.prompt, glob_pattern=create_deep_rag.glob_pattern, citation_mode=create_deep_rag.citation_mode, @@ -1092,6 +1093,7 @@ async def test_create_resume_trigger_create_batch_transform( mock_start_batch_transform.assert_called_once_with( name=create_batch_transform.name, index_name=create_batch_transform.index_name, + index_id=None, prompt=create_batch_transform.prompt, output_columns=create_batch_transform.output_columns, storage_bucket_folder_path_prefix=create_batch_transform.storage_bucket_folder_path_prefix, diff --git a/tests/sdk/services/test_context_grounding_service.py b/tests/sdk/services/test_context_grounding_service.py index a746d9e39..f58a954cb 100644 --- a/tests/sdk/services/test_context_grounding_service.py +++ b/tests/sdk/services/test_context_grounding_service.py @@ -1921,3 +1921,123 @@ async def test_download_batch_transform_result_async_creates_nested_directories( assert destination.exists() assert destination.read_bytes() == b"col1,col2\nval1,val2" assert destination.parent.exists() + + @pytest.mark.anyio + async def test_start_batch_transform_ephemeral_async( + self, + httpx_mock: HTTPXMock, + service: ContextGroundingService, + base_url: str, + org: str, + tenant: str, + version: str, + ) -> None: + httpx_mock.add_response( + url=f"{base_url}{org}{tenant}/ecs_/v2/indexes/ephemeral-index-id/createBatchRag", + status_code=200, + json={ + "id": "new-batch-transform-id", + "lastBatchRagStatus": "Queued", + "errorMessage": None, + }, + ) + + output_columns = [ + BatchTransformOutputColumn( + name="summary", + description="A summary of the document", + ) + ] + + response = await service.start_batch_transform_ephemeral_async( + name="my-ephemeral-batch-transform", + prompt="Summarize all documents", + output_columns=output_columns, + storage_bucket_folder_path_prefix="data", + enable_web_search_grounding=False, + index_id="ephemeral-index-id", + ) + + assert isinstance(response, BatchTransformCreationResponse) + assert response.id == "new-batch-transform-id" + assert response.last_batch_rag_status == "Queued" + + sent_requests = httpx_mock.get_requests() + if sent_requests is None: + raise Exception("No request was sent") + + assert sent_requests[0].method == "POST" + assert ( + f"{base_url}{org}{tenant}/ecs_/v2/indexes/ephemeral-index-id/createBatchRag" + in str(sent_requests[0].url) + ) + + request_data = json.loads(sent_requests[0].content) + assert request_data["name"] == "my-ephemeral-batch-transform" + assert request_data["prompt"] == "Summarize all documents" + assert request_data["targetFileGlobPattern"] == "data/*" + assert request_data["useWebSearchGrounding"] is False + + assert HEADER_USER_AGENT in sent_requests[0].headers + assert ( + sent_requests[0].headers[HEADER_USER_AGENT] + == f"UiPath.Python.Sdk/UiPath.Python.Sdk.Activities.ContextGroundingService.start_batch_transform_ephemeral_async/{version}" + ) + + @pytest.mark.anyio + async def test_start_deep_rag_ephemeral_async( + self, + httpx_mock: HTTPXMock, + service: ContextGroundingService, + base_url: str, + org: str, + tenant: str, + version: str, + ) -> None: + httpx_mock.add_response( + url=f"{base_url}{org}{tenant}/ecs_/v2/indexes/ephemeral-index-id/createDeepRag?$select=id,lastDeepRagStatus,createdDate", + status_code=200, + json={ + "id": "new-deep-rag-task-id", + "lastDeepRagStatus": "Queued", + "createdDate": "2024-01-15T10:30:00Z", + }, + ) + + response = await service.start_deep_rag_ephemeral_async( + name="my-ephemeral-deep-rag-task", + prompt="Summarize all documents related to financial reports", + glob_pattern="*.pdf", + citation_mode=CitationMode.INLINE, + index_id="ephemeral-index-id", + ) + + assert isinstance(response, DeepRagCreationResponse) + assert response.id == "new-deep-rag-task-id" + assert response.last_deep_rag_status == "Queued" + assert response.created_date == "2024-01-15T10:30:00Z" + + sent_requests = httpx_mock.get_requests() + if sent_requests is None: + raise Exception("No request was sent") + + assert sent_requests[0].method == "POST" + assert ( + f"{base_url}{org}{tenant}/ecs_/v2/indexes/ephemeral-index-id/createDeepRag" + in str(sent_requests[0].url) + ) + + request_data = json.loads(sent_requests[0].content) + assert request_data["name"] == "my-ephemeral-deep-rag-task" + assert ( + request_data["prompt"] + == "Summarize all documents related to financial reports" + ) + assert request_data["globPattern"] == "*.pdf" + assert request_data["citationMode"] == "Inline" + + assert HEADER_USER_AGENT in sent_requests[0].headers + assert ( + sent_requests[0].headers[HEADER_USER_AGENT] + == f"UiPath.Python.Sdk/UiPath.Python.Sdk.Activities.ContextGroundingService.start_deep_rag_ephemeral_async/{version}" + )