Stagehand

from typing import TypedDict
from playwright.async_api import Page
from pydantic import BaseModel
from intuned_runtime import attempt_store, get_ai_gateway_config
from stagehand import AsyncStagehand
from stagehand.types.model_config_param import ModelConfigParam
from stagehand.types.session_start_params import Browser, BrowserLaunchOptions


class Params(TypedDict):
    category: str | None


class BookDetails(BaseModel):
    title: str
    price: str
    rating: str | None = None
    availability: str | None = None


class BooksResponse(BaseModel):
    books: list[BookDetails]


MAX_PAGES = 10


async def automation(page: Page, params: Params, **_kwargs):
    base_url, api_key = get_ai_gateway_config()
    cdp_url = attempt_store.get("cdp_url")

    model_name = "openai/gpt-5-mini"
    model_config: ModelConfigParam = {
        "model_name": model_name,
        "api_key": api_key,
        "base_url": base_url,
        "provider": "openai",
    }

    # Initialize Stagehand with act/extract/observe capabilities
    client = AsyncStagehand(
        server="local",
        model_api_key=api_key,
        local_ready_timeout_s=30.0,
    )
    launch_options: BrowserLaunchOptions = {"headless": False}
    if cdp_url is not None:
        launch_options["cdp_url"] = str(cdp_url)
    browser: Browser = {"type": "local", "launch_options": launch_options}
    session = await client.sessions.start(
        model_name=model_name,
        browser=browser,
    )
    session_id = session.data.session_id

    await page.set_viewport_size({"width": 1280, "height": 800})

    category = params.get("category")
    all_books: list[BookDetails] = []

    try:
        await client.sessions.navigate(
            id=session_id,
            url="https://books.toscrape.com",
        )

        # Navigate to category if specified using observe and act
        if category:
            await client.sessions.observe(
                id=session_id,
                instruction=f'the "{category}" category link in the sidebar',
                options={"model": model_config},
            )
            await client.sessions.act(
                id=session_id,
                input=f'Click on the "{category}" category link in the sidebar',
                options={"model": model_config},
            )

        # Collect books from all pages
        for page_num in range(1, MAX_PAGES + 1):
            # Extract all book details from the current page
            result = await client.sessions.extract(
                id=session_id,
                instruction="Extract all books on this page including title, price, rating and availability for each book",
                options={"model": model_config},
                schema={
                    "type": "object",
                    "properties": {
                        "books": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "title": {"type": "string"},
                                    "price": {"type": "string"},
                                    "rating": {"type": "string"},
                                    "availability": {"type": "string"},
                                },
                                "required": ["title", "price"],
                            },
                        }
                    },
                    "required": ["books"],
                },
            )

            books_data = BooksResponse.model_validate(result.data.result)
            all_books.extend(books_data.books)

            # Check if there's a next page and navigate to it
            if page_num < MAX_PAGES:
                try:
                    next_button = await client.sessions.observe(
                        id=session_id,
                        instruction='the "next" button or link to go to the next page',
                        options={"model": model_config},
                    )
                    if not next_button.data.result:
                        break
                    await client.sessions.act(
                        id=session_id,
                        input='Click the "next" button to go to the next page',
                        options={"model": model_config},
                    )
                except Exception:
                    break

    finally:
        await client.sessions.end(session_id)

    return BooksResponse(books=all_books)

import z from "zod";
import { Stagehand } from "@browserbasehq/stagehand";
import type { Page, BrowserContext } from "playwright";
import { attemptStore, getAiGatewayConfig } from "@intuned/runtime";

interface Params {
  category?: string;
}

const bookDetailsSchema = z.object({
  books: z.array(
    z.object({
      title: z.string(),
      price: z.string(),
      rating: z.string().optional(),
      availability: z.string().optional(),
    })
  ),
});

type BookDetails = z.infer<typeof bookDetailsSchema>["books"][number];
type BooksResponse = z.infer<typeof bookDetailsSchema>;

const MAX_PAGES = 10;

async function getWebSocketUrl(cdpUrl: string): Promise<string> {
  if (cdpUrl.includes("ws://") || cdpUrl.includes("wss://")) {
    return cdpUrl;
  }
  const versionUrl = cdpUrl.endsWith("/")
    ? `${cdpUrl}json/version`
    : `${cdpUrl}/json/version`;
  const response = await fetch(versionUrl);
  const data = await response.json();
  return data.webSocketDebuggerUrl;
}

export default async function handler(
  { category }: Params,
  page: Page,
  _context: BrowserContext
) {
  const { baseUrl, apiKey } = await getAiGatewayConfig();
  const cdpUrl = attemptStore.get("cdpUrl") as string;
  const webSocketUrl = await getWebSocketUrl(cdpUrl);

  // Initialize Stagehand with observe/act/extract capabilities
  const stagehand = new Stagehand({
    env: "LOCAL",
    localBrowserLaunchOptions: {
      cdpUrl: webSocketUrl,
      viewport: { width: 1280, height: 800 },
    },
    model: {
      modelName: "openai/gpt-5-mini",
      apiKey,
      baseURL: baseUrl,
    },
  });
  await stagehand.init();

  await page.setViewportSize({ width: 1280, height: 800 });

  const allBooks: BookDetails[] = [];

  try {
    await page.goto("https://books.toscrape.com");

    // Navigate to category if specified using observe and act
    if (category) {
      await stagehand.observe(`the "${category}" category link in the sidebar`);
      await stagehand.act(`Click on the "${category}" category link in the sidebar`);
    }

    // Collect books from all pages
    for (let pageNum = 1; pageNum <= MAX_PAGES; pageNum++) {
      // Extract all book details from the current page
      const result = await stagehand.extract(
        "Extract all books visible on the page with their complete details",
        bookDetailsSchema
      );
      allBooks.push(...(result.books ?? []));

      // Check if there's a next page and navigate to it
      if (pageNum < MAX_PAGES) {
        try {
          const nextButton = await stagehand.observe(
            'the "next" button or link to go to the next page'
          );
          if (!nextButton || nextButton.length === 0) {
            break;
          }
          await stagehand.act('Click the "next" button to go to the next page');
        } catch (e) {
          break;
        }
      }
    }
  } finally {
    await stagehand.close();
  }

  return { books: allBooks } satisfies BooksResponse;
}