feat: add static args support (#431)

Author: andreiancuta-uipath

src/uipath_langchain/_utils/_jsonpath.py

@@ -0,0 +1,40 @@
+"""JSONPath utilities."""
+
+from jsonpath_ng import (  # type: ignore[import-untyped]
+    Child,
+    Fields,
+    JSONPath,
+    Root,
+    Slice,
+    parse,
+)
+
+
+def parse_jsonpath_segments(json_path: str) -> list[str]:
+    """Parse JSON path $['a']['b'] into list of segments ['a', 'b']."""
+    jsonpath_expr = parse(json_path)
+    return _extract_segments(jsonpath_expr)
+
+
+def _extract_segments(expr: JSONPath) -> list[str]:
+    """Recursively extract segments from a JSONPath expression tree.
+
+    Args:
+        expr: The JSONPath expression node (Root, Child, Fields, etc.)
+    """
+    parts: list[str] = []
+
+    match expr:
+        case Root():
+            pass
+        case Fields():
+            # Leaf node, add it to the path
+            parts.extend(expr.fields)
+        case Child():
+            # Child node, walk left then right to maintain order
+            parts.extend(_extract_segments(expr.left))
+            parts.extend(_extract_segments(expr.right))
+        case Slice(start=None, end=None, step=None):
+            parts.append("*")
+
+    return parts

src/uipath_langchain/agent/react/agent.py

@@ -108,9 +108,10 @@ def create_agent(
     llm_node = create_llm_node(
         model,
         llm_tools,
-        config.is_conversational,
-        config.llm_messages_limit,
-        config.thinking_messages_limit,
+        input_schema=input_schema,
+        is_conversational=config.is_conversational,
+        llm_messages_limit=config.llm_messages_limit,
+        thinking_messages_limit=config.thinking_messages_limit,
     )
     llm_with_guardrails_subgraph = create_llm_guardrails_subgraph(
         (AgentGraphNode.LLM, llm_node), guardrails, input_schema=input_schema

src/uipath_langchain/agent/react/llm_node.py

@@ -1,19 +1,27 @@
 """LLM node for ReAct Agent graph."""
 
-from typing import Literal, Sequence
+from typing import Literal, Sequence, TypeVar
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage, ToolCall
 from langchain_core.tools import BaseTool
+from pydantic import BaseModel
 from uipath.runtime.errors import UiPathErrorCategory, UiPathErrorCode
 
+from uipath_langchain.agent.tools.static_args import (
+    apply_static_argument_properties_to_schema,
+)
+from uipath_langchain.agent.tools.structured_tool_with_argument_properties import (
+    StructuredToolWithArgumentProperties,
+)
+
 from ..exceptions import AgentTerminationException
 from .constants import (
     DEFAULT_MAX_CONSECUTIVE_THINKING_MESSAGES,
     DEFAULT_MAX_LLM_MESSAGES,
 )
 from .types import FLOW_CONTROL_TOOLS, AgentGraphState
-from .utils import count_consecutive_thinking_messages
+from .utils import count_consecutive_thinking_messages, extract_input_data_from_state
 
 OPENAI_COMPATIBLE_CHAT_MODELS = (
     "UiPathChatOpenAI",
@@ -48,9 +56,14 @@ def _filter_control_flow_tool_calls(
     return [tc for tc in tool_calls if tc.get("name") not in FLOW_CONTROL_TOOLS]
 
 
+StateT = TypeVar("StateT", bound=AgentGraphState)
+InputT = TypeVar("InputT", bound=BaseModel)
+
+
 def create_llm_node(
     model: BaseChatModel,
-    tools: Sequence[BaseTool] | None = None,
+    tools: Sequence[BaseTool],
+    input_schema: type[InputT] | None = None,
     is_conversational: bool = False,
     llm_messages_limit: int = DEFAULT_MAX_LLM_MESSAGES,
     thinking_messages_limit: int = DEFAULT_MAX_CONSECUTIVE_THINKING_MESSAGES,
@@ -69,12 +82,10 @@ def create_llm_node(
             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):
+    async def llm_node(state: StateT):
         messages: list[AnyMessage] = state.messages
-
         agent_ai_messages = sum(1 for msg in messages if isinstance(msg, AIMessage))
         if agent_ai_messages >= llm_messages_limit:
             raise AgentTerminationException(
@@ -86,6 +97,11 @@ async def llm_node(state: AgentGraphState):
 
         consecutive_thinking_messages = count_consecutive_thinking_messages(messages)
 
+        static_schema_tools = _apply_tool_argument_properties(
+            bindable_tools, state, input_schema
+        )
+        base_llm = model.bind_tools(static_schema_tools)
+
         if (
             not is_conversational
             and bindable_tools
@@ -111,3 +127,19 @@ async def llm_node(state: AgentGraphState):
         return {"messages": [response]}
 
     return llm_node
+
+
+def _apply_tool_argument_properties(
+    tools: list[BaseTool],
+    state: StateT,
+    input_schema: type[InputT] | None = None,
+) -> list[BaseTool]:
+    """Apply dynamic schema modifications to tools based on their argument_properties."""
+
+    agent_input = extract_input_data_from_state(state, input_schema or type(state))
+    return [
+        apply_static_argument_properties_to_schema(tool, agent_input)
+        if isinstance(tool, StructuredToolWithArgumentProperties)
+        else tool
+        for tool in tools
+    ]

src/uipath_langchain/agent/react/utils.py

@@ -34,6 +34,30 @@ def resolve_output_model(
     return END_EXECUTION_TOOL.args_schema
 
 
+def extract_input_data_from_state(
+    state: BaseModel | dict[str, Any],
+    input_model: type[BaseModel],
+) -> dict[str, Any]:
+    """Extract only input schema fields from graph state, filtering out internal fields.
+
+    This prevents internal LangGraph state fields (messages, termination, agent_outcome, etc.)
+    from leaking into template interpolation.
+
+    Args:
+        state: The combined agent graph state (InnerAgentGraphState = AgentGraphState + input_schema).
+               At runtime, this is a dynamically created class that inherits from both.
+        input_model: The input schema model defining allowed fields
+
+    Returns:
+        Dictionary containing only graph input arguments defined in the agent's input_schema
+    """
+    if isinstance(state, BaseModel):
+        graph_state = state.model_dump()
+    else:
+        graph_state = state
+    return input_model.model_validate(graph_state, from_attributes=True).model_dump()
+
+
 def count_consecutive_thinking_messages(messages: Sequence[BaseMessage]) -> int:
     """Count consecutive AIMessages without tool calls at end of message history."""
     if not messages:

src/uipath_langchain/agent/tools/integration_tool.py

@@ -3,20 +3,27 @@
 import copy
 from typing import Any
 
+from jsonschema_pydantic_converter import transform as create_model
+from langchain.tools import BaseTool
+from langchain_core.messages import ToolCall
 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 uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
-from uipath_langchain.agent.tools.tool_node import ToolWrapperMixin
+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 (
+    ToolWrapperMixin,
+    ToolWrapperReturnType,
+)
 
 from .structured_tool_with_output_type import StructuredToolWithOutputType
 from .utils import sanitize_dict_for_serialization, sanitize_tool_name
 
 
-class StructuredToolWithStaticArgs(StructuredToolWithOutputType, ToolWrapperMixin):
+class StructuredToolWithWrapper(StructuredToolWithOutputType, ToolWrapperMixin):
     pass
 
 
@@ -167,13 +174,15 @@ async def integration_tool_fn(**kwargs: Any):
 
         return result
 
-    from uipath_langchain.agent.wrappers.static_args_wrapper import (
-        get_static_args_wrapper,
-    )
-
-    wrapper = get_static_args_wrapper(resource)
+    async def integration_tool_wrapper(
+        tool: BaseTool,
+        call: ToolCall,
+        state: AgentGraphState,
+    ) -> ToolWrapperReturnType:
+        modified_args = handle_static_args(resource, state, call["args"])
+        return await tool.ainvoke(modified_args)
 
-    tool = StructuredToolWithStaticArgs(
+    tool = StructuredToolWithWrapper(
         name=tool_name,
         description=resource.description,
         args_schema=input_model,
@@ -184,6 +193,6 @@ async def integration_tool_fn(**kwargs: Any):
             "display_name": resource.name,
         },
     )
-    tool.set_tool_wrappers(awrapper=wrapper)
+    tool.set_tool_wrappers(awrapper=integration_tool_wrapper)
 
     return tool

src/uipath_langchain/agent/tools/internal_tools/analyze_files_tool.py

@@ -3,7 +3,8 @@
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
-from langchain_core.tools import StructuredTool
+from langchain_core.messages.tool import ToolCall
+from langchain_core.tools import BaseTool, StructuredTool
 from uipath.agent.models.agent import (
     AgentInternalToolResourceConfig,
 )
@@ -12,11 +13,16 @@
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 from uipath_langchain.agent.react.llm_with_files import FileInfo, llm_call_with_files
-from uipath_langchain.agent.tools.structured_tool_with_output_type import (
-    StructuredToolWithOutputType,
+from uipath_langchain.agent.react.types import AgentGraphState
+from uipath_langchain.agent.tools.static_args import handle_static_args
+from uipath_langchain.agent.tools.structured_tool_with_argument_properties import (
+    StructuredToolWithArgumentProperties,
+)
+from uipath_langchain.agent.tools.tool_node import (
+    ToolWrapperReturnType,
 )
-from uipath_langchain.agent.tools.tool_node import ToolWrapperMixin
 from uipath_langchain.agent.tools.utils import sanitize_tool_name
+from uipath_langchain.agent.wrappers import get_job_attachment_wrapper
 
 ANALYZE_FILES_SYSTEM_MESSAGE = (
     "Process the provided files to complete the given task. "
@@ -25,17 +31,9 @@
 )
 
 
-class AnalyzeFileTool(StructuredToolWithOutputType, ToolWrapperMixin):
-    pass
-
-
 def create_analyze_file_tool(
     resource: AgentInternalToolResourceConfig, llm: BaseChatModel
 ) -> StructuredTool:
-    from uipath_langchain.agent.wrappers.job_attachment_wrapper import (
-        get_job_attachment_wrapper,
-    )
-
     tool_name = sanitize_tool_name(resource.name)
     input_model = create_model(resource.input_schema)
     output_model = create_model(resource.output_schema)
@@ -69,15 +67,25 @@ async def tool_fn(**kwargs: Any):
         result = await llm_call_with_files(messages, files, llm)
         return result
 
-    wrapper = get_job_attachment_wrapper(output_type=output_model)
-    tool = AnalyzeFileTool(
+    job_attachment_wrapper = get_job_attachment_wrapper(output_type=output_model)
+
+    async def analyze_file_tool_wrapper(
+        tool: BaseTool,
+        call: ToolCall,
+        state: AgentGraphState,
+    ) -> ToolWrapperReturnType:
+        call["args"] = handle_static_args(resource, state, call["args"])
+        return await job_attachment_wrapper(tool, call, state)
+
+    tool = StructuredToolWithArgumentProperties(
         name=tool_name,
         description=resource.description,
         args_schema=input_model,
         coroutine=tool_fn,
         output_type=output_model,
+        argument_properties=resource.argument_properties,
     )
-    tool.set_tool_wrappers(awrapper=wrapper)
+    tool.set_tool_wrappers(awrapper=analyze_file_tool_wrapper)
     return tool
 
 

src/uipath_langchain/agent/tools/process_tool.py

@@ -2,26 +2,28 @@
 
 from typing import Any
 
+from jsonschema_pydantic_converter import transform as create_model
+from langchain.tools import BaseTool
+from langchain_core.messages import ToolCall
 from langchain_core.tools import StructuredTool
 from langgraph.types import interrupt
 from uipath.agent.models.agent import AgentProcessToolResourceConfig
 from uipath.eval.mocks import mockable
 from uipath.platform.common import InvokeProcess
 
-from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
-from uipath_langchain.agent.wrappers.job_attachment_wrapper import (
-    get_job_attachment_wrapper,
+from uipath_langchain.agent.react.types import AgentGraphState
+from uipath_langchain.agent.tools.static_args import handle_static_args
+from uipath_langchain.agent.tools.structured_tool_with_argument_properties import (
+    StructuredToolWithArgumentProperties,
 )
+from uipath_langchain.agent.tools.tool_node import (
+    ToolWrapperReturnType,
+)
+from uipath_langchain.agent.wrappers import get_job_attachment_wrapper
 
-from .structured_tool_with_output_type import StructuredToolWithOutputType
-from .tool_node import ToolWrapperMixin
 from .utils import sanitize_tool_name
 
 
-class ProcessTool(StructuredToolWithOutputType, ToolWrapperMixin):
-    pass
-
-
 def create_process_tool(resource: AgentProcessToolResourceConfig) -> StructuredTool:
     """Uses interrupt() to suspend graph execution until process completes (handled by runtime)."""
     tool_name: str = sanitize_tool_name(resource.name)
@@ -48,8 +50,17 @@ async def process_tool_fn(**kwargs: Any):
             )
         )
 
-    wrapper = get_job_attachment_wrapper(output_type=output_model)
-    tool = ProcessTool(
+    job_attachment_wrapper = get_job_attachment_wrapper(output_type=output_model)
+
+    async def process_tool_wrapper(
+        tool: BaseTool,
+        call: ToolCall,
+        state: AgentGraphState,
+    ) -> ToolWrapperReturnType:
+        call["args"] = handle_static_args(resource, state, call["args"])
+        return await job_attachment_wrapper(tool, call, state)
+
+    tool = StructuredToolWithArgumentProperties(
         name=tool_name,
         description=resource.description,
         args_schema=input_model,
@@ -60,6 +71,8 @@ async def process_tool_fn(**kwargs: Any):
             "display_name": process_name,
             "folder_path": folder_path,
         },
+        argument_properties=resource.argument_properties,
     )
-    tool.set_tool_wrappers(awrapper=wrapper)
+    tool.set_tool_wrappers(awrapper=process_tool_wrapper)
+
     return tool