From 988e7db0484e273a328733ea861887f807fc0920 Mon Sep 17 00:00:00 2001
From: radu-mocanu <radu.mocanu@uipath.com>
Date: Wed, 22 Jul 2026 10:41:20 +0300
Subject: [PATCH] fix: type memory recall graph state

---
 pyproject.toml                                  |  2 +-
 src/uipath_langchain/agent/react/memory_node.py |  4 ++--
 tests/agent/react/test_memory_node.py           | 11 ++++++++---
 uv.lock                                         |  2 +-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 93e91d9fc..92290ce6f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.14.14"
+version = "0.14.15"
 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"
diff --git a/src/uipath_langchain/agent/react/memory_node.py b/src/uipath_langchain/agent/react/memory_node.py
index da2400a4a..c44f2c217 100644
--- a/src/uipath_langchain/agent/react/memory_node.py
+++ b/src/uipath_langchain/agent/react/memory_node.py
@@ -20,7 +20,7 @@
     SearchSettings,
 )
 
-from .types import MemoryConfig
+from .types import AgentGraphState, MemoryConfig
 from .utils import extract_input_data_from_state
 
 logger = logging.getLogger(__name__)
@@ -49,7 +49,7 @@ def create_memory_recall_node(
         An async callable suitable for ``builder.add_node()``.
     """
 
-    async def memory_recall_node(state: Any) -> dict[str, Any]:
+    async def memory_recall_node(state: AgentGraphState) -> dict[str, Any]:
         input_model = input_schema if input_schema is not None else BaseModel
         input_arguments = extract_input_data_from_state(state, input_model)
         if not input_arguments:
diff --git a/tests/agent/react/test_memory_node.py b/tests/agent/react/test_memory_node.py
index be94cc2cc..26c1bf2ca 100644
--- a/tests/agent/react/test_memory_node.py
+++ b/tests/agent/react/test_memory_node.py
@@ -17,8 +17,10 @@
 )
 from uipath_langchain.agent.react.types import (
     AgentGraphNode,
+    AgentGraphState,
     MemoryConfig,
 )
+from uipath_langchain.agent.react.utils import create_state_with_input
 
 
 class _TopicInput(BaseModel):
@@ -34,9 +36,12 @@ class _TopicLevelInput(BaseModel):
     level: str = Field(default="")
 
 
-def _make_state(**user_fields: Any) -> dict[str, Any]:
-    """Create a state dict with standard internal fields + user fields."""
-    return {"messages": [], "inner_state": {}, **user_fields}
+def _make_state(**user_fields: Any) -> AgentGraphState:
+    """Create a graph state with standard internal fields and user fields."""
+    state_model = create_state_with_input(_TopicInput)
+    return state_model.model_validate(
+        {"messages": [], "inner_state": {}, **user_fields}
+    )
 
 
 class TestBuildSearchFields:
diff --git a/uv.lock b/uv.lock
index d3ac0c391..dea534808 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4498,7 +4498,7 @@ wheels = [
 
 [[package]]
 name = "uipath-langchain"
-version = "0.14.14"
+version = "0.14.15"
 source = { editable = "." }
 dependencies = [
     { name = "a2a-sdk" },