From 67207b97768a6f0f38abd0b25ffadd40825074ac Mon Sep 17 00:00:00 2001 From: Asaf Korem Date: Fri, 18 Oct 2024 18:02:26 +0300 Subject: [PATCH] fix(StepPerformer): collect new context on retry. --- src/actions/StepPerformer.test.ts | 2 +- src/actions/StepPerformer.ts | 138 ++++++++++++++-------------- src/integration tests/index.test.ts | 4 +- 3 files changed, 73 insertions(+), 71 deletions(-) diff --git a/src/actions/StepPerformer.test.ts b/src/actions/StepPerformer.test.ts index 58a6f89..af0fbc2 100644 --- a/src/actions/StepPerformer.test.ts +++ b/src/actions/StepPerformer.test.ts @@ -238,7 +238,7 @@ describe('StepPerformer', () => { mockPromptHandler.runPrompt.mockRejectedValueOnce(error); mockPromptHandler.runPrompt.mockRejectedValueOnce(retryError); - await expect(stepPerformer.perform(intent)).rejects.toThrow(error); + await expect(stepPerformer.perform(intent)).rejects.toThrow(retryError); expect(mockPromptCreator.createPrompt).toHaveBeenCalledTimes(2); expect(mockPromptHandler.runPrompt).toHaveBeenCalledTimes(2); expect(mockCodeEvaluator.evaluate).not.toHaveBeenCalled(); diff --git a/src/actions/StepPerformer.ts b/src/actions/StepPerformer.ts index 648c57d..1489c66 100644 --- a/src/actions/StepPerformer.ts +++ b/src/actions/StepPerformer.ts @@ -5,7 +5,7 @@ import {CodeEvaluationResult, PreviousStep, PromptHandler} from '@/types'; import * as fs from 'fs'; import * as path from 'path'; import * as crypto from 'crypto'; -import {extractCodeBlock} from "@/utils/extractCodeBlock"; +import {extractCodeBlock} from '@/utils/extractCodeBlock'; export class StepPerformer { private cache: Map = new Map(); @@ -51,85 +51,87 @@ export class StepPerformer { } } - async perform(step: string, previous: PreviousStep[] = []): Promise { - // todo: replace with the user's logger - console.log("\x1b[90m%s\x1b[0m%s", "Copilot performing: ", `"${step}"`); - - // Load cache before every operation - this.loadCacheFromFile(); - - const snapshot = this.promptHandler.isSnapshotImageSupported() ? await this.snapshotManager.captureSnapshotImage() : undefined; + private async captureSnapshotAndViewHierarchy() { + const snapshot = this.promptHandler.isSnapshotImageSupported() + ? await this.snapshotManager.captureSnapshotImage() + : undefined; const viewHierarchy = await this.snapshotManager.captureViewHierarchyString(); - const isSnapshotImageAttached = - snapshot != null && this.promptHandler.isSnapshotImageSupported(); - - const cacheKey = this.generateCacheKey(step, previous, viewHierarchy); + const isSnapshotImageAttached = snapshot != null && this.promptHandler.isSnapshotImageSupported(); - let code: string | undefined = undefined; + return { snapshot, viewHierarchy, isSnapshotImageAttached }; + } - try { - if (this.cache.has(cacheKey)) { - code = this.cache.get(cacheKey); - } else { - const prompt = this.promptCreator.createPrompt( - step, - viewHierarchy, - isSnapshotImageAttached, - previous, - ); - - const promptResult = await this.promptHandler.runPrompt(prompt, snapshot); - code = extractCodeBlock(promptResult); - - this.cache.set(cacheKey, code); - this.saveCacheToFile(); - } + private async generateCode( + step: string, + previous: PreviousStep[], + snapshot: any, + viewHierarchy: string, + isSnapshotImageAttached: boolean, + ): Promise { + const cacheKey = this.generateCacheKey(step, previous, viewHierarchy); - if (!code) { - throw new Error('Failed to generate code from intent'); - } + if (this.cache.has(cacheKey)) { + return this.cache.get(cacheKey); + } else { + const prompt = this.promptCreator.createPrompt(step, viewHierarchy, isSnapshotImageAttached, previous); + const promptResult = await this.promptHandler.runPrompt(prompt, snapshot); + const code = extractCodeBlock(promptResult); - return await this.codeEvaluator.evaluate(code, this.context); - } catch (error) { - console.log("\x1b[33m%s\x1b[0m", "Failed to evaluate the code, Copilot is retrying..."); - - // Extend 'previous' array with the failure message as the result - const result = code - ? `Failed to evaluate "${step}", tried with generated code: "${code}". Validate the code against the APIs and hierarchy and let's try a different approach. If can't, return a code that throws a descriptive error.` - : `Failed to perform "${step}", could not generate prompt result. Let's try a different approach. If can't, return a code that throws a descriptive error.`; - - const newPrevious = [...previous, { - step, - code: code ?? 'undefined', - result - }]; - - const retryPrompt = this.promptCreator.createPrompt( - step, - viewHierarchy, - isSnapshotImageAttached, - newPrevious, - ); + this.cache.set(cacheKey, code); + this.saveCacheToFile(); - try { - const retryPromptResult = await this.promptHandler.runPrompt(retryPrompt, snapshot); - code = extractCodeBlock(retryPromptResult); + return code; + } + } - const result = await this.codeEvaluator.evaluate(code, this.context); + async perform(step: string, previous: PreviousStep[] = [], attempts: number = 2): Promise { + // TODO: replace with the user's logger + console.log('\x1b[90m%s\x1b[0m%s', 'Copilot performing:', `"${step}"`); - // Cache the result under the _original_ cache key - this.cache.set(cacheKey, code); - this.saveCacheToFile(); + this.loadCacheFromFile(); - return result; - } catch (retryError) { - // Log the retry error - console.error('Retry failed:', retryError); + let lastError: any = null; + let lastCode: string | undefined; - // Throw the original error if retry fails - throw error; + for (let attempt = 1; attempt <= attempts; attempt++) { + try { + console.log('\x1b[90m%s\x1b[0m', `Attempt ${attempt} for step: "${step}"`); + + // Capture updated snapshot and view hierarchy on each attempt + const { snapshot, viewHierarchy, isSnapshotImageAttached } = await this.captureSnapshotAndViewHierarchy(); + + const code = await this.generateCode(step, previous, snapshot, viewHierarchy, isSnapshotImageAttached); + lastCode = code; + + if (!code) { + throw new Error('Failed to generate code from intent'); + } + + return await this.codeEvaluator.evaluate(code, this.context); + } catch (error) { + lastError = error; + console.log('\x1b[33m%s\x1b[0m', `Attempt ${attempt} failed for step "${step}": ${error instanceof Error ? error.message : error}`); + + if (attempt < attempts) { + console.log('\x1b[33m%s\x1b[0m', 'Copilot is retrying...'); + + const resultMessage = lastCode + ? `Failed to evaluate "${step}", tried with generated code: "${lastCode}". Validate the code against the APIs and hierarchy and let's try a different approach. If can't, return a code that throws a descriptive error.` + : `Failed to perform "${step}", could not generate prompt result. Let's try a different approach. If can't, return a code that throws a descriptive error.`; + + previous = [ + ...previous, + { + step, + code: lastCode ?? 'undefined', + result: resultMessage, + }, + ]; + } } } + + throw lastError; } } diff --git a/src/integration tests/index.test.ts b/src/integration tests/index.test.ts index 37a78ac..6187140 100644 --- a/src/integration tests/index.test.ts +++ b/src/integration tests/index.test.ts @@ -156,8 +156,8 @@ describe('Copilot Integration Tests', () => { )).rejects.toThrow('Username field not found'); expect(mockPromptHandler.runPrompt).toHaveBeenCalledTimes(3); - expect(mockFrameworkDriver.captureSnapshotImage).toHaveBeenCalledTimes(2); - expect(mockFrameworkDriver.captureViewHierarchyString).toHaveBeenCalledTimes(2); + expect(mockFrameworkDriver.captureSnapshotImage).toHaveBeenCalledTimes(3); + expect(mockFrameworkDriver.captureViewHierarchyString).toHaveBeenCalledTimes(3); }); });