feat(stream): add method for filtering reasoning chunks from AI streaming#724
feat(stream): add method for filtering reasoning chunks from AI streaming#724adarshba wants to merge 1 commit into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughBaseProvider now exposes its StreamHandler to subclasses; StreamHandler gains createFilteredFullStream to emit only text content from full/text streams; GoogleVertexProvider uses the new filtered full-stream API; tests added for filtering, errors, and fallback behavior. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant GoogleVertexProvider as Provider
participant StreamHandler
participant FullStreamSource as FullStream
participant TextStreamSource as TextStream
Client->>Provider: executeStream(request)
Provider->>StreamHandler: createFilteredFullStream({ fullStream?, textStream? })
alt fullStream available
StreamHandler->>FullStream: iterate chunks
else fallback to textStream
StreamHandler->>TextStream: iterate strings
end
loop per chunk
FullStream-->>StreamHandler: chunk (text-delta | string | reasoning | error | tool | source)
alt text-delta or string
StreamHandler-->>Provider: yield { content: "..." } %%#{"fill":"#E6F4EA"}
else non-text reasoning/tool/source
StreamHandler-->>StreamHandler: skip chunk
else error chunk
StreamHandler-->>Provider: throw StreamingError %%#{"fill":"#FFEFE6"}
end
end
Provider-->>Client: streamed content events
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
c55b290 to
dfb0f18
Compare
There was a problem hiding this comment.
Pull request overview
This PR adds a new streaming method to filter out reasoning/thinking chunks from AI SDK's fullStream, preventing chain-of-thought (CoT) tokens from leaking to end users when using models with extended thinking capabilities like Gemini's experimental reasoning models.
Key Changes:
- Added
createFilteredFullStream()method to StreamHandler that filters reasoning-related chunk types - Modified Google Vertex provider to use the new filtered stream method instead of
createTextStream() - Changed streamHandler visibility from private to protected in BaseProvider to allow provider access
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| src/lib/core/modules/StreamHandler.ts | Added new createFilteredFullStream() method to filter reasoning/thinking chunks from fullStream while preserving text-delta chunks |
| src/lib/providers/googleVertex.ts | Updated streaming implementation to use createFilteredFullStream() instead of createTextStream() for reasoning model support |
| src/lib/core/baseProvider.ts | Changed streamHandler from private to protected visibility to enable provider-level access |
| pre-commit.sh | Commented out npm test command in pre-commit hook |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
3faa759 to
1c9f27d
Compare
|
@coderabbitai fullreview |
🧩 Analysis chain🏁 Script executed: Length of output: 37752 🏁 Script executed: Length of output: 6045 🏁 Script executed: Length of output: 67367 🏁 Script executed: Length of output: 967 🏁 Script executed: Length of output: 43 🏁 Script executed: Length of output: 15693 🏁 Script executed: Length of output: 10188 🏁 Script executed: Length of output: 435 🏁 Script executed: Length of output: 128 🏁 Script executed: Length of output: 2263 @adarshba I've completed a comprehensive review of PR #724. The implementation is solid and well-tested, but I've identified some consistency and documentation concerns that should be addressed. ✅ Strengths
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/lib/core/modules/StreamHandler.ts (1)
90-97: Enhance JSDoc with usage guidance and filtered chunk types.The implementation is solid and past review concerns have been addressed (error handling, type safety, test coverage). However, the JSDoc could be more helpful by clarifying:
- When to use this method vs
createTextStream- Complete list of filtered chunk types
- Example use case
📝 Enhanced JSDoc suggestion
/** - * Create filtered stream from AI SDK fullStream + * Create filtered stream from AI SDK fullStream or textStream * - * Filters out reasoning/thinking chunks from the stream. - * Used as safety net to prevent chain-of-thought tokens from leaking to users. + * Filters out internal/reasoning chunks from AI SDK's fullStream to prevent + * chain-of-thought tokens from leaking to end users. Falls back to textStream + * when fullStream is unavailable. * - * @param result - AI SDK stream result + * **When to use:** + * - Use this for reasoning-capable models (e.g., o1, Gemini 2.0 Flash Thinking) + * - Use `createTextStream` for simple text-only streaming without reasoning chunks + * + * **Filtered chunk types:** + * - reasoning, reasoning-signature, redacted-reasoning + * - source, tool-call, tool-result + * + * **Passed-through chunk types:** + * - text-delta (yielded as content) + * - string chunks (yielded as content) + * - error chunks (thrown as Error) + * + * @param result - AI SDK stream result with fullStream and/or textStream + * @returns AsyncGenerator yielding { content: string } chunks + * @throws Error when error chunk is encountered in stream + * + * @example + * ```ts + * const stream = this.streamHandler.createFilteredFullStream(result); + * for await (const chunk of stream) { + * console.log(chunk.content); // Only user-visible text content + * } + * ``` */
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/lib/core/baseProvider.tssrc/lib/core/modules/StreamHandler.tssrc/lib/providers/googleVertex.tstest/unit/stream/reasoning-filter.test.ts
🧰 Additional context used
📓 Path-based instructions (2)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.{ts,tsx}: Maintain strict TypeScript type safety across all modules with comprehensive type definitions organized by domain to avoid circular dependencies
Use ErrorFactory for creating typed errors throughout the application
Wrap async operations with withTimeout utility for timeout handling
Files:
src/lib/core/baseProvider.tssrc/lib/providers/googleVertex.tssrc/lib/core/modules/StreamHandler.tstest/unit/stream/reasoning-filter.test.ts
**/providers/*.ts
📄 CodeRabbit inference engine (CLAUDE.md)
Providers must extend a base provider or implement the provider interface and register in ProviderRegistry.registerAllProviders() with provider name, factory function, default model, and aliases
Files:
src/lib/providers/googleVertex.ts
🧠 Learnings (5)
📓 Common learnings
Learnt from: vigneshJuspay
Repo: juspay/neurolink PR: 0
File: :0-0
Timestamp: 2025-12-15T18:35:37.783Z
Learning: In juspay/neurolink TTS implementation (PR #691), the new StreamChunk discriminated union type introduced in TTS-019 will be integrated with StreamResult.stream during the actual TTS streaming implementation PR (TTS-020/TTS-021), not in the type-definition PR. This phased approach keeps type updates and implementation changes atomic.
📚 Learning: 2025-09-17T17:55:15.261Z
Learnt from: RajuSudhar
Repo: juspay/neurolink PR: 173
File: src/lib/index.ts:16-16
Timestamp: 2025-09-17T17:55:15.261Z
Learning: In src/lib/types/providers.ts, ProviderConfig was renamed to AIModelProviderConfig to deduplicate type names, as there was an existing ProviderConfig type that better suited the "ProviderConfig" name. This was an intentional breaking change for better type organization.
Applied to files:
src/lib/core/baseProvider.ts
📚 Learning: 2025-12-10T12:24:51.147Z
Learnt from: CR
Repo: juspay/neurolink PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-10T12:24:51.147Z
Learning: Applies to **/types/index.ts : Add new provider names to the AIProviderName enum in src/lib/types/index.ts when adding a new provider
Applied to files:
src/lib/core/baseProvider.ts
📚 Learning: 2025-12-10T12:24:51.147Z
Learnt from: CR
Repo: juspay/neurolink PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-10T12:24:51.147Z
Learning: Applies to **/utils/messageBuilder.ts : Message construction must be handled through MessageBuilder in src/lib/utils/messageBuilder.ts, which handles text, images, PDFs, and CSV files with provider-specific adapters
Applied to files:
src/lib/core/baseProvider.ts
📚 Learning: 2025-09-01T22:58:39.149Z
Learnt from: sudharsan-juspay
Repo: juspay/neurolink PR: 140
File: src/lib/core/types.ts:198-203
Timestamp: 2025-09-01T22:58:39.149Z
Learning: In src/lib/core/types.ts, StreamOptions (imported from streamTypes.js) and StreamingOptions are intentionally different types with different use cases. StreamingOptions is for unified AI requests with multiple provider configurations, while StreamOptions is for individual streaming operations.
Applied to files:
src/lib/core/modules/StreamHandler.ts
🧬 Code graph analysis (2)
src/lib/core/baseProvider.ts (1)
src/lib/core/modules/StreamHandler.ts (1)
StreamHandler(32-200)
test/unit/stream/reasoning-filter.test.ts (1)
src/lib/core/modules/StreamHandler.ts (1)
StreamHandler(32-200)
🔇 Additional comments (4)
src/lib/core/baseProvider.ts (1)
64-64: LGTM: Safe visibility change to enable subclass access.Changing
streamHandlerfrom private to protected is a sound design decision that allows provider subclasses to access StreamHandler methods (like the newcreateFilteredFullStream) while maintaining encapsulation. This is a non-breaking change that aligns with the composition pattern used throughout BaseProvider.test/unit/stream/reasoning-filter.test.ts (1)
1-317: Excellent test coverage addressing prior review feedback.This comprehensive test suite thoroughly validates the
createFilteredFullStreammethod:
- ✅ All filtered chunk types (reasoning, reasoning-signature, redacted-reasoning, source, tool-call, tool-result)
- ✅ Pass-through chunks (text-delta, strings)
- ✅ Error handling (error chunks with and without messages)
- ✅ Fallback behavior (textStream when fullStream unavailable)
- ✅ Edge cases (empty streams, mixed chunk types)
The test structure is clear, mocks are appropriate, and assertions are precise. This directly addresses the prior review comment about lacking test coverage.
src/lib/core/modules/StreamHandler.ts (1)
105-133: Implementation is correct and addresses all prior review feedback.The filtering logic properly:
- ✅ Handles error chunks by throwing (addresses prior comment)
- ✅ Uses type definitions for safety (addresses prior comment)
- ✅ Filters reasoning/internal chunks to prevent CoT token leakage
- ✅ Passes through user-visible content (text-delta, strings)
- ✅ Falls back to textStream when fullStream unavailable
- ✅ Has comprehensive test coverage (in reasoning-filter.test.ts)
The implementation is production-ready.
src/lib/providers/googleVertex.ts (1)
986-987: LGTM: Correct implementation of filtered streaming for Gemini adaptive thinking models.The switch to
this.streamHandler.createFilteredFullStream(result)correctly implements reasoning-chunk filtering to prevent chain-of-thought token leakage. This is appropriate for googleVertex.ts since it supports Gemini models with adaptive thinking capabilities.Note on other providers: While googleAiStudio can also use Gemini 3 Pro (which has adaptive thinking), it currently relies on
createTextStreamrather thancreateFilteredFullStream. Whether this requires updating depends on whether googleAiStudio'sstreamText()result receives afullStreamcontaining reasoning chunks. If googleAiStudio only receives atextStream, no filtering would be needed at the stream handler level. A brief audit of googleAiStudio's actual stream result structure would clarify whether similar filtering is necessary.
1c9f27d to
cefe8cb
Compare
| * } | ||
| * ``` | ||
| */ | ||
| createFilteredFullStream(result: { |
There was a problem hiding this comment.
@adarshba Will this stop emitting events?
How can you make sure all the events are sent as much as possible because the events are required?
There was a problem hiding this comment.
the method only filters the text stream, tool events are captured via onStepFinish callback so tool events remain untouched.
|
@coderabbitai fullreview Also check if the comments I have added are valid. |
🧩 Analysis chain🏁 Script executed: Length of output: 2160 🏁 Script executed: Length of output: 4967 🏁 Script executed: Length of output: 8179 🏁 Script executed: Length of output: 5794 @murdore, I've analyzed your comment and can confirm that your concern about events not being emitted is not valid. Here's why: Analysis of Event Emission vs. Chunk FilteringThe key insight is that event handlers and chunk filtering operate at different layers: 1. Event Handlers Are Registered at AI SDK Level (Line 905-959)const result = await streamText({
model: model,
messages: messages,
experimental_telemetry: this.getStreamTelemetryConfig(options),
onError: (event) => { /* ... */ }, // Line 919 - AI SDK calls this
onFinish: (event) => { /* ... */ }, // Line 931 - AI SDK calls this
onChunk: () => { chunkCount++; }, // Line 942 - AI SDK calls this
onStepFinish: ({ toolCalls, toolResults }) => { /* ... */ } // Line 946
});These event handlers are passed as configuration to 2. Chunk Filtering Happens Downstream (Line 963)const transformedStream = this.streamHandler.createFilteredFullStream(result);The 3. Why Events Are Not AffectedThe event handlers ( 4. Evidence from the CodeThe
All of these will continue to fire normally because they're part of the AI SDK's internal processing, not the consumer-facing stream transformation. ConclusionYour concern about events being stopped is not applicable in this implementation. The events are emitted at the AI SDK level and will fire for all chunks (including reasoning chunks), while The filtering is purely cosmetic from an event perspective - all internal telemetry, error handling, and callbacks remain fully functional. Now proceeding with the full review you requested. 🧠 Learnings used✅ Actions performedFull review triggered. |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/lib/core/modules/StreamHandler.ts (1)
134-157: Type guards ensure safe chunk processing.The inline type guards (
hasType,isTextDelta,isError) andgetErrorMessageutility provide proper type narrowing and safe error handling, addressing the past review concern about unsafe type casts.Optional: Consider reusing
getErrorMessageandisErrorfromsrc/lib/types/common.ts(lines 91-93, 110-121) for consistency across the codebase, though the current inline implementation is also acceptable.🔎 Optional refactor to use existing utilities
+import { getErrorMessage as getCommonErrorMessage, isError as isCommonError } from "../../types/common.js"; + /** * StreamHandler class - Handles streaming operations for AI providers */ export class StreamHandler { ... createFilteredFullStream(result: { fullStream?: AsyncIterable<unknown>; textStream?: AsyncIterable<string>; }): AsyncGenerator<{ content: string }> { // Type guards const hasType = (val: unknown): val is { type: string } => val !== null && typeof val === "object" && "type" in val && typeof (val as Record<string, unknown>).type === "string"; const isTextDelta = (val: { type: string; }): val is { type: "text-delta"; textDelta: string } => val.type === "text-delta" && "textDelta" in val && typeof (val as Record<string, unknown>).textDelta === "string"; const isError = (val: { type: string; }): val is { type: "error"; error: unknown } => val.type === "error" && "error" in val; - - const getErrorMessage = (error: unknown): string => { - if (error !== null && typeof error === "object" && "message" in error) { - const msg = (error as Record<string, unknown>).message; - return typeof msg === "string" ? msg : "Unknown error"; - } - return "Unknown error"; - }; return (async function* () { if (result.fullStream) { for await (const chunk of result.fullStream) { if (typeof chunk === "string") { yield { content: chunk }; } else if (hasType(chunk)) { if (isError(chunk)) { throw new Error( - `Streaming error: ${getErrorMessage(chunk.error)}`, + `Streaming error: ${getCommonErrorMessage(chunk.error)}`, ); } else if (isTextDelta(chunk)) { yield { content: chunk.textDelta }; } } } } else if (result.textStream) { for await (const chunk of result.textStream) { yield { content: chunk }; } } })(); }
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/lib/core/baseProvider.tssrc/lib/core/modules/StreamHandler.tssrc/lib/providers/googleVertex.tstest/unit/stream/reasoning-filter.test.ts
🧰 Additional context used
📓 Path-based instructions (2)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.{ts,tsx}: Maintain strict TypeScript type safety across all modules with comprehensive type definitions organized by domain to avoid circular dependencies
Use ErrorFactory for creating typed errors throughout the application
Wrap async operations with withTimeout utility for timeout handling
Files:
test/unit/stream/reasoning-filter.test.tssrc/lib/core/baseProvider.tssrc/lib/core/modules/StreamHandler.tssrc/lib/providers/googleVertex.ts
**/providers/*.ts
📄 CodeRabbit inference engine (CLAUDE.md)
Providers must extend a base provider or implement the provider interface and register in ProviderRegistry.registerAllProviders() with provider name, factory function, default model, and aliases
Files:
src/lib/providers/googleVertex.ts
🧠 Learnings (4)
📓 Common learnings
Learnt from: vigneshJuspay
Repo: juspay/neurolink PR: 0
File: :0-0
Timestamp: 2025-12-15T18:35:37.783Z
Learning: In juspay/neurolink TTS implementation (PR #691), the new StreamChunk discriminated union type introduced in TTS-019 will be integrated with StreamResult.stream during the actual TTS streaming implementation PR (TTS-020/TTS-021), not in the type-definition PR. This phased approach keeps type updates and implementation changes atomic.
📚 Learning: 2025-09-17T17:55:15.261Z
Learnt from: RajuSudhar
Repo: juspay/neurolink PR: 173
File: src/lib/index.ts:16-16
Timestamp: 2025-09-17T17:55:15.261Z
Learning: In src/lib/types/providers.ts, ProviderConfig was renamed to AIModelProviderConfig to deduplicate type names, as there was an existing ProviderConfig type that better suited the "ProviderConfig" name. This was an intentional breaking change for better type organization.
Applied to files:
src/lib/core/baseProvider.tssrc/lib/providers/googleVertex.ts
📚 Learning: 2025-12-10T12:24:51.147Z
Learnt from: CR
Repo: juspay/neurolink PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-10T12:24:51.147Z
Learning: Applies to **/utils/messageBuilder.ts : Message construction must be handled through MessageBuilder in src/lib/utils/messageBuilder.ts, which handles text, images, PDFs, and CSV files with provider-specific adapters
Applied to files:
src/lib/core/baseProvider.ts
📚 Learning: 2025-09-01T22:58:39.149Z
Learnt from: sudharsan-juspay
Repo: juspay/neurolink PR: 140
File: src/lib/core/types.ts:198-203
Timestamp: 2025-09-01T22:58:39.149Z
Learning: In src/lib/core/types.ts, StreamOptions (imported from streamTypes.js) and StreamingOptions are intentionally different types with different use cases. StreamingOptions is for unified AI requests with multiple provider configurations, while StreamOptions is for individual streaming operations.
Applied to files:
src/lib/core/modules/StreamHandler.tssrc/lib/providers/googleVertex.ts
🧬 Code graph analysis (3)
test/unit/stream/reasoning-filter.test.ts (1)
src/lib/core/modules/StreamHandler.ts (1)
StreamHandler(32-248)
src/lib/core/baseProvider.ts (1)
src/lib/core/modules/StreamHandler.ts (1)
StreamHandler(32-248)
src/lib/core/modules/StreamHandler.ts (1)
src/lib/types/common.ts (2)
isError(92-94)getErrorMessage(111-122)
🔇 Additional comments (9)
src/lib/core/baseProvider.ts (1)
64-64: LGTM! Clean encapsulation improvement.Changing
streamHandlerfromprivatetoprotectedappropriately allows subclasses like GoogleVertexProvider to access the newcreateFilteredFullStreammethod while maintaining encapsulation. Thereadonlymodifier ensures subclasses cannot reassign the handler.src/lib/providers/googleVertex.ts (1)
986-987: Integration looks good, but consider googleAiStudio consistency.The switch to
createFilteredFullStreamcorrectly filters reasoning chunks for GoogleVertex. Based on PR objectives, googleAiStudio.ts still usescreateTextStreamand should also be updated to usecreateFilteredFullStreamfor consistency, as both use Google's reasoning-capable models.Per PR objectives summary: "Inconsistent provider adoption (High): GoogleVertex was updated, but googleAiStudio.ts still uses createTextStream. Recommendation: update googleAiStudio to use this.streamHandler.createFilteredFullStream(result) for consistency."
test/unit/stream/reasoning-filter.test.ts (5)
25-151: Excellent coverage of filtered chunk types.All six filtered chunk types (
reasoning,reasoning-signature,redacted-reasoning,source,tool-call,tool-result) are thoroughly tested with clear assertions verifying they're excluded whiletext-deltachunks pass through.
153-197: Pass-through behavior well tested.Tests validate both typed
text-deltachunks and plain string chunks pass through correctly, ensuring the filter doesn't drop valid content.
199-231: Error handling properly validated.Tests cover both error chunks with specific messages and unknown errors, verifying the stream throws appropriately. This addresses the past review concern about error chunk handling.
233-285: Fallback and edge cases covered.Tests validate the fallback to
textStreamwhenfullStreamis unavailable and handle empty stream scenarios gracefully, ensuring robust behavior.
287-336: Comprehensive test suite with clear documentation.The mixed chunk test validates realistic scenarios where filtered and valid chunks are interleaved. The test coverage summary at the end clearly documents all validated behaviors, making the test file easy to maintain.
src/lib/core/modules/StreamHandler.ts (2)
90-129: Excellent documentation addressing key concerns.The JSDoc comprehensively explains:
- Purpose: prevent chain-of-thought token leakage
- When to use: Google Vertex with reasoning-capable models
- Filtered vs. passed-through chunk types
- Fallback behavior to textStream
- Practical example
This addresses the PR objective about missing documentation and the past review comment about expanding "CoT" abbreviation.
159-181: Clean filtering logic with proper fallback behavior.The async generator correctly:
- Prioritizes
fullStreamovertextStream- Filters reasoning/thinking chunks while passing through text-delta
- Throws on error chunks to surface failures
- Falls back gracefully to
textStreamwhen neededRegarding event emission (addressing murdore's question): Event hooks (
onChunk,onFinish,onError,onStepFinish) are registered in the provider'sstreamText()call (see googleVertex.ts lines 919-964) and are invoked by the AI SDK before chunks reach this filter. This filter only affects what the consumer receives from the stream, not internal SDK event emission. Per PR objectives, this was validated by the reviewer.
cefe8cb to
55597b0
Compare
Pull Request
Description
Add createFilteredFullStream() to StreamHandler that filters out
reasoning/thinking chunks (reasoning, reasoning-signature, redacted-reasoning)
from AI SDK fullStream. This prevents chain-of-thought tokens from leaking
to users when using models with extended thinking capabilities.
Type of Change
Related Issues
Changes Made
AI Provider Impact
Component Impact
Testing
Test Environment
Performance Impact
Breaking Changes
Screenshots/Demo
Simulator.Screen.Recording.-.iPhone.17.Pro.-.2025-12-26.at.18.37.38.mov
Checklist
Additional Notes
Use this method:
createFilteredFullStreamfor reasoning models in the providersSummary by CodeRabbit
New Features
Improvements
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.