make ai sdk native

Author: sameelarif

evals/index.eval.ts

@@ -34,12 +34,12 @@ import { StagehandEvalError } from "@/types/stagehandErrors";
 import { CustomOpenAIClient } from "@/examples/external_clients/customOpenAI";
 import OpenAI from "openai";
 import { initStagehand } from "./initStagehand";
-import { AISdkClient } from "@/examples/external_clients/aisdk";
 import { google } from "@ai-sdk/google";
 import { anthropic } from "@ai-sdk/anthropic";
 import { groq } from "@ai-sdk/groq";
 import { cerebras } from "@ai-sdk/cerebras";
 import { openai } from "@ai-sdk/openai";
+import { AISdkClient } from "@/lib/llm/aisdk";
 dotenv.config();
 
 /**

examples/ai_sdk_example.ts

@@ -1,15 +1,12 @@
-import { openai } from "@ai-sdk/openai";
 import { Stagehand } from "@/dist";
-import { AISdkClient } from "./external_clients/aisdk";
 import StagehandConfig from "@/stagehand.config";
+import { openai } from "@ai-sdk/openai";
 import { z } from "zod";
 
 async function example() {
   const stagehand = new Stagehand({
     ...StagehandConfig,
-    llmClient: new AISdkClient({
-      model: openai("gpt-4o"),
-    }),
+    modelName: openai("gpt-4o"),
   });
 
   await stagehand.init();

examples/external_clients/aisdk.ts

@@ -45,6 +45,7 @@ import {
   MissingEnvironmentVariableError,
   UnsupportedModelError,
 } from "../types/stagehandErrors";
+import { LanguageModel } from "ai";
 
 dotenv.config({ path: ".env" });
 
@@ -384,7 +385,7 @@ export class Stagehand {
   public llmClient: LLMClient;
   public readonly userProvidedInstructions?: string;
   private usingAPI: boolean;
-  private modelName: AvailableModel;
+  private modelName: AvailableModel | LanguageModel;
   public apiClient: StagehandAPI | undefined;
   public readonly waitForCaptchaSolves: boolean;
   private localBrowserLaunchOptions?: LocalBrowserLaunchOptions;
@@ -656,18 +657,23 @@ export class Stagehand {
         projectId: this.projectId,
         logger: this.logger,
       });
+
       const modelApiKey =
+        // @ts-expect-error - this is a temporary fix to allow the modelName to be a LanguageModel
         LLMProvider.getModelProvider(this.modelName) === "openai"
           ? process.env.OPENAI_API_KEY || this.llmClient.clientOptions.apiKey
-          : LLMProvider.getModelProvider(this.modelName) === "anthropic"
+          : // @ts-expect-error - this is a temporary fix to allow the modelName to be a LanguageModel
+            LLMProvider.getModelProvider(this.modelName) === "anthropic"
             ? process.env.ANTHROPIC_API_KEY ||
               this.llmClient.clientOptions.apiKey
-            : LLMProvider.getModelProvider(this.modelName) === "google"
+            : // @ts-expect-error - this is a temporary fix to allow the modelName to be a LanguageModel
+              LLMProvider.getModelProvider(this.modelName) === "google"
               ? process.env.GOOGLE_API_KEY ||
                 this.llmClient.clientOptions.apiKey
               : undefined;
 
       const { sessionId } = await this.apiClient.init({
+        // @ts-expect-error - this is a temporary fix to allow the modelName to be a LanguageModel
         modelName: this.modelName,
         modelApiKey: modelApiKey,
         domSettleTimeoutMs: this.domSettleTimeoutMs,

lib/index.ts

@@ -1,20 +1,35 @@
+import {
+  UnsupportedModelError,
+  UnsupportedModelProviderError,
+} from "@/types/stagehandErrors";
+import { LanguageModel } from "ai";
 import { LogLine } from "../../types/log";
 import {
   AvailableModel,
   ClientOptions,
   ModelProvider,
 } from "../../types/model";
 import { LLMCache } from "../cache/LLMCache";
+import { AISdkClient } from "./aisdk";
 import { AnthropicClient } from "./AnthropicClient";
 import { CerebrasClient } from "./CerebrasClient";
 import { GoogleClient } from "./GoogleClient";
 import { GroqClient } from "./GroqClient";
 import { LLMClient } from "./LLMClient";
 import { OpenAIClient } from "./OpenAIClient";
-import {
-  UnsupportedModelError,
-  UnsupportedModelProviderError,
-} from "@/types/stagehandErrors";
+
+function modelToProvider(
+  modelName: AvailableModel | LanguageModel,
+): ModelProvider {
+  if (typeof modelName === "string") {
+    const provider = modelToProviderMap[modelName];
+    if (!provider) {
+      throw new UnsupportedModelError(Object.keys(modelToProviderMap));
+    }
+    return provider;
+  }
+  return "aisdk";
+}
 
 const modelToProviderMap: { [key in AvailableModel]: ModelProvider } = {
   "gpt-4.1": "openai",
@@ -81,53 +96,63 @@ export class LLMProvider {
   }
 
   getClient(
-    modelName: AvailableModel,
+    modelName: AvailableModel | LanguageModel,
     clientOptions?: ClientOptions,
   ): LLMClient {
-    const provider = modelToProviderMap[modelName];
+    const provider = modelToProvider(modelName);
     if (!provider) {
       throw new UnsupportedModelError(Object.keys(modelToProviderMap));
     }
 
+    if (provider === "aisdk") {
+      return new AISdkClient({
+        model: modelName as LanguageModel,
+        logger: this.logger,
+        enableCaching: this.enableCaching,
+        cache: this.cache,
+      });
+    }
+
+    const availableModel = modelName as AvailableModel;
     switch (provider) {
       case "openai":
         return new OpenAIClient({
           logger: this.logger,
           enableCaching: this.enableCaching,
           cache: this.cache,
-          modelName,
+          modelName: availableModel,
           clientOptions,
         });
       case "anthropic":
         return new AnthropicClient({
           logger: this.logger,
           enableCaching: this.enableCaching,
           cache: this.cache,
-          modelName,
+          modelName: availableModel,
           clientOptions,
         });
       case "cerebras":
         return new CerebrasClient({
           logger: this.logger,
           enableCaching: this.enableCaching,
           cache: this.cache,
-          modelName,
+          modelName: availableModel,
           clientOptions,
         });
       case "groq":
         return new GroqClient({
           logger: this.logger,
           enableCaching: this.enableCaching,
           cache: this.cache,
-          modelName,
+          modelName: availableModel,
           clientOptions,
         });
       case "google":
         return new GoogleClient({
           logger: this.logger,
           enableCaching: this.enableCaching,
           cache: this.cache,
-          modelName,
+          modelName: availableModel,
           clientOptions,
         });
       default: