add batch rag

Author: CalebMartinUiPath

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."""
 

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.

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,