coco-framework 0.1.0__py3-none-any.whl

coco_graph/pregel/protocol.py

@@ -0,0 +1,288 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from collections.abc import AsyncIterator, Callable, Iterator, Sequence
+from typing import Any, Generic, Literal, overload
+
+from coco_core.runnables import Runnable, RunnableConfig
+from coco_core.runnables.graph import Graph as DrawableGraph
+from typing_extensions import Self
+
+from coco_graph.types import (
+    All,
+    Command,
+    GraphOutput,
+    StateSnapshot,
+    StateUpdate,
+    StreamMode,
+    StreamPart,
+)
+from coco_graph.typing import ContextT, InputT, OutputT, StateT
+
+__all__ = ("PregelProtocol", "StreamProtocol")
+
+
+class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, OutputT]):
+    @abstractmethod
+    def with_config(
+        self, config: RunnableConfig | None = None, **kwargs: Any
+    ) -> Self: ...
+
+    @abstractmethod
+    def get_graph(
+        self,
+        config: RunnableConfig | None = None,
+        *,
+        xray: int | bool = False,
+    ) -> DrawableGraph: ...
+
+    @abstractmethod
+    async def aget_graph(
+        self,
+        config: RunnableConfig | None = None,
+        *,
+        xray: int | bool = False,
+    ) -> DrawableGraph: ...
+
+    @abstractmethod
+    def get_state(
+        self, config: RunnableConfig, *, subgraphs: bool = False
+    ) -> StateSnapshot: ...
+
+    @abstractmethod
+    async def aget_state(
+        self, config: RunnableConfig, *, subgraphs: bool = False
+    ) -> StateSnapshot: ...
+
+    @abstractmethod
+    def get_state_history(
+        self,
+        config: RunnableConfig,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[StateSnapshot]: ...
+
+    @abstractmethod
+    def aget_state_history(
+        self,
+        config: RunnableConfig,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[StateSnapshot]: ...
+
+    @abstractmethod
+    def bulk_update_state(
+        self,
+        config: RunnableConfig,
+        updates: Sequence[Sequence[StateUpdate]],
+    ) -> RunnableConfig: ...
+
+    @abstractmethod
+    async def abulk_update_state(
+        self,
+        config: RunnableConfig,
+        updates: Sequence[Sequence[StateUpdate]],
+    ) -> RunnableConfig: ...
+
+    @abstractmethod
+    def update_state(
+        self,
+        config: RunnableConfig,
+        values: dict[str, Any] | Any | None,
+        as_node: str | None = None,
+    ) -> RunnableConfig: ...
+
+    @abstractmethod
+    async def aupdate_state(
+        self,
+        config: RunnableConfig,
+        values: dict[str, Any] | Any | None,
+        as_node: str | None = None,
+    ) -> RunnableConfig: ...
+
+    @overload
+    @abstractmethod
+    def stream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        version: Literal["v2"],
+    ) -> Iterator[StreamPart[StateT, OutputT]]: ...
+
+    @overload
+    @abstractmethod
+    def stream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        version: Literal["v1"] = ...,
+    ) -> Iterator[dict[str, Any] | Any]: ...
+
+    @abstractmethod
+    def stream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        version: Literal["v1", "v2"] = "v1",
+    ) -> Iterator[dict[str, Any] | Any]: ...
+
+    @overload
+    @abstractmethod
+    def astream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        version: Literal["v2"],
+    ) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
+
+    @overload
+    @abstractmethod
+    def astream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        version: Literal["v1"] = ...,
+    ) -> AsyncIterator[dict[str, Any] | Any]: ...
+
+    @abstractmethod
+    def astream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        version: Literal["v1", "v2"] = "v1",
+    ) -> AsyncIterator[dict[str, Any] | Any]: ...
+
+    @overload
+    @abstractmethod
+    def invoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        version: Literal["v2"],
+    ) -> GraphOutput[OutputT]: ...
+
+    @overload
+    @abstractmethod
+    def invoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        version: Literal["v1"] = ...,
+    ) -> dict[str, Any] | Any: ...
+
+    @abstractmethod
+    def invoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        version: Literal["v1", "v2"] = "v1",
+    ) -> dict[str, Any] | Any: ...
+
+    @overload
+    @abstractmethod
+    async def ainvoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        version: Literal["v2"],
+    ) -> GraphOutput[OutputT]: ...
+
+    @overload
+    @abstractmethod
+    async def ainvoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        version: Literal["v1"] = ...,
+    ) -> dict[str, Any] | Any: ...
+
+    @abstractmethod
+    async def ainvoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        version: Literal["v1", "v2"] = "v1",
+    ) -> dict[str, Any] | Any: ...
+
+
+StreamChunk = tuple[tuple[str, ...], str, Any]
+
+
+class StreamProtocol:
+    __slots__ = ("modes", "__call__")
+
+    modes: set[StreamMode]
+
+    __call__: Callable[[StreamChunk], None]
+
+    def __init__(
+        self,
+        __call__: Callable[[StreamChunk], None],
+        modes: set[StreamMode],
+    ) -> None:
+        self.__call__ = __call__
+        self.modes = modes

coco_graph/pregel/types.py

@@ -0,0 +1,38 @@
+"""Re-export types moved to coco_graph.types"""
+
+from coco_graph.types import (
+    All,
+    CachePolicy,
+    PregelExecutableTask,
+    PregelTask,
+    RetryPolicy,
+    StateSnapshot,
+    StateUpdate,
+    StreamMode,
+    StreamWriter,
+    default_retry_on,
+)
+
+__all__ = [
+    "All",
+    "StateUpdate",
+    "CachePolicy",
+    "PregelExecutableTask",
+    "PregelTask",
+    "RetryPolicy",
+    "StateSnapshot",
+    "StreamMode",
+    "StreamWriter",
+    "default_retry_on",
+]
+
+from warnings import warn
+
+from coco_graph.warnings import CocoGraphDeprecatedSinceV10
+
+warn(
+    "Importing from coco_graph.pregel.types is deprecated. "
+    "Please use 'from coco_graph.types import ...' instead.",
+    CocoGraphDeprecatedSinceV10,
+    stacklevel=2,
+)

coco_graph/runtime.py

@@ -0,0 +1,350 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass, field, replace
+from typing import Any, Generic, Protocol, cast, runtime_checkable
+
+from coco_graph.store.base import BaseStore
+from typing_extensions import TypedDict, Unpack
+
+from coco_graph._internal._constants import CONF, CONFIG_KEY_RUNTIME
+from coco_graph.config import get_config
+from coco_graph.types import _DC_KWARGS, StreamWriter
+from coco_graph.typing import ContextT
+
+__all__ = (
+    "BaseUser",
+    "ExecutionInfo",
+    "RunControl",
+    "Runtime",
+    "ServerInfo",
+    "get_runtime",
+)
+
+
+@runtime_checkable
+class BaseUser(Protocol):
+    """The authenticated user a server attaches to a run.
+
+    Defined here rather than imported from the upstream platform client, which
+    this framework does not depend on. Same shape: attributes and dict access.
+    """
+
+    @property
+    def is_authenticated(self) -> bool:
+        """Whether the user is authenticated."""
+        ...
+
+    @property
+    def display_name(self) -> str:
+        """The display name of the user."""
+        ...
+
+    @property
+    def identity(self) -> str:
+        """The unique identifier for the user."""
+        ...
+
+    @property
+    def permissions(self) -> Sequence[str]:
+        """The permissions associated with the user."""
+        ...
+
+    def __getitem__(self, key: str) -> Any:
+        """Get a key from the user."""
+        ...
+
+    def __contains__(self, key: object) -> bool:
+        """Check if a key exists on the user."""
+        ...
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over the user's keys."""
+        ...
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionInfo:
+    """Read-only execution info/metadata for the execution of current thread/run/node."""
+
+    checkpoint_id: str
+    """The checkpoint ID for the current execution."""
+
+    checkpoint_ns: str
+    """The checkpoint namespace for the current execution."""
+
+    task_id: str
+    """The task ID for the current execution."""
+
+    thread_id: str | None = None
+    """The thread ID for the current execution.
+
+    None when running without a checkpointer (i.e., no persistence)."""
+
+    run_id: str | None = None
+    """The run ID for the current execution.
+
+    None when `run_id` is not provided in the RunnableConfig."""
+
+    node_attempt: int = 1
+    """Current node execution attempt number (1-indexed)."""
+
+    node_first_attempt_time: float | None = None
+    """Unix timestamp (seconds) for when the first attempt started."""
+
+    def patch(self, **overrides: Any) -> ExecutionInfo:
+        """Return a new execution info object with selected fields replaced."""
+        return replace(self, **overrides)
+
+
+@dataclass(frozen=True, slots=True)
+class ServerInfo:
+    """Metadata injected by CocoGraph Server. None when running open-source CocoGraph without LangSmith deployments."""
+
+    assistant_id: str
+    """The assistant ID for the current execution."""
+
+    graph_id: str
+    """The graph ID for the current execution."""
+
+    user: BaseUser | None = None
+    """The authenticated user, if any.
+
+    This implements the `BaseUser` protocol above,
+    which supports both attribute access (e.g. `user.identity`) and dict-like
+    access (e.g. `user["identity"]`).
+    """
+
+
+class RunControl:
+    """Run-scoped control surface for cooperative draining.
+
+    Intended for a single graph run. Create a fresh `RunControl` per run;
+    reusing a control after `request_drain()` leaves it drained.
+
+    Safe to call from any thread: the drain request is represented by a
+    single attribute write, so no lock is needed for this signal.
+    If more mutable state is added here, add synchronization.
+    """
+
+    __slots__ = ("_drain_reason",)
+
+    def __init__(self) -> None:
+        self._drain_reason: str | None = None
+
+    def request_drain(self, reason: str = "shutdown") -> None:
+        self._drain_reason = reason
+
+    @property
+    def drain_requested(self) -> bool:
+        return self._drain_reason is not None
+
+    @property
+    def drain_reason(self) -> str | None:
+        return self._drain_reason
+
+
+def _no_op_stream_writer(_: Any) -> None: ...
+
+
+def _no_op_heartbeat() -> None: ...
+
+
+class _RuntimeOverrides(TypedDict, Generic[ContextT], total=False):
+    context: ContextT
+    store: BaseStore | None
+    stream_writer: StreamWriter
+    heartbeat: Callable[[], None]
+    previous: Any
+    execution_info: ExecutionInfo
+    server_info: ServerInfo | None
+    control: RunControl | None
+
+
+@dataclass(**_DC_KWARGS)
+class Runtime(Generic[ContextT]):
+    """Convenience class that bundles run-scoped context and other runtime utilities.
+
+    This class is injected into graph nodes and middleware. It provides access to
+    `context`, `store`, `stream_writer`, `previous`, and `execution_info`.
+
+    !!! note "Accessing `config`"
+
+        `Runtime` does not include `config`. To access `RunnableConfig`, you can inject
+        it directly by adding a `config: RunnableConfig` parameter to your node function
+        (recommended), or use `get_config()` from `coco_graph.config`.
+
+    !!! note
+        `ToolRuntime` (from `coco_graph.prebuilt`) is a subclass that provides similar
+        functionality but is designed specifically for tools. It shares `context`, `store`,
+        and `stream_writer` with `Runtime`, and adds tool-specific attributes like `config`,
+        `state`, and `tool_call_id`.
+
+    !!! version-added "Added in version v0.6.0"
+
+    Example:
+
+    ```python
+    from typing import TypedDict
+    from coco_graph.graph import StateGraph
+    from dataclasses import dataclass
+    from coco_graph.runtime import Runtime
+    from coco_graph.store.memory import InMemoryStore
+
+
+    @dataclass
+    class Context:  # (1)!
+        user_id: str
+
+
+    class State(TypedDict, total=False):
+        response: str
+
+
+    store = InMemoryStore()  # (2)!
+    store.put(("users",), "user_123", {"name": "Alice"})
+
+
+    def personalized_greeting(state: State, runtime: Runtime[Context]) -> State:
+        '''Generate personalized greeting using runtime context and store.'''
+        user_id = runtime.context.user_id  # (3)!
+        name = "unknown_user"
+        if runtime.store:
+            if memory := runtime.store.get(("users",), user_id):
+                name = memory.value["name"]
+
+        response = f"Hello {name}! Nice to see you again."
+        return {"response": response}
+
+
+    graph = (
+        StateGraph(state_schema=State, context_schema=Context)
+        .add_node("personalized_greeting", personalized_greeting)
+        .set_entry_point("personalized_greeting")
+        .set_finish_point("personalized_greeting")
+        .compile(store=store)
+    )
+
+    result = graph.invoke({}, context=Context(user_id="user_123"))
+    print(result)
+    # > {'response': 'Hello Alice! Nice to see you again.'}
+    ```
+
+    1. Define a schema for the runtime context.
+    2. Create a store to persist memories and other information.
+    3. Use the runtime context to access the `user_id`.
+    """
+
+    context: ContextT = field(default=None)  # type: ignore[assignment]
+    """Static context for the graph run, like `user_id`, `db_conn`, etc.
+
+    Can also be thought of as 'run dependencies'."""
+
+    store: BaseStore | None = field(default=None)
+    """Store for the graph run, enabling persistence and memory."""
+
+    stream_writer: StreamWriter = field(default=_no_op_stream_writer)
+    """Function that writes to the custom stream."""
+
+    heartbeat: Callable[[], None] = field(default=_no_op_heartbeat)
+    """Record progress for the current node's `idle_timeout`.
+
+    Call this from inside long-running work that does not naturally emit
+    writes, stream chunks, child tasks, or Coco callback events, to
+    prevent the node from being treated as idle. It is also the only
+    progress signal honored under `TimeoutPolicy(refresh_on="heartbeat")`.
+    Outside an idle-timed attempt this is a no-op.
+    """
+
+    previous: Any = field(default=None)
+    """The previous return value for the given thread.
+
+    Only available with the functional API when a checkpointer is provided.
+    """
+
+    execution_info: ExecutionInfo | None = field(default=None)
+    """Read-only execution information/metadata for the current node run.
+
+    None before task preparation populates it."""
+
+    server_info: ServerInfo | None = field(default=None)
+    """Metadata injected by CocoGraph Server. None when running open-source CocoGraph without LangSmith deployments."""
+
+    control: RunControl | None = field(default=None)
+    """Run-scoped control plane for cooperative draining.
+
+    Populated automatically during graph runs. None outside an active
+    graph runtime.
+    """
+
+    def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
+        """Merge two runtimes together.
+
+        If a value is not provided in the other runtime, the value from the current runtime is used.
+        """
+        return Runtime(
+            context=other.context or self.context,
+            store=other.store or self.store,
+            stream_writer=other.stream_writer
+            if other.stream_writer is not _no_op_stream_writer
+            else self.stream_writer,
+            heartbeat=other.heartbeat
+            if other.heartbeat is not _no_op_heartbeat
+            else self.heartbeat,
+            previous=self.previous if other.previous is None else other.previous,
+            execution_info=other.execution_info or self.execution_info,
+            server_info=other.server_info or self.server_info,
+            control=other.control or self.control,
+        )
+
+    def override(
+        self, **overrides: Unpack[_RuntimeOverrides[ContextT]]
+    ) -> Runtime[ContextT]:
+        """Replace the runtime with a new runtime with the given overrides."""
+        return replace(self, **overrides)
+
+    def patch_execution_info(self, **overrides: Any) -> Runtime[ContextT]:
+        """Return a new runtime with selected execution_info fields replaced."""
+        if self.execution_info is None:
+            msg = "Cannot patch execution_info before it has been set"
+            raise RuntimeError(msg)
+        return replace(
+            self,
+            execution_info=self.execution_info.patch(**overrides),
+        )
+
+    @property
+    def drain_requested(self) -> bool:
+        return self.control.drain_requested if self.control is not None else False
+
+    @property
+    def drain_reason(self) -> str | None:
+        return self.control.drain_reason if self.control is not None else None
+
+
+DEFAULT_RUNTIME = Runtime(
+    context=None,
+    store=None,
+    stream_writer=_no_op_stream_writer,
+    heartbeat=_no_op_heartbeat,
+    previous=None,
+    execution_info=None,
+    control=None,
+)
+
+
+def get_runtime(context_schema: type[ContextT] | None = None) -> Runtime[ContextT]:
+    """Get the runtime for the current graph run.
+
+    Args:
+        context_schema: Optional schema used for type hinting the return type of the runtime.
+
+    Returns:
+        The runtime for the current graph run.
+    """
+
+    # TODO: in an ideal world, we would have a context manager for
+    # the runtime that's independent of the config. this will follow
+    # from the removal of the configurable packing
+    runtime = cast(Runtime[ContextT], get_config()[CONF].get(CONFIG_KEY_RUNTIME))
+    return runtime