coco-framework 0.1.0__py3-none-any.whl

coco_graph/graph/ui.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from typing import Any, Literal, cast
+from uuid import uuid4
+
+from coco_core.messages import AnyMessage
+from typing_extensions import TypedDict
+
+from coco_graph.config import get_config, get_stream_writer
+from coco_graph.constants import CONF
+
+__all__ = (
+    "UIMessage",
+    "RemoveUIMessage",
+    "AnyUIMessage",
+    "push_ui_message",
+    "delete_ui_message",
+    "ui_message_reducer",
+)
+
+
+class UIMessage(TypedDict):
+    """A message type for UI updates in CocoGraph.
+
+    This TypedDict represents a UI message that can be sent to update the UI state.
+    It contains information about the UI component to render and its properties.
+
+    Attributes:
+        type: Literal type indicating this is a UI message.
+        id: Unique identifier for the UI message.
+        name: Name of the UI component to render.
+        props: Properties to pass to the UI component.
+        metadata: Additional metadata about the UI message.
+    """
+
+    type: Literal["ui"]
+    id: str
+    name: str
+    props: dict[str, Any]
+    metadata: dict[str, Any]
+
+
+class RemoveUIMessage(TypedDict):
+    """A message type for removing UI components in CocoGraph.
+
+    This TypedDict represents a message that can be sent to remove a UI component
+    from the current state.
+
+    Attributes:
+        type: Literal type indicating this is a remove-ui message.
+        id: Unique identifier of the UI message to remove.
+    """
+
+    type: Literal["remove-ui"]
+    id: str
+
+
+AnyUIMessage = UIMessage | RemoveUIMessage
+
+
+def push_ui_message(
+    name: str,
+    props: dict[str, Any],
+    *,
+    id: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    message: AnyMessage | None = None,
+    state_key: str | None = "ui",
+    merge: bool = False,
+) -> UIMessage:
+    """Push a new UI message to update the UI state.
+
+    This function creates and sends a UI message that will be rendered in the UI.
+    It also updates the graph state with the new UI message.
+
+    Args:
+        name: Name of the UI component to render.
+        props: Properties to pass to the UI component.
+        id: Optional unique identifier for the UI message.
+            If not provided, a random UUID will be generated.
+        metadata: Optional additional metadata about the UI message.
+        message: Optional message object to associate with the UI message.
+        state_key: Key in the graph state where the UI messages are stored.
+        merge: Whether to merge props with existing UI message (True) or replace
+            them (False).
+
+    Returns:
+        The created UI message.
+
+    Example:
+        ```python
+        push_ui_message(
+            name="component-name",
+            props={"content": "Hello world"},
+        )
+        ```
+
+    """
+    from coco_graph._internal._constants import CONFIG_KEY_SEND
+
+    writer = get_stream_writer()
+    config = get_config()
+
+    message_id = None
+    if message:
+        if isinstance(message, dict) and "id" in message:
+            message_id = message.get("id")
+        elif hasattr(message, "id"):
+            message_id = message.id
+
+    evt: UIMessage = {
+        "type": "ui",
+        "id": id or str(uuid4()),
+        "name": name,
+        "props": props,
+        "metadata": {
+            "merge": merge,
+            "run_id": config.get("run_id", None),
+            "tags": config.get("tags", None),
+            "name": config.get("run_name", None),
+            **(metadata or {}),
+            **({"message_id": message_id} if message_id else {}),
+        },
+    }
+
+    writer(evt)
+    if state_key:
+        config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
+
+    return evt
+
+
+def delete_ui_message(id: str, *, state_key: str = "ui") -> RemoveUIMessage:
+    """Delete a UI message by ID from the UI state.
+
+    This function creates and sends a message to remove a UI component from the current state.
+    It also updates the graph state to remove the UI message.
+
+    Args:
+        id: Unique identifier of the UI component to remove.
+        state_key: Key in the graph state where the UI messages are stored. Defaults to "ui".
+
+    Returns:
+        The remove UI message.
+
+    Example:
+        ```python
+        delete_ui_message("message-123")
+        ```
+
+    """
+    from coco_graph._internal._constants import CONFIG_KEY_SEND
+
+    writer = get_stream_writer()
+    config = get_config()
+
+    evt: RemoveUIMessage = {"type": "remove-ui", "id": id}
+
+    writer(evt)
+    config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
+
+    return evt
+
+
+def ui_message_reducer(
+    left: list[AnyUIMessage] | AnyUIMessage,
+    right: list[AnyUIMessage] | AnyUIMessage,
+) -> list[AnyUIMessage]:
+    """Merge two lists of UI messages, supporting removing UI messages.
+
+    This function combines two lists of UI messages, handling both regular UI messages
+    and `remove-ui` messages. When a `remove-ui` message is encountered, it removes any
+    UI message with the matching ID from the current state.
+
+    Args:
+        left: First list of UI messages or single UI message.
+        right: Second list of UI messages or single UI message.
+
+    Returns:
+        Combined list of UI messages with removals applied.
+
+    Example:
+        ```python
+        messages = ui_message_reducer(
+            [{"type": "ui", "id": "1", "name": "Chat", "props": {}}],
+            {"type": "remove-ui", "id": "1"},
+        )
+        ```
+
+    """
+    if not isinstance(left, list):
+        left = [left]
+
+    if not isinstance(right, list):
+        right = [right]
+
+    # merge messages
+    merged = left.copy()
+    merged_by_id = {m.get("id"): i for i, m in enumerate(merged)}
+    ids_to_remove = set()
+
+    for msg in right:
+        msg_id = msg.get("id")
+
+        if (existing_idx := merged_by_id.get(msg_id)) is not None:
+            if msg.get("type") == "remove-ui":
+                ids_to_remove.add(msg_id)
+            else:
+                ids_to_remove.discard(msg_id)
+
+                if cast(UIMessage, msg).get("metadata", {}).get("merge", False):
+                    prev_msg = merged[existing_idx]
+                    msg = msg.copy()
+                    msg["props"] = {**prev_msg["props"], **msg["props"]}
+
+                merged[existing_idx] = msg
+        else:
+            if msg.get("type") == "remove-ui":
+                raise ValueError(
+                    f"Attempting to delete an UI message with an ID that doesn't exist ('{msg_id}')"
+                )
+
+            merged_by_id[msg_id] = len(merged)
+            merged.append(msg)
+
+    merged = [m for m in merged if m.get("id") not in ids_to_remove]
+    return merged

coco_graph/managed/__init__.py

@@ -0,0 +1,3 @@
+from coco_graph.managed.is_last_step import IsLastStep, RemainingSteps
+
+__all__ = ("IsLastStep", "RemainingSteps")

coco_graph/managed/base.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from inspect import isclass
+from typing import (
+    Any,
+    Generic,
+    TypeGuard,
+    TypeVar,
+)
+
+from coco_graph._internal._scratchpad import PregelScratchpad
+
+V = TypeVar("V")
+U = TypeVar("U")
+
+__all__ = ("ManagedValueSpec", "ManagedValueMapping")
+
+
+class ManagedValue(ABC, Generic[V]):
+    @staticmethod
+    @abstractmethod
+    def get(scratchpad: PregelScratchpad) -> V: ...
+
+
+ManagedValueSpec = type[ManagedValue]
+
+
+def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
+    return isclass(value) and issubclass(value, ManagedValue)
+
+
+ManagedValueMapping = dict[str, ManagedValueSpec]

coco_graph/managed/is_last_step.py

@@ -0,0 +1,24 @@
+from typing import Annotated
+
+from coco_graph._internal._scratchpad import PregelScratchpad
+from coco_graph.managed.base import ManagedValue
+
+__all__ = ("IsLastStep", "RemainingStepsManager")
+
+
+class IsLastStepManager(ManagedValue[bool]):
+    @staticmethod
+    def get(scratchpad: PregelScratchpad) -> bool:
+        return scratchpad.step == scratchpad.stop - 1
+
+
+IsLastStep = Annotated[bool, IsLastStepManager]
+
+
+class RemainingStepsManager(ManagedValue[int]):
+    @staticmethod
+    def get(scratchpad: PregelScratchpad) -> int:
+        return scratchpad.stop - scratchpad.step
+
+
+RemainingSteps = Annotated[int, RemainingStepsManager]

coco_graph/prebuilt/__init__.py

@@ -0,0 +1,23 @@
+"""coco_graph.prebuilt exposes a higher-level API for creating and executing agents and tools."""
+
+from coco_graph.prebuilt._tool_call_transformer import ToolCallTransformer
+from coco_graph.prebuilt.chat_agent_executor import create_react_agent
+from coco_graph.prebuilt.tool_node import (
+    InjectedState,
+    InjectedStore,
+    ToolNode,
+    ToolRuntime,
+    tools_condition,
+)
+from coco_graph.prebuilt.tool_validator import ValidationNode
+
+__all__ = [
+    "create_react_agent",
+    "ToolNode",
+    "ToolCallTransformer",
+    "tools_condition",
+    "ValidationNode",
+    "InjectedState",
+    "InjectedStore",
+    "ToolRuntime",
+]

coco_graph/prebuilt/_tool_call_stream.py

@@ -0,0 +1,117 @@
+"""In-process handle for a single tool call's streaming execution.
+
+Mirrors the shape of `ChatModelStream` from trustcoco-core but simpler —
+a tool has one output channel, no content-block multiplexing. Populated
+by `ToolCallTransformer` as `tool-started` / `tool-output-delta` /
+`tool-finished` / `tool-error` events flow in on the `tools` channel.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from typing import Any
+
+from coco_graph.stream.stream_channel import StreamChannel
+
+
+class ToolCallStream:
+    """Scoped view of a single tool call's lifecycle.
+
+    Yielded on `run.tool_calls` once per `tool-started` event. Fields
+    are populated as events arrive:
+
+    - `tool_call_id`, `tool_name`, `input`: stable from the start event.
+    - `output_deltas`: a `StreamChannel` of delta chunks. Iterate (sync or
+      async) to consume partial output in arrival order.
+    - `output`: terminal payload from `tool-finished`, or `None` if the
+      call failed or is still in flight.
+    - `error`: terminal error string from `tool-error`, or `None` if the
+      call succeeded or is still in flight.
+    - `completed`: True once a terminal event (`tool-finished` or
+      `tool-error`) has been observed.
+
+    `ToolCallStream` is not meant to be constructed by end users — it's
+    produced by `ToolCallTransformer` as events flow through the mux.
+    """
+
+    def __init__(
+        self,
+        tool_call_id: str,
+        tool_name: str,
+        input: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize a fresh handle for a tool call.
+
+        Args:
+            tool_call_id: The `tool_call_id` from the AIMessage.
+            tool_name: The tool's name.
+            input: The tool's input arguments (as reported by
+                `on_tool_start`), or `None` if none were captured.
+        """
+        self.tool_call_id = tool_call_id
+        self.tool_name = tool_name
+        self.input = input
+        self._output_deltas: StreamChannel[Any] = StreamChannel()
+        self.output: Any = None
+        self.error: str | None = None
+        self.completed = False
+
+    @property
+    def output_deltas(self) -> StreamChannel[Any]:
+        """The channel of streamed `tool-output-delta` payloads.
+
+        Iterate (sync or async depending on how the run was started)
+        to consume partial output in arrival order. The log closes when
+        the tool finishes or errors.
+        """
+        return self._output_deltas
+
+    def _bind(self, *, is_async: bool) -> None:
+        """Bind the deltas log to sync or async iteration.
+
+        Called by `ToolCallTransformer` when constructing this handle so
+        the log matches the enclosing mux's mode.
+        """
+        self._output_deltas._bind(is_async=is_async)
+
+    def _push_delta(self, delta: Any) -> None:
+        self._output_deltas.push(delta)
+
+    def _finish(self, output: Any) -> None:
+        self.output = output
+        self.completed = True
+        self._output_deltas.close()
+
+    def _fail(self, message: str) -> None:
+        self.error = message
+        self.completed = True
+        self._output_deltas.close()
+
+    def __iter__(self) -> Iterator[Any]:
+        """Iterate delta chunks synchronously.
+
+        Equivalent to `iter(self.output_deltas)`. Raises `TypeError` if
+        the underlying log is bound to async mode.
+        """
+        return iter(self._output_deltas)
+
+    def __aiter__(self) -> AsyncIterator[Any]:
+        """Iterate delta chunks asynchronously.
+
+        Equivalent to `aiter(self.output_deltas)`. Raises `TypeError`
+        if the underlying log is bound to sync mode.
+        """
+        return self._output_deltas.__aiter__()
+
+    def __repr__(self) -> str:
+        status = (
+            "completed"
+            if self.completed and self.error is None
+            else "failed"
+            if self.completed
+            else "running"
+        )
+        return (
+            f"ToolCallStream(tool_call_id={self.tool_call_id!r}, "
+            f"tool_name={self.tool_name!r}, status={status})"
+        )

coco_graph/prebuilt/_tool_call_transformer.py

@@ -0,0 +1,165 @@
+"""Transformer that projects `tools` channel events into `ToolCallStream`s."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from coco_core.messages import ToolMessage
+from coco_graph.stream._types import ProtocolEvent, StreamTransformer
+from coco_graph.stream.stream_channel import StreamChannel
+
+from coco_graph.prebuilt._tool_call_stream import ToolCallStream
+
+
+def _is_serialized_tool_message(value: Any) -> bool:
+    """Detect a serialized Coco `ToolMessage` payload.
+
+    Example:
+        {
+            "lc": 1,
+            "type": "constructor",
+            "id": ["coco_core", "messages", "ToolMessage"],
+            "kwargs": {"content": "raw tool result", "tool_call_id": "call_1"},
+        }
+    """
+    return (
+        isinstance(value, dict)
+        and value.get("type") == "constructor"
+        and isinstance(value.get("id"), list)
+        and value["id"][-1] == "ToolMessage"
+    )
+
+
+def _normalize_tool_output(output: Any) -> Any:
+    if isinstance(output, ToolMessage):
+        return output.content
+    if _is_serialized_tool_message(output):
+        kwargs = output.get("kwargs")
+        if isinstance(kwargs, dict):
+            return kwargs.get("content")
+    return output
+
+
+class ToolCallTransformer(StreamTransformer):
+    """Project `tools` channel events into `ToolCallStream` handles.
+
+    Each `tool-started` event spawns a `ToolCallStream`, pushed onto
+    `run.tool_calls`. Subsequent `tool-output-delta` events append to
+    that stream's deltas log; `tool-finished` and `tool-error` close it.
+
+    Native transformer — the `tool_calls` projection is exposed as a
+    direct attribute on the run stream.
+
+    A nameless `StreamChannel[ToolCallStream]` is used (no protocol
+    auto-forwarding) because the live handles are not serializable and
+    should not be injected into the main event log. Wire consumers
+    subscribe to the `tools` channel instead, where the raw protocol
+    events flow through untouched by this transformer (`process`
+    returns `True`).
+
+    Registered explicitly by users at compile time via
+    `builder.compile(transformers=[ToolCallTransformer])` — not a
+    default built-in, so the `tools` channel is user-opt-in.
+    """
+
+    _native = True
+    required_stream_modes = ("tools",)
+
+    def __init__(self, scope: tuple[str, ...] = ()) -> None:
+        super().__init__(scope)
+        self._log: StreamChannel[ToolCallStream] = StreamChannel()
+        self._active: dict[str, ToolCallStream] = {}
+        self._is_async = False
+        self._pump_fn: Callable[[], bool] | None = None
+        self._apump_fn: Callable[[], Awaitable[bool]] | None = None
+
+    def init(self) -> dict[str, Any]:
+        return {"tool_calls": self._log}
+
+    def _bind_pump(self, fn: Callable[[], bool]) -> None:
+        """Wire the sync pull callback onto this transformer.
+
+        Called by `StreamMux.bind_pump`. Stored so each new
+        `ToolCallStream` created by `process` can wire its deltas log
+        for pump-driven iteration.
+        """
+        self._pump_fn = fn
+        self._is_async = False
+
+    def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
+        """Async counterpart to `_bind_pump`."""
+        self._apump_fn = fn
+        self._is_async = True
+
+    def _new_stream(
+        self,
+        tool_call_id: str,
+        tool_name: str,
+        tool_input: dict[str, Any] | None,
+    ) -> ToolCallStream:
+        stream = ToolCallStream(tool_call_id, tool_name, tool_input)
+        stream._bind(is_async=self._is_async)
+        if self._apump_fn is not None:
+            stream._output_deltas._arequest_more = self._apump_fn
+        if self._pump_fn is not None:
+            stream._output_deltas._request_more = self._pump_fn
+        return stream
+
+    def process(self, event: ProtocolEvent) -> bool:
+        if event["method"] != "tools":
+            return True
+
+        # Only project events emitted at this transformer's scope. Subgraph
+        # events still flow through the parent's mux (the parent's main
+        # event log keeps them) but they belong to the child mini-mux's
+        # `tool_calls` projection, not the parent's.
+        if tuple(event["params"]["namespace"]) != self.scope:
+            return True
+
+        data = event["params"]["data"]
+        tool_call_id = data.get("tool_call_id")
+        if tool_call_id is None:
+            return True
+        event_type = data.get("event")
+
+        stream: ToolCallStream | None
+        if event_type == "tool-started":
+            stream = self._new_stream(
+                tool_call_id,
+                data.get("tool_name", ""),
+                data.get("input"),
+            )
+            self._active[tool_call_id] = stream
+            self._log.push(stream)
+        elif event_type == "tool-output-delta":
+            stream = self._active.get(tool_call_id)
+            if stream is not None:
+                stream._push_delta(data.get("delta"))
+        elif event_type == "tool-finished":
+            stream = self._active.pop(tool_call_id, None)
+            if stream is not None:
+                stream._finish(_normalize_tool_output(data.get("output")))
+        elif event_type == "tool-error":
+            stream = self._active.pop(tool_call_id, None)
+            if stream is not None:
+                stream._fail(data.get("message", ""))
+
+        # Pass-through — wire consumers subscribe to the `tools` channel
+        # directly and reconstruct handles client-side.
+        return True
+
+    def finalize(self) -> None:
+        """Close any still-active tool streams left open at run end."""
+        for stream in self._active.values():
+            if not stream.completed:
+                stream._finish(None)
+        self._active.clear()
+
+    def fail(self, err: BaseException) -> None:
+        """Fail any still-active tool streams when the run errors."""
+        message = str(err)
+        for stream in self._active.values():
+            if not stream.completed:
+                stream._fail(message)
+        self._active.clear()