Merge pull request #388 from UiPath/feat/isolate-inner-state

feat: isolate inner schema to avoid name collisions

Author: andreitava-uip

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.3.1"
+version = "0.3.2"
 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/exceptions/__init__.py

@@ -1,4 +1,7 @@
-from .exceptions import AgentNodeRoutingException, AgentTerminationException
+from .exceptions import (
+    AgentNodeRoutingException,
+    AgentTerminationException,
+)
 
 __all__ = [
     "AgentNodeRoutingException",

src/uipath_langchain/agent/react/agent.py

@@ -36,14 +36,14 @@
 
 
 def create_state_with_input(input_schema: Type[InputT]):
-    InnerAgentGraphState = type(
-        "InnerAgentGraphState",
+    CompleteAgentGraphState = type(
+        "CompleteAgentGraphState",
         (AgentGraphState, input_schema),
         {},
     )
 
-    cast(type[BaseModel], InnerAgentGraphState).model_rebuild()
-    return InnerAgentGraphState
+    cast(type[BaseModel], CompleteAgentGraphState).model_rebuild()
+    return CompleteAgentGraphState
 
 
 def create_agent(
@@ -85,12 +85,12 @@ def create_agent(
     )
     terminate_node = create_terminate_node(output_schema)
 
-    InnerAgentGraphState = create_state_with_input(
+    CompleteAgentGraphState = create_state_with_input(
         input_schema if input_schema is not None else BaseModel
     )
 
     builder: StateGraph[AgentGraphState, None, InputT, OutputT] = StateGraph(
-        InnerAgentGraphState, input_schema=input_schema, output_schema=output_schema
+        CompleteAgentGraphState, input_schema=input_schema, output_schema=output_schema
     )
     init_with_guardrails_subgraph = create_agent_init_guardrails_subgraph(
         (AgentGraphNode.GUARDED_INIT, init_node),

src/uipath_langchain/agent/react/init_node.py

@@ -15,7 +15,7 @@ def create_init_node(
     | Callable[[Any], Sequence[SystemMessage | HumanMessage]],
     input_schema: type[BaseModel] | None,
 ):
-    def graph_state_init(state: Any):
+    def graph_state_init(state: Any) -> Any:
         if callable(messages):
             resolved_messages = messages(state)
         else:
@@ -29,7 +29,9 @@ def graph_state_init(state: Any):
 
         return {
             "messages": list(resolved_messages),
-            "job_attachments": job_attachments_dict,
+            "inner_state": {
+                "job_attachments": job_attachments_dict,
+            },
         }
 
     return graph_state_init

src/uipath_langchain/agent/react/reducers.py

@@ -0,0 +1,90 @@
+"""Dict-like object reducers for merging state with field-specific reducers."""
+
+from typing import Any
+
+from pydantic import BaseModel
+from uipath.platform.attachments import Attachment
+
+
+def add_job_attachments(
+    left: dict[str, Attachment], right: dict[str, Attachment]
+) -> dict[str, Attachment]:
+    """Merge attachment dictionaries, with right values taking precedence.
+
+    This reducer function merges two dictionaries of attachments by UUID string.
+    If the same UUID exists in both dictionaries, the value from 'right' takes precedence.
+
+    Args:
+        left: Existing dictionary of attachments keyed by UUID string
+        right: New dictionary of attachments to merge
+
+    Returns:
+        Merged dictionary with right values overriding left values for duplicate keys
+    """
+    if not right:
+        return left
+
+    if not left:
+        return right
+
+    return {**left, **right}
+
+
+def merge_objects(left: Any, right: Any) -> Any:
+    """Merge a Pydantic model with another model or dict, with right values taking precedence.
+
+    Applies field-specific reducers from annotation metadata when merging values.
+
+    Args:
+        left: Existing Pydantic BaseModel instance
+        right: New Pydantic BaseModel instance or dict to merge
+
+    Returns:
+        New Pydantic model instance with merged values
+
+    Raises:
+        TypeError: If left is not a Pydantic BaseModel or right is not a BaseModel or dict
+    """
+    if not right:
+        return left
+
+    if not left:
+        return right
+
+    # validate input types
+    if not isinstance(left, BaseModel):
+        raise TypeError("Left object must be a Pydantic BaseModel")
+
+    if not isinstance(right, (BaseModel, dict)):
+        raise TypeError("Right object must be a Pydantic BaseModel or dict")
+
+    model_fields = type(left).model_fields
+    merged_values = {}
+
+    for field_name in model_fields:
+        merged_values[field_name] = getattr(left, field_name)
+
+    for field_name in model_fields:
+        if isinstance(right, BaseModel):
+            if hasattr(right, field_name):
+                right_value = getattr(right, field_name)
+            else:
+                continue  # field not present in right
+        else:
+            # right is dict
+            if field_name not in right:
+                continue  # field not present in right
+            right_value = right[field_name]
+
+        field_info = model_fields[field_name]
+        left_value = merged_values[field_name]
+
+        # apply reducer if defined
+        if field_info.metadata and callable(field_info.metadata[0]):
+            reducer_func = field_info.metadata[0]
+            merged_values[field_name] = reducer_func(left_value, right_value)
+        else:
+            merged_values[field_name] = right_value
+
+    # return new model instance with merged values
+    return type(left)(**merged_values)

src/uipath_langchain/agent/react/terminate_node.py

@@ -57,8 +57,8 @@ def create_terminate_node(
     """
 
     def terminate_node(state: AgentGraphState):
-        if state.termination:
-            _handle_agent_termination(state.termination)
+        if state.inner_state.termination:
+            _handle_agent_termination(state.inner_state.termination)
 
         last_message = state.messages[-1]
         if not isinstance(last_message, AIMessage):

src/uipath_langchain/agent/react/types.py

@@ -6,7 +6,7 @@
 from pydantic import BaseModel, Field
 from uipath.platform.attachments import Attachment
 
-from uipath_langchain.agent.react.utils import add_job_attachments
+from uipath_langchain.agent.react.reducers import add_job_attachments, merge_objects
 
 
 class AgentTerminationSource(StrEnum):
@@ -21,12 +21,18 @@ class AgentTermination(BaseModel):
     detail: str = ""
 
 
+class InnerAgentGraphState(BaseModel):
+    job_attachments: Annotated[dict[str, Attachment], add_job_attachments] = {}
+    termination: AgentTermination | None = None
+
+
 class AgentGraphState(BaseModel):
     """Agent Graph state for standard loop execution."""
 
     messages: Annotated[list[AnyMessage], add_messages] = []
-    job_attachments: Annotated[dict[str, Attachment], add_job_attachments] = {}
-    termination: AgentTermination | None = None
+    inner_state: Annotated[InnerAgentGraphState, merge_objects] = Field(
+        default_factory=InnerAgentGraphState
+    )
 
 
 class AgentGuardrailsGraphState(AgentGraphState):

src/uipath_langchain/agent/react/utils.py

@@ -5,7 +5,6 @@
 from langchain_core.messages import AIMessage, BaseMessage
 from pydantic import BaseModel
 from uipath.agent.react import END_EXECUTION_TOOL
-from uipath.platform.attachments import Attachment
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
 
@@ -49,27 +48,3 @@ def count_consecutive_thinking_messages(messages: Sequence[BaseMessage]) -> int:
         count += 1
 
     return count
-
-
-def add_job_attachments(
-    left: dict[str, Attachment], right: dict[str, Attachment]
-) -> dict[str, Attachment]:
-    """Merge attachment dictionaries, with right values taking precedence.
-
-    This reducer function merges two dictionaries of attachments by UUID string.
-    If the same UUID exists in both dictionaries, the value from 'right' takes precedence.
-
-    Args:
-        left: Existing dictionary of attachments keyed by UUID string
-        right: New dictionary of attachments to merge
-
-    Returns:
-        Merged dictionary with right values overriding left values for duplicate keys
-    """
-    if not right:
-        return left
-
-    if not left:
-        return right
-
-    return {**left, **right}

src/uipath_langchain/agent/tools/escalation_tool.py

@@ -114,10 +114,12 @@ async def escalation_wrapper(
                             tool_call_id=call["id"],
                         )
                     ],
-                    "termination": {
-                        "source": AgentTerminationSource.ESCALATION,
-                        "title": termination_title,
-                        "detail": output_detail,
+                    "inner_state": {
+                        "termination": {
+                            "source": AgentTerminationSource.ESCALATION,
+                            "title": termination_title,
+                            "detail": output_detail,
+                        }
                     },
                 },
                 goto=AgentGraphNode.TERMINATE,

src/uipath_langchain/agent/wrappers/job_attachment_wrapper.py

@@ -51,7 +51,7 @@ async def job_attachment_wrapper(
             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
+                paths, input_args, state.inner_state.job_attachments, errors
             )
 
             if errors: