Pass LLM logger in createChatCompletion (#385)

* migrate to new options syntax

* fix

* input logger in extract/observe

* changeset

* patch -> minor

* import and override stagehandconfig

---------

Co-authored-by: Anirudh Kamath <[email protected]>

Author: sameelarif

Diff for: .changeset/soft-snails-lick.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": minor
+---
+
+Moved the LLMClient logger paremeter to the createChatCompletion method options.

Diff for: examples/external_client.ts

@@ -1,20 +1,15 @@
-import { type ConstructorParams, Stagehand } from "../lib";
+import { Stagehand } from "../lib";
 import { z } from "zod";
 import { OllamaClient } from "./external_clients/ollama";
-
-const StagehandConfig: ConstructorParams = {
-  env: "BROWSERBASE",
-  apiKey: process.env.BROWSERBASE_API_KEY,
-  projectId: process.env.BROWSERBASE_PROJECT_ID,
-  verbose: 1,
-  llmClient: new OllamaClient({
-    modelName: "llama3.2",
-  }),
-  debugDom: true,
-};
+import StagehandConfig from "./stagehand.config";
 
 async function example() {
-  const stagehand = new Stagehand(StagehandConfig);
+  const stagehand = new Stagehand({
+    ...StagehandConfig,
+    llmClient: new OllamaClient({
+      modelName: "llama3.2",
+    }),
+  });
 
   await stagehand.init();
   await stagehand.page.goto("https://news.ycombinator.com");

Diff for: examples/external_clients/ollama.ts

@@ -1,14 +1,5 @@
 import OpenAI, { type ClientOptions } from "openai";
 import { zodResponseFormat } from "openai/helpers/zod";
-import type { LLMCache } from "../../lib/cache/LLMCache";
-import { validateZodSchema } from "../../lib/utils";
-import {
-  type ChatCompletionOptions,
-  type ChatMessage,
-  LLMClient,
-} from "../../lib/llm/LLMClient";
-import type { LogLine } from "../../types/log";
-import type { AvailableModel } from "../../types/model";
 import type {
   ChatCompletion,
   ChatCompletionAssistantMessageParam,
@@ -19,23 +10,28 @@ import type {
   ChatCompletionSystemMessageParam,
   ChatCompletionUserMessageParam,
 } from "openai/resources/chat";
+import type { LLMCache } from "../../lib/cache/LLMCache";
+import {
+  type ChatMessage,
+  CreateChatCompletionOptions,
+  LLMClient,
+} from "../../lib/llm/LLMClient";
+import { validateZodSchema } from "../../lib/utils";
+import type { AvailableModel } from "../../types/model";
 
 export class OllamaClient extends LLMClient {
   public type = "ollama" as const;
   private client: OpenAI;
   private cache: LLMCache | undefined;
-  public logger: (message: LogLine) => void;
   private enableCaching: boolean;
   public clientOptions: ClientOptions;
 
   constructor({
-    logger,
     enableCaching = false,
     cache = undefined,
     modelName = "llama3.2",
     clientOptions,
   }: {
-    logger?: (message: LogLine) => void;
     enableCaching?: boolean;
     cache?: LLMCache;
     modelName?: string;
@@ -47,16 +43,16 @@ export class OllamaClient extends LLMClient {
       baseURL: clientOptions?.baseURL || "http://localhost:11434/v1",
       apiKey: "ollama",
     });
-    this.logger = logger;
     this.cache = cache;
     this.enableCaching = enableCaching;
     this.modelName = modelName as AvailableModel;
   }
 
-  async createChatCompletion<T = ChatCompletion>(
-    options: ChatCompletionOptions,
+  async createChatCompletion<T = ChatCompletion>({
+    options,
     retries = 3,
-  ): Promise<T> {
+    logger,
+  }: CreateChatCompletionOptions): Promise<T> {
     const { image, requestId, ...optionsWithoutImageAndRequestId } = options;
 
     // TODO: Implement vision support
@@ -66,7 +62,7 @@ export class OllamaClient extends LLMClient {
       );
     }
 
-    this.logger({
+    logger({
       category: "ollama",
       message: "creating chat completion",
       level: 1,
@@ -122,7 +118,7 @@ export class OllamaClient extends LLMClient {
       );
 
       if (cachedResponse) {
-        this.logger({
+        logger({
           category: "llm_cache",
           message: "LLM cache hit - returning cached response",
           level: 1,
@@ -140,7 +136,7 @@ export class OllamaClient extends LLMClient {
         return cachedResponse;
       }
 
-      this.logger({
+      logger({
         category: "llm_cache",
         message: "LLM cache miss - no cached response found",
         level: 1,
@@ -168,7 +164,7 @@ export class OllamaClient extends LLMClient {
       model: this.modelName,
     };
 
-    this.logger({
+    logger({
       category: "ollama",
       message: "creating chat completion",
       level: 1,
@@ -257,7 +253,7 @@ export class OllamaClient extends LLMClient {
 
     const response = await this.client.chat.completions.create(body);
 
-    this.logger({
+    logger({
       category: "ollama",
       message: "response",
       level: 1,
@@ -279,7 +275,11 @@ export class OllamaClient extends LLMClient {
 
       if (!validateZodSchema(options.response_model.schema, parsedData)) {
         if (retries > 0) {
-          return this.createChatCompletion(options, retries - 1);
+          return this.createChatCompletion({
+            options,
+            logger,
+            retries: retries - 1,
+          });
         }
 
         throw new Error("Invalid response schema");
@@ -299,7 +299,7 @@ export class OllamaClient extends LLMClient {
     }
 
     if (this.enableCaching) {
-      this.logger({
+      logger({
         category: "llm_cache",
         message: "caching response",
         level: 1,

Diff for: lib/handlers/extractHandler.ts

@@ -306,6 +306,7 @@ export class StagehandExtractHandler {
       chunksTotal: 1,
       llmClient,
       requestId,
+      logger: this.logger,
     });
 
     const {
@@ -434,6 +435,7 @@ export class StagehandExtractHandler {
       chunksTotal: chunks.length,
       requestId,
       isUsingTextExtract: false,
+      logger: this.logger,
     });
 
     const {

Diff for: lib/handlers/observeHandler.ts

@@ -120,6 +120,7 @@ export class StagehandObserveHandler {
       llmClient,
       image: annotatedScreenshot,
       requestId,
+      logger: this.logger,
     });
 
     const elementsWithSelectors = observationResponse.elements.map(

Diff for: lib/index.ts

@@ -372,9 +372,6 @@ export class Stagehand {
         this.llmClient = undefined;
       }
     }
-    if (this.llmClient && !this.llmClient.logger) {
-      this.llmClient.logger = this.logger;
-    }
     this.domSettleTimeoutMs = domSettleTimeoutMs ?? 30_000;
     this.headless = headless ?? false;
     this.browserbaseSessionCreateParams = browserbaseSessionCreateParams;