Add caching

.gitignore

@@ -56,4 +56,4 @@ screenshots/
 
 # tarball
 *.tar.gz
-*.tgz
+*.tgz

packages/shortest/src/ai/client.ts

@@ -2,6 +2,7 @@ import Anthropic from "@anthropic-ai/sdk";
 import pc from "picocolors";
 import { BrowserTool } from "../browser/core/browser-tool";
 import { AIConfig } from "../types/ai";
+import { CacheAction, CacheStep } from "../types/cache";
 import { SYSTEM_PROMPT } from "./prompts";
 import { AITools } from "./tools";
 
@@ -64,6 +65,8 @@ export class AIClient {
     _toolOutputCallback?: (name: string, input: any) => void,
   ) {
     const messages: Anthropic.Beta.Messages.BetaMessageParam[] = [];
+    // temp cache store
+    const pendingCache: Partial<{ steps?: CacheStep[] }> = {};
 
     // Log the conversation
     if (this.debugMode) {
@@ -112,18 +115,59 @@ export class AIClient {
 
         // Check for tool use
         if (response.stop_reason === "tool_use") {
-          const toolResults = response.content
-            .filter((block) => block.type === "tool_use")
-            .map((block) => {
-              const toolBlock =
-                block as Anthropic.Beta.Messages.BetaToolUseBlock;
+          const toolBlocks: Anthropic.Beta.Messages.BetaToolUseBlock[] =
+            response.content.filter((block) => block.type === "tool_use");
+
+          const toolResults = toolBlocks.map((toolBlock) => {
+            return {
+              toolBlock,
+
+              result: browserTool.execute(toolBlock.input as any),
+            };
+          });
+
+          const results = await Promise.all(toolResults.map((t) => t.result));
+
+          const getExtras = async (
+            toolBlock: Anthropic.Beta.Messages.BetaToolUseBlock,
+          ) => {
+            let extras: any = {};
+
+            // @ts-expect-error Incorrect interface on our side leads to this error
+            // @see https://docs.anthropic.com/en/docs/build-with-claude/computer-use#computer-tool:~:text=%2C%0A%20%20%20%20%20%20%20%20%7D%2C-,%22coordinate%22,-%3A%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22description
+            if (toolBlock.input.coordinate) {
+              // @ts-expect-error
+              const [x, y] = toolBlock.input.coordinate;
+
+              const componentStr =
+                await browserTool.getNormalizedComponentStringByCoords(x, y);
+
+              extras = { componentStr };
+            }
+
+            return extras;
+          };
+
+          const newCacheSteps = await Promise.all(
+            toolBlocks.map(async (_toolBlock, i) => {
+              const extras = await getExtras(toolBlocks[i]);
+
               return {
-                toolBlock,
-                result: browserTool.execute(toolBlock.input as any),
+                action: toolBlocks[i] as CacheAction,
+                reasoning: response.content.map(
+                  (block) => (block as any).text,
+                )[0],
+                result: results[i].output || null,
+                extras,
+                timestamp: Date.now(),
               };
-            });
+            }),
+          );
 
-          const results = await Promise.all(toolResults.map((t) => t.result));
+          pendingCache.steps = [
+            ...(pendingCache.steps || []),
+            ...(newCacheSteps || []),
+          ];
 
           // Log tool results
           if (this.debugMode) {
@@ -162,6 +206,7 @@ export class AIClient {
           return {
             messages,
             finalResponse: response,
+            pendingCache,
           };
         }
       } catch (error: any) {

packages/shortest/src/browser/core/browser-tool.ts

@@ -754,4 +754,95 @@ export class BrowserTool extends BaseBrowserTool {
       if (trail) trail.style.display = "none";
     });
   }
+
+  /**
+   * Retrieves normalized component string by X and Y coordinates
+   * This is primarily used to determine change in UI
+   * Playwright currently does not support such functionality
+   * @see https://github.com/microsoft/playwright/issues/13273
+   */
+  async getNormalizedComponentStringByCoords(x: number, y: number) {
+    return await this.getPage().evaluate(
+      ({ x, y, allowedAttr }) => {
+        const elem = document.elementFromPoint(x, y);
+        if (elem) {
+          // todo: test func below
+          const clone = elem.cloneNode(true) as HTMLElement;
+
+          /**
+           * Gets deepest nested child node
+           * If several nodes are on the same depth, the first node would be returned
+           */
+          function getDeepestChildNode(element: Element): HTMLElement {
+            let deepestChild = element.cloneNode(true) as HTMLElement;
+            let maxDepth = 0;
+
+            function traverse(node: any, depth: number) {
+              if (depth > maxDepth) {
+                maxDepth = depth;
+                deepestChild = node;
+              }
+
+              Array.from(node.children).forEach((child) => {
+                traverse(child, depth + 1);
+              });
+            }
+
+            traverse(deepestChild, 0);
+            return deepestChild;
+          }
+
+          const deepestNode = getDeepestChildNode(clone);
+
+          // get several parents if present
+          const node = deepestNode.parentElement
+            ? deepestNode.parentElement.parentElement
+              ? deepestNode.parentElement.parentElement
+              : deepestNode.parentElement
+            : deepestNode;
+
+          /**
+           * Recursively delete attributes from Nodes
+           */
+          function cleanAttributesRecursively(
+            element: Element,
+            options: { exceptions: string[] },
+          ) {
+            Array.from(element.attributes).forEach((attr) => {
+              if (!options.exceptions.includes(attr.name)) {
+                element.removeAttribute(attr.name);
+              }
+            });
+
+            Array.from(element.children).forEach((child) => {
+              cleanAttributesRecursively(child, options);
+            });
+          }
+
+          cleanAttributesRecursively(node, {
+            exceptions: allowedAttr,
+          });
+
+          // trim and remove white spaces
+          return node.outerHTML.trim().replace(/\s+/g, " ");
+        } else {
+          return "";
+        }
+      },
+      {
+        x,
+        y,
+        allowedAttr: [
+          "type",
+          "name",
+          "placeholder",
+          "aria-label",
+          "role",
+          "title",
+          "alt",
+          "d", // for <path> tags
+        ],
+      },
+    );
+  }
 }