Update caching logic

Author: Natim

src/openai/_models.py

@@ -19,6 +19,7 @@
     override,
     runtime_checkable,
 )
+from weakref import WeakKeyDictionary
 
 import pydantic
 from pydantic.fields import FieldInfo
@@ -77,6 +78,8 @@
 
 ReprArgs = Sequence[Tuple[Optional[str], Any]]
 
+_DISCRIMINATOR_CACHE: "WeakKeyDictionary[type, DiscriminatorDetails]" = WeakKeyDictionary()
+
 
 @runtime_checkable
 class _ConfigProtocol(Protocol):
@@ -593,11 +596,6 @@ def construct_type(*, value: object, type_: object, metadata: Optional[List[Any]
     return value
 
 
-@runtime_checkable
-class CachedDiscriminatorType(Protocol):
-    __discriminator__: DiscriminatorDetails
-
-
 class DiscriminatorDetails:
     field_name: str
     """The name of the discriminator field in the variant class, e.g.
@@ -640,8 +638,9 @@ def __init__(
 
 
 def _build_discriminated_union_meta(*, union: type, meta_annotations: tuple[Any, ...]) -> DiscriminatorDetails | None:
-    if isinstance(union, CachedDiscriminatorType):
-        return union.__discriminator__
+    cached_discriminator = _DISCRIMINATOR_CACHE.get(union)
+    if cached_discriminator is not None:
+        return cached_discriminator
 
     discriminator_field_name: str | None = None
 
@@ -694,7 +693,7 @@ def _build_discriminated_union_meta(*, union: type, meta_annotations: tuple[Any,
         discriminator_field=discriminator_field_name,
         discriminator_alias=discriminator_alias,
     )
-    cast(CachedDiscriminatorType, Annotated[union, details])
+    _DISCRIMINATOR_CACHE[union] = details
     return details