From 22f5367af73fa5aa79ba0c378564d1d23bb08d9e Mon Sep 17 00:00:00 2001
From: gbaian10 <34255899+gbaian10@users.noreply.github.com>
Date: Fri, 23 Aug 2024 00:32:21 +0800
Subject: [PATCH] langgraph: fix add_node input schema error (#1332)

---
 libs/langgraph/langgraph/graph/state.py | 12 ++++---
 libs/langgraph/tests/test_state.py      | 42 ++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index d846bc8d3..3ff07b15c 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -1,3 +1,4 @@
+import inspect
 import logging
 import typing
 import warnings
@@ -330,10 +331,13 @@ class StateGraph(Graph):
                 hints := get_type_hints(action.__call__) or get_type_hints(action)
             ):
                 if input is None:
-                    input_hint = hints[list(hints.keys())[0]]
-                    if isinstance(input_hint, type) and get_type_hints(input_hint):
-                        input = input_hint
-        except TypeError:
+                    first_parameter_name = next(
+                        iter(inspect.signature(action).parameters.keys())
+                    )
+                    if input_hint := hints.get(first_parameter_name):
+                        if isinstance(input_hint, type) and get_type_hints(input_hint):
+                            input = input_hint
+        except (TypeError, StopIteration):
             pass
         if input is not None:
             self._add_schema(input)
diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py
index ecf35eb2a..4088e38c9 100644
--- a/libs/langgraph/tests/test_state.py
+++ b/libs/langgraph/tests/test_state.py
@@ -2,10 +2,11 @@ from typing import Annotated as Annotated2
 from typing import Any
 
 import pytest
+from langchain_core.runnables import RunnableConfig
 from pydantic.v1 import BaseModel
 from typing_extensions import Annotated, TypedDict
 
-from langgraph.graph.state import _warn_invalid_state_schema
+from langgraph.graph.state import StateGraph, _warn_invalid_state_schema
 
 
 class State(BaseModel):
@@ -46,3 +47,42 @@ def test_doesnt_warn_valid_schema(schema: Any):
     # Assert the function does not raise a warning
     with pytest.warns(None):
         _warn_invalid_state_schema(schema)
+
+
+def test_state_schema_with_type_hint():
+    class InputState(TypedDict):
+        question: str
+
+    class OutputState(TypedDict):
+        input_state: InputState
+
+    def complete_hint(state: InputState) -> OutputState:
+        return {"input_state": state}
+
+    def miss_first_hint(state, config: RunnableConfig) -> OutputState:
+        return {"input_state": state}
+
+    def only_return_hint(state, config) -> OutputState:
+        return {"input_state": state}
+
+    def miss_all_hint(state, config):
+        return {"input_state": state}
+
+    graph = StateGraph(input=InputState, output=OutputState)
+    actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
+
+    for action in actions:
+        graph.add_node(action)
+
+    graph.set_entry_point(actions[0].__name__)
+    for i in range(len(actions) - 1):
+        graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
+    graph.set_finish_point(actions[-1].__name__)
+
+    graph = graph.compile()
+
+    input_state = InputState(question="Hello World!")
+    output_state = OutputState(input_state=input_state)
+    for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
+        node_name = actions[i].__name__
+        assert c[node_name] == output_state