🛸 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

[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?

[hizbul-inument]

I am getting some extra text at the beginning.
Chat model is starting...Plan: I will greet the user and ask for their name and location.
Action: ```json
[
{
"tool_name": "directly_answer",
"parameters": {}
}
]

Relevant Documents: None
Cited Documents: None
Answer: Hello! I'm Mak, how can I help you today? 
@Odrec did you get any solution?

[eyurtsev]

@hizbul-inument / @Odrec there isn't enough context in your questions to be able to answer them.

You're probably using an agent that isn't relying on native tool calling capabilities of chat models, which was fine for state of the art 1.5 years ago, but definitely isn't a good idea these days. Please review the quick start tutorials and concepts associated with tool calling in the main documentation.

https://python.langchain.com/docs/concepts/#functiontool-calling

[hizbul-inument]

Here is further context: https://stackoverflow.com/questions/79065125/how-to-get-langchain-stream-data-without-logs

[hizbul-inument]

@eyurtsev using create_cohere_react_agent agent. cohere saying in documents logs may come with extra text, but i am getting it during streaming. Ref: https://docs.cohere.com/v2/docs/implementing-a-multi-step-agent-with-langchain

[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.

[Sanzid88888]

#22227
@eyurtsev can you please check this

[mahimairaja]

Is it possible to track the details like total_tokens and total_cost when using astream_events?

[eyurtsev]

If the model adds the information into the standard usage metadata field then yes.

from langchain_openai import ChatOpenAI

model = ChatOpenAI(stream_usage=True) # <-- With OpenAI you have to opt-in

events = [event async for event in  model.astream_events("hello", version="v2")]
last_event = events[-1]
ai_message = last_event['data']['output']
ai_message.usage_metadata

If it's not in that field, but it's in the raw response (LLMResult), you can always attach a callback handler and extract the information from model response.