codekingpro/portable-devtools

1"""Transformer that projects `tools` channel events into `ToolCallStream`s."""2 3from __future__ import annotations4 5from collections.abc import Awaitable, Callable6from typing import Any7 8from langgraph.stream._types import ProtocolEvent, StreamTransformer9from langgraph.stream.stream_channel import StreamChannel10 11from langgraph.prebuilt._tool_call_stream import ToolCallStream12 13 14class ToolCallTransformer(StreamTransformer):15    """Project `tools` channel events into `ToolCallStream` handles.16 17    Each `tool-started` event spawns a `ToolCallStream`, pushed onto18    `run.tool_calls`. Subsequent `tool-output-delta` events append to19    that stream's deltas log; `tool-finished` and `tool-error` close it.20 21    Native transformer — the `tool_calls` projection is exposed as a22    direct attribute on the run stream.23 24    A nameless `StreamChannel[ToolCallStream]` is used (no protocol25    auto-forwarding) because the live handles are not serializable and26    should not be injected into the main event log. Wire consumers27    subscribe to the `tools` channel instead, where the raw protocol28    events flow through untouched by this transformer (`process`29    returns `True`).30 31    Registered explicitly by users at compile time via32    `builder.compile(transformers=[ToolCallTransformer])` — not a33    default built-in, so the `tools` channel is user-opt-in.34    """35 36    _native = True37    required_stream_modes = ("tools",)38 39    def __init__(self, scope: tuple[str, ...] = ()) -> None:40        super().__init__(scope)41        self._log: StreamChannel[ToolCallStream] = StreamChannel()42        self._active: dict[str, ToolCallStream] = {}43        self._is_async = False44        self._pump_fn: Callable[[], bool] | None = None45        self._apump_fn: Callable[[], Awaitable[bool]] | None = None46 47    def init(self) -> dict[str, Any]:48        return {"tool_calls": self._log}49 50    def _bind_pump(self, fn: Callable[[], bool]) -> None:51        """Wire the sync pull callback onto this transformer.52 53        Called by `StreamMux.bind_pump`. Stored so each new54        `ToolCallStream` created by `process` can wire its deltas log55        for pump-driven iteration.56        """57        self._pump_fn = fn58        self._is_async = False59 60    def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:61        """Async counterpart to `_bind_pump`."""62        self._apump_fn = fn63        self._is_async = True64 65    def _new_stream(66        self,67        tool_call_id: str,68        tool_name: str,69        tool_input: dict[str, Any] | None,70    ) -> ToolCallStream:71        stream = ToolCallStream(tool_call_id, tool_name, tool_input)72        stream._bind(is_async=self._is_async)73        if self._apump_fn is not None:74            stream._output_deltas._arequest_more = self._apump_fn75        if self._pump_fn is not None:76            stream._output_deltas._request_more = self._pump_fn77        return stream78 79    def process(self, event: ProtocolEvent) -> bool:80        if event["method"] != "tools":81            return True82 83        # Only project events emitted at this transformer's scope. Subgraph84        # events still flow through the parent's mux (the parent's main85        # event log keeps them) but they belong to the child mini-mux's86        # `tool_calls` projection, not the parent's.87        if tuple(event["params"]["namespace"]) != self.scope:88            return True89 90        data = event["params"]["data"]91        tool_call_id = data.get("tool_call_id")92        if tool_call_id is None:93            return True94        event_type = data.get("event")95 96        stream: ToolCallStream | None97        if event_type == "tool-started":98            stream = self._new_stream(99                tool_call_id,100                data.get("tool_name", ""),101                data.get("input"),102            )103            self._active[tool_call_id] = stream104            self._log.push(stream)105        elif event_type == "tool-output-delta":106            stream = self._active.get(tool_call_id)107            if stream is not None:108                stream._push_delta(data.get("delta"))109        elif event_type == "tool-finished":110            stream = self._active.pop(tool_call_id, None)111            if stream is not None:112                stream._finish(data.get("output"))113        elif event_type == "tool-error":114            stream = self._active.pop(tool_call_id, None)115            if stream is not None:116                stream._fail(data.get("message", ""))117 118        # Pass-through — wire consumers subscribe to the `tools` channel119        # directly and reconstruct handles client-side.120        return True121 122    def finalize(self) -> None:123        """Close any still-active tool streams left open at run end."""124        for stream in self._active.values():125            if not stream.completed:126                stream._finish(None)127        self._active.clear()128 129    def fail(self, err: BaseException) -> None:130        """Fail any still-active tool streams when the run errors."""131        message = str(err)132        for stream in self._active.values():133            if not stream.completed:134                stream._fail(message)135        self._active.clear()136