Update CerebrasClient to use ai-sdk instead of OpenAI

Author: clayostroff

lib/index.ts

@@ -1046,4 +1046,5 @@ export * from "../types/stagehandApiErrors";
 export * from "../types/stagehandErrors";
 export * from "./llm/LLMClient";
 export * from "./llm/aisdk";
+export { CerebrasClient } from "./llm/CerebrasClient";
 export { connectToMCPServer };

lib/llm/CerebrasClient.ts

@@ -1,24 +1,15 @@
-import OpenAI from "openai";
-import type { ClientOptions } from "openai";
-import { zodToJsonSchema } from "zod-to-json-schema";
 import { LogLine } from "../../types/log";
-import { AvailableModel } from "../../types/model";
+import { AvailableModel, ClientOptions } from "../../types/model";
 import { LLMCache } from "../cache/LLMCache";
-import {
-  ChatMessage,
-  CreateChatCompletionOptions,
-  LLMClient,
-  LLMResponse,
-} from "./LLMClient";
-import { CreateChatCompletionResponseError } from "@/types/stagehandErrors";
+import { AISdkClient } from "./aisdk";
+import { LLMClient, CreateChatCompletionOptions, LLMResponse } from "./LLMClient";
+import { createCerebras } from "@ai-sdk/cerebras";
+import { LanguageModel } from "ai";
 
 export class CerebrasClient extends LLMClient {
   public type = "cerebras" as const;
-  private client: OpenAI;
-  private cache: LLMCache | undefined;
-  private enableCaching: boolean;
-  public clientOptions: ClientOptions;
   public hasVision = false;
+  private aisdkClient: AISdkClient;
 
   constructor({
     enableCaching = false,
@@ -36,308 +27,46 @@ export class CerebrasClient extends LLMClient {
   }) {
     super(modelName, userProvidedInstructions);
 
-    // Create OpenAI client with the base URL set to Cerebras API
-    this.client = new OpenAI({
-      baseURL: "https://api.cerebras.ai/v1",
-      apiKey: clientOptions?.apiKey || process.env.CEREBRAS_API_KEY,
-      ...clientOptions,
-    });
-
-    this.cache = cache;
-    this.enableCaching = enableCaching;
-    this.modelName = modelName;
-    this.clientOptions = clientOptions;
-  }
-
-  async createChatCompletion<T = LLMResponse>({
-    options,
-    retries,
-    logger,
-  }: CreateChatCompletionOptions): Promise<T> {
-    const optionsWithoutImage = { ...options };
-    delete optionsWithoutImage.image;
+    // Transform model name to remove cerebras- prefix
+    const cerebrasModelName = modelName.startsWith("cerebras-")
+      ? modelName.split("cerebras-")[1]
+      : modelName;
 
-    logger({
-      category: "cerebras",
-      message: "creating chat completion",
-      level: 2,
-      auxiliary: {
-        options: {
-          value: JSON.stringify(optionsWithoutImage),
-          type: "object",
-        },
-      },
+    // Create Cerebras provider with API key
+    const cerebrasProvider = createCerebras({
+      apiKey: (clientOptions?.apiKey as string) || process.env.CEREBRAS_API_KEY,
     });
 
-    // Try to get cached response
-    const cacheOptions = {
-      model: this.modelName.split("cerebras-")[1],
-      messages: options.messages,
-      temperature: options.temperature,
-      response_model: options.response_model,
-      tools: options.tools,
-      retries: retries,
-    };
-
-    if (this.enableCaching) {
-      const cachedResponse = await this.cache.get<T>(
-        cacheOptions,
-        options.requestId,
-      );
-      if (cachedResponse) {
-        logger({
-          category: "llm_cache",
-          message: "LLM cache hit - returning cached response",
-          level: 1,
-          auxiliary: {
-            cachedResponse: {
-              value: JSON.stringify(cachedResponse),
-              type: "object",
-            },
-            requestId: {
-              value: options.requestId,
-              type: "string",
-            },
-            cacheOptions: {
-              value: JSON.stringify(cacheOptions),
-              type: "object",
-            },
-          },
-        });
-        return cachedResponse as T;
-      }
-    }
-
-    // Format messages for Cerebras API (using OpenAI format)
-    const formattedMessages = options.messages.map((msg: ChatMessage) => {
-      const baseMessage = {
-        content:
-          typeof msg.content === "string"
-            ? msg.content
-            : Array.isArray(msg.content) &&
-                msg.content.length > 0 &&
-                "text" in msg.content[0]
-              ? msg.content[0].text
-              : "",
-      };
-
-      // Cerebras only supports system, user, and assistant roles
-      if (msg.role === "system") {
-        return { ...baseMessage, role: "system" as const };
-      } else if (msg.role === "assistant") {
-        return { ...baseMessage, role: "assistant" as const };
-      } else {
-        // Default to user for any other role
-        return { ...baseMessage, role: "user" as const };
-      }
-    });
-
-    // Format tools if provided
-    let tools = options.tools?.map((tool) => ({
-      type: "function" as const,
-      function: {
-        name: tool.name,
-        description: tool.description,
-        parameters: {
-          type: "object",
-          properties: tool.parameters.properties,
-          required: tool.parameters.required,
-        },
-      },
-    }));
-
-    // Add response model as a tool if provided
-    if (options.response_model) {
-      const jsonSchema = zodToJsonSchema(options.response_model.schema) as {
-        properties?: Record<string, unknown>;
-        required?: string[];
-      };
-      const schemaProperties = jsonSchema.properties || {};
-      const schemaRequired = jsonSchema.required || [];
-
-      const responseTool = {
-        type: "function" as const,
-        function: {
-          name: "print_extracted_data",
-          description:
-            "Prints the extracted data based on the provided schema.",
-          parameters: {
-            type: "object",
-            properties: schemaProperties,
-            required: schemaRequired,
-          },
-        },
-      };
-
-      tools = tools ? [...tools, responseTool] : [responseTool];
-    }
-
-    try {
-      // Use OpenAI client with Cerebras API
-      const apiResponse = await this.client.chat.completions.create({
-        model: this.modelName.split("cerebras-")[1],
-        messages: [
-          ...formattedMessages,
-          // Add explicit instruction to return JSON if we have a response model
-          ...(options.response_model
-            ? [
-                {
-                  role: "system" as const,
-                  content: `IMPORTANT: Your response must be valid JSON that matches this schema: ${JSON.stringify(
-                    options.response_model.schema,
-                  )}`,
-                },
-              ]
-            : []),
-        ],
-        temperature: options.temperature || 0.7,
-        max_tokens: options.maxTokens,
-        tools: tools,
-        tool_choice: options.tool_choice || "auto",
-      });
-
-      // Format the response to match the expected LLMResponse format
-      const response: LLMResponse = {
-        id: apiResponse.id,
-        object: "chat.completion",
-        created: Date.now(),
-        model: this.modelName.split("cerebras-")[1],
-        choices: [
-          {
-            index: 0,
-            message: {
-              role: "assistant",
-              content: apiResponse.choices[0]?.message?.content || null,
-              tool_calls: apiResponse.choices[0]?.message?.tool_calls || [],
-            },
-            finish_reason: apiResponse.choices[0]?.finish_reason || "stop",
-          },
-        ],
-        usage: {
-          prompt_tokens: apiResponse.usage?.prompt_tokens || 0,
-          completion_tokens: apiResponse.usage?.completion_tokens || 0,
-          total_tokens: apiResponse.usage?.total_tokens || 0,
-        },
-      };
-
-      logger({
-        category: "cerebras",
-        message: "response",
-        level: 2,
-        auxiliary: {
-          response: {
-            value: JSON.stringify(response),
-            type: "object",
-          },
-          requestId: {
-            value: options.requestId,
-            type: "string",
-          },
-        },
-      });
-
-      // If we have no response model, just return the entire LLMResponse
-      if (!options.response_model) {
-        if (this.enableCaching) {
-          await this.cache.set(cacheOptions, response, options.requestId);
-        }
-        return response as T;
-      }
-
-      // If we have a response model, parse JSON from tool calls or content
-      const toolCall = response.choices[0]?.message?.tool_calls?.[0];
-      if (toolCall?.function?.arguments) {
-        try {
-          const result = JSON.parse(toolCall.function.arguments);
-          const finalResponse = {
-            data: result,
-            usage: response.usage,
-          };
-          if (this.enableCaching) {
-            await this.cache.set(
-              cacheOptions,
-              finalResponse,
-              options.requestId,
-            );
-          }
-          return finalResponse as T;
-        } catch (e) {
-          logger({
-            category: "cerebras",
-            message: "failed to parse tool call arguments as JSON, retrying",
-            level: 0,
-            auxiliary: {
-              error: {
-                value: e.message,
-                type: "string",
-              },
-            },
-          });
+    // Get the specific model from the provider
+    const cerebrasModel = cerebrasProvider(cerebrasModelName);
+
+    this.aisdkClient = new AISdkClient({
+      model: cerebrasModel as unknown as LanguageModel,
+      logger: (message: LogLine) => {
+        // Transform log messages to use cerebras category
+        const transformedMessage = {
+          ...message,
+          category:
+            message.category === "aisdk" ? "cerebras" : message.category,
+        };
+        // Call the original logger if it exists
+        if (
+          typeof (this as unknown as { logger?: (message: LogLine) => void })
+            .logger === "function"
+        ) {
+          (this as unknown as { logger: (message: LogLine) => void }).logger(
+            transformedMessage,
+          );
         }
-      }
-
-      // If we have content but no tool calls, try to parse the content as JSON
-      const content = response.choices[0]?.message?.content;
-      if (content) {
-        try {
-          const jsonMatch = content.match(/\{[\s\S]*\}/);
-          if (jsonMatch) {
-            const result = JSON.parse(jsonMatch[0]);
-            const finalResponse = {
-              data: result,
-              usage: response.usage,
-            };
-            if (this.enableCaching) {
-              await this.cache.set(
-                cacheOptions,
-                finalResponse,
-                options.requestId,
-              );
-            }
-            return finalResponse as T;
-          }
-        } catch (e) {
-          logger({
-            category: "cerebras",
-            message: "failed to parse content as JSON",
-            level: 0,
-            auxiliary: {
-              error: {
-                value: e.message,
-                type: "string",
-              },
-            },
-          });
-        }
-      }
-
-      // If we still haven't found valid JSON and have retries left, try again
-      if (!retries || retries < 5) {
-        return this.createChatCompletion({
-          options,
-          logger,
-          retries: (retries ?? 0) + 1,
-        });
-      }
+      },
+      enableCaching,
+      cache,
+    });
+  }
 
-      throw new CreateChatCompletionResponseError("Invalid response schema");
-    } catch (error) {
-      logger({
-        category: "cerebras",
-        message: "error creating chat completion",
-        level: 0,
-        auxiliary: {
-          error: {
-            value: error.message,
-            type: "string",
-          },
-          requestId: {
-            value: options.requestId,
-            type: "string",
-          },
-        },
-      });
-      throw error;
-    }
+  async createChatCompletion<T = LLMResponse>(
+    options: CreateChatCompletionOptions,
+  ): Promise<T> {
+    return this.aisdkClient.createChatCompletion<T>(options);
   }
 }