Avoid repeated runtime calls to get_type_hints (#4888)

libs/langgraph/langgraph/graph/state.py

@@ -72,7 +72,11 @@ from langgraph.pregel.write import (
 )
 from langgraph.store.base import BaseStore
 from langgraph.types import All, CachePolicy, Checkpointer, Command, RetryPolicy, Send
-from langgraph.utils.fields import get_field_default, get_update_as_tuples
+from langgraph.utils.fields import (
+    get_cached_annotated_keys,
+    get_field_default,
+    get_update_as_tuples,
+)
 from langgraph.utils.pydantic import create_model
 from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
 
@@ -885,7 +889,7 @@ class CompiledStateGraph(Pregel):
                     else:
                         updates.extend(_get_updates(i) or ())
                 return updates
-            elif (t := type(input)) and get_type_hints(t):
+            elif (t := type(input)) and get_cached_annotated_keys(t):
                 return get_update_as_tuples(input, output_keys)
             else:
                 msg = create_error_message(

libs/langgraph/langgraph/types.py

@@ -14,7 +14,6 @@ from typing import (
     TypeVar,
     Union,
     cast,
-    get_type_hints,
 )
 
 from langchain_core.runnables import Runnable, RunnableConfig
@@ -23,7 +22,7 @@ from xxhash import xxh3_128_hexdigest
 
 from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
 from langgraph.utils.cache import default_cache_key
-from langgraph.utils.fields import get_update_as_tuples
+from langgraph.utils.fields import get_cached_annotated_keys, get_update_as_tuples
 
 if TYPE_CHECKING:
     from langgraph.pregel.protocol import PregelProtocol
@@ -350,8 +349,8 @@ class Command(Generic[N], ToolOutputMixin):
             for t in self.update
         ):
             return self.update
-        elif hints := get_type_hints(type(self.update)):
-            return get_update_as_tuples(self.update, tuple(hints.keys()))
+        elif keys := get_cached_annotated_keys(type(self.update)):
+            return get_update_as_tuples(self.update, keys)
         elif self.update is not None:
             return [("__root__", self.update)]
         else:

libs/langgraph/langgraph/utils/fields.py

@@ -1,4 +1,6 @@
 import dataclasses
+import types
+import weakref
 from collections.abc import Generator, Sequence
 from typing import Annotated, Any, Optional, Union, get_type_hints
 
@@ -178,3 +180,24 @@ def get_update_as_tuples(input: Any, keys: Sequence[str]) -> list[tuple[str, Any
             or (keep is not None and k in keep)
         )
     ]
+
+
+ANNOTATED_KEYS_CACHE: weakref.WeakKeyDictionary[type[Any], tuple[str, ...]] = (
+    weakref.WeakKeyDictionary()
+)
+
+
+def get_cached_annotated_keys(obj: type[Any]) -> tuple[str, ...]:
+    """Return cached annotated keys for a Python class."""
+    if obj in ANNOTATED_KEYS_CACHE:
+        return ANNOTATED_KEYS_CACHE[obj]
+    if isinstance(obj, type):
+        keys: list[str] = []
+        for base in reversed(obj.__mro__):
+            ann = base.__dict__.get("__annotations__")
+            if ann is None or isinstance(ann, types.GetSetDescriptorType):
+                continue
+            keys.extend(ann.keys())
+        return ANNOTATED_KEYS_CACHE.setdefault(obj, tuple(keys))
+    else:
+        raise TypeError(f"Expected a type, got {type(obj)}. ")