feat: ✨ Add gpt-oss to the bedrock provider

Author: chezsmithy

core/llm/llms/Bedrock.ts

@@ -2,12 +2,12 @@ import {
   BedrockRuntimeClient,
   ContentBlock,
   ConversationRole,
+  ConverseCommand,
   ConverseStreamCommand,
-  ConverseStreamCommandOutput,
   ImageFormat,
   InvokeModelCommand,
   Message,
-  ToolConfiguration,
+  ToolConfiguration
 } from "@aws-sdk/client-bedrock-runtime";
 import { fromNodeProviderChain } from "@aws-sdk/credential-providers";
 
@@ -91,47 +91,33 @@ class Bedrock extends BaseLLM {
     signal: AbortSignal,
     options: CompletionOptions,
   ): AsyncGenerator<ChatMessage> {
-    const credentials = await this._getCredentials();
-    const client = new BedrockRuntimeClient({
-      region: this.region,
-      endpoint: this.apiBase,
-      credentials: {
-        accessKeyId: credentials.accessKeyId,
-        secretAccessKey: credentials.secretAccessKey,
-        sessionToken: credentials.sessionToken || "",
-      },
-    });
-
-    let config_headers =
-      this.requestOptions && this.requestOptions.headers
-        ? this.requestOptions.headers
-        : {};
-    // AWS SigV4 requires strict canonicalization of headers.
-    // DO NOT USE "_" in your header name. It will return an error like below.
-    // "The request signature we calculated does not match the signature you provided."
+    if (options.stream !== false) {
+      yield* this._streamChatStreaming(messages, signal, options);
+    } else {
+      yield* this._streamChatNonStreaming(messages, signal, options);
+    }
+  }
 
-    client.middlewareStack.add(
-      (next) => async (args: any) => {
-        args.request.headers = {
-          ...args.request.headers,
-          ...config_headers,
-        };
-        return next(args);
-      },
-      {
-        step: "build",
-      },
-    );
+  /**
+   * Handles streaming chat using ConverseStreamCommand
+   */
+  private async *_streamChatStreaming(
+    messages: ChatMessage[],
+    signal: AbortSignal,
+    options: CompletionOptions,
+  ): AsyncGenerator<ChatMessage> {
+    const client = await this._createBedrockClient();
+    this._addClientMiddleware(client);
 
     const input = this._generateConverseInput(messages, {
       ...options,
       stream: true,
     });
-    const command = new ConverseStreamCommand(input);
 
-    const response = (await client.send(command, {
+    const command = new ConverseStreamCommand(input);
+    const response = await client.send(command, {
       abortSignal: signal,
-    })) as ConverseStreamCommandOutput;
+    });
 
     if (!response?.stream) {
       throw new Error("No stream received from Bedrock API");
@@ -158,17 +144,17 @@ class Bedrock extends BaseLLM {
               role: "assistant",
               content: chunk.contentBlockDelta.delta.text,
             };
-            continue;
+            continue; // Continue parsing the stream
           }
 
-          // Handle text content
+          // Handle reasoning text content
           if ((chunk.contentBlockDelta.delta as any).reasoningContent?.text) {
             yield {
               role: "thinking",
               content: (chunk.contentBlockDelta.delta as any).reasoningContent
                 .text,
             };
-            continue;
+            continue; // Continue parsing the stream
           }
 
           // Handle signature for thinking
@@ -178,7 +164,7 @@ class Bedrock extends BaseLLM {
               content: "",
               signature: delta.reasoningContent.signature,
             };
-            continue;
+            continue; // Continue parsing the stream
           }
 
           // Handle redacted thinking
@@ -188,7 +174,7 @@ class Bedrock extends BaseLLM {
               content: "",
               redactedThinking: delta.redactedReasoning.data,
             };
-            continue;
+            continue; // Continue parsing the stream
           }
 
           if (
@@ -201,7 +187,7 @@ class Bedrock extends BaseLLM {
             }
             this._currentToolResponse.input +=
               chunk.contentBlockDelta.delta.toolUse.input;
-            continue;
+            continue; // Continue parsing the stream
           }
         }
 
@@ -213,7 +199,7 @@ class Bedrock extends BaseLLM {
               content: "",
               redactedThinking: start.redactedReasoning.data,
             };
-            continue;
+            continue; // Continue parsing the stream
           }
 
           const toolUse = chunk.contentBlockStart.start.toolUse;
@@ -224,7 +210,7 @@ class Bedrock extends BaseLLM {
               input: "",
             };
           }
-          continue;
+          continue; // Continue parsing the stream
         }
 
         if (chunk.contentBlockStop) {
@@ -245,7 +231,7 @@ class Bedrock extends BaseLLM {
             };
             this._currentToolResponse = null;
           }
-          continue;
+          continue; // Continue parsing the stream
         }
       }
     } catch (error: unknown) {
@@ -255,6 +241,133 @@ class Bedrock extends BaseLLM {
     }
   }
 
+  /**
+   * Handles non-streaming chat using ConverseCommand
+   */
+  private async *_streamChatNonStreaming(
+    messages: ChatMessage[],
+    signal: AbortSignal,
+    options: CompletionOptions,
+  ): AsyncGenerator<ChatMessage> {
+    const client = await this._createBedrockClient();
+    this._addClientMiddleware(client);
+
+    const input = this._generateConverseInput(messages, {
+      ...options,
+      stream: false,
+    });
+
+    const command = new ConverseCommand(input);
+    const response = await client.send(command, {
+      abortSignal: signal,
+    });
+
+    // Reset cache metrics for new request
+    this._promptCachingMetrics = {
+      cacheReadInputTokens: 0,
+      cacheWriteInputTokens: 0,
+    };
+
+    try {
+      if (response.output?.message?.content) {
+        for (const contentBlock of response.output.message.content) {
+          if (contentBlock.text) {
+            yield {
+              role: "assistant",
+              content: contentBlock.text,
+            };
+          }
+
+          if ((contentBlock as any).reasoningContent) {
+            const reasoningContent = (contentBlock as any).reasoningContent;
+            if (reasoningContent.reasoningText) {
+              yield {
+                role: "thinking",
+                content: reasoningContent.reasoningText.text || "",
+                signature: reasoningContent.reasoningText.signature,
+              };
+            }
+            if (reasoningContent.redactedContent) {
+              yield {
+                role: "thinking",
+                content: "",
+                redactedThinking: reasoningContent.redactedContent,
+              };
+            }
+          }
+
+          if (contentBlock.toolUse) {
+            yield {
+              role: "assistant",
+              content: "",
+              toolCalls: [
+                {
+                  id: contentBlock.toolUse.toolUseId,
+                  type: "function",
+                  function: {
+                    name: contentBlock.toolUse.name,
+                    arguments: JSON.stringify(contentBlock.toolUse.input || {}),
+                  },
+                },
+              ],
+            };
+          }
+        }
+      }
+
+      // Handle usage metadata if available
+      if (response.usage) {
+        console.log(`${JSON.stringify(response.usage)}`);
+      }
+    } catch (error: unknown) {
+      // Clean up state and let the original error bubble up to the retry decorator
+      this._currentToolResponse = null;
+      throw error;
+    }
+  }
+
+  /**
+   * Creates and configures a Bedrock Runtime Client
+   */
+  private async _createBedrockClient(): Promise<BedrockRuntimeClient> {
+    const credentials = await this._getCredentials();
+    return new BedrockRuntimeClient({
+      region: this.region,
+      endpoint: this.apiBase,
+      credentials: {
+        accessKeyId: credentials.accessKeyId,
+        secretAccessKey: credentials.secretAccessKey,
+        sessionToken: credentials.sessionToken || "",
+      },
+    });
+  }
+
+  /**
+   * Adds middleware to the Bedrock client for custom headers
+   */
+  private _addClientMiddleware(client: BedrockRuntimeClient): void {
+    const config_headers =
+      this.requestOptions && this.requestOptions.headers
+        ? this.requestOptions.headers
+        : {};
+    // AWS SigV4 requires strict canonicalization of headers.
+    // DO NOT USE "_" in your header name. It will return an error like below.
+    // "The request signature we calculated does not match the signature you provided."
+
+    client.middlewareStack.add(
+      (next) => async (args: any) => {
+        args.request.headers = {
+          ...args.request.headers,
+          ...config_headers,
+        };
+        return next(args);
+      },
+      {
+        step: "build",
+      },
+    );
+  }
+
   /**
    * Generates the input payload for the Bedrock Converse API
    * @param messages - Array of chat messages

core/llm/toolSupport.ts

@@ -124,6 +124,7 @@ export const PROVIDER_TOOL_SUPPORT: Record<string, (model: string) => boolean> =
           "nova-pro",
           "nova-micro",
           "nova-premier",
+          "gpt-oss",
         ].some((part) => model.toLowerCase().includes(part))
       ) {
         return true;