Merge branch 'main' into fix-escalation-tool-calls-extraction

Author: ctiliescuuipath

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.34"
+version = "0.1.35"
 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"

src/uipath_langchain/agent/react/guardrails/guardrails_subgraph.py

@@ -1,9 +1,9 @@
 from functools import partial
-from typing import Any, Callable, Sequence
+from typing import Any, Callable, Mapping, Sequence
 
+from langgraph._internal._runnable import RunnableCallable
 from langgraph.constants import END, START
 from langgraph.graph import StateGraph
-from langgraph.prebuilt import ToolNode
 from uipath.platform.guardrails import (
     BaseGuardrail,
     BuiltInValidatorGuardrail,
@@ -221,13 +221,13 @@ def create_llm_guardrails_subgraph(
 
 
 def create_tools_guardrails_subgraph(
-    tool_nodes: dict[str, ToolNode],
+    tool_nodes: Mapping[str, RunnableCallable],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
-) -> dict[str, ToolNode]:
+) -> dict[str, RunnableCallable]:
     """Create tool nodes with guardrails.
     Args:
     """
-    result: dict[str, ToolNode] = {}
+    result: dict[str, RunnableCallable] = {}
     for tool_name, tool_node in tool_nodes.items():
         subgraph = create_tool_guardrails_subgraph(
             (tool_name, tool_node),

src/uipath_langchain/agent/tools/__init__.py

@@ -7,7 +7,7 @@
 from .tool_factory import (
     create_tools_from_resources,
 )
-from .tool_node import create_tool_node
+from .tool_node import ToolWrapperMixin, UiPathToolNode, create_tool_node
 
 __all__ = [
     "create_tools_from_resources",
@@ -16,4 +16,6 @@
     "create_process_tool",
     "create_integration_tool",
     "create_mcp_tools",
+    "UiPathToolNode",
+    "ToolWrapperMixin",
 ]

src/uipath_langchain/agent/tools/integration_tool.py

@@ -4,16 +4,21 @@
 from typing import Any
 
 from jsonschema_pydantic_converter import transform as create_model
-from langchain.tools import ToolRuntime
 from langchain_core.tools import StructuredTool
 from uipath.agent.models.agent import AgentIntegrationToolResourceConfig
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
 from uipath.platform.connections import ActivityMetadata, ActivityParameterLocationInfo
 
-from .static_args import handle_static_args
+from uipath_langchain.agent.tools.tool_node import ToolWrapperMixin
+from uipath_langchain.agent.wrappers.static_args_wrapper import get_static_args_wrapper
+
 from .structured_tool_with_output_type import StructuredToolWithOutputType
-from .utils import sanitize_tool_name
+from .utils import sanitize_dict_for_serialization, sanitize_tool_name
+
+
+class StructuredToolWithStaticArgs(StructuredToolWithOutputType, ToolWrapperMixin):
+    pass
 
 
 def remove_asterisk_from_properties(fields: dict[str, Any]) -> dict[str, Any]:
@@ -150,32 +155,27 @@ def create_integration_tool(
         input_schema=input_model.model_json_schema(),
         output_schema=output_model.model_json_schema(),
     )
-    async def integration_tool_fn(runtime: ToolRuntime, **kwargs: Any):
+    async def integration_tool_fn(**kwargs: Any):
         try:
-            # we manually validating here and not passing input_model to StructuredTool
-            # because langchain itself will block their own injected arguments (like runtime) if the model is strict
-            val_args = input_model.model_validate(kwargs)
-            args = handle_static_args(
-                resource=resource,
-                runtime=runtime,
-                input_args=val_args.model_dump(),
-            )
             result = await sdk.connections.invoke_activity_async(
                 activity_metadata=activity_metadata,
                 connection_id=connection_id,
-                activity_input=args,
+                activity_input=sanitize_dict_for_serialization(kwargs),
             )
         except Exception:
             raise
 
         return result
 
-    tool = StructuredToolWithOutputType(
+    wrapper = get_static_args_wrapper(resource)
+
+    tool = StructuredToolWithStaticArgs(
         name=tool_name,
         description=resource.description,
-        args_schema=resource.input_schema,
+        args_schema=input_model,
         coroutine=integration_tool_fn,
         output_type=output_model,
     )
+    tool.set_tool_wrappers(awrapper=wrapper)
 
     return tool

src/uipath_langchain/agent/tools/static_args.py

@@ -1,25 +1,27 @@
 """Handles static arguments for tool calls."""
 
-from typing import Any, Dict, List
+from typing import Any
 
 from jsonpath_ng import parse  # type: ignore[import-untyped]
-from langchain.tools import ToolRuntime
+from pydantic import BaseModel
 from uipath.agent.models.agent import (
     AgentIntegrationToolParameter,
     AgentIntegrationToolResourceConfig,
     BaseAgentResourceConfig,
 )
 
+from .utils import sanitize_dict_for_serialization
+
 
 def resolve_static_args(
     resource: BaseAgentResourceConfig,
-    agent_input: Dict[str, Any],
-) -> Dict[str, Any]:
+    agent_input: dict[str, Any],
+) -> dict[str, Any]:
     """Resolves static arguments for a given resource with a given input.
 
     Args:
         resource: The agent resource configuration.
-        input: Othe input arguments passed to the agent.
+        input: The input arguments passed to the agent.
 
     Returns:
         A dictionary of expanded arguments to be used in the tool call.
@@ -35,9 +37,9 @@ def resolve_static_args(
 
 
 def resolve_integration_static_args(
-    parameters: List[AgentIntegrationToolParameter],
-    agent_input: Dict[str, Any],
-) -> Dict[str, Any]:
+    parameters: list[AgentIntegrationToolParameter],
+    agent_input: dict[str, Any],
+) -> dict[str, Any]:
     """Resolves static arguments for an integration tool resource.
 
     Args:
@@ -48,7 +50,7 @@ def resolve_integration_static_args(
         A dictionary of expanded static arguments for the integration tool.
     """
 
-    static_args: Dict[str, Any] = {}
+    static_args: dict[str, Any] = {}
     for param in parameters:
         value = None
 
@@ -75,34 +77,10 @@ def resolve_integration_static_args(
     return static_args
 
 
-def sanitize_for_serialization(args: Dict[str, Any]) -> Dict[str, Any]:
-    """Convert Pydantic models in args to dicts."""
-    converted_args: Dict[str, Any] = {}
-    for key, value in args.items():
-        # handle Pydantic model
-        if hasattr(value, "model_dump"):
-            converted_args[key] = value.model_dump()
-
-        elif isinstance(value, list):
-            # handle list of Pydantic models
-            converted_list = []
-            for item in value:
-                if hasattr(item, "model_dump"):
-                    converted_list.append(item.model_dump())
-                else:
-                    converted_list.append(item)
-            converted_args[key] = converted_list
-
-        # handle regular value or unexpected type
-        else:
-            converted_args[key] = value
-    return converted_args
-
-
 def apply_static_args(
-    static_args: Dict[str, Any],
-    kwargs: Dict[str, Any],
-) -> Dict[str, Any]:
+    static_args: dict[str, Any],
+    kwargs: dict[str, Any],
+) -> dict[str, Any]:
     """Applies static arguments to the given input arguments.
 
     Args:
@@ -113,7 +91,7 @@ def apply_static_args(
         Merged input arguments with static arguments applied.
     """
 
-    sanitized_args = sanitize_for_serialization(kwargs)
+    sanitized_args = sanitize_dict_for_serialization(kwargs)
     for json_path, value in static_args.items():
         expr = parse(json_path)
         expr.update_or_create(sanitized_args, value)
@@ -122,8 +100,8 @@ def apply_static_args(
 
 
 def handle_static_args(
-    resource: BaseAgentResourceConfig, runtime: ToolRuntime, input_args: Dict[str, Any]
-) -> Dict[str, Any]:
+    resource: BaseAgentResourceConfig, state: BaseModel, input_args: dict[str, Any]
+) -> dict[str, Any]:
     """Resolves and applies static arguments for a tool call.
     Args:
         resource: The agent resource configuration.
@@ -133,6 +111,6 @@ def handle_static_args(
         A dictionary of input arguments with static arguments applied.
     """
 
-    static_args = resolve_static_args(resource, dict(runtime.state))
+    static_args = resolve_static_args(resource, dict(state))
     merged_args = apply_static_args(static_args, input_args)
     return merged_args

src/uipath_langchain/agent/tools/tool_node.py

@@ -1,22 +1,150 @@
 """Tool node factory wiring directly to LangGraph's ToolNode."""
 
 from collections.abc import Sequence
+from inspect import signature
+from typing import Any, Awaitable, Callable, Literal
 
+from langchain_core.messages.ai import AIMessage
+from langchain_core.messages.tool import ToolCall, ToolMessage
 from langchain_core.tools import BaseTool
-from langgraph.prebuilt import ToolNode
+from langgraph._internal._runnable import RunnableCallable
+from langgraph.types import Command
+from pydantic import BaseModel
 
+# the type safety can be improved with generics
+ToolWrapperType = Callable[
+    [BaseTool, ToolCall, Any], dict[str, Any] | Command[Any] | None
+]
+AsyncToolWrapperType = Callable[
+    [BaseTool, ToolCall, Any],
+    Awaitable[dict[str, Any] | Command[Any] | None],
+]
+OutputType = dict[Literal["messages"], list[ToolMessage]] | Command[Any] | None
 
-def create_tool_node(tools: Sequence[BaseTool]) -> dict[str, ToolNode]:
+
+class UiPathToolNode(RunnableCallable):
+    """
+    A ToolNode that can be used in a React agent graph.
+    It extracts the tool call from the state messages and invokes the tool.
+    It supports optional synchronous and asynchronous wrappers for custom processing.
+    Generic over the state model.
+    Args:
+        tool: The tool to invoke.
+        wrapper: An optional synchronous wrapper for custom processing.
+        awrapper: An optional asynchronous wrapper for custom processing.
+
+    Returns:
+        A dict with ToolMessage or a Command.
+    """
+
+    def __init__(
+        self,
+        tool: BaseTool,
+        wrapper: ToolWrapperType | None = None,
+        awrapper: AsyncToolWrapperType | None = None,
+    ):
+        super().__init__(func=self._func, afunc=self._afunc, name=tool.name)
+        self.tool = tool
+        self.wrapper = wrapper
+        self.awrapper = awrapper
+
+    def _func(self, state: Any) -> OutputType:
+        call = self._extract_tool_call(state)
+        if call is None:
+            return None
+        if self.wrapper:
+            filtered_state = self._filter_state(state, self.wrapper)
+            result = self.wrapper(self.tool, call, filtered_state)
+        else:
+            result = self.tool.invoke(call["args"])
+
+        return self._process_result(call, result)
+
+    async def _afunc(self, state: Any) -> OutputType:
+        call = self._extract_tool_call(state)
+        if call is None:
+            return None
+        if self.awrapper:
+            filtered_state = self._filter_state(state, self.awrapper)
+            result = await self.awrapper(self.tool, call, filtered_state)
+        else:
+            result = await self.tool.ainvoke(call["args"])
+
+        return self._process_result(call, result)
+
+    def _extract_tool_call(self, state: Any) -> ToolCall | None:
+        """Extract the tool call from the state messages."""
+
+        if not hasattr(state, "messages"):
+            raise ValueError("State does not have messages key")
+
+        last_message = state.messages[-1]
+        if not isinstance(last_message, AIMessage):
+            raise ValueError("Last message in message stack is not an AIMessage.")
+
+        for tool_call in last_message.tool_calls:
+            if tool_call["name"] == self.tool.name:
+                return tool_call
+        return None
+
+    def _process_result(
+        self, call: ToolCall, result: dict[str, Any] | Command[Any] | None
+    ) -> OutputType:
+        """Process the tool result into a message format or return a Command."""
+        if isinstance(result, Command):
+            return result
+        else:
+            message = ToolMessage(
+                content=str(result), name=call["name"], tool_call_id=call["id"]
+            )
+            return {"messages": [message]}
+
+    def _filter_state(
+        self, state: Any, wrapper: ToolWrapperType | AsyncToolWrapperType
+    ) -> BaseModel:
+        """Filter the state to the expected model type."""
+        model_type = list(signature(wrapper).parameters.values())[2].annotation
+        if not issubclass(model_type, BaseModel):
+            raise ValueError(
+                "Wrapper state parameter must be a pydantic BaseModel subclass."
+            )
+        return model_type.model_validate(state, from_attributes=True)
+
+
+class ToolWrapperMixin:
+    wrapper: ToolWrapperType | None = None
+    awrapper: AsyncToolWrapperType | None = None
+
+    def set_tool_wrappers(
+        self,
+        wrapper: ToolWrapperType | None = None,
+        awrapper: AsyncToolWrapperType | None = None,
+    ) -> None:
+        """Define wrappers for the tool execution."""
+        self.wrapper = wrapper
+        self.awrapper = awrapper
+
+
+def create_tool_node(tools: Sequence[BaseTool]) -> dict[str, UiPathToolNode]:
     """Create individual ToolNode for each tool.
 
     Args:
         tools: Sequence of tools to create nodes for.
+        agentState: The type of the agent state model.
 
     Returns:
-        Dict mapping tool.name -> ToolNode([tool]).
+        Dict mapping tool.name -> ReactToolNode([tool]).
         Each tool gets its own dedicated node for middleware composition.
 
     Note:
         handle_tool_errors=False delegates error handling to LangGraph's error boundary.
     """
-    return {tool.name: ToolNode([tool], handle_tool_errors=False) for tool in tools}
+    dict_mapping: dict[str, UiPathToolNode] = {}
+    for tool in tools:
+        if isinstance(tool, ToolWrapperMixin):
+            dict_mapping[tool.name] = UiPathToolNode(
+                tool, wrapper=tool.wrapper, awrapper=tool.awrapper
+            )
+        else:
+            dict_mapping[tool.name] = UiPathToolNode(tool, wrapper=None, awrapper=None)
+    return dict_mapping