Streaming Responses

If the application LLM supports streaming, the NeMo Guardrails library can stream tokens as well. Streaming is automatically enabled when you use the stream_async() method - no configuration is required.

For information about configuring streaming with output guardrails, refer to the following:

• For configuration, refer to Output Rail Streaming.
• For sample Python client code, refer to Tutorials.

Usage

Chat CLI

You can enable streaming when launching the NeMo Guardrails library chat CLI by using the --streaming option:

$
nemoguardrails chat --config=examples/configs/streaming --streaming

Python API

You can use the streaming directly from the python API in two ways:

1. Simple: receive just the chunks (tokens).
2. Full: receive both the chunks as they are generated and the full response at the end.

For the simple usage, you need to call the stream_async method on the LLMRails instance:

1
from nemoguardrails import LLMRails
2
3
app = LLMRails(config)
4
5
history = [{"role": "user", "content": "What is the capital of France?"}]
6
7
async for chunk in app.stream_async(messages=history):
8
    print(f"CHUNK: {chunk}")
9
    # Or do something else with the token

For the full usage, you need to provide a StreamingHandler instance to the generate_async method on the LLMRails instance:

1
from nemoguardrails import LLMRails
2
from nemoguardrails.streaming import StreamingHandler
3
4
app = LLMRails(config)
5
6
history = [{"role": "user", "content": "What is the capital of France?"}]
7
8
streaming_handler = StreamingHandler()
9
10
async def process_tokens():
11
    async for chunk in streaming_handler:
12
        print(f"CHUNK: {chunk}")
13
        # Or do something else with the token
14
15
asyncio.create_task(process_tokens())
16
17
result = await app.generate_async(
18
    messages=history, streaming_handler=streaming_handler
19
)
20
print(result)

Warning: Using StreamingHandler directly is deprecated and will be removed in a future release. Use stream_async() instead.

Using External Async Token Generators

You can also provide your own async generator that yields tokens, which is useful when:

• You want to use a different LLM provider that has its own streaming API.
• You have pre-generated responses that you want to stream through guardrails.
• You want to implement custom token generation logic.
• You want to test your output rails or its config in streaming mode on predefined responses without actually relying on an actual LLM generation.

To use an external generator, pass it to the generator parameter of stream_async:

1
from nemoguardrails import LLMRails
2
from typing import AsyncIterator
3
4
app = LLMRails(config)
5
6
async def my_token_generator() -> AsyncIterator[str]:
7
    # This could be from OpenAI API, Anthropic API, or any other LLM API that already has a streaming token generator. Mocking the stream here, for a simple example.
8
    tokens = ["Hello", " ", "world", "!"]
9
    for token in tokens:
10
        yield token
11
12
messages = [{"role": "user", "content": "The most famous program ever written is"}]
13
14
# use the external generator with guardrails
15
async for chunk in app.stream_async(
16
    messages=messages,
17
    generator=my_token_generator()
18
):
19
    print(f"CHUNK: {chunk}")

When using an external generator:

• The internal LLM generation is completely bypassed.
• Output rails are still applied to the LLM responses returned by the external generator, if configured.
• The generator should yield string tokens.

Example with a real LLM API:

1
async def openai_streaming_generator(messages) -> AsyncIterator[str]:
2
    """Example using OpenAI's streaming API."""
3
    import openai
4
5
    stream = await openai.ChatCompletion.create(
6
        model="gpt-4o",
7
        messages=messages,
8
        stream=True
9
    )
10
11
    # Yield tokens as they arrive
12
    async for chunk in stream:
13
        if chunk.choices[0].delta.content:
14
            yield chunk.choices[0].delta.content
15
16
config = RailsConfig.from_path("config/with_output_rails")
17
app = LLMRails(config)
18
19
async for chunk in app.stream_async(
20
    messages=[{"role": "user", "content": "Tell me a story"}],
21
    generator=openai_streaming_generator(messages)
22
):
23
    # output rails will be applied to these chunks
24
    print(chunk, end="", flush=True)

This feature enables seamless integration of the NeMo Guardrails library with any streaming LLM or token source while maintaining all the safety features of output rails.

Streaming Metadata

When using stream_async(), you can receive per-chunk metadata (e.g., token usage, finish reason) by setting include_metadata=True:

1
async for chunk in rails.stream_async(messages=messages, include_metadata=True):
2
    print(chunk)

With include_metadata=True, each chunk is a dict with a mandatory "text" key. The final chunk also includes a "metadata" key containing response_metadata (finish reason, model name) and usage_metadata (token counts):

1
{"text": "Hello"}
2
{"text": "!"}
3
{"text": "", "metadata": {
4
    "response_metadata": {"finish_reason": "stop", "model_name": "gpt-4o"},
5
    "usage_metadata": {"input_tokens": 75, "output_tokens": 9, "total_tokens": 84}
6
}}

Without include_metadata, chunks are plain strings (default behavior).

Token Usage Tracking

Token usage statistics are available when streaming responses, depending on provider support. When the provider does not return token usage statistics, the final chunk’s metadata will contain response_metadata and usage_metadata set to None.

Accessing Token Usage Information

You can access token usage statistics through the detailed logging capabilities of the NeMo Guardrails library. Use the log generation option to capture comprehensive information about LLM calls, including token usage:

1
response = rails.generate(messages=messages, options={
2
    "log": {
3
        "llm_calls": True,
4
        "activated_rails": True
5
    }
6
})
7
8
for llm_call in response.log.llm_calls:
9
    print(f"Task: {llm_call.task}")
10
    print(f"Total tokens: {llm_call.total_tokens}")
11
    print(f"Prompt tokens: {llm_call.prompt_tokens}")
12
    print(f"Completion tokens: {llm_call.completion_tokens}")

Alternatively, you can use the explain() method to get a summary of token usage:

1
info = rails.explain()
2
info.print_llm_calls_summary()

For more information about streaming token usage support across different providers, refer to the LangChain documentation on token usage tracking. For detailed information about accessing generation logs and token usage, see Generation Options: Detailed Logging Information and Logging.

Streaming for LLMs Deployed Using HuggingFacePipeline

We also support streaming for LLMs deployed using HuggingFacePipeline. One example is provided in the HF Pipeline Dolly configuration.

To use streaming for HF Pipeline LLMs, you need to create an nemoguardrails.integrations.langchain.providers.huggingface.AsyncTextIteratorStreamer streamer object, add it to the kwargs of the pipeline and to the model_kwargs of the HuggingFacePipelineCompatible object.

1
from nemoguardrails.integrations.langchain.providers.huggingface import AsyncTextIteratorStreamer
2
3
# instantiate tokenizer object required by LLM
4
streamer = AsyncTextIteratorStreamer(tokenizer, skip_prompt=True)
5
params = {"temperature": 0.01, "max_new_tokens": 100, "streamer": streamer}
6
7
pipe = pipeline(
8
    # all other parameters
9
    **params,
10
)
11
12
llm = HuggingFacePipelineCompatible(pipeline=pipe, model_kwargs=params)