🛸 Streaming: RFC Adding astream_event to all Runnable objects to help with streaming use cases

[eyurtsev]

Hi everyone!

We want to improve the streaming experience in LangChain. We're considering adding a astream_event method to the Runnable interface. The code below is from the following PR and has not yet been committed into langchain-core.

Does this API work for your use case? Anything you'd like to see done better or differently? We would love to hear any and all feedback.

Gist

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI


prompt = ChatPromptTemplate.from_messages([('system', 'You are the cat equivalent of 007.'), ('human', '{question}')])

prompt.invoke({"question": "hello"})

model = ChatOpenAI(temperature=0)
chain = prompt | model

idx = 0
async for event in chain.astream_events({"question": "tell me a secret"}):
    print(event)
    idx += 1 
    if idx > 8:
        break

First 8 events:

{'event': 'on_chain_start', 'run_id': '7f63fc12-6c18-4496-a34c-7eeef652d5c9', 'name': 'RunnableSequence', 'tags': [], 'metadata': {}, 'data': {'input': {'question': 'tell me a secret'}}}
{'event': 'on_prompt_start', 'name': 'ChatPromptTemplate', 'run_id': '427876bc-1670-4092-adb3-b1d0d17f619b', 'tags': ['seq:step:1'], 'metadata': {}, 'data': {'input': {'question': 'tell me a secret'}}}
{'event': 'on_prompt_end', 'name': 'ChatPromptTemplate', 'run_id': '427876bc-1670-4092-adb3-b1d0d17f619b', 'tags': ['seq:step:1'], 'metadata': {}, 'data': {'input': {'question': 'tell me a secret'}, 'output': ChatPromptValue(messages=[SystemMessage(content='You are the cat equivalent of 007.'), HumanMessage(content='tell me a secret')])}}
{'event': 'on_chat_model_start', 'name': 'ChatOpenAI', 'run_id': '3a290b8c-69a9-4e2b-bff0-9aa0b683b9cb', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'input': {'messages': [[SystemMessage(content='You are the cat equivalent of 007.'), HumanMessage(content='tell me a secret')]]}}}
{'event': 'on_chain_stream', 'run_id': '7f63fc12-6c18-4496-a34c-7eeef652d5c9', 'tags': [], 'metadata': {}, 'name': 'RunnableSequence', 'data': {'chunk': AIMessageChunk(content='')}}
{'event': 'on_chat_model_stream', 'name': 'ChatOpenAI', 'run_id': '3a290b8c-69a9-4e2b-bff0-9aa0b683b9cb', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': AIMessageChunk(content='')}}
{'event': 'on_chain_stream', 'run_id': '7f63fc12-6c18-4496-a34c-7eeef652d5c9', 'tags': [], 'metadata': {}, 'name': 'RunnableSequence', 'data': {'chunk': AIMessageChunk(content='As')}}
{'event': 'on_chat_model_stream', 'name': 'ChatOpenAI', 'run_id': '3a290b8c-69a9-4e2b-bff0-9aa0b683b9cb', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': AIMessageChunk(content='As')}}
{'event': 'on_chain_stream', 'run_id': '7f63fc12-6c18-4496-a34c-7eeef652d5c9', 'tags': [], 'metadata': {}, 'name': 'RunnableSequence', 'data': {'chunk': AIMessageChunk(content=' the')}}

Example Notebook

Please check out this notebook.

The notebook shows how to get streaming working from LLMs used within tools.

Event Hooks Reference

Here is a reference table that shows some events that might be emitted by the various Runnable objects.
Definitions for some of the Runnable are included after the table.

⚠️ When streaming the inputs for the runnable will not be available until the input stream has been entirely consumed This means that the inputs will be available at for the corresponding end hook rather than start event.

event	name	chunk	input	output
on_chat_model_start	[model name]		{"messages": [[SystemMessage, HumanMessage]]}
on_chat_model_stream	[model name]	AIMessageChunk(content="hello")
on_chat_model_end	[model name]		{"messages": [[SystemMessage, HumanMessage]]}	{"generations": [...], "llm_output": None, ...}
on_llm_start	[model name]		{'input': 'hello'}
on_llm_stream	[model name]	'Hello'
on_llm_end	[model name]		'Hello human!'
on_chain_start	format_docs
on_chain_stream	format_docs	"hello world!, goodbye world!"
on_chain_end	format_docs		[Document(...)]	"hello world!, goodbye world!"
on_tool_start	some_tool		{"x": 1, "y": "2"}
on_tool_stream	some_tool	{"x": 1, "y": "2"}
on_tool_end	some_tool			{"x": 1, "y": "2"}
on_retriever_start	[retriever name]		{"query": "hello"}
on_retriever_chunk	[retriever name]	{documents: [...]}
on_retriever_end	[retriever name]		{"query": "hello"}	{documents: [...]}
on_prompt_start	[template_name]		{"question": "hello"}
on_prompt_end	[template_name]		{"question": "hello"}	ChatPromptValue(messages: [SystemMessage, ...])

Here are declarations associated with the events shown above:

format_docs:

def format_docs(docs: List[Document]) -> str:
    '''Format the docs.'''
    return ", ".join([doc.page_content for doc in docs])

format_docs = RunnableLambda(format_docs)

some_tool:

@tool
def some_tool(x: int, y: str) -> dict:
    '''Some_tool.'''
    return {"x": x, "y": y}

prompt:

template = ChatPromptTemplate.from_messages(
    [("system", "You are Cat Agent 007"), ("human", "{question}")]
).with_config({"run_name": "my_template", "tags": ["my_template"]})

[eyurtsev]

This code is merged on master. If you see anything failing or want changes to the schema API, please post here!

[danilopeixoto]

Hi @eyurtsev,
Thanks for enabling streaming for LangChain agents and tools.

Do you know if the current implementation could work with RemoteRunnable?

[eyurtsev]

Not supported yet, but here's the progress so far: langchain-ai/langserve#413. Will be released soonish

[eyurtsev]

https://github.com/langchain-ai/langserve/releases/tag/v0.0.40 -- OK released.

Please open an issue in LangServe if you bump into any problems with the LangServe implementation. Serialization is a bit difficult in this case if you're using any custom data types. I'd recommend making sure that the client code is only requesting event sets that it cares about to reduce amount of data set over the wire + reduce any potential of having serialization issues appearing from steps in the chain that you don't care about

[danilopeixoto]

Thank you! I had the chance to test it out, and it is working exceptionally well for my specific use case (basic chat with agents and tools).

[eyurtsev]

Documentation is available here:

General Docs: https://python.langchain.com/docs/expression_language/streaming
Docs for AgentExecutor: https://python.langchain.com/docs/modules/agents/how_to/streaming
Docs for LangGraph: https://github.com/langchain-ai/langgraph/blob/main/examples/streaming-tokens.ipynb
YouTube: https://youtube.com/watch?v=ZcEMLz27sL4

[baskaryan]

Update to LangGraph docs: langchain-ai/langgraph#67

[SavvasMohito]

what version of python are you on?

@hwchase17 3.10. Is this only supported on a later version?

[eyurtsev]

@SavvasMohito it's supported for earlier python as well, but you need to propagate the config explicitly (so that callbacks are propagated). But we will try to release a fix within a few days so you won't need to manually propagate callbacks in earlier python versions

[SavvasMohito]

I upgraded to python 3.12 and it works fine now. Thanks for the help @eyurtsev :)

[Odrec]

When I'm trying event streaming following the documentation using the AgentExecutor with a web search tool. In this part:

    if kind == "on_chat_model_stream":
        content = event["data"]["chunk"].content
        if content:
            # Empty content in the context of OpenAI means
            # that the model is asking for a tool to be invoked.
            # So we only print non-empty content
            print(content, end="|")

I'm getting this output at the beginning:

```|json|
|{
|   | "|action|":| "|Final| Answer|",
|   | "|action|_input|":|Once| in| a| quaint|

So instead of just the text answer (Once| in| a| quaint|) it comes wrapped in a string following some json format? which makes parsing during streaming pretty difficult. What I am doing wrong?

[adreo00]

@eyurtsev Linking to this discussion, as per your request.
#18694

[alexgg278]

For my use case I would like to stream only the final Answer of the Agent. However, my astream_events returns all kind of events by the "on_chat_model_stream", any clue why I am not getting this event that I would need for streaming the final generated output?

@eyurtsev

[eyurtsev]

Follow this tutorial: https://python.langchain.com/docs/modules/agents/how_to/streaming#custom-streaming-with-events

It will show how to stream the final output from the agent.

astream events provides rich output, you can filter the output based on what you need for your use case

[Aekansh-Ak]

model = AutoModelForCausalLM.from_pretrained(model_id,
#load_in_8bit=True, # set these options if your GPU supports them!
device_map='auto',
torch_dtype=torch.float16,
#low_cpu_mem_usage=False
)

pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_length=21000,
temperature=0.3,
top_p=0.95,
repetition_penalty=1.15,
)

llm = HuggingFacePipeline(pipeline=pipe)

rag_chain = (
{"context": retriever, "question": RunnablePassthrough()}
| prompt
| llm
| StrOutputParser()
)

async for chunk in rag_chain.astream("How are you?"):
print(chunk, end="", flush=True)

Returns error-:

async for chunk in rag_chain.astream("How are you?"):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

SyntaxError: 'async for' outside async function

[eyurtsev]

@Aekansh-Ak this is not an error. Please refer to other online resources on how to use async programming in python.