@browserbasehq/stagehand 1.4.0 → 1.5.0

package/lib/handlers/extractHandler.ts

@@ -0,0 +1,179 @@
+import { LLMProvider } from "../llm/LLMProvider";
+import { Stagehand } from "../index";
+import { z } from "zod";
+import { LogLine } from "../../types/log";
+import { extract } from "../inference";
+import { LLMClient } from "../llm/LLMClient";
+
+export class StagehandExtractHandler {
+  private readonly stagehand: Stagehand;
+
+  private readonly logger: (logLine: LogLine) => void;
+  private readonly waitForSettledDom: (
+    domSettleTimeoutMs?: number,
+  ) => Promise<void>;
+  private readonly startDomDebug: () => Promise<void>;
+  private readonly cleanupDomDebug: () => Promise<void>;
+  private readonly llmProvider: LLMProvider;
+  private readonly llmClient: LLMClient;
+  private readonly verbose: 0 | 1 | 2;
+
+  constructor({
+    stagehand,
+    logger,
+    waitForSettledDom,
+    startDomDebug,
+    cleanupDomDebug,
+    llmProvider,
+    llmClient,
+    verbose,
+  }: {
+    stagehand: Stagehand;
+    logger: (message: {
+      category?: string;
+      message: string;
+      level?: number;
+      auxiliary?: { [key: string]: { value: string; type: string } };
+    }) => void;
+    waitForSettledDom: (domSettleTimeoutMs?: number) => Promise<void>;
+    startDomDebug: () => Promise<void>;
+    cleanupDomDebug: () => Promise<void>;
+    llmProvider: LLMProvider;
+    llmClient: LLMClient;
+    verbose: 0 | 1 | 2;
+  }) {
+    this.stagehand = stagehand;
+    this.logger = logger;
+    this.waitForSettledDom = waitForSettledDom;
+    this.startDomDebug = startDomDebug;
+    this.cleanupDomDebug = cleanupDomDebug;
+    this.llmProvider = llmProvider;
+    this.llmClient = llmClient;
+    this.verbose = verbose;
+  }
+
+  public async extract<T extends z.AnyZodObject>({
+    instruction,
+    schema,
+    progress = "",
+    content = {},
+    chunksSeen = [],
+    llmClient,
+    requestId,
+    domSettleTimeoutMs,
+  }: {
+    instruction: string;
+    schema: T;
+    progress?: string;
+    content?: z.infer<T>;
+    chunksSeen?: Array<number>;
+    llmClient: LLMClient;
+    requestId?: string;
+    domSettleTimeoutMs?: number;
+  }): Promise<z.infer<T>> {
+    this.logger({
+      category: "extraction",
+      message: "starting extraction",
+      level: 1,
+      auxiliary: {
+        instruction: {
+          value: instruction,
+          type: "string",
+        },
+      },
+    });
+
+    await this.waitForSettledDom(domSettleTimeoutMs);
+    await this.startDomDebug();
+    const { outputString, chunk, chunks } = await this.stagehand.page.evaluate(
+      (chunksSeen?: number[]) => window.processDom(chunksSeen ?? []),
+      chunksSeen,
+    );
+
+    this.logger({
+      category: "extraction",
+      message: "received output from processDom.",
+      auxiliary: {
+        chunk: {
+          value: chunk.toString(),
+          type: "integer",
+        },
+        chunks_left: {
+          value: (chunks.length - chunksSeen.length).toString(),
+          type: "integer",
+        },
+        chunks_total: {
+          value: chunks.length.toString(),
+          type: "integer",
+        },
+      },
+    });
+
+    const extractionResponse = await extract({
+      instruction,
+      progress,
+      previouslyExtractedContent: content,
+      domElements: outputString,
+      schema,
+      llmClient,
+      chunksSeen: chunksSeen.length,
+      chunksTotal: chunks.length,
+      requestId,
+    });
+
+    const {
+      metadata: { progress: newProgress, completed },
+      ...output
+    } = extractionResponse;
+    await this.cleanupDomDebug();
+
+    this.logger({
+      category: "extraction",
+      message: "received extraction response",
+      auxiliary: {
+        extraction_response: {
+          value: JSON.stringify(extractionResponse),
+          type: "object",
+        },
+      },
+    });
+
+    chunksSeen.push(chunk);
+
+    if (completed || chunksSeen.length === chunks.length) {
+      this.logger({
+        category: "extraction",
+        message: "got response",
+        auxiliary: {
+          extraction_response: {
+            value: JSON.stringify(extractionResponse),
+            type: "object",
+          },
+        },
+      });
+
+      return output;
+    } else {
+      this.logger({
+        category: "extraction",
+        message: "continuing extraction",
+        auxiliary: {
+          extraction_response: {
+            value: JSON.stringify(extractionResponse),
+            type: "object",
+          },
+        },
+      });
+      await this.waitForSettledDom(domSettleTimeoutMs);
+      return this.extract({
+        instruction,
+        schema,
+        progress: newProgress,
+        content: output,
+        chunksSeen,
+        llmClient,
+        domSettleTimeoutMs,
+      });
+    }
+  }
+}

package/lib/handlers/observeHandler.ts

@@ -0,0 +1,170 @@
+import { LogLine } from "../../types/log";
+import { Stagehand } from "../index";
+import { observe } from "../inference";
+import { LLMClient } from "../llm/LLMClient";
+import { LLMProvider } from "../llm/LLMProvider";
+import { generateId } from "../utils";
+import { ScreenshotService } from "../vision";
+
+export class StagehandObserveHandler {
+  private readonly stagehand: Stagehand;
+  private readonly logger: (logLine: LogLine) => void;
+  private readonly waitForSettledDom: (
+    domSettleTimeoutMs?: number,
+  ) => Promise<void>;
+  private readonly startDomDebug: () => Promise<void>;
+  private readonly cleanupDomDebug: () => Promise<void>;
+  private readonly llmProvider: LLMProvider;
+  private readonly verbose: 0 | 1 | 2;
+  private readonly llmClient: LLMClient;
+  private observations: {
+    [key: string]: {
+      result: { selector: string; description: string }[];
+      instruction: string;
+    };
+  };
+
+  constructor({
+    stagehand,
+    logger,
+    waitForSettledDom,
+    startDomDebug,
+    cleanupDomDebug,
+    llmProvider,
+    verbose,
+    llmClient,
+  }: {
+    stagehand: Stagehand;
+    logger: (logLine: LogLine) => void;
+    waitForSettledDom: (domSettleTimeoutMs?: number) => Promise<void>;
+    startDomDebug: () => Promise<void>;
+    cleanupDomDebug: () => Promise<void>;
+    llmProvider: LLMProvider;
+    verbose: 0 | 1 | 2;
+    llmClient: LLMClient;
+  }) {
+    this.stagehand = stagehand;
+    this.logger = logger;
+    this.waitForSettledDom = waitForSettledDom;
+    this.startDomDebug = startDomDebug;
+    this.cleanupDomDebug = cleanupDomDebug;
+    this.llmProvider = llmProvider;
+    this.verbose = verbose;
+    this.llmClient = llmClient;
+    this.observations = {};
+  }
+
+  private async _recordObservation(
+    instruction: string,
+    result: { selector: string; description: string }[],
+  ): Promise<string> {
+    const id = generateId(instruction);
+
+    this.observations[id] = { result, instruction };
+
+    return id;
+  }
+
+  public async observe({
+    instruction,
+    useVision,
+    fullPage,
+    llmClient,
+    requestId,
+    domSettleTimeoutMs,
+  }: {
+    instruction: string;
+    useVision: boolean;
+    fullPage: boolean;
+    llmClient: LLMClient;
+    requestId?: string;
+    domSettleTimeoutMs?: number;
+  }): Promise<{ selector: string; description: string }[]> {
+    if (!instruction) {
+      instruction = `Find elements that can be used for any future actions in the page. These may be navigation links, related pages, section/subsection links, buttons, or other interactive elements. Be comprehensive: if there are multiple elements that may be relevant for future actions, return all of them.`;
+    }
+    this.logger({
+      category: "observation",
+      message: "starting observation",
+      level: 1,
+      auxiliary: {
+        instruction: {
+          value: instruction,
+          type: "string",
+        },
+      },
+    });
+
+    await this.waitForSettledDom(domSettleTimeoutMs);
+    await this.startDomDebug();
+    let { outputString, selectorMap } = await this.stagehand.page.evaluate(
+      (fullPage: boolean) =>
+        fullPage ? window.processAllOfDom() : window.processDom([]),
+      fullPage,
+    );
+
+    let annotatedScreenshot: Buffer | undefined;
+    if (useVision === true) {
+      if (!llmClient.hasVision) {
+        this.logger({
+          category: "observation",
+          message: "Model does not support vision. Skipping vision processing.",
+          level: 1,
+          auxiliary: {
+            model: {
+              value: llmClient.modelName,
+              type: "string",
+            },
+          },
+        });
+      } else {
+        const screenshotService = new ScreenshotService(
+          this.stagehand.page,
+          selectorMap,
+          this.verbose,
+          this.logger,
+        );
+
+        annotatedScreenshot =
+          await screenshotService.getAnnotatedScreenshot(fullPage);
+        outputString = "n/a. use the image to find the elements.";
+      }
+    }
+
+    const observationResponse = await observe({
+      instruction,
+      domElements: outputString,
+      llmClient,
+      image: annotatedScreenshot,
+      requestId,
+    });
+
+    const elementsWithSelectors = observationResponse.elements.map(
+      (element) => {
+        const { elementId, ...rest } = element;
+
+        return {
+          ...rest,
+          selector: `xpath=${selectorMap[elementId][0]}`,
+        };
+      },
+    );
+
+    await this.cleanupDomDebug();
+
+    this.logger({
+      category: "observation",
+      message: "found elements",
+      level: 1,
+      auxiliary: {
+        elements: {
+          value: JSON.stringify(elementsWithSelectors),
+          type: "object",
+        },
+      },
+    });
+
+    await this._recordObservation(instruction, elementsWithSelectors);
+    return elementsWithSelectors;
+  }
+}