codekingpro/portable-devtools

1from __future__ import annotations2 3from typing import Any, Literal, cast4from uuid import uuid45 6from langchain_core.messages import AnyMessage7from typing_extensions import TypedDict8 9from langgraph.config import get_config, get_stream_writer10from langgraph.constants import CONF11 12__all__ = (13    "UIMessage",14    "RemoveUIMessage",15    "AnyUIMessage",16    "push_ui_message",17    "delete_ui_message",18    "ui_message_reducer",19)20 21 22class UIMessage(TypedDict):23    """A message type for UI updates in LangGraph.24 25    This TypedDict represents a UI message that can be sent to update the UI state.26    It contains information about the UI component to render and its properties.27 28    Attributes:29        type: Literal type indicating this is a UI message.30        id: Unique identifier for the UI message.31        name: Name of the UI component to render.32        props: Properties to pass to the UI component.33        metadata: Additional metadata about the UI message.34    """35 36    type: Literal["ui"]37    id: str38    name: str39    props: dict[str, Any]40    metadata: dict[str, Any]41 42 43class RemoveUIMessage(TypedDict):44    """A message type for removing UI components in LangGraph.45 46    This TypedDict represents a message that can be sent to remove a UI component47    from the current state.48 49    Attributes:50        type: Literal type indicating this is a remove-ui message.51        id: Unique identifier of the UI message to remove.52    """53 54    type: Literal["remove-ui"]55    id: str56 57 58AnyUIMessage = UIMessage | RemoveUIMessage59 60 61def push_ui_message(62    name: str,63    props: dict[str, Any],64    *,65    id: str | None = None,66    metadata: dict[str, Any] | None = None,67    message: AnyMessage | None = None,68    state_key: str | None = "ui",69    merge: bool = False,70) -> UIMessage:71    """Push a new UI message to update the UI state.72 73    This function creates and sends a UI message that will be rendered in the UI.74    It also updates the graph state with the new UI message.75 76    Args:77        name: Name of the UI component to render.78        props: Properties to pass to the UI component.79        id: Optional unique identifier for the UI message.80            If not provided, a random UUID will be generated.81        metadata: Optional additional metadata about the UI message.82        message: Optional message object to associate with the UI message.83        state_key: Key in the graph state where the UI messages are stored.84        merge: Whether to merge props with existing UI message (True) or replace85            them (False).86 87    Returns:88        The created UI message.89 90    Example:91        ```python92        push_ui_message(93            name="component-name",94            props={"content": "Hello world"},95        )96        ```97 98    """99    from langgraph._internal._constants import CONFIG_KEY_SEND100 101    writer = get_stream_writer()102    config = get_config()103 104    message_id = None105    if message:106        if isinstance(message, dict) and "id" in message:107            message_id = message.get("id")108        elif hasattr(message, "id"):109            message_id = message.id110 111    evt: UIMessage = {112        "type": "ui",113        "id": id or str(uuid4()),114        "name": name,115        "props": props,116        "metadata": {117            "merge": merge,118            "run_id": config.get("run_id", None),119            "tags": config.get("tags", None),120            "name": config.get("run_name", None),121            **(metadata or {}),122            **({"message_id": message_id} if message_id else {}),123        },124    }125 126    writer(evt)127    if state_key:128        config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])129 130    return evt131 132 133def delete_ui_message(id: str, *, state_key: str = "ui") -> RemoveUIMessage:134    """Delete a UI message by ID from the UI state.135 136    This function creates and sends a message to remove a UI component from the current state.137    It also updates the graph state to remove the UI message.138 139    Args:140        id: Unique identifier of the UI component to remove.141        state_key: Key in the graph state where the UI messages are stored. Defaults to "ui".142 143    Returns:144        The remove UI message.145 146    Example:147        ```python148        delete_ui_message("message-123")149        ```150 151    """152    from langgraph._internal._constants import CONFIG_KEY_SEND153 154    writer = get_stream_writer()155    config = get_config()156 157    evt: RemoveUIMessage = {"type": "remove-ui", "id": id}158 159    writer(evt)160    config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])161 162    return evt163 164 165def ui_message_reducer(166    left: list[AnyUIMessage] | AnyUIMessage,167    right: list[AnyUIMessage] | AnyUIMessage,168) -> list[AnyUIMessage]:169    """Merge two lists of UI messages, supporting removing UI messages.170 171    This function combines two lists of UI messages, handling both regular UI messages172    and `remove-ui` messages. When a `remove-ui` message is encountered, it removes any173    UI message with the matching ID from the current state.174 175    Args:176        left: First list of UI messages or single UI message.177        right: Second list of UI messages or single UI message.178 179    Returns:180        Combined list of UI messages with removals applied.181 182    Example:183        ```python184        messages = ui_message_reducer(185            [{"type": "ui", "id": "1", "name": "Chat", "props": {}}],186            {"type": "remove-ui", "id": "1"},187        )188        ```189 190    """191    if not isinstance(left, list):192        left = [left]193 194    if not isinstance(right, list):195        right = [right]196 197    # merge messages198    merged = left.copy()199    merged_by_id = {m.get("id"): i for i, m in enumerate(merged)}200    ids_to_remove = set()201 202    for msg in right:203        msg_id = msg.get("id")204 205        if (existing_idx := merged_by_id.get(msg_id)) is not None:206            if msg.get("type") == "remove-ui":207                ids_to_remove.add(msg_id)208            else:209                ids_to_remove.discard(msg_id)210 211                if cast(UIMessage, msg).get("metadata", {}).get("merge", False):212                    prev_msg = merged[existing_idx]213                    msg = msg.copy()214                    msg["props"] = {**prev_msg["props"], **msg["props"]}215 216                merged[existing_idx] = msg217        else:218            if msg.get("type") == "remove-ui":219                raise ValueError(220                    f"Attempting to delete an UI message with an ID that doesn't exist ('{msg_id}')"221                )222 223            merged_by_id[msg_id] = len(merged)224            merged.append(msg)225 226    merged = [m for m in merged if m.get("id") not in ids_to_remove]227    return merged228