BaseChatModel.astream() does not pass the run_manager object to the BaseChatModel._astream() method

[olokshyn]

Checked other resources

• I added a very descriptive title to this issue.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.
• I am sure that this is a bug in LangChain rather than my code.
• The bug is not resolved by updating to the latest stable version of LangChain (or the specific integration package).

Example Code

The bug was introduced in langchain/libs/core/langchain_core/language_models/chat_models.py (link to master) since v0.1.14.

First, note the definition of the BaseChatModel._astream() method (link to master):

    async def _astream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:

It accepts the optional run_manager: Optional[AsyncCallbackManagerForLLMRun] parameter.

That's how BaseChatModel.astream() calls BaseChatModel._astream() (link to master), note that it never passes the run_manager parameter:

        (run_manager,) = await callback_manager.on_chat_model_start(
            dumpd(self),
            [messages],
            invocation_params=params,
            options=options,
            name=config.get("run_name"),
            run_id=config.pop("run_id", None),
            batch_size=1,
        )

        generation: Optional[ChatGenerationChunk] = None
        try:
            async for chunk in self._astream(
                messages,
                stop=stop,
                **kwargs,
            ):

That's how BaseChatModel.astream() used to call BaseChatModel._astream() in v0.1.13, note that the run_manager object used to be passed properly:

        (run_manager,) = await callback_manager.on_chat_model_start(
            dumpd(self),
            [messages],
            invocation_params=params,
            options=options,
            name=config.get("run_name"),
            run_id=config.pop("run_id", None),
            batch_size=1,
        )

        generation: Optional[ChatGenerationChunk] = None
        try:
            async for chunk in self._astream(
                messages,
                stop=stop,
                run_manager=run_manager,
                **kwargs,
            ):

As stated in the current docs, BaseChatMode._astream() method can be overridden by users, so preserving the same API within minor version updates is important. That's why I consider this a bug.

Error Message and Stack Trace (if applicable)

No response

Description

I am using a RunnableConfig object to pass the chat ID through the whole pipeline for logging purposes:

async for event in my_chain.astream(
    input={"question": ...},
    config=RunnableConfig(metadata={"chat_id": chat.id}),
):

Then I try to get the chat ID back in my custom model derived from BaseChatModel like this:

    async def _astream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        if not run_manager.metadata.get("chat_id"):
            raise ValueError("chat_id is required to extract the logger")
        chat_id = run_manager.metadata["chat_id"]
        ...

This code worked perfectly fine in v0.1.13.
Starting with v.0.1.14 and up to the current master, the run_manager object is never passed inside the BaseChatModel._astream() method. Hence, my code always sees run_manager == None and fails with the ValueError exception.

System Info

LangChain libs versions:

langchain==0.1.17
langchain-community==0.0.36
langchain-core==0.1.48
langchain-openai==0.1.6
langchain-text-splitters==0.0.1

Platform: MacOS
Python: 3.11.7

System Information
------------------
> OS:  Darwin
> OS Version:  Darwin Kernel Version 22.6.0: Mon Feb 19 19:45:09 PST 2024; root:xnu-8796.141.3.704.6~1/RELEASE_ARM64_T6000
> Python Version:  3.11.7 (main, Dec 15 2023, 12:09:04) [Clang 14.0.6 ]

Package Information
-------------------
> langchain_core: 0.1.48
> langchain: 0.1.17
> langchain_community: 0.0.36
> langsmith: 0.1.49
> langchain_openai: 0.1.6
> langchain_text_splitters: 0.0.1

Packages not installed (Not Necessarily a Problem)
--------------------------------------------------
The following packages were not found:

> langgraph
> langserve

[eyurtsev]

Correct fix here is to likely propagate config to the _ainvoke interface. Same issue appears with BaseRetriever interface where the method that's meant to be overridden by user doesn't accept config