Add ability to stream custom events from a tool

[daniilr]

Description

For some specific long-running tool runs it is useful to update a user with with current status. Right now tools can only either return complete result or throw error/retry.

Example use case: backup tool can yield status updates: {'status': 'Generating backup', 'progress': 0.2}, {'status': 'Uploading to s3', 'progress': 0.5} before it returns final result (link to s3).

I don't have a good idea of how to achieve that and I also don't like the way it is handled in LangChain.

References

LangChain

https://docs.langchain.com/oss/python/langgraph/streaming#tool
In LangChain get_stream_writer is accessible via tools:

from langchain.tools import tool
from langgraph.config import get_stream_writer

@tool
def query_database(query: str) -> str:
    """Query the database."""
    # Access the stream writer to send custom data
    writer = get_stream_writer()  
    # Emit a custom key-value pair (e.g., progress update)
    writer({"data": "Retrieved 0/100 records", "type": "progress"})  
    # perform query
    # Emit another custom key-value pair
    writer({"data": "Retrieved 100/100 records", "type": "progress"})
    return "some-answer"

graph = ... # define a graph that uses this tool

# Set stream_mode="custom" to receive the custom data in the stream
for chunk in graph.stream(inputs, stream_mode="custom"):
    print(chunk)

I am not sure if this is a good design.

Vercel AI

https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#preliminary-tool-results

In Vercel AI tools can be async generators. The last thing yielded is returned to LLM:

tool({
  description: 'Get the current weather.',
  inputSchema: z.object({
    location: z.string(),
  }),
  async *execute({ location }) {
    yield {
      status: 'loading' as const,
      text: `Getting weather for ${location}`,
      weather: undefined,
    };

    await new Promise(resolve => setTimeout(resolve, 3000));

    const temperature = 72 + Math.floor(Math.random() * 21) - 10;

    yield {
      status: 'success' as const,
      text: `The weather in ${location} is ${temperature}°F`,
      temperature,
    };
  },
});

[DouweM]

Duplicate of #2382, already being implemented in #3114, should land later this month :)

[daniilr]

Thanks @DouweM, this already looks like a good design!