add tests

Author: CalebMartinUiPath

src/uipath/platform/common/interrupt_models.py

@@ -4,6 +4,10 @@
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
+from uipath.platform.context_grounding.context_grounding_index import (
+    ContextGroundingIndex,
+)
+
 from ..action_center import Task
 from ..context_grounding import (
     BatchTransformCreationResponse,
@@ -85,16 +89,16 @@ class WaitDeepRag(BaseModel):
 
 
 class CreateEphemeralIndex(BaseModel):
-    """Model representing a Jit Index task creation."""
+    """Model representing a Ephemeral Index task creation."""
 
     usage: str
     attachments: list[str]
 
 
 class WaitEphemeralIndex(BaseModel):
-    """Model representing a wait Jit Index task."""
+    """Model representing a wait Ephemeral Index task."""
 
-    id: str
+    index: ContextGroundingIndex
 
 
 class CreateBatchTransform(BaseModel):

src/uipath/platform/context_grounding/context_grounding_index.py

@@ -44,9 +44,9 @@ class ContextGroundingIndex(BaseModel):
         extra="allow",
     )
 
-    @field_serializer("last_ingested", "last_queried", when_used="json")
+    @field_serializer("last_ingested", "last_queried")
     def serialize_datetime(self, value):
-        """Serialize datetime fields to ISO 8601 format for JSON output."""
+        """Serialize datetime fields to ISO 8601 format."""
         if isinstance(value, datetime):
             return value.isoformat() if value else None
         return value

src/uipath/platform/resume_triggers/_protocol.py

@@ -526,14 +526,14 @@ 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:
-                raise Exception("Failed to start ephemeral index")
+                raise Exception("Failed to create ephemeral index")
             resume_trigger.item_key = ephemeral_index.id
 
     async def _handle_batch_rag_job_trigger(

tests/cli/test_hitl.py

@@ -25,13 +25,21 @@
     WaitJob,
     WaitTask,
 )
+from uipath.platform.common.interrupt_models import (
+    CreateEphemeralIndex,
+    WaitEphemeralIndex,
+)
 from uipath.platform.context_grounding import (
     BatchTransformCreationResponse,
     BatchTransformOutputColumn,
     Citation,
     CitationMode,
     DeepRagCreationResponse,
     DeepRagStatus,
+    IndexStatus,
+)
+from uipath.platform.context_grounding.context_grounding_index import (
+    ContextGroundingIndex,
 )
 from uipath.platform.orchestrator import (
     Job,
@@ -517,6 +525,95 @@ async def test_read_batch_rag_trigger_pending(
                 reader = UiPathResumeTriggerReader()
                 await reader.read_trigger(resume_trigger)
 
+    @pytest.mark.anyio
+    async def test_read_ephemeral_index_trigger_successful(
+        self,
+        setup_test_env: None,
+    ) -> None:
+        """Test reading a successful ephemeral index trigger."""
+        index_id = "test-ephemeral-index-id"
+        index_data = {
+            "id": index_id,
+            "name": "test-index",
+            "lastIngestionStatus": IndexStatus.SUCCESSFUL.value,
+        }
+
+        mock_retrieve_by_id = AsyncMock(return_value=index_data)
+
+        with patch(
+            "uipath.platform.context_grounding._context_grounding_service.ContextGroundingService.retrieve_by_id",
+            new=mock_retrieve_by_id,
+        ):
+            resume_trigger = UiPathResumeTrigger(
+                trigger_type=UiPathResumeTriggerType.EPHEMERAL_INDEX,
+                item_key=index_id,
+            )
+            reader = UiPathResumeTriggerReader()
+            result = await reader.read_trigger(resume_trigger)
+
+            assert isinstance(result, ContextGroundingIndex)
+            assert result.id == index_id
+            mock_retrieve_by_id.assert_called_once_with(index_id)
+
+    @pytest.mark.anyio
+    async def test_read_ephemeral_index_trigger_pending(
+        self,
+        setup_test_env: None,
+    ) -> None:
+        """Test reading a pending ephemeral index trigger raises pending error."""
+        from uipath.core.errors import UiPathPendingTriggerError
+
+        index_id = "test-ephemeral-index-id"
+        index_data = {
+            "id": index_id,
+            "name": "test-index",
+            "lastIngestionStatus": IndexStatus.IN_PROGRESS.value,
+        }
+
+        mock_retrieve_by_id = AsyncMock(return_value=index_data)
+
+        with patch(
+            "uipath.platform.context_grounding._context_grounding_service.ContextGroundingService.retrieve_by_id",
+            new=mock_retrieve_by_id,
+        ):
+            resume_trigger = UiPathResumeTrigger(
+                trigger_type=UiPathResumeTriggerType.EPHEMERAL_INDEX,
+                item_key=index_id,
+            )
+
+            with pytest.raises(UiPathPendingTriggerError):
+                reader = UiPathResumeTriggerReader()
+                await reader.read_trigger(resume_trigger)
+
+    @pytest.mark.anyio
+    async def test_read_ephemeral_index_trigger_failed(
+        self,
+        setup_test_env: None,
+    ) -> None:
+        """Test reading a failed ephemeral index trigger raises faulted error."""
+        index_id = "test-ephemeral-index-id"
+        index_data = {
+            "id": index_id,
+            "name": "test-index",
+            "lastIngestionStatus": IndexStatus.FAILED.value,
+        }
+
+        mock_retrieve_by_id = AsyncMock(return_value=index_data)
+
+        with patch(
+            "uipath.platform.context_grounding._context_grounding_service.ContextGroundingService.retrieve_by_id",
+            new=mock_retrieve_by_id,
+        ):
+            resume_trigger = UiPathResumeTrigger(
+                trigger_type=UiPathResumeTriggerType.EPHEMERAL_INDEX,
+                item_key=index_id,
+            )
+
+            with pytest.raises(UiPathFaultedTriggerError) as exc_info:
+                reader = UiPathResumeTriggerReader()
+                await reader.read_trigger(resume_trigger)
+            assert exc_info.value.args[0] == ErrorCategory.USER
+
 
 class TestHitlProcessor:
     """Tests for the HitlProcessor class."""
@@ -778,6 +875,62 @@ async def test_create_resume_trigger_wait_batch_transform(
         assert resume_trigger.trigger_type == UiPathResumeTriggerType.BATCH_RAG
         assert resume_trigger.item_key == batch_transform_id
 
+    @pytest.mark.anyio
+    async def test_create_resume_trigger_create_ephemeral_index(
+        self,
+        setup_test_env: None,
+    ) -> None:
+        """Test creating a resume trigger for CreateEphemeralIndex."""
+        index_id = "test-ephemeral-index-id"
+        attachments = ["attachment-uuid-1", "attachment-uuid-2"]
+        create_ephemeral_index = CreateEphemeralIndex(
+            usage="DeepRAG",
+            attachments=attachments,
+        )
+
+        mock_index = ContextGroundingIndex(
+            id=index_id,
+            name="ephemeral-index",
+            last_ingestion_status=IndexStatus.QUEUED.value,
+        )
+        mock_create_ephemeral_index = AsyncMock(return_value=mock_index)
+
+        with patch(
+            "uipath.platform.context_grounding._context_grounding_service.ContextGroundingService.create_ephemeral_index",
+            new=mock_create_ephemeral_index,
+        ):
+            processor = UiPathResumeTriggerCreator()
+            resume_trigger = await processor.create_trigger(create_ephemeral_index)
+
+            assert resume_trigger is not None
+            assert resume_trigger.trigger_type == UiPathResumeTriggerType.EPHEMERAL_INDEX
+            assert resume_trigger.item_key == index_id
+            mock_create_ephemeral_index.assert_called_once_with(
+                usage=create_ephemeral_index.usage,
+                attachments=create_ephemeral_index.attachments,
+            )
+
+    @pytest.mark.anyio
+    async def test_create_resume_trigger_wait_ephemeral_index(
+        self,
+        setup_test_env: None,
+    ) -> None:
+        """Test creating a resume trigger for WaitEphemeralIndex."""
+        index_id = "test-ephemeral-index-id"
+        ephemeral_index = ContextGroundingIndex(
+            id=index_id,
+            name="ephemeral-index",
+            last_ingestion_status=IndexStatus.IN_PROGRESS.value,
+        )
+        wait_ephemeral_index = WaitEphemeralIndex(index=ephemeral_index)
+
+        processor = UiPathResumeTriggerCreator()
+        resume_trigger = await processor.create_trigger(wait_ephemeral_index)
+
+        assert resume_trigger is not None
+        assert resume_trigger.trigger_type == UiPathResumeTriggerType.EPHEMERAL_INDEX
+        assert resume_trigger.item_key == index_id
+
 
 class TestDocumentExtractionModels:
     """Tests for document extraction models."""

tests/sdk/services/test_context_grounding_service.py

@@ -1760,6 +1760,61 @@ def test_download_batch_transform_result_creates_nested_directories(
         assert destination.read_bytes() == b"col1,col2\nval1,val2"
         assert destination.parent.exists()
 
+    def test_create_ephemeral_index(
+        self,
+        httpx_mock: HTTPXMock,
+        service: ContextGroundingService,
+        base_url: str,
+        org: str,
+        tenant: str,
+        version: str,
+    ) -> None:
+        import uuid
+
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/indexes/createephemeral",
+            status_code=200,
+            json={
+                "id": "ephemeral-index-id",
+                "name": "ephemeral-index",
+                "lastIngestionStatus": "Queued",
+            },
+        )
+
+        attachment_ids = [uuid.uuid4(), uuid.uuid4()]
+        index = service.create_ephemeral_index(
+            usage="DeepRAG",
+            attachments=attachment_ids,
+        )
+
+        assert isinstance(index, ContextGroundingIndex)
+        assert index.id == "ephemeral-index-id"
+        assert index.name == "ephemeral-index"
+        assert index.last_ingestion_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 (
+            sent_requests[0].url
+            == f"{base_url}{org}{tenant}/ecs_/v2/indexes/createephemeral"
+        )
+
+        request_data = json.loads(sent_requests[0].content)
+        assert request_data["usage"] == "DeepRAG"
+        assert "dataSource" in request_data
+        assert request_data["dataSource"]["attachments"] == [
+            str(att) for att in attachment_ids
+        ]
+
+        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.create_ephemeral_index/{version}"
+        )
+
     @pytest.mark.anyio
     async def test_download_batch_transform_result_async_creates_nested_directories(
         self,