xpath -> extract (#483)

* add xpath to AXNode type

* rebase main

* allow ObserveResult in page.extract()

* pass observation to textExtract

* update extract handler

* update function signatures

* observe -> extract

* prettier

* parameterize scrolling back to top

* rm debug code

* rm logging

* rm comments

* revert

* rm comments

* rm logging

* revert unnecessary changes in extractHandler.ts

* delete handler.ts

* xpath -> node function

* shared chunk collection function

* abstract class instead of interface/shared function

* documentation

* dont scroll to bottom unless necessary

* block: start when scrolling into view

* scroll all the way to the bottom when needed

* align bottom edge for scroll into view

* subtract 25% viewport height from scroll location

* add scrollTo param

* evals

* add scrollTo param in collectDomChunks

* changeset

* scrollTo = true in processDom

* accept xpath string instead of ObserveResult

* update changeset

* xpath -> selector rename

* rm observeResult from page.extract

Author: seanmcguire12

Diff for: .changeset/cuddly-bulldogs-juggle.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": minor
+---
+
+you can now do targetted extraction by passing an xpath string into extract. This limits the dom processing step to a target element, reducing tokens and increasing speed.

Diff for: evals/evals.config.json

@@ -247,6 +247,14 @@
     {
       "name": "observe_iframes2",
       "categories": ["observe"]
+    },
+    {
+      "name": "extract_hamilton_weather",
+      "categories": ["text_extract"]
+    },
+    {
+      "name": "extract_regulations_table",
+      "categories": ["text_extract"]
     }
   ]
 }

Diff for: evals/tasks/extract_hamilton_weather.ts

@@ -0,0 +1,79 @@
+import { EvalFunction } from "@/types/evals";
+import { initStagehand } from "@/evals/initStagehand";
+import { z } from "zod";
+
+export const extract_hamilton_weather: EvalFunction = async ({
+  modelName,
+  logger,
+  useTextExtract,
+}) => {
+  const { stagehand, initResponse } = await initStagehand({
+    modelName,
+    logger,
+  });
+
+  const { debugUrl, sessionUrl } = initResponse;
+
+  try {
+    await stagehand.page.goto("https://hamilton-weather.surge.sh/");
+    const xpath =
+      "/html/body[1]/div[5]/main[1]/article[1]/div[6]/div[2]/div[1]/table[1]";
+
+    const weatherData = await stagehand.page.extract({
+      instruction: "extract the weather data for Sun, Feb 23 at 11PM",
+      schema: z.object({
+        temperature: z.string(),
+        weather_description: z.string(),
+        wind: z.string(),
+        humidity: z.string(),
+        barometer: z.string(),
+        visibility: z.string(),
+      }),
+      modelName,
+      useTextExtract,
+      selector: xpath,
+    });
+
+    // Define the expected weather data
+    const expectedWeatherData = {
+      temperature: "27 °F",
+      weather_description: "Light snow. Overcast.",
+      wind: "6 mph",
+      humidity: "93%",
+      barometer: '30.07 "Hg',
+      visibility: "10 mi",
+    };
+
+    // Check that every field matches the expected value
+    const isWeatherCorrect =
+      weatherData.temperature === expectedWeatherData.temperature &&
+      weatherData.weather_description ===
+        expectedWeatherData.weather_description &&
+      weatherData.wind === expectedWeatherData.wind &&
+      weatherData.humidity === expectedWeatherData.humidity &&
+      weatherData.barometer === expectedWeatherData.barometer &&
+      weatherData.visibility === expectedWeatherData.visibility;
+
+    await stagehand.close();
+
+    return {
+      _success: isWeatherCorrect,
+      weatherData,
+      debugUrl,
+      sessionUrl,
+      logs: logger.getLogs(),
+    };
+  } catch (error) {
+    console.error("Error or timeout occurred:", error);
+
+    await stagehand.close();
+
+    return {
+      _success: false,
+      error: JSON.parse(JSON.stringify(error, null, 2)),
+      debugUrl,
+      sessionUrl,
+      logs: logger.getLogs(),
+    };
+  }
+};

Diff for: evals/tasks/extract_regulations_table.ts

@@ -0,0 +1,96 @@
+import { EvalFunction } from "@/types/evals";
+import { initStagehand } from "@/evals/initStagehand";
+import { z } from "zod";
+
+export const extract_regulations_table: EvalFunction = async ({
+  modelName,
+  logger,
+  useTextExtract,
+}) => {
+  const { stagehand, initResponse } = await initStagehand({
+    modelName,
+    logger,
+  });
+
+  const { debugUrl, sessionUrl } = initResponse;
+
+  try {
+    await stagehand.page.goto(
+      "https://www.ncc.gov.ng/technical-regulation/standards/numbering",
+    );
+
+    const xpath =
+      "/html/body/div[2]/section[4]/div/div/div[1]/main/div[2]/div/div/div/div/div/div/div[2]/div[2]/div[3]/div[1]";
+
+    const allottees = await stagehand.page.extract({
+      instruction:
+        "Extract ALL of the Allottees and their corresponding name, area, and area code.",
+      schema: z.object({
+        allottee_list: z.array(
+          z.object({
+            allottee_name: z.string(),
+            area: z.string(),
+            area_code: z.string(),
+            access_code: z.string(),
+          }),
+        ),
+      }),
+      modelName,
+      useTextExtract,
+      selector: xpath,
+    });
+
+    // Define the expected weather data
+    const allottees_expected_first = {
+      allottee_name: "101 Communications Limited",
+      area: "Lagos",
+      area_code: "0201",
+      access_code: "249",
+    };
+
+    const allottees_expected_last = {
+      allottee_name: "21st Century Technologies Limited",
+      area: "Lagos",
+      area_code: "0201",
+      access_code: "278",
+    };
+
+    const expected_length = 10;
+
+    const allotteeList = allottees.allottee_list;
+
+    // Check that the first entry, last entry, and total number match expectations
+    const isFirstCorrect =
+      JSON.stringify(allotteeList[0]) ===
+      JSON.stringify(allottees_expected_first);
+    const isLastCorrect =
+      JSON.stringify(allotteeList[allotteeList.length - 1]) ===
+      JSON.stringify(allottees_expected_last);
+    const isLengthCorrect = allotteeList.length === expected_length;
+
+    const isRegulationsCorrect =
+      isFirstCorrect && isLastCorrect && isLengthCorrect;
+
+    await stagehand.close();
+
+    return {
+      _success: isRegulationsCorrect,
+      regulationsData: allottees,
+      debugUrl,
+      sessionUrl,
+      logs: logger.getLogs(),
+    };
+  } catch (error) {
+    console.error("Error or timeout occurred:", error);
+
+    await stagehand.close();
+
+    return {
+      _success: false,
+      error: JSON.parse(JSON.stringify(error, null, 2)),
+      debugUrl,
+      sessionUrl,
+      logs: logger.getLogs(),
+    };
+  }
+};

Diff for: lib/StagehandPage.ts

@@ -246,8 +246,10 @@ export class StagehandPage {
             };
           }
           if (prop === "extract") {
-            return async (options: ExtractOptions<z.AnyZodObject>) => {
-              return this.extract(options);
+            return async (
+              instructionOrOptions: string | ExtractOptions<z.AnyZodObject>,
+            ) => {
+              return this.extract(instructionOrOptions);
             };
           }
           if (prop === "observe") {
@@ -566,8 +568,17 @@ export class StagehandPage {
       modelClientOptions,
       domSettleTimeoutMs,
       useTextExtract,
+      selector,
     } = options;
 
+    // Throw a NotImplementedError if the user passed in an `xpath`
+    // and `useTextExtract` is false
+    if (selector && useTextExtract !== true) {
+      throw new Error(
+        "NotImplementedError: Passing an xpath into extract is only supported when `useTextExtract: true`.",
+      );
+    }
+
     if (this.api) {
       return this.api.extract<T>(options);
     }
@@ -605,6 +616,7 @@ export class StagehandPage {
         requestId,
         domSettleTimeoutMs,
         useTextExtract,
+        selector,
       })
       .catch((e) => {
         this.stagehand.log({

Diff for: lib/dom/DomChunk.ts

@@ -0,0 +1,6 @@
+export interface DomChunk {
+  startOffset: number;
+  endOffset: number;
+  outputString: string;
+  selectorMap: Record<number, string[]>;
+}

Diff for: lib/dom/ElementContainer.ts

@@ -1,8 +1,33 @@
 import { StagehandContainer } from "./StagehandContainer";
 
-export class ElementContainer implements StagehandContainer {
-  constructor(private el: HTMLElement) {}
+/**
+ * The ElementContainer class is a container implementation for a specific
+ * HTML element.
+ *
+ * Unlike `GlobalPageContainer`, which manages the entire page,
+ * this class focuses on one particular `HTMLElement`. Operations
+ * such as `scrollTo` and `scrollIntoView` apply to that element
+ * rather than `window`.
+ */
+export class ElementContainer extends StagehandContainer {
+  /**
+   * Creates an instance of `ElementContainer` tied to a specific element.
+   * @param el - The scrollable `HTMLElement` that this container controls.
+   */
+  constructor(private el: HTMLElement) {
+    super();
+  }
+
+  public getRootElement(): HTMLElement {
+    return this.el;
+  }
 
+  /**
+   * Retrieves the height of the visible viewport within this element
+   * (`el.clientHeight`).
+   *
+   * @returns The visible (client) height of the element, in pixels.
+   */
   public getViewportHeight(): number {
     return this.el.clientHeight;
   }
@@ -11,12 +36,50 @@ export class ElementContainer implements StagehandContainer {
     return this.el.scrollHeight;
   }
 
+  /**
+   * Returns the element's current vertical scroll offset.
+   */
+  public getScrollPosition(): number {
+    return this.el.scrollTop;
+  }
+
+  /**
+   * Smoothly scrolls this element to the specified vertical offset, and
+   * waits for the scrolling to complete.
+   *
+   * @param offset - The scroll offset (in pixels) from the top of the element.
+   * @returns A promise that resolves once scrolling is finished.
+   */
   public async scrollTo(offset: number): Promise<void> {
-    await new Promise((resolve) => setTimeout(resolve, 1500));
-    this.el.scrollTo({ top: offset, left: 0, behavior: "smooth" });
+    this.el.scrollTo({ top: offset, behavior: "smooth" });
+    await this.waitForScrollEnd();
+  }
+
+  /**
+   * Scrolls this element so that the given `element` is visible, or
+   * scrolls to the top if none is provided. Smoothly animates the scroll
+   * and waits until it finishes.
+   *
+   * @param element - The child element to bring into view. If omitted, scrolls to top.
+   * @returns A promise that resolves once scrolling completes.
+   */
+  public async scrollIntoView(element?: HTMLElement): Promise<void> {
+    if (!element) {
+      this.el.scrollTo({ top: 0, behavior: "smooth" });
+    } else {
+      element.scrollIntoView({ behavior: "smooth", block: "end" });
+    }
     await this.waitForScrollEnd();
   }
 
+  /**
+   * Internal helper that waits until scrolling in this element has
+   * fully stopped. It listens for scroll events on the element,
+   * resetting a short timer every time a scroll occurs, and resolves
+   * once there's no scroll for ~100ms.
+   *
+   * @returns A promise that resolves when scrolling has finished.
+   */
   private async waitForScrollEnd(): Promise<void> {
     return new Promise<void>((resolve) => {
       let scrollEndTimer: number;