uipath-langchain 0.1.28__py3-none-any.whl → 0.3.1__py3-none-any.whl

uipath_langchain/agent/tools/tool_node.py

@@ -1,22 +1,151 @@
 """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.runnables.config import RunnableConfig
 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, config: RunnableConfig | None = None) -> 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, config: RunnableConfig | None = None
+    ) -> 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

uipath_langchain/agent/tools/utils.py

@@ -1,6 +1,7 @@
 """Tool-related utility functions."""
 
 import re
+from typing import Any
 
 
 def sanitize_tool_name(name: str) -> str:
@@ -9,3 +10,33 @@ def sanitize_tool_name(name: str) -> str:
     sanitized_tool_name = re.sub(r"[^a-zA-Z0-9_-]", "", trim_whitespaces)
     sanitized_tool_name = sanitized_tool_name[:64]
     return sanitized_tool_name
+
+
+def sanitize_dict_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())
+                elif hasattr(item, "value"):
+                    converted_list.append(item.value)
+                else:
+                    converted_list.append(item)
+            converted_args[key] = converted_list
+
+        # handle enum-like objects with value attribute
+        elif hasattr(value, "value"):
+            converted_args[key] = value.value
+
+        # handle regular value or unexpected type
+        else:
+            converted_args[key] = value
+    return converted_args

uipath_langchain/agent/wrappers/__init__.py

@@ -0,0 +1,6 @@
+"""Wrappers to add behavior to tools while keeping them graph agnostic."""
+
+from .job_attachment_wrapper import get_job_attachment_wrapper
+from .static_args_wrapper import get_static_args_wrapper
+
+__all__ = ["get_static_args_wrapper", "get_job_attachment_wrapper"]

uipath_langchain/agent/wrappers/job_attachment_wrapper.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from langchain_core.messages.tool import ToolCall
+from langchain_core.tools import BaseTool
+from langgraph.types import Command
+from pydantic import BaseModel
+
+from uipath_langchain.agent.react.job_attachments import (
+    get_job_attachment_paths,
+    replace_job_attachment_ids,
+)
+from uipath_langchain.agent.react.types import AgentGraphState
+from uipath_langchain.agent.tools.tool_node import AsyncToolWrapperType
+
+
+def get_job_attachment_wrapper() -> AsyncToolWrapperType:
+    """Create a tool wrapper that validates and replaces job attachment IDs with full attachment objects.
+
+    This wrapper extracts job attachment paths from the tool's schema, validates that all
+    referenced attachments exist in the agent state, and replaces attachment IDs with complete
+    attachment objects before invoking the tool.
+
+    Args:
+        resource: The agent tool resource configuration
+
+    Returns:
+        An async tool wrapper function that handles job attachment validation and replacement
+    """
+
+    async def job_attachment_wrapper(
+        tool: BaseTool,
+        call: ToolCall,
+        state: AgentGraphState,
+    ) -> dict[str, Any] | Command[Any] | None:
+        """Validate and replace job attachments in tool arguments before invocation.
+
+        Args:
+            tool: The tool to wrap
+            call: The tool call containing arguments
+            state: The agent graph state containing job attachments
+
+        Returns:
+            Tool invocation result, or error dict if attachment validation fails
+        """
+        input_args = call["args"]
+        modified_input_args = input_args
+
+        if isinstance(tool.args_schema, type) and issubclass(
+            tool.args_schema, BaseModel
+        ):
+            errors: list[str] = []
+            paths = get_job_attachment_paths(tool.args_schema)
+            modified_input_args = replace_job_attachment_ids(
+                paths, input_args, state.job_attachments, errors
+            )
+
+            if errors:
+                return {"error": "\n".join(errors)}
+
+        return await tool.ainvoke(modified_input_args)
+
+    return job_attachment_wrapper

uipath_langchain/agent/wrappers/static_args_wrapper.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from langchain_core.messages.tool import ToolCall
+from langchain_core.tools import BaseTool
+from langgraph.types import Command
+from uipath.agent.models.agent import BaseAgentResourceConfig
+
+from uipath_langchain.agent.react.types import AgentGraphState
+from uipath_langchain.agent.tools.static_args import handle_static_args
+from uipath_langchain.agent.tools.tool_node import AsyncToolWrapperType
+
+
+def get_static_args_wrapper(
+    resource: BaseAgentResourceConfig,
+) -> AsyncToolWrapperType:
+    """Returns an asynchronous tool wrapper that applies static arguments.
+
+    Args:
+        resource: The agent resource configuration.
+
+    Returns:
+        An asynchronous tool wrapper function.
+    """
+
+    async def static_args_wrapper(
+        tool: BaseTool,
+        call: ToolCall,
+        state: AgentGraphState,
+    ) -> dict[str, Any] | Command[Any] | None:
+        input_args = call["args"]
+        merged_args = handle_static_args(resource, state, input_args)
+        return await tool.ainvoke(merged_args)
+
+    return static_args_wrapper

uipath_langchain/chat/__init__.py

@@ -1,10 +1,14 @@
 from .mapper import UiPathChatMessagesMapper
 from .models import UiPathAzureChatOpenAI, UiPathChat
 from .openai import UiPathChatOpenAI
+from .supported_models import BedrockModels, GeminiModels, OpenAIModels
 
 __all__ = [
     "UiPathChat",
     "UiPathAzureChatOpenAI",
     "UiPathChatOpenAI",
     "UiPathChatMessagesMapper",
+    "OpenAIModels",
+    "BedrockModels",
+    "GeminiModels",
 ]

uipath_langchain/chat/bedrock.py

@@ -48,10 +48,14 @@ class AwsBedrockCompletionsPassthroughClient:
         model: str,
         token: str,
         api_flavor: str,
+        agenthub_config: Optional[str] = None,
+        byo_connection_id: Optional[str] = None,
     ):
         self.model = model
         self.token = token
         self.api_flavor = api_flavor
+        self.agenthub_config = agenthub_config
+        self.byo_connection_id = byo_connection_id
         self._vendor = "awsbedrock"
         self._url: Optional[str] = None
 
@@ -101,6 +105,10 @@ class AwsBedrockCompletionsPassthroughClient:
             "X-UiPath-Streaming-Enabled": streaming,
         }
 
+        if self.agenthub_config:
+            headers["X-UiPath-AgentHub-Config"] = self.agenthub_config
+        if self.byo_connection_id:
+            headers["X-UiPath-LlmGateway-ByoIsConnectionId"] = self.byo_connection_id
         job_key = os.getenv("UIPATH_JOB_KEY")
         process_key = os.getenv("UIPATH_PROCESS_KEY")
         if job_key:
@@ -118,6 +126,8 @@ class UiPathChatBedrockConverse(ChatBedrockConverse):
         tenant_id: Optional[str] = None,
         token: Optional[str] = None,
         model_name: str = BedrockModels.anthropic_claude_haiku_4_5,
+        agenthub_config: Optional[str] = None,
+        byo_connection_id: Optional[str] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -141,6 +151,8 @@ class UiPathChatBedrockConverse(ChatBedrockConverse):
             model=model_name,
             token=token,
             api_flavor="converse",
+            agenthub_config=agenthub_config,
+            byo_connection_id=byo_connection_id,
         )
 
         client = passthrough_client.get_client()
@@ -156,6 +168,8 @@ class UiPathChatBedrock(ChatBedrock):
         tenant_id: Optional[str] = None,
         token: Optional[str] = None,
         model_name: str = BedrockModels.anthropic_claude_haiku_4_5,
+        agenthub_config: Optional[str] = None,
+        byo_connection_id: Optional[str] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -179,6 +193,8 @@ class UiPathChatBedrock(ChatBedrock):
             model=model_name,
             token=token,
             api_flavor="invoke",
+            agenthub_config=agenthub_config,
+            byo_connection_id=byo_connection_id,
         )
 
         client = passthrough_client.get_client()

uipath_langchain/chat/mapper.py

@@ -41,6 +41,7 @@ class UiPathChatMessagesMapper:
     def __init__(self):
         """Initialize the mapper with empty state."""
         self.tool_call_to_ai_message: dict[str, str] = {}
+        self.current_message: AIMessageChunk
         self.seen_message_ids: set[str] = set()
 
     def _extract_text(self, content: Any) -> str:
@@ -141,7 +142,7 @@ class UiPathChatMessagesMapper:
     def map_event(
         self,
         message: BaseMessage,
-    ) -> UiPathConversationMessageEvent | None:
+    ) -> list[UiPathConversationMessageEvent] | None:
         """Convert LangGraph BaseMessage (chunk or full) into a UiPathConversationMessageEvent.
 
         Args:
@@ -168,16 +169,45 @@ class UiPathChatMessagesMapper:
 
             # Check if this is the last chunk by examining chunk_position
             if message.chunk_position == "last":
+                events: list[UiPathConversationMessageEvent] = []
+
+                # Loop through all content_blocks in current_message and create toolCallStart events for each tool_call_chunk
+                if self.current_message and self.current_message.content_blocks:
+                    for block in self.current_message.content_blocks:
+                        if block.get("type") == "tool_call_chunk":
+                            tool_chunk_block = cast(ToolCallChunk, block)
+                            tool_call_id = tool_chunk_block.get("id")
+                            tool_name = tool_chunk_block.get("name")
+                            tool_args = tool_chunk_block.get("args")
+
+                            if tool_call_id:
+                                tool_event = UiPathConversationMessageEvent(
+                                    message_id=message.id,
+                                    tool_call=UiPathConversationToolCallEvent(
+                                        tool_call_id=tool_call_id,
+                                        start=UiPathConversationToolCallStartEvent(
+                                            tool_name=tool_name,
+                                            timestamp=timestamp,
+                                            input=UiPathInlineValue(inline=tool_args),
+                                        ),
+                                    ),
+                                )
+                                events.append(tool_event)
+
+                # Create the final event for the message
                 msg_event.end = UiPathConversationMessageEndEvent(timestamp=timestamp)
                 msg_event.content_part = UiPathConversationContentPartEvent(
                     content_part_id=f"chunk-{message.id}-0",
                     end=UiPathConversationContentPartEndEvent(),
                 )
-                return msg_event
+                events.append(msg_event)
+
+                return events
 
             # For every new message_id, start a new message
             if message.id not in self.seen_message_ids:
                 self.seen_message_ids.add(message.id)
+                self.current_message = message
                 msg_event.start = UiPathConversationMessageStartEvent(
                     role="assistant", timestamp=timestamp
                 )
@@ -200,7 +230,6 @@ class UiPathChatMessagesMapper:
                             content_part_id=f"chunk-{message.id}-0",
                             chunk=UiPathConversationContentPartChunkEvent(
                                 data=text,
-                                content_part_sequence=0,
                             ),
                         )
 
@@ -210,19 +239,10 @@ class UiPathChatMessagesMapper:
                         tool_call_id = tool_chunk_block.get("id")
                         if tool_call_id:
                             # Track tool_call_id -> ai_message_id mapping
-                            self.tool_call_to_ai_message[str(tool_call_id)] = message.id
-
-                        args = tool_chunk_block.get("args") or ""
+                            self.tool_call_to_ai_message[tool_call_id] = message.id
 
-                        msg_event.content_part = UiPathConversationContentPartEvent(
-                            content_part_id=f"chunk-{message.id}-0",
-                            chunk=UiPathConversationContentPartChunkEvent(
-                                data=args,
-                                content_part_sequence=0,
-                            ),
-                        )
-                        # Continue so that multiple tool_call_chunks in the same block list
-                        # are handled correctly
+                        # Accumulate the message chunk
+                        self.current_message = self.current_message + message
                         continue
 
             # Fallback: raw string content on the chunk (rare when using content_blocks)
@@ -231,7 +251,6 @@ class UiPathChatMessagesMapper:
                     content_part_id=f"content-{message.id}",
                     chunk=UiPathConversationContentPartChunkEvent(
                         data=message.content,
-                        content_part_sequence=0,
                     ),
                 )
 
@@ -241,7 +260,7 @@ class UiPathChatMessagesMapper:
                 or msg_event.tool_call
                 or msg_event.end
             ):
-                return msg_event
+                return [msg_event]
 
             return None
 
@@ -275,35 +294,34 @@ class UiPathChatMessagesMapper:
                     # Keep as string if not valid JSON
                     pass
 
-            return UiPathConversationMessageEvent(
-                message_id=result_message_id or str(uuid4()),
-                tool_call=UiPathConversationToolCallEvent(
-                    tool_call_id=message.tool_call_id,
-                    start=UiPathConversationToolCallStartEvent(
-                        tool_name=message.name,
-                        arguments=None,
-                        timestamp=timestamp,
+            return [
+                UiPathConversationMessageEvent(
+                    message_id=result_message_id or str(uuid4()),
+                    tool_call=UiPathConversationToolCallEvent(
+                        tool_call_id=message.tool_call_id,
+                        end=UiPathConversationToolCallEndEvent(
+                            timestamp=timestamp,
+                            output=UiPathInlineValue(inline=content_value),
+                        ),
                     ),
-                    end=UiPathConversationToolCallEndEvent(
-                        timestamp=timestamp,
-                        output=UiPathInlineValue(inline=content_value),
-                    ),
-                ),
-            )
+                )
+            ]
 
         # --- Fallback for other BaseMessage types ---
         text_content = self._extract_text(message.content)
-        return UiPathConversationMessageEvent(
-            message_id=message.id,
-            start=UiPathConversationMessageStartEvent(
-                role="assistant", timestamp=timestamp
-            ),
-            content_part=UiPathConversationContentPartEvent(
-                content_part_id=f"cp-{message.id}",
-                chunk=UiPathConversationContentPartChunkEvent(data=text_content),
-            ),
-            end=UiPathConversationMessageEndEvent(),
-        )
+        return [
+            UiPathConversationMessageEvent(
+                message_id=message.id,
+                start=UiPathConversationMessageStartEvent(
+                    role="assistant", timestamp=timestamp
+                ),
+                content_part=UiPathConversationContentPartEvent(
+                    content_part_id=f"cp-{message.id}",
+                    chunk=UiPathConversationContentPartChunkEvent(data=text_content),
+                ),
+                end=UiPathConversationMessageEndEvent(),
+            )
+        ]
 
 
 __all__ = ["UiPathChatMessagesMapper"]