feat(bedrock): support streaming (#128)

Author: Kludex

gateway/src/api/bedrock.ts

@@ -2,10 +2,10 @@
  * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html
  */
 
-import type { ConverseRequest, ConverseResponse } from '@aws-sdk/client-bedrock-runtime'
-import { BaseAPI } from './base'
+import type { ConverseRequest, ConverseResponse, ConverseStreamOutput } from '@aws-sdk/client-bedrock-runtime'
+import { BaseAPI, type ExtractedRequest, type ExtractedResponse, type ExtractorConfig } from './base'
 
-export class ConverseAPI extends BaseAPI<ConverseRequest, ConverseResponse> {
+export class ConverseAPI extends BaseAPI<ConverseRequest, ConverseResponse, ConverseStreamOutput> {
   defaultBaseUrl = 'https://bedrock-runtime.us-east-1.amazonaws.com'
 
   requestStopSequences = (requestBody: ConverseRequest): string[] | undefined => {
@@ -28,6 +28,38 @@ export class ConverseAPI extends BaseAPI<ConverseRequest, ConverseResponse> {
   responseId = (_responseBody: ConverseResponse): string | undefined => {
     return undefined
   }
+
+  // SafeExtractor implementation
+  requestExtractors: ExtractorConfig<ConverseRequest, ExtractedRequest> = {
+    requestModel: (requestBody: ConverseRequest) => {
+      this.extractedRequest.requestModel = requestBody.modelId
+    },
+    maxTokens: (requestBody: ConverseRequest) => {
+      this.extractedRequest.maxTokens = requestBody.inferenceConfig?.maxTokens
+    },
+    temperature: (requestBody: ConverseRequest) => {
+      this.extractedRequest.temperature = requestBody.inferenceConfig?.temperature
+    },
+    topP: (requestBody: ConverseRequest) => {
+      this.extractedRequest.topP = requestBody.inferenceConfig?.topP
+    },
+    stopSequences: (requestBody: ConverseRequest) => {
+      this.extractedRequest.stopSequences = requestBody.inferenceConfig?.stopSequences
+    },
+  }
+
+  responseExtractors: ExtractorConfig<ConverseResponse, ExtractedResponse> = {}
+
+  chunkExtractors: ExtractorConfig<ConverseStreamOutput, ExtractedResponse> = {
+    usage: (chunk: ConverseStreamOutput) => {
+      if ('usage' in chunk) {
+        this.extractedResponse.usage = this.extractUsage(chunk)
+      }
+    },
+    responseModel: (_chunk: ConverseStreamOutput) => {
+      this.extractedResponse.responseModel = this.requestModel
+    },
+  }
 }
 
 // TODO(Marcelo): Add input/output messages extraction.

gateway/src/providers/bedrock.ts

@@ -26,9 +26,9 @@ export class BedrockProvider extends DefaultProviderProxy {
     } catch (_error) {
       return { error: 'invalid request JSON' }
     }
-    const m = this.inferModel(this.restOfPath)
-    if (m) {
-      return { requestBodyText, requestBodyData, requestModel: m[1] }
+    const model = this.inferModel(this.restOfPath)
+    if (model) {
+      return { requestBodyText, requestBodyData, requestModel: model }
     }
     return { error: 'unable to find model in path' }
   }

gateway/src/providers/default.ts

@@ -6,6 +6,7 @@ import {
   type Provider as UsageProvider,
 } from '@pydantic/genai-prices'
 import * as logfire from '@pydantic/logfire-api'
+import { EventStreamCodec } from '@smithy/eventstream-codec'
 import { createParser, type EventSourceMessage } from 'eventsource-parser'
 
 import type { GatewayOptions } from '..'
@@ -385,12 +386,14 @@ export class DefaultProviderProxy {
 
     // IMPORTANT: Start consuming BOTH streams immediately to prevent tee() from buffering
     // The tee() requires both streams to be consumed concurrently, otherwise it will buffer
-
-    // Tee stream: one for client, one for processing
     const [responseStream, processingStream] = response.body.tee()
 
-    // Parse SSE events from processing stream
-    const events = this.parseSSE(processingStream)
+    let events: AsyncIterable<JsonData>
+    if (responseHeaders.get('content-type')?.toLowerCase().startsWith('application/vnd.amazon.eventstream')) {
+      events = this.parseAmazonEventStream(processingStream)
+    } else {
+      events = this.parseSSE(processingStream)
+    }
 
     // @ts-expect-error: TODO(Marcelo): Fix this type error.
     const extractionPromise = this.processChunks(modelAPI, events)
@@ -474,9 +477,42 @@ export class DefaultProviderProxy {
     }
   }
 
+  protected async *parseAmazonEventStream(stream: ReadableStream<Uint8Array>): AsyncIterable<JsonData> {
+    const encoder = new TextEncoder()
+    const codec = new EventStreamCodec((str) => str, encoder.encode)
+    const decoder = new TextDecoder()
+    let buffer = new Uint8Array(0)
+
+    for await (const chunk of stream) {
+      // Append incoming chunk to buffer since messages can span multiple network chunks
+      const combined = new Uint8Array(buffer.length + chunk.length)
+      combined.set(buffer, 0)
+      combined.set(chunk, buffer.length)
+      buffer = combined
+
+      // Extract complete messages from buffer (eventstream format: 4-byte length prefix + message data)
+      while (buffer.length >= 4) {
+        const messageLength = new DataView(buffer.buffer, buffer.byteOffset).getUint32(0, false)
+        if (buffer.length < messageLength) break
+
+        try {
+          const message = codec.decode(buffer.subarray(0, messageLength))
+          if (message.body?.length > 0) {
+            yield JSON.parse(decoder.decode(message.body))
+          }
+          buffer = buffer.subarray(messageLength)
+        } catch (error) {
+          logfire.reportError('Error parsing Amazon EventStream', error as Error)
+          break
+        }
+      }
+    }
+  }
+
   protected isStreaming(responseHeaders: Headers, requestBodyData: JsonData): boolean {
     return (
       responseHeaders.get('content-type')?.toLowerCase().startsWith('text/event-stream') ||
+      responseHeaders.get('content-type')?.toLowerCase().startsWith('application/vnd.amazon.eventstream') ||
       ('stream' in requestBodyData && requestBodyData.stream === true)
     )
   }

gateway/test/providers/bedrock.spec.ts

@@ -1,3 +1,4 @@
+import { EventStreamCodec } from '@smithy/eventstream-codec'
 import { describe, expect } from 'vitest'
 import { deserializeRequest } from '../otel'
 import { test } from '../setup'
@@ -24,4 +25,50 @@ describe('bedrock', () => {
     expect(otelBatch, 'otelBatch length not 1').toHaveLength(1)
     expect(deserializeRequest(otelBatch[0]!)).toMatchSnapshot('span')
   })
+
+  test('bedrock/stream', async ({ gateway }) => {
+    const { fetch, otelBatch } = gateway
+
+    const result = await fetch('https://example.com/converse/model/amazon.nova-micro-v1%3A0/converse-stream', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json', Authorization: 'healthy', 'x-vcr-filename': 'stream' },
+      body: JSON.stringify({
+        modelId: 'amazon.nova-premier-v1:0',
+        system: [{ text: 'You are a helpful assistant.' }],
+        messages: [{ role: 'user', content: [{ text: 'What is the capital of France?' }] }],
+      }),
+    })
+    const chunks: string[] = []
+    for await (const chunk of parseEventStream(result.body!)) {
+      chunks.push(chunk)
+    }
+    expect(chunks).toMatchSnapshot('chunks')
+    expect(otelBatch, 'otelBatch length not 1').toHaveLength(1)
+    expect(deserializeRequest(otelBatch[0]!)).toMatchSnapshot('span')
+  })
 })
+
+async function* parseEventStream(stream: ReadableStream<Uint8Array>): AsyncIterable<string> {
+  const encoder = new TextEncoder()
+  const codec = new EventStreamCodec((str) => str, encoder.encode)
+  const decoder = new TextDecoder()
+  let buffer = new Uint8Array(0)
+
+  for await (const chunk of stream) {
+    const combined = new Uint8Array(buffer.length + chunk.length)
+    combined.set(buffer, 0)
+    combined.set(chunk, buffer.length)
+    buffer = combined
+
+    while (buffer.length >= 4) {
+      const messageLength = new DataView(buffer.buffer, buffer.byteOffset).getUint32(0, false)
+      if (buffer.length < messageLength) break
+
+      const message = codec.decode(buffer.subarray(0, messageLength))
+      if (message.body?.length > 0) {
+        yield JSON.parse(decoder.decode(message.body))
+      }
+      buffer = buffer.subarray(messageLength)
+    }
+  }
+}