Checked other resources

• This is a feature request, not a bug report or usage question.
• I added a clear and descriptive title that summarizes the feature request.
• I used the GitHub search to find a similar feature request and didn't find it.
• I checked the LangChain documentation and API reference to see if this feature already exists.
• This is not related to the langchain-community package.

Package (Required)

• langchain
• langchain-openai
• langchain-anthropic
• langchain-classic
• langchain-core
• langchain-model-profiles
• langchain-tests
• langchain-text-splitters
• langchain-chroma
• langchain-deepseek
• langchain-exa
• langchain-fireworks
• langchain-groq
• langchain-huggingface
• langchain-mistralai
• langchain-nomic
• langchain-ollama
• langchain-perplexity
• langchain-qdrant
• langchain-xai
• Other / not sure / general

Feature Description

When using astream with updates streaming mode, I would like to be able to identify the final answer / response from the system to stream to the client of the agent streaming API endpoint.

Code:

async for stream_mode, data in self._agent.astream( # The output of each step in the graph. The output shape depends on the stream_mode.
    {"messages": [{"role": "user", "content": message}]},
    stream_mode = modes, # Use this to stream all values in the state after each step.
    config = config, # This is needed by Checkpointer
    subgraphs = subgraphs
):
    for source, update in data.items():
        if source == "model":
            await output_queue.put(update)

I use RunnableConfig["callbacks"] to stream tokens:

class TokenQueueStreamingHandler(AsyncCallbackHandler):
    """LangChain callback handler for streaming LLM tokens to an asyncio queue."""
    def __init__(self, queue: Queue):
        self.queue = queue

    async def on_llm_new_token(self, token: str, **kwargs) -> None:
        if token:
            await self.queue.put(token)

The kwargs has the following structure and value:

# kwargs: {'chunk': ChatGenerationChunk(text='We', message=AIMessageChunk(content='We', additional_kwargs={}, response_metadata={}, tool_calls=[], invalid_tool_calls=[], tool_call_chunks=[])), 
#          'run_id': UUID('019c028e-19a4-7213-92d0-26535292aaca'), 
#          'parent_run_id': UUID('019c028e-1983-7f42-8648-8f02f450018c'), 
#          'tags': ['seq:step:1'], 
#          'verbose': False}

When I use the values streaming mode, I can save the messages into a list/array and then at the end of the async, I return the last message which is the final answer to user's query. However, streaming the LLM tokens is a challenge for this use case as there is no way to tell the context of the tokens spit out by the LLM and which LLM.

Observations:

1. additional_kwargs={}, response_metadata={}, tool_calls=[], invalid_tool_calls=[], tool_call_chunks=[] these values are not populated in the kwargs at all.
2. tags': ['seq:step:1'] always carries the same value.

Use Case

Use RunnableConfig["callbacks"] to selectively stream tokens only of the final AIMessage to the user request.

Proposed Solution

For inspiration, ollama chat has a thinking field in the stream which is used to show or hide reasoning traces before the final answer arrives. https://docs.ollama.com/capabilities/streaming

Alternatives Considered

No response

Additional Context

No response