diff --git a/src/uipath/agent/models/agent.py b/src/uipath/agent/models/agent.py index bdb5e05d5..d966533e2 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.""" @@ -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.""" @@ -505,11 +558,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 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" ) diff --git a/src/uipath/platform/common/interrupt_models.py b/src/uipath/platform/common/interrupt_models.py index 6754c48e2..6a8005914 100644 --- a/src/uipath/platform/common/interrupt_models.py +++ b/src/uipath/platform/common/interrupt_models.py @@ -75,6 +75,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 31cadcc0a..283764a57 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. @@ -846,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: @@ -859,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, @@ -898,6 +897,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 +912,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, diff --git a/src/uipath/platform/resume_triggers/_protocol.py b/src/uipath/platform/resume_triggers/_protocol.py index b9af806d4..6da49aefd 100644 --- a/src/uipath/platform/resume_triggers/_protocol.py +++ b/src/uipath/platform/resume_triggers/_protocol.py @@ -514,6 +514,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,