Issue #22: Add Non-Streaming Mode Support to BedrockModelProvider

Author: chaynabors

src/models/__tests__/bedrock.test.ts

@@ -31,30 +31,49 @@ function setupMockSend(streamGenerator: () => AsyncGenerator<unknown>): void {
 }
 
 // Mock the AWS SDK
-vi.mock('@aws-sdk/client-bedrock-runtime', () => {
-  const mockSend = vi.fn(
-    async (): Promise<{ stream: AsyncIterable<unknown> }> => ({
-      stream: (async function* (): AsyncGenerator<unknown> {
-        yield { messageStart: { role: 'assistant' } }
-        yield { contentBlockStart: { contentBlockIndex: 0 } }
-        yield { contentBlockDelta: { delta: { text: 'Hello' }, contentBlockIndex: 0 } }
-        yield { contentBlockStop: { contentBlockIndex: 0 } }
-        yield { messageStop: { stopReason: 'end_turn' } }
-        yield {
-          metadata: {
-            usage: {
-              inputTokens: 10,
-              outputTokens: 5,
-              totalTokens: 15,
-            },
-            metrics: {
-              latencyMs: 100,
+vi.mock('@aws-sdk/client-bedrock-runtime', async () => {
+  // Mock command classes that the code under test will instantiate
+  const ConverseStreamCommand = vi.fn()
+  const ConverseCommand = vi.fn()
+
+  const mockSend = vi.fn(async (command: unknown) => {
+    // Check which constructor was used to create the command object
+    if (command instanceof ConverseStreamCommand) {
+      // Return a streaming response
+      return {
+        stream: (async function* (): AsyncGenerator<unknown> {
+          yield { messageStart: { role: 'assistant' } }
+          yield { contentBlockStart: { contentBlockIndex: 0 } }
+          yield { contentBlockDelta: { delta: { text: 'Hello' }, contentBlockIndex: 0 } }
+          yield { contentBlockStop: { contentBlockIndex: 0 } }
+          yield { messageStop: { stopReason: 'end_turn' } }
+          yield {
+            metadata: {
+              usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
+              metrics: { latencyMs: 100 },
             },
+          }
+        })(),
+      }
+    }
+
+    if (command instanceof ConverseCommand) {
+      // Return a non-streaming (full) response for the non-streaming API
+      return {
+        output: {
+          message: {
+            role: 'assistant',
+            content: [{ text: 'Hello' }],
           },
-        }
-      })(),
-    })
-  )
+        },
+        stopReason: 'end_turn',
+        usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
+        metrics: { latencyMs: 100 },
+      }
+    }
+
+    throw new Error('Unhandled command type in mock')
+  })
 
   // Create a mock ValidationException class
   class MockValidationException extends Error {
@@ -68,7 +87,8 @@ vi.mock('@aws-sdk/client-bedrock-runtime', () => {
     BedrockRuntimeClient: vi.fn().mockImplementation(() => ({
       send: mockSend,
     })),
-    ConverseStreamCommand: vi.fn(),
+    ConverseStreamCommand,
+    ConverseCommand,
     ValidationException: MockValidationException,
   }
 })
@@ -400,7 +420,49 @@ describe('BedrockModel', () => {
   })
 
   describe('stream', () => {
-    it('yields and validate events', async () => {
+    it('yields and validates nonstreaming events', async () => {
+      const provider = new BedrockModel({ stream: false })
+      const messages: Message[] = [{ role: 'user', content: [{ type: 'textBlock', text: 'Hello' }] }]
+
+      const events = await collectEvents(provider.stream(messages))
+
+      expect(events).toStrictEqual([
+        {
+          role: 'assistant',
+          type: 'modelMessageStartEvent',
+        },
+        {
+          type: 'modelContentBlockStartEvent',
+        },
+        {
+          delta: {
+            text: 'Hello',
+            type: 'textDelta',
+          },
+          type: 'modelContentBlockDeltaEvent',
+        },
+        {
+          type: 'modelContentBlockStopEvent',
+        },
+        {
+          stopReason: 'endTurn',
+          type: 'modelMessageStopEvent',
+        },
+        {
+          metrics: {
+            latencyMs: 100,
+          },
+          type: 'modelMetadataEvent',
+          usage: {
+            inputTokens: 10,
+            outputTokens: 5,
+            totalTokens: 15,
+          },
+        },
+      ])
+    })
+
+    it('yields and validates streaming events', async () => {
       const provider = new BedrockModel()
       const messages: Message[] = [{ role: 'user', content: [{ type: 'textBlock', text: 'Hello' }] }]
 

src/models/bedrock.ts

@@ -25,6 +25,8 @@ import {
   type ConverseStreamMetadataEvent as BedrockConverseStreamMetadataEvent,
   ContentBlockDelta,
   type ToolConfiguration,
+  ConverseCommand,
+  type ConverseCommandOutput,
 } from '@aws-sdk/client-bedrock-runtime'
 import type { Model, BaseModelConfig, StreamOptions } from '../models/model'
 import type { Message, ContentBlock } from '../types/messages'
@@ -136,6 +138,16 @@ export interface BedrockModelConfig extends BaseModelConfig {
    * @see https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/client/bedrock-runtime/command/ConverseStreamCommand/
    */
   additionalArgs?: JSONValue
+
+  /**
+   * Whether or not to stream responses from the model.
+   *
+   * This will use the ConverseStream API instead of the Converse API.
+   *
+   * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
+   * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStream.html
+   */
+  stream?: boolean
 }
 
 /**
@@ -308,19 +320,27 @@ export class BedrockModel implements Model<BedrockModelConfig, BedrockRuntimeCli
       // Format the request for Bedrock
       const request = this._formatRequest(messages, options)
 
-      // Create and send the command
-      const command = new ConverseStreamCommand(request)
-      const response = await this._client.send(command)
-
-      // Stream the response
-      if (response.stream) {
-        for await (const chunk of response.stream) {
-          // Map Bedrock events to SDK events
-          const events = this._mapBedrockEventToSDKEvents(chunk)
-          for (const event of events) {
-            yield event
+      if (this._config.stream !== false) {
+        // Create and send the command
+        const command = new ConverseStreamCommand(request)
+        const response = await this._client.send(command)
+
+        // Stream the response
+        if (response.stream) {
+          for await (const chunk of response.stream) {
+            // Map Bedrock events to SDK events
+            const events = this._mapStreamedBedrockEventToSDKEvent(chunk)
+            for (const event of events) {
+              yield event
+            }
           }
         }
+      } else {
+        const command = new ConverseCommand(request)
+        const response = await this._client.send(command)
+        for (const event of this._mapBedrockEventToSDKEvent(response)) {
+          yield event
+        }
       }
     } catch (error) {
       const err = error as Error
@@ -497,13 +517,139 @@ export class BedrockModel implements Model<BedrockModelConfig, BedrockRuntimeCli
     }
   }
 
+  private _mapBedrockEventToSDKEvent(event: ConverseCommandOutput): ModelStreamEvent[] {
+    const events: ModelStreamEvent[] = []
+
+    // Message start
+    const output = ensureDefined(event.output, 'event.output')
+    const message = ensureDefined(output.message, 'output.message')
+    const role = ensureDefined(message.role, 'message.role')
+    events.push({
+      type: 'modelMessageStartEvent',
+      role,
+    })
+
+    // Match on content blocks
+    const content = ensureDefined(message.content, 'message.content')
+    content.forEach((block, index) => {
+      if (block.text) {
+        events.push({
+          type: 'modelContentBlockStartEvent',
+        })
+
+        events.push({
+          type: 'modelContentBlockDeltaEvent',
+          delta: {
+            type: 'textDelta',
+            text: block.text,
+          },
+        })
+
+        events.push({
+          type: 'modelContentBlockStopEvent',
+        })
+      } else if (block.toolUse) {
+        events.push({
+          type: 'modelContentBlockStartEvent',
+          contentBlockIndex: index,
+          start: {
+            type: 'toolUseStart',
+            name: ensureDefined(block.toolUse.name, 'toolUse.name'),
+            toolUseId: ensureDefined(block.toolUse.toolUseId, 'toolUse.toolUseId'),
+          },
+        })
+
+        events.push({
+          type: 'modelContentBlockDeltaEvent',
+          contentBlockIndex: index,
+          delta: {
+            type: 'toolUseInputDelta',
+            input: JSON.stringify(ensureDefined(block.toolUse.input, 'toolUse.input')),
+          },
+        })
+
+        events.push({
+          type: 'modelContentBlockStopEvent',
+          contentBlockIndex: index,
+        })
+      } else if (block.reasoningContent) {
+        const reasoningText = ensureDefined(block.reasoningContent.reasoningText, 'reasoningContent.reasoningText')
+        events.push({
+          type: 'modelContentBlockDeltaEvent',
+          contentBlockIndex: index,
+          delta: {
+            type: 'reasoningDelta',
+            text: ensureDefined(reasoningText.text, 'reasoningText.text'),
+          },
+        })
+
+        if (reasoningText.signature) {
+          events.push({
+            type: 'modelContentBlockDeltaEvent',
+            contentBlockIndex: index,
+            delta: {
+              type: 'reasoningDelta',
+              signature: reasoningText.signature,
+            },
+          })
+        }
+
+        events.push({
+          type: 'modelContentBlockStopEvent',
+          contentBlockIndex: index,
+        })
+      }
+    })
+
+    const stopReasonRaw = ensureDefined(event.stopReason, 'event.stopReason') as string
+    let mappedStopReason: string
+
+    if (stopReasonRaw in STOP_REASON_MAP) {
+      mappedStopReason = STOP_REASON_MAP[stopReasonRaw as keyof typeof STOP_REASON_MAP]
+    } else {
+      console.warn(`Unknown stop reason: "${stopReasonRaw}". Converting to camelCase: "${snakeToCamel(stopReasonRaw)}"`)
+      mappedStopReason = snakeToCamel(stopReasonRaw) // Assumes snakeToCamel utility exists
+    }
+
+    // Adjust for tool_use, which is sometimes incorrectly reported as end_turn
+    if (mappedStopReason === 'endTurn' && event.output?.message?.content?.some((block) => 'toolUse' in block)) {
+      mappedStopReason = 'toolUse'
+      console.warn(`Adjusting stop reason from 'end_turn' to 'tool_use' due to tool use in content blocks.`)
+    }
+
+    events.push({
+      type: 'modelMessageStopEvent',
+      stopReason: mappedStopReason,
+    })
+
+    const usage = ensureDefined(event.usage, 'output.usage')
+    const metadataEvent: ModelStreamEvent = {
+      type: 'modelMetadataEvent',
+      usage: {
+        inputTokens: ensureDefined(usage.inputTokens, 'usage.inputTokens'),
+        outputTokens: ensureDefined(usage.outputTokens, 'usage.outputTokens'),
+        totalTokens: ensureDefined(usage.totalTokens, 'usage.totalTokens'),
+      },
+    }
+
+    if (event.metrics) {
+      metadataEvent.metrics = {
+        latencyMs: ensureDefined(event.metrics.latencyMs, 'metrics.latencyMs'),
+      }
+    }
+
+    events.push(metadataEvent)
+
+    return events
+  }
+
   /**
    * Maps a Bedrock event to SDK streaming events.
    *
    * @param chunk - Bedrock event chunk
    * @returns Array of SDK streaming events
    */
-  private _mapBedrockEventToSDKEvents(chunk: ConverseStreamOutput): ModelStreamEvent[] {
+  private _mapStreamedBedrockEventToSDKEvent(chunk: ConverseStreamOutput): ModelStreamEvent[] {
     const events: ModelStreamEvent[] = []
 
     // Extract the event type key