From ffbcdd1eccdf6dd48e14122fcebe76eda65a3da2 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 12 Mar 2025 18:29:55 -0700
Subject: [PATCH] weakref

---
 libs/langgraph/langgraph/graph/state.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index a6f997d25..228aa87a3 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -2,6 +2,7 @@ import inspect
 import logging
 import typing
 import warnings
+import weakref
 from functools import partial
 from inspect import isclass, isfunction, ismethod, signature
 from types import FunctionType
@@ -951,14 +952,18 @@ def _pick_mapper(
 
 
 class _SchemaCoercionMapper:
-    _cache: dict[tuple[Type[Any], int], "_SchemaCoercionMapper"] = {}
+    _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "_SchemaCoercionMapper"]] = (
+        weakref.WeakKeyDictionary()
+    )
 
     def __new__(cls, schema: Type[Any], max_depth: int = 5) -> "_SchemaCoercionMapper":
-        key = (schema, max_depth)
-        if key in cls._cache:
-            return cls._cache[key]
+        if schema not in cls._cache:
+            cls._cache[schema] = {}
+        if max_depth in cls._cache[schema]:
+            return cls._cache[schema][max_depth]
+
         inst = super().__new__(cls)
-        cls._cache[key] = inst
+        cls._cache[schema][max_depth] = inst
         return inst
 
     def __init__(self, schema: Type[Any], max_depth: int = 5):