254 lines
7.1 KiB
Python
254 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import logging
|
|
import sys
|
|
import types
|
|
from collections import deque
|
|
from enum import Enum
|
|
from typing import (
|
|
Annotated,
|
|
Any,
|
|
Literal,
|
|
Union,
|
|
get_args,
|
|
get_origin,
|
|
get_type_hints,
|
|
)
|
|
|
|
from langchain_core import messages as lc_messages
|
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
from pydantic import BaseModel
|
|
from typing_extensions import NotRequired, Required, is_typeddict
|
|
|
|
try:
|
|
from langgraph.checkpoint.serde._msgpack import ( # noqa: F401
|
|
STRICT_MSGPACK_ENABLED,
|
|
)
|
|
except ImportError:
|
|
STRICT_MSGPACK_ENABLED = False
|
|
|
|
_warned_allowlist_unsupported = False
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _supports_checkpointer_allowlist() -> bool:
|
|
return hasattr(BaseCheckpointSaver, "with_allowlist")
|
|
|
|
|
|
_SUPPORTS_ALLOWLIST = _supports_checkpointer_allowlist()
|
|
|
|
|
|
def apply_checkpointer_allowlist(
|
|
checkpointer: Any, allowlist: set[tuple[str, ...]] | None
|
|
) -> Any:
|
|
if not checkpointer or allowlist is None or checkpointer in (True, False):
|
|
return checkpointer
|
|
if not _SUPPORTS_ALLOWLIST:
|
|
global _warned_allowlist_unsupported
|
|
if not _warned_allowlist_unsupported:
|
|
logger.warning(
|
|
"Checkpointer does not support with_allowlist; strict msgpack "
|
|
"allowlist will be skipped."
|
|
)
|
|
_warned_allowlist_unsupported = True
|
|
return checkpointer
|
|
return checkpointer.with_allowlist(allowlist)
|
|
|
|
|
|
def curated_core_allowlist() -> set[tuple[str, ...]]:
|
|
allowlist: set[tuple[str, ...]] = set()
|
|
for name in (
|
|
"BaseMessage",
|
|
"BaseMessageChunk",
|
|
"HumanMessage",
|
|
"HumanMessageChunk",
|
|
"AIMessage",
|
|
"AIMessageChunk",
|
|
"SystemMessage",
|
|
"SystemMessageChunk",
|
|
"ChatMessage",
|
|
"ChatMessageChunk",
|
|
"ToolMessage",
|
|
"ToolMessageChunk",
|
|
"FunctionMessage",
|
|
"FunctionMessageChunk",
|
|
"RemoveMessage",
|
|
):
|
|
cls = getattr(lc_messages, name, None)
|
|
if cls is None:
|
|
continue
|
|
allowlist.add((cls.__module__, cls.__name__))
|
|
|
|
return allowlist
|
|
|
|
|
|
def build_serde_allowlist(
|
|
*,
|
|
schemas: list[type[Any]] | None = None,
|
|
channels: dict[str, Any] | None = None,
|
|
) -> set[tuple[str, ...]]:
|
|
allowlist = curated_core_allowlist()
|
|
if schemas:
|
|
schemas = [schema for schema in schemas if schema is not None]
|
|
return allowlist | collect_allowlist_from_schemas(
|
|
schemas=schemas,
|
|
channels=channels,
|
|
)
|
|
|
|
|
|
def collect_allowlist_from_schemas(
|
|
*,
|
|
schemas: list[type[Any]] | None = None,
|
|
channels: dict[str, Any] | None = None,
|
|
) -> set[tuple[str, ...]]:
|
|
allowlist: set[tuple[str, ...]] = set()
|
|
seen: set[Any] = set()
|
|
seen_ids: set[int] = set()
|
|
|
|
if schemas:
|
|
for schema in schemas:
|
|
_collect_from_type(schema, allowlist, seen, seen_ids)
|
|
|
|
if channels:
|
|
for channel in channels.values():
|
|
value_type = getattr(channel, "ValueType", None)
|
|
if value_type is not None:
|
|
_collect_from_type(value_type, allowlist, seen, seen_ids)
|
|
update_type = getattr(channel, "UpdateType", None)
|
|
if update_type is not None:
|
|
_collect_from_type(update_type, allowlist, seen, seen_ids)
|
|
|
|
return allowlist
|
|
|
|
|
|
def _collect_from_type(
|
|
typ: Any,
|
|
allowlist: set[tuple[str, ...]],
|
|
seen: set[Any],
|
|
seen_ids: set[int],
|
|
) -> None:
|
|
if _already_seen(typ, seen, seen_ids):
|
|
return
|
|
|
|
if typ is Any or typ is None:
|
|
return
|
|
|
|
if typ is Literal:
|
|
return
|
|
|
|
if isinstance(typ, types.UnionType):
|
|
for arg in typ.__args__:
|
|
_collect_from_type(arg, allowlist, seen, seen_ids)
|
|
return
|
|
|
|
origin = get_origin(typ)
|
|
if origin is Union:
|
|
for arg in get_args(typ):
|
|
_collect_from_type(arg, allowlist, seen, seen_ids)
|
|
return
|
|
if origin is Annotated or origin in (Required, NotRequired):
|
|
args = get_args(typ)
|
|
if args:
|
|
_collect_from_type(args[0], allowlist, seen, seen_ids)
|
|
return
|
|
|
|
if origin is Literal:
|
|
return
|
|
|
|
if origin in (list, set, tuple, dict, deque, frozenset):
|
|
for arg in get_args(typ):
|
|
_collect_from_type(arg, allowlist, seen, seen_ids)
|
|
return
|
|
|
|
if hasattr(typ, "__supertype__"):
|
|
_collect_from_type(typ.__supertype__, allowlist, seen, seen_ids)
|
|
return
|
|
|
|
if is_typeddict(typ):
|
|
for field_type in _safe_get_type_hints(typ).values():
|
|
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
|
return
|
|
|
|
if _is_pydantic_model(typ):
|
|
allowlist.add((typ.__module__, typ.__name__))
|
|
field_types = _safe_get_type_hints(typ)
|
|
if field_types:
|
|
for field_type in field_types.values():
|
|
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
|
else:
|
|
for field_type in _pydantic_field_types(typ):
|
|
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
|
return
|
|
|
|
if dataclasses.is_dataclass(typ):
|
|
if typ_name := getattr(typ, "__name__", None):
|
|
allowlist.add((typ.__module__, typ_name))
|
|
field_types = _safe_get_type_hints(typ)
|
|
if field_types:
|
|
for field_type in field_types.values():
|
|
_collect_from_type(field_type, allowlist, seen, seen_ids)
|
|
else:
|
|
for field in dataclasses.fields(typ):
|
|
_collect_from_type(field.type, allowlist, seen, seen_ids)
|
|
return
|
|
|
|
if isinstance(typ, type) and issubclass(typ, Enum):
|
|
allowlist.add((typ.__module__, typ.__name__))
|
|
return
|
|
|
|
|
|
def _already_seen(typ: Any, seen: set[Any], seen_ids: set[int]) -> bool:
|
|
try:
|
|
if typ in seen:
|
|
return True
|
|
seen.add(typ)
|
|
return False
|
|
except TypeError:
|
|
typ_id = id(typ)
|
|
if typ_id in seen_ids:
|
|
return True
|
|
seen_ids.add(typ_id)
|
|
return False
|
|
|
|
|
|
def _safe_get_type_hints(typ: Any) -> dict[str, Any]:
|
|
try:
|
|
module = sys.modules.get(getattr(typ, "__module__", ""))
|
|
globalns = module.__dict__ if module else None
|
|
localns = dict(vars(typ)) if hasattr(typ, "__dict__") else None
|
|
return get_type_hints(
|
|
typ, globalns=globalns, localns=localns, include_extras=True
|
|
)
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _is_pydantic_model(typ: Any) -> bool:
|
|
if not isinstance(typ, type):
|
|
return False
|
|
if issubclass(typ, BaseModel):
|
|
return True
|
|
try:
|
|
from pydantic.v1 import BaseModel as BaseModelV1
|
|
except Exception:
|
|
return False
|
|
return issubclass(typ, BaseModelV1)
|
|
|
|
|
|
def _pydantic_field_types(typ: type[Any]) -> list[Any]:
|
|
if hasattr(typ, "model_fields"):
|
|
return [
|
|
field.annotation
|
|
for field in typ.model_fields.values()
|
|
if getattr(field, "annotation", None) is not None
|
|
]
|
|
if hasattr(typ, "__fields__"):
|
|
return [
|
|
field.outer_type_
|
|
for field in typ.__fields__.values()
|
|
if getattr(field, "outer_type_", None) is not None
|
|
]
|
|
return []
|