codekingpro/portable-devtools

1from __future__ import annotations2 3import uuid4import warnings5from collections.abc import Callable, Sequence6from functools import partial7from typing import (8    Annotated,9    Any,10    Literal,11    cast,12)13 14from langchain_core.messages import (15    AnyMessage,16    BaseMessage,17    BaseMessageChunk,18    MessageLikeRepresentation,19    RemoveMessage,20    convert_to_messages,21    message_chunk_to_message,22)23from typing_extensions import TypedDict, deprecated24 25from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP26from langgraph.graph.state import StateGraph27from langgraph.warnings import LangGraphDeprecatedSinceV1028 29__all__ = (30    "add_messages",31    "MessagesState",32    "MessageGraph",33    "REMOVE_ALL_MESSAGES",34)35 36Messages = list[MessageLikeRepresentation] | MessageLikeRepresentation37 38REMOVE_ALL_MESSAGES = "__remove_all__"39 40 41def _add_messages_wrapper(func: Callable) -> Callable[[Messages, Messages], Messages]:42    def _add_messages(43        left: Messages | None = None, right: Messages | None = None, **kwargs: Any44    ) -> Messages | Callable[[Messages, Messages], Messages]:45        if left is not None and right is not None:46            return func(left, right, **kwargs)47        elif left is not None or right is not None:48            msg = (49                f"Must specify non-null arguments for both 'left' and 'right'. Only "50                f"received: '{'left' if left else 'right'}'."51            )52            raise ValueError(msg)53        else:54            return partial(func, **kwargs)55 56    _add_messages.__doc__ = func.__doc__57    return cast(Callable[[Messages, Messages], Messages], _add_messages)58 59 60@_add_messages_wrapper61def add_messages(62    left: Messages,63    right: Messages,64    *,65    format: Literal["langchain-openai"] | None = None,66) -> Messages:67    """Merges two lists of messages, updating existing messages by ID.68 69    By default, this ensures the state is "append-only", unless the70    new message has the same ID as an existing message.71 72    Args:73        left: The base list of `Messages`.74        right: The list of `Messages` (or single `Message`) to merge75            into the base list.76        format: The format to return messages in. If `None` then `Messages` will be77            returned as is. If `langchain-openai` then `Messages` will be returned as78            `BaseMessage` objects with their contents formatted to match OpenAI message79            format, meaning contents can be string, `'text'` blocks, or `'image_url'` blocks80            and tool responses are returned as their own `ToolMessage` objects.81 82            !!! important "Requirement"83 84                Must have `langchain-core>=0.3.11` installed to use this feature.85 86    Returns:87        A new list of messages with the messages from `right` merged into `left`.88        If a message in `right` has the same ID as a message in `left`, the89            message from `right` will replace the message from `left`.90 91    Example: Basic usage92        ```python93        from langchain_core.messages import AIMessage, HumanMessage94 95        msgs1 = [HumanMessage(content="Hello", id="1")]96        msgs2 = [AIMessage(content="Hi there!", id="2")]97        add_messages(msgs1, msgs2)98        # [HumanMessage(content='Hello', id='1'), AIMessage(content='Hi there!', id='2')]99        ```100 101    Example: Overwrite existing message102        ```python103        msgs1 = [HumanMessage(content="Hello", id="1")]104        msgs2 = [HumanMessage(content="Hello again", id="1")]105        add_messages(msgs1, msgs2)106        # [HumanMessage(content='Hello again', id='1')]107        ```108 109    Example: Use in a StateGraph110        ```python111        from typing import Annotated112        from typing_extensions import TypedDict113        from langgraph.graph import StateGraph114 115 116        class State(TypedDict):117            messages: Annotated[list, add_messages]118 119 120        builder = StateGraph(State)121        builder.add_node("chatbot", lambda state: {"messages": [("assistant", "Hello")]})122        builder.set_entry_point("chatbot")123        builder.set_finish_point("chatbot")124        graph = builder.compile()125        graph.invoke({})126        # {'messages': [AIMessage(content='Hello', id=...)]}127        ```128 129    Example: Use OpenAI message format130        ```python131        from typing import Annotated132        from typing_extensions import TypedDict133        from langgraph.graph import StateGraph, add_messages134 135 136        class State(TypedDict):137            messages: Annotated[list, add_messages(format="langchain-openai")]138 139 140        def chatbot_node(state: State) -> list:141            return {142                "messages": [143                    {144                        "role": "user",145                        "content": [146                            {147                                "type": "text",148                                "text": "Here's an image:",149                                "cache_control": {"type": "ephemeral"},150                            },151                            {152                                "type": "image",153                                "source": {154                                    "type": "base64",155                                    "media_type": "image/jpeg",156                                    "data": "1234",157                                },158                            },159                        ],160                    },161                ]162            }163 164 165        builder = StateGraph(State)166        builder.add_node("chatbot", chatbot_node)167        builder.set_entry_point("chatbot")168        builder.set_finish_point("chatbot")169        graph = builder.compile()170        graph.invoke({"messages": []})171        # {172        #     'messages': [173        #         HumanMessage(174        #             content=[175        #                 {"type": "text", "text": "Here's an image:"},176        #                 {177        #                     "type": "image_url",178        #                     "image_url": {"url": "data:image/jpeg;base64,1234"},179        #                 },180        #             ],181        #         ),182        #     ]183        # }184        ```185 186    """187    remove_all_idx = None188    # coerce to list189    if not isinstance(left, list):190        left = [left]  # type: ignore[assignment]191    if not isinstance(right, list):192        right = [right]  # type: ignore[assignment]193    # coerce to message194    left = [195        message_chunk_to_message(cast(BaseMessageChunk, m))196        for m in convert_to_messages(left)197    ]198    right = [199        message_chunk_to_message(cast(BaseMessageChunk, m))200        for m in convert_to_messages(right)201    ]202    # assign missing ids203    for m in left:204        if m.id is None:205            m.id = str(uuid.uuid4())206    for idx, m in enumerate(right):207        if m.id is None:208            m.id = str(uuid.uuid4())209        if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:210            remove_all_idx = idx211 212    if remove_all_idx is not None:213        return right[remove_all_idx + 1 :]214 215    # merge216    merged = left.copy()217    merged_by_id = {m.id: i for i, m in enumerate(merged)}218    ids_to_remove = set()219    for m in right:220        if (existing_idx := merged_by_id.get(m.id)) is not None:221            if isinstance(m, RemoveMessage):222                ids_to_remove.add(m.id)223            else:224                ids_to_remove.discard(m.id)225                merged[existing_idx] = m226        else:227            if isinstance(m, RemoveMessage):228                raise ValueError(229                    f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"230                )231 232            merged_by_id[m.id] = len(merged)233            merged.append(m)234    merged = [m for m in merged if m.id not in ids_to_remove]235 236    if format == "langchain-openai":237        merged = _format_messages(merged)238    elif format:239        msg = f"Unrecognized {format=}. Expected one of 'langchain-openai', None."240        raise ValueError(msg)241    else:242        pass243 244    return merged245 246 247def _messages_delta_reducer(248    state: list[AnyMessage], writes: list[list[AnyMessage]]249) -> list[AnyMessage]:250    """**Experimental.** Batch reducer for use with `DeltaChannel`.251 252    Processes all writes in one pass — dedup by ID, `RemoveMessage`253    tombstoning — without calling `add_messages`.254 255    This reducer is batching-invariant, as required by `DeltaChannel`:256    `reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.257 258    Raw dict / string / tuple inputs are coerced to typed `BaseMessage`259    objects so that HTTP-driven graphs work without a separate coercion260    step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,261    unknown-id `RemoveMessage` errors, missing-id UUID assignment, and262    `BaseMessageChunk` conversion are not handled here.263 264    Example::265 266        from typing import Annotated267        from langgraph.channels.delta import DeltaChannel268        from langgraph.graph.message import _messages_delta_reducer269 270        class State(TypedDict):271            messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]272    """273 274    # Each write is either a list of message-likes or a single message-like275    # (BaseMessage / dict / str / tuple). Only lists flatten; everything276    # else is one message.277    flat: list[Any] = []278    for w in writes:279        if isinstance(w, list):280            flat.extend(w)281        else:282            flat.append(w)283    # Steady state: the reducer's own output is already typed, so skip284    # `convert_to_messages` on state when the first element is a BaseMessage.285    # Only raw input (initial dicts, deserialized blobs) hits the slow path.286    if state and isinstance(state[0], BaseMessage):287        state_msgs = state288    else:289        state_msgs = cast("list[AnyMessage]", convert_to_messages(state))290    msgs = cast("list[AnyMessage]", convert_to_messages(flat))291 292    index: dict[str, int] = {293        m.id: i for i, m in enumerate(state_msgs) if m.id is not None294    }295    result: list[AnyMessage | None] = list(state_msgs)296    for msg in msgs:297        mid = msg.id298        if mid is None:299            result.append(msg)300        elif isinstance(msg, RemoveMessage):301            if mid in index:302                result[index[mid]] = None303                del index[mid]304        elif mid in index:305            result[index[mid]] = msg306        else:307            index[mid] = len(result)308            result.append(msg)309    return [m for m in result if m is not None]310 311 312@deprecated(313    "MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",314    category=None,315)316class MessageGraph(StateGraph):317    """A StateGraph where every node receives a list of messages as input and returns one or more messages as output.318 319    MessageGraph is a subclass of StateGraph whose entire state is a single, append-only* list of messages.320    Each node in a MessageGraph takes a list of messages as input and returns zero or more321    messages as output. The `add_messages` function is used to merge the output messages from each node322    into the existing list of messages in the graph's state.323 324    Examples:325        ```pycon326        >>> from langgraph.graph.message import MessageGraph327        ...328        >>> builder = MessageGraph()329        >>> builder.add_node("chatbot", lambda state: [("assistant", "Hello!")])330        >>> builder.set_entry_point("chatbot")331        >>> builder.set_finish_point("chatbot")332        >>> builder.compile().invoke([("user", "Hi there.")])333        [HumanMessage(content="Hi there.", id='...'), AIMessage(content="Hello!", id='...')]334        ```335 336        ```pycon337        >>> from langchain_core.messages import AIMessage, HumanMessage, ToolMessage338        >>> from langgraph.graph.message import MessageGraph339        ...340        >>> builder = MessageGraph()341        >>> builder.add_node(342        ...     "chatbot",343        ...     lambda state: [344        ...         AIMessage(345        ...             content="Hello!",346        ...             tool_calls=[{"name": "search", "id": "123", "args": {"query": "X"}}],347        ...         )348        ...     ],349        ... )350        >>> builder.add_node(351        ...     "search", lambda state: [ToolMessage(content="Searching...", tool_call_id="123")]352        ... )353        >>> builder.set_entry_point("chatbot")354        >>> builder.add_edge("chatbot", "search")355        >>> builder.set_finish_point("search")356        >>> builder.compile().invoke([HumanMessage(content="Hi there. Can you search for X?")])357        {'messages': [HumanMessage(content="Hi there. Can you search for X?", id='b8b7d8f4-7f4d-4f4d-9c1d-f8b8d8f4d9c1'),358                     AIMessage(content="Hello!", id='f4d9c1d8-8d8f-4d9c-b8b7-d8f4f4d9c1d8'),359                     ToolMessage(content="Searching...", id='d8f4f4d9-c1d8-4f4d-b8b7-d8f4f4d9c1d8', tool_call_id="123")]}360        ```361    """362 363    def __init__(self) -> None:364        warnings.warn(365            "MessageGraph is deprecated in LangGraph v1.0.0, to be removed in v2.0.0. Please use StateGraph with a `messages` key instead.",366            category=LangGraphDeprecatedSinceV10,367            stacklevel=2,368        )369        super().__init__(Annotated[list[AnyMessage], add_messages])  # type: ignore[arg-type]370 371 372class MessagesState(TypedDict):373    messages: Annotated[list[AnyMessage], add_messages]374 375 376def _format_messages(messages: Sequence[BaseMessage]) -> list[BaseMessage]:377    try:378        from langchain_core.messages import convert_to_openai_messages379    except ImportError:380        msg = (381            "Must have langchain-core>=0.3.11 installed to use automatic message "382            "formatting (format='langchain-openai'). Please update your langchain-core "383            "version or remove the 'format' flag. Returning un-formatted "384            "messages."385        )386        warnings.warn(msg)387        return list(messages)388    else:389        return convert_to_messages(convert_to_openai_messages(messages))390 391 392def push_message(393    message: MessageLikeRepresentation | BaseMessageChunk,394    *,395    state_key: str | None = "messages",396) -> AnyMessage:397    """Write a message manually to the `messages` / `messages-tuple` stream mode.398 399    Will automatically write to the channel specified in the `state_key` unless `state_key` is `None`.400    """401 402    from langchain_core.callbacks.base import (403        BaseCallbackHandler,404        BaseCallbackManager,405    )406 407    from langgraph.config import get_config408    from langgraph.pregel._messages import StreamMessagesHandler409 410    config = get_config()411    message = next(x for x in convert_to_messages([message]))412 413    if message.id is None:414        raise ValueError("Message ID is required")415 416    if isinstance(config["callbacks"], BaseCallbackManager):417        manager = config["callbacks"]418        handlers = manager.handlers419    elif isinstance(config["callbacks"], list) and all(420        isinstance(x, BaseCallbackHandler) for x in config["callbacks"]421    ):422        handlers = config["callbacks"]423 424    if stream_handler := next(425        (x for x in handlers if isinstance(x, StreamMessagesHandler)), None426    ):427        metadata = config["metadata"]428        message_meta = (429            tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),430            metadata,431        )432        stream_handler._emit(message_meta, message, dedupe=False)433 434    if state_key:435        config[CONF][CONFIG_KEY_SEND]([(state_key, message)])436 437    return message438