change to structured extract, better prompts for tools

Author: Kylejeong2

src/sessionManager.ts

@@ -6,7 +6,6 @@ import { AvailableModel, Stagehand } from "@browserbasehq/stagehand";
 import type { Config } from "../config.js"; 
 import type { Cookie } from "playwright-core";
 
-// Define the type for a session object
 export type BrowserSession = {
   browser: Browser;
   page: Page;

src/tools/extract.ts

@@ -4,14 +4,37 @@ import type { Context } from "../context.js";
 import type { ToolActionResult } from "../context.js";
 
 const ExtractInputSchema = z.object({
-  random_string: z.string().optional().describe("Dummy parameter for no-parameter tools"),
+  instruction: z.string().describe(
+    "The specific instruction for what information to extract from the current page. " +
+    "Be as detailed and specific as possible about what you want to extract. For example: " +
+    "'Extract all product names and prices from the listing page' or 'Get the article title, " +
+    "author, and publication date from this blog post'. The more specific your instruction, " +
+    "the better the extraction results will be. Avoid vague instructions like 'get everything' " +
+    "or 'extract the data'. Instead, be explicit about the exact elements, text, or information you need."
+  ),
+  schema: z.string().describe(
+    "A JSON schema string that defines the exact structure and format of the data you want to extract. " +
+    "This schema should be a valid JSON string that describes the expected output format. For example: " +
+    "'{\"type\": \"object\", \"properties\": {\"title\": {\"type\": \"string\"}, \"price\": {\"type\": \"number\"}}}' " +
+    "or '{\"type\": \"array\", \"items\": {\"type\": \"object\", \"properties\": {\"name\": {\"type\": \"string\"}, " +
+    "\"description\": {\"type\": \"string\"}}}}'. The schema helps ensure the extracted data is properly " +
+    "formatted and structured. If the schema is invalid JSON, extraction will proceed without schema validation. " +
+    "Use this to specify exactly how you want the extracted information organized and typed."
+  ),
 });
 
 type ExtractInput = z.infer<typeof ExtractInputSchema>;
 
 const extractSchema: ToolSchema<typeof ExtractInputSchema> = {
   name: "stagehand_extract",
-  description: "Extracts all of the text from the current page.",
+  description: 
+    "Extracts structured information and text content from the current web page based on specific instructions " +
+    "and a defined schema. This tool is ideal for scraping data, gathering information, or pulling specific " +
+    "content from web pages. Use this tool when you need to get text content, data, or information from a page " +
+    "rather than interacting with elements. For interactive elements like buttons, forms, or clickable items, " +
+    "use the observe tool instead. The extraction works best when you provide clear, specific instructions " +
+    "about what to extract and a well-defined JSON schema for the expected output format. This ensures " +
+    "the extracted data is properly structured and usable.",
   inputSchema: ExtractInputSchema,
 };
 
@@ -21,40 +44,25 @@ async function handleExtract(
 ): Promise<ToolResult> {
   const action = async (): Promise<ToolActionResult> => {
     try {
-      const page = await context.getActivePage();
-      if (!page) {
-        throw new Error("No active page available");
-      }
+      const stagehand = await context.getStagehand();
+      
+      let parsedSchema = null;
+        try {
+          parsedSchema = JSON.parse(params.schema);
+        } catch (error) {
+          throw new Error(`Invalid schema format: ${error instanceof Error ? error.message : 'Unknown error'}`);
+        }
 
-      const bodyText = await page.evaluate(() => document.body.innerText);
-
-      const content = bodyText
-        .split("\n")
-        .map((line) => line.trim())
-        .filter((line) => {
-          if (!line) return false;
-
-          if (
-            (line.includes("{") && line.includes("}")) ||
-            line.includes("@keyframes") || // Remove CSS animations
-            line.match(/^\.[a-zA-Z0-9_-]+\s*{/) || // Remove CSS lines starting with .className {
-            line.match(/^[a-zA-Z-]+:[a-zA-Z0-9%\s\(\)\.,-]+;$/) // Remove lines like "color: blue;" or "margin: 10px;"
-          ) {
-            return false;
-          }
-          return true;
-        })
-        .map((line) => {
-          return line.replace(/\\u([0-9a-fA-F]{4})/g, (_, hex) =>
-            String.fromCharCode(parseInt(hex, 16))
-          );
-        });
+      const extraction = await stagehand.page.extract({
+        instruction: params.instruction,
+        schema: parsedSchema // If schema not properly formatted, will just extract without given schema
+      });
 
       return {
         content: [
           {
             type: "text",
-            text: `Extracted content:\n${content.join("\n")}`,
+            text: `Extracted content:\n${extraction.join("\n")}`,
           },
         ],
       };

src/tools/observe.ts

@@ -5,7 +5,18 @@ import type { ToolActionResult } from "../context.js";
 
 const ObserveInputSchema = z.object({
   instruction: z.string().describe(
-    "Instruction for observation (e.g., 'find the login button'). This instruction must be extremely specific."
+    "Detailed instruction for what specific elements or components to observe on the web page. " +
+    "This instruction must be extremely specific and descriptive. For example: 'Find the red login button " +
+    "in the top right corner', 'Locate the search input field with placeholder text', or 'Identify all " +
+    "clickable product cards on the page'. The more specific and detailed your instruction, the better " +
+    "the observation results will be. Avoid generic instructions like 'find buttons' or 'see elements'. " +
+    "Instead, describe the visual characteristics, location, text content, or functionality of the elements " +
+    "you want to observe. This tool is designed to help you identify interactive elements that you can " +
+    "later use with the act tool for performing actions like clicking, typing, or form submission."
+  ),
+  returnAction: z.boolean().optional().describe(
+    "Whether to return the action to perform on the element. If true, the action will be returned as a string. " +
+    "If false, the action will not be returned."
   ),
 });
 
@@ -14,7 +25,14 @@ type ObserveInput = z.infer<typeof ObserveInputSchema>;
 const observeSchema: ToolSchema<typeof ObserveInputSchema> = {
   name: "stagehand_observe",
   description:
-    "Observes elements on the web page. Use this tool to observe elements that you can later use in an action. Use observe instead of extract when dealing with actionable (interactable) elements rather than text. More often than not, you'll want to use extract instead of observe when dealing with scraping or extracting structured text.",
+    "Observes and identifies specific interactive elements on the current web page that can be used for subsequent actions. " +
+    "This tool is specifically designed for finding actionable (interactable) elements such as buttons, links, form fields, " +
+    "dropdowns, checkboxes, and other UI components that you can interact with. Use this tool when you need to locate " +
+    "elements before performing actions with the act tool. DO NOT use this tool for extracting text content or data - " +
+    "use the extract tool instead for that purpose. The observe tool returns detailed information about the identified " +
+    "elements including their properties, location, and interaction capabilities. This information can then be used " +
+    "to craft precise actions. The more specific your observation instruction, the more accurate the element identification " +
+    "will be. Think of this as your 'eyes' on the page to find exactly what you need to interact with.",
   inputSchema: ObserveInputSchema,
 };
 
@@ -28,7 +46,7 @@ async function handleObserve(
 
       const observations = await stagehand.page.observe({
         instruction: params.instruction,
-        returnAction: false,
+        returnAction: params.returnAction,
       });
 
       return {

src/tools/screenshot.ts

@@ -5,7 +5,7 @@ import type { ToolActionResult } from "../context.js";
 import { screenshots } from "../resources.js";
 
 const ScreenshotInputSchema = z.object({
-  random_string: z.string().optional().describe("Dummy parameter for no-parameter tools"),
+  name: z.string().optional().describe("The name of the screenshot"),
 });
 
 type ScreenshotInput = z.infer<typeof ScreenshotInputSchema>;
@@ -34,7 +34,11 @@ async function handleScreenshot(
 
       // Convert buffer to base64 string and store in memory
       const screenshotBase64 = screenshotBuffer.toString("base64");
-      const name = `screenshot-${new Date()
+      const name = 
+      `screenshot-${params.name}-${new Date()
+        .toISOString()
+        .replace(/:/g, "-")}` || 
+      `screenshot-${new Date()
         .toISOString()
         .replace(/:/g, "-")}`;
       screenshots.set(name, screenshotBase64);