Feat: Output Rails Streaming

[Pouyanpi]

Description

Example:

from nemoguardrails import RailsConfig, LLMRails

config = RailsConfig.from_path("./examples/configs/abc_streaming")

rails = LLMRails(config)

messages = [{"role": "user", "content": "what can you do"}]

async for chunk in rails.stream_async(messages=messages):
    print(f"CHUNK: {chunk}")

or using chat cli

nemoguardrails chat --config="./examples/configs/abc_streaming"

Todo

• finalize streaming config parameter names
• test functioanlity on microservice ✅
• docs: output rails streaming #976 by @mikemckiernan
• SSE
• add config example with Nemoguard NIMs?
• test multi output rails with Nemoguard and self check output

[Pouyanpi]

requires #965

[github-actions]

Documentation preview

https://nvidia.github.io/NeMo-Guardrails/review/pr-966

[mikemckiernan]

Unsolicited feedback. LMK if you'd prefer that I not provide unsolicited feedback.

[Pouyanpi]

Thank you @mikemckiernan 👍🏻 just don't review guardrails/library changes here, their PR is #965 .

I'll push changes wrt review soon, thx!

[cparisien]

This is looking good to me. Let's make sure we are clear on the "hierarchy" of the "streaming: true" configs and the default behaviour.
Also I'd like to make sure I have a clear understanding of which rails will support this behaviour now, and which won't.

[cparisien]

Elaborate a bit on the description. e.g., how the look back size is how many tokens/characters in total are sent to the rail each time as context. Also give the units -- is this counted in characters?

[Pouyanpi]

size is always number of tokens (words).

I need to change the param names and also default value. These were used for testing

window_size is number of tokens that we apply the output rails on, so when the streaming starts it corresponds to first #window_size tokens, then in the next iteration(s) we use #look_back_size tokens from the previous iteration(s) but still respect the window_size.

[Pouyanpi]

@trebedea need your expert intuition here, what do you think is a good default here (ignore the defaults I've set)

[Pouyanpi]

Basically, we have 2 parameters:

1. The size of the token group processed at one time: This defines how many tokens are in each block that the output rails operate on.
2. The amount of previous tokens reused in the next iteration: This defines the carryover or context from the previous block that is blended into the current processing batch.

decision:
chunk_size instead of window_size
context_size instead of look_back_size

chunk_size and context_size are both clear and abstract, communicating exactly what they do: one determines the size of the token block and the other determines how much of the previouss block is reused.

[Pouyanpi]

fyr, @mikemckiernan @cparisien

[mikemckiernan]

I like those two field names--and I'm extremely picky. FYI: https://nvidia.github.io/NeMo-Guardrails/review/pr-976/user-guides/configuration-guide.html#enabling-streaming-output

[cparisien]

Thanks I'm good with these new names too.

[cparisien]

Not clear what this means?

[Pouyanpi]

stream chunk of tokens first, then apply output rails on them (pros: lower perceived TTFT)

or

apply the output rails on that chunk and then stream (pros: if the chunk is blocked then we don't stream blocked content)

is it clear?

[cparisien]

Thanks that's clear, I expect we'll learn from early users how they experience it.

[Pouyanpi]

Also I'd like to make sure I have a clear understanding of which rails will support this behaviour now, and which won't.

• All the output rails that are defined in the library
• any user defined rail that uses @action

[Pouyanpi]

Some thoughts on default values:

chuk_size: 256
context_size: 64

A chunk of about 200–256 tokens tends to be large enough to encapsulate a meaningful “unit” of generated text. This provides the rails with enough information? At the same time, it isn’t so large that increases TTFT latency of first (if stream_first: False) and subsequent iterations.

Summary

• Chunk Size (200 tokens):
  • Pros: Captures a complete, meaningful unit of text for effective rail evaluation.
  • Cons: May introduce slight delay if waiting for the full chunk (when stream_first: False).
• Context Size (50 tokens):
  • Pros: Ensures continuity across chunks with minimal impact on latency.
  • Cons: Too small a context might miss cross-chunk violations, but 25% overlap is generally a good compromise.

Input Length	Chunk Size	Context Size	Rails Invocation
512	256	64	3
600	256	64	3
256	256	64	1
1024	256	64	5
1024	256	32	5
1024	256	32	5
1024	128	32	11
512	128	32	5

cc @cparisien

[Pouyanpi]

• remove abc_streaming example from configs

[Pouyanpi]

tested ✅

 ...
 rails:
  input:
    flows:
      - topic safety check input $model=topic_control
  output:
    flows:
      - content safety check output $model=content_safety
      - self check output
    chunk_size: 10
  dialog:
    single_call:
      enabled: False

streaming: True

[Pouyanpi]

in chat cli we have following behavior:

stream_first: True

stream_first: False

for the stream_async interface

stream_first: True

async for chunk in rails.stream_async(messages=messages):
   print(f"CHUNK:{chunk}")

CHUNK:I
....
CHUNK: Is
CHUNK: there
CHUNK: anything
CHUNK: specific
CHUNK: you
CHUNK: would
CHUNK: like
CHUNK: to
CHUNK: know?
CHUNK:{"event": "ABORT", "data": {"reason": "Blocked by self check output rails."}}

stream_first: False

async for chunk in rails.stream_async(messages=messages):
   print(f"CHUNK:{chunk}")
  
CHUNK:{"event": "ABORT", "data": {"reason": "Blocked by self check output rails."}}

**The SSE is meant for the microservice. **

what do you think? @cparisien @mikemckiernan

[Pouyanpi]

in chat cli we have following behavior:

stream_first: True

stream_first: False

Here I can remove the output rails name if it is client facing and they should not be aware. Or we can have a -debug mode that also includes the output rails name.

[mikemckiernan]

CHUNK:{"event": "ABORT", "data": {"reason": "Blocked by self check output rails."}}

**The SSE is meant for the microservice. **

what do you think? @cparisien @mikemckiernan

You read my mind--I meant to ask the question "How does the client know if the partial response was unsafe and how does the client check?" I have no experience in this area, but the concept of an event has appeal. (We can cover this offline...but does the client check each chunk if it's a dict or str? If dict, it's bad news?) No strong feeling on making the rail name optional. Couldn't the client filter it? Would they want to log it (maybe logged elsewhere...)?

[Pouyanpi]

Nice!

(We can cover this offline...but does the client check each chunk if it's a dict or str? If dict, it's bad news?)

it is a json str so not a dict but can be converted to one if needed.

No strong feeling on making the rail name optional. Couldn't the client filter it? Would they want to log it (maybe logged elsewhere...)?

Agreed! We include the name 👍🏻

@mikemckiernan