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

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

@@ -1,24 +1,31 @@
 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.core.guardrails import DeterministicGuardrail
 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},
@@ -110,7 +117,7 @@ def _create_guardrails_subgraph(
             ExecutionStage.POST_EXECUTION,
             node_factory,
             END,
-            inner_node,
+            inner_name,
         )
         subgraph.add_edge(inner_name, first_post_exec_guardrail_node)
     else:
@@ -197,10 +204,21 @@ def create_llm_guardrails_subgraph(
     llm_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
 ):
+    """Create a guarded LLM node.
+
+    Args:
+        llm_node: Tuple of (node_name, node_callable) for the LLM node.
+        guardrails: Optional sequence of (guardrail, action) tuples.
+
+    Returns:
+        Either the original node callable (if no applicable guardrails) or a compiled
+        LangGraph subgraph that enforces the configured guardrails.
+    """
     applicable_guardrails = [
         (guardrail, _)
         for (guardrail, _) in (guardrails or [])
         if GuardrailScope.LLM in guardrail.selector.scopes
+        and not isinstance(guardrail, DeterministicGuardrail)
     ]
     if applicable_guardrails is None or len(applicable_guardrails) == 0:
         return llm_node[1]
@@ -215,13 +233,19 @@ 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]:
-    """Create tool nodes with guardrails.
+) -> dict[str, RunnableCallable]:
+    """Create tool nodes with guardrails applied.
     Args:
+        tool_nodes: Mapping of tool name to a LangGraph `ToolNode`.
+        guardrails: Optional sequence of (guardrail, action) tuples.
+
+    Returns:
+        A mapping of tool name to either the original `ToolNode` or a compiled subgraph
+        that enforces the matching tool guardrails.
     """
-    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),
@@ -232,37 +256,107 @@ def create_tools_guardrails_subgraph(
     return result
 
 
-def create_agent_guardrails_subgraph(
-    agent_node: tuple[str, Any],
+def create_agent_init_guardrails_subgraph(
+    init_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
-    execution_stage: ExecutionStage,
-):
-    """Create a subgraph for AGENT-scoped guardrails that applies checks at the specified stage.
+) -> Any:
+    """Create a subgraph for the INIT node and apply AGENT guardrails after INIT.
+
+    This subgraph intentionally **runs the INIT node first** (so it can seed/normalize
+    the agent state), and then evaluates guardrails as **PRE_EXECUTION**. This lets
+    guardrails intended to run "before agent execution" validate the post-init state.
+
+    Args:
+        init_node: Tuple of (node_name, node_callable) for the INIT node.
+        guardrails: Optional sequence of (guardrail, action) tuples.
 
-    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.
+    Returns:
+        Either the original node callable (if no applicable guardrails) or a compiled
+        LangGraph subgraph that runs INIT then enforces PRE_EXECUTION AGENT guardrails.
     """
     applicable_guardrails = [
         (guardrail, _)
         for (guardrail, _) in (guardrails or [])
         if GuardrailScope.AGENT in guardrail.selector.scopes
+        and not isinstance(guardrail, DeterministicGuardrail)
     ]
+    applicable_guardrails = _filter_guardrails_by_stage(
+        applicable_guardrails, ExecutionStage.PRE_EXECUTION
+    )
     if applicable_guardrails is None or len(applicable_guardrails) == 0:
-        return agent_node[1]
+        return init_node[1]
 
-    return _create_guardrails_subgraph(
-        main_inner_node=agent_node,
+    inner_name, inner_node = init_node
+    subgraph = StateGraph(AgentGuardrailsGraphState)
+    subgraph.add_node(inner_name, inner_node)
+    subgraph.add_edge(START, inner_name)
+
+    first_guardrail_node = _build_guardrail_node_chain(
+        subgraph=subgraph,
         guardrails=applicable_guardrails,
         scope=GuardrailScope.AGENT,
-        execution_stages=[execution_stage],
-        node_factory=create_agent_guardrail_node,
+        execution_stage=ExecutionStage.PRE_EXECUTION,
+        node_factory=create_agent_init_guardrail_node,
+        next_node=END,
+        guarded_node_name=inner_name,
     )
+    subgraph.add_edge(inner_name, first_guardrail_node)
+    return subgraph.compile()
+
+
+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
+        and not isinstance(guardrail, DeterministicGuardrail)
+    ]
+    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,
 ):
+    """Create a guarded tool node.
+
+    Args:
+        tool_node: Tuple of (tool_name, tool_node_callable).
+        guardrails: Optional sequence of (guardrail, action) tuples.
+
+    Returns:
+        Either the original tool node callable (if no matching guardrails) or a compiled
+        LangGraph subgraph that enforces the matching tool guardrails.
+    """
     tool_name, _ = tool_node
     applicable_guardrails = [
         (guardrail, action)

uipath_langchain/agent/react/init_node.py

@@ -3,11 +3,17 @@
 from typing import Any, Callable, Sequence
 
 from langchain_core.messages import HumanMessage, SystemMessage
+from pydantic import BaseModel
+
+from .job_attachments import (
+    get_job_attachments,
+)
 
 
 def create_init_node(
     messages: Sequence[SystemMessage | HumanMessage]
     | Callable[[Any], Sequence[SystemMessage | HumanMessage]],
+    input_schema: type[BaseModel] | None,
 ):
     def graph_state_init(state: Any):
         if callable(messages):
@@ -15,6 +21,15 @@ def create_init_node(
         else:
             resolved_messages = messages
 
-        return {"messages": list(resolved_messages)}
+        schema = input_schema if input_schema is not None else BaseModel
+        job_attachments = get_job_attachments(schema, state)
+        job_attachments_dict = {
+            str(att.id): att for att in job_attachments if att.id is not None
+        }
+
+        return {
+            "messages": list(resolved_messages),
+            "job_attachments": job_attachments_dict,
+        }
 
     return graph_state_init

uipath_langchain/agent/react/job_attachments.py

@@ -0,0 +1,125 @@
+"""Job attachment utilities for ReAct Agent."""
+
+import copy
+import uuid
+from typing import Any
+
+from jsonpath_ng import parse  # type: ignore[import-untyped]
+from pydantic import BaseModel
+from uipath.platform.attachments import Attachment
+
+from .json_utils import extract_values_by_paths, get_json_paths_by_type
+
+
+def get_job_attachments(
+    schema: type[BaseModel],
+    data: dict[str, Any] | BaseModel,
+) -> list[Attachment]:
+    """Extract job attachments from data based on schema and convert to Attachment objects.
+
+    Args:
+        schema: The Pydantic model class defining the data structure
+        data: The data object (dict or Pydantic model) to extract attachments from
+
+    Returns:
+        List of Attachment objects
+    """
+    job_attachment_paths = get_job_attachment_paths(schema)
+    job_attachments = extract_values_by_paths(data, job_attachment_paths)
+
+    result = []
+    for attachment in job_attachments:
+        result.append(Attachment.model_validate(attachment, from_attributes=True))
+
+    return result
+
+
+def get_job_attachment_paths(model: type[BaseModel]) -> list[str]:
+    """Get JSONPath expressions for all job attachment fields in a Pydantic model.
+
+    Args:
+        model: The Pydantic model class to analyze
+
+    Returns:
+        List of JSONPath expressions pointing to job attachment fields
+    """
+    return get_json_paths_by_type(model, "Job_attachment")
+
+
+def replace_job_attachment_ids(
+    json_paths: list[str],
+    tool_args: dict[str, Any],
+    state: dict[str, Attachment],
+    errors: list[str],
+) -> dict[str, Any]:
+    """Replace job attachment IDs in tool_args with full attachment objects from state.
+
+    For each JSON path, this function finds matching objects in tool_args and
+    replaces them with corresponding attachment objects from state. The matching
+    is done by looking up the object's 'ID' field in the state dictionary.
+
+    If an ID is not a valid UUID or is not present in state, an error message
+    is added to the errors list.
+
+    Args:
+        json_paths: List of JSONPath expressions (e.g., ["$.attachment", "$.attachments[*]"])
+        tool_args: The dictionary containing tool arguments to modify
+        state: Dictionary mapping attachment UUID strings to Attachment objects
+        errors: List to collect error messages for invalid or missing IDs
+
+    Returns:
+        Modified copy of tool_args with attachment IDs replaced by full objects
+
+    Example:
+        >>> state = {
+        ...     "123e4567-e89b-12d3-a456-426614174000": Attachment(id="123e4567-e89b-12d3-a456-426614174000", name="file1.pdf"),
+        ...     "223e4567-e89b-12d3-a456-426614174001": Attachment(id="223e4567-e89b-12d3-a456-426614174001", name="file2.pdf")
+        ... }
+        >>> tool_args = {
+        ...     "attachment": {"ID": "123"},
+        ...     "other_field": "value"
+        ... }
+        >>> paths = ['$.attachment']
+        >>> errors = []
+        >>> replace_job_attachment_ids(paths, tool_args, state, errors)
+        {'attachment': {'ID': '123', 'name': 'file1.pdf', ...}, 'other_field': 'value'}
+    """
+    result = copy.deepcopy(tool_args)
+
+    for json_path in json_paths:
+        expr = parse(json_path)
+        matches = expr.find(result)
+
+        for match in matches:
+            current_value = match.value
+
+            if isinstance(current_value, dict) and "ID" in current_value:
+                attachment_id_str = str(current_value["ID"])
+
+                try:
+                    uuid.UUID(attachment_id_str)
+                except (ValueError, AttributeError):
+                    errors.append(
+                        _create_job_attachment_error_message(attachment_id_str)
+                    )
+                    continue
+
+                if attachment_id_str in state:
+                    replacement_value = state[attachment_id_str]
+                    match.full_path.update(
+                        result, replacement_value.model_dump(by_alias=True, mode="json")
+                    )
+                else:
+                    errors.append(
+                        _create_job_attachment_error_message(attachment_id_str)
+                    )
+
+    return result
+
+
+def _create_job_attachment_error_message(attachment_id_str: str) -> str:
+    return (
+        f"Could not find JobAttachment with ID='{attachment_id_str}'. "
+        f"Try invoking the tool again and please make sure that you pass "
+        f"valid JobAttachment IDs associated with existing JobAttachments in the current context."
+    )

uipath_langchain/agent/react/json_utils.py

@@ -0,0 +1,183 @@
+import sys
+from typing import Any, ForwardRef, Union, get_args, get_origin
+
+from jsonpath_ng import parse  # type: ignore[import-untyped]
+from pydantic import BaseModel
+
+
+def get_json_paths_by_type(model: type[BaseModel], type_name: str) -> list[str]:
+    """Get JSONPath expressions for all fields that reference a specific type.
+
+    This function recursively traverses nested Pydantic models to find all paths
+    that lead to fields of the specified type.
+
+    Args:
+        model: A Pydantic model class
+        type_name: The name of the type to search for (e.g., "Job_attachment")
+
+    Returns:
+        List of JSONPath expressions using standard JSONPath syntax.
+        For array fields, uses [*] to indicate all array elements.
+
+    Example:
+        >>> schema = {
+        ...     "type": "object",
+        ...     "properties": {
+        ...         "attachment": {"$ref": "#/definitions/job-attachment"},
+        ...         "attachments": {
+        ...             "type": "array",
+        ...             "items": {"$ref": "#/definitions/job-attachment"}
+        ...         }
+        ...     },
+        ...     "definitions": {
+        ...         "job-attachment": {"type": "object", "properties": {"id": {"type": "string"}}}
+        ...     }
+        ... }
+        >>> model = transform(schema)
+        >>> _get_json_paths_by_type(model, "Job_attachment")
+        ['$.attachment', '$.attachments[*]']
+    """
+
+    def _recursive_search(
+        current_model: type[BaseModel], current_path: str
+    ) -> list[str]:
+        """Recursively search for fields of the target type."""
+        json_paths = []
+
+        target_type = _get_target_type(current_model, type_name)
+        matches_type = _create_type_matcher(type_name, target_type)
+
+        for field_name, field_info in current_model.model_fields.items():
+            annotation = field_info.annotation
+
+            if current_path:
+                field_path = f"{current_path}.{field_name}"
+            else:
+                field_path = f"$.{field_name}"
+
+            annotation = _unwrap_optional(annotation)
+            origin = get_origin(annotation)
+
+            if matches_type(annotation):
+                json_paths.append(field_path)
+                continue
+
+            if origin is list:
+                args = get_args(annotation)
+                if args:
+                    list_item_type = args[0]
+                    if matches_type(list_item_type):
+                        json_paths.append(f"{field_path}[*]")
+                        continue
+
+                    if _is_pydantic_model(list_item_type):
+                        nested_paths = _recursive_search(
+                            list_item_type, f"{field_path}[*]"
+                        )
+                        json_paths.extend(nested_paths)
+                        continue
+
+            if _is_pydantic_model(annotation):
+                nested_paths = _recursive_search(annotation, field_path)
+                json_paths.extend(nested_paths)
+
+        return json_paths
+
+    return _recursive_search(model, "")
+
+
+def extract_values_by_paths(
+    obj: dict[str, Any] | BaseModel, json_paths: list[str]
+) -> list[Any]:
+    """Extract values from an object using JSONPath expressions.
+
+    Args:
+        obj: The object (dict or Pydantic model) to extract values from
+        json_paths: List of JSONPath expressions. **Paths are assumed to be disjoint**
+                    (non-overlapping). If paths overlap, duplicate values will be returned.
+
+    Returns:
+        List of all extracted values (flattened)
+
+    Example:
+        >>> obj = {
+        ...     "attachment": {"id": "123"},
+        ...     "attachments": [{"id": "456"}, {"id": "789"}]
+        ... }
+        >>> paths = ['$.attachment', '$.attachments[*]']
+        >>> _extract_values_by_paths(obj, paths)
+        [{'id': '123'}, {'id': '456'}, {'id': '789'}]
+    """
+    data = obj.model_dump() if isinstance(obj, BaseModel) else obj
+
+    results = []
+    for json_path in json_paths:
+        expr = parse(json_path)
+        matches = expr.find(data)
+        results.extend([match.value for match in matches])
+
+    return results
+
+
+def _get_target_type(model: type[BaseModel], type_name: str) -> Any:
+    """Get the target type from the model's module.
+
+    Args:
+        model: A Pydantic model class
+        type_name: The name of the type to search for
+
+    Returns:
+        The target type if found, None otherwise
+    """
+    model_module = sys.modules.get(model.__module__)
+    if model_module and hasattr(model_module, type_name):
+        return getattr(model_module, type_name)
+    return None
+
+
+def _create_type_matcher(type_name: str, target_type: Any) -> Any:
+    """Create a function that checks if an annotation matches the target type.
+
+    Args:
+        type_name: The name of the type to match
+        target_type: The actual type object (can be None)
+
+    Returns:
+        A function that takes an annotation and returns True if it matches
+    """
+
+    def matches_type(annotation: Any) -> bool:
+        """Check if an annotation matches the target type name."""
+        if isinstance(annotation, ForwardRef):
+            return annotation.__forward_arg__ == type_name
+        if isinstance(annotation, str):
+            return annotation == type_name
+        if hasattr(annotation, "__name__") and annotation.__name__ == type_name:
+            return True
+        if target_type is not None and annotation is target_type:
+            return True
+        return False
+
+    return matches_type
+
+
+def _unwrap_optional(annotation: Any) -> Any:
+    """Unwrap Optional/Union types to get the underlying type.
+
+    Args:
+        annotation: The type annotation to unwrap
+
+    Returns:
+        The unwrapped type, or the original if not Optional/Union
+    """
+    origin = get_origin(annotation)
+    if origin is Union:
+        args = get_args(annotation)
+        non_none_args = [arg for arg in args if arg is not type(None)]
+        if non_none_args:
+            return non_none_args[0]
+    return annotation
+
+
+def _is_pydantic_model(annotation: Any) -> bool:
+    return isinstance(annotation, type) and issubclass(annotation, BaseModel)

uipath_langchain/agent/react/jsonschema_pydantic_converter.py

@@ -0,0 +1,76 @@
+import inspect
+import sys
+from types import ModuleType
+from typing import Any, Type, get_args, get_origin
+
+from jsonschema_pydantic_converter import transform_with_modules
+from pydantic import BaseModel
+
+# Shared pseudo-module for all dynamically created types
+# This allows get_type_hints() to resolve forward references
+_DYNAMIC_MODULE_NAME = "jsonschema_pydantic_converter._dynamic"
+
+
+def _get_or_create_dynamic_module() -> ModuleType:
+    """Get or create the shared pseudo-module for dynamic types."""
+    if _DYNAMIC_MODULE_NAME not in sys.modules:
+        pseudo_module = ModuleType(_DYNAMIC_MODULE_NAME)
+        pseudo_module.__doc__ = (
+            "Shared module for dynamically generated Pydantic models from JSON schemas"
+        )
+        sys.modules[_DYNAMIC_MODULE_NAME] = pseudo_module
+    return sys.modules[_DYNAMIC_MODULE_NAME]
+
+
+def create_model(
+    schema: dict[str, Any],
+) -> Type[BaseModel]:
+    model, namespace = transform_with_modules(schema)
+    corrected_namespace: dict[str, Any] = {}
+
+    def collect_types(annotation: Any) -> None:
+        """Recursively collect all BaseModel types from an annotation."""
+        # Unwrap generic types like List, Optional, etc.
+        origin = get_origin(annotation)
+        if origin is not None:
+            for arg in get_args(annotation):
+                collect_types(arg)
+
+        elif inspect.isclass(annotation) and issubclass(annotation, BaseModel):
+            # Find the original name for this type from the namespace
+            for type_name, type_def in namespace.items():
+                # Match by class name since rebuild may create new instances
+                if (
+                    hasattr(annotation, "__name__")
+                    and hasattr(type_def, "__name__")
+                    and annotation.__name__ == type_def.__name__
+                ):
+                    # Store the actual annotation type, not the old namespace one
+                    annotation.__name__ = type_name
+                    corrected_namespace[type_name] = annotation
+                    break
+
+    # Collect all types from field annotations
+    for field_info in model.model_fields.values():
+        collect_types(field_info.annotation)
+
+    # Get the shared pseudo-module and populate it with this schema's types
+    # This ensures that forward references can be resolved by get_type_hints()
+    # when the model is used with external libraries (e.g., LangGraph)
+    pseudo_module = _get_or_create_dynamic_module()
+
+    # Populate the pseudo-module with all types from the namespace
+    # Use the original names so forward references resolve correctly
+    for type_name, type_def in corrected_namespace.items():
+        setattr(pseudo_module, type_name, type_def)
+
+    setattr(pseudo_module, model.__name__, model)
+
+    # Update the model's __module__ to point to the shared pseudo-module
+    model.__module__ = _DYNAMIC_MODULE_NAME
+
+    # Update the __module__ of all generated types in the namespace
+    for type_def in corrected_namespace.values():
+        if inspect.isclass(type_def) and issubclass(type_def, BaseModel):
+            type_def.__module__ = _DYNAMIC_MODULE_NAME
+    return model