Skip to content

feat: use CreateDeepRagRaw and WaitEphemeralIndexRaw #665

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
GabrielVasilescu04 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: use CreateDeepRagRaw and WaitEphemeralIndexRaw #665

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ requires-python = ">=3.11"
dependencies = [
"uipath>=2.10.0, <2.11.0",
"uipath-core>=0.5.2, <0.6.0",
"uipath-platform>=0.0.8, <0.1.0",
"uipath-platform==0.0.13.dev1014125197",
"uipath-runtime>=0.9.1, <0.10.0",
"langgraph>=1.0.0, <2.0.0",
"langchain-core>=1.2.11, <2.0.0",
Expand Down Expand Up @@ -71,6 +71,9 @@ dev = [
"rust-just>=1.39.0",
]

[tool.uv.sources]
uipath-platform = { index = "testpypi" }

[tool.hatch.build.targets.wheel]
packages = ["src/uipath_langchain"]

Expand Down
16 changes: 13 additions & 3 deletions src/uipath_langchain/agent/tools/context_tool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,12 @@
)
from uipath.eval.mocks import mockable
from uipath.platform import UiPath
from uipath.platform.common import CreateBatchTransform, CreateDeepRag, UiPathConfig
from uipath.platform.common import CreateBatchTransform, CreateDeepRagRaw, UiPathConfig
from uipath.platform.context_grounding import (
BatchTransformOutputColumn,
CitationMode,
DeepRagContent,
DeepRagStatus,
)
from uipath.runtime.errors import UiPathErrorCategory

Expand Down Expand Up @@ -200,7 +201,7 @@ async def context_tool_fn(query: Optional[str] = None) -> dict[str, Any]:

@durable_interrupt
async def create_deep_rag():
return CreateDeepRag(
return CreateDeepRagRaw(
name=f"task-{uuid.uuid4()}",
index_name=index_name,
prompt=actual_prompt,
Expand All @@ -209,7 +210,16 @@ async def create_deep_rag():
glob_pattern=glob_pattern,
)

return await create_deep_rag()
result = await create_deep_rag()
if result.last_deep_rag_status == DeepRagStatus.FAILED:
return result.failure_reason

if result.content:
content = result.content.model_dump()
content["deepRagId"] = result.id
return content

return {"status": result.last_deep_rag_status, "__internal": "NO_CONTENT"}

return StructuredToolWithOutputType(
name=tool_name,
Expand Down
21 changes: 17 additions & 4 deletions src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,12 @@
)
from uipath.eval.mocks import mockable
from uipath.platform import UiPath
from uipath.platform.common import CreateDeepRag, WaitEphemeralIndex
from uipath.platform.common import CreateDeepRagRaw, WaitEphemeralIndexRaw
from uipath.platform.context_grounding import (
CitationMode,
DeepRagStatus,
EphemeralIndexUsage,
IndexStatus,
)
from uipath.platform.context_grounding.context_grounding_index import (
ContextGroundingIndex,
Expand Down Expand Up @@ -125,7 +127,7 @@ async def create_ephemeral_index():
)
)
if ephemeral_index.in_progress_ingestion():
return WaitEphemeralIndex(index=ephemeral_index)
return WaitEphemeralIndexRaw(index=ephemeral_index)
return ReadyEphemeralIndex(index=ephemeral_index)

index_result = await create_ephemeral_index()
Expand All @@ -134,9 +136,12 @@ async def create_ephemeral_index():
else:
ephemeral_index = index_result

if ephemeral_index.last_ingestion_status == IndexStatus.FAILED:
return ephemeral_index.last_ingestion_failure_reason

@durable_interrupt
async def create_deeprag():
return CreateDeepRag(
return CreateDeepRagRaw(
name=f"task-{uuid.uuid4()}",
index_name=ephemeral_index.name,
index_id=ephemeral_index.id,
Expand All @@ -147,7 +152,15 @@ async def create_deeprag():

result = await create_deeprag()

return result
if result.last_deep_rag_status == DeepRagStatus.FAILED:
return result.failure_reason

if result.content:
content = result.content.model_dump()
content["deepRagId"] = result.id
return content

return {"status": result.last_deep_rag_status, "__internal": "NO_CONTENT"}

return await invoke_deeprag(**kwargs)

Expand Down
194 changes: 188 additions & 6 deletions tests/agent/tools/internal_tools/test_deeprag_tool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
DeepRagFileExtension,
DeepRagFileExtensionSetting,
)
from uipath.platform.context_grounding import DeepRagResponse, DeepRagStatus, IndexStatus
from uipath.platform.context_grounding.context_grounding import DeepRagContent
from uipath.platform.context_grounding.context_grounding_index import (
ContextGroundingIndex,
)
Expand Down Expand Up @@ -152,8 +154,16 @@ async def test_create_deeprag_tool_static_query_index_ready(

# Index is ready → ReadyEphemeralIndex skips interrupt() (no scratchpad in tests).
# Only create_deeprag calls interrupt().
deeprag_id = str(uuid.uuid4())
mock_interrupt.side_effect = [
{"text": "Deep RAG analysis result"},
DeepRagResponse(
id=deeprag_id,
name="test-deeprag",
created_date="2024-01-01",
last_deep_rag_status=DeepRagStatus.SUCCESSFUL,
content=DeepRagContent(text="Deep RAG analysis result", citations=[]),
failure_reason=None,
),
]

mock_wrapper = Mock()
Expand All @@ -175,7 +185,7 @@ async def test_create_deeprag_tool_static_query_index_ready(
result = await tool.coroutine(attachment=mock_attachment)

# Verify result
assert result == {"text": "Deep RAG analysis result"}
assert result == {"text": "Deep RAG analysis result", "citations": [], "deepRagId": deeprag_id}

# Verify ephemeral index was created
mock_uipath.context_grounding.create_ephemeral_index_async.assert_called_once()
Expand Down Expand Up @@ -228,9 +238,17 @@ async def test_create_deeprag_tool_static_query_wait_for_ingestion(
)

# First interrupt returns completed index, second returns DeepRAG result
deeprag_id = str(uuid.uuid4())
mock_interrupt.side_effect = [
mock_index_complete,
{"text": "Deep RAG analysis after waiting"},
DeepRagResponse(
id=deeprag_id,
name="test-deeprag",
created_date="2024-01-01",
last_deep_rag_status=DeepRagStatus.SUCCESSFUL,
content=DeepRagContent(text="Deep RAG analysis after waiting", citations=[]),
failure_reason=None,
),
]

mock_wrapper = Mock()
Expand All @@ -248,7 +266,7 @@ async def test_create_deeprag_tool_static_query_wait_for_ingestion(
result = await tool.coroutine(attachment=mock_attachment)

# Verify result
assert result == {"text": "Deep RAG analysis after waiting"}
assert result == {"text": "Deep RAG analysis after waiting", "citations": [], "deepRagId": deeprag_id}

# Verify interrupt was called twice (WaitEphemeralIndex + CreateDeepRag)
assert mock_interrupt.call_count == 2
Expand Down Expand Up @@ -286,8 +304,16 @@ async def test_create_deeprag_tool_dynamic_query(
)

# Index is ready → ReadyEphemeralIndex skips interrupt(). Only create_deeprag fires.
deeprag_id = str(uuid.uuid4())
mock_interrupt.side_effect = [
{"content": "Dynamic query result"},
DeepRagResponse(
id=deeprag_id,
name="test-deeprag",
created_date="2024-01-01",
last_deep_rag_status=DeepRagStatus.SUCCESSFUL,
content=DeepRagContent(text="Dynamic query result", citations=[]),
failure_reason=None,
),
]

mock_wrapper = Mock()
Expand All @@ -307,7 +333,7 @@ async def test_create_deeprag_tool_dynamic_query(
)

# Verify result
assert result == {"content": "Dynamic query result"}
assert result == {"text": "Dynamic query result", "citations": [], "deepRagId": deeprag_id}

@patch(
"uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
Expand Down Expand Up @@ -345,6 +371,162 @@ async def test_create_deeprag_tool_missing_query_dynamic(
with pytest.raises(ValueError, match="Query is required for DeepRAG tool"):
await tool.coroutine(attachment=mock_attachment)

@patch(
"uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
)
@patch("uipath_langchain.agent.tools.internal_tools.deeprag_tool.UiPath")
@patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
@patch(
"uipath_langchain.agent.tools.internal_tools.deeprag_tool.mockable",
lambda **kwargs: lambda f: f,
)
async def test_invoke_returns_error_message_on_failed_deeprag(
self,
mock_interrupt,
mock_uipath_class,
mock_get_wrapper,
resource_config_static,
mock_llm,
):
"""Test that tool returns failure_reason string when DeepRAG processing fails."""
mock_uipath = AsyncMock()
mock_uipath_class.return_value = mock_uipath

mock_index = ContextGroundingIndex(
id=str(uuid.uuid4()),
name="ephemeral-index-123",
last_ingestion_status=IndexStatus.SUCCESSFUL,
)
mock_uipath.context_grounding.create_ephemeral_index_async = AsyncMock(
return_value=mock_index
)

failed_deep_rag = DeepRagResponse(
id=str(uuid.uuid4()),
name="test-deeprag",
created_date="2024-01-01",
last_deep_rag_status=DeepRagStatus.FAILED,
content=None,
failure_reason="DeepRAG processing failed due to an internal error",
)

# Index is ready (no interrupt for index), DeepRAG fails
mock_interrupt.side_effect = [failed_deep_rag]

mock_wrapper = Mock()
mock_get_wrapper.return_value = mock_wrapper

tool = create_deeprag_tool(resource_config_static, mock_llm)

mock_attachment = MockAttachment(
ID=str(uuid.uuid4()), FullName="test.pdf", MimeType="application/pdf"
)

assert tool.coroutine is not None
result = await tool.coroutine(attachment=mock_attachment)

assert result == "DeepRAG processing failed due to an internal error"
assert mock_interrupt.call_count == 1

@patch(
"uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
)
@patch("uipath_langchain.agent.tools.internal_tools.deeprag_tool.UiPath")
@patch(
"uipath_langchain.agent.tools.internal_tools.deeprag_tool.mockable",
lambda **kwargs: lambda f: f,
)
async def test_invoke_returns_error_message_on_failed_ephemeral_index(
self,
mock_uipath_class,
mock_get_wrapper,
resource_config_static,
mock_llm,
):
"""Test that tool returns failure reason when ephemeral index fails immediately."""
mock_uipath = AsyncMock()
mock_uipath_class.return_value = mock_uipath

mock_index = ContextGroundingIndex(
id=str(uuid.uuid4()),
name="ephemeral-index-123",
last_ingestion_status=IndexStatus.FAILED,
last_ingestion_failure_reason="Ingestion failed due to unsupported file format",
)
mock_uipath.context_grounding.create_ephemeral_index_async = AsyncMock(
return_value=mock_index
)

mock_wrapper = Mock()
mock_get_wrapper.return_value = mock_wrapper

tool = create_deeprag_tool(resource_config_static, mock_llm)

mock_attachment = MockAttachment(
ID=str(uuid.uuid4()), FullName="test.pdf", MimeType="application/pdf"
)

assert tool.coroutine is not None
result = await tool.coroutine(attachment=mock_attachment)

assert result == "Ingestion failed due to unsupported file format"

@patch(
"uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
)
@patch("uipath_langchain.agent.tools.internal_tools.deeprag_tool.UiPath")
@patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
@patch(
"uipath_langchain.agent.tools.internal_tools.deeprag_tool.mockable",
lambda **kwargs: lambda f: f,
)
async def test_invoke_returns_error_message_on_failed_ephemeral_index_after_wait(
self,
mock_interrupt,
mock_uipath_class,
mock_get_wrapper,
resource_config_static,
mock_llm,
):
"""Test that tool returns failure reason when ephemeral index fails after waiting."""
mock_uipath = AsyncMock()
mock_uipath_class.return_value = mock_uipath

pending_id = str(uuid.uuid4())
mock_index_pending = ContextGroundingIndex(
id=pending_id,
name="ephemeral-index-456",
last_ingestion_status=IndexStatus.IN_PROGRESS,
)
mock_uipath.context_grounding.create_ephemeral_index_async = AsyncMock(
return_value=mock_index_pending
)

mock_index_failed = {
"id": pending_id,
"name": mock_index_pending.name,
"last_ingestion_status": "Failed",
"last_ingestion_failure_reason": "Ingestion failed during processing",
}

# First (and only) interrupt returns the failed index; DeepRAG is never reached
mock_interrupt.side_effect = [mock_index_failed]

mock_wrapper = Mock()
mock_get_wrapper.return_value = mock_wrapper

tool = create_deeprag_tool(resource_config_static, mock_llm)

mock_attachment = MockAttachment(
ID=str(uuid.uuid4()), FullName="test.pdf", MimeType="application/pdf"
)

assert tool.coroutine is not None
result = await tool.coroutine(attachment=mock_attachment)

assert result == "Ingestion failed during processing"
assert mock_interrupt.call_count == 1

@patch(
"uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper"
)
Expand Down
Loading