Incorrect listeners parameters for Runnable.with_listeners() and .map()

skozlovf commented 2 months ago

Checked other resources

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

Example Code

from langchain_core.runnables import RunnableLambda
from langchain_core.tracers.schemas import Run

def fake_chain(inputs: dict) -> dict:
    return {**inputs, "key": "extra"}

def on_start(run: Run):
    print("on_start:", run.inputs)

def on_end(run: Run):
    print("on_end:  ", run.outputs)

chain = RunnableLambda(fake_chain).with_listeners(on_end=on_end, on_start=on_start)
chain = chain.map()

data = [{"name": "one"}, {"name": "two"}]
out = chain.invoke(data, config={"max_concurrency": 1})
print("result:  ", out)

max_concurrency is added for simplicity.

Error Message and Stack Trace (if applicable)

No response

Description

I want to store fake_chain output using listeners. with_listeners() allows to hook only top level runnable (according to its docstring). But run object is incorrect if use map().

I expect to see

on_start: {'name': 'one'}
on_start: {'name': 'two'}
on_end:   {'name': 'one', 'key': 'extra'}
on_end:   {'name': 'two', 'key': 'extra'}
result:   [{'name': 'one', 'key': 'extra'}, {'name': 'two', 'key': 'extra'}]

but get

on_start: {'name': 'one'}
on_start: {'name': 'one'}  # <!
on_end:   {'name': 'one', 'key': 'extra'}
on_end:   {'name': 'one', 'key': 'extra'}   # <!
result:   [{'name': 'one', 'key': 'extra'}, {'name': 'two', 'key': 'extra'}]

I didn't dive deeper, but smth wrong happens in the RunnableBindingBase.batch() -> _merge_configs() (a guess).

System Info

$ pip freeze | grep langchain
langchain==0.1.16               
langchain-anthropic==0.1.4      
langchain-community==0.0.33     
langchain-core==0.1.43          
langchain-google-genai==0.0.11  
langchain-google-vertexai==0.1.2
langchain-groq==0.0.1           
langchain-openai==0.0.8         
langchain-text-splitters==0.0.1

platform: linux python: 3.11.8