azure/bedrock api integration

CHANGELOG.md

@@ -233,15 +233,13 @@
   We're thrilled to announce the release of Stagehand 2.0, bringing significant improvements to make browser automation more powerful, faster, and easier to use than ever before.
 
   ### 🚀 New Features
-
   - **Introducing `stagehand.agent`**: A powerful new way to integrate SOTA Computer use models or Browserbase's [Open Operator](https://operator.browserbase.com) into Stagehand with one line of code! Perfect for multi-step workflows and complex interactions. [Learn more](https://docs.stagehand.dev/concepts/agent)
   - **Lightning-fast `act` and `extract`**: Major performance improvements to make your automations run significantly faster.
   - **Enhanced Logging**: Better visibility into what's happening during automation with improved logging and debugging capabilities.
   - **Comprehensive Documentation**: A completely revamped documentation site with better examples, guides, and best practices.
   - **Improved Error Handling**: More descriptive errors and better error recovery to help you debug issues faster.
 
   ### 🛠️ Developer Experience
-
   - **Better TypeScript Support**: Enhanced type definitions and better IDE integration
   - **Better Error Messages**: Clearer, more actionable error messages to help you debug faster
   - **Improved Caching**: More reliable action caching for better performance
@@ -502,7 +500,6 @@
 - [#316](https://github.com/browserbase/stagehand/pull/316) [`902e633`](https://github.com/browserbase/stagehand/commit/902e633e126a58b80b757ea0ecada01a7675a473) Thanks [@kamath](https://github.com/kamath)! - rename browserbaseResumeSessionID -> browserbaseSessionID
 
 - [#296](https://github.com/browserbase/stagehand/pull/296) [`f11da27`](https://github.com/browserbase/stagehand/commit/f11da27a20409c240ceeea2003d520f676def61a) Thanks [@kamath](https://github.com/kamath)! - - Deprecate fields in `init` in favor of constructor options
-
   - Deprecate `initFromPage` in favor of `browserbaseResumeSessionID` in constructor
   - Rename `browserBaseSessionCreateParams` -> `browserbaseSessionCreateParams`
 

lib/StagehandPage.ts

@@ -739,6 +739,7 @@ ${scriptContent} \
             const result = await this.api.act({
               ...observeResult,
               frameId: this.rootFrameId,
+              modelClientOptions: this.stagehand["modelClientOptions"],
             });
             this.stagehand.addToHistory("act", observeResult, result);
             return result;
@@ -836,7 +837,10 @@ ${scriptContent} \
       if (!instructionOrOptions) {
         let result: ExtractResult<T>;
         if (this.api) {
-          result = await this.api.extract<T>({ frameId: this.rootFrameId });
+          result = await this.api.extract<T>({
+            frameId: this.rootFrameId,
+            modelClientOptions: this.stagehand["modelClientOptions"],
+          });
         } else {
           result = await this.extractHandler.extract();
         }

lib/a11y/utils.ts

@@ -183,7 +183,8 @@ export async function buildBackendIdMaps(
           if (n.contentDocument && locate(n.contentDocument)) return true;
           return false;
         } else {
-          if (n.backendNodeId === backendNodeId) return (iframeNode = n), true;
+          if (n.backendNodeId === backendNodeId)
+            return ((iframeNode = n), true);
           return (
             (n.children?.some(locate) ?? false) ||
             (n.contentDocument ? locate(n.contentDocument) : false)

lib/api.ts

@@ -48,7 +48,6 @@ export class StagehandAPI {
 
   async init({
     modelName,
-    modelApiKey,
     domSettleTimeoutMs,
     verbose,
     debugDom,
@@ -59,11 +58,6 @@ export class StagehandAPI {
     browserbaseSessionCreateParams,
     browserbaseSessionID,
   }: StartSessionParams): Promise<StartSessionResult> {
-    if (!modelApiKey) {
-      throw new StagehandAPIError("modelApiKey is required");
-    }
-    this.modelApiKey = modelApiKey;
-
     const region = browserbaseSessionCreateParams?.region;
     if (region && region !== "us-west-2") {
       return { sessionId: browserbaseSessionID ?? null, available: false };
@@ -186,10 +180,19 @@ export class StagehandAPI {
     const queryString = urlParams.toString();
     const url = `/sessions/${this.sessionId}/${method}${queryString ? `?${queryString}` : ""}`;
 
-    const response = await this.request(url, {
-      method: "POST",
-      body: JSON.stringify(args),
-    });
+    // Extract modelClientOptions from args if present
+    const modelClientOptions = (
+      args as { modelClientOptions?: Record<string, unknown> }
+    )?.modelClientOptions;
+
+    const response = await this.request(
+      url,
+      {
+        method: "POST",
+        body: JSON.stringify(args),
+      },
+      modelClientOptions,
+    );
 
     if (!response.ok) {
       const errorBody = await response.text();
@@ -248,6 +251,7 @@ export class StagehandAPI {
   private async request(
     path: string,
     options: RequestInit = {},
+    modelClientOptions?: Record<string, unknown>,
   ): Promise<Response> {
     const defaultHeaders: Record<string, string> = {
       "x-bb-api-key": this.apiKey,
@@ -261,6 +265,12 @@ export class StagehandAPI {
       "x-sdk-version": STAGEHAND_VERSION,
     };
 
+    // Add modelClientOptions as a header if provided
+    if (modelClientOptions) {
+      defaultHeaders["x-model-client-options"] =
+        JSON.stringify(modelClientOptions);
+    }
+
     if (options.method === "POST" && options.body) {
       defaultHeaders["Content-Type"] = "application/json";
     }

lib/index.ts

@@ -45,7 +45,11 @@ import { LLMProvider } from "./llm/LLMProvider";
 import { StagehandLogger } from "./logger";
 import { connectToMCPServer } from "./mcp/connection";
 import { resolveTools } from "./mcp/utils";
-import { isRunningInBun, loadApiKeyFromEnv } from "./utils";
+import {
+  isRunningInBun,
+  loadApiKeyFromEnv,
+  loadBedrockClientOptions,
+} from "./utils";
 
 dotenv.config({ path: ".env" });
 
@@ -587,28 +591,38 @@ export class Stagehand {
     if (!modelClientOptions?.apiKey) {
       // If no API key is provided, try to load it from the environment
       if (LLMProvider.getModelProvider(this.modelName) === "aisdk") {
-        modelApiKey = loadApiKeyFromEnv(
-          this.modelName.split("/")[0],
-          this.logger,
-        );
+        const provider = this.modelName.split("/")[0];
+
+        // Special handling for Amazon Bedrock's complex authentication
+        if (provider === "bedrock") {
+          const bedrockOptions = loadBedrockClientOptions(this.logger);
+          this.modelClientOptions = {
+            ...modelClientOptions,
+            ...bedrockOptions,
+          };
+        } else {
+          // Standard single API key handling for other AISDK providers
+          modelApiKey = loadApiKeyFromEnv(provider, this.logger);
+          this.modelClientOptions = {
+            ...modelClientOptions,
+            apiKey: modelApiKey,
+          };
+        }
       } else {
         // Temporary add for legacy providers
         modelApiKey =
           LLMProvider.getModelProvider(this.modelName) === "openai"
-            ? process.env.OPENAI_API_KEY ||
-              this.llmClient?.clientOptions?.apiKey
+            ? process.env.OPENAI_API_KEY
             : LLMProvider.getModelProvider(this.modelName) === "anthropic"
-              ? process.env.ANTHROPIC_API_KEY ||
-                this.llmClient?.clientOptions?.apiKey
+              ? process.env.ANTHROPIC_API_KEY
               : LLMProvider.getModelProvider(this.modelName) === "google"
-                ? process.env.GOOGLE_API_KEY ||
-                  this.llmClient?.clientOptions?.apiKey
+                ? process.env.GOOGLE_API_KEY
                 : undefined;
+        this.modelClientOptions = {
+          ...modelClientOptions,
+          apiKey: modelApiKey,
+        };
       }
-      this.modelClientOptions = {
-        ...modelClientOptions,
-        apiKey: modelApiKey,
-      };
     } else {
       this.modelClientOptions = modelClientOptions;
     }
@@ -756,7 +770,7 @@ export class Stagehand {
         logger: this.logger,
       });
 
-      const modelApiKey = this.modelClientOptions?.apiKey;
+      const modelApiKey = this.modelClientOptions?.apiKey as string;
       const { sessionId, available } = await this.apiClient.init({
         modelName: this.modelName,
         modelApiKey: modelApiKey,

lib/llm/AnthropicClient.ts

@@ -1,3 +1,4 @@
+import { CreateChatCompletionResponseError } from "@/types/stagehandErrors";
 import Anthropic, { ClientOptions } from "@anthropic-ai/sdk";
 import {
   ImageBlockParam,
@@ -14,14 +15,13 @@ import {
   LLMClient,
   LLMResponse,
 } from "./LLMClient";
-import { CreateChatCompletionResponseError } from "@/types/stagehandErrors";
 
 export class AnthropicClient extends LLMClient {
   public type = "anthropic" as const;
   private client: Anthropic;
   private cache: LLMCache | undefined;
   private enableCaching: boolean;
-  public clientOptions: ClientOptions;
+  public clientOptions?: ClientOptions;
 
   constructor({
     enableCaching = false,

lib/llm/CerebrasClient.ts

@@ -1,5 +1,6 @@
-import OpenAI from "openai";
+import { CreateChatCompletionResponseError } from "@/types/stagehandErrors";
 import type { ClientOptions } from "openai";
+import OpenAI from "openai";
 import { zodToJsonSchema } from "zod-to-json-schema";
 import { LogLine } from "../../types/log";
 import { AvailableModel } from "../../types/model";
@@ -10,14 +11,13 @@ import {
   LLMClient,
   LLMResponse,
 } from "./LLMClient";
-import { CreateChatCompletionResponseError } from "@/types/stagehandErrors";
 
 export class CerebrasClient extends LLMClient {
   public type = "cerebras" as const;
   private client: OpenAI;
   private cache: LLMCache | undefined;
   private enableCaching: boolean;
-  public clientOptions: ClientOptions;
+  public clientOptions?: ClientOptions;
   public hasVision = false;
 
   constructor({

lib/llm/GoogleClient.ts

@@ -83,7 +83,7 @@ export class GoogleClient extends LLMClient {
       clientOptions.apiKey = loadApiKeyFromEnv("google_legacy", logger);
     }
     this.clientOptions = clientOptions;
-    this.client = new GoogleGenAI({ apiKey: clientOptions.apiKey });
+    this.client = new GoogleGenAI({ apiKey: clientOptions.apiKey as string });
     this.cache = cache;
     this.enableCaching = enableCaching;
     this.modelName = modelName;

lib/llm/LLMClient.ts

@@ -12,7 +12,7 @@ import {
 } from "ai";
 import { ZodType } from "zod/v3";
 import { LogLine } from "../../types/log";
-import { AvailableModel, ClientOptions } from "../../types/model";
+import { AvailableModel } from "../../types/model";
 
 export interface ChatMessage {
   role: "system" | "user" | "assistant";
@@ -100,7 +100,6 @@ export abstract class LLMClient {
   public type: "openai" | "anthropic" | "cerebras" | "groq" | (string & {});
   public modelName: AvailableModel | (string & {});
   public hasVision: boolean;
-  public clientOptions: ClientOptions;
   public userProvidedInstructions?: string;
 
   constructor(modelName: AvailableModel, userProvidedInstructions?: string) {

lib/llm/LLMProvider.ts

@@ -1,8 +1,22 @@
+import { AISDKCustomProvider, AISDKProvider } from "@/types/llm";
 import {
   UnsupportedAISDKModelProviderError,
   UnsupportedModelError,
   UnsupportedModelProviderError,
 } from "@/types/stagehandErrors";
+import { anthropic, createAnthropic } from "@ai-sdk/anthropic";
+import { azure, createAzure } from "@ai-sdk/azure";
+import { bedrock, createAmazonBedrock } from "@ai-sdk/amazon-bedrock";
+import { cerebras, createCerebras } from "@ai-sdk/cerebras";
+import { createDeepSeek, deepseek } from "@ai-sdk/deepseek";
+import { createGoogleGenerativeAI, google } from "@ai-sdk/google";
+import { createGroq, groq } from "@ai-sdk/groq";
+import { createMistral, mistral } from "@ai-sdk/mistral";
+import { createOpenAI, openai } from "@ai-sdk/openai";
+import { createPerplexity, perplexity } from "@ai-sdk/perplexity";
+import { createTogetherAI, togetherai } from "@ai-sdk/togetherai";
+import { createXai, xai } from "@ai-sdk/xai";
+import { ollama } from "ollama-ai-provider";
 import { LogLine } from "../../types/log";
 import {
   AvailableModel,
@@ -17,19 +31,6 @@ import { GoogleClient } from "./GoogleClient";
 import { GroqClient } from "./GroqClient";
 import { LLMClient } from "./LLMClient";
 import { OpenAIClient } from "./OpenAIClient";
-import { openai, createOpenAI } from "@ai-sdk/openai";
-import { anthropic, createAnthropic } from "@ai-sdk/anthropic";
-import { google, createGoogleGenerativeAI } from "@ai-sdk/google";
-import { xai, createXai } from "@ai-sdk/xai";
-import { azure, createAzure } from "@ai-sdk/azure";
-import { groq, createGroq } from "@ai-sdk/groq";
-import { cerebras, createCerebras } from "@ai-sdk/cerebras";
-import { togetherai, createTogetherAI } from "@ai-sdk/togetherai";
-import { mistral, createMistral } from "@ai-sdk/mistral";
-import { deepseek, createDeepSeek } from "@ai-sdk/deepseek";
-import { perplexity, createPerplexity } from "@ai-sdk/perplexity";
-import { ollama } from "ollama-ai-provider";
-import { AISDKProvider, AISDKCustomProvider } from "@/types/llm";
 
 const AISDKProviders: Record<string, AISDKProvider> = {
   openai,
@@ -38,6 +39,7 @@ const AISDKProviders: Record<string, AISDKProvider> = {
   xai,
   azure,
   groq,
+  bedrock,
   cerebras,
   togetherai,
   mistral,
@@ -52,6 +54,7 @@ const AISDKProvidersWithAPIKey: Record<string, AISDKCustomProvider> = {
   xai: createXai,
   azure: createAzure,
   groq: createGroq,
+  bedrock: createAmazonBedrock,
   cerebras: createCerebras,
   togetherai: createTogetherAI,
   mistral: createMistral,
@@ -97,23 +100,18 @@ const modelToProviderMap: { [key in AvailableModel]: ModelProvider } = {
 export function getAISDKLanguageModel(
   subProvider: string,
   subModelName: string,
-  apiKey?: string,
-  baseURL?: string,
+  modelClientOptions?: ClientOptions,
 ) {
-  if (apiKey) {
+  if (modelClientOptions && Object.keys(modelClientOptions).length > 0) {
     const creator = AISDKProvidersWithAPIKey[subProvider];
     if (!creator) {
       throw new UnsupportedAISDKModelProviderError(
         subProvider,
         Object.keys(AISDKProvidersWithAPIKey),
       );
     }
-    // Create the provider instance with the API key and baseURL if provided
-    const providerConfig: { apiKey: string; baseURL?: string } = { apiKey };
-    if (baseURL) {
-      providerConfig.baseURL = baseURL;
-    }
-    const provider = creator(providerConfig);
+    // Create the provider instance with the custom configuration options
+    const provider = creator(modelClientOptions as Record<string, unknown>);
     // Get the specific model from the provider
     return provider(subModelName);
   } else {
@@ -170,8 +168,7 @@ export class LLMProvider {
       const languageModel = getAISDKLanguageModel(
         subProvider,
         subModelName,
-        clientOptions?.apiKey,
-        clientOptions?.baseURL,
+        clientOptions,
       );
 
       return new AISdkClient({

lib/llm/OpenAIClient.ts

@@ -1,3 +1,8 @@
+import {
+  CreateChatCompletionResponseError,
+  StagehandError,
+  ZodSchemaValidationError,
+} from "@/types/stagehandErrors";
 import OpenAI, { ClientOptions } from "openai";
 import { zodResponseFormat } from "openai/helpers/zod";
 import {
@@ -21,11 +26,6 @@ import {
   LLMClient,
   LLMResponse,
 } from "./LLMClient";
-import {
-  CreateChatCompletionResponseError,
-  StagehandError,
-  ZodSchemaValidationError,
-} from "@/types/stagehandErrors";
 
 export class OpenAIClient extends LLMClient {
   public type = "openai" as const;