coco-framework 0.1.0__py3-none-any.whl

coco_graph/cache/sqlite/__init__.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import asyncio
+import datetime
+import sqlite3
+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 SqliteCache(BaseCache[ValueT]):
+    """File-based cache using SQLite."""
+
+    def __init__(
+        self,
+        *,
+        path: str,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        """Initialize the cache with a file path."""
+        super().__init__(serde=serde)
+        # SQLite backing store
+        self._conn = sqlite3.connect(
+            path,
+            check_same_thread=False,
+        )
+        # Serialize access to the shared connection across threads
+        self._lock = threading.RLock()
+        # Better concurrency & atomicity
+        self._conn.execute("PRAGMA journal_mode=WAL;")
+        # Schema: key -> (expiry, encoding, value)
+        self._conn.execute(
+            """CREATE TABLE IF NOT EXISTS cache (
+                ns TEXT,
+                key TEXT,
+                expiry REAL,
+                encoding TEXT NOT NULL,
+                val BLOB NOT NULL,
+                PRIMARY KEY (ns, key)
+            )"""
+        )
+        self._conn.commit()
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        with self._lock, self._conn:
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            if not keys:
+                return {}
+            placeholders = ",".join("(?, ?)" for _ in keys)
+            params: list[str] = []
+            for ns_tuple, key in keys:
+                params.extend((",".join(ns_tuple), key))
+            cursor = self._conn.execute(
+                f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})",
+                tuple(params),
+            )
+            values: dict[FullKey, ValueT] = {}
+            rows = cursor.fetchall()
+            for ns, key, expiry, encoding, raw in rows:
+                if expiry is not None and now > expiry:
+                    # purge expired entry
+                    self._conn.execute(
+                        "DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key)
+                    )
+                    continue
+                values[(tuple(ns.split(",")), key)] = self.serde.loads_typed(
+                    (encoding, raw)
+                )
+            return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return await asyncio.to_thread(self.get, keys)
+
+    def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys and TTLs."""
+        with self._lock, self._conn:
+            now = datetime.datetime.now(datetime.timezone.utc)
+            for key, (value, ttl) in mapping.items():
+                if ttl is not None:
+                    delta = datetime.timedelta(seconds=ttl)
+                    expiry: float | None = (now + delta).timestamp()
+                else:
+                    expiry = None
+                encoding, raw = self.serde.dumps_typed(value)
+                self._conn.execute(
+                    "INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)",
+                    (",".join(key[0]), key[1], expiry, encoding, raw),
+                )
+
+    async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys and TTLs."""
+        await asyncio.to_thread(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."""
+        with self._lock, self._conn:
+            if namespaces is None:
+                self._conn.execute("DELETE FROM cache")
+            else:
+                placeholders = ",".join("?" for _ in namespaces)
+                self._conn.execute(
+                    f"DELETE FROM cache WHERE (ns) IN ({placeholders})",
+                    tuple(",".join(key) for key in namespaces),
+                )
+
+    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."""
+        await asyncio.to_thread(self.clear, namespaces)
+
+    def __del__(self) -> None:
+        try:
+            self._conn.close()
+        except Exception:
+            pass

coco_graph/callbacks.py

@@ -0,0 +1,394 @@
+"""Graph lifecycle callback interfaces and event payloads.
+
+This module defines the public callback surface for observing CocoGraph-specific
+lifecycle transitions such as interrupt and resume.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, Literal, TypeAlias, TypeVar
+from uuid import UUID
+
+from coco_core.callbacks import BaseCallbackHandler, BaseCallbackManager
+from coco_core.callbacks.manager import ahandle_event, handle_event
+from coco_core.runnables import RunnableConfig
+
+from coco_graph.types import Interrupt
+
+__all__ = (
+    "GraphCallbackHandler",
+    "GraphInterruptEvent",
+    "GraphLifecycleEvent",
+    "GraphLifecycleStatus",
+    "GraphResumeEvent",
+    "get_async_graph_callback_manager_for_config",
+    "get_sync_graph_callback_manager_for_config",
+)
+
+
+GraphLifecycleStatus: TypeAlias = Literal[
+    "input",
+    "pending",
+    "done",
+    "interrupt_before",
+    "interrupt_after",
+    "out_of_steps",
+]
+"""Allowed lifecycle statuses reported in graph lifecycle callback events."""
+
+
+@dataclass(frozen=True)
+class GraphInterruptEvent:
+    """Graph lifecycle event emitted when execution pauses for interrupts."""
+
+    run_id: UUID | None
+    """Run id for the current graph execution, if available."""
+
+    status: GraphLifecycleStatus
+    """Loop status when the interrupt was captured."""
+
+    checkpoint_id: str
+    """Checkpoint id associated with the interrupted execution."""
+
+    checkpoint_ns: tuple[str, ...]
+    """Checkpoint namespace path for the current graph or subgraph."""
+
+    interrupts: tuple[Interrupt, ...]
+    """Interrupt payloads that caused the graph to pause."""
+
+
+@dataclass(frozen=True)
+class GraphResumeEvent:
+    """Graph lifecycle event emitted when execution resumes from a checkpoint."""
+
+    run_id: UUID | None
+    """Run id for the current graph execution, if available."""
+
+    status: GraphLifecycleStatus
+    """Loop status when the resume was captured."""
+
+    checkpoint_id: str
+    """Checkpoint id the graph resumed from."""
+
+    checkpoint_ns: tuple[str, ...]
+    """Checkpoint namespace path for the current graph or subgraph."""
+
+
+GraphLifecycleEvent: TypeAlias = GraphInterruptEvent | GraphResumeEvent
+"""Union of all public graph lifecycle callback event payloads.
+
+Use this alias when a callback or helper can receive either interrupt or resume
+lifecycle events.
+"""
+
+
+class GraphCallbackHandler(BaseCallbackHandler):
+    """Base class for graph-level lifecycle callbacks.
+
+    Subclass this handler to observe graph lifecycle transitions that are
+    specific to CocoGraph execution, rather than generic Coco runnable
+    callbacks.
+
+    Instances can be passed through `config["callbacks"]` when invoking a
+    graph. Only handlers that inherit from `GraphCallbackHandler` receive these
+    lifecycle events.
+    """
+
+    def on_interrupt(self, event: GraphInterruptEvent) -> Any:
+        """Run when graph execution pauses due to one or more interrupts.
+
+        Args:
+            event: Interrupt lifecycle event payload.
+        """
+
+    def on_resume(self, event: GraphResumeEvent) -> Any:
+        """Run when graph execution resumes from a persisted checkpoint.
+
+        Args:
+            event: Resume lifecycle event payload.
+        """
+
+
+_MISSING = object()
+
+
+def _filter_graph_handlers(
+    handlers: list[BaseCallbackHandler],
+) -> list[GraphCallbackHandler]:
+    return [h for h in handlers if isinstance(h, GraphCallbackHandler)]
+
+
+def _init_base_manager(
+    manager: BaseCallbackManager,
+    handlers: Sequence[GraphCallbackHandler] | None,
+    inheritable_handlers: Sequence[GraphCallbackHandler] | None,
+    parent_run_id: UUID | None,
+    *,
+    tags: list[str] | None,
+    inheritable_tags: list[str] | None,
+    metadata: dict[str, Any] | None,
+    inheritable_metadata: dict[str, Any] | None,
+    run_id: UUID | None,
+) -> None:
+    base_handlers: list[BaseCallbackHandler] = []
+    base_inheritable_handlers: list[BaseCallbackHandler] = []
+    if handlers is not None:
+        base_handlers.extend(handlers)
+    if inheritable_handlers is not None:
+        base_inheritable_handlers.extend(inheritable_handlers)
+    BaseCallbackManager.__init__(
+        manager,
+        handlers=base_handlers,
+        inheritable_handlers=base_inheritable_handlers,
+        parent_run_id=parent_run_id,
+        tags=tags,
+        inheritable_tags=inheritable_tags,
+        metadata=metadata,
+        inheritable_metadata=inheritable_metadata,
+    )
+    manager.run_id = run_id  # type: ignore[attr-defined]
+
+
+def _configure_graph_callbacks(
+    cls: type[_GraphManagerT],
+    callbacks: object | None,
+    *,
+    run_id: UUID | None,
+) -> _GraphManagerT:
+    if callbacks is None:
+        return cls(run_id=run_id)
+    if isinstance(callbacks, cls):
+        return callbacks.copy(run_id=run_id)
+    if isinstance(callbacks, (_GraphCallbackManager, _AsyncGraphCallbackManager)):
+        # Cross-type: extract handlers into the requested cls.
+        return cls(
+            handlers=_filter_graph_handlers(callbacks.handlers),
+            inheritable_handlers=_filter_graph_handlers(callbacks.inheritable_handlers),
+            parent_run_id=callbacks.parent_run_id,
+            tags=callbacks.tags.copy(),
+            inheritable_tags=callbacks.inheritable_tags.copy(),
+            metadata=callbacks.metadata.copy(),
+            inheritable_metadata=callbacks.inheritable_metadata.copy(),
+            run_id=run_id,
+        )
+    if isinstance(callbacks, BaseCallbackManager):
+        return cls(
+            handlers=_filter_graph_handlers(callbacks.handlers),
+            inheritable_handlers=_filter_graph_handlers(callbacks.inheritable_handlers),
+            parent_run_id=callbacks.parent_run_id,
+            tags=callbacks.tags.copy(),
+            inheritable_tags=callbacks.inheritable_tags.copy(),
+            metadata=callbacks.metadata.copy(),
+            inheritable_metadata=callbacks.inheritable_metadata.copy(),
+            run_id=run_id,
+        )
+    if isinstance(callbacks, GraphCallbackHandler):
+        return cls((callbacks,), run_id=run_id)
+    if isinstance(callbacks, (str, bytes)) or not isinstance(callbacks, Sequence):
+        raise TypeError("callbacks must be a handler, sequence, or manager")
+    return cls(_filter_graph_handlers(list(callbacks)), run_id=run_id)
+
+
+def _copy_graph_manager(
+    manager: _GraphCallbackManager | _AsyncGraphCallbackManager,
+    cls: type[_GraphManagerT],
+    run_id: UUID | None | object,
+) -> _GraphManagerT:
+    resolved_run_id: UUID | None
+    if run_id is _MISSING:
+        resolved_run_id = manager.run_id
+    else:
+        if run_id is not None and not isinstance(run_id, UUID):
+            raise TypeError("run_id must be a UUID or None")
+        resolved_run_id = run_id
+
+    return cls(
+        handlers=_filter_graph_handlers(manager.handlers),
+        inheritable_handlers=_filter_graph_handlers(manager.inheritable_handlers),
+        parent_run_id=manager.parent_run_id,
+        tags=manager.tags.copy(),
+        inheritable_tags=manager.inheritable_tags.copy(),
+        metadata=manager.metadata.copy(),
+        inheritable_metadata=manager.inheritable_metadata.copy(),
+        run_id=resolved_run_id,
+    )
+
+
+class _GraphCallbackManager(BaseCallbackManager):
+    """Sync dispatcher for graph lifecycle events."""
+
+    run_id: UUID | None
+
+    def __init__(
+        self,
+        handlers: Sequence[GraphCallbackHandler] | None = None,
+        inheritable_handlers: Sequence[GraphCallbackHandler] | None = None,
+        parent_run_id: UUID | None = None,
+        *,
+        tags: list[str] | None = None,
+        inheritable_tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        inheritable_metadata: dict[str, Any] | None = None,
+        run_id: UUID | None = None,
+    ) -> None:
+        _init_base_manager(
+            self,
+            handlers,
+            inheritable_handlers,
+            parent_run_id,
+            tags=tags,
+            inheritable_tags=inheritable_tags,
+            metadata=metadata,
+            inheritable_metadata=inheritable_metadata,
+            run_id=run_id,
+        )
+
+    def copy(
+        self,
+        *,
+        run_id: UUID | None | object = _MISSING,
+    ) -> _GraphCallbackManager:
+        return _copy_graph_manager(self, _GraphCallbackManager, run_id)
+
+    @classmethod
+    def configure(
+        cls,
+        callbacks: object | None = None,
+        *,
+        run_id: UUID | None = None,
+    ) -> _GraphCallbackManager:
+        return _configure_graph_callbacks(cls, callbacks, run_id=run_id)
+
+    def on_interrupt(self, event: GraphInterruptEvent) -> None:
+        handle_event(
+            self.handlers,
+            "on_interrupt",
+            None,
+            event,
+        )
+
+    def on_resume(self, event: GraphResumeEvent) -> None:
+        handle_event(
+            self.handlers,
+            "on_resume",
+            None,
+            event,
+        )
+
+
+class _AsyncGraphCallbackManager(BaseCallbackManager):
+    """Async dispatcher for graph lifecycle events."""
+
+    run_id: UUID | None
+
+    @property
+    def is_async(self) -> bool:
+        """Return whether the manager is async."""
+        return True
+
+    def __init__(
+        self,
+        handlers: Sequence[GraphCallbackHandler] | None = None,
+        inheritable_handlers: Sequence[GraphCallbackHandler] | None = None,
+        parent_run_id: UUID | None = None,
+        *,
+        tags: list[str] | None = None,
+        inheritable_tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        inheritable_metadata: dict[str, Any] | None = None,
+        run_id: UUID | None = None,
+    ) -> None:
+        _init_base_manager(
+            self,
+            handlers,
+            inheritable_handlers,
+            parent_run_id,
+            tags=tags,
+            inheritable_tags=inheritable_tags,
+            metadata=metadata,
+            inheritable_metadata=inheritable_metadata,
+            run_id=run_id,
+        )
+
+    def copy(
+        self,
+        *,
+        run_id: UUID | None | object = _MISSING,
+    ) -> _AsyncGraphCallbackManager:
+        return _copy_graph_manager(self, _AsyncGraphCallbackManager, run_id)
+
+    @classmethod
+    def configure(
+        cls,
+        callbacks: object | None = None,
+        *,
+        run_id: UUID | None = None,
+    ) -> _AsyncGraphCallbackManager:
+        return _configure_graph_callbacks(cls, callbacks, run_id=run_id)
+
+    async def on_interrupt(self, event: GraphInterruptEvent) -> None:
+        await ahandle_event(
+            self.handlers,
+            "on_interrupt",
+            None,
+            event,
+        )
+
+    async def on_resume(self, event: GraphResumeEvent) -> None:
+        await ahandle_event(
+            self.handlers,
+            "on_resume",
+            None,
+            event,
+        )
+
+
+_GraphManagerT = TypeVar(
+    "_GraphManagerT", _GraphCallbackManager, _AsyncGraphCallbackManager
+)
+
+GraphCallbacks: TypeAlias = (
+    _GraphCallbackManager
+    | _AsyncGraphCallbackManager
+    | BaseCallbackManager
+    | GraphCallbackHandler
+    | Sequence[BaseCallbackHandler]
+    | Sequence[GraphCallbackHandler]
+    | None
+)
+
+
+def get_sync_graph_callback_manager_for_config(
+    config: RunnableConfig,
+    *,
+    run_id: UUID | None = None,
+) -> _GraphCallbackManager:
+    """Build a sync graph lifecycle callback manager from a runnable config.
+
+    This helper filters `config["callbacks"]` down to handlers that inherit
+    from `GraphCallbackHandler` and binds the provided `run_id` onto the
+    returned manager.
+    """
+    return _GraphCallbackManager.configure(
+        config.get("callbacks"),
+        run_id=run_id,
+    )
+
+
+def get_async_graph_callback_manager_for_config(
+    config: RunnableConfig,
+    *,
+    run_id: UUID | None = None,
+) -> _AsyncGraphCallbackManager:
+    """Build an async graph lifecycle callback manager from a runnable config.
+
+    This helper filters `config["callbacks"]` down to handlers that inherit
+    from `GraphCallbackHandler` and binds the provided `run_id` onto the
+    returned manager.
+    """
+    return _AsyncGraphCallbackManager.configure(
+        config.get("callbacks"),
+        run_id=run_id,
+    )

coco_graph/channels/__init__.py

@@ -0,0 +1,29 @@
+from coco_graph.channels.any_value import AnyValue
+from coco_graph.channels.base import BaseChannel
+from coco_graph.channels.binop import BinaryOperatorAggregate
+from coco_graph.channels.delta import DeltaChannel
+from coco_graph.channels.ephemeral_value import EphemeralValue
+from coco_graph.channels.last_value import LastValue, LastValueAfterFinish
+from coco_graph.channels.named_barrier_value import (
+    NamedBarrierValue,
+    NamedBarrierValueAfterFinish,
+)
+from coco_graph.channels.topic import Topic
+from coco_graph.channels.untracked_value import UntrackedValue
+
+__all__ = (
+    # base
+    "BaseChannel",
+    # value types
+    "AnyValue",
+    "LastValue",
+    "LastValueAfterFinish",
+    "UntrackedValue",
+    "EphemeralValue",
+    "BinaryOperatorAggregate",
+    "DeltaChannel",
+    "NamedBarrierValue",
+    "NamedBarrierValueAfterFinish",
+    # topics
+    "Topic",
+)

coco_graph/channels/any_value.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Generic
+
+from typing_extensions import Self
+
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel, Value
+from coco_graph.errors import EmptyChannelError
+
+__all__ = ("AnyValue",)
+
+
+class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the last value received, assumes that if multiple values are
+    received, they are all equal."""
+
+    __slots__ = ("typ", "value")
+
+    value: Value | Any
+
+    def __init__(self, typ: Any, key: str = "") -> None:
+        super().__init__(typ, key)
+        self.value = MISSING
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, AnyValue)
+
+    @property
+    def ValueType(self) -> type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.key)
+        empty.value = self.value
+        return empty
+
+    def from_checkpoint(self, checkpoint: Value) -> Self:
+        empty = self.__class__(self.typ, self.key)
+        if checkpoint is not MISSING:
+            empty.value = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        if len(values) == 0:
+            if self.value is MISSING:
+                return False
+            else:
+                self.value = MISSING
+                return True
+
+        self.value = values[-1]
+        return True
+
+    def get(self) -> Value:
+        if self.value is MISSING:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING
+
+    def checkpoint(self) -> Value:
+        return self.value