codekingpro/portable-devtools

1from __future__ import annotations2 3from collections.abc import Iterable, Iterator, Mapping, Sequence4from dataclasses import asdict5from typing import Any6from uuid import UUID7 8from langchain_core.runnables import RunnableConfig9from langgraph.checkpoint.base import CheckpointMetadata, PendingWrite10 11from langgraph._internal._config import patch_checkpoint_map12from langgraph._internal._constants import (13    CONF,14    CONFIG_KEY_CHECKPOINT_NS,15    ERROR,16    INTERRUPT,17    NS_END,18    NS_SEP,19    RETURN,20)21from langgraph._internal._typing import MISSING22from langgraph.channels.base import BaseChannel23from langgraph.constants import TAG_HIDDEN24from langgraph.pregel._io import read_channels25from langgraph.types import (26    CheckpointPayload,27    PregelExecutableTask,28    PregelTask,29    StateSnapshot,30    TaskPayload,31    TaskResultPayload,32)33 34TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")35 36 37def map_debug_tasks(tasks: Iterable[PregelExecutableTask]) -> Iterator[TaskPayload]:38    """Produce "task" events for stream_mode=debug."""39    for task in tasks:40        if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):41            continue42 43        yield {44            "id": task.id,45            "name": task.name,46            "input": task.input,47            "triggers": task.triggers,48        }49 50 51def is_multiple_channel_write(value: Any) -> bool:52    """Return True if the payload already wraps multiple writes from the same channel."""53    return (54        isinstance(value, dict)55        and "$writes" in value56        and isinstance(value["$writes"], list)57    )58 59 60def map_task_result_writes(writes: Sequence[tuple[str, Any]]) -> dict[str, Any]:61    """Folds task writes into a result dict and aggregates multiple writes to the same channel.62 63    If the channel contains a single write, we record the write in the result dict as `{channel: write}`64    If the channel contains multiple writes, we record the writes in the result dict as `{channel: {'$writes': [write1, write2, ...]}}`"""65 66    result: dict[str, Any] = {}67    for channel, value in writes:68        existing = result.get(channel)69 70        if existing is not None:71            channel_writes = (72                existing["$writes"]73                if is_multiple_channel_write(existing)74                else [existing]75            )76            channel_writes.append(value)77            result[channel] = {"$writes": channel_writes}78        else:79            result[channel] = value80    return result81 82 83def map_debug_task_results(84    task_tup: tuple[PregelExecutableTask, Sequence[tuple[str, Any]]],85    stream_keys: str | Sequence[str],86) -> Iterator[TaskResultPayload]:87    """Produce "task_result" events for stream_mode=debug."""88    stream_channels_list = (89        [stream_keys] if isinstance(stream_keys, str) else stream_keys90    )91    task, writes = task_tup92    yield {93        "id": task.id,94        "name": task.name,95        "error": next((w[1] for w in writes if w[0] == ERROR), None),96        "result": map_task_result_writes(97            [w for w in writes if w[0] in stream_channels_list or w[0] == RETURN]98        ),99        "interrupts": [100            asdict(v)101            for w in writes102            if w[0] == INTERRUPT103            for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])104        ],105    }106 107 108def rm_pregel_keys(config: RunnableConfig | None) -> RunnableConfig | None:109    """Remove pregel-specific keys from the config."""110    if config is None:111        return config112    return {113        "configurable": {114            k: v115            for k, v in config.get("configurable", {}).items()116            if not k.startswith("__pregel_")117        }118    }119 120 121def map_debug_checkpoint(122    config: RunnableConfig,123    channels: Mapping[str, BaseChannel],124    stream_channels: str | Sequence[str],125    metadata: CheckpointMetadata,126    tasks: Iterable[PregelExecutableTask],127    pending_writes: list[PendingWrite],128    parent_config: RunnableConfig | None,129    output_keys: str | Sequence[str],130) -> Iterator[CheckpointPayload]:131    """Produce "checkpoint" events for stream_mode=debug."""132 133    parent_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")134    task_states: dict[str, RunnableConfig | StateSnapshot] = {}135 136    for task in tasks:137        if not task.subgraphs:138            continue139 140        # assemble checkpoint_ns for this task141        task_ns = f"{task.name}{NS_END}{task.id}"142        if parent_ns:143            task_ns = f"{parent_ns}{NS_SEP}{task_ns}"144 145        # set config as signal that subgraph checkpoints exist146        task_states[task.id] = {147            CONF: {148                "thread_id": config[CONF]["thread_id"],149                CONFIG_KEY_CHECKPOINT_NS: task_ns,150            }151        }152 153    yield {154        "config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),155        "parent_config": rm_pregel_keys(patch_checkpoint_map(parent_config, metadata)),156        "values": read_channels(channels, stream_channels),157        "metadata": metadata,158        "next": [t.name for t in tasks],159        "tasks": [160            {161                "id": t.id,162                "name": t.name,163                "error": t.error,164                "state": t.state,165            }166            if t.error167            else {168                "id": t.id,169                "name": t.name,170                "result": t.result,171                "interrupts": tuple(asdict(i) for i in t.interrupts),172                "state": t.state,173            }174            if t.result175            else {176                "id": t.id,177                "name": t.name,178                "interrupts": tuple(asdict(i) for i in t.interrupts),179                "state": t.state,180            }181            for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)182        ],183    }184 185 186def tasks_w_writes(187    tasks: Iterable[PregelTask | PregelExecutableTask],188    pending_writes: list[PendingWrite] | None,189    states: dict[str, RunnableConfig | StateSnapshot] | None,190    output_keys: str | Sequence[str],191) -> tuple[PregelTask, ...]:192    """Apply writes / subgraph states to tasks to be returned in a StateSnapshot."""193    pending_writes = pending_writes or []194    out: list[PregelTask] = []195    for task in tasks:196        rtn = next(197            (198                val199                for tid, chan, val in pending_writes200                if tid == task.id and chan == RETURN201            ),202            MISSING,203        )204        task_error = next(205            (exc for tid, n, exc in pending_writes if tid == task.id and n == ERROR),206            None,207        )208        task_interrupts = tuple(209            v210            for tid, n, vv in pending_writes211            if tid == task.id and n == INTERRUPT212            for v in (vv if isinstance(vv, Sequence) else [vv])213        )214 215        task_writes = [216            (chan, val)217            for tid, chan, val in pending_writes218            if tid == task.id and chan not in (ERROR, INTERRUPT, RETURN)219        ]220 221        if rtn is not MISSING:222            task_result = rtn223        elif isinstance(output_keys, str):224            # unwrap single channel writes to just the write value225            filtered_writes = [226                (chan, val) for chan, val in task_writes if chan == output_keys227            ]228            mapped_writes = map_task_result_writes(filtered_writes)229            task_result = mapped_writes.get(str(output_keys)) if mapped_writes else None230        else:231            if isinstance(output_keys, str):232                output_keys = [output_keys]233            # map task result writes to the desired output channels234            # repeateed writes to the same channel are aggregated into: {'$writes': [write1, write2, ...]}235            filtered_writes = [236                (chan, val) for chan, val in task_writes if chan in output_keys237            ]238            mapped_writes = map_task_result_writes(filtered_writes)239            task_result = mapped_writes if filtered_writes else {}240 241        has_writes = rtn is not MISSING or any(242            w[0] == task.id and w[1] not in (ERROR, INTERRUPT) for w in pending_writes243        )244 245        out.append(246            PregelTask(247                task.id,248                task.name,249                task.path,250                task_error,251                task_interrupts,252                states.get(task.id) if states else None,253                task_result if has_writes else None,254            )255        )256    return tuple(out)257 258 259COLOR_MAPPING = {260    "black": "0;30",261    "red": "0;31",262    "green": "0;32",263    "yellow": "0;33",264    "blue": "0;34",265    "magenta": "0;35",266    "cyan": "0;36",267    "white": "0;37",268    "gray": "1;30",269}270 271 272def get_colored_text(text: str, color: str) -> str:273    """Get colored text."""274    return f"\033[1;3{COLOR_MAPPING[color]}m{text}\033[0m"275 276 277def get_bolded_text(text: str) -> str:278    """Get bolded text."""279    return f"\033[1m{text}\033[0m"280