Skip to content

Commit

Permalink
fix(StepPerformer): collect new context on retry.
Browse files Browse the repository at this point in the history
  • Loading branch information
@asafkorem
asafkorem committed Oct 18, 2024
1 parent 127ad41 commit 67207b9
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 71 deletions.
2 changes: 1 addition & 1 deletion src/actions/StepPerformer.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
138 changes: 70 additions & 68 deletions src/actions/StepPerformer.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, any> = new Map();
Expand Down Expand Up @@ -51,85 +51,87 @@ export class StepPerformer {
}
}

async perform(step: string, previous: PreviousStep[] = []): Promise<CodeEvaluationResult> {
// 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<string> {
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<CodeEvaluationResult> {
// 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;
}
}
4 changes: 2 additions & 2 deletions src/integration tests/index.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});

Expand Down

0 comments on commit 67207b9

Please sign in to comment.