@@ -18,6 +18,7 @@ type Accumulator struct {
1818 streamAccumulators sync.Map // Track accumulators by request ID (atomic)
1919
2020 chatStreamChunkPool sync.Pool // Pool for reusing StreamChunk structs
21+ responsesStreamChunkPool sync.Pool // Pool for reusing ResponsesStreamChunk structs
2122 audioStreamChunkPool sync.Pool // Pool for reusing AudioStreamChunk structs
2223 transcriptionStreamChunkPool sync.Pool // Pool for reusing TranscriptionStreamChunk structs
2324
@@ -80,14 +81,32 @@ func (a *Accumulator) putTranscriptionStreamChunk(chunk *TranscriptionStreamChun
8081 a .transcriptionStreamChunkPool .Put (chunk )
8182}
8283
84+ // getResponsesStreamChunk gets a responses stream chunk from the pool
85+ func (a * Accumulator ) getResponsesStreamChunk () * ResponsesStreamChunk {
86+ return a .responsesStreamChunkPool .Get ().(* ResponsesStreamChunk )
87+ }
88+
89+ // putResponsesStreamChunk returns a responses stream chunk to the pool
90+ func (a * Accumulator ) putResponsesStreamChunk (chunk * ResponsesStreamChunk ) {
91+ chunk .Timestamp = time.Time {}
92+ chunk .StreamResponse = nil
93+ chunk .Cost = nil
94+ chunk .SemanticCacheDebug = nil
95+ chunk .ErrorDetails = nil
96+ chunk .FinishReason = nil
97+ chunk .TokenUsage = nil
98+ a .responsesStreamChunkPool .Put (chunk )
99+ }
100+
83101// CreateStreamAccumulator creates a new stream accumulator for a request
84102func (a * Accumulator ) createStreamAccumulator (requestID string ) * StreamAccumulator {
85103 sc := & StreamAccumulator {
86- RequestID : requestID ,
87- ChatStreamChunks : make ([]* ChatStreamChunk , 0 ),
88- IsComplete : false ,
89- Timestamp : time .Now (),
90- Object : "" ,
104+ RequestID : requestID ,
105+ ChatStreamChunks : make ([]* ChatStreamChunk , 0 ),
106+ ResponsesStreamChunks : make ([]* ResponsesStreamChunk , 0 ),
107+ IsComplete : false ,
108+ Timestamp : time .Now (),
109+ Object : "" ,
91110 }
92111 a .streamAccumulators .Store (requestID , sc )
93112 return sc
@@ -174,6 +193,30 @@ func (a *Accumulator) addAudioStreamChunk(requestID string, chunk *AudioStreamCh
174193 return nil
175194}
176195
196+ // addResponsesStreamChunk adds a responses stream chunk to the stream accumulator
197+ func (a * Accumulator ) addResponsesStreamChunk (requestID string , chunk * ResponsesStreamChunk , object string , isFinalChunk bool ) error {
198+ accumulator := a .getOrCreateStreamAccumulator (requestID )
199+ // Lock the accumulator
200+ accumulator .mu .Lock ()
201+ defer accumulator .mu .Unlock ()
202+ if accumulator .StartTimestamp .IsZero () {
203+ accumulator .StartTimestamp = chunk .Timestamp
204+ }
205+ // Store object type once (from first chunk)
206+ if accumulator .Object == "" && object != "" {
207+ accumulator .Object = object
208+ }
209+ // Add chunk to the list (chunks arrive in order)
210+ accumulator .ResponsesStreamChunks = append (accumulator .ResponsesStreamChunks , chunk )
211+ // Check if this is the final chunk
212+ // Set FinalTimestamp when either FinishReason is present or token usage exists
213+ // This handles both normal completion chunks and usage-only last chunks
214+ if isFinalChunk {
215+ accumulator .FinalTimestamp = chunk .Timestamp
216+ }
217+ return nil
218+ }
219+
177220// cleanupStreamAccumulator removes the stream accumulator for a request
178221func (a * Accumulator ) cleanupStreamAccumulator (requestID string ) {
179222 if accumulator , exists := a .streamAccumulators .Load (requestID ); exists {
@@ -182,6 +225,9 @@ func (a *Accumulator) cleanupStreamAccumulator(requestID string) {
182225 for _ , chunk := range acc .ChatStreamChunks {
183226 a .putChatStreamChunk (chunk )
184227 }
228+ for _ , chunk := range acc .ResponsesStreamChunks {
229+ a .putResponsesStreamChunk (chunk )
230+ }
185231 for _ , chunk := range acc .AudioStreamChunks {
186232 a .putAudioStreamChunk (chunk )
187233 }
@@ -263,7 +309,7 @@ func (a *Accumulator) appendContentToMessage(message *schemas.ChatMessage, newCo
263309}
264310
265311// ProcessStreamingResponse processes a streaming response
266- // It handles both audio and chat streaming responses
312+ // It handles chat, audio, and responses streaming responses
267313func (a * Accumulator ) ProcessStreamingResponse (ctx * context.Context , result * schemas.BifrostResponse , bifrostErr * schemas.BifrostError ) (* ProcessedStreamResponse , error ) {
268314 // Check if this is a streaming response
269315 if result == nil {
@@ -272,6 +318,8 @@ func (a *Accumulator) ProcessStreamingResponse(ctx *context.Context, result *sch
272318 requestType := result .ExtraFields .RequestType
273319 isAudioStreaming := requestType == schemas .SpeechStreamRequest || requestType == schemas .TranscriptionStreamRequest
274320 isChatStreaming := requestType == schemas .ChatCompletionStreamRequest || requestType == schemas .TextCompletionStreamRequest
321+ isResponsesStreaming := requestType == schemas .ResponsesStreamRequest
322+
275323 if isChatStreaming {
276324 // Handle text-based streaming with ordered accumulation
277325 return a .processChatStreamingResponse (ctx , result , bifrostErr )
@@ -283,6 +331,9 @@ func (a *Accumulator) ProcessStreamingResponse(ctx *context.Context, result *sch
283331 if requestType == schemas .SpeechStreamRequest {
284332 return a .processAudioStreamingResponse (ctx , result , bifrostErr )
285333 }
334+ } else if isResponsesStreaming {
335+ // Handle responses streaming with responses accumulation
336+ return a .processResponsesStreamingResponse (ctx , result , bifrostErr )
286337 }
287338 return nil , fmt .Errorf ("request type missing/invalid for accumulator" )
288339}
@@ -295,6 +346,9 @@ func (a *Accumulator) Cleanup() {
295346 for _ , chunk := range accumulator .ChatStreamChunks {
296347 a .chatStreamChunkPool .Put (chunk )
297348 }
349+ for _ , chunk := range accumulator .ResponsesStreamChunks {
350+ a .responsesStreamChunkPool .Put (chunk )
351+ }
298352 for _ , chunk := range accumulator .TranscriptionStreamChunks {
299353 a .transcriptionStreamChunkPool .Put (chunk )
300354 }
@@ -360,6 +414,11 @@ func NewAccumulator(pricingManager *pricing.PricingManager, logger schemas.Logge
360414 return & ChatStreamChunk {}
361415 },
362416 },
417+ responsesStreamChunkPool : sync.Pool {
418+ New : func () any {
419+ return & ResponsesStreamChunk {}
420+ },
421+ },
363422 audioStreamChunkPool : sync.Pool {
364423 New : func () any {
365424 return & AudioStreamChunk {}
@@ -381,6 +440,7 @@ func NewAccumulator(pricingManager *pricing.PricingManager, logger schemas.Logge
381440 // Prewarm the pools for better performance at startup
382441 for range 1000 {
383442 a .chatStreamChunkPool .Put (& ChatStreamChunk {})
443+ a .responsesStreamChunkPool .Put (& ResponsesStreamChunk {})
384444 a .audioStreamChunkPool .Put (& AudioStreamChunk {})
385445 a .transcriptionStreamChunkPool .Put (& TranscriptionStreamChunk {})
386446 }
0 commit comments