uipath-langchain 0.1.24__py3-none-any.whl → 0.1.34__py3-none-any.whl

uipath_langchain/agent/react/agent.py

@@ -9,9 +9,14 @@ from langgraph.graph import StateGraph
 from pydantic import BaseModel
 from uipath.platform.guardrails import BaseGuardrail
 
-from ..guardrails import create_llm_guardrails_subgraph
 from ..guardrails.actions import GuardrailAction
 from ..tools import create_tool_node
+from .guardrails.guardrails_subgraph import (
+    create_agent_init_guardrails_subgraph,
+    create_agent_terminate_guardrails_subgraph,
+    create_llm_guardrails_subgraph,
+    create_tools_guardrails_subgraph,
+)
 from .init_node import (
     create_init_node,
 )
@@ -19,7 +24,7 @@ from .llm_node import (
     create_llm_node,
 )
 from .router import (
-    route_agent,
+    create_route_agent,
 )
 from .terminate_node import (
     create_terminate_node,
@@ -53,7 +58,7 @@ def create_agent(
     config: AgentGraphConfig | None = None,
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None = None,
 ) -> StateGraph[AgentGraphState, None, InputT, OutputT]:
-    """Build agent graph with INIT -> AGENT(subgraph) <-> TOOLS loop, terminated by control flow tools.
+    """Build agent graph with INIT -> AGENT (subgraph) <-> TOOLS loop, terminated by control flow tools.
 
     The AGENT node is a subgraph that runs:
     - before-agent guardrail middlewares
@@ -73,6 +78,9 @@ def create_agent(
 
     init_node = create_init_node(messages)
     tool_nodes = create_tool_node(agent_tools)
+    tool_nodes_with_guardrails = create_tools_guardrails_subgraph(
+        tool_nodes, guardrails
+    )
     terminate_node = create_terminate_node(output_schema)
 
     InnerAgentGraphState = create_state_with_input(
@@ -82,23 +90,32 @@ def create_agent(
     builder: StateGraph[AgentGraphState, None, InputT, OutputT] = StateGraph(
         InnerAgentGraphState, input_schema=input_schema, output_schema=output_schema
     )
-    builder.add_node(AgentGraphNode.INIT, init_node)
+    init_with_guardrails_subgraph = create_agent_init_guardrails_subgraph(
+        (AgentGraphNode.GUARDED_INIT, init_node),
+        guardrails,
+    )
+    builder.add_node(AgentGraphNode.INIT, init_with_guardrails_subgraph)
 
-    for tool_name, tool_node in tool_nodes.items():
+    for tool_name, tool_node in tool_nodes_with_guardrails.items():
         builder.add_node(tool_name, tool_node)
 
-    builder.add_node(AgentGraphNode.TERMINATE, terminate_node)
+    terminate_with_guardrails_subgraph = create_agent_terminate_guardrails_subgraph(
+        (AgentGraphNode.GUARDED_TERMINATE, terminate_node),
+        guardrails,
+    )
+    builder.add_node(AgentGraphNode.TERMINATE, terminate_with_guardrails_subgraph)
 
     builder.add_edge(START, AgentGraphNode.INIT)
 
-    llm_node = create_llm_node(model, llm_tools)
+    llm_node = create_llm_node(model, llm_tools, config.thinking_messages_limit)
     llm_with_guardrails_subgraph = create_llm_guardrails_subgraph(
         (AgentGraphNode.LLM, llm_node), guardrails
     )
     builder.add_node(AgentGraphNode.AGENT, llm_with_guardrails_subgraph)
     builder.add_edge(AgentGraphNode.INIT, AgentGraphNode.AGENT)
 
-    tool_node_names = list(tool_nodes.keys())
+    tool_node_names = list(tool_nodes_with_guardrails.keys())
+    route_agent = create_route_agent(config.thinking_messages_limit)
     builder.add_conditional_edges(
         AgentGraphNode.AGENT,
         route_agent,

uipath_langchain/agent/react/constants.py

@@ -1,2 +1 @@
-# Agent routing configuration
-MAX_SUCCESSIVE_COMPLETIONS = 1
+MAX_CONSECUTIVE_THINKING_MESSAGES = 0

uipath_langchain/agent/{guardrails → react/guardrails}/guardrails_subgraph.py

@@ -1,22 +1,30 @@
+from functools import partial
 from typing import Any, Callable, Sequence
 
 from langgraph.constants import END, START
 from langgraph.graph import StateGraph
+from langgraph.prebuilt import ToolNode
 from uipath.platform.guardrails import (
     BaseGuardrail,
     BuiltInValidatorGuardrail,
     GuardrailScope,
 )
 
-from uipath_langchain.agent.guardrails.types import ExecutionStage
-
-from .actions.base_action import GuardrailAction, GuardrailActionNode
-from .guardrail_nodes import (
-    create_agent_guardrail_node,
+from uipath_langchain.agent.guardrails.actions.base_action import (
+    GuardrailAction,
+    GuardrailActionNode,
+)
+from uipath_langchain.agent.guardrails.guardrail_nodes import (
+    create_agent_init_guardrail_node,
+    create_agent_terminate_guardrail_node,
     create_llm_guardrail_node,
     create_tool_guardrail_node,
 )
-from .types import AgentGuardrailsGraphState
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+from uipath_langchain.agent.react.types import (
+    AgentGraphState,
+    AgentGuardrailsGraphState,
+)
 
 _VALIDATOR_ALLOWED_STAGES = {
     "prompt_injection": {ExecutionStage.PRE_EXECUTION},
@@ -90,6 +98,7 @@ def _create_guardrails_subgraph(
             ExecutionStage.PRE_EXECUTION,
             node_factory,
             inner_name,
+            inner_name,
         )
         subgraph.add_edge(START, first_pre_exec_guardrail_node)
     else:
@@ -107,6 +116,7 @@ def _create_guardrails_subgraph(
             ExecutionStage.POST_EXECUTION,
             node_factory,
             END,
+            inner_node,
         )
         subgraph.add_edge(inner_name, first_post_exec_guardrail_node)
     else:
@@ -130,6 +140,7 @@ def _build_guardrail_node_chain(
         GuardrailActionNode,
     ],
     next_node: str,
+    guarded_node_name: str,
 ) -> str:
     """Recursively build a chain of guardrail nodes in reverse order.
 
@@ -157,7 +168,10 @@ def _build_guardrail_node_chain(
     remaining_guardrails = guardrails[:-1]
 
     fail_node_name, fail_node = action.action_node(
-        guardrail=guardrail, scope=scope, execution_stage=execution_stage
+        guardrail=guardrail,
+        scope=scope,
+        execution_stage=execution_stage,
+        guarded_component_name=guarded_node_name,
     )
 
     # Create the guardrail evaluation node.
@@ -179,6 +193,7 @@ def _build_guardrail_node_chain(
         execution_stage,
         node_factory,
         guardrail_node_name,
+        guarded_node_name,
     )
 
     return previous_node_name
@@ -193,6 +208,9 @@ def create_llm_guardrails_subgraph(
         for (guardrail, _) in (guardrails or [])
         if GuardrailScope.LLM in guardrail.selector.scopes
     ]
+    if applicable_guardrails is None or len(applicable_guardrails) == 0:
+        return llm_node[1]
+
     return _create_guardrails_subgraph(
         main_inner_node=llm_node,
         guardrails=applicable_guardrails,
@@ -202,30 +220,84 @@ def create_llm_guardrails_subgraph(
     )
 
 
-def create_agent_guardrails_subgraph(
-    agent_node: tuple[str, Any],
+def create_tools_guardrails_subgraph(
+    tool_nodes: dict[str, ToolNode],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
-    execution_stage: ExecutionStage,
-):
-    """Create a subgraph for AGENT-scoped guardrails that applies checks at the specified stage.
-
-    This is intended for wrapping nodes like INIT or TERMINATE, where guardrails should run
-    either before (pre-execution) or after (post-execution) the node logic.
+) -> dict[str, ToolNode]:
+    """Create tool nodes with guardrails.
+    Args:
     """
+    result: dict[str, ToolNode] = {}
+    for tool_name, tool_node in tool_nodes.items():
+        subgraph = create_tool_guardrails_subgraph(
+            (tool_name, tool_node),
+            guardrails,
+        )
+        result[tool_name] = subgraph
+
+    return result
+
+
+def create_agent_init_guardrails_subgraph(
+    init_node: tuple[str, Any],
+    guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+):
+    """Create a subgraph for INIT node that applies guardrails on the state messages."""
     applicable_guardrails = [
         (guardrail, _)
         for (guardrail, _) in (guardrails or [])
         if GuardrailScope.AGENT in guardrail.selector.scopes
     ]
+    if applicable_guardrails is None or len(applicable_guardrails) == 0:
+        return init_node[1]
+
     return _create_guardrails_subgraph(
-        main_inner_node=agent_node,
+        main_inner_node=init_node,
         guardrails=applicable_guardrails,
         scope=GuardrailScope.AGENT,
-        execution_stages=[execution_stage],
-        node_factory=create_agent_guardrail_node,
+        execution_stages=[ExecutionStage.POST_EXECUTION],
+        node_factory=create_agent_init_guardrail_node,
     )
 
 
+def create_agent_terminate_guardrails_subgraph(
+    terminate_node: tuple[str, Any],
+    guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+):
+    """Create a subgraph for TERMINATE node that applies guardrails on the agent result."""
+    node_name, node_func = terminate_node
+
+    def terminate_wrapper(state: Any) -> dict[str, Any]:
+        # Call original terminate node
+        result = node_func(state)
+        # Store result in state
+        return {"agent_result": result, "messages": state.messages}
+
+    applicable_guardrails = [
+        (guardrail, _)
+        for (guardrail, _) in (guardrails or [])
+        if GuardrailScope.AGENT in guardrail.selector.scopes
+    ]
+    if applicable_guardrails is None or len(applicable_guardrails) == 0:
+        return terminate_node[1]
+
+    subgraph = _create_guardrails_subgraph(
+        main_inner_node=(node_name, terminate_wrapper),
+        guardrails=applicable_guardrails,
+        scope=GuardrailScope.AGENT,
+        execution_stages=[ExecutionStage.POST_EXECUTION],
+        node_factory=create_agent_terminate_guardrail_node,
+    )
+
+    async def run_terminate_subgraph(
+        state: AgentGraphState,
+    ) -> dict[str, Any]:
+        result_state = await subgraph.ainvoke(state)
+        return result_state["agent_result"]
+
+    return run_terminate_subgraph
+
+
 def create_tool_guardrails_subgraph(
     tool_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
@@ -238,10 +310,13 @@ def create_tool_guardrails_subgraph(
         and guardrail.selector.match_names is not None
         and tool_name in guardrail.selector.match_names
     ]
+    if applicable_guardrails is None or len(applicable_guardrails) == 0:
+        return tool_node[1]
+
     return _create_guardrails_subgraph(
         main_inner_node=tool_node,
         guardrails=applicable_guardrails,
         scope=GuardrailScope.TOOL,
         execution_stages=[ExecutionStage.PRE_EXECUTION, ExecutionStage.POST_EXECUTION],
-        node_factory=create_tool_guardrail_node,
+        node_factory=partial(create_tool_guardrail_node, tool_name=tool_name),
     )

uipath_langchain/agent/react/llm_node.py

@@ -1,34 +1,65 @@
-"""LLM node implementation for LangGraph."""
+"""LLM node for ReAct Agent graph."""
 
-from typing import Sequence
+from typing import Literal, Sequence
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage
 from langchain_core.tools import BaseTool
 
-from .constants import MAX_SUCCESSIVE_COMPLETIONS
+from .constants import MAX_CONSECUTIVE_THINKING_MESSAGES
 from .types import AgentGraphState
-from .utils import count_successive_completions
+from .utils import count_consecutive_thinking_messages
+
+OPENAI_COMPATIBLE_CHAT_MODELS = (
+    "UiPathChatOpenAI",
+    "AzureChatOpenAI",
+    "ChatOpenAI",
+    "UiPathChat",
+    "UiPathAzureChatOpenAI",
+)
+
+
+def _get_required_tool_choice_by_model(
+    model: BaseChatModel,
+) -> Literal["required", "any"]:
+    """Get the appropriate tool_choice value to enforce tool usage based on model type.
+
+    "required" - OpenAI compatible required tool_choice value
+    "any" - Vertex and Bedrock parameter for required tool_choice value
+    """
+    model_class_name = model.__class__.__name__
+    if model_class_name in OPENAI_COMPATIBLE_CHAT_MODELS:
+        return "required"
+    return "any"
 
 
 def create_llm_node(
     model: BaseChatModel,
     tools: Sequence[BaseTool] | None = None,
+    thinking_messages_limit: int = MAX_CONSECUTIVE_THINKING_MESSAGES,
 ):
-    """Invoke LLM with tools and dynamically control tool_choice based on successive completions.
+    """Create LLM node with dynamic tool_choice enforcement.
 
-    When successive completions reach the limit, tool_choice is set to "required" to force
-    the LLM to use a tool and prevent infinite reasoning loops.
+    Controls when to force tool usage based on consecutive thinking steps
+    to prevent infinite loops and ensure progress.
+
+    Args:
+        model: The chat model to use
+        tools: Available tools to bind
+        thinking_messages_limit: Max consecutive LLM responses without tool calls
+            before enforcing tool usage. 0 = force tools every time.
     """
     bindable_tools = list(tools) if tools else []
     base_llm = model.bind_tools(bindable_tools) if bindable_tools else model
+    tool_choice_required_value = _get_required_tool_choice_by_model(model)
 
     async def llm_node(state: AgentGraphState):
         messages: list[AnyMessage] = state.messages
 
-        successive_completions = count_successive_completions(messages)
-        if successive_completions >= MAX_SUCCESSIVE_COMPLETIONS:
-            llm = base_llm.bind(tool_choice="required")
+        consecutive_thinking_messages = count_consecutive_thinking_messages(messages)
+
+        if bindable_tools and consecutive_thinking_messages >= thinking_messages_limit:
+            llm = base_llm.bind(tool_choice=tool_choice_required_value)
         else:
             llm = base_llm
 

uipath_langchain/agent/react/router.py

@@ -6,9 +6,8 @@ from langchain_core.messages import AIMessage, AnyMessage, ToolCall
 from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
 
 from ..exceptions import AgentNodeRoutingException
-from .constants import MAX_SUCCESSIVE_COMPLETIONS
 from .types import AgentGraphNode, AgentGraphState
-from .utils import count_successive_completions
+from .utils import count_consecutive_thinking_messages
 
 FLOW_CONTROL_TOOLS = [END_EXECUTION_TOOL.name, RAISE_ERROR_TOOL.name]
 
@@ -48,50 +47,62 @@ def __validate_last_message_is_AI(messages: list[AnyMessage]) -> AIMessage:
     return last_message
 
 
-def route_agent(
-    state: AgentGraphState,
-) -> list[str] | Literal[AgentGraphNode.AGENT, AgentGraphNode.TERMINATE]:
-    """Route after agent: handles all routing logic including control flow detection.
+def create_route_agent(thinking_messages_limit: int = 0):
+    """Create a routing function configured with thinking_messages_limit.
 
-    Routing logic:
-    1. If multiple tool calls exist, filter out control flow tools (EndExecution, RaiseError)
-    2. If control flow tool(s) remain, route to TERMINATE
-    3. If regular tool calls remain, route to specific tool nodes (return list of tool names)
-    4. If no tool calls, handle successive completions
+    Args:
+        thinking_messages_limit: Max consecutive thinking messages before error
 
     Returns:
-        - list[str]: Tool node names for parallel execution
-        - AgentGraphNode.AGENT: For successive completions
-        - AgentGraphNode.TERMINATE: For control flow termination
-
-    Raises:
-        AgentNodeRoutingException: When encountering unexpected state (empty messages, non-AIMessage, or excessive completions)
+        Routing function for LangGraph conditional edges
     """
-    messages = state.messages
-    last_message = __validate_last_message_is_AI(messages)
 
-    tool_calls = list(last_message.tool_calls) if last_message.tool_calls else []
-    tool_calls = __filter_control_flow_tool_calls(tool_calls)
+    def route_agent(
+        state: AgentGraphState,
+    ) -> list[str] | Literal[AgentGraphNode.AGENT, AgentGraphNode.TERMINATE]:
+        """Route after agent: handles all routing logic including control flow detection.
+
+        Routing logic:
+        1. If multiple tool calls exist, filter out control flow tools (EndExecution, RaiseError)
+        2. If control flow tool(s) remain, route to TERMINATE
+        3. If regular tool calls remain, route to specific tool nodes (return list of tool names)
+        4. If no tool calls, handle consecutive completions
+
+        Returns:
+            - list[str]: Tool node names for parallel execution
+            - AgentGraphNode.AGENT: For consecutive completions
+            - AgentGraphNode.TERMINATE: For control flow termination
+
+        Raises:
+            AgentNodeRoutingException: When encountering unexpected state (empty messages, non-AIMessage, or excessive completions)
+        """
+        messages = state.messages
+        last_message = __validate_last_message_is_AI(messages)
 
-    if tool_calls and __has_control_flow_tool(tool_calls):
-        return AgentGraphNode.TERMINATE
+        tool_calls = list(last_message.tool_calls) if last_message.tool_calls else []
+        tool_calls = __filter_control_flow_tool_calls(tool_calls)
 
-    if tool_calls:
-        return [tc["name"] for tc in tool_calls]
+        if tool_calls and __has_control_flow_tool(tool_calls):
+            return AgentGraphNode.TERMINATE
 
-    successive_completions = count_successive_completions(messages)
+        if tool_calls:
+            return [tc["name"] for tc in tool_calls]
+
+        consecutive_thinking_messages = count_consecutive_thinking_messages(messages)
+
+        if consecutive_thinking_messages > thinking_messages_limit:
+            raise AgentNodeRoutingException(
+                f"Agent exceeded consecutive completions limit without producing tool calls "
+                f"(completions: {consecutive_thinking_messages}, max: {thinking_messages_limit}). "
+                f"This should not happen as tool_choice='required' is enforced at the limit."
+            )
+
+        if last_message.content:
+            return AgentGraphNode.AGENT
 
-    if successive_completions > MAX_SUCCESSIVE_COMPLETIONS:
         raise AgentNodeRoutingException(
-            f"Agent exceeded successive completions limit without producing tool calls "
-            f"(completions: {successive_completions}, max: {MAX_SUCCESSIVE_COMPLETIONS}). "
-            f"This should not happen as tool_choice='required' is enforced at the limit."
+            f"Agent produced empty response without tool calls "
+            f"(completions: {consecutive_thinking_messages}, has_content: False)"
         )
 
-    if last_message.content:
-        return AgentGraphNode.AGENT
-
-    raise AgentNodeRoutingException(
-        f"Agent produced empty response without tool calls "
-        f"(completions: {successive_completions}, has_content: False)"
-    )
+    return route_agent

uipath_langchain/agent/react/types.py

@@ -1,5 +1,5 @@
 from enum import StrEnum
-from typing import Annotated
+from typing import Annotated, Any, Optional
 
 from langchain_core.messages import AnyMessage
 from langgraph.graph.message import add_messages
@@ -25,15 +25,29 @@ class AgentGraphState(BaseModel):
     termination: AgentTermination | None = None
 
 
+class AgentGuardrailsGraphState(AgentGraphState):
+    """Agent Guardrails Graph state for guardrail subgraph."""
+
+    guardrail_validation_result: Optional[str] = None
+    agent_result: Optional[dict[str, Any]] = None
+
+
 class AgentGraphNode(StrEnum):
     INIT = "init"
+    GUARDED_INIT = "guarded-init"
     AGENT = "agent"
     LLM = "llm"
     TOOLS = "tools"
     TERMINATE = "terminate"
+    GUARDED_TERMINATE = "guarded-terminate"
 
 
 class AgentGraphConfig(BaseModel):
     recursion_limit: int = Field(
         default=50, ge=1, description="Maximum recursion limit for the agent graph"
     )
+    thinking_messages_limit: int = Field(
+        default=0,
+        ge=0,
+        description="Max consecutive thinking messages before enforcing tool usage. 0 = force tools every time.",
+    )

uipath_langchain/agent/react/utils.py

@@ -28,7 +28,7 @@ def resolve_output_model(
     return END_EXECUTION_TOOL.args_schema
 
 
-def count_successive_completions(messages: Sequence[BaseMessage]) -> int:
+def count_consecutive_thinking_messages(messages: Sequence[BaseMessage]) -> int:
     """Count consecutive AIMessages without tool calls at end of message history."""
     if not messages:
         return 0

uipath_langchain/agent/tools/__init__.py

@@ -2,6 +2,7 @@
 
 from .context_tool import create_context_tool
 from .integration_tool import create_integration_tool
+from .mcp_tool import create_mcp_tools
 from .process_tool import create_process_tool
 from .tool_factory import (
     create_tools_from_resources,
@@ -14,4 +15,5 @@ __all__ = [
     "create_context_tool",
     "create_process_tool",
     "create_integration_tool",
+    "create_mcp_tools",
 ]

uipath_langchain/agent/tools/mcp_tool.py

@@ -0,0 +1,86 @@
+import asyncio
+import os
+from collections import Counter, defaultdict
+from contextlib import AsyncExitStack, asynccontextmanager
+from itertools import chain
+
+import httpx
+from langchain_core.tools import BaseTool
+from langchain_mcp_adapters.tools import load_mcp_tools
+from mcp import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+from uipath._utils._ssl_context import get_httpx_client_kwargs
+from uipath.agent.models.agent import AgentMcpResourceConfig
+
+
+def _deduplicate_tools(tools: list[BaseTool]) -> list[BaseTool]:
+    """Deduplicate tools by appending numeric suffix to duplicate names."""
+    counts = Counter(tool.name for tool in tools)
+    seen: defaultdict[str, int] = defaultdict(int)
+
+    for tool in tools:
+        if counts[tool.name] > 1:
+            seen[tool.name] += 1
+            tool.name = f"{tool.name}_{seen[tool.name]}"
+
+    return tools
+
+
+def _filter_tools(tools: list[BaseTool], cfg: AgentMcpResourceConfig) -> list[BaseTool]:
+    """Filter tools to only include those in available_tools."""
+    allowed = {t.name for t in cfg.available_tools}
+    return [t for t in tools if t.name in allowed]
+
+
+@asynccontextmanager
+async def create_mcp_tools(
+    config: AgentMcpResourceConfig | list[AgentMcpResourceConfig],
+    max_concurrency: int = 5,
+):
+    """Connect to UiPath MCP server(s) and yield LangChain-compatible tools."""
+    if not (base_url := os.getenv("UIPATH_URL")):
+        raise ValueError("UIPATH_URL environment variable is not set")
+    if not (access_token := os.getenv("UIPATH_ACCESS_TOKEN")):
+        raise ValueError("UIPATH_ACCESS_TOKEN environment variable is not set")
+
+    configs = config if isinstance(config, list) else [config]
+    enabled = [c for c in configs if c.is_enabled is not False]
+
+    if not enabled:
+        yield []
+        return
+
+    base_url = base_url.rstrip("/")
+    semaphore = asyncio.Semaphore(max_concurrency)
+
+    default_client_kwargs = get_httpx_client_kwargs()
+    client_kwargs = {
+        **default_client_kwargs,
+        "headers": {"Authorization": f"Bearer {access_token}"},
+        "timeout": httpx.Timeout(60),
+    }
+
+    async def init_session(
+        session: ClientSession, cfg: AgentMcpResourceConfig
+    ) -> list[BaseTool]:
+        async with semaphore:
+            await session.initialize()
+            tools = await load_mcp_tools(session)
+            return _filter_tools(tools, cfg)
+
+    async def create_session(
+        stack: AsyncExitStack, cfg: AgentMcpResourceConfig
+    ) -> ClientSession:
+        url = f"{base_url}/agenthub_/mcp/{cfg.folder_path}/{cfg.slug}"
+        http_client = await stack.enter_async_context(
+            httpx.AsyncClient(**client_kwargs)
+        )
+        read, write, _ = await stack.enter_async_context(
+            streamable_http_client(url=url, http_client=http_client)
+        )
+        return await stack.enter_async_context(ClientSession(read, write))
+
+    async with AsyncExitStack() as stack:
+        sessions = [(await create_session(stack, cfg), cfg) for cfg in enabled]
+        results = await asyncio.gather(*[init_session(s, cfg) for s, cfg in sessions])
+        yield _deduplicate_tools(list(chain.from_iterable(results)))

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",
 ]