wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
97 lines
2.7 KiB
Python
97 lines
2.7 KiB
Python
"""Capability detection for checkpointer implementations."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING
|
|
|
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
|
|
class Capability(str, Enum):
|
|
"""Capabilities that a checkpointer may support."""
|
|
|
|
PUT = "put"
|
|
PUT_WRITES = "put_writes"
|
|
GET_TUPLE = "get_tuple"
|
|
LIST = "list"
|
|
DELETE_THREAD = "delete_thread"
|
|
DELETE_FOR_RUNS = "delete_for_runs"
|
|
COPY_THREAD = "copy_thread"
|
|
PRUNE = "prune"
|
|
DELTA_CHANNEL_HISTORY = "delta_channel_history"
|
|
|
|
|
|
# Capabilities that every checkpointer must support.
|
|
BASE_CAPABILITIES = frozenset(
|
|
{
|
|
Capability.PUT,
|
|
Capability.PUT_WRITES,
|
|
Capability.GET_TUPLE,
|
|
Capability.LIST,
|
|
Capability.DELETE_THREAD,
|
|
}
|
|
)
|
|
|
|
# Capabilities that are optional extensions.
|
|
EXTENDED_CAPABILITIES = frozenset(
|
|
{
|
|
Capability.DELETE_FOR_RUNS,
|
|
Capability.COPY_THREAD,
|
|
Capability.PRUNE,
|
|
Capability.DELTA_CHANNEL_HISTORY,
|
|
}
|
|
)
|
|
|
|
ALL_CAPABILITIES = BASE_CAPABILITIES | EXTENDED_CAPABILITIES
|
|
|
|
# Maps capability to the async method name on BaseCheckpointSaver (or subclass).
|
|
_CAPABILITY_METHOD_MAP: dict[Capability, str] = {
|
|
Capability.PUT: "aput",
|
|
Capability.PUT_WRITES: "aput_writes",
|
|
Capability.GET_TUPLE: "aget_tuple",
|
|
Capability.LIST: "alist",
|
|
Capability.DELETE_THREAD: "adelete_thread",
|
|
Capability.DELETE_FOR_RUNS: "adelete_for_runs",
|
|
Capability.COPY_THREAD: "acopy_thread",
|
|
Capability.PRUNE: "aprune",
|
|
Capability.DELTA_CHANNEL_HISTORY: "aget_delta_channel_history",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DetectedCapabilities:
|
|
"""Result of capability detection for a checkpointer type."""
|
|
|
|
detected: frozenset[Capability]
|
|
missing: frozenset[Capability]
|
|
|
|
@classmethod
|
|
def from_instance(cls, saver: BaseCheckpointSaver) -> DetectedCapabilities:
|
|
"""Detect capabilities from a checkpointer instance."""
|
|
inner_type = type(saver)
|
|
detected: set[Capability] = set()
|
|
|
|
for cap, method_name in _CAPABILITY_METHOD_MAP.items():
|
|
if _is_overridden(inner_type, method_name):
|
|
detected.add(cap)
|
|
|
|
detected_fs = frozenset(detected)
|
|
return cls(
|
|
detected=detected_fs,
|
|
missing=ALL_CAPABILITIES - detected_fs,
|
|
)
|
|
|
|
|
|
def _is_overridden(inner_type: type, method: str) -> bool:
|
|
"""Check if *method* on *inner_type* differs from the base class default."""
|
|
base = getattr(BaseCheckpointSaver, method, None)
|
|
impl = getattr(inner_type, method, None)
|
|
if base is None or impl is None:
|
|
return impl is not None
|
|
return impl is not base
|