coco-framework 0.1.0__py3-none-any.whl

coco_graph/_internal/_scratchpad.py

@@ -0,0 +1,19 @@
+import dataclasses
+from collections.abc import Callable
+from typing import Any
+
+from coco_graph.types import _DC_KWARGS
+
+
+@dataclasses.dataclass(**_DC_KWARGS)
+class PregelScratchpad:
+    step: int
+    stop: int
+    # call
+    call_counter: Callable[[], int]
+    # interrupt
+    interrupt_counter: Callable[[], int]
+    get_null_resume: Callable[[bool], Any]
+    resume: list[Any]
+    # subgraph
+    subgraph_counter: Callable[[], int]

coco_graph/_internal/_serde.py

@@ -0,0 +1,253 @@
+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 coco_core import messages as lc_messages
+from coco_graph.checkpoint.base import BaseCheckpointSaver
+from pydantic import BaseModel
+from typing_extensions import NotRequired, Required, is_typeddict
+
+try:
+    from coco_graph.checkpoint.serde._msgpack import (
+        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 []

coco_graph/_internal/_timeout.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import Literal
+
+from coco_graph.types import TimeoutPolicy
+
+_SYNC_TIMEOUT_PREFIX = (
+    "Node timeouts are only supported for async nodes because sync Python "
+    "execution cannot be safely cancelled in-process."
+)
+
+
+def coerce_timeout_policy(
+    value: float | timedelta | TimeoutPolicy | None,
+) -> TimeoutPolicy | None:
+    """Normalize a timeout value to positive-second policy fields."""
+    return TimeoutPolicy.coerce(value)
+
+
+def sync_timeout_unsupported(
+    name: str, *, kind: Literal["Node", "Task"] = "Node"
+) -> ValueError:
+    """Build the canonical error for using `timeout` with a sync target."""
+    return ValueError(f"{_SYNC_TIMEOUT_PREFIX} {kind} {name!r} is sync.")

coco_graph/_internal/_typing.py

@@ -0,0 +1,54 @@
+"""Private typing utilities for CocoGraph."""
+
+from __future__ import annotations
+
+from dataclasses import Field
+from typing import Any, ClassVar, Protocol, TypeAlias
+
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+
+class TypedDictLikeV1(Protocol):
+    """Protocol to represent types that behave like TypedDicts
+
+    Version 1: using `ClassVar` for keys."""
+
+    __required_keys__: ClassVar[frozenset[str]]
+    __optional_keys__: ClassVar[frozenset[str]]
+
+
+class TypedDictLikeV2(Protocol):
+    """Protocol to represent types that behave like TypedDicts
+
+    Version 2: not using `ClassVar` for keys."""
+
+    __required_keys__: frozenset[str]
+    __optional_keys__: frozenset[str]
+
+
+class DataclassLike(Protocol):
+    """Protocol to represent types that behave like dataclasses.
+
+    Inspired by the private _DataclassT from dataclasses that uses a similar protocol as a bound."""
+
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
+
+
+StateLike: TypeAlias = TypedDictLikeV1 | TypedDictLikeV2 | DataclassLike | BaseModel
+"""Type alias for state-like types.
+
+It can either be a `TypedDict`, `dataclass`, or Pydantic `BaseModel`.
+Note: we cannot use either `TypedDict` or `dataclass` directly due to limitations in type checking.
+"""
+
+MISSING = object()
+"""Unset sentinel value."""
+
+
+class DeprecatedKwargs(TypedDict):
+    """TypedDict to use for extra keyword arguments, enabling type checking warnings for deprecated arguments."""
+
+
+EMPTY_SEQ: tuple[str, ...] = tuple()
+"""An empty sequence of strings."""

coco_graph/cache/base/__init__.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from typing import Generic, TypeVar
+
+from coco_graph.checkpoint.serde.base import SerializerProtocol
+from coco_graph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+ValueT = TypeVar("ValueT")
+Namespace = tuple[str, ...]
+FullKey = tuple[Namespace, str]
+
+
+class BaseCache(ABC, Generic[ValueT]):
+    """Base class for a cache."""
+
+    serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=False)
+
+    def __init__(self, *, serde: SerializerProtocol | None = None) -> None:
+        """Initialize the cache with a serializer."""
+        self.serde = serde or self.serde
+
+    @abstractmethod
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+
+    @abstractmethod
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+
+    @abstractmethod
+    def set(self, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys and TTLs."""
+
+    @abstractmethod
+    async def aset(self, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys and TTLs."""
+
+    @abstractmethod
+    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+
+    @abstractmethod
+    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Asynchronously delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""

coco_graph/cache/memory/__init__.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import datetime
+import threading
+from collections.abc import Mapping, Sequence
+
+from coco_graph.cache.base import BaseCache, FullKey, Namespace, ValueT
+from coco_graph.checkpoint.serde.base import SerializerProtocol
+
+
+class InMemoryCache(BaseCache[ValueT]):
+    def __init__(self, *, serde: SerializerProtocol | None = None):
+        super().__init__(serde=serde)
+        self._cache: dict[Namespace, dict[str, tuple[str, bytes, float | None]]] = {}
+        self._lock = threading.RLock()
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        with self._lock:
+            if not keys:
+                return {}
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            values: dict[FullKey, ValueT] = {}
+            for ns_tuple, key in keys:
+                ns = Namespace(ns_tuple)
+                if ns in self._cache and key in self._cache[ns]:
+                    enc, val, expiry = self._cache[ns][key]
+                    if expiry is None or now < expiry:
+                        values[(ns, key)] = self.serde.loads_typed((enc, val))
+                    else:
+                        del self._cache[ns][key]
+            return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return self.get(keys)
+
+    def set(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys."""
+        with self._lock:
+            now = datetime.datetime.now(datetime.timezone.utc)
+            for (ns, key), (value, ttl) in keys.items():
+                if ttl is not None:
+                    delta = datetime.timedelta(seconds=ttl)
+                    expiry: float | None = (now + delta).timestamp()
+                else:
+                    expiry = None
+                if ns not in self._cache:
+                    self._cache[ns] = {}
+                self._cache[ns][key] = (
+                    *self.serde.dumps_typed(value),
+                    expiry,
+                )
+
+    async def aset(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys."""
+        self.set(keys)
+
+    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        with self._lock:
+            if namespaces is None:
+                self._cache.clear()
+            else:
+                for ns in namespaces:
+                    if ns in self._cache:
+                        del self._cache[ns]
+
+    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Asynchronously delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        self.clear(namespaces)

coco_graph/cache/redis/__init__.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from coco_graph.cache.base import BaseCache, FullKey, Namespace, ValueT
+from coco_graph.checkpoint.serde.base import SerializerProtocol
+
+
+class RedisCache(BaseCache[ValueT]):
+    """Redis-based cache implementation with TTL support."""
+
+    def __init__(
+        self,
+        redis: Any,
+        *,
+        serde: SerializerProtocol | None = None,
+        prefix: str = "coco_graph:cache:",
+    ) -> None:
+        """Initialize the cache with a Redis client.
+
+        Args:
+            redis: Redis client instance (sync or async)
+            serde: Serializer to use for values
+            prefix: Key prefix for all cached values
+        """
+        super().__init__(serde=serde)
+        self.redis = redis
+        self.prefix = prefix
+
+    def _make_key(self, ns: Namespace, key: str) -> str:
+        """Create a Redis key from namespace and key."""
+        ns_str = ":".join(ns) if ns else ""
+        return f"{self.prefix}{ns_str}:{key}" if ns_str else f"{self.prefix}{key}"
+
+    def _parse_key(self, redis_key: str) -> tuple[Namespace, str]:
+        """Parse a Redis key back to namespace and key."""
+        if not redis_key.startswith(self.prefix):
+            raise ValueError(
+                f"Key {redis_key} does not start with prefix {self.prefix}"
+            )
+
+        remaining = redis_key[len(self.prefix) :]
+        if ":" in remaining:
+            parts = remaining.split(":")
+            key = parts[-1]
+            ns_parts = parts[:-1]
+            return (tuple(ns_parts), key)
+        else:
+            return (tuple(), remaining)
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        if not keys:
+            return {}
+
+        # Build Redis keys
+        redis_keys = [self._make_key(ns, key) for ns, key in keys]
+
+        # Get values from Redis using MGET
+        try:
+            raw_values = self.redis.mget(redis_keys)
+        except Exception:
+            # If Redis is unavailable, return empty dict
+            return {}
+
+        values: dict[FullKey, ValueT] = {}
+        for i, raw_value in enumerate(raw_values):
+            if raw_value is not None:
+                try:
+                    # Deserialize the value
+                    encoding, data = raw_value.split(b":", 1)
+                    values[keys[i]] = self.serde.loads_typed((encoding.decode(), data))
+                except Exception:
+                    # Skip corrupted entries
+                    continue
+
+        return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return self.get(keys)
+
+    def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys and TTLs."""
+        if not mapping:
+            return
+
+        # Use pipeline for efficient batch operations
+        pipe = self.redis.pipeline()
+
+        for (ns, key), (value, ttl) in mapping.items():
+            redis_key = self._make_key(ns, key)
+            encoding, data = self.serde.dumps_typed(value)
+
+            # Store as "encoding:data" format
+            serialized_value = f"{encoding}:".encode() + data
+
+            if ttl is not None:
+                pipe.setex(redis_key, ttl, serialized_value)
+            else:
+                pipe.set(redis_key, serialized_value)
+
+        try:
+            pipe.execute()
+        except Exception:
+            # Silently fail if Redis is unavailable
+            pass
+
+    async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys and TTLs."""
+        self.set(mapping)
+
+    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        try:
+            if namespaces is None:
+                # Clear all keys with our prefix
+                pattern = f"{self.prefix}*"
+                keys = self.redis.keys(pattern)
+                if keys:
+                    self.redis.delete(*keys)
+            else:
+                # Clear specific namespaces
+                keys_to_delete = []
+                for ns in namespaces:
+                    ns_str = ":".join(ns) if ns else ""
+                    pattern = (
+                        f"{self.prefix}{ns_str}:*" if ns_str else f"{self.prefix}*"
+                    )
+                    keys = self.redis.keys(pattern)
+                    keys_to_delete.extend(keys)
+
+                if keys_to_delete:
+                    self.redis.delete(*keys_to_delete)
+        except Exception:
+            # Silently fail if Redis is unavailable
+            pass
+
+    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Asynchronously delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        self.clear(namespaces)