langgraph[patch]: fix annotations in inherited tool schemas (#2236)

fix #2220

.github/workflows/_test_langgraph.yml

@@ -21,7 +21,7 @@ jobs:
           - "latest"
         include:
           - python-version: "3.11"
-            core-version: ">=0.2.39,<0.3.0"
+            core-version: ">=0.2.42,<0.3.0"
 
     defaults:
       run:

libs/langgraph/langgraph/prebuilt/tool_executor.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Sequence, Union, cast
+from typing import Any, Callable, Sequence, Union
 
 from langchain_core.load.serializable import Serializable
 from langchain_core.runnables import RunnableConfig
@@ -101,8 +101,7 @@ class ToolExecutor(RunnableCallable):
     ) -> None:
         super().__init__(self._execute, afunc=self._aexecute, trace=False)
         tools_ = [
-            tool if isinstance(tool, BaseTool) else cast(BaseTool, create_tool(tool))
-            for tool in tools
+            tool if isinstance(tool, BaseTool) else create_tool(tool) for tool in tools
         ]
         self.tools = tools_
         self.tool_map = {t.name: t for t in tools_}

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -34,6 +34,7 @@ from langchain_core.runnables.config import (
 from langchain_core.runnables.utils import Input
 from langchain_core.tools import BaseTool, InjectedToolArg
 from langchain_core.tools import tool as create_tool
+from langchain_core.tools.base import get_all_basemodel_annotations
 from typing_extensions import Annotated, get_args, get_origin
 
 from langgraph.errors import GraphInterrupt
@@ -200,7 +201,7 @@ class ToolNode(RunnableCallable):
         self.messages_key = messages_key
         for tool_ in tools:
             if not isinstance(tool_, BaseTool):
-                tool_ = cast(BaseTool, create_tool(tool_))
+                tool_ = create_tool(tool_)
             self.tools_by_name[tool_.name] = tool_
             self.tool_to_state_args[tool_.name] = _get_state_args(tool_)
             self.tool_to_store_arg[tool_.name] = _get_store_arg(tool_)
@@ -659,7 +660,7 @@ def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
     full_schema = tool.get_input_schema()
     tool_args_to_state_fields: Dict = {}
 
-    for name, type_ in full_schema.__annotations__.items():
+    for name, type_ in get_all_basemodel_annotations(full_schema).items():
         injections = [
             type_arg
             for type_arg in get_args(type_)
@@ -683,7 +684,7 @@ def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
 
 def _get_store_arg(tool: BaseTool) -> Optional[str]:
     full_schema = tool.get_input_schema()
-    for name, type_ in full_schema.__annotations__.items():
+    for name, type_ in get_all_basemodel_annotations(full_schema).items():
         injections = [
             type_arg
             for type_arg in get_args(type_)

libs/langgraph/langgraph/pregel/remote.py

@@ -214,7 +214,7 @@ class RemoteGraph(PregelProtocol):
                     interrupts=tuple(interrupts),
                     state=self._create_state_snapshot(task["state"])
                     if task["state"]
-                    else {"configurable": task["checkpoint"]}
+                    else cast(RunnableConfig, {"configurable": task["checkpoint"]})
                     if task["checkpoint"]
                     else None,
                     result=task.get("result"),

libs/langgraph/pyproject.toml

@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
 
 [tool.poetry.dependencies]
 python = ">=3.9.0,<4.0"
-langchain-core = ">=0.2.39,<0.4"
+langchain-core = "^0.2.42 || ^0.3.14"
 langgraph-checkpoint = "^2.0.0"
 langgraph-sdk = "^0.1.32"
 

libs/langgraph/tests/test_prebuilt.py

@@ -50,6 +50,7 @@ from langgraph.prebuilt.tool_node import (
     TOOL_CALL_ERROR_TEMPLATE,
     InjectedState,
     InjectedStore,
+    _get_state_args,
     _infer_handled_types,
 )
 from langgraph.store.base import BaseStore
@@ -1332,3 +1333,18 @@ async def test_return_direct() -> None:
             id=result["messages"][3].id,
         ),
     ]
+
+
+def test__get_state_args() -> None:
+    class Schema1(BaseModel):
+        a: Annotated[str, InjectedState]
+
+    class Schema2(Schema1):
+        b: Annotated[int, InjectedState("bar")]
+
+    @dec_tool(args_schema=Schema2)
+    def foo(a: str, b: int) -> float:
+        """return"""
+        return 0.0
+
+    assert _get_state_args(foo) == {"a": None, "b": "bar"}