feat(iorails): Add streaming usage - #2198
Conversation
github-actions
Bot
added
status: needs triage
tgasser-nv
added
status: triaged
1e1fa8a to
5686b6e
Compare
|
@coderabbitai Review this PR |
|
@greptile-apps Review this PR |
|
✅ Action performedReview finished.
|
📝 WalkthroughWalkthroughStreaming now preserves provider-specific chunk fields, HTTP response headers, and usage metadata through model parsing and IORails streaming. Metadata is attached to streamed deltas, while usage-only metadata is folded into the terminal frame. ChangesStreaming Metadata
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Upstream
participant ModelEngine
participant IORails
participant StreamingHandler
Upstream->>ModelEngine: SSE chunks and response headers
ModelEngine->>IORails: LLMResponseChunk with provider_metadata
IORails->>StreamingHandler: Text delta with metadata
IORails->>StreamingHandler: END_OF_STREAM with pending usage
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemoguardrails/guardrails/iorails.py`:
- Around line 158-162: Update the streamed metadata construction in
nemoguardrails/guardrails/iorails.py lines 158-162 to include non-None
reasoning_tokens and cached_tokens alongside the existing usage fields,
preserving them on both content and terminal frames. Add assertions for these
optional fields in tests/guardrails/test_iorails_streaming.py lines 965-983 and
lines 1072-1084, covering the content-bearing and terminal frames respectively.
In `@nemoguardrails/guardrails/model_engine.py`:
- Around line 213-215: Restrict streamed provider metadata to explicit safe
allowlists: update the chunk metadata construction in
nemoguardrails/guardrails/model_engine.py:213-215 to retain only approved
non-standard fields, filter response headers at model_engine.py:705-707 to
approved diagnostics, and merge only sanitized metadata at
model_engine.py:777-780; update nemoguardrails/guardrails/iorails.py:155-156 to
avoid forwarding raw metadata. Add assertions in
tests/guardrails/test_model_engine.py:840-886 and
tests/guardrails/test_iorails_streaming.py:1018-1035 confirming sensitive
fields, credentials, Set-Cookie, and excluded headers never reach public
metadata or IORails frames.
- Around line 56-59: Remove all newly added explanatory comments and docstrings
at the listed sites: model_engine.py lines 56-59, 89-90, 212, 679-680, 702-704,
and 774-775; test_model_engine.py lines 62-66, 801, 821, 841, and 864;
iorails.py lines 147-152, 1384-1387, and 1390-1400; and
test_iorails_streaming.py lines 530-534, 938, 942, 963, 987, 1005, 1019,
1039-1044, and 1070. Preserve the surrounding implementation and existing
documentation, removing only the additions identified by the review.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4bf0aad0-e544-4137-912c-38aa7e464a60
📒 Files selected for processing (4)
nemoguardrails/guardrails/iorails.pynemoguardrails/guardrails/model_engine.pytests/guardrails/test_iorails_streaming.pytests/guardrails/test_model_engine.py
Greptile SummaryThis PR adds streaming usage and provider metadata (response headers and non-standard body fields such as
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/iorails.py | Adds _stream_chunk_metadata helper and threads usage/provider metadata through the main-LLM streaming loop; correctly folds usage-only terminal chunks into the END_OF_STREAM frame. |
| nemoguardrails/guardrails/model_engine.py | Adds _STANDARD_CHUNK_KEYS constant, extracts non-standard body fields as provider_metadata in _parse_chat_completion_chunk, attaches lowercased response headers to every yielded chunk, and defaults include_usage=True in stream_call. |
| tests/guardrails/test_iorails_streaming.py | Adds TestStreamAsyncMetadata with 6 focused test cases covering all new metadata-surfacing behaviours; helper _build_sse_streaming_mock extended with headers parameter. |
| tests/guardrails/test_model_engine.py | Adds 4 new stream_call tests: include_usage default, caller-provided stream_options override, response headers in provider_metadata, and non-standard body key (nvext) surfacing. |
Sequence Diagram
sequenceDiagram
participant Client
participant IORails
participant StreamingHandler
participant ModelEngine
Client->>IORails: "stream_async(include_metadata=True)"
IORails->>ModelEngine: "stream_call(stream_options={include_usage:True})"
loop SSE chunks
ModelEngine-->>IORails: LLMResponseChunk(delta_content, provider_metadata, usage?)
alt chunk has delta_content
IORails->>IORails: _stream_chunk_metadata(chunk) → chunk_metadata
IORails->>StreamingHandler: push_chunk(delta_content, chunk_metadata)
StreamingHandler-->>Client: "{"text": "...", "metadata": {provider_metadata, usage?}}"
else chunk has usage only (no delta_content)
IORails->>IORails: "pending_usage_metadata = _stream_chunk_metadata(chunk)"
else reasoning-only / empty chunk
Note over IORails: provider_metadata dropped (no frame emitted)
end
end
IORails->>StreamingHandler: push_chunk(END_OF_STREAM, pending_usage_metadata)
StreamingHandler-->>Client: "{"text": "", "metadata": {provider_metadata, usage, response_metadata, usage_metadata}}"
Reviews (6): Last reviewed commit: "Trim out obvious comments" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR threads streaming token usage and provider metadata (HTTP response headers, non-standard body fields like
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/model_engine.py | Adds per-stream response-header capture, non-standard body-key extraction into provider_metadata, and default include_usage=True; logic is sound with one minor note on the shared response_headers dict reference. |
| nemoguardrails/guardrails/iorails.py | New _stream_chunk_metadata() helper and updated streaming loop correctly route per-chunk metadata and defer terminal usage to the END_OF_STREAM frame; reasoning_tokens/cached_tokens are silently dropped from serialized usage. |
| tests/guardrails/test_iorails_streaming.py | Adds six new test cases in TestStreamAsyncMetadata covering usage-only, content+usage, provider_metadata, plain-string passthrough, response-header, and no-empty-frame behaviours; coverage is thorough. |
| tests/guardrails/test_model_engine.py | Adds four unit tests for stream_call covering default include_usage injection, caller override, response-header attachment, and non-standard body-key surfacing; all straightforward and correct. |
Sequence Diagram
sequenceDiagram
participant Caller
participant IORails
participant StreamingHandler
participant ModelEngine
Caller->>IORails: "stream_async(messages, include_metadata=True)"
IORails->>ModelEngine: "stream_model_call(main, messages, stream_options={include_usage:True})"
loop SSE chunks
ModelEngine-->>ModelEngine: parse chunk body → LLMResponseChunk
ModelEngine-->>ModelEngine: merge response_headers into provider_metadata
ModelEngine-->>IORails: yield LLMResponseChunk
alt chunk has delta_content
IORails->>StreamingHandler: "push_chunk(text, {provider_metadata, [usage]})"
StreamingHandler-->>Caller: "{text, metadata: {provider_metadata, [usage]}}"
else chunk has usage only (terminal)
IORails-->>IORails: store as pending_usage_metadata
end
end
IORails->>StreamingHandler: push_chunk(END_OF_STREAM, pending_usage_metadata)
StreamingHandler-->>Caller: "{text empty, metadata: {usage, provider_metadata, response_metadata, usage_metadata}}"
Prompt To Fix All With AI
Fix the following 2 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 2
nemoguardrails/guardrails/iorails.py:157-162
`reasoning_tokens` and `cached_tokens` are silently dropped when serializing usage into the streaming metadata dict. Both fields are populated by `_parse_usage` for OpenAI-compatible reasoning models (via `completion_tokens_details` / `prompt_tokens_details`), so callers that rely on them for cost accounting will see zeros in the streaming path even when the non-streaming `UsageInfo` carries real values.
```suggestion
if chunk.usage:
usage_dict = {
"input_tokens": chunk.usage.input_tokens,
"output_tokens": chunk.usage.output_tokens,
"total_tokens": chunk.usage.total_tokens,
}
if chunk.usage.reasoning_tokens is not None:
usage_dict["reasoning_tokens"] = chunk.usage.reasoning_tokens
if chunk.usage.cached_tokens is not None:
usage_dict["cached_tokens"] = chunk.usage.cached_tokens
metadata["usage"] = usage_dict
```
### Issue 2 of 2
nemoguardrails/guardrails/model_engine.py:679
The comment has "parity" duplicated — "for parity with LLMRails parity" — worth cleaning up before merge.
```suggestion
# Request token usage on the terminal stream chunk for parity with LLMRails.
```
Reviews (2): Last reviewed commit: "Combine last content-chunk and usage-onl..." | Re-trigger Greptile
|
@greptile-apps There are two summary reports with different scores on this PR: #2198 (comment) and #2198 (comment) . Create a new summary/score comment based on latest PR and I'll hide the other two as outdated. |
2 participants
Description
Stacked PR: This is stacked on top of #2178, which added GenerationOptions and GenerationResponse to the non-streaming IORails path. This achieved parity with LLMRails (minus the Colang-specific fields). This is already merged to
develop.This small PR connects the streaming usage chunks through to the Guardrails client.
Related Issue(s)
Verification
Pre-commit
Unit-test
Integration-test with LLMRails vs IORails comparison (using compare_stream_metadata.py)
iorails_stream_frames.json
llmrails_stream_frames.json
Integration test with Chat
AI Assistance
Checklist
Summary by CodeRabbit
New Features
Bug Fixes