1"""Module that contains tests for runnable.astream_events API."""23import asyncio4import inspect5import sys6import uuid7from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Sequence8from functools import partial9from itertools import cycle10from typing import (11    Any,12    cast,13)1415import pytest16from blockbuster import BlockBuster17from pydantic import BaseModel18from typing_extensions import override1920from langchain_core.callbacks import CallbackManagerForRetrieverRun, Callbacks21from langchain_core.callbacks.manager import (22    adispatch_custom_event,23)24from langchain_core.chat_history import BaseChatMessageHistory25from langchain_core.documents import Document26from langchain_core.language_models import FakeStreamingListLLM, GenericFakeChatModel27from langchain_core.messages import (28    AIMessage,29    AIMessageChunk,30    BaseMessage,31    HumanMessage,32    SystemMessage,33)34from langchain_core.prompt_values import ChatPromptValue35from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder36from langchain_core.retrievers import BaseRetriever37from langchain_core.runnables import (38    ConfigurableField,39    Runnable,40    RunnableConfig,41    RunnableGenerator,42    RunnableLambda,43    chain,44    ensure_config,45)46from langchain_core.runnables.config import (47    get_async_callback_manager_for_config,48)49from langchain_core.runnables.history import RunnableWithMessageHistory50from langchain_core.runnables.schema import StreamEvent51from langchain_core.runnables.utils import Addable52from langchain_core.tools import tool53from langchain_core.utils.aiter import aclosing54from tests.unit_tests.runnables.test_runnable_events_v1 import (55    _assert_events_equal_allow_superset_metadata,56)57from tests.unit_tests.stubs import _any_id_ai_message, _any_id_ai_message_chunk5859# The v2 event tests include a compatibility case for `RunnableWithMessageHistory`,60# so constructing that deprecated wrapper is expected in this module.61pytestmark = pytest.mark.filterwarnings(62    "ignore:RunnableWithMessageHistory is deprecated. Use LangGraph's built-in "63    "persistence instead.:"64    "langchain_core._api.deprecation.LangChainDeprecationWarning"65)666768def _with_nulled_run_id(events: Sequence[StreamEvent]) -> list[StreamEvent]:69    """Removes the run IDs from events."""70    for event in events:71        assert "run_id" in event, f"Event {event} does not have a run_id."72        assert "parent_ids" in event, f"Event {event} does not have parent_ids."73        assert isinstance(event["run_id"], str), (74            f"Event {event} run_id is not a string."75        )76        assert isinstance(event["parent_ids"], list), (77            f"Event {event} parent_ids is not a list."78        )7980    return cast(81        "list[StreamEvent]",82        [{**event, "run_id": "", "parent_ids": []} for event in events],83    )848586async def _collect_events(87    events: AsyncIterator[StreamEvent], *, with_nulled_ids: bool = True88) -> list[StreamEvent]:89    """Collect the events and remove the run ids."""90    materialized_events = [event async for event in events]9192    if with_nulled_ids:93        events_ = _with_nulled_run_id(materialized_events)94    else:95        events_ = materialized_events96    for event in events_:97        event["tags"] = sorted(event["tags"])98    return events_99100101async def test_event_stream_with_simple_function_tool() -> None:102    """Test the event stream with a function and tool."""103104    def foo(x: int) -> dict[str, int]:105        """Foo."""106        _ = x107        return {"x": 5}108109    @tool110    def get_docs(x: int) -> list[Document]:111        """Hello Doc."""112        _ = x113        return [Document(page_content="hello")]114115    chain = RunnableLambda(foo) | get_docs116    events = await _collect_events(chain.astream_events({}, version="v2"))117    _assert_events_equal_allow_superset_metadata(118        events,119        [120            {121                "event": "on_chain_start",122                "run_id": "",123                "parent_ids": [],124                "name": "RunnableSequence",125                "tags": [],126                "metadata": {},127                "data": {"input": {}},128            },129            {130                "event": "on_chain_start",131                "name": "foo",132                "run_id": "",133                "parent_ids": [],134                "tags": ["seq:step:1"],135                "metadata": {},136                "data": {},137            },138            {139                "event": "on_chain_stream",140                "name": "foo",141                "run_id": "",142                "parent_ids": [],143                "tags": ["seq:step:1"],144                "metadata": {},145                "data": {"chunk": {"x": 5}},146            },147            {148                "event": "on_chain_end",149                "name": "foo",150                "run_id": "",151                "parent_ids": [],152                "tags": ["seq:step:1"],153                "metadata": {},154                "data": {"input": {}, "output": {"x": 5}},155            },156            {157                "event": "on_tool_start",158                "name": "get_docs",159                "run_id": "",160                "parent_ids": [],161                "tags": ["seq:step:2"],162                "metadata": {},163                "data": {"input": {"x": 5}},164            },165            {166                "event": "on_tool_end",167                "name": "get_docs",168                "run_id": "",169                "parent_ids": [],170                "tags": ["seq:step:2"],171                "metadata": {},172                "data": {"input": {"x": 5}, "output": [Document(page_content="hello")]},173            },174            {175                "event": "on_chain_stream",176                "run_id": "",177                "parent_ids": [],178                "tags": [],179                "metadata": {},180                "name": "RunnableSequence",181                "data": {"chunk": [Document(page_content="hello")]},182            },183            {184                "event": "on_chain_end",185                "name": "RunnableSequence",186                "run_id": "",187                "parent_ids": [],188                "tags": [],189                "metadata": {},190                "data": {"output": [Document(page_content="hello")]},191            },192        ],193    )194195196async def test_event_stream_with_single_lambda() -> None:197    """Test the event stream with a tool."""198199    def reverse(s: str) -> str:200        """Reverse a string."""201        return s[::-1]202203    chain = RunnableLambda(func=reverse)204205    events = await _collect_events(chain.astream_events("hello", version="v2"))206    _assert_events_equal_allow_superset_metadata(207        events,208        [209            {210                "data": {"input": "hello"},211                "event": "on_chain_start",212                "metadata": {},213                "name": "reverse",214                "run_id": "",215                "parent_ids": [],216                "tags": [],217            },218            {219                "data": {"chunk": "olleh"},220                "event": "on_chain_stream",221                "metadata": {},222                "name": "reverse",223                "run_id": "",224                "parent_ids": [],225                "tags": [],226            },227            {228                "data": {"output": "olleh"},229                "event": "on_chain_end",230                "metadata": {},231                "name": "reverse",232                "run_id": "",233                "parent_ids": [],234                "tags": [],235            },236        ],237    )238239240async def test_event_stream_with_triple_lambda() -> None:241    def reverse(s: str) -> str:242        """Reverse a string."""243        return s[::-1]244245    r = RunnableLambda(func=reverse)246247    chain = (248        r.with_config({"run_name": "1"})249        | r.with_config({"run_name": "2"})250        | r.with_config({"run_name": "3"})251    )252    events = await _collect_events(chain.astream_events("hello", version="v2"))253    _assert_events_equal_allow_superset_metadata(254        events,255        [256            {257                "data": {"input": "hello"},258                "event": "on_chain_start",259                "metadata": {},260                "name": "RunnableSequence",261                "run_id": "",262                "parent_ids": [],263                "tags": [],264            },265            {266                "data": {},267                "event": "on_chain_start",268                "metadata": {},269                "name": "1",270                "run_id": "",271                "parent_ids": [],272                "tags": ["seq:step:1"],273            },274            {275                "data": {"chunk": "olleh"},276                "event": "on_chain_stream",277                "metadata": {},278                "name": "1",279                "run_id": "",280                "parent_ids": [],281                "tags": ["seq:step:1"],282            },283            {284                "data": {},285                "event": "on_chain_start",286                "metadata": {},287                "name": "2",288                "run_id": "",289                "parent_ids": [],290                "tags": ["seq:step:2"],291            },292            {293                "data": {"input": "hello", "output": "olleh"},294                "event": "on_chain_end",295                "metadata": {},296                "name": "1",297                "run_id": "",298                "parent_ids": [],299                "tags": ["seq:step:1"],300            },301            {302                "data": {"chunk": "hello"},303                "event": "on_chain_stream",304                "metadata": {},305                "name": "2",306                "run_id": "",307                "parent_ids": [],308                "tags": ["seq:step:2"],309            },310            {311                "data": {},312                "event": "on_chain_start",313                "metadata": {},314                "name": "3",315                "run_id": "",316                "parent_ids": [],317                "tags": ["seq:step:3"],318            },319            {320                "data": {"input": "olleh", "output": "hello"},321                "event": "on_chain_end",322                "metadata": {},323                "name": "2",324                "run_id": "",325                "parent_ids": [],326                "tags": ["seq:step:2"],327            },328            {329                "data": {"chunk": "olleh"},330                "event": "on_chain_stream",331                "metadata": {},332                "name": "3",333                "run_id": "",334                "parent_ids": [],335                "tags": ["seq:step:3"],336            },337            {338                "data": {"chunk": "olleh"},339                "event": "on_chain_stream",340                "metadata": {},341                "name": "RunnableSequence",342                "run_id": "",343                "parent_ids": [],344                "tags": [],345            },346            {347                "data": {"input": "hello", "output": "olleh"},348                "event": "on_chain_end",349                "metadata": {},350                "name": "3",351                "run_id": "",352                "parent_ids": [],353                "tags": ["seq:step:3"],354            },355            {356                "data": {"output": "olleh"},357                "event": "on_chain_end",358                "metadata": {},359                "name": "RunnableSequence",360                "run_id": "",361                "parent_ids": [],362                "tags": [],363            },364        ],365    )366367368async def test_event_stream_exception() -> None:369    def step(name: str, err: str | None, val: str) -> str:370        if err:371            raise ValueError(err)372        return val + name[-1]373374    chain = (375        RunnableLambda(partial(step, "step1", None))376        | RunnableLambda(partial(step, "step2", "ERR"))377        | RunnableLambda(partial(step, "step3", None))378    )379380    with pytest.raises(ValueError, match="ERR"):381        await _collect_events(chain.astream_events("X", version="v2"))382383384async def test_event_stream_with_triple_lambda_test_filtering() -> None:385    """Test filtering based on tags / names."""386387    def reverse(s: str) -> str:388        """Reverse a string."""389        return s[::-1]390391    r = RunnableLambda(func=reverse)392393    chain = (394        r.with_config({"run_name": "1"})395        | r.with_config({"run_name": "2", "tags": ["my_tag"]})396        | r.with_config({"run_name": "3", "tags": ["my_tag"]})397    )398    events = await _collect_events(399        chain.astream_events("hello", include_names=["1"], version="v2")400    )401    _assert_events_equal_allow_superset_metadata(402        events,403        [404            {405                "data": {"input": "hello"},406                "event": "on_chain_start",407                "metadata": {},408                "name": "1",409                "run_id": "",410                "parent_ids": [],411                "tags": ["seq:step:1"],412            },413            {414                "data": {"chunk": "olleh"},415                "event": "on_chain_stream",416                "metadata": {},417                "name": "1",418                "run_id": "",419                "parent_ids": [],420                "tags": ["seq:step:1"],421            },422            {423                "data": {"output": "olleh"},424                "event": "on_chain_end",425                "metadata": {},426                "name": "1",427                "run_id": "",428                "parent_ids": [],429                "tags": ["seq:step:1"],430            },431        ],432    )433434    events = await _collect_events(435        chain.astream_events(436            "hello", include_tags=["my_tag"], exclude_names=["2"], version="v2"437        )438    )439    _assert_events_equal_allow_superset_metadata(440        events,441        [442            {443                "data": {"input": "hello"},444                "event": "on_chain_start",445                "metadata": {},446                "name": "3",447                "run_id": "",448                "parent_ids": [],449                "tags": ["my_tag", "seq:step:3"],450            },451            {452                "data": {"chunk": "olleh"},453                "event": "on_chain_stream",454                "metadata": {},455                "name": "3",456                "run_id": "",457                "parent_ids": [],458                "tags": ["my_tag", "seq:step:3"],459            },460            {461                "data": {"output": "olleh"},462                "event": "on_chain_end",463                "metadata": {},464                "name": "3",465                "run_id": "",466                "parent_ids": [],467                "tags": ["my_tag", "seq:step:3"],468            },469        ],470    )471472473async def test_event_stream_with_lambdas_from_lambda() -> None:474    as_lambdas = RunnableLambda[Any, dict[str, str]](475        lambda _: {"answer": "goodbye"}476    ).with_config({"run_name": "my_lambda"})477    events = await _collect_events(478        as_lambdas.astream_events({"question": "hello"}, version="v2")479    )480    _assert_events_equal_allow_superset_metadata(481        events,482        [483            {484                "data": {"input": {"question": "hello"}},485                "event": "on_chain_start",486                "metadata": {},487                "name": "my_lambda",488                "run_id": "",489                "parent_ids": [],490                "tags": [],491            },492            {493                "data": {"chunk": {"answer": "goodbye"}},494                "event": "on_chain_stream",495                "metadata": {},496                "name": "my_lambda",497                "run_id": "",498                "parent_ids": [],499                "tags": [],500            },501            {502                "data": {"output": {"answer": "goodbye"}},503                "event": "on_chain_end",504                "metadata": {},505                "name": "my_lambda",506                "run_id": "",507                "parent_ids": [],508                "tags": [],509            },510        ],511    )512513514async def test_astream_events_from_model() -> None:515    """Test the output of a model."""516    infinite_cycle = cycle([AIMessage(content="hello world!")])517    # When streaming GenericFakeChatModel breaks AIMessage into chunks based on spaces518    model = (519        GenericFakeChatModel(messages=infinite_cycle)520        .with_config(521            {522                "metadata": {"a": "b"},523                "tags": ["my_model"],524                "run_name": "my_model",525            }526        )527        .bind(stop="<stop_token>")528    )529    events = await _collect_events(model.astream_events("hello", version="v2"))530    _assert_events_equal_allow_superset_metadata(531        events,532        [533            {534                "data": {"input": "hello"},535                "event": "on_chat_model_start",536                "metadata": {537                    "a": "b",538                    "ls_model_type": "chat",539                    "ls_stop": "<stop_token>",540                },541                "name": "my_model",542                "run_id": "",543                "parent_ids": [],544                "tags": ["my_model"],545            },546            {547                "data": {548                    "chunk": _any_id_ai_message_chunk(549                        content="hello",550                    )551                },552                "event": "on_chat_model_stream",553                "metadata": {554                    "a": "b",555                    "ls_model_type": "chat",556                    "ls_stop": "<stop_token>",557                },558                "name": "my_model",559                "run_id": "",560                "parent_ids": [],561                "tags": ["my_model"],562            },563            {564                "data": {"chunk": _any_id_ai_message_chunk(content=" ")},565                "event": "on_chat_model_stream",566                "metadata": {567                    "a": "b",568                    "ls_model_type": "chat",569                    "ls_stop": "<stop_token>",570                },571                "name": "my_model",572                "run_id": "",573                "parent_ids": [],574                "tags": ["my_model"],575            },576            {577                "data": {578                    "chunk": _any_id_ai_message_chunk(579                        content="world!", chunk_position="last"580                    )581                },582                "event": "on_chat_model_stream",583                "metadata": {584                    "a": "b",585                    "ls_model_type": "chat",586                    "ls_stop": "<stop_token>",587                },588                "name": "my_model",589                "run_id": "",590                "parent_ids": [],591                "tags": ["my_model"],592            },593            {594                "data": {595                    "output": _any_id_ai_message_chunk(596                        content="hello world!", chunk_position="last"597                    ),598                },599                "event": "on_chat_model_end",600                "metadata": {601                    "a": "b",602                    "ls_model_type": "chat",603                    "ls_stop": "<stop_token>",604                },605                "name": "my_model",606                "run_id": "",607                "parent_ids": [],608                "tags": ["my_model"],609            },610        ],611    )612613614async def test_astream_with_model_in_chain() -> None:615    """Scenarios with model when it is not the only runnable in the chain."""616    infinite_cycle = cycle([AIMessage(content="hello world!")])617    # When streaming GenericFakeChatModel breaks AIMessage into chunks based on spaces618    model = (619        GenericFakeChatModel(messages=infinite_cycle)620        .with_config(621            {622                "metadata": {"a": "b"},623                "tags": ["my_model"],624                "run_name": "my_model",625            }626        )627        .bind(stop="<stop_token>")628    )629630    @RunnableLambda631    def i_dont_stream(value: Any, config: RunnableConfig) -> Any:632        return model.invoke(value, config if sys.version_info >= (3, 11) else None)633634    events = await _collect_events(i_dont_stream.astream_events("hello", version="v2"))635    _assert_events_equal_allow_superset_metadata(636        events,637        [638            {639                "data": {"input": "hello"},640                "event": "on_chain_start",641                "metadata": {},642                "name": "i_dont_stream",643                "run_id": "",644                "parent_ids": [],645                "tags": [],646            },647            {648                "data": {"input": {"messages": [[HumanMessage(content="hello")]]}},649                "event": "on_chat_model_start",650                "metadata": {651                    "a": "b",652                    "ls_model_type": "chat",653                    "ls_stop": "<stop_token>",654                },655                "name": "my_model",656                "run_id": "",657                "parent_ids": [],658                "tags": ["my_model"],659            },660            {661                "data": {662                    "chunk": _any_id_ai_message_chunk(663                        content="hello",664                    )665                },666                "event": "on_chat_model_stream",667                "metadata": {668                    "a": "b",669                    "ls_model_type": "chat",670                    "ls_stop": "<stop_token>",671                },672                "name": "my_model",673                "run_id": "",674                "parent_ids": [],675                "tags": ["my_model"],676            },677            {678                "data": {"chunk": _any_id_ai_message_chunk(content=" ")},679                "event": "on_chat_model_stream",680                "metadata": {681                    "a": "b",682                    "ls_model_type": "chat",683                    "ls_stop": "<stop_token>",684                },685                "name": "my_model",686                "run_id": "",687                "parent_ids": [],688                "tags": ["my_model"],689            },690            {691                "data": {692                    "chunk": _any_id_ai_message_chunk(693                        content="world!", chunk_position="last"694                    )695                },696                "event": "on_chat_model_stream",697                "metadata": {698                    "a": "b",699                    "ls_model_type": "chat",700                    "ls_stop": "<stop_token>",701                },702                "name": "my_model",703                "run_id": "",704                "parent_ids": [],705                "tags": ["my_model"],706            },707            {708                "data": {709                    "input": {"messages": [[HumanMessage(content="hello")]]},710                    "output": _any_id_ai_message(content="hello world!"),711                },712                "event": "on_chat_model_end",713                "metadata": {714                    "a": "b",715                    "ls_model_type": "chat",716                    "ls_stop": "<stop_token>",717                },718                "name": "my_model",719                "run_id": "",720                "parent_ids": [],721                "tags": ["my_model"],722            },723            {724                "data": {"chunk": _any_id_ai_message(content="hello world!")},725                "event": "on_chain_stream",726                "metadata": {},727                "name": "i_dont_stream",728                "run_id": "",729                "parent_ids": [],730                "tags": [],731            },732            {733                "data": {"output": _any_id_ai_message(content="hello world!")},734                "event": "on_chain_end",735                "metadata": {},736                "name": "i_dont_stream",737                "run_id": "",738                "parent_ids": [],739                "tags": [],740            },741        ],742    )743744    @RunnableLambda745    async def ai_dont_stream(value: Any, config: RunnableConfig) -> Any:746        return await model.ainvoke(747            value, config if sys.version_info >= (3, 11) else None748        )749750    events = await _collect_events(ai_dont_stream.astream_events("hello", version="v2"))751    _assert_events_equal_allow_superset_metadata(752        events,753        [754            {755                "data": {"input": "hello"},756                "event": "on_chain_start",757                "metadata": {},758                "name": "ai_dont_stream",759                "run_id": "",760                "parent_ids": [],761                "tags": [],762            },763            {764                "data": {"input": {"messages": [[HumanMessage(content="hello")]]}},765                "event": "on_chat_model_start",766                "metadata": {767                    "a": "b",768                    "ls_model_type": "chat",769                    "ls_stop": "<stop_token>",770                },771                "name": "my_model",772                "run_id": "",773                "parent_ids": [],774                "tags": ["my_model"],775            },776            {777                "data": {778                    "chunk": _any_id_ai_message_chunk(779                        content="hello",780                    )781                },782                "event": "on_chat_model_stream",783                "metadata": {784                    "a": "b",785                    "ls_model_type": "chat",786                    "ls_stop": "<stop_token>",787                },788                "name": "my_model",789                "run_id": "",790                "parent_ids": [],791                "tags": ["my_model"],792            },793            {794                "data": {"chunk": _any_id_ai_message_chunk(content=" ")},795                "event": "on_chat_model_stream",796                "metadata": {797                    "a": "b",798                    "ls_model_type": "chat",799                    "ls_stop": "<stop_token>",800                },801                "name": "my_model",802                "run_id": "",803                "parent_ids": [],804                "tags": ["my_model"],805            },806            {807                "data": {808                    "chunk": _any_id_ai_message_chunk(809                        content="world!", chunk_position="last"810                    )811                },812                "event": "on_chat_model_stream",813                "metadata": {814                    "a": "b",815                    "ls_model_type": "chat",816                    "ls_stop": "<stop_token>",817                },818                "name": "my_model",819                "run_id": "",820                "parent_ids": [],821                "tags": ["my_model"],822            },823            {824                "data": {825                    "input": {"messages": [[HumanMessage(content="hello")]]},826                    "output": _any_id_ai_message(content="hello world!"),827                },828                "event": "on_chat_model_end",829                "metadata": {830                    "a": "b",831                    "ls_model_type": "chat",832                    "ls_stop": "<stop_token>",833                },834                "name": "my_model",835                "run_id": "",836                "parent_ids": [],837                "tags": ["my_model"],838            },839            {840                "data": {"chunk": _any_id_ai_message(content="hello world!")},841                "event": "on_chain_stream",842                "metadata": {},843                "name": "ai_dont_stream",844                "run_id": "",845                "parent_ids": [],846                "tags": [],847            },848            {849                "data": {"output": _any_id_ai_message(content="hello world!")},850                "event": "on_chain_end",851                "metadata": {},852                "name": "ai_dont_stream",853                "run_id": "",854                "parent_ids": [],855                "tags": [],856            },857        ],858    )859860861async def test_event_stream_with_simple_chain() -> None:862    """Test as event stream."""863    template = ChatPromptTemplate.from_messages(864        [865            ("system", "You are Cat Agent 007"),866            ("human", "{question}"),867        ]868    ).with_config({"run_name": "my_template", "tags": ["my_template"]})869870    infinite_cycle = cycle(871        [872            AIMessage(content="hello world!", id="ai1"),873            AIMessage(content="goodbye world!", id="ai2"),874        ]875    )876    # When streaming GenericFakeChatModel breaks AIMessage into chunks based on spaces877    model = (878        GenericFakeChatModel(messages=infinite_cycle)879        .with_config(880            {881                "metadata": {"a": "b"},882                "tags": ["my_model"],883                "run_name": "my_model",884            }885        )886        .bind(stop="<stop_token>")887    )888889    chain = (template | model).with_config(890        {891            "metadata": {"foo": "bar"},892            "tags": ["my_chain"],893            "run_name": "my_chain",894        }895    )896897    events = await _collect_events(898        chain.astream_events({"question": "hello"}, version="v2")899    )900    _assert_events_equal_allow_superset_metadata(901        events,902        [903            {904                "data": {"input": {"question": "hello"}},905                "event": "on_chain_start",906                "metadata": {"foo": "bar"},907                "name": "my_chain",908                "run_id": "",909                "parent_ids": [],910                "tags": ["my_chain"],911            },912            {913                "data": {"input": {"question": "hello"}},914                "event": "on_prompt_start",915                "metadata": {"foo": "bar"},916                "name": "my_template",917                "run_id": "",918                "parent_ids": [],919                "tags": ["my_chain", "my_template", "seq:step:1"],920            },921            {922                "data": {923                    "input": {"question": "hello"},924                    "output": ChatPromptValue(925                        messages=[926                            SystemMessage(content="You are Cat Agent 007"),927                            HumanMessage(content="hello"),928                        ]929                    ),930                },931                "event": "on_prompt_end",932                "metadata": {"foo": "bar"},933                "name": "my_template",934                "run_id": "",935                "parent_ids": [],936                "tags": ["my_chain", "my_template", "seq:step:1"],937            },938            {939                "data": {940                    "input": {941                        "messages": [942                            [943                                SystemMessage(content="You are Cat Agent 007"),944                                HumanMessage(content="hello"),945                            ]946                        ]947                    }948                },949                "event": "on_chat_model_start",950                "metadata": {951                    "a": "b",952                    "foo": "bar",953                    "ls_model_type": "chat",954                    "ls_stop": "<stop_token>",955                },956                "name": "my_model",957                "run_id": "",958                "parent_ids": [],959                "tags": ["my_chain", "my_model", "seq:step:2"],960            },961            {962                "data": {963                    "chunk": AIMessageChunk(964                        content="hello",965                        id="ai1",966                    )967                },968                "event": "on_chat_model_stream",969                "metadata": {970                    "a": "b",971                    "foo": "bar",972                    "ls_model_type": "chat",973                    "ls_stop": "<stop_token>",974                },975                "name": "my_model",976                "run_id": "",977                "parent_ids": [],978                "tags": ["my_chain", "my_model", "seq:step:2"],979            },980            {981                "data": {982                    "chunk": AIMessageChunk(983                        content="hello",984                        id="ai1",985                    )986                },987                "event": "on_chain_stream",988                "metadata": {"foo": "bar"},989                "name": "my_chain",990                "run_id": "",991                "parent_ids": [],992                "tags": ["my_chain"],993            },994            {995                "data": {"chunk": AIMessageChunk(content=" ", id="ai1")},996                "event": "on_chat_model_stream",997                "metadata": {998                    "a": "b",999                    "foo": "bar",1000                    "ls_model_type": "chat",1001                    "ls_stop": "<stop_token>",1002                },1003                "name": "my_model",1004                "run_id": "",1005                "parent_ids": [],1006                "tags": ["my_chain", "my_model", "seq:step:2"],1007            },1008            {1009                "data": {"chunk": AIMessageChunk(content=" ", id="ai1")},1010                "event": "on_chain_stream",1011                "metadata": {"foo": "bar"},1012                "name": "my_chain",1013                "run_id": "",1014                "parent_ids": [],1015                "tags": ["my_chain"],1016            },1017            {1018                "data": {1019                    "chunk": AIMessageChunk(1020                        content="world!", id="ai1", chunk_position="last"1021                    )1022                },1023                "event": "on_chat_model_stream",1024                "metadata": {1025                    "a": "b",1026                    "foo": "bar",1027                    "ls_model_type": "chat",1028                    "ls_stop": "<stop_token>",1029                },1030                "name": "my_model",1031                "run_id": "",1032                "parent_ids": [],1033                "tags": ["my_chain", "my_model", "seq:step:2"],1034            },1035            {1036                "data": {1037                    "chunk": AIMessageChunk(1038                        content="world!", id="ai1", chunk_position="last"1039                    )1040                },1041                "event": "on_chain_stream",1042                "metadata": {"foo": "bar"},1043                "name": "my_chain",1044                "run_id": "",1045                "parent_ids": [],1046                "tags": ["my_chain"],1047            },1048            {1049                "data": {1050                    "input": {1051                        "messages": [1052                            [1053                                SystemMessage(content="You are Cat Agent 007"),1054                                HumanMessage(content="hello"),1055                            ]1056                        ]1057                    },1058                    "output": AIMessageChunk(1059                        content="hello world!", id="ai1", chunk_position="last"1060                    ),1061                },1062                "event": "on_chat_model_end",1063                "metadata": {1064                    "a": "b",1065                    "foo": "bar",1066                    "ls_model_type": "chat",1067                    "ls_stop": "<stop_token>",1068                },1069                "name": "my_model",1070                "run_id": "",1071                "parent_ids": [],1072                "tags": ["my_chain", "my_model", "seq:step:2"],1073            },1074            {1075                "data": {1076                    "output": AIMessageChunk(1077                        content="hello world!", id="ai1", chunk_position="last"1078                    )1079                },1080                "event": "on_chain_end",1081                "metadata": {"foo": "bar"},1082                "name": "my_chain",1083                "run_id": "",1084                "parent_ids": [],1085                "tags": ["my_chain"],1086            },1087        ],1088    )108910901091async def test_event_streaming_with_tools() -> None:1092    """Test streaming events with different tool definitions."""10931094    @tool1095    def parameterless() -> str:1096        """A tool that does nothing."""1097        return "hello"10981099    @tool1100    def with_callbacks(callbacks: Callbacks) -> str:1101        """A tool that does nothing."""1102        _ = callbacks1103        return "world"11041105    @tool1106    def with_parameters(x: int, y: str) -> dict[str, Any]:1107        """A tool that does nothing."""1108        return {"x": x, "y": y}11091110    @tool1111    def with_parameters_and_callbacks(1112        x: int, y: str, callbacks: Callbacks1113    ) -> dict[str, Any]:1114        """A tool that does nothing."""1115        _ = callbacks1116        return {"x": x, "y": y}11171118    events = await _collect_events(parameterless.astream_events({}, version="v2"))1119    _assert_events_equal_allow_superset_metadata(1120        events,1121        [1122            {1123                "data": {"input": {}},1124                "event": "on_tool_start",1125                "metadata": {},1126                "name": "parameterless",1127                "run_id": "",1128                "parent_ids": [],1129                "tags": [],1130            },1131            {1132                "data": {"output": "hello"},1133                "event": "on_tool_end",1134                "metadata": {},1135                "name": "parameterless",1136                "run_id": "",1137                "parent_ids": [],1138                "tags": [],1139            },1140        ],1141    )1142    events = await _collect_events(with_callbacks.astream_events({}, version="v2"))1143    _assert_events_equal_allow_superset_metadata(1144        events,1145        [1146            {1147                "data": {"input": {}},1148                "event": "on_tool_start",1149                "metadata": {},1150                "name": "with_callbacks",1151                "run_id": "",1152                "parent_ids": [],1153                "tags": [],1154            },1155            {1156                "data": {"output": "world"},1157                "event": "on_tool_end",1158                "metadata": {},1159                "name": "with_callbacks",1160                "run_id": "",1161                "parent_ids": [],1162                "tags": [],1163            },1164        ],1165    )1166    events = await _collect_events(1167        with_parameters.astream_events({"x": 1, "y": "2"}, version="v2")1168    )1169    _assert_events_equal_allow_superset_metadata(1170        events,1171        [1172            {1173                "data": {"input": {"x": 1, "y": "2"}},1174                "event": "on_tool_start",1175                "metadata": {},1176                "name": "with_parameters",1177                "run_id": "",1178                "parent_ids": [],1179                "tags": [],1180            },1181            {1182                "data": {"output": {"x": 1, "y": "2"}},1183                "event": "on_tool_end",1184                "metadata": {},1185                "name": "with_parameters",1186                "run_id": "",1187                "parent_ids": [],1188                "tags": [],1189            },1190        ],1191    )11921193    events = await _collect_events(1194        with_parameters_and_callbacks.astream_events({"x": 1, "y": "2"}, version="v2")1195    )1196    _assert_events_equal_allow_superset_metadata(1197        events,1198        [1199            {1200                "data": {"input": {"x": 1, "y": "2"}},1201                "event": "on_tool_start",1202                "metadata": {},1203                "name": "with_parameters_and_callbacks",1204                "run_id": "",1205                "parent_ids": [],1206                "tags": [],1207            },1208            {1209                "data": {"output": {"x": 1, "y": "2"}},1210                "event": "on_tool_end",1211                "metadata": {},1212                "name": "with_parameters_and_callbacks",1213                "run_id": "",1214                "parent_ids": [],1215                "tags": [],1216            },1217        ],1218    )121912201221class HardCodedRetriever(BaseRetriever):1222    documents: list[Document]12231224    @override1225    def _get_relevant_documents(1226        self, query: str, *, run_manager: CallbackManagerForRetrieverRun1227    ) -> list[Document]:1228        return self.documents122912301231async def test_event_stream_with_retriever() -> None:1232    """Test the event stream with a retriever."""1233    retriever = HardCodedRetriever(1234        documents=[1235            Document(1236                page_content="hello world!",1237                metadata={"foo": "bar"},1238            ),1239            Document(1240                page_content="goodbye world!",1241                metadata={"food": "spare"},1242            ),1243        ]1244    )1245    events = await _collect_events(1246        retriever.astream_events({"query": "hello"}, version="v2")1247    )1248    _assert_events_equal_allow_superset_metadata(1249        events,1250        [1251            {1252                "data": {1253                    "input": {"query": "hello"},1254                },1255                "event": "on_retriever_start",1256                "metadata": {},1257                "name": "HardCodedRetriever",1258                "run_id": "",1259                "parent_ids": [],1260                "tags": [],1261            },1262            {1263                "data": {1264                    "output": [1265                        Document(page_content="hello world!", metadata={"foo": "bar"}),1266                        Document(1267                            page_content="goodbye world!", metadata={"food": "spare"}1268                        ),1269                    ]1270                },1271                "event": "on_retriever_end",1272                "metadata": {},1273                "name": "HardCodedRetriever",1274                "run_id": "",1275                "parent_ids": [],1276                "tags": [],1277            },1278        ],1279    )128012811282async def test_event_stream_with_retriever_and_formatter() -> None:1283    """Test the event stream with a retriever."""1284    retriever = HardCodedRetriever(1285        documents=[1286            Document(1287                page_content="hello world!",1288                metadata={"foo": "bar"},1289            ),1290            Document(1291                page_content="goodbye world!",1292                metadata={"food": "spare"},1293            ),1294        ]1295    )12961297    def format_docs(docs: list[Document]) -> str:1298        """Format the docs."""1299        return ", ".join([doc.page_content for doc in docs])13001301    chain = retriever | format_docs1302    events = await _collect_events(chain.astream_events("hello", version="v2"))1303    _assert_events_equal_allow_superset_metadata(1304        events,1305        [1306            {1307                "data": {"input": "hello"},1308                "event": "on_chain_start",1309                "metadata": {},1310                "name": "RunnableSequence",1311                "run_id": "",1312                "parent_ids": [],1313                "tags": [],1314            },1315            {1316                "data": {"input": {"query": "hello"}},1317                "event": "on_retriever_start",1318                "metadata": {},1319                "name": "HardCodedRetriever",1320                "run_id": "",1321                "parent_ids": [],1322                "tags": ["seq:step:1"],1323            },1324            {1325                "data": {1326                    "input": {"query": "hello"},1327                    "output": [1328                        Document(page_content="hello world!", metadata={"foo": "bar"}),1329                        Document(1330                            page_content="goodbye world!", metadata={"food": "spare"}1331                        ),1332                    ],1333                },1334                "event": "on_retriever_end",1335                "metadata": {},1336                "name": "HardCodedRetriever",1337                "run_id": "",1338                "parent_ids": [],1339                "tags": ["seq:step:1"],1340            },1341            {1342                "data": {},1343                "event": "on_chain_start",1344                "metadata": {},1345                "name": "format_docs",1346                "run_id": "",1347                "parent_ids": [],1348                "tags": ["seq:step:2"],1349            },1350            {1351                "data": {"chunk": "hello world!, goodbye world!"},1352                "event": "on_chain_stream",1353                "metadata": {},1354                "name": "format_docs",1355                "run_id": "",1356                "parent_ids": [],1357                "tags": ["seq:step:2"],1358            },1359            {1360                "data": {"chunk": "hello world!, goodbye world!"},1361                "event": "on_chain_stream",1362                "metadata": {},1363                "name": "RunnableSequence",1364                "run_id": "",1365                "parent_ids": [],1366                "tags": [],1367            },1368            {1369                "data": {1370                    "input": [1371                        Document(page_content="hello world!", metadata={"foo": "bar"}),1372                        Document(1373                            page_content="goodbye world!", metadata={"food": "spare"}1374                        ),1375                    ],1376                    "output": "hello world!, goodbye world!",1377                },1378                "event": "on_chain_end",1379                "metadata": {},1380                "name": "format_docs",1381                "run_id": "",1382                "parent_ids": [],1383                "tags": ["seq:step:2"],1384            },1385            {1386                "data": {"output": "hello world!, goodbye world!"},1387                "event": "on_chain_end",1388                "metadata": {},1389                "name": "RunnableSequence",1390                "run_id": "",1391                "parent_ids": [],1392                "tags": [],1393            },1394        ],1395    )139613971398async def test_event_stream_on_chain_with_tool() -> None:1399    """Test the event stream with a tool."""14001401    @tool1402    def concat(a: str, b: str) -> str:1403        """A tool that does nothing."""1404        return a + b14051406    def reverse(s: str) -> str:1407        """Reverse a string."""1408        return s[::-1]14091410    chain = concat | reverse14111412    events = await _collect_events(1413        chain.astream_events({"a": "hello", "b": "world"}, version="v2")1414    )1415    _assert_events_equal_allow_superset_metadata(1416        events,1417        [1418            {1419                "data": {"input": {"a": "hello", "b": "world"}},1420                "event": "on_chain_start",1421                "metadata": {},1422                "name": "RunnableSequence",1423                "run_id": "",1424                "parent_ids": [],1425                "tags": [],1426            },1427            {1428                "data": {"input": {"a": "hello", "b": "world"}},1429                "event": "on_tool_start",1430                "metadata": {},1431                "name": "concat",1432                "run_id": "",1433                "parent_ids": [],1434                "tags": ["seq:step:1"],1435            },1436            {1437                "data": {"input": {"a": "hello", "b": "world"}, "output": "helloworld"},1438                "event": "on_tool_end",1439                "metadata": {},1440                "name": "concat",1441                "run_id": "",1442                "parent_ids": [],1443                "tags": ["seq:step:1"],1444            },1445            {1446                "data": {},1447                "event": "on_chain_start",1448                "metadata": {},1449                "name": "reverse",1450                "run_id": "",1451                "parent_ids": [],1452                "tags": ["seq:step:2"],1453            },1454            {1455                "data": {"chunk": "dlrowolleh"},1456                "event": "on_chain_stream",1457                "metadata": {},1458                "name": "reverse",1459                "run_id": "",1460                "parent_ids": [],1461                "tags": ["seq:step:2"],1462            },1463            {1464                "data": {"chunk": "dlrowolleh"},1465                "event": "on_chain_stream",1466                "metadata": {},1467                "name": "RunnableSequence",1468                "run_id": "",1469                "parent_ids": [],1470                "tags": [],1471            },1472            {1473                "data": {"input": "helloworld", "output": "dlrowolleh"},1474                "event": "on_chain_end",1475                "metadata": {},1476                "name": "reverse",1477                "run_id": "",1478                "parent_ids": [],1479                "tags": ["seq:step:2"],1480            },1481            {1482                "data": {"output": "dlrowolleh"},1483                "event": "on_chain_end",1484                "metadata": {},1485                "name": "RunnableSequence",1486                "run_id": "",1487                "parent_ids": [],1488                "tags": [],1489            },1490        ],1491    )149214931494@pytest.mark.xfail(reason="Fix order of callback invocations in RunnableSequence")1495async def test_chain_ordering() -> None:1496    """Test the event stream with a tool."""14971498    def foo(a: str) -> str:1499        return a15001501    def bar(a: str) -> str:1502        return a15031504    chain = RunnableLambda(foo) | RunnableLambda(bar)1505    iterable = chain.astream_events("q", version="v2")15061507    events = []15081509    try:1510        for _ in range(10):1511            next_chunk = await anext(iterable)1512            events.append(next_chunk)1513    except Exception:1514        pass15151516    events = _with_nulled_run_id(events)1517    for event in events:1518        event["tags"] = sorted(event["tags"])15191520    _assert_events_equal_allow_superset_metadata(1521        events,1522        [1523            {1524                "data": {"input": "q"},1525                "event": "on_chain_start",1526                "metadata": {},1527                "name": "RunnableSequence",1528                "run_id": "",1529                "parent_ids": [],1530                "tags": [],1531            },1532            {1533                "data": {},1534                "event": "on_chain_start",1535                "metadata": {},1536                "name": "foo",1537                "run_id": "",1538                "parent_ids": [],1539                "tags": ["seq:step:1"],1540            },1541            {1542                "data": {"chunk": "q"},1543                "event": "on_chain_stream",1544                "metadata": {},1545                "name": "foo",1546                "run_id": "",1547                "parent_ids": [],1548                "tags": ["seq:step:1"],1549            },1550            {1551                "data": {"input": "q", "output": "q"},1552                "event": "on_chain_end",1553                "metadata": {},1554                "name": "foo",1555                "run_id": "",1556                "parent_ids": [],1557                "tags": ["seq:step:1"],1558            },1559            {1560                "data": {},1561                "event": "on_chain_start",1562                "metadata": {},1563                "name": "bar",1564                "run_id": "",1565                "parent_ids": [],1566                "tags": ["seq:step:2"],1567            },1568            {1569                "data": {"chunk": "q"},1570                "event": "on_chain_stream",1571                "metadata": {},1572                "name": "bar",1573                "run_id": "",1574                "parent_ids": [],1575                "tags": ["seq:step:2"],1576            },1577            {1578                "data": {"chunk": "q"},1579                "event": "on_chain_stream",1580                "metadata": {},1581                "name": "RunnableSequence",1582                "run_id": "",1583                "parent_ids": [],1584                "tags": [],1585            },1586            {1587                "data": {"input": "q", "output": "q"},1588                "event": "on_chain_end",1589                "metadata": {},1590                "name": "bar",1591                "run_id": "",1592                "parent_ids": [],1593                "tags": ["seq:step:2"],1594            },1595            {1596                "data": {"output": "q"},1597                "event": "on_chain_end",1598                "metadata": {},1599                "name": "RunnableSequence",1600                "run_id": "",1601                "parent_ids": [],1602                "tags": [],1603            },1604        ],1605    )160616071608async def test_event_stream_with_retry() -> None:1609    """Test the event stream with a tool."""16101611    def success(_: str) -> str:1612        return "success"16131614    def fail(_: str) -> None:1615        """Simple func."""1616        msg = "fail"1617        raise ValueError(msg)16181619    chain = RunnableLambda(success) | RunnableLambda(fail).with_retry(1620        stop_after_attempt=1,1621    )1622    iterable = chain.astream_events("q", version="v2")16231624    events = []16251626    try:1627        for _ in range(10):1628            next_chunk = await anext(iterable)1629            events.append(next_chunk)1630    except Exception:1631        pass16321633    events = _with_nulled_run_id(events)1634    for event in events:1635        event["tags"] = sorted(event["tags"])16361637    _assert_events_equal_allow_superset_metadata(1638        events,1639        [1640            {1641                "data": {"input": "q"},1642                "event": "on_chain_start",1643                "metadata": {},1644                "name": "RunnableSequence",1645                "run_id": "",1646                "parent_ids": [],1647                "tags": [],1648            },1649            {1650                "data": {},1651                "event": "on_chain_start",1652                "metadata": {},1653                "name": "success",1654                "run_id": "",1655                "parent_ids": [],1656                "tags": ["seq:step:1"],1657            },1658            {1659                "data": {"chunk": "success"},1660                "event": "on_chain_stream",1661                "metadata": {},1662                "name": "success",1663                "run_id": "",1664                "parent_ids": [],1665                "tags": ["seq:step:1"],1666            },1667            {1668                "data": {},1669                "event": "on_chain_start",1670                "metadata": {},1671                "name": "fail",1672                "run_id": "",1673                "parent_ids": [],1674                "tags": ["seq:step:2"],1675            },1676            {1677                "data": {"input": "q", "output": "success"},1678                "event": "on_chain_end",1679                "metadata": {},1680                "name": "success",1681                "run_id": "",1682                "parent_ids": [],1683                "tags": ["seq:step:1"],1684            },1685        ],1686    )168716881689async def test_with_llm() -> None:1690    """Test with regular llm."""1691    prompt = ChatPromptTemplate.from_messages(1692        [1693            ("system", "You are Cat Agent 007"),1694            ("human", "{question}"),1695        ]1696    ).with_config({"run_name": "my_template", "tags": ["my_template"]})1697    llm = FakeStreamingListLLM(responses=["abc"])16981699    chain = prompt | llm1700    events = await _collect_events(1701        chain.astream_events({"question": "hello"}, version="v2")1702    )1703    _assert_events_equal_allow_superset_metadata(1704        events,1705        [1706            {1707                "data": {"input": {"question": "hello"}},1708                "event": "on_chain_start",1709                "metadata": {},1710                "name": "RunnableSequence",1711                "run_id": "",1712                "parent_ids": [],1713                "tags": [],1714            },1715            {1716                "data": {"input": {"question": "hello"}},1717                "event": "on_prompt_start",1718                "metadata": {},1719                "name": "my_template",1720                "run_id": "",1721                "parent_ids": [],1722                "tags": ["my_template", "seq:step:1"],1723            },1724            {1725                "data": {1726                    "input": {"question": "hello"},1727                    "output": ChatPromptValue(1728                        messages=[1729                            SystemMessage(content="You are Cat Agent 007"),1730                            HumanMessage(content="hello"),1731                        ]1732                    ),1733                },1734                "event": "on_prompt_end",1735                "metadata": {},1736                "name": "my_template",1737                "run_id": "",1738                "parent_ids": [],1739                "tags": ["my_template", "seq:step:1"],1740            },1741            {1742                "data": {1743                    "input": {1744                        "prompts": ["System: You are Cat Agent 007\nHuman: hello"]1745                    }1746                },1747                "event": "on_llm_start",1748                "metadata": {},1749                "name": "FakeStreamingListLLM",1750                "run_id": "",1751                "parent_ids": [],1752                "tags": ["seq:step:2"],1753            },1754            {1755                "data": {1756                    "input": {1757                        "prompts": ["System: You are Cat Agent 007\nHuman: hello"]1758                    },1759                    "output": {1760                        "generations": [1761                            [1762                                {1763                                    "generation_info": None,1764                                    "text": "abc",1765                                    "type": "Generation",1766                                }1767                            ]1768                        ],1769                        "llm_output": None,1770                    },1771                },1772                "event": "on_llm_end",1773                "metadata": {},1774                "name": "FakeStreamingListLLM",1775                "run_id": "",1776                "parent_ids": [],1777                "tags": ["seq:step:2"],1778            },1779            {1780                "data": {"chunk": "a"},1781                "event": "on_chain_stream",1782                "metadata": {},1783                "name": "RunnableSequence",1784                "run_id": "",1785                "parent_ids": [],1786                "tags": [],1787            },1788            {1789                "data": {"chunk": "b"},1790                "event": "on_chain_stream",1791                "metadata": {},1792                "name": "RunnableSequence",1793                "run_id": "",1794                "parent_ids": [],1795                "tags": [],1796            },1797            {1798                "data": {"chunk": "c"},1799                "event": "on_chain_stream",1800                "metadata": {},1801                "name": "RunnableSequence",1802                "run_id": "",1803                "parent_ids": [],1804                "tags": [],1805            },1806            {1807                "data": {"output": "abc"},1808                "event": "on_chain_end",1809                "metadata": {},1810                "name": "RunnableSequence",1811                "run_id": "",1812                "parent_ids": [],1813                "tags": [],1814            },1815        ],1816    )181718181819async def test_runnable_each() -> None:1820    """Test runnable each astream_events."""18211822    async def add_one(x: int) -> int:1823        return x + 118241825    add_one_map = RunnableLambda(add_one).map()1826    assert await add_one_map.ainvoke([1, 2, 3]) == [2, 3, 4]18271828    with pytest.raises(NotImplementedError):1829        _ = [_ async for _ in add_one_map.astream_events([1, 2, 3], version="v2")]183018311832async def test_events_astream_config() -> None:1833    """Test that astream events support accepting config."""1834    infinite_cycle = cycle([AIMessage(content="hello world!", id="ai1")])1835    good_world_on_repeat = cycle([AIMessage(content="Goodbye world", id="ai2")])1836    model = GenericFakeChatModel(messages=infinite_cycle).configurable_fields(1837        messages=ConfigurableField(1838            id="messages",1839            name="Messages",1840            description="Messages return by the LLM",1841        )1842    )18431844    model_02 = model.with_config({"configurable": {"messages": good_world_on_repeat}})1845    assert model_02.invoke("hello") == AIMessage(content="Goodbye world", id="ai2")18461847    events = await _collect_events(model_02.astream_events("hello", version="v2"))1848    _assert_events_equal_allow_superset_metadata(1849        events,1850        [1851            {1852                "data": {"input": "hello"},1853                "event": "on_chat_model_start",1854                "metadata": {"ls_model_type": "chat"},1855                "name": "GenericFakeChatModel",1856                "run_id": "",1857                "parent_ids": [],1858                "tags": [],1859            },1860            {1861                "data": {1862                    "chunk": AIMessageChunk(1863                        content="Goodbye",1864                        id="ai2",1865                    )1866                },1867                "event": "on_chat_model_stream",1868                "metadata": {"ls_model_type": "chat"},1869                "name": "GenericFakeChatModel",1870                "run_id": "",1871                "parent_ids": [],1872                "tags": [],1873            },1874            {1875                "data": {"chunk": AIMessageChunk(content=" ", id="ai2")},1876                "event": "on_chat_model_stream",1877                "metadata": {"ls_model_type": "chat"},1878                "name": "GenericFakeChatModel",1879                "run_id": "",1880                "parent_ids": [],1881                "tags": [],1882            },1883            {1884                "data": {1885                    "chunk": AIMessageChunk(1886                        content="world", id="ai2", chunk_position="last"1887                    )1888                },1889                "event": "on_chat_model_stream",1890                "metadata": {"ls_model_type": "chat"},1891                "name": "GenericFakeChatModel",1892                "run_id": "",1893                "parent_ids": [],1894                "tags": [],1895            },1896            {1897                "data": {1898                    "output": AIMessageChunk(1899                        content="Goodbye world", id="ai2", chunk_position="last"1900                    ),1901                },1902                "event": "on_chat_model_end",1903                "metadata": {"ls_model_type": "chat"},1904                "name": "GenericFakeChatModel",1905                "run_id": "",1906                "parent_ids": [],1907                "tags": [],1908            },1909        ],1910    )191119121913async def test_runnable_with_message_history() -> None:1914    class InMemoryHistory(BaseChatMessageHistory, BaseModel):1915        """In memory implementation of chat message history."""19161917        # Attention: for the tests use an Any type to work-around a pydantic issue1918        # where it re-instantiates a list, so mutating the list doesn't end up mutating1919        # the content in the store!19201921        # Using Any type here rather than list[BaseMessage] due to pydantic issue!1922        messages: Any19231924        def add_message(self, message: BaseMessage) -> None:1925            """Add a self-created message to the store."""1926            self.messages.append(message)19271928        def clear(self) -> None:1929            self.messages = []19301931    # Here we use a global variable to store the chat message history.1932    # This will make it easier to inspect it to see the underlying results.1933    store: dict[str, list[BaseMessage]] = {}19341935    def get_by_session_id(session_id: str) -> BaseChatMessageHistory:1936        """Get a chat message history."""1937        if session_id not in store:1938            store[session_id] = []1939        return InMemoryHistory(messages=store[session_id])19401941    infinite_cycle = cycle(1942        [1943            AIMessage(content="hello", id="ai3"),1944            AIMessage(content="world", id="ai4"),1945        ]1946    )19471948    prompt = ChatPromptTemplate.from_messages(1949        [1950            ("system", "You are a cat"),1951            MessagesPlaceholder(variable_name="history"),1952            ("human", "{question}"),1953        ]1954    )1955    model = GenericFakeChatModel(messages=infinite_cycle)19561957    chain = prompt | model1958    with_message_history = RunnableWithMessageHistory(1959        chain,1960        get_session_history=get_by_session_id,1961        input_messages_key="question",1962        history_messages_key="history",1963    )19641965    # patch with_message_history._get_output_messages to listen for errors1966    # so we can raise them in this main thread1967    raised_errors = []19681969    def collect_errors(fn: Callable[..., Any]) -> Callable[..., Any]:1970        nonlocal raised_errors19711972        def _get_output_messages(*args: Any, **kwargs: Any) -> Any:1973            try:1974                return fn(*args, **kwargs)1975            except Exception as e:1976                raised_errors.append(e)1977                raise19781979        return _get_output_messages19801981    old_ref = with_message_history._get_output_messages1982    with_message_history.__dict__["_get_output_messages"] = collect_errors(old_ref)1983    await with_message_history.with_config(1984        {"configurable": {"session_id": "session-123"}}1985    ).ainvoke({"question": "hello"})19861987    assert store == {1988        "session-123": [1989            HumanMessage(content="hello"),1990            AIMessage(content="hello", id="ai3"),1991        ]1992    }19931994    await asyncio.to_thread(1995        with_message_history.with_config(1996            {"configurable": {"session_id": "session-123"}}1997        ).invoke,1998        {"question": "meow"},1999    )2000    assert store == {