1import asyncio2import re3import sys4import time5import uuid6import warnings7from collections.abc import (8    AsyncIterator,9    Awaitable,10    Callable,11    Iterator,12    Sequence,13)14from functools import partial15from operator import itemgetter16from typing import Any, cast17from uuid import UUID1819import pytest20from freezegun import freeze_time21from packaging import version22from pydantic import BaseModel, Field, ValidationError23from pydantic.v1 import BaseModel as BaseModelV124from pydantic.v1 import Field as FieldV125from pydantic.v1 import ValidationError as ValidationErrorV126from pytest_mock import MockerFixture27from syrupy.assertion import SnapshotAssertion28from typing_extensions import TypedDict, override2930from langchain_core.callbacks import BaseCallbackHandler31from langchain_core.callbacks.manager import (32    AsyncCallbackManagerForRetrieverRun,33    CallbackManagerForRetrieverRun,34    atrace_as_chain_group,35    trace_as_chain_group,36)37from langchain_core.documents import Document38from langchain_core.language_models import (39    FakeListChatModel,40    FakeListLLM,41    FakeStreamingListLLM,42)43from langchain_core.language_models.fake_chat_models import GenericFakeChatModel44from langchain_core.load import dumpd, dumps45from langchain_core.load.load import loads46from langchain_core.messages import AIMessageChunk, HumanMessage, SystemMessage47from langchain_core.messages.base import BaseMessage48from langchain_core.output_parsers import (49    BaseOutputParser,50    CommaSeparatedListOutputParser,51    StrOutputParser,52)53from langchain_core.outputs.chat_generation import ChatGeneration54from langchain_core.outputs.llm_result import LLMResult55from langchain_core.prompt_values import ChatPromptValue, StringPromptValue56from langchain_core.prompts import (57    ChatPromptTemplate,58    HumanMessagePromptTemplate,59    MessagesPlaceholder,60    PromptTemplate,61    SystemMessagePromptTemplate,62)63from langchain_core.retrievers import BaseRetriever64from langchain_core.runnables import (65    AddableDict,66    ConfigurableField,67    ConfigurableFieldMultiOption,68    ConfigurableFieldSingleOption,69    RouterRunnable,70    Runnable,71    RunnableAssign,72    RunnableBinding,73    RunnableBranch,74    RunnableConfig,75    RunnableGenerator,76    RunnableLambda,77    RunnableParallel,78    RunnablePassthrough,79    RunnablePick,80    RunnableSequence,81    add,82    chain,83)84from langchain_core.runnables.base import RunnableMap, RunnableSerializable85from langchain_core.runnables.utils import Input, Output86from langchain_core.tools import BaseTool, tool87from langchain_core.tracers import (88    BaseTracer,89    ConsoleCallbackHandler,90    Run,91    RunLog,92    RunLogPatch,93)94from langchain_core.tracers._compat import pydantic_copy95from langchain_core.tracers.context import collect_runs96from langchain_core.utils.pydantic import (97    PYDANTIC_VERSION,98    TypeBaseModel,99    model_validate,100)101from langchain_core.version import VERSION102from tests.unit_tests.pydantic_utils import (103    _normalize_schema,104    _schema,105    skip_if_no_pydantic_v1,106)107from tests.unit_tests.stubs import AnyStr, _any_id_ai_message, _any_id_ai_message_chunk108109# Several tests assert the legacy `RunLog` / `RunLogPatch` output produced by110# `astream_log`, which cannot be replaced by `astream` without losing coverage.111pytestmark = pytest.mark.filterwarnings(112    "ignore:astream_log is deprecated. Use astream instead.:"113    "langchain_core._api.deprecation.LangChainDeprecationWarning"114)115116PYDANTIC_VERSION_AT_LEAST_29 = version.parse("2.9") <= PYDANTIC_VERSION117PYDANTIC_VERSION_AT_LEAST_210 = version.parse("2.10") <= PYDANTIC_VERSION118119120def _normalize_lc_version(value: str) -> str:121    return value.replace(122        f"'langchain-core': '{VERSION}'",123        "'langchain-core': '<version>'",124    )125126127class FakeTracer(BaseTracer):128    """Fake tracer that records LangChain execution.129130    It replaces run IDs with deterministic UUIDs for snapshotting.131    """132133    def __init__(self) -> None:134        """Initialize the tracer."""135        super().__init__()136        self.runs: list[Run] = []137        self.uuids_map: dict[UUID, UUID] = {}138        self.uuids_generator = (139            UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000)140        )141142    def _replace_uuid(self, uuid: UUID) -> UUID:143        if uuid not in self.uuids_map:144            self.uuids_map[uuid] = next(self.uuids_generator)145        return self.uuids_map[uuid]146147    def _replace_message_id(self, maybe_message: Any) -> Any:148        if isinstance(maybe_message, BaseMessage):149            maybe_message.id = str(next(self.uuids_generator))150        if isinstance(maybe_message, ChatGeneration):151            maybe_message.message.id = str(next(self.uuids_generator))152        if isinstance(maybe_message, LLMResult):153            for i, gen_list in enumerate(maybe_message.generations):154                for j, gen in enumerate(gen_list):155                    maybe_message.generations[i][j] = self._replace_message_id(gen)156        if isinstance(maybe_message, dict):157            for k, v in maybe_message.items():158                maybe_message[k] = self._replace_message_id(v)159        if isinstance(maybe_message, list):160            for i, v in enumerate(maybe_message):161                maybe_message[i] = self._replace_message_id(v)162163        return maybe_message164165    def _copy_run(self, run: Run) -> Run:166        if run.dotted_order:167            levels = run.dotted_order.split(".")168            processed_levels = []169            for level in levels:170                timestamp, run_id = level.split("Z")171                new_run_id = self._replace_uuid(UUID(run_id))172                processed_level = f"{timestamp}Z{new_run_id}"173                processed_levels.append(processed_level)174            new_dotted_order = ".".join(processed_levels)175        else:176            new_dotted_order = None177        update_dict = {178            "id": self._replace_uuid(run.id),179            "parent_run_id": (180                self.uuids_map[run.parent_run_id] if run.parent_run_id else None181            ),182            "child_runs": [self._copy_run(child) for child in run.child_runs],183            "trace_id": self._replace_uuid(run.trace_id) if run.trace_id else None,184            "dotted_order": new_dotted_order,185            "inputs": self._replace_message_id(run.inputs),186            "outputs": self._replace_message_id(run.outputs),187        }188        return pydantic_copy(run, update=update_dict)189190    def _persist_run(self, run: Run) -> None:191        """Persist a run."""192        self.runs.append(self._copy_run(run))193194    def flattened_runs(self) -> list[Run]:195        q = [*self.runs]196        result = []197        while q:198            parent = q.pop()199            result.append(parent)200            if parent.child_runs:201                q.extend(parent.child_runs)202        return result203204    @property205    def run_ids(self) -> list[uuid.UUID | None]:206        runs = self.flattened_runs()207        uuids_map = {v: k for k, v in self.uuids_map.items()}208        return [uuids_map.get(r.id) for r in runs]209210211class FakeRunnable(Runnable[str, int]):212    @override213    def invoke(214        self,215        input: str,216        config: RunnableConfig | None = None,217        **kwargs: Any,218    ) -> int:219        return len(input)220221222class FakeRunnableSerializable(RunnableSerializable[str, int]):223    hello: str = ""224225    @override226    def invoke(227        self,228        input: str,229        config: RunnableConfig | None = None,230        **kwargs: Any,231    ) -> int:232        return len(input)233234235class FakeRetriever(BaseRetriever):236    @override237    def _get_relevant_documents(238        self, query: str, *, run_manager: CallbackManagerForRetrieverRun239    ) -> list[Document]:240        return [Document(page_content="foo"), Document(page_content="bar")]241242    @override243    async def _aget_relevant_documents(244        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun245    ) -> list[Document]:246        return [Document(page_content="foo"), Document(page_content="bar")]247248249@pytest.mark.skipif(250    PYDANTIC_VERSION_AT_LEAST_210,251    reason=(252        "Only test with most recent version of pydantic. "253        "Pydantic introduced small fixes to generated JSONSchema on minor versions."254    ),255)256def test_schemas(snapshot: SnapshotAssertion) -> None:257    fake = FakeRunnable()  # str -> int258259    assert fake.get_input_jsonschema() == {260        "title": "FakeRunnableInput",261        "type": "string",262    }263    assert fake.get_output_jsonschema() == {264        "title": "FakeRunnableOutput",265        "type": "integer",266    }267    assert fake.get_config_jsonschema(include=["tags", "metadata", "run_name"]) == {268        "properties": {269            "metadata": {270                "default": None,271                "title": "Metadata",272                "type": "object",273            },274            "run_name": {"default": None, "title": "Run Name", "type": "string"},275            "tags": {276                "default": None,277                "items": {"type": "string"},278                "title": "Tags",279                "type": "array",280            },281        },282        "title": "FakeRunnableConfig",283        "type": "object",284    }285286    fake_bound = FakeRunnable().bind(a="b")  # str -> int287288    assert fake_bound.get_input_jsonschema() == {289        "title": "FakeRunnableInput",290        "type": "string",291    }292    assert fake_bound.get_output_jsonschema() == {293        "title": "FakeRunnableOutput",294        "type": "integer",295    }296297    fake_w_fallbacks = FakeRunnable().with_fallbacks((fake,))  # str -> int298299    assert fake_w_fallbacks.get_input_jsonschema() == {300        "title": "FakeRunnableInput",301        "type": "string",302    }303    assert fake_w_fallbacks.get_output_jsonschema() == {304        "title": "FakeRunnableOutput",305        "type": "integer",306    }307308    def typed_lambda_impl(x: str) -> int:309        return len(x)310311    typed_lambda = RunnableLambda(typed_lambda_impl)  # str -> int312313    assert typed_lambda.get_input_jsonschema() == {314        "title": "typed_lambda_impl_input",315        "type": "string",316    }317    assert typed_lambda.get_output_jsonschema() == {318        "title": "typed_lambda_impl_output",319        "type": "integer",320    }321322    async def typed_async_lambda_impl(x: str) -> int:323        return len(x)324325    typed_async_lambda = RunnableLambda(typed_async_lambda_impl)  # str -> int326327    assert typed_async_lambda.get_input_jsonschema() == {328        "title": "typed_async_lambda_impl_input",329        "type": "string",330    }331    assert typed_async_lambda.get_output_jsonschema() == {332        "title": "typed_async_lambda_impl_output",333        "type": "integer",334    }335336    fake_ret = FakeRetriever()  # str -> list[Document]337338    assert fake_ret.get_input_jsonschema() == {339        "title": "FakeRetrieverInput",340        "type": "string",341    }342    assert _normalize_schema(fake_ret.get_output_jsonschema()) == {343        "$defs": {344            "Document": {345                "description": "Class for storing a piece of text and "346                "associated metadata.\n"347                "\n"348                "!!! note\n"349                "\n"350                "    `Document` is for **retrieval workflows**, not chat I/O. For "351                "sending text\n"352                "    to an LLM in a conversation, use message types from "353                "`langchain.messages`.\n"354                "\n"355                "Example:\n"356                "    ```python\n"357                "    from langchain_core.documents import Document\n"358                "\n"359                "    document = Document(\n"360                '        page_content="Hello, world!", '361                'metadata={"source": "https://example.com"}\n'362                "    )\n"363                "    ```",364                "properties": {365                    "id": {366                        "anyOf": [{"type": "string"}, {"type": "null"}],367                        "default": None,368                        "title": "Id",369                    },370                    "metadata": {"title": "Metadata", "type": "object"},371                    "page_content": {"title": "Page Content", "type": "string"},372                    "type": {373                        "const": "Document",374                        "default": "Document",375                        "title": "Type",376                    },377                },378                "required": ["page_content"],379                "title": "Document",380                "type": "object",381            }382        },383        "items": {"$ref": "#/$defs/Document"},384        "title": "FakeRetrieverOutput",385        "type": "array",386    }387388    fake_llm = FakeListLLM(responses=["a"])  # str -> list[list[str]]389390    assert _schema(fake_llm.input_schema) == snapshot(name="fake_llm_input_schema")391    assert _schema(fake_llm.output_schema) == {392        "title": "FakeListLLMOutput",393        "type": "string",394    }395396    fake_chat = FakeListChatModel(responses=["a"])  # str -> list[list[str]]397398    assert _schema(fake_chat.input_schema) == snapshot(name="fake_chat_input_schema")399    assert _schema(fake_chat.output_schema) == snapshot(name="fake_chat_output_schema")400401    chat_prompt = ChatPromptTemplate.from_messages(402        [403            MessagesPlaceholder(variable_name="history"),404            ("human", "Hello, how are you?"),405        ]406    )407408    assert _normalize_schema(chat_prompt.get_input_jsonschema()) == snapshot(409        name="chat_prompt_input_schema"410    )411    assert _normalize_schema(chat_prompt.get_output_jsonschema()) == snapshot(412        name="chat_prompt_output_schema"413    )414415    prompt = PromptTemplate.from_template("Hello, {name}!")416417    assert prompt.get_input_jsonschema() == {418        "title": "PromptInput",419        "type": "object",420        "properties": {"name": {"title": "Name", "type": "string"}},421        "required": ["name"],422    }423    assert _schema(prompt.output_schema) == snapshot(name="prompt_output_schema")424425    prompt_mapper = PromptTemplate.from_template("Hello, {name}!").map()426427    assert _normalize_schema(prompt_mapper.get_input_jsonschema()) == {428        "$defs": {429            "PromptInput": {430                "properties": {"name": {"title": "Name", "type": "string"}},431                "required": ["name"],432                "title": "PromptInput",433                "type": "object",434            }435        },436        "default": None,437        "items": {"$ref": "#/$defs/PromptInput"},438        "title": "RunnableEach<PromptTemplate>Input",439        "type": "array",440    }441    assert _schema(prompt_mapper.output_schema) == snapshot(442        name="prompt_mapper_output_schema"443    )444445    list_parser = CommaSeparatedListOutputParser()446447    assert _schema(list_parser.input_schema) == snapshot(448        name="list_parser_input_schema"449    )450    assert _schema(list_parser.output_schema) == {451        "title": "CommaSeparatedListOutputParserOutput",452        "type": "array",453        "items": {"type": "string"},454    }455456    seq = prompt | fake_llm | list_parser457458    assert seq.get_input_jsonschema() == {459        "title": "PromptInput",460        "type": "object",461        "properties": {"name": {"title": "Name", "type": "string"}},462        "required": ["name"],463    }464    assert seq.get_output_jsonschema() == {465        "type": "array",466        "items": {"type": "string"},467        "title": "CommaSeparatedListOutputParserOutput",468    }469470    router = RouterRunnable[Any]({})471472    assert _schema(router.input_schema) == {473        "$ref": "#/definitions/RouterInput",474        "definitions": {475            "RouterInput": {476                "description": "Router input.",477                "properties": {478                    "input": {"title": "Input"},479                    "key": {"title": "Key", "type": "string"},480                },481                "required": ["key", "input"],482                "title": "RouterInput",483                "type": "object",484            }485        },486        "title": "RouterRunnableInput",487    }488    assert router.get_output_jsonschema() == {"title": "RouterRunnableOutput"}489490    seq_w_map = (491        prompt492        | fake_llm493        | {494            "original": RunnablePassthrough(input_type=str),495            "as_list": list_parser,496            "length": typed_lambda_impl,497        }498    )499500    assert seq_w_map.get_input_jsonschema() == {501        "title": "PromptInput",502        "type": "object",503        "properties": {"name": {"title": "Name", "type": "string"}},504        "required": ["name"],505    }506    assert seq_w_map.get_output_jsonschema() == {507        "title": "RunnableParallel<original,as_list,length>Output",508        "type": "object",509        "properties": {510            "original": {"title": "Original", "type": "string"},511            "length": {"title": "Length", "type": "integer"},512            "as_list": {513                "title": "As List",514                "type": "array",515                "items": {"type": "string"},516            },517        },518        "required": ["original", "as_list", "length"],519    }520521    # Add a test for schema of runnable assign522    def foo(x: int) -> int:523        return x524525    foo_ = RunnableLambda(foo)526527    assert foo_.assign(bar=lambda _: "foo").get_output_jsonschema() == {528        "properties": {"bar": {"title": "Bar"}, "root": {"title": "Root"}},529        "required": ["root", "bar"],530        "title": "RunnableAssignOutput",531        "type": "object",532    }533534535def test_passthrough_assign_schema() -> None:536    retriever = FakeRetriever()  # str -> list[Document]537    prompt = PromptTemplate.from_template("{context} {question}")538    fake_llm = FakeListLLM(responses=["a"])  # str -> list[list[str]]539540    seq_w_assign = (541        RunnablePassthrough.assign(context=itemgetter("question") | retriever)542        | prompt543        | fake_llm544    )545546    assert seq_w_assign.get_input_jsonschema() == {547        "properties": {"question": {"title": "Question", "type": "string"}},548        "title": "RunnableSequenceInput",549        "type": "object",550        "required": ["question"],551    }552    assert seq_w_assign.get_output_jsonschema() == {553        "title": "FakeListLLMOutput",554        "type": "string",555    }556557    invalid_seq_w_assign = (558        RunnablePassthrough.assign(context=itemgetter("question") | retriever)559        | fake_llm  # type: ignore[operator]560    )561562    # fallback to RunnableAssign.input_schema if next runnable doesn't have563    # expected dict input_schema564    assert invalid_seq_w_assign.get_input_jsonschema() == {565        "properties": {"question": {"title": "Question"}},566        "title": "RunnableParallel<context>Input",567        "type": "object",568        "required": ["question"],569    }570571572def test_lambda_schemas(snapshot: SnapshotAssertion) -> None:573    first_lambda = lambda x: x["hello"]  # noqa: E731574    assert RunnableLambda(first_lambda).get_input_jsonschema() == {575        "title": "RunnableLambdaInput",576        "type": "object",577        "properties": {"hello": {"title": "Hello"}},578        "required": ["hello"],579    }580581    second_lambda = lambda x, y: (x["hello"], x["bye"], y["bah"])  # noqa: E731582    assert RunnableLambda(second_lambda).get_input_jsonschema() == {583        "title": "RunnableLambdaInput",584        "type": "object",585        "properties": {"hello": {"title": "Hello"}, "bye": {"title": "Bye"}},586        "required": ["bye", "hello"],587    }588589    def get_value(value):  # type: ignore[no-untyped-def] # noqa: ANN001,ANN202590        return value["variable_name"]591592    assert RunnableLambda(get_value).get_input_jsonschema() == {593        "title": "get_value_input",594        "type": "object",595        "properties": {"variable_name": {"title": "Variable Name"}},596        "required": ["variable_name"],597    }598599    async def aget_value(value):  # type: ignore[no-untyped-def] # noqa: ANN001,ANN202600        return (value["variable_name"], value.get("another"))601602    assert RunnableLambda(aget_value).get_input_jsonschema() == {603        "title": "aget_value_input",604        "type": "object",605        "properties": {606            "another": {"title": "Another"},607            "variable_name": {"title": "Variable Name"},608        },609        "required": ["another", "variable_name"],610    }611612    async def aget_values(value):  # type: ignore[no-untyped-def] # noqa: ANN001,ANN202613        return {614            "hello": value["variable_name"],615            "bye": value["variable_name"],616            "byebye": value["yo"],617        }618619    assert RunnableLambda(aget_values).get_input_jsonschema() == {620        "title": "aget_values_input",621        "type": "object",622        "properties": {623            "variable_name": {"title": "Variable Name"},624            "yo": {"title": "Yo"},625        },626        "required": ["variable_name", "yo"],627    }628629    class InputType(TypedDict):630        variable_name: str631        yo: int632633    class OutputType(TypedDict):634        hello: str635        bye: str636        byebye: int637638    async def aget_values_typed(value: InputType) -> OutputType:639        return {640            "hello": value["variable_name"],641            "bye": value["variable_name"],642            "byebye": value["yo"],643        }644645    assert _normalize_schema(646        RunnableLambda(aget_values_typed).get_input_jsonschema()647    ) == _normalize_schema(648        {649            "$defs": {650                "InputType": {651                    "properties": {652                        "variable_name": {653                            "title": "Variable Name",654                            "type": "string",655                        },656                        "yo": {"title": "Yo", "type": "integer"},657                    },658                    "required": ["variable_name", "yo"],659                    "title": "InputType",660                    "type": "object",661                }662            },663            "allOf": [{"$ref": "#/$defs/InputType"}],664            "title": "aget_values_typed_input",665        }666    )667668    if PYDANTIC_VERSION_AT_LEAST_29:669        assert _normalize_schema(670            RunnableLambda(aget_values_typed).get_output_jsonschema()671        ) == snapshot(name="schema8")672673674def test_with_types_with_type_generics() -> None:675    """Verify that with_types works if we use things like list[int]."""676677    def foo(x: int) -> None:678        """Add one to the input."""679        raise NotImplementedError680681    # Try specifying some682    RunnableLambda(foo).with_types(683        output_type=list[int],  # type: ignore[arg-type]684        input_type=list[int],  # type: ignore[arg-type]685    )686    RunnableLambda(foo).with_types(687        output_type=Sequence[int],  # type: ignore[arg-type]688        input_type=Sequence[int],  # type: ignore[arg-type]689    )690691692def test_schema_with_itemgetter() -> None:693    """Test runnable with itemgetter."""694    foo = RunnableLambda(itemgetter("hello"))695    assert _schema(foo.input_schema) == {696        "properties": {"hello": {"title": "Hello"}},697        "required": ["hello"],698        "title": "RunnableLambdaInput",699        "type": "object",700    }701    prompt = ChatPromptTemplate.from_template("what is {language}?")702    chain = {"language": itemgetter("language")} | prompt703    assert _schema(chain.input_schema) == {704        "properties": {"language": {"title": "Language"}},705        "required": ["language"],706        "title": "RunnableParallel<language>Input",707        "type": "object",708    }709710711def test_schema_complex_seq() -> None:712    prompt1 = ChatPromptTemplate.from_template("what is the city {person} is from?")713    prompt2 = ChatPromptTemplate.from_template(714        "what country is the city {city} in? respond in {language}"715    )716717    model = FakeListChatModel(responses=[""])718719    chain1 = RunnableSequence[dict[str, Any], str](720        prompt1, model, StrOutputParser(), name="city_chain"721    )722723    assert chain1.name == "city_chain"724725    chain2 = (726        {"city": chain1, "language": itemgetter("language")}727        | prompt2728        | model729        | StrOutputParser()730    )731732    assert chain2.get_input_jsonschema() == {733        "title": "RunnableParallel<city,language>Input",734        "type": "object",735        "properties": {736            "person": {"title": "Person", "type": "string"},737            "language": {"title": "Language"},738        },739        "required": ["person", "language"],740    }741742    assert chain2.get_output_jsonschema() == {743        "title": "StrOutputParserOutput",744        "type": "string",745    }746747    assert chain2.with_types(input_type=str).get_input_jsonschema() == {748        "title": "RunnableSequenceInput",749        "type": "string",750    }751752    assert chain2.with_types(input_type=int).get_output_jsonschema() == {753        "title": "StrOutputParserOutput",754        "type": "string",755    }756757    class InputType(BaseModel):758        person: str759760    assert chain2.with_types(input_type=InputType).get_input_jsonschema() == {761        "title": "InputType",762        "type": "object",763        "properties": {"person": {"title": "Person", "type": "string"}},764        "required": ["person"],765    }766767768def test_configurable_fields(snapshot: SnapshotAssertion) -> None:769    fake_llm = FakeListLLM(responses=["a"])  # str -> list[list[str]]770771    assert fake_llm.invoke("...") == "a"772773    fake_llm_configurable = fake_llm.configurable_fields(774        responses=ConfigurableField(775            id="llm_responses",776            name="LLM Responses",777            description="A list of fake responses for this LLM",778        )779    )780781    assert fake_llm_configurable.invoke("...") == "a"782783    if PYDANTIC_VERSION_AT_LEAST_29:784        assert _normalize_schema(785            fake_llm_configurable.get_config_jsonschema()786        ) == snapshot(name="schema2")787788    fake_llm_configured = fake_llm_configurable.with_config(789        configurable={"llm_responses": ["b"]}790    )791792    assert fake_llm_configured.invoke("...") == "b"793794    prompt = PromptTemplate.from_template("Hello, {name}!")795796    assert prompt.invoke({"name": "John"}) == StringPromptValue(text="Hello, John!")797798    prompt_configurable = prompt.configurable_fields(799        template=ConfigurableField(800            id="prompt_template",801            name="Prompt Template",802            description="The prompt template for this chain",803        )804    )805806    assert prompt_configurable.invoke({"name": "John"}) == StringPromptValue(807        text="Hello, John!"808    )809810    if PYDANTIC_VERSION_AT_LEAST_29:811        assert _normalize_schema(812            prompt_configurable.get_config_jsonschema()813        ) == snapshot(name="schema3")814815    prompt_configured = prompt_configurable.with_config(816        configurable={"prompt_template": "Hello, {name}! {name}!"}817    )818819    assert prompt_configured.invoke({"name": "John"}) == StringPromptValue(820        text="Hello, John! John!"821    )822823    assert prompt_configurable.with_config(824        configurable={"prompt_template": "Hello {name} in {lang}"}825    ).get_input_jsonschema() == {826        "title": "PromptInput",827        "type": "object",828        "properties": {829            "lang": {"title": "Lang", "type": "string"},830            "name": {"title": "Name", "type": "string"},831        },832        "required": ["lang", "name"],833    }834835    chain_configurable = prompt_configurable | fake_llm_configurable | StrOutputParser()836837    assert chain_configurable.invoke({"name": "John"}) == "a"838839    if PYDANTIC_VERSION_AT_LEAST_29:840        assert _normalize_schema(841            chain_configurable.get_config_jsonschema()842        ) == snapshot(name="schema4")843844    assert (845        chain_configurable.with_config(846            configurable={847                "prompt_template": "A very good morning to you, {name} {lang}!",848                "llm_responses": ["c"],849            }850        ).invoke({"name": "John", "lang": "en"})851        == "c"852    )853854    assert chain_configurable.with_config(855        configurable={856            "prompt_template": "A very good morning to you, {name} {lang}!",857            "llm_responses": ["c"],858        }859    ).get_input_jsonschema() == {860        "title": "PromptInput",861        "type": "object",862        "properties": {863            "lang": {"title": "Lang", "type": "string"},864            "name": {"title": "Name", "type": "string"},865        },866        "required": ["lang", "name"],867    }868869    chain_with_map_configurable = prompt_configurable | {870        "llm1": fake_llm_configurable | StrOutputParser(),871        "llm2": fake_llm_configurable | StrOutputParser(),872        "llm3": fake_llm.configurable_fields(873            responses=ConfigurableField("other_responses")874        )875        | StrOutputParser(),876    }877878    assert chain_with_map_configurable.invoke({"name": "John"}) == {879        "llm1": "a",880        "llm2": "a",881        "llm3": "a",882    }883884    if PYDANTIC_VERSION_AT_LEAST_29:885        assert _normalize_schema(886            chain_with_map_configurable.get_config_jsonschema()887        ) == snapshot(name="schema5")888889    assert chain_with_map_configurable.with_config(890        configurable={891            "prompt_template": "A very good morning to you, {name}!",892            "llm_responses": ["c"],893            "other_responses": ["d"],894        }895    ).invoke({"name": "John"}) == {"llm1": "c", "llm2": "c", "llm3": "d"}896897898def test_configurable_alts_factory() -> None:899    fake_llm = FakeListLLM(responses=["a"]).configurable_alternatives(900        ConfigurableField(id="llm", name="LLM"),901        chat=partial(FakeListLLM, responses=["b"]),902    )903904    assert fake_llm.invoke("...") == "a"905906    assert fake_llm.with_config(configurable={"llm": "chat"}).invoke("...") == "b"907908909def test_configurable_fields_prefix_keys(snapshot: SnapshotAssertion) -> None:910    fake_chat = FakeListChatModel(responses=["b"]).configurable_fields(911        responses=ConfigurableFieldMultiOption(912            id="responses",913            name="Chat Responses",914            options={915                "hello": "A good morning to you!",916                "bye": "See you later!",917                "helpful": "How can I help you?",918            },919            default=["hello", "bye"],920        ),921        # (sleep is a configurable field in FakeListChatModel)922        sleep=ConfigurableField(923            id="chat_sleep",924            is_shared=True,925        ),926    )927    fake_llm = (928        FakeListLLM(responses=["a"])929        .configurable_fields(930            responses=ConfigurableField(931                id="responses",932                name="LLM Responses",933                description="A list of fake responses for this LLM",934            )935        )936        .configurable_alternatives(937            ConfigurableField(id="llm", name="LLM"),938            chat=fake_chat | StrOutputParser(),939            prefix_keys=True,940        )941    )942    prompt = PromptTemplate.from_template("Hello, {name}!").configurable_fields(943        template=ConfigurableFieldSingleOption(944            id="prompt_template",945            name="Prompt Template",946            description="The prompt template for this chain",947            options={948                "hello": "Hello, {name}!",949                "good_morning": "A very good morning to you, {name}!",950            },951            default="hello",952        )953    )954955    chain = prompt | fake_llm956957    if PYDANTIC_VERSION_AT_LEAST_29:958        assert _normalize_schema(_schema(chain.config_schema())) == snapshot(959            name="schema6"960        )961962963def test_configurable_fields_example(snapshot: SnapshotAssertion) -> None:964    fake_chat = FakeListChatModel(responses=["b"]).configurable_fields(965        responses=ConfigurableFieldMultiOption(966            id="chat_responses",967            name="Chat Responses",968            options={969                "hello": "A good morning to you!",970                "bye": "See you later!",971                "helpful": "How can I help you?",972            },973            default=["hello", "bye"],974        )975    )976    fake_llm = (977        FakeListLLM(responses=["a"])978        .configurable_fields(979            responses=ConfigurableField(980                id="llm_responses",981                name="LLM Responses",982                description="A list of fake responses for this LLM",983            )984        )985        .configurable_alternatives(986            ConfigurableField(id="llm", name="LLM"),987            chat=fake_chat | StrOutputParser(),988        )989    )990991    prompt = PromptTemplate.from_template("Hello, {name}!").configurable_fields(992        template=ConfigurableFieldSingleOption(993            id="prompt_template",994            name="Prompt Template",995            description="The prompt template for this chain",996            options={997                "hello": "Hello, {name}!",998                "good_morning": "A very good morning to you, {name}!",999            },1000            default="hello",1001        )1002    )10031004    # deduplication of configurable fields1005    chain_configurable = prompt | fake_llm | (lambda x: {"name": x}) | prompt | fake_llm10061007    assert chain_configurable.invoke({"name": "John"}) == "a"10081009    if PYDANTIC_VERSION_AT_LEAST_29:1010        assert _normalize_schema(1011            chain_configurable.get_config_jsonschema()1012        ) == snapshot(name="schema7")10131014    assert (1015        chain_configurable.with_config(configurable={"llm": "chat"}).invoke(1016            {"name": "John"}1017        )1018        == "A good morning to you!"1019    )10201021    assert (1022        chain_configurable.with_config(1023            configurable={"llm": "chat", "chat_responses": ["helpful"]}1024        ).invoke({"name": "John"})1025        == "How can I help you?"1026    )102710281029def test_passthrough_tap(mocker: MockerFixture) -> None:1030    fake = FakeRunnable()1031    mock = mocker.Mock()10321033    seq = RunnablePassthrough[Any](mock) | fake | RunnablePassthrough[Any](mock)10341035    assert seq.invoke("hello", my_kwarg="value") == 51036    assert mock.call_args_list == [1037        mocker.call("hello", my_kwarg="value"),1038        mocker.call(5),1039    ]1040    mock.reset_mock()10411042    assert seq.batch(["hello", "byebye"], my_kwarg="value") == [5, 6]1043    assert len(mock.call_args_list) == 41044    for call in [1045        mocker.call("hello", my_kwarg="value"),1046        mocker.call("byebye", my_kwarg="value"),1047        mocker.call(5),1048        mocker.call(6),1049    ]:1050        assert call in mock.call_args_list1051    mock.reset_mock()10521053    assert seq.batch(["hello", "byebye"], my_kwarg="value", return_exceptions=True) == [1054        5,1055        6,1056    ]1057    assert len(mock.call_args_list) == 41058    for call in [1059        mocker.call("hello", my_kwarg="value"),1060        mocker.call("byebye", my_kwarg="value"),1061        mocker.call(5),1062        mocker.call(6),1063    ]:1064        assert call in mock.call_args_list1065    mock.reset_mock()10661067    assert sorted(1068        a1069        for a in seq.batch_as_completed(1070            ["hello", "byebye"], my_kwarg="value", return_exceptions=True1071        )1072    ) == [1073        (0, 5),1074        (1, 6),1075    ]1076    assert len(mock.call_args_list) == 41077    for call in [1078        mocker.call("hello", my_kwarg="value"),1079        mocker.call("byebye", my_kwarg="value"),1080        mocker.call(5),1081        mocker.call(6),1082    ]:1083        assert call in mock.call_args_list1084    mock.reset_mock()10851086    assert list(1087        seq.stream("hello", {"metadata": {"key": "value"}}, my_kwarg="value")1088    ) == [5]1089    assert mock.call_args_list == [1090        mocker.call("hello", my_kwarg="value"),1091        mocker.call(5),1092    ]1093    mock.reset_mock()109410951096async def test_passthrough_tap_async(mocker: MockerFixture) -> None:1097    fake = FakeRunnable()1098    mock = mocker.Mock()10991100    seq = RunnablePassthrough[Any](mock) | fake | RunnablePassthrough[Any](mock)11011102    assert await seq.ainvoke("hello", my_kwarg="value") == 51103    assert mock.call_args_list == [1104        mocker.call("hello", my_kwarg="value"),1105        mocker.call(5),1106    ]1107    mock.reset_mock()11081109    assert await seq.abatch(["hello", "byebye"], my_kwarg="value") == [5, 6]1110    assert len(mock.call_args_list) == 41111    for call in [1112        mocker.call("hello", my_kwarg="value"),1113        mocker.call("byebye", my_kwarg="value"),1114        mocker.call(5),1115        mocker.call(6),1116    ]:1117        assert call in mock.call_args_list1118    mock.reset_mock()11191120    assert await seq.abatch(1121        ["hello", "byebye"], my_kwarg="value", return_exceptions=True1122    ) == [1123        5,1124        6,1125    ]1126    assert len(mock.call_args_list) == 41127    for call in [1128        mocker.call("hello", my_kwarg="value"),1129        mocker.call("byebye", my_kwarg="value"),1130        mocker.call(5),1131        mocker.call(6),1132    ]:1133        assert call in mock.call_args_list1134    mock.reset_mock()11351136    assert sorted(1137        [1138            a1139            async for a in seq.abatch_as_completed(1140                ["hello", "byebye"], my_kwarg="value", return_exceptions=True1141            )1142        ]1143    ) == [1144        (0, 5),1145        (1, 6),1146    ]1147    assert len(mock.call_args_list) == 41148    for call in [1149        mocker.call("hello", my_kwarg="value"),1150        mocker.call("byebye", my_kwarg="value"),1151        mocker.call(5),1152        mocker.call(6),1153    ]:1154        assert call in mock.call_args_list1155    mock.reset_mock()11561157    assert [1158        part1159        async for part in seq.astream(1160            "hello", {"metadata": {"key": "value"}}, my_kwarg="value"1161        )1162    ] == [5]1163    assert mock.call_args_list == [1164        mocker.call("hello", my_kwarg="value"),1165        mocker.call(5),1166    ]116711681169async def test_with_config_metadata_passthrough(mocker: MockerFixture) -> None:1170    fake = FakeRunnableSerializable()1171    spy = mocker.spy(fake.__class__, "invoke")1172    fakew = fake.configurable_fields(hello=ConfigurableField(id="hello", name="Hello"))11731174    assert (1175        fakew.with_config(tags=["a-tag"]).invoke(1176            "hello",1177            {1178                "configurable": {"hello": "there", "__secret_key": "nahnah"},1179                "metadata": {"bye": "now"},1180            },1181        )1182        == 51183    )1184    assert spy.call_args_list[0].args[1:] == (1185        "hello",1186        {1187            "tags": ["a-tag"],1188            "callbacks": None,1189            "recursion_limit": 25,1190            "configurable": {"hello": "there", "__secret_key": "nahnah"},1191            "metadata": {"bye": "now"},1192        },1193    )1194    spy.reset_mock()119511961197def test_with_config(mocker: MockerFixture) -> None:1198    fake = FakeRunnable()1199    spy = mocker.spy(fake, "invoke")12001201    assert fake.with_config(tags=["a-tag"]).invoke("hello") == 51202    assert spy.call_args_list == [1203        mocker.call(1204            "hello",1205            {"tags": ["a-tag"], "metadata": {}, "configurable": {}},1206        ),1207    ]1208    spy.reset_mock()12091210    fake_1 = RunnablePassthrough[Any]()1211    fake_2 = RunnablePassthrough[Any]()1212    spy_seq_step = mocker.spy(fake_1.__class__, "invoke")12131214    sequence = fake_1.with_config(tags=["a-tag"]) | fake_2.with_config(1215        tags=["b-tag"], max_concurrency=51216    )1217    assert sequence.invoke("hello") == "hello"1218    assert len(spy_seq_step.call_args_list) == 21219    for i, call in enumerate(spy_seq_step.call_args_list):1220        assert call.args[1] == "hello"1221        if i == 0:1222            assert call.args[2].get("tags") == ["a-tag"]1223            assert call.args[2].get("max_concurrency") is None1224        else:1225            assert call.args[2].get("tags") == ["b-tag"]1226            assert call.args[2].get("max_concurrency") == 51227    mocker.stop(spy_seq_step)12281229    assert [1230        *fake.with_config(tags=["a-tag"]).stream(1231            "hello", {"metadata": {"key": "value"}}1232        )1233    ] == [5]1234    assert spy.call_args_list == [1235        mocker.call(1236            "hello",1237            {"tags": ["a-tag"], "metadata": {"key": "value"}, "configurable": {}},1238        ),1239    ]1240    spy.reset_mock()12411242    assert fake.with_config(recursion_limit=5).batch(1243        ["hello", "wooorld"], [{"tags": ["a-tag"]}, {"metadata": {"key": "value"}}]1244    ) == [5, 7]12451246    assert len(spy.call_args_list) == 21247    for i, call in enumerate(1248        sorted(spy.call_args_list, key=lambda x: 0 if x.args[0] == "hello" else 1)1249    ):1250        assert call.args[0] == ("hello" if i == 0 else "wooorld")1251        if i == 0:1252            assert call.args[1].get("recursion_limit") == 51253            assert call.args[1].get("tags") == ["a-tag"]1254            assert call.args[1].get("metadata") == {}1255        else:1256            assert call.args[1].get("recursion_limit") == 51257            assert call.args[1].get("tags") == []1258            assert call.args[1].get("metadata") == {"key": "value"}12591260    spy.reset_mock()12611262    assert sorted(1263        c1264        for c in fake.with_config(recursion_limit=5).batch_as_completed(1265            ["hello", "wooorld"],1266            [{"tags": ["a-tag"]}, {"metadata": {"key": "value"}}],1267        )1268    ) == [(0, 5), (1, 7)]12691270    assert len(spy.call_args_list) == 21271    for i, call in enumerate(1272        sorted(spy.call_args_list, key=lambda x: 0 if x.args[0] == "hello" else 1)1273    ):1274        assert call.args[0] == ("hello" if i == 0 else "wooorld")1275        if i == 0:1276            assert call.args[1].get("recursion_limit") == 51277            assert call.args[1].get("tags") == ["a-tag"]1278            assert call.args[1].get("metadata") == {}1279        else:1280            assert call.args[1].get("recursion_limit") == 51281            assert call.args[1].get("tags") == []1282            assert call.args[1].get("metadata") == {"key": "value"}12831284    spy.reset_mock()12851286    assert fake.with_config(metadata={"a": "b"}).batch(1287        ["hello", "wooorld"], {"tags": ["a-tag"]}1288    ) == [5, 7]1289    assert len(spy.call_args_list) == 21290    for i, call in enumerate(spy.call_args_list):1291        assert call.args[0] == ("hello" if i == 0 else "wooorld")1292        assert call.args[1].get("tags") == ["a-tag"]1293        assert call.args[1].get("metadata") == {"a": "b"}1294    spy.reset_mock()12951296    assert sorted(1297        c for c in fake.batch_as_completed(["hello", "wooorld"], {"tags": ["a-tag"]})1298    ) == [(0, 5), (1, 7)]1299    assert len(spy.call_args_list) == 21300    for i, call in enumerate(spy.call_args_list):1301        assert call.args[0] == ("hello" if i == 0 else "wooorld")1302        assert call.args[1].get("tags") == ["a-tag"]130313041305async def test_with_config_async(mocker: MockerFixture) -> None:1306    fake = FakeRunnable()1307    spy = mocker.spy(fake, "invoke")13081309    handler = ConsoleCallbackHandler()1310    assert (1311        await fake.with_config(metadata={"a": "b"}).ainvoke(1312            "hello", config={"callbacks": [handler]}1313        )1314        == 51315    )1316    assert spy.call_args_list == [1317        mocker.call(1318            "hello",1319            {1320                "callbacks": [handler],1321                "metadata": {"a": "b"},1322                "configurable": {},1323                "tags": [],1324            },1325        ),1326    ]1327    spy.reset_mock()13281329    assert [1330        part async for part in fake.with_config(metadata={"a": "b"}).astream("hello")1331    ] == [5]1332    assert spy.call_args_list == [1333        mocker.call("hello", {"metadata": {"a": "b"}, "tags": [], "configurable": {}}),1334    ]1335    spy.reset_mock()13361337    assert await fake.with_config(recursion_limit=5, tags=["c"]).abatch(1338        ["hello", "wooorld"], {"metadata": {"key": "value"}}1339    ) == [1340        5,1341        7,1342    ]1343    assert sorted(spy.call_args_list) == [1344        mocker.call(1345            "hello",1346            {1347                "metadata": {"key": "value"},1348                "tags": ["c"],1349                "callbacks": None,1350                "recursion_limit": 5,1351                "configurable": {},1352            },1353        ),1354        mocker.call(1355            "wooorld",1356            {1357                "metadata": {"key": "value"},1358                "tags": ["c"],1359                "callbacks": None,1360                "recursion_limit": 5,1361                "configurable": {},1362            },1363        ),1364    ]1365    spy.reset_mock()13661367    assert sorted(1368        [1369            c1370            async for c in fake.with_config(1371                recursion_limit=5, tags=["c"]1372            ).abatch_as_completed(["hello", "wooorld"], {"metadata": {"key": "value"}})1373        ]1374    ) == [1375        (0, 5),1376        (1, 7),1377    ]1378    assert len(spy.call_args_list) == 21379    first_call = next(call for call in spy.call_args_list if call.args[0] == "hello")1380    assert first_call == mocker.call(1381        "hello",1382        {1383            "metadata": {"key": "value"},1384            "tags": ["c"],1385            "callbacks": None,1386            "recursion_limit": 5,1387            "configurable": {},1388        },1389    )1390    second_call = next(call for call in spy.call_args_list if call.args[0] == "wooorld")1391    assert second_call == mocker.call(1392        "wooorld",1393        {1394            "metadata": {"key": "value"},1395            "tags": ["c"],1396            "callbacks": None,1397            "recursion_limit": 5,1398            "configurable": {},1399        },1400    )140114021403def test_default_method_implementations(mocker: MockerFixture) -> None:1404    fake = FakeRunnable()1405    spy = mocker.spy(fake, "invoke")14061407    assert fake.invoke("hello", {"tags": ["a-tag"]}) == 51408    assert spy.call_args_list == [1409        mocker.call("hello", {"tags": ["a-tag"]}),1410    ]1411    spy.reset_mock()14121413    assert [*fake.stream("hello", {"metadata": {"key": "value"}})] == [5]1414    assert spy.call_args_list == [1415        mocker.call("hello", {"metadata": {"key": "value"}}),1416    ]1417    spy.reset_mock()14181419    assert fake.batch(1420        ["hello", "wooorld"], [{"tags": ["a-tag"]}, {"metadata": {"key": "value"}}]1421    ) == [5, 7]14221423    assert len(spy.call_args_list) == 21424    for call in spy.call_args_list:1425        call_arg = call.args[0]14261427        if call_arg == "hello":1428            assert call_arg == "hello"1429            assert call.args[1].get("tags") == ["a-tag"]1430            assert call.args[1].get("metadata") == {}1431        else:1432            assert call_arg == "wooorld"1433            assert call.args[1].get("tags") == []1434            assert call.args[1].get("metadata") == {"key": "value"}14351436    spy.reset_mock()14371438    assert fake.batch(["hello", "wooorld"], {"tags": ["a-tag"]}) == [5, 7]1439    assert len(spy.call_args_list) == 21440    assert {call.args[0] for call in spy.call_args_list} == {"hello", "wooorld"}1441    for call in spy.call_args_list:1442        assert call.args[1].get("tags") == ["a-tag"]1443        assert call.args[1].get("metadata") == {}144414451446async def test_default_method_implementations_async(mocker: MockerFixture) -> None:1447    fake = FakeRunnable()1448    spy = mocker.spy(fake, "invoke")14491450    assert await fake.ainvoke("hello", config={"callbacks": []}) == 51451    assert spy.call_args_list == [1452        mocker.call("hello", {"callbacks": []}),1453    ]1454    spy.reset_mock()14551456    assert [part async for part in fake.astream("hello")] == [5]1457    assert spy.call_args_list == [1458        mocker.call("hello", None),1459    ]1460    spy.reset_mock()14611462    assert await fake.abatch(["hello", "wooorld"], {"metadata": {"key": "value"}}) == [1463        5,1464        7,1465    ]1466    assert {call.args[0] for call in spy.call_args_list} == {"hello", "wooorld"}1467    for call in spy.call_args_list:1468        assert call.args[1] == {1469            "metadata": {"key": "value"},1470            "tags": [],1471            "callbacks": None,1472            "recursion_limit": 25,1473            "configurable": {},1474        }147514761477def test_prompt() -> None:1478    prompt = ChatPromptTemplate.from_messages(1479        messages=[1480            SystemMessage(content="You are a nice assistant."),1481            HumanMessagePromptTemplate.from_template("{question}"),1482        ]1483    )1484    expected = ChatPromptValue(1485        messages=[1486            SystemMessage(content="You are a nice assistant."),1487            HumanMessage(content="What is your name?"),1488        ]1489    )14901491    assert prompt.invoke({"question": "What is your name?"}) == expected14921493    assert prompt.batch(1494        [1495            {"question": "What is your name?"},1496            {"question": "What is your favorite color?"},1497        ]1498    ) == [1499        expected,1500        ChatPromptValue(1501            messages=[1502                SystemMessage(content="You are a nice assistant."),1503                HumanMessage(content="What is your favorite color?"),1504            ]1505        ),1506    ]15071508    assert [*prompt.stream({"question": "What is your name?"})] == [expected]150915101511async def test_prompt_async() -> None:1512    prompt = ChatPromptTemplate.from_messages(1513        messages=[1514            SystemMessage(content="You are a nice assistant."),1515            HumanMessagePromptTemplate.from_template("{question}"),1516        ]1517    )1518    expected = ChatPromptValue(1519        messages=[1520            SystemMessage(content="You are a nice assistant."),1521            HumanMessage(content="What is your name?"),1522        ]1523    )15241525    assert await prompt.ainvoke({"question": "What is your name?"}) == expected15261527    assert await prompt.abatch(1528        [1529            {"question": "What is your name?"},1530            {"question": "What is your favorite color?"},1531        ]1532    ) == [1533        expected,1534        ChatPromptValue(1535            messages=[1536                SystemMessage(content="You are a nice assistant."),1537                HumanMessage(content="What is your favorite color?"),1538            ]1539        ),1540    ]15411542    assert [1543        part async for part in prompt.astream({"question": "What is your name?"})1544    ] == [expected]15451546    stream_log = [1547        part async for part in prompt.astream_log({"question": "What is your name?"})1548    ]15491550    assert len(stream_log[0].ops) == 11551    assert stream_log[0].ops[0]["op"] == "replace"1552    assert stream_log[0].ops[0]["path"] == ""1553    assert stream_log[0].ops[0]["value"]["logs"] == {}1554    assert stream_log[0].ops[0]["value"]["final_output"] is None1555    assert stream_log[0].ops[0]["value"]["streamed_output"] == []1556    assert isinstance(stream_log[0].ops[0]["value"]["id"], str)15571558    assert stream_log[1:] == [1559        RunLogPatch(1560            {"op": "add", "path": "/streamed_output/-", "value": expected},1561            {1562                "op": "replace",1563                "path": "/final_output",1564                "value": ChatPromptValue(1565                    messages=[1566                        SystemMessage(content="You are a nice assistant."),1567                        HumanMessage(content="What is your name?"),1568                    ]1569                ),1570            },1571        ),1572    ]15731574    stream_log_state = [1575        part1576        async for part in prompt.astream_log(1577            {"question": "What is your name?"}, diff=False1578        )1579    ]15801581    # remove random id1582    stream_log[0].ops[0]["value"]["id"] = "00000000-0000-0000-0000-000000000000"1583    stream_log_state[-1].ops[0]["value"]["id"] = "00000000-0000-0000-0000-000000000000"1584    stream_log_state[-1].state["id"] = "00000000-0000-0000-0000-000000000000"15851586    # assert output with diff=False matches output with diff=True1587    assert stream_log_state[-1].ops == [op for chunk in stream_log for op in chunk.ops]1588    assert stream_log_state[-1] == RunLog(1589        *[op for chunk in stream_log for op in chunk.ops],1590        state={1591            "final_output": ChatPromptValue(1592                messages=[1593                    SystemMessage(content="You are a nice assistant."),1594                    HumanMessage(content="What is your name?"),1595                ]1596            ),1597            "id": "00000000-0000-0000-0000-000000000000",1598            "logs": {},1599            "streamed_output": [1600                ChatPromptValue(1601                    messages=[1602                        SystemMessage(content="You are a nice assistant."),1603                        HumanMessage(content="What is your name?"),1604                    ]1605                )1606            ],1607            "type": "prompt",1608            "name": "ChatPromptTemplate",1609        },1610    )16111612    # nested inside trace_with_chain_group16131614    async with atrace_as_chain_group("a_group") as manager:1615        stream_log_nested = [1616            part1617            async for part in prompt.astream_log(1618                {"question": "What is your name?"}, config={"callbacks": manager}1619            )1620        ]16211622    assert len(stream_log_nested[0].ops) == 11623    assert stream_log_nested[0].ops[0]["op"] == "replace"1624    assert stream_log_nested[0].ops[0]["path"] == ""1625    assert stream_log_nested[0].ops[0]["value"]["logs"] == {}1626    assert stream_log_nested[0].ops[0]["value"]["final_output"] is None1627    assert stream_log_nested[0].ops[0]["value"]["streamed_output"] == []1628    assert isinstance(stream_log_nested[0].ops[0]["value"]["id"], str)16291630    assert stream_log_nested[1:] == [1631        RunLogPatch(1632            {"op": "add", "path": "/streamed_output/-", "value": expected},1633            {1634                "op": "replace",1635                "path": "/final_output",1636                "value": ChatPromptValue(1637                    messages=[1638                        SystemMessage(content="You are a nice assistant."),1639                        HumanMessage(content="What is your name?"),1640                    ]1641                ),1642            },1643        ),1644    ]164516461647def test_prompt_template_params() -> None:1648    prompt = ChatPromptTemplate.from_template(1649        "Respond to the following question: {question}"1650    )1651    result = prompt.invoke(1652        {1653            "question": "test",1654            "topic": "test",1655        }1656    )1657    assert result == ChatPromptValue(1658        messages=[HumanMessage(content="Respond to the following question: test")]1659    )16601661    with pytest.raises(KeyError):1662        prompt.invoke({})166316641665def test_with_listeners(mocker: MockerFixture) -> None:1666    prompt = (1667        SystemMessagePromptTemplate.from_template("You are a nice assistant.")1668        + "{question}"1669    )1670    chat = FakeListChatModel(responses=["foo"])16711672    chain = prompt | chat16731674    mock_start = mocker.Mock()1675    mock_end = mocker.Mock()16761677    chain.with_listeners(on_start=mock_start, on_end=mock_end).invoke(1678        {"question": "Who are you?"}1679    )16801681    assert mock_start.call_count == 11682    assert mock_start.call_args[0][0].name == "RunnableSequence"1683    assert mock_end.call_count == 116841685    mock_start.reset_mock()1686    mock_end.reset_mock()16871688    with trace_as_chain_group("hello") as manager:1689        chain.with_listeners(on_start=mock_start, on_end=mock_end).invoke(1690            {"question": "Who are you?"}, {"callbacks": manager}1691        )16921693    assert mock_start.call_count == 11694    assert mock_start.call_args[0][0].name == "RunnableSequence"1695    assert mock_end.call_count == 1169616971698async def test_with_listeners_async(mocker: MockerFixture) -> None:1699    prompt = (1700        SystemMessagePromptTemplate.from_template("You are a nice assistant.")1701        + "{question}"1702    )1703    chat = FakeListChatModel(responses=["foo"])17041705    chain = prompt | chat17061707    mock_start = mocker.Mock()1708    mock_end = mocker.Mock()17091710    await chain.with_listeners(on_start=mock_start, on_end=mock_end).ainvoke(1711        {"question": "Who are you?"}1712    )17131714    assert mock_start.call_count == 11715    assert mock_start.call_args[0][0].name == "RunnableSequence"1716    assert mock_end.call_count == 117171718    mock_start.reset_mock()1719    mock_end.reset_mock()17201721    async with atrace_as_chain_group("hello") as manager:1722        await chain.with_listeners(on_start=mock_start, on_end=mock_end).ainvoke(1723            {"question": "Who are you?"}, {"callbacks": manager}1724        )17251726    assert mock_start.call_count == 11727    assert mock_start.call_args[0][0].name == "RunnableSequence"1728    assert mock_end.call_count == 1172917301731def test_with_listener_propagation(mocker: MockerFixture) -> None:1732    prompt = (1733        SystemMessagePromptTemplate.from_template("You are a nice assistant.")1734        + "{question}"1735    )1736    chat = FakeListChatModel(responses=["foo"])1737    chain = prompt | chat1738    mock_start = mocker.Mock()1739    mock_end = mocker.Mock()1740    chain_with_listeners = chain.with_listeners(on_start=mock_start, on_end=mock_end)17411742    chain_with_listeners.with_retry().invoke({"question": "Who are you?"})17431744    assert mock_start.call_count == 11745    assert mock_start.call_args[0][0].name == "RunnableSequence"1746    assert mock_end.call_count == 117471748    mock_start.reset_mock()1749    mock_end.reset_mock()17501751    chain_with_listeners.invoke({"question": "Who are you?"})17521753    assert mock_start.call_count == 11754    assert mock_start.call_args[0][0].name == "RunnableSequence"1755    assert mock_end.call_count == 117561757    mock_start.reset_mock()1758    mock_end.reset_mock()17591760    chain_with_listeners.with_config({"tags": ["foo"]}).invoke(1761        {"question": "Who are you?"}1762    )17631764    assert mock_start.call_count == 11765    assert mock_start.call_args[0][0].name == "RunnableSequence"1766    assert mock_end.call_count == 117671768    mock_start.reset_mock()1769    mock_end.reset_mock()17701771    chain_with_listeners.bind(stop=["foo"]).invoke({"question": "Who are you?"})17721773    assert mock_start.call_count == 11774    assert mock_start.call_args[0][0].name == "RunnableSequence"1775    assert mock_end.call_count == 117761777    mock_start.reset_mock()1778    mock_end.reset_mock()17791780    mock_start_inner = mocker.Mock()1781    mock_end_inner = mocker.Mock()17821783    chain_with_listeners.with_listeners(1784        on_start=mock_start_inner, on_end=mock_end_inner1785    ).invoke({"question": "Who are you?"})17861787    assert mock_start.call_count == 11788    assert mock_start.call_args[0][0].name == "RunnableSequence"1789    assert mock_end.call_count == 11790    assert mock_start_inner.call_count == 11791    assert mock_start_inner.call_args[0][0].name == "RunnableSequence"1792    assert mock_end_inner.call_count == 1179317941795@freeze_time("2023-01-01")1796@pytest.mark.usefixtures("deterministic_uuids")1797def test_prompt_with_chat_model(1798    mocker: MockerFixture,1799    snapshot: SnapshotAssertion,1800) -> None:1801    prompt = (1802        SystemMessagePromptTemplate.from_template("You are a nice assistant.")1803        + "{question}"1804    )1805    chat = FakeListChatModel(responses=["foo"])18061807    chain = prompt | chat18081809    assert _normalize_lc_version(repr(chain)) == snapshot1810    assert isinstance(chain, RunnableSequence)1811    assert chain.first == prompt1812    assert chain.middle == []1813    assert chain.last == chat1814    assert _normalize_lc_version(dumps(chain, pretty=True)) == snapshot18151816    # Test invoke1817    prompt_spy = mocker.spy(prompt.__class__, "invoke")1818    chat_spy = mocker.spy(chat.__class__, "invoke")1819    tracer = FakeTracer()1820    assert chain.invoke(1821        {"question": "What is your name?"}, {"callbacks": [tracer]}1822    ) == _any_id_ai_message(content="foo")1823    assert prompt_spy.call_args.args[1] == {"question": "What is your name?"}1824    assert chat_spy.call_args.args[1] == ChatPromptValue(1825        messages=[1826            SystemMessage(content="You are a nice assistant."),1827            HumanMessage(content="What is your name?"),1828        ]1829    )18301831    assert tracer.runs == snapshot18321833    mocker.stop(prompt_spy)1834    mocker.stop(chat_spy)18351836    # Test batch1837    prompt_spy = mocker.spy(prompt.__class__, "batch")1838    chat_spy = mocker.spy(chat.__class__, "batch")1839    tracer = FakeTracer()1840    assert chain.batch(1841        [1842            {"question": "What is your name?"},1843            {"question": "What is your favorite color?"},1844        ],1845        {"callbacks": [tracer]},1846    ) == [1847        _any_id_ai_message(content="foo"),1848        _any_id_ai_message(content="foo"),1849    ]1850    assert prompt_spy.call_args.args[1] == [1851        {"question": "What is your name?"},1852        {"question": "What is your favorite color?"},1853    ]1854    assert chat_spy.call_args.args[1] == [1855        ChatPromptValue(1856            messages=[1857                SystemMessage(content="You are a nice assistant."),1858                HumanMessage(content="What is your name?"),1859            ]1860        ),1861        ChatPromptValue(1862            messages=[1863                SystemMessage(content="You are a nice assistant."),1864                HumanMessage(content="What is your favorite color?"),1865            ]1866        ),1867    ]1868    assert (1869        len(1870            [1871                r1872                for r in tracer.runs1873                if r.parent_run_id is None and len(r.child_runs) == 21874            ]1875        )1876        == 21877    ), "Each of 2 outer runs contains exactly two inner runs (1 prompt, 1 chat)"1878    mocker.stop(prompt_spy)1879    mocker.stop(chat_spy)18801881    # Test stream1882    prompt_spy = mocker.spy(prompt.__class__, "invoke")1883    chat_spy = mocker.spy(chat.__class__, "stream")1884    tracer = FakeTracer()1885    assert [1886        *chain.stream({"question": "What is your name?"}, {"callbacks": [tracer]})1887    ] == [1888        _any_id_ai_message_chunk(content="f"),1889        _any_id_ai_message_chunk(content="o"),1890        _any_id_ai_message_chunk(content="o", chunk_position="last"),1891    ]1892    assert prompt_spy.call_args.args[1] == {"question": "What is your name?"}1893    assert chat_spy.call_args.args[1] == ChatPromptValue(1894        messages=[1895            SystemMessage(content="You are a nice assistant."),1896            HumanMessage(content="What is your name?"),1897        ]1898    )189919001901@freeze_time("2023-01-01")1902@pytest.mark.usefixtures("deterministic_uuids")1903async def test_prompt_with_chat_model_async(1904    mocker: MockerFixture,1905    snapshot: SnapshotAssertion,1906) -> None:1907    prompt = (1908        SystemMessagePromptTemplate.from_template("You are a nice assistant.")1909        + "{question}"1910    )1911    chat = FakeListChatModel(responses=["foo"])19121913    chain = prompt | chat19141915    assert _normalize_lc_version(repr(chain)) == snapshot1916    assert isinstance(chain, RunnableSequence)1917    assert chain.first == prompt1918    assert chain.middle == []1919    assert chain.last == chat1920    assert _normalize_lc_version(dumps(chain, pretty=True)) == snapshot19211922    # Test invoke1923    prompt_spy = mocker.spy(prompt.__class__, "ainvoke")1924    chat_spy = mocker.spy(chat.__class__, "ainvoke")1925    tracer = FakeTracer()1926    assert await chain.ainvoke(1927        {"question": "What is your name?"}, {"callbacks": [tracer]}1928    ) == _any_id_ai_message(content="foo")1929    assert prompt_spy.call_args.args[1] == {"question": "What is your name?"}1930    assert chat_spy.call_args.args[1] == ChatPromptValue(1931        messages=[1932            SystemMessage(content="You are a nice assistant."),1933            HumanMessage(content="What is your name?"),1934        ]1935    )19361937    assert tracer.runs == snapshot19381939    mocker.stop(prompt_spy)1940    mocker.stop(chat_spy)19411942    # Test batch1943    prompt_spy = mocker.spy(prompt.__class__, "abatch")1944    chat_spy = mocker.spy(chat.__class__, "abatch")1945    tracer = FakeTracer()1946    assert await chain.abatch(1947        [1948            {"question": "What is your name?"},1949            {"question": "What is your favorite color?"},1950        ],1951        {"callbacks": [tracer]},1952    ) == [1953        _any_id_ai_message(content="foo"),1954        _any_id_ai_message(content="foo"),1955    ]1956    assert prompt_spy.call_args.args[1] == [1957        {"question": "What is your name?"},1958        {"question": "What is your favorite color?"},1959    ]1960    assert chat_spy.call_args.args[1] == [1961        ChatPromptValue(1962            messages=[1963                SystemMessage(content="You are a nice assistant."),1964                HumanMessage(content="What is your name?"),1965            ]1966        ),1967        ChatPromptValue(1968            messages=[1969                SystemMessage(content="You are a nice assistant."),1970                HumanMessage(content="What is your favorite color?"),1971            ]1972        ),1973    ]1974    assert (1975        len(1976            [1977                r1978                for r in tracer.runs1979                if r.parent_run_id is None and len(r.child_runs) == 21980            ]1981        )1982        == 21983    ), "Each of 2 outer runs contains exactly two inner runs (1 prompt, 1 chat)"1984    mocker.stop(prompt_spy)1985    mocker.stop(chat_spy)19861987    # Test stream1988    prompt_spy = mocker.spy(prompt.__class__, "ainvoke")1989    chat_spy = mocker.spy(chat.__class__, "astream")1990    tracer = FakeTracer()1991    assert [1992        a1993        async for a in chain.astream(1994            {"question": "What is your name?"}, {"callbacks": [tracer]}1995        )1996    ] == [1997        _any_id_ai_message_chunk(content="f"),1998        _any_id_ai_message_chunk(content="o"),1999        _any_id_ai_message_chunk(content="o", chunk_position="last"),2000    ]