From 0fc3e3602c0a198e5960f7dc5ad39ca586ef8e53 Mon Sep 17 00:00:00 2001
From: Luc Gagan <lucgagan@gmail.com>
Date: Tue, 14 Nov 2023 10:10:35 -0600
Subject: [PATCH] feat: add locator_evaluate (#1)

---
 src/createActions.ts | 31 +++++++++++++++++++++++++++++++
 src/prompt.ts        |  1 +
 tests/auto.spec.ts   |  9 +++++++++
 3 files changed, 41 insertions(+)

diff --git a/src/createActions.ts b/src/createActions.ts
index 6f99ca8..aa4ff46 100644
--- a/src/createActions.ts
+++ b/src/createActions.ts
@@ -50,6 +50,37 @@ export const createActions = (
         },
       },
     },
+    locator_evaluate: {
+      function: async (args: { pageFunction: string; elementId: string }) => {
+        return {
+          result: await getLocator(args.elementId).evaluate(
+            args.pageFunction
+          ),
+        };
+      },
+      description: 'Execute JavaScript code in the page, taking the matching element as an argument.',
+      name: "locator_evaluate",
+      parameters: {
+        type: 'object',
+        properties: {
+          elementId: {
+            type: "string",
+          },
+          pageFunction: {
+            type: 'string',
+            description: 'Function to be evaluated in the page context, e.g. node => node.innerText',
+          },
+        },
+      },
+      parse: (args: string) => {
+        return z
+          .object({
+            elementId: z.string(),
+            pageFunction: z.string(),
+          })
+          .parse(JSON.parse(args));
+      }
+    },
     locator_getAttribute: {
       function: async (args: { attributeName: string; elementId: string }) => {
         return {
diff --git a/src/prompt.ts b/src/prompt.ts
index 6629501..7833a3d 100644
--- a/src/prompt.ts
+++ b/src/prompt.ts
@@ -11,6 +11,7 @@ export const prompt = (message: TaskMessage) => {
 
 * When creating CSS selectors, ensure they are unique and specific enough to select only one element, even if there are multiple elements of the same type (like multiple h1 elements).
 * Avoid using generic tags like 'h1' alone. Instead, combine them with other attributes or structural relationships to form a unique selector.
+* You must not derive data from the page if you are able to do so by using one of the provided functions, e.g. locator_evaluate.
 
 Webpage snapshot:
 
diff --git a/tests/auto.spec.ts b/tests/auto.spec.ts
index 37eda75..eacd945 100644
--- a/tests/auto.spec.ts
+++ b/tests/auto.spec.ts
@@ -9,6 +9,15 @@ test("executes query", async ({ page }) => {
   expect(headerText).toBe("Hello, Rayrun!");
 });
 
+test("executes query using locator_evaluate", async ({ page }) => {
+  await page.goto("/");
+
+  const headerText = await auto("get the first letter of the header text", { page, test });
+
+  // TODO assert that we are using locator_evaluate to get the first letter
+  expect(headerText).toBe("H");
+});
+
 test("executes action", async ({ page }) => {
   await page.goto("/");