feat: add context agent tools

Author: radu-mocanu

pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.44"
+version = "0.1.45"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
 dependencies = [
-    "uipath>=2.2.44, <2.3.0",
+    "uipath>=2.2.45,<2.3.0",
     "langgraph>=1.0.0, <2.0.0",
     "langchain-core>=1.0.0, <2.0.0",
     "aiosqlite==0.21.0",

src/uipath_langchain/agent/react/agent.py

@@ -77,6 +77,7 @@ def create_agent(
     llm_tools: list[BaseTool] = [*agent_tools, *flow_control_tools]
 
     init_node = create_init_node(messages, input_schema)
+
     tool_nodes = create_tool_node(agent_tools)
     tool_nodes_with_guardrails = create_tools_guardrails_subgraph(
         tool_nodes, guardrails

src/uipath_langchain/agent/tools/__init__.py

@@ -1,6 +1,7 @@
 """Tool creation and management for LowCode agents."""
 
 from .context_tool import create_context_tool
+from .escalation_tool import create_escalation_tool
 from .integration_tool import create_integration_tool
 from .mcp_tool import create_mcp_tools
 from .process_tool import create_process_tool
@@ -15,6 +16,7 @@
     "create_context_tool",
     "create_process_tool",
     "create_integration_tool",
+    "create_escalation_tool",
     "create_mcp_tools",
     "UiPathToolNode",
     "ToolWrapperMixin",

src/uipath_langchain/agent/tools/context_tool.py

@@ -1,12 +1,24 @@
 """Context tool creation for semantic index retrieval."""
 
+import uuid
 from typing import Any
 
 from langchain_core.documents import Document
 from langchain_core.tools import StructuredTool
+from langgraph.types import interrupt
 from pydantic import BaseModel, Field
-from uipath.agent.models.agent import AgentContextResourceConfig
+from uipath.agent.models.agent import (
+    AgentContextResourceConfig,
+    AgentContextRetrievalMode,
+)
 from uipath.eval.mocks import mockable
+from uipath.platform.common import CreateBatchTransform, CreateDeepRag
+from uipath.platform.context_grounding import (
+    BatchTransformOutputColumn,
+    BatchTransformResponse,
+    CitationMode,
+    DeepRagResponse,
+)
 
 from uipath_langchain.retrievers import ContextGroundingRetriever
 
@@ -16,6 +28,18 @@
 
 def create_context_tool(resource: AgentContextResourceConfig) -> StructuredTool:
     tool_name = sanitize_tool_name(resource.name)
+    retrieval_mode = resource.settings.retrieval_mode.lower()
+    if retrieval_mode == AgentContextRetrievalMode.DEEP_RAG.value.lower():
+        return handle_deep_rag(tool_name, resource)
+    elif retrieval_mode == AgentContextRetrievalMode.BATCH_TRANSFORM.value.lower():
+        return handle_batch_transform(tool_name, resource)
+    else:
+        return handle_semantic_search(tool_name, resource)
+
+
+def handle_semantic_search(
+    tool_name: str, resource: AgentContextResourceConfig
+) -> StructuredTool:
     retriever = ContextGroundingRetriever(
         index_name=resource.index_name,
         folder_path=resource.folder_path,
@@ -52,3 +76,128 @@ async def context_tool_fn(query: str) -> dict[str, Any]:
         coroutine=context_tool_fn,
         output_type=output_model,
     )
+
+
+def handle_deep_rag(
+    tool_name: str, resource: AgentContextResourceConfig
+) -> StructuredTool:
+    ensure_valid_fields(resource)
+    # needed for type checking
+    assert resource.settings.query is not None
+    assert resource.settings.query.value is not None
+
+    index_name = resource.index_name
+    prompt = resource.settings.query.value
+    if not resource.settings.citation_mode:
+        raise ValueError("Citation mode is required for Deep RAG")
+    citation_mode = CitationMode(resource.settings.citation_mode.value)
+
+    input_model = None
+    output_model = DeepRagResponse
+
+    @mockable(
+        name=resource.name,
+        description=resource.description,
+        input_schema=input_model,
+        output_schema=output_model.model_json_schema(),
+        example_calls=[],  # Examples cannot be provided for context.
+    )
+    async def context_tool_fn() -> dict[str, Any]:
+        # TODO: add glob pattern support
+        return interrupt(
+            CreateDeepRag(
+                name=f"task-{uuid.uuid4()}",
+                index_name=index_name,
+                prompt=prompt,
+                citation_mode=citation_mode,
+            )
+        )
+
+    return StructuredToolWithOutputType(
+        name=tool_name,
+        description=resource.description,
+        args_schema=input_model,
+        coroutine=context_tool_fn,
+        output_type=output_model,
+    )
+
+
+def handle_batch_transform(
+    tool_name: str, resource: AgentContextResourceConfig
+) -> StructuredTool:
+    ensure_valid_fields(resource)
+
+    # needed for type checking
+    assert resource.settings.query is not None
+    assert resource.settings.query.value is not None
+
+    index_name = resource.index_name
+    prompt = resource.settings.query.value
+
+    index_folder_path = resource.folder_path
+    if not resource.settings.web_search_grounding:
+        raise ValueError("Web search grounding field is required for Batch Transform")
+    enable_web_search_grounding = (
+        resource.settings.web_search_grounding.value.lower() == "enabled"
+    )
+
+    batch_transform_output_columns: list[BatchTransformOutputColumn] = []
+    if (output_columns := resource.settings.output_columns) is None or not len(
+        output_columns
+    ):
+        raise ValueError(
+            "Batch transform requires at least one output column to be specified in settings.output_columns"
+        )
+
+    for column in output_columns:
+        batch_transform_output_columns.append(
+            BatchTransformOutputColumn(
+                name=column.name,
+                description=column.description,
+            )
+        )
+
+    class BatchTransformSchemaModel(BaseModel):
+        destination_path: str = Field(
+            ...,
+            description="The relative file path destination for the modified csv file",
+        )
+
+    input_model = BatchTransformSchemaModel
+    output_model = BatchTransformResponse
+
+    @mockable(
+        name=resource.name,
+        description=resource.description,
+        input_schema=input_model.model_json_schema(),
+        output_schema=output_model.model_json_schema(),
+        example_calls=[],  # Examples cannot be provided for context.
+    )
+    async def context_tool_fn(destination_path: str) -> dict[str, Any]:
+        # TODO: storage_bucket_folder_path_prefix  support
+        return interrupt(
+            CreateBatchTransform(
+                name=f"task-{uuid.uuid4()}",
+                index_name=index_name,
+                prompt=prompt,
+                destination_path=destination_path,
+                index_folder_path=index_folder_path,
+                enable_web_search_grounding=enable_web_search_grounding,
+                output_columns=batch_transform_output_columns,
+            )
+        )
+
+    return StructuredToolWithOutputType(
+        name=tool_name,
+        description=resource.description,
+        args_schema=input_model,
+        coroutine=context_tool_fn,
+        output_type=output_model,
+    )
+
+
+def ensure_valid_fields(resource_config: AgentContextResourceConfig):
+    if not resource_config.settings.query:
+        raise ValueError("Query object is required")
+    if not resource_config.settings.query.value:
+        raise ValueError("Query prompt is required")