coco-framework 0.1.0__py3-none-any.whl

coco_graph/pregel/_write.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import (
+    Any,
+    NamedTuple,
+    TypeVar,
+    cast,
+)
+
+from coco_core.runnables import Runnable, RunnableConfig
+
+from coco_graph._internal._constants import CONF, CONFIG_KEY_SEND, TASKS
+from coco_graph._internal._runnable import RunnableCallable
+from coco_graph._internal._typing import MISSING
+from coco_graph.errors import InvalidUpdateError
+from coco_graph.types import Send
+
+TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
+R = TypeVar("R", bound=Runnable)
+
+SKIP_WRITE = object()
+PASSTHROUGH = object()
+
+
+class ChannelWriteEntry(NamedTuple):
+    channel: str
+    """Channel name to write to."""
+    value: Any = PASSTHROUGH
+    """Value to write, or PASSTHROUGH to use the input."""
+    skip_none: bool = False
+    """Whether to skip writing if the value is None."""
+    mapper: Callable | None = None
+    """Function to transform the value before writing."""
+
+
+class ChannelWriteTupleEntry(NamedTuple):
+    mapper: Callable[[Any], Sequence[tuple[str, Any]] | None]
+    """Function to extract tuples from value."""
+    value: Any = PASSTHROUGH
+    """Value to write, or PASSTHROUGH to use the input."""
+    static: Sequence[tuple[str, Any, str | None]] | None = None
+    """Optional, declared writes for static analysis."""
+
+
+class ChannelWrite(RunnableCallable):
+    """Implements the logic for sending writes to CONFIG_KEY_SEND.
+    Can be used as a runnable or as a static method to call imperatively."""
+
+    writes: list[ChannelWriteEntry | ChannelWriteTupleEntry | Send]
+    """Sequence of write entries or Send objects to write."""
+
+    def __init__(
+        self,
+        writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send],
+        *,
+        tags: Sequence[str] | None = None,
+    ):
+        super().__init__(
+            func=self._write,
+            afunc=self._awrite,
+            name=None,
+            tags=tags,
+            trace=False,
+        )
+        self.writes = cast(
+            list[ChannelWriteEntry | ChannelWriteTupleEntry | Send], writes
+        )
+
+    def get_name(self, suffix: str | None = None, *, name: str | None = None) -> str:
+        if not name:
+            name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else '...' if isinstance(w, ChannelWriteTupleEntry) else w.node for w in self.writes)}>"
+        return super().get_name(suffix, name=name)
+
+    def _write(self, input: Any, config: RunnableConfig) -> None:
+        writes = [
+            ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
+            if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
+            else ChannelWriteTupleEntry(write.mapper, input)
+            if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH
+            else write
+            for write in self.writes
+        ]
+        self.do_write(
+            config,
+            writes,
+        )
+        return input
+
+    async def _awrite(self, input: Any, config: RunnableConfig) -> None:
+        writes = [
+            ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
+            if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
+            else ChannelWriteTupleEntry(write.mapper, input)
+            if isinstance(write, ChannelWriteTupleEntry) and write.value is PASSTHROUGH
+            else write
+            for write in self.writes
+        ]
+        self.do_write(
+            config,
+            writes,
+        )
+        return input
+
+    @staticmethod
+    def do_write(
+        config: RunnableConfig,
+        writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send],
+        allow_passthrough: bool = True,
+    ) -> None:
+        # validate
+        for w in writes:
+            if isinstance(w, ChannelWriteEntry):
+                if w.channel == TASKS:
+                    raise InvalidUpdateError(
+                        "Cannot write to the reserved channel TASKS"
+                    )
+                if w.value is PASSTHROUGH and not allow_passthrough:
+                    raise InvalidUpdateError("PASSTHROUGH value must be replaced")
+            if isinstance(w, ChannelWriteTupleEntry):
+                if w.value is PASSTHROUGH and not allow_passthrough:
+                    raise InvalidUpdateError("PASSTHROUGH value must be replaced")
+        # if we want to persist writes found before hitting a ParentCommand
+        # can move this to a finally block
+        write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
+        write(_assemble_writes(writes))
+
+    @staticmethod
+    def is_writer(runnable: Runnable) -> bool:
+        """Used by PregelNode to distinguish between writers and other runnables."""
+        return (
+            isinstance(runnable, ChannelWrite)
+            or getattr(runnable, "_is_channel_writer", MISSING) is not MISSING
+        )
+
+    @staticmethod
+    def get_static_writes(
+        runnable: Runnable,
+    ) -> Sequence[tuple[str, Any, str | None]] | None:
+        """Used to get conditional writes a writer declares for static analysis."""
+        if isinstance(runnable, ChannelWrite):
+            return [
+                w
+                for entry in runnable.writes
+                if isinstance(entry, ChannelWriteTupleEntry) and entry.static
+                for w in entry.static
+            ] or None
+        elif writes := getattr(runnable, "_is_channel_writer", MISSING):
+            if writes is not MISSING:
+                writes = cast(
+                    Sequence[tuple[ChannelWriteEntry | Send, str | None]],
+                    writes,
+                )
+                entries = [e for e, _ in writes]
+                labels = [la for _, la in writes]
+                return [(*t, la) for t, la in zip(_assemble_writes(entries), labels)]
+
+    @staticmethod
+    def register_writer(
+        runnable: R,
+        static: Sequence[tuple[ChannelWriteEntry | Send, str | None]] | None = None,
+    ) -> R:
+        """Used to mark a runnable as a writer, so that it can be detected by is_writer.
+        Instances of ChannelWrite are automatically marked as writers.
+        Optionally, a list of declared writes can be passed for static analysis."""
+        # using object.__setattr__ to work around objects that override __setattr__
+        # eg. pydantic models and dataclasses
+        object.__setattr__(runnable, "_is_channel_writer", static)
+        return runnable
+
+
+def _assemble_writes(
+    writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send],
+) -> list[tuple[str, Any]]:
+    """Assembles the writes into a list of tuples."""
+    tuples: list[tuple[str, Any]] = []
+    for w in writes:
+        if isinstance(w, Send):
+            tuples.append((TASKS, w))
+        elif isinstance(w, ChannelWriteTupleEntry):
+            if ww := w.mapper(w.value):
+                tuples.extend(ww)
+        elif isinstance(w, ChannelWriteEntry):
+            value = w.mapper(w.value) if w.mapper is not None else w.value
+            if value is SKIP_WRITE:
+                continue
+            if w.skip_none and value is None:
+                continue
+            tuples.append((w.channel, value))
+        else:
+            raise ValueError(f"Invalid write entry: {w}")
+    return tuples

coco_graph/pregel/debug.py

@@ -0,0 +1,302 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import asdict
+from typing import Any
+from uuid import UUID
+
+from coco_core.runnables import RunnableConfig
+from coco_graph.checkpoint.base import (
+    EXCLUDED_METADATA_KEYS,
+    CheckpointMetadata,
+    PendingWrite,
+)
+
+from coco_graph._internal._config import filter_to_user_tags, patch_checkpoint_map
+from coco_graph._internal._constants import (
+    CONF,
+    CONFIG_KEY_CHECKPOINT_NS,
+    ERROR,
+    INTERRUPT,
+    NS_END,
+    NS_SEP,
+    RETURN,
+)
+from coco_graph._internal._typing import MISSING
+from coco_graph.channels.base import BaseChannel
+from coco_graph.constants import TAG_HIDDEN
+from coco_graph.pregel._io import read_channels
+from coco_graph.types import (
+    CheckpointPayload,
+    PregelExecutableTask,
+    PregelTask,
+    StateSnapshot,
+    TaskPayload,
+    TaskResultPayload,
+)
+
+TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
+
+
+def map_debug_tasks(tasks: Iterable[PregelExecutableTask]) -> Iterator[TaskPayload]:
+    """Produce "task" events for stream_mode=debug."""
+    for task in tasks:
+        if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):
+            continue
+
+        payload: TaskPayload = {
+            "id": task.id,
+            "name": task.name,
+            "input": task.input,
+            "triggers": task.triggers,
+        }
+        # Forward user-meaningful metadata only — drop coco_graph's internal
+        # framework keys (coco_graph_node/step/triggers/path/checkpoint_ns,
+        # thread_id, ...), which are redundant with the task's own fields and
+        # namespace. Keys like `lc_agent_name`, `ls_integration`, and any
+        # user-supplied metadata ride along. Filtered config tags are folded in
+        # under `tags`, mirroring the messages stream handler. (The comprehension
+        # also yields a fresh dict, so mutating `md` doesn't touch task.config.)
+        if task.config is not None:
+            md = {
+                k: v
+                for k, v in (task.config.get("metadata") or {}).items()
+                if k not in EXCLUDED_METADATA_KEYS
+            }
+            filtered_tags = filter_to_user_tags(task.config.get("tags"))
+            if filtered_tags is not None:
+                md["tags"] = filtered_tags
+            if md:
+                payload["metadata"] = md
+        yield payload
+
+
+def is_multiple_channel_write(value: Any) -> bool:
+    """Return True if the payload already wraps multiple writes from the same channel."""
+    return (
+        isinstance(value, dict)
+        and "$writes" in value
+        and isinstance(value["$writes"], list)
+    )
+
+
+def map_task_result_writes(writes: Sequence[tuple[str, Any]]) -> dict[str, Any]:
+    """Folds task writes into a result dict and aggregates multiple writes to the same channel.
+
+    If the channel contains a single write, we record the write in the result dict as `{channel: write}`
+    If the channel contains multiple writes, we record the writes in the result dict as `{channel: {'$writes': [write1, write2, ...]}}`"""
+
+    result: dict[str, Any] = {}
+    for channel, value in writes:
+        existing = result.get(channel)
+
+        if existing is not None:
+            channel_writes = (
+                existing["$writes"]
+                if is_multiple_channel_write(existing)
+                else [existing]
+            )
+            channel_writes.append(value)
+            result[channel] = {"$writes": channel_writes}
+        else:
+            result[channel] = value
+    return result
+
+
+def map_debug_task_results(
+    task_tup: tuple[PregelExecutableTask, Sequence[tuple[str, Any]]],
+    stream_keys: str | Sequence[str],
+) -> Iterator[TaskResultPayload]:
+    """Produce "task_result" events for stream_mode=debug."""
+    stream_channels_list = (
+        [stream_keys] if isinstance(stream_keys, str) else stream_keys
+    )
+    task, writes = task_tup
+    yield {
+        "id": task.id,
+        "name": task.name,
+        "error": next((w[1] for w in writes if w[0] == ERROR), None),
+        "result": map_task_result_writes(
+            [w for w in writes if w[0] in stream_channels_list or w[0] == RETURN]
+        ),
+        "interrupts": [
+            asdict(v)
+            for w in writes
+            if w[0] == INTERRUPT
+            for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
+        ],
+    }
+
+
+def rm_pregel_keys(config: RunnableConfig | None) -> RunnableConfig | None:
+    """Remove pregel-specific keys from the config."""
+    if config is None:
+        return config
+    return {
+        "configurable": {
+            k: v
+            for k, v in config.get("configurable", {}).items()
+            if not k.startswith("__pregel_")
+        }
+    }
+
+
+def map_debug_checkpoint(
+    config: RunnableConfig,
+    channels: Mapping[str, BaseChannel],
+    stream_channels: str | Sequence[str],
+    metadata: CheckpointMetadata,
+    tasks: Iterable[PregelExecutableTask],
+    pending_writes: list[PendingWrite],
+    parent_config: RunnableConfig | None,
+    output_keys: str | Sequence[str],
+) -> Iterator[CheckpointPayload]:
+    """Produce "checkpoint" events for stream_mode=debug."""
+
+    parent_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
+    task_states: dict[str, RunnableConfig | StateSnapshot] = {}
+
+    for task in tasks:
+        if not task.subgraphs:
+            continue
+
+        # assemble checkpoint_ns for this task
+        task_ns = f"{task.name}{NS_END}{task.id}"
+        if parent_ns:
+            task_ns = f"{parent_ns}{NS_SEP}{task_ns}"
+
+        # set config as signal that subgraph checkpoints exist
+        task_states[task.id] = {
+            CONF: {
+                "thread_id": config[CONF]["thread_id"],
+                CONFIG_KEY_CHECKPOINT_NS: task_ns,
+            }
+        }
+
+    yield {
+        "config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),
+        "parent_config": rm_pregel_keys(patch_checkpoint_map(parent_config, metadata)),
+        "values": read_channels(channels, stream_channels),
+        "metadata": metadata,
+        "next": [t.name for t in tasks],
+        "tasks": [
+            {
+                "id": t.id,
+                "name": t.name,
+                "error": t.error,
+                "state": t.state,
+            }
+            if t.error
+            else {
+                "id": t.id,
+                "name": t.name,
+                "result": t.result,
+                "interrupts": tuple(asdict(i) for i in t.interrupts),
+                "state": t.state,
+            }
+            if t.result
+            else {
+                "id": t.id,
+                "name": t.name,
+                "interrupts": tuple(asdict(i) for i in t.interrupts),
+                "state": t.state,
+            }
+            for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)
+        ],
+    }
+
+
+def tasks_w_writes(
+    tasks: Iterable[PregelTask | PregelExecutableTask],
+    pending_writes: list[PendingWrite] | None,
+    states: dict[str, RunnableConfig | StateSnapshot] | None,
+    output_keys: str | Sequence[str],
+) -> tuple[PregelTask, ...]:
+    """Apply writes / subgraph states to tasks to be returned in a StateSnapshot."""
+    pending_writes = pending_writes or []
+    out: list[PregelTask] = []
+    for task in tasks:
+        rtn = next(
+            (
+                val
+                for tid, chan, val in pending_writes
+                if tid == task.id and chan == RETURN
+            ),
+            MISSING,
+        )
+        task_error = next(
+            (exc for tid, n, exc in pending_writes if tid == task.id and n == ERROR),
+            None,
+        )
+        task_interrupts = tuple(
+            v
+            for tid, n, vv in pending_writes
+            if tid == task.id and n == INTERRUPT
+            for v in (vv if isinstance(vv, Sequence) else [vv])
+        )
+
+        task_writes = [
+            (chan, val)
+            for tid, chan, val in pending_writes
+            if tid == task.id and chan not in (ERROR, INTERRUPT, RETURN)
+        ]
+
+        if rtn is not MISSING:
+            task_result = rtn
+        elif isinstance(output_keys, str):
+            # unwrap single channel writes to just the write value
+            filtered_writes = [
+                (chan, val) for chan, val in task_writes if chan == output_keys
+            ]
+            mapped_writes = map_task_result_writes(filtered_writes)
+            task_result = mapped_writes.get(str(output_keys)) if mapped_writes else None
+        else:
+            if isinstance(output_keys, str):
+                output_keys = [output_keys]
+            # map task result writes to the desired output channels
+            # repeateed writes to the same channel are aggregated into: {'$writes': [write1, write2, ...]}
+            filtered_writes = [
+                (chan, val) for chan, val in task_writes if chan in output_keys
+            ]
+            mapped_writes = map_task_result_writes(filtered_writes)
+            task_result = mapped_writes if filtered_writes else {}
+
+        has_writes = rtn is not MISSING or any(
+            w[0] == task.id and w[1] not in (ERROR, INTERRUPT) for w in pending_writes
+        )
+
+        out.append(
+            PregelTask(
+                task.id,
+                task.name,
+                task.path,
+                task_error,
+                task_interrupts,
+                states.get(task.id) if states else None,
+                task_result if has_writes else None,
+            )
+        )
+    return tuple(out)
+
+
+COLOR_MAPPING = {
+    "black": "0;30",
+    "red": "0;31",
+    "green": "0;32",
+    "yellow": "0;33",
+    "blue": "0;34",
+    "magenta": "0;35",
+    "cyan": "0;36",
+    "white": "0;37",
+    "gray": "1;30",
+}
+
+
+def get_colored_text(text: str, color: str) -> str:
+    """Get colored text."""
+    return f"\033[1;3{COLOR_MAPPING[color]}m{text}\033[0m"
+
+
+def get_bolded_text(text: str) -> str:
+    """Get bolded text."""
+    return f"\033[1m{text}\033[0m"