feat: support custom selector

package.json

@@ -1,6 +1,6 @@
 {
   "name": "web-wand",
-  "version": "2.0.0",
+  "version": "2.0.2",
   "description": "chrome extension",
   "repository": {
     "type": "git",

package.lib.json

@@ -1,6 +1,6 @@
 {
   "name": "web-wand-lib",
-  "version": "2.0.0-beta-1",
+  "version": "2.0.2",
   "description": "Helper library for Web Wand",
   "repository": {
     "type": "git",

src/helpers/parseResponse.test.ts

@@ -1,62 +1,65 @@
-import { it, expect, describe } from '@jest/globals';
-import { parseResponse } from './parseResponse';
+import { it, expect, describe } from "@jest/globals";
+import { parseResponse } from "./parseResponse";
 
-describe('parseResponse', () => {
-  it('should parse a response', () => {
+describe("parseResponse", () => {
+  it("should parse a response", () => {
     expect(
       parseResponse(
-        `<Thought>Click the "Sign Up" button</Thought>\n<Action>click(123)</Action>`
-      )
+        `<Thought>Click the "Sign Up" button</Thought>\n<Action>click(123)</Action>`,
+        false,
+      ),
     ).toEqual({
       thought: 'Click the "Sign Up" button',
-      action: 'click(123)',
+      action: "click(123)",
       parsedAction: {
-        name: 'click',
+        name: "click",
         args: {
           elementId: 123,
         },
       },
     });
   });
 
-  it('should return an error if the thought is not found', () => {
-    expect(parseResponse(`<Action>click(123)</Action>`)).toEqual({
-      error: 'Invalid response: Thought not found in the model response.',
+  it("should return an error if the thought is not found", () => {
+    expect(parseResponse(`<Action>click(123)</Action>`, false)).toEqual({
+      error: "Invalid response: Thought not found in the model response.",
     });
   });
 
-  it('should return an error if the action is not found', () => {
+  it("should return an error if the action is not found", () => {
     expect(
-      parseResponse(`<Thought>Click the "Sign Up" button</Thought>`)
+      parseResponse(`<Thought>Click the "Sign Up" button</Thought>`, true),
     ).toEqual({
-      error: 'Invalid response: Action not found in the model response.',
+      error: "Invalid response: Action not found in the model response.",
     });
   });
 
-  it('should return an error if the action is invalid', () => {
+  it("should return an error if the action is invalid", () => {
     expect(
       parseResponse(
-        `<Thought>Click the "Sign Up" button</Thought>\n<Action>click(123, 456)</Action>`
-      )
+        `<Thought>Click the "Sign Up" button</Thought>\n<Action>click(123, 456)</Action>`,
+        true,
+      ),
     ).toEqual({
       error:
         'Invalid number of arguments: Expected 1 for action "click", but got 2.',
     });
   });
 
-  it('should parse a response with multiple arguments', () => {
+  it("should parse a response with multiple arguments", () => {
     expect(
       parseResponse(
-        `<Thought>Click the "Sign Up" button</Thought>\n<Action>setValue(123, "hello")</Action>`
-      )
+        `<Thought>Click the "Sign Up" button</Thought>\n<Action>setValue(123, "hello")</Action>`,
+        true,
+      ),
     ).toEqual({
       thought: 'Click the "Sign Up" button',
       action: 'setValue(123, "hello")',
       parsedAction: {
-        name: 'setValue',
+        name: "setValue",
         args: {
           elementId: 123,
-          value: 'hello',
+          value: "hello",
         },
       },
     });
@@ -65,13 +68,14 @@ describe('parseResponse', () => {
   it("Should call the 'finish' action", () => {
     expect(
       parseResponse(
-        `<Thought>Click the "Sign Up" button</Thought>\n<Action>finish()</Action>`
-      )
+        `<Thought>Click the "Sign Up" button</Thought>\n<Action>finish()</Action>`,
+        true,
+      ),
     ).toEqual({
       thought: 'Click the "Sign Up" button',
-      action: 'finish()',
+      action: "finish()",
       parsedAction: {
-        name: 'finish',
+        name: "finish",
         args: {},
       },
     });

src/helpers/parseResponse.ts

@@ -1,4 +1,8 @@
-import { ActionPayload, availableActions } from './availableActions';
+import {
+  ActionPayload,
+  availableActions,
+  availableActionsVision,
+} from "./availableActions";
 
 export type ParsedResponseSuccess = {
   thought: string;
@@ -21,9 +25,9 @@ export function extractJsonFromMarkdown(input: string): string[] {
   let match;
   while ((match = regex.exec(input)) !== null) {
     // If 'json' is specified, add the content to the results array
-    if (match[1] === 'json') {
+    if (match[1] === "json") {
       results.push(match[2]);
-    } else if (match[2].startsWith('{')) {
+    } else if (match[2].startsWith("{")) {
       results.push(match[2]);
     }
   }
@@ -36,8 +40,8 @@ function parseFunctionCall(callString: string) {
   const matches = callString.match(functionPattern);
 
   if (!matches) {
-    console.error('Input does not match a function call pattern.', callString);
-    throw new Error('Input does not match a function call pattern.');
+    console.error("Input does not match a function call pattern.", callString);
+    throw new Error("Input does not match a function call pattern.");
   }
 
   const [, name, argsPart] = matches;
@@ -69,27 +73,30 @@ function parseFunctionCall(callString: string) {
   return { name, args };
 }
 
-export function parseResponse(text: string): ParsedResponse {
+export function parseResponse(
+  text: string,
+  isVisionModel: boolean,
+): ParsedResponse {
   let action;
   try {
     action = JSON.parse(text);
   } catch (_e) {
     try {
       action = JSON.parse(extractJsonFromMarkdown(text)[0]);
     } catch (_e) {
-      throw new Error('Response does not contain valid JSON.');
+      throw new Error("Response does not contain valid JSON.");
     }
   }
 
   if (!action.thought) {
     return {
-      error: 'Invalid response: Thought not found in the model response.',
+      error: "Invalid response: Thought not found in the model response.",
     };
   }
 
   if (!action.action) {
     return {
-      error: 'Invalid response: Action not found in the model response.',
+      error: "Invalid response: Action not found in the model response.",
     };
   }
 
@@ -99,9 +106,16 @@ export function parseResponse(text: string): ParsedResponse {
   const { name: actionName, args: argsArray } = parseFunctionCall(actionString);
   console.log(actionName, argsArray);
 
-  const availableAction = availableActions.find(
-    (action) => action.name === actionName
-  );
+  let availableAction = null;
+  if (isVisionModel) {
+    availableAction = availableActionsVision.find(
+      (action) => action.name === actionName,
+    );
+  } else {
+    availableAction = availableActions.find(
+      (action) => action.name === actionName,
+    );
+  }
 
   if (!availableAction) {
     return {