coco-framework 0.1.0__py3-none-any.whl

coco_graph/pregel/_tools.py

@@ -0,0 +1,268 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable, Iterator
+from contextvars import ContextVar, Token
+from typing import Any, TypeVar, cast
+from uuid import UUID
+
+from coco_core.callbacks import BaseCallbackHandler
+
+from coco_graph._internal._constants import NS_SEP
+from coco_graph.constants import TAG_NOSTREAM
+from coco_graph.pregel.protocol import StreamChunk
+
+try:
+    from coco_core.tracers._streaming import _StreamingCallbackHandler
+except ImportError:
+    _StreamingCallbackHandler = object  # type: ignore[assignment,misc]
+
+
+T = TypeVar("T")
+
+ToolCallWriter = Callable[[Any], None]
+"""A closure bound to a single tool call that emits `tool-output-delta` events."""
+
+_tool_call_writer: ContextVar[ToolCallWriter | None] = ContextVar(
+    "coco_graph_tool_call_writer", default=None
+)
+"""ContextVar holding the writer for the currently-executing tool call.
+
+Set by `StreamToolCallHandler.on_tool_start` and reset on end/error.
+Read by `ToolRuntime.emit_output_delta` (in `coco_graph.prebuilt`).
+"""
+
+
+class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
+    """Callback handler that emits tool-call lifecycle events on the stream.
+
+    Fires on Coco's `on_tool_*` callbacks and pushes to the `tools`
+    stream mode. Emits `tool-started` / `tool-output-delta` /
+    `tool-finished` / `tool-error` payloads keyed by `tool_call_id`.
+
+    While a tool is executing, this handler sets `_tool_call_writer` to a
+    closure bound to that call's namespace and `tool_call_id`.
+    `ToolRuntime.emit_output_delta` reads that ContextVar so tool bodies
+    can stream partial output without threading the writer through their
+    own signature.
+
+    Attached by `Pregel.stream` / `astream` when `"tools"` is in
+    `stream_modes`. `run_inline = True` keeps event ordering
+    deterministic.
+    """
+
+    run_inline = True
+
+    def __init__(
+        self,
+        stream: Callable[[StreamChunk], None],
+        subgraphs: bool,
+        *,
+        parent_ns: tuple[str, ...] | None = None,
+    ) -> None:
+        """Configure the handler to stream tool-call events.
+
+        Args:
+            stream: Callable that accepts a `StreamChunk` tuple
+                `(namespace, mode, payload)` and enqueues it.
+            subgraphs: Whether to emit events from tools called inside
+                nested subgraphs. When False, only tools at the
+                handler's own scope (`parent_ns`) emit.
+            parent_ns: Namespace where the handler was attached.
+                Mirrors the `StreamMessagesHandler` escape hatch:
+                tools whose containing namespace equals `parent_ns`
+                still emit even with `subgraphs=False`, so a node that
+                explicitly streams a subgraph with `stream_mode="tools"`
+                sees its own tools.
+        """
+        self.stream = stream
+        self.subgraphs = subgraphs
+        self.parent_ns = parent_ns
+        # run_id → (namespace, tool_call_id, ContextVar token)
+        # `on_tool_end` does not receive `tool_call_id` in kwargs, so
+        # we correlate by `run_id` which is present on every callback.
+        self._run_to_call: dict[
+            UUID, tuple[tuple[str, ...], str, Token[ToolCallWriter | None]]
+        ] = {}
+
+    def _ns_for_emit(
+        self,
+        metadata: dict[str, Any] | None,
+        tags: list[str] | None,
+    ) -> tuple[str, ...] | None:
+        """Resolve the namespace this tool call should emit at, or `None` to skip.
+
+        Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
+        derivation: parses `coco_graph_checkpoint_ns` (which ends with
+        the `node_name:task_id` of the calling node), drops that
+        trailing segment, and returns the containing subgraph's own
+        namespace. Returns `None` when the call should be silently
+        suppressed:
+
+        - `metadata` is missing — handler is attached to a context
+          without Pregel routing info.
+        - `TAG_NOSTREAM` is in `tags` — caller explicitly opted out.
+        - Tool runs in a subgraph (`len(ns) > 0`) and the handler was
+          attached with `subgraphs=False` and a different `parent_ns`
+          than the call's containing subgraph.
+        """
+        if not metadata:
+            return None
+        if tags and TAG_NOSTREAM in tags:
+            return None
+        nskey = metadata.get("coco_graph_checkpoint_ns")
+        if not nskey:
+            ns: tuple[str, ...] = ()
+        else:
+            ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
+        if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
+            return None
+        return ns
+
+    def _start(
+        self,
+        serialized: dict[str, Any] | None,
+        input_str: str,
+        *,
+        run_id: UUID,
+        metadata: dict[str, Any] | None,
+        tags: list[str] | None,
+        inputs: dict[str, Any] | None,
+        kwargs: dict[str, Any],
+    ) -> None:
+        ns = self._ns_for_emit(metadata, tags)
+        if ns is None:
+            return
+        tool_call_id = cast("str | None", kwargs.get("tool_call_id")) or str(run_id)
+        tool_name = (
+            (serialized or {}).get("name")
+            or cast("str | None", kwargs.get("name"))
+            or ""
+        )
+
+        def writer(delta: Any) -> None:
+            self.stream(
+                (
+                    ns,
+                    "tools",
+                    {
+                        "event": "tool-output-delta",
+                        "tool_call_id": tool_call_id,
+                        "delta": delta,
+                    },
+                )
+            )
+
+        token = _tool_call_writer.set(writer)
+        self._run_to_call[run_id] = (ns, tool_call_id, token)
+
+        payload: dict[str, Any] = {
+            "event": "tool-started",
+            "tool_call_id": tool_call_id,
+            "tool_name": tool_name,
+        }
+        if inputs is not None:
+            payload["input"] = inputs
+        self.stream((ns, "tools", payload))
+
+    def _end(self, output: Any, *, run_id: UUID) -> None:
+        info = self._run_to_call.pop(run_id, None)
+        if info is None:
+            return
+        ns, tool_call_id, token = info
+        self._reset_writer(token)
+        self.stream(
+            (
+                ns,
+                "tools",
+                {
+                    "event": "tool-finished",
+                    "tool_call_id": tool_call_id,
+                    "output": output,
+                },
+            )
+        )
+
+    def _error(self, error: BaseException, *, run_id: UUID) -> None:
+        info = self._run_to_call.pop(run_id, None)
+        if info is None:
+            return
+        ns, tool_call_id, token = info
+        self._reset_writer(token)
+        self.stream(
+            (
+                ns,
+                "tools",
+                {
+                    "event": "tool-error",
+                    "tool_call_id": tool_call_id,
+                    "message": str(error),
+                },
+            )
+        )
+
+    def tap_output_aiter(
+        self, run_id: UUID, output: AsyncIterator[T]
+    ) -> AsyncIterator[T]:
+        """Pass-through — required by the `_StreamingCallbackHandler` protocol."""
+        return output
+
+    def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
+        """Pass-through — sync counterpart to `tap_output_aiter`."""
+        return output
+
+    @staticmethod
+    def _reset_writer(token: Token[ToolCallWriter | None]) -> None:
+        # Token is invalid if `on_tool_end` runs in a different context
+        # than `on_tool_start` (e.g. coco may hand off to a thread
+        # worker without copying the context). Swallow that case; the
+        # ContextVar lifetime is bounded by the enclosing task anyway.
+        try:
+            _tool_call_writer.reset(token)
+        except ValueError:
+            pass
+
+    # ------------------------------------------------------------------
+    # Sync callbacks
+    # ------------------------------------------------------------------
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._start(
+            serialized,
+            input_str,
+            run_id=run_id,
+            metadata=metadata,
+            tags=tags,
+            inputs=inputs,
+            kwargs=kwargs,
+        )
+
+    def on_tool_end(
+        self,
+        output: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._end(output, run_id=run_id)
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._error(error, run_id=run_id)

coco_graph/pregel/_utils.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import dis
+import re
+from collections.abc import Callable, Sequence
+from functools import partial
+from types import CodeType, FunctionType
+from typing import Any
+
+from coco_core.runnables import (
+    Runnable,
+    RunnableLambda,
+    RunnableParallel,
+    RunnableSequence,
+)
+from coco_core.runnables.base import RunnableBindingBase
+from coco_core.runnables.config import run_in_executor
+from coco_graph.checkpoint.base import ChannelVersions
+
+from coco_graph._internal._runnable import RunnableCallable, RunnableSeq
+from coco_graph._internal._timeout import sync_timeout_unsupported
+from coco_graph.pregel.protocol import PregelProtocol
+
+_SEQUENCE_TYPES = (RunnableSeq, RunnableSequence)
+
+
+def get_new_channel_versions(
+    previous_versions: ChannelVersions, current_versions: ChannelVersions
+) -> ChannelVersions:
+    """Get subset of current_versions that are newer than previous_versions."""
+    if previous_versions:
+        version_type = type(next(iter(current_versions.values()), None))
+        null_version = version_type()  # type: ignore[misc]
+        new_versions = {
+            k: v
+            for k, v in current_versions.items()
+            if v > previous_versions.get(k, null_version)  # type: ignore[operator]
+        }
+    else:
+        new_versions = current_versions
+
+    return new_versions
+
+
+def find_subgraph_pregel(candidate: Runnable) -> PregelProtocol | None:
+    from coco_graph.pregel import Pregel
+
+    candidates: list[Runnable] = [candidate]
+
+    for c in candidates:
+        if (
+            isinstance(c, PregelProtocol)
+            # subgraphs that disabled checkpointing are not considered
+            and (not isinstance(c, Pregel) or c.checkpointer is not False)
+        ):
+            return c
+        elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq):
+            candidates.extend(c.steps)
+        elif isinstance(c, RunnableLambda):
+            candidates.extend(c.deps)
+        elif isinstance(c, RunnableCallable):
+            if c.func is not None:
+                candidates.extend(
+                    nl.__self__ if hasattr(nl, "__self__") else nl
+                    for nl in get_function_nonlocals(c.func)
+                )
+            elif c.afunc is not None:
+                candidates.extend(
+                    nl.__self__ if hasattr(nl, "__self__") else nl
+                    for nl in get_function_nonlocals(c.afunc)
+                )
+
+    return None
+
+
+def _sequence_steps(runnable: Runnable) -> Sequence[Runnable] | None:
+    if isinstance(runnable, _SEQUENCE_TYPES):
+        return runnable.steps
+    return None
+
+
+def _parallel_steps(runnable: Runnable) -> Sequence[Runnable] | None:
+    if isinstance(runnable, RunnableParallel):
+        return tuple(runnable.steps__.values())
+    return None
+
+
+def _has_method_override(runnable: Runnable, method_name: str) -> bool:
+    method = getattr(type(runnable), method_name, None)
+    return method is not None and method is not getattr(Runnable, method_name)
+
+
+def _is_executor_backed_afunc(afunc: Callable[..., Any] | None) -> bool:
+    return isinstance(afunc, partial) and afunc.func is run_in_executor
+
+
+def _has_native_async(runnable: Runnable) -> bool:
+    if isinstance(runnable, RunnableCallable):
+        return runnable.afunc is not None and not _is_executor_backed_afunc(
+            runnable.afunc
+        )
+    if isinstance(runnable, RunnableLambda):
+        return bool(getattr(runnable, "afunc", False))
+    return _has_method_override(runnable, "ainvoke")
+
+
+def _runnable_has_native_async(runnable: Runnable) -> bool:
+    """Return whether a runnable can be idle-timed without known sync code.
+
+    For custom runnable subclasses, an `ainvoke` override is treated as the
+    async contract. We do not introspect whether that implementation delegates
+    to blocking work internally — e.g. a subclass whose `ainvoke` calls
+    `asyncio.to_thread(self.invoke, ...)` will pass this check but the wrapped
+    sync work is still uncancellable. Idle-timeout enforcement on such a
+    runnable will fire `NodeTimeoutError` correctly, but the background thread
+    will keep running until its sync work returns.
+    """
+
+    while isinstance(runnable, RunnableBindingBase):
+        runnable = runnable.bound
+    steps = _sequence_steps(runnable)
+    if steps is None:
+        steps = _parallel_steps(runnable)
+    if steps is not None:
+        return all(_runnable_has_native_async(step) for step in steps)
+    # Raw callables and the common composition wrappers created by graph
+    # builders fall through here. We do not exhaustively unwrap every Runnable
+    # wrapper — wrappers that provide `ainvoke` are treated as owning the async
+    # contract.
+    return _has_native_async(runnable)
+
+
+def validate_timeout_supported(runnable: Runnable, *, name: str) -> None:
+    if not _runnable_has_native_async(runnable):
+        raise sync_timeout_unsupported(name)
+
+
+# Values treated as dead ends when deciding whether to walk a function's
+# bytecode. A container can hold a graph, but `find_subgraph_pregel` does not
+# look inside one, so skipping it costs nothing while that holds. Matched by
+# exact type, since a subclass of a builtin can carry attributes.
+_LEAF_TYPES = frozenset(
+    {
+        int,
+        float,
+        complex,
+        bool,
+        str,
+        bytes,
+        bytearray,
+        list,
+        tuple,
+        dict,
+        set,
+        frozenset,
+        type(None),
+    }
+)
+
+
+def get_function_nonlocals(func: Callable) -> list[Any]:
+    """Get the values a function reaches from outside its own scope.
+
+    Args:
+        func: The function to check.
+
+    Returns:
+        Every captured cell value, the globals the function names, and each
+        value along an attribute path it loads. Over-approximates: a value can
+        come back without the function reaching it at runtime.
+    """
+    func = getattr(func, "__func__", func)  # bound method -> function
+    wrapped = getattr(func, "__wrapped__", None)
+    if callable(wrapped):
+        func = getattr(wrapped, "__func__", wrapped)
+    if not isinstance(func, FunctionType):
+        return []
+    code = func.__code__
+
+    cells: dict[str, Any] = {}
+    for name, cell in zip(code.co_freevars, func.__closure__ or ()):
+        try:
+            cells[name] = cell.cell_contents
+        except ValueError:
+            continue  # empty cell: a recursive def not yet bound
+
+    # Every captured value counts, referenced or not: over-declaring costs an
+    # introspection entry, under-declaring drops the subgraph's checkpoints and
+    # stream events. Checking each cell against the bytecode would cost more and
+    # only trade the cheap error for the expensive one.
+    values: list[Any] = list(cells.values())
+    global_ns = func.__globals__
+    globals_ = {name: global_ns[name] for name in code.co_names if name in global_ns}
+    if all(type(v) in _LEAF_TYPES for v in (*cells.values(), *globals_.values())):
+        return values
+
+    # Nested code objects hold the references made by inner defs, lambdas and
+    # comprehensions, which resolve against the namespaces gathered above.
+    codes = [code]
+    for c in codes:
+        codes.extend(k for k in c.co_consts if isinstance(k, CodeType))
+        value: Any = None
+        for instruction in dis.get_instructions(c):
+            opname = instruction.opname
+            if opname == "LOAD_GLOBAL":
+                value = globals_.get(instruction.argval)
+            elif opname == "LOAD_DEREF":
+                value = cells.get(instruction.argval)
+            elif opname in ("LOAD_ATTR", "LOAD_METHOD"):
+                value = getattr(value, instruction.argval, None)
+            else:
+                value = None  # anything else ends the chain: `a, b.c` is not `a.c`
+                continue
+            if value is not None:
+                values.append(value)
+    return values
+
+
+def is_xxh3_128_hexdigest(value: str) -> bool:
+    """Check if the given string matches the format of xxh3_128_hexdigest."""
+    return bool(re.fullmatch(r"[0-9a-f]{32}", value))

coco_graph/pregel/_validate.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from coco_graph._internal._constants import RESERVED
+from coco_graph.channels.base import BaseChannel
+from coco_graph.managed.base import ManagedValueMapping
+from coco_graph.pregel._read import PregelNode
+from coco_graph.types import All
+
+
+def validate_graph(
+    nodes: Mapping[str, PregelNode],
+    channels: dict[str, BaseChannel],
+    managed: ManagedValueMapping,
+    input_channels: str | Sequence[str],
+    output_channels: str | Sequence[str],
+    stream_channels: str | Sequence[str] | None,
+    interrupt_after_nodes: All | Sequence[str],
+    interrupt_before_nodes: All | Sequence[str],
+) -> None:
+    for chan in channels:
+        if chan in RESERVED:
+            raise ValueError(f"Channel name '{chan}' is reserved")
+    for name in managed:
+        if name in RESERVED:
+            raise ValueError(f"Managed name '{name}' is reserved")
+
+    subscribed_channels = set[str]()
+    for name, node in nodes.items():
+        if name in RESERVED:
+            raise ValueError(f"Node name '{name}' is reserved")
+        if isinstance(node, PregelNode):
+            subscribed_channels.update(node.triggers)
+            if isinstance(node.channels, str):
+                if node.channels not in channels:
+                    raise ValueError(
+                        f"Node {name} reads channel '{node.channels}' "
+                        f"not in known channels: '{repr(sorted(channels))[:100]}'"
+                    )
+            else:
+                for chan in node.channels:
+                    if chan not in channels and chan not in managed:
+                        raise ValueError(
+                            f"Node {name} reads channel '{chan}' "
+                            f"not in known channels: '{repr(sorted(channels))[:100]}'"
+                        )
+        else:
+            raise TypeError(
+                f"Invalid node type {type(node)}, expected PregelNode or NodeBuilder"
+            )
+
+    for chan in subscribed_channels:
+        if chan not in channels:
+            raise ValueError(
+                f"Subscribed channel '{chan}' not "
+                f"in known channels: '{repr(sorted(channels))[:100]}'"
+            )
+
+    if isinstance(input_channels, str):
+        if input_channels not in channels:
+            raise ValueError(
+                f"Input channel '{input_channels}' not "
+                f"in known channels: '{repr(sorted(channels))[:100]}'"
+            )
+        if input_channels not in subscribed_channels:
+            raise ValueError(
+                f"Input channel {input_channels} is not subscribed to by any node"
+            )
+    else:
+        for chan in input_channels:
+            if chan not in channels:
+                raise ValueError(
+                    f"Input channel '{chan}' not in '{repr(sorted(channels))[:100]}'"
+                )
+        if all(chan not in subscribed_channels for chan in input_channels):
+            raise ValueError(
+                f"None of the input channels {input_channels} "
+                f"are subscribed to by any node"
+            )
+
+    all_output_channels = set[str]()
+    if isinstance(output_channels, str):
+        all_output_channels.add(output_channels)
+    else:
+        all_output_channels.update(output_channels)
+    if isinstance(stream_channels, str):
+        all_output_channels.add(stream_channels)
+    elif stream_channels is not None:
+        all_output_channels.update(stream_channels)
+
+    for chan in all_output_channels:
+        if chan not in channels:
+            raise ValueError(
+                f"Output channel '{chan}' not "
+                f"in known channels: '{repr(sorted(channels))[:100]}'"
+            )
+
+    if interrupt_after_nodes != "*":
+        for n in interrupt_after_nodes:
+            if n not in nodes:
+                raise ValueError(f"Node {n} not in nodes")
+    if interrupt_before_nodes != "*":
+        for n in interrupt_before_nodes:
+            if n not in nodes:
+                raise ValueError(f"Node {n} not in nodes")
+
+
+def validate_keys(
+    keys: str | Sequence[str] | None,
+    channels: Mapping[str, Any],
+) -> None:
+    if isinstance(keys, str):
+        if keys not in channels:
+            raise ValueError(f"Key {keys} not in channels")
+    elif keys is not None:
+        for chan in keys:
+            if chan not in channels:
+                raise ValueError(f"Key {chan} not in channels")