How to Implement Token-by-Token Streaming with LangChain (Complete Guide)

How to Implement Token-by-Token Streaming with LangChain (Complete Guide)

Imagine you are asking a smart AI a question. Instead of waiting a long time for the full answer to appear all at once, wouldn’t it be much nicer to see the AI typing its response word by word, just like a person? This is exactly what token-by-token streaming helps you achieve. It makes your applications feel faster and much more interactive for anyone using them.

This guide will show you how to implement token-by-token streaming with LangChain. We will explore different methods, from simple built-in tools to creating your own advanced solutions. You will learn how to make your AI interactions smooth and engaging.

Understanding Token Streaming: What Is It?

Token streaming means getting the AI’s answer piece by piece, as soon as each part is ready. Instead of waiting for the entire message, you receive tiny chunks of text. These chunks are called “tokens.”

A token can be a single word, a part of a word, or even punctuation. When you see text appearing gradually, you are seeing token-by-token streaming in action. This approach significantly improves how users experience your AI applications.

Why Token-by-Token Streaming is Super Important

Imagine chatting with a friend who pauses for a long time, then suddenly sends a huge message. It can feel a bit slow and unresponsive. Now, think about a friend who types and sends messages line by line as they think.

This second way feels much more natural and engaging. Token-by-token streaming with LangChain offers this same benefit for your AI tools. It makes the AI seem like it’s thinking and responding in real-time.

Better User Experience

Seeing immediate output keeps users engaged. They don’t have to stare at a blank screen, wondering if the AI is working. This instant feedback creates a much more satisfying interaction.

Perceived Speed

Even if the total time to get the full answer is the same, receiving it gradually feels faster. Users appreciate not having to wait for a big block of text to appear at once. It helps manage expectations during longer generation times.

Real-time Interactions

For applications like chatbots, code generators, or content creation tools, real-time output is key. It allows users to react to the AI’s response as it’s being generated. They can even stop the generation if they see it going in the wrong direction.

Getting Started with LangChain

Before we dive into streaming, let’s make sure you have LangChain installed. You can easily install it using a simple command. This will set up all the necessary parts for our examples.

First, open your terminal or command prompt. Then, type the following command to get LangChain ready on your computer.

pip install langchain langchain-community langchain-openai

We’ll also need an OpenAI API key for our examples. Make sure you have one ready and set it as an environment variable. This allows LangChain to connect to the powerful AI models.

import os
os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY" # Replace with your actual key

LangChain’s Built-in Streaming: StreamingStdOutCallbackHandler

LangChain provides a very easy way to start with token-by-token streaming. It’s called StreamingStdOutCallbackHandler. This special tool takes each token as it arrives and immediately prints it to your console. It’s a fantastic way to see streaming in action without much effort.

Using this handler is quite simple. You just need to tell your LangChain model to use it when you make a call. Let’s look at how you can set this up.

from langchain_openai import ChatOpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

# Set up our chat model
llm = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,  # This is important: tells the model to stream!
    callbacks=[StreamingStdOutCallbackHandler()] # Add our streaming handler
)

# Ask a question and see the tokens stream to your console
print("--- Streaming with StreamingStdOutCallbackHandler ---")
llm.invoke("Write a short poem about a friendly robot helping a lost cat.")

When you run this code, you will see the poem appear word by word in your terminal. Each piece of text prints as soon as it’s generated by the AI model. This demonstrates the core idea of StreamingStdOutCallbackHandler usage very clearly.

This handler is great for quick tests and seeing the streaming effect. However, it’s quite basic; it only prints to the console. For more control, like displaying tokens in a web browser or saving them, we’ll need something more advanced.

The Need for Custom Callback Handlers

While StreamingStdOutCallbackHandler is helpful, it has limitations. It only prints to your terminal, which isn’t useful for many real-world applications. Imagine wanting to show the streaming text on a website or save it to a file. The built-in handler can’t do that.

This is where custom callback handlers come into play. LangChain lets you create your own special tools to manage how tokens are handled. You can decide what happens with each token as it arrives.

Custom handlers give you complete control over the token processing logic. You can store tokens in a list, send them over a network connection, or update a user interface. This flexibility is crucial for building interactive and dynamic AI applications.

You can learn more about the different types of callbacks in LangChain by visiting our guide on Understanding LangChain Callbacks. It provides a deep dive into the callback system.

Building Your First Custom Callback Handler (Synchronous)

Let’s build a simple custom callback handler. This handler will collect all the tokens in a list and then print the full message at the end. It’s a good first step to understand how they work.

We’ll create a class that inherits from LangChain’s BaseCallbackHandler. This class will have special methods that LangChain calls at different points during the AI’s thought process. The key method for streaming is on_llm_new_token.

from langchain.callbacks.base import BaseCallbackHandler
from langchain_openai import ChatOpenAI

class MyCustomSyncHandler(BaseCallbackHandler):
    def __init__(self):
        self.tokens = []

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        """Runs on new LLM token. Only available when streaming is enabled."""
        self.tokens.append(token)
        print(f"Received token: '{token}'") # We'll still print to see it happening

    def on_llm_end(self, response, **kwargs) -> None:
        """Runs when LLM ends running."""
        full_response = "".join(self.tokens)
        print(f"\n--- Full response gathered by custom handler: ---\n{full_response}")
        self.tokens = [] # Clear for next use

# Set up our chat model with the custom handler
llm = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[MyCustomSyncHandler()]
)

print("\n--- Streaming with MyCustomSyncHandler ---")
llm.invoke("Explain the concept of photosynthesis in a simple way.")

In this example, on_llm_new_token gets called every time a new token is generated. We add it to our tokens list. The on_llm_end method is called once the AI has finished its entire response. Here, we join all the collected tokens to form the complete message and print it. This simple handler allows you to capture and process each token.

Handling Asynchronous Streaming with AsyncCallbackHandler

Many modern applications are built to be asynchronous, meaning they can do many things at once without waiting. For these applications, we need an AsyncCallbackHandler implementation. LangChain provides AsyncCallbackHandler for this exact purpose. It allows your streaming logic to run without blocking the rest of your program.

This is crucial for web servers or user interfaces that need to remain responsive while waiting for AI responses. Instead of on_llm_new_token, you’ll use on_llm_new_token but with the async keyword. All the callback methods in your async handler should be async def.

Let’s create an asynchronous version of our custom handler. This version will still collect tokens, but it uses asynchronous methods. This prepares us for more complex async applications, like those built with FastAPI or websockets.

from langchain.callbacks.base import AsyncCallbackHandler
from langchain_openai import ChatOpenAI
import asyncio

class MyCustomAsyncHandler(AsyncCallbackHandler):
    def __init__(self):
        self.tokens = []

    async def on_llm_new_token(self, token: str, **kwargs) -> None:
        """Async runs on new LLM token."""
        self.tokens.append(token)
        print(f"Async Received token: '{token}'")
        # Imagine sending this token to a websocket or updating a UI asynchronously
        await asyncio.sleep(0.01) # Simulate some async work

    async def on_llm_end(self, response, **kwargs) -> None:
        """Async runs when LLM ends running."""
        full_response = "".join(self.tokens)
        print(f"\n--- Full response gathered by custom async handler: ---\n{full_response}")
        self.tokens = [] # Clear for next use

# Set up our chat model with the async custom handler
llm_async = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[MyCustomAsyncHandler()]
)

async def main_async_streaming():
    print("\n--- Streaming with MyCustomAsyncHandler ---")
    await llm_async.ainvoke("Write a very short, humorous haiku about a sleepy cat.")

# Run the async function
asyncio.run(main_async_streaming())

Notice the use of await llm_async.ainvoke() instead of llm.invoke(). This is how you call asynchronous LangChain methods. The asyncio.run() function is used to execute the main asynchronous function. This setup ensures that your application stays non-blocking while processing tokens.

You can even combine synchronous and asynchronous handlers if needed. LangChain is flexible in how you stack your callback logic. Just ensure your main invocation matches the type of handler (sync with invoke, async with ainvoke).

Deep Dive: Token Processing Logic

When you receive tokens from the AI, they aren’t always perfect words. Sometimes, a token might be part of a word, or it might contain special characters. Understanding this token processing logic is key to displaying smooth and correct output.

Your custom handler needs to be smart about how it manages these incoming tokens. You might want to collect several tokens before displaying them to form complete words or sentences. This is where buffering comes in handy.

Buffering Strategies

Buffering strategies involve collecting multiple tokens before showing them to the user. Why do this? Imagine if an AI sends “hello” as two tokens: “he” and “llo”. If you display “he” then “llo”, it looks fine. But what if it sends “app” then “le”? Displaying “app” then “le” might not look as natural as waiting for “apple”.

By buffering, you can gather tokens until you have a complete word, a punctuation mark, or a certain number of characters. Then, you release the buffered text. This makes the output appear more natural and less choppy.

Here’s a simple buffering example:

from langchain.callbacks.base import BaseCallbackHandler
from langchain_openai import ChatOpenAI
import time

class BufferedCallbackHandler(BaseCallbackHandler):
    def __init__(self, buffer_size=10):
        self.buffer = []
        self.buffer_size = buffer_size

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.buffer.append(token)
        if len(self.buffer) >= self.buffer_size or token.endswith(('.', '!', '?', '\n')):
            self._flush_buffer()

    def _flush_buffer(self):
        if self.buffer:
            text_to_print = "".join(self.buffer)
            print(text_to_print, end="", flush=True)
            self.buffer = []

    def on_llm_end(self, response, **kwargs) -> None:
        self._flush_buffer() # Ensure any remaining tokens are printed
        print("\n--- End of buffered stream ---")

llm_buffered = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[BufferedCallbackHandler(buffer_size=5)] # Flush every 5 tokens or on punctuation
)

print("\n--- Streaming with BufferedCallbackHandler ---")
llm_buffered.invoke("Tell me a brief story about a detective cat solving the mystery of the missing fish.")

In this BufferedCallbackHandler, we print tokens when the buffer reaches a certain size or when a punctuation mark is seen. This improves readability. The flush=True in print() is important to make sure the text appears immediately.

Handling Partial Tokens

Sometimes, especially with complex languages or when dealing with character encodings, a token might not be a complete readable character. For example, a multi-byte character (like some emojis or non-English letters) might be split across two tokens. This is rare with common models like OpenAI’s gpt-3.5-turbo for English, but it’s good to be aware of.

Most modern AI models and LangChain handle basic handling partial tokens quite well for common use cases. If you encounter strange characters or broken output, it might be due to encoding issues.

The common way to manage this is to append tokens to a buffer and only decode/display them when a full, valid character sequence is formed. For most cases, simply joining the str tokens works fine because the models are usually trained to output valid Unicode sequences.

If you were building a very low-level system, you might need to manage bytes and decode them carefully. For typical LangChain applications, simply joining the string tokens is sufficient. The example above handles most cases effectively by just appending and joining string tokens.

Advanced Streaming Techniques

Once you’ve mastered the basics, you can explore more advanced ways to enhance your token-by-token streaming langchain setup. These techniques offer more control and richer interactions.

Streaming Metadata

Sometimes, you might want to send more than just the token text during streaming. Streaming metadata refers to extra information that comes along with each token or at different stages of the generation. This could include things like the current word count, sentiment of the text so far, or special instructions for the UI.

While LangChain’s on_llm_new_token primarily provides the token string, you can enrich your callback handler to infer or attach metadata. For example, you could maintain a running count of words.

Here’s an idea of how you might include some simple metadata logic:

import json

class MetadataStreamingHandler(AsyncCallbackHandler):
    def __init__(self):
        self.full_response = ""
        self.word_count = 0

    async def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.full_response += token
        self.word_count = len(self.full_response.split()) # Simple word count
        
        metadata = {
            "token": token,
            "current_word_count": self.word_count,
            "timestamp": time.time()
        }
        # In a real app, you'd send this JSON over a websocket
        print(f"Metadata stream: {json.dumps(metadata)}")
        await asyncio.sleep(0.01) # Simulate async sending

    async def on_llm_end(self, response, **kwargs) -> None:
        print(f"\n--- Final Word Count: {self.word_count} ---")
        self.full_response = ""
        self.word_count = 0

llm_metadata = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[MetadataStreamingHandler()]
)

async def main_metadata_streaming():
    print("\n--- Streaming with MetadataStreamingHandler ---")
    await llm_metadata.ainvoke("Describe the process of a seed growing into a tree, focusing on key stages.")

asyncio.run(main_metadata_streaming())

This handler doesn’t actually receive metadata from the LLM in on_llm_new_token (as the token argument is just the string). Instead, it generates metadata based on the incoming token. You can then use this metadata in your frontend application. For truly receiving metadata from the LLM, you might need specific LLM integrations that support it, or you can build more complex parsing logic into your handler.

Token Counting During Streaming

Keeping track of how many tokens have been generated is useful for several reasons. It helps you understand the length of the response in real-time, manage costs (as many models charge per token), and implement features like “X words generated so far.” This is token counting during streaming.

The easiest way to do token counting during streaming is within your custom callback handler. Every time on_llm_new_token is called, you simply increment a counter. LangChain models often return tokens that are roughly equivalent to words or sub-words, making this count quite representative.

class TokenCounterHandler(BaseCallbackHandler):
    def __init__(self):
        self.token_count = 0
        self.character_count = 0

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.token_count += 1
        self.character_count += len(token)
        print(f"Token: '{token}' (Tokens so far: {self.token_count}, Chars: {self.character_count})", end="\r", flush=True)

    def on_llm_end(self, response, **kwargs) -> None:
        print(f"\n--- Streaming finished. Total tokens: {self.token_count}, Total characters: {self.character_count} ---")
        self.token_count = 0
        self.character_count = 0

llm_token_count = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[TokenCounterHandler()]
)

print("\n--- Streaming with TokenCounterHandler ---")
llm_token_count.invoke("Write a creative description for a magical forest where trees whisper secrets.")

The \r character in the print statement makes the output overwrite the current line, creating a dynamic counter effect. This is a neat trick for terminal-based progress updates. You can easily adapt this to update a counter on a web page.

Throttling Token Output

Sometimes the AI generates tokens incredibly fast, perhaps too fast for your application to display gracefully, or for a user to read comfortably. Throttling token output means slowing down the rate at which tokens are processed or displayed. This can be useful for improving readability or managing system load.

You can implement throttling by adding a small delay in your on_llm_new_token method. For asynchronous handlers, you would use await asyncio.sleep(). For synchronous handlers, you’d use time.sleep(). Be careful not to make the delay too long, or it defeats the purpose of streaming.

import time

class ThrottlingHandler(AsyncCallbackHandler):
    def __init__(self, delay_per_token: float = 0.05):
        self.delay_per_token = delay_per_token
        self.buffer = []

    async def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.buffer.append(token)
        # Flush the buffer immediately but with a delay for each character
        for char in token:
            print(char, end="", flush=True)
            await asyncio.sleep(self.delay_per_token)

    async def on_llm_end(self, response, **kwargs) -> None:
        print("\n--- Throttled stream finished ---")
        self.buffer = []

llm_throttled = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[ThrottlingHandler(delay_per_token=0.08)] # 80ms delay per char
)

async def main_throttled_streaming():
    print("\n--- Streaming with ThrottlingHandler ---")
    await llm_throttled.ainvoke("Describe a vivid dream you had last night. Be detailed.")

asyncio.run(main_throttled_streaming())

This ThrottlingHandler shows characters one by one with a small delay, simulating a “typing” effect. You could also throttle by delaying after a certain number of tokens or complete words. This helps control the flow of information.

Putting It All Together: A Complete Streaming Example

Now let’s combine several of these concepts into one robust token-by-token streaming langchain example. We’ll create an asynchronous handler that buffers tokens, counts them, and then outputs them in a more controlled manner, ready for a web application.

This example will demonstrate:

• Asynchronous callback handling.
• Buffering tokens for more natural word-by-word output.
• Counting tokens.
• Simulating sending data to a client (e.g., via websockets).

from langchain.callbacks.base import AsyncCallbackHandler
from langchain_openai import ChatOpenAI
import asyncio
import json
import time

class CompleteStreamingHandler(AsyncCallbackHandler):
    def __init__(self):
        self.current_buffer = ""
        self.total_tokens = 0
        self.complete_response = ""
        self.messages_to_client = [] # Simulate messages to send to a client

    async def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.current_buffer += token
        self.total_tokens += 1
        self.complete_response += token

        # Logic to decide when to "send" a chunk to the client
        # Here, we send when we see a space, punctuation, or a new line,
        # or if the buffer gets too big.
        if " " in self.current_buffer or \
           any(p in self.current_buffer for p in ".,!?:;\n") or \
           len(self.current_buffer) > 10: # Max buffer size before forcing a flush
            await self._send_chunk_to_client()

        await asyncio.sleep(0.005) # Small non-blocking delay to simulate network latency

    async def _send_chunk_to_client(self):
        if not self.current_buffer:
            return

        # Prepare the data to send to the client (e.g., via websocket)
        message_data = {
            "type": "token_chunk",
            "content": self.current_buffer,
            "total_tokens": self.total_tokens,
            "timestamp": time.time()
        }
        self.messages_to_client.append(message_data)
        
        # In a real application, you would send this via websocket.
        # For this example, we'll just print it.
        print(f"[{self.total_tokens}] SENT: '{self.current_buffer}'", end="", flush=True)

        self.current_buffer = "" # Clear buffer after sending

    async def on_llm_end(self, response, **kwargs) -> None:
        # Send any remaining buffer content
        await self._send_chunk_to_client()
        
        final_message = {
            "type": "end_of_stream",
            "final_tokens": self.total_tokens,
            "full_response": self.complete_response,
            "timestamp": time.time()
        }
        self.messages_to_client.append(final_message)
        print(f"\n\n--- Stream Complete ---")
        print(f"Total tokens processed: {self.total_tokens}")
        print(f"Full response: {self.complete_response}")
        print(f"Simulated client messages: {json.dumps(self.messages_to_client, indent=2)}")

        # Reset for next invocation
        self.current_buffer = ""
        self.total_tokens = 0
        self.complete_response = ""
        self.messages_to_client = []


llm_complete_stream = ChatOpenAI(
    model="gpt-3.5-turbo",
    streaming=True,
    callbacks=[CompleteStreamingHandler()]
)

async def main_complete_streaming_example():
    print("\n--- Running Complete Streaming Example ---")
    await llm_complete_stream.ainvoke(
        "Write a detailed paragraph about the benefits of learning to code early in life. "
        "Include aspects like problem-solving skills, creativity, and future opportunities."
    )

# Run the complete example
asyncio.run(main_complete_streaming_example())

This CompleteStreamingHandler is a powerful template. It handles token processing logic to group tokens into meaningful chunks. It accurately performs token counting during streaming and stores streaming metadata (like total tokens and timestamp) along with each chunk. The messages_to_client list simulates sending these structured messages to a frontend, which could then display them dynamically. This is a practical example of how to implement token-by-token streaming langchain in a real-world scenario. You can imagine integrating this with a WebSocket server to push updates to a browser.

Common Challenges and Troubleshooting

Implementing token-by-token streaming can sometimes throw a few curveballs. Knowing what to look out for can save you a lot of time. Here are some common challenges you might face and how to troubleshoot them.

Not Seeing Any Streaming Output

• streaming=True Missing: This is the most common mistake. Ensure you explicitly set streaming=True when you initialize your ChatOpenAI or LLM model. Without it, the model won’t send tokens incrementally.
• No Callback Handler: If you don’t provide any callback handler, LangChain won’t know what to do with the streaming tokens. Even StreamingStdOutCallbackHandler will show you output.
• Asynchronous Context: If you’re using an AsyncCallbackHandler, make sure you’re calling the AI model with await llm.ainvoke() (or achat()/agenerate()). Using llm.invoke() with an async handler won’t work correctly.

Choppy or Inconsistent Output

• Buffering Issues: If tokens appear too rapidly or are cut off, revisit your buffering strategies. You might be flushing the buffer too often or not often enough. Try adjusting the buffer_size in your custom handler.
• Encoding Problems: While rare with modern LLMs and LangChain, ensure your environment and any custom decoding logic are handling UTF-8 correctly. Strange characters often point to encoding.

Performance Issues

• Excessive Logging/Processing: If your on_llm_new_token method does a lot of heavy processing, it can slow down the entire stream. Keep this method as lean as possible.
• Over-Throttling: If you implement throttling token output, ensure the delay isn’t too long. A very high time.sleep() or asyncio.sleep() value will make the streaming painfully slow.

Integrations with Frontend (WebSockets, SSE)

• Network Latency: Even if your backend is fast, network delays between your server and the client can affect perceived streaming speed. Optimize your network communication.
• Frontend Rendering: Ensure your frontend JavaScript is efficient at appending new chunks of text to the UI. Frequent DOM manipulations can be slow. Consider using virtual DOM libraries or batching updates.

By keeping these points in mind, you can effectively debug and optimize your token-by-token streaming langchain implementations. Always start with the simplest working example and build up complexity gradually.

Best Practices for Token-by-Token Streaming

To make your token-by-token streaming langchain robust and user-friendly, follow these best practices:

1. Always Enable Streaming: Explicitly set streaming=True in your LLM configuration. This is the foundation for all streaming features.
2. Use Asynchronous Handlers for Web Apps: If your application is a web server (like FastAPI, Flask with async) or a desktop UI, prefer AsyncCallbackHandler. This ensures your application stays responsive.
3. Implement Robust Buffering: Don’t just print raw tokens. Use buffering strategies to collect tokens until you have meaningful chunks. This could be full words, sentences (ending with punctuation), or based on a character limit.
4. Handle Edge Cases in on_llm_end: Always ensure that any remaining buffered tokens are flushed in your on_llm_end method. This guarantees no text is left un-displayed at the end of the response.
5. Keep Callback Logic Lean: The on_llm_new_token method is called frequently. Avoid heavy computations, complex database operations, or long-running tasks inside it. If you need complex processing, consider offloading it to a separate worker or message queue.
6. Provide Visual Feedback: In a UI, don’t just show the text appearing. Add a blinking cursor, a “typing…” indicator, or other subtle animations. This enhances the perceived responsiveness.
7. Implement Throttling (If Needed): If the AI generates text too fast for human readability, consider throttling token output. This can make the experience more natural, but balance it carefully to avoid making it feel slow.
8. Track Metrics: Use token counting during streaming to keep track of response length. This is useful for cost management, display limits, and analytical purposes.
9. Consider streaming metadata: If your frontend needs more than just text, think about what streaming metadata (like word count, status, or timestamps) you can generate and send along with the tokens.
10. Test Thoroughly: Test your streaming implementation with various prompts, especially long ones, to ensure it handles everything correctly, including partial tokens and buffer flushing.
11. Graceful Error Handling: Implement on_llm_error in your callback handler to catch and manage any issues that occur during LLM generation. This prevents your application from crashing.

By adhering to these practices, you can build powerful, efficient, and user-friendly token-by-token streaming langchain applications. Your users will appreciate the dynamic and interactive experience.

Conclusion

You’ve learned how to implement token-by-token streaming langchain from the ground up. We started with the simple StreamingStdOutCallbackHandler and moved on to creating your own powerful custom callback handlers. You now understand token processing logic, buffering strategies, and even handling partial tokens.

You also explored advanced topics like streaming metadata, token counting during streaming, and throttling token output. With this comprehensive guide, you are well-equipped to integrate dynamic, real-time AI responses into your applications. Your users will love the faster, more engaging experience that token-by-token streaming provides.

Start building your next interactive AI application today, knowing you can deliver responses just like a human is typing them!