From d661d528b21625a0bd93b828cfd204317f4fb9ae Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 30 May 2025 16:30:53 -0700
Subject: [PATCH] Avoid repeated runtime calls to get_type_hints

- Replace get_type_hints logic with much simpler implementation which only collects annotated keys (not the annotations themselves)
- Cache annotations in a WeakKeyDictionary
---
 libs/langgraph/langgraph/graph/state.py  |  8 ++++++--
 libs/langgraph/langgraph/types.py        |  7 +++----
 libs/langgraph/langgraph/utils/fields.py | 23 +++++++++++++++++++++++
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index c1b54f7a2..907d158b0 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/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(
diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py
index 0d02bac38..d6db21db9 100644
--- a/libs/langgraph/langgraph/types.py
+++ b/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:
diff --git a/libs/langgraph/langgraph/utils/fields.py b/libs/langgraph/langgraph/utils/fields.py
index e13c697d8..4f8cd78cc 100644
--- a/libs/langgraph/langgraph/utils/fields.py
+++ b/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)}. ")