From 74ac2734d497545a16829523b5718847ff2bc7bf Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Mon, 23 Sep 2024 14:32:47 -0700 Subject: [PATCH 1/8] If one step fails, try repeating the previous step --- packages/ai/.gitignore | 1 + packages/ai/package.json | 1 + packages/ai/src/LLMStep.ts | 4 +++- packages/ai/src/scripts/createSquiggle.ts | 4 ++++ packages/ai/src/scripts/editSquiggle.ts | 6 ++++++ packages/ai/src/workflows/SquiggleWorkflow.ts | 15 +++++++-------- packages/ai/src/workflows/Workflow.ts | 8 ++++++++ pnpm-lock.yaml | 4 +++- 8 files changed, 33 insertions(+), 10 deletions(-) diff --git a/packages/ai/.gitignore b/packages/ai/.gitignore index 9b1c8b133c..f18f97e98e 100644 --- a/packages/ai/.gitignore +++ b/packages/ai/.gitignore @@ -1 +1,2 @@ /dist +.env \ No newline at end of file diff --git a/packages/ai/package.json b/packages/ai/package.json index 629c154851..dfa4aa9d27 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -26,6 +26,7 @@ "axios": "^1.7.2", "chalk": "^5.3.0", "clsx": "^2.1.1", + "dotenv": "^16.4.5", "next": "14.2.4", "openai": "^4.56.1", "ts-node": "^10.9.2", diff --git a/packages/ai/src/LLMStep.ts b/packages/ai/src/LLMStep.ts index c532f9463a..02db56a1ad 100644 --- a/packages/ai/src/LLMStep.ts +++ b/packages/ai/src/LLMStep.ts @@ -77,6 +77,7 @@ export class LLMStepTemplate { export class LLMStepInstance { public id: string; + public inputs: Inputs; private logger: Logger; private conversationMessages: Message[] = []; public llmMetricsList: LlmMetrics[] = []; @@ -87,11 +88,12 @@ export class LLMStepInstance { constructor( public readonly template: LLMStepTemplate, public readonly workflow: Workflow, - private readonly inputs: Inputs + inputs: Inputs ) { this.startTime = Date.now(); this.id = crypto.randomUUID(); this.logger = new Logger(); + this.inputs = inputs; } getLogs(): TimestampedLogEntry[] { diff --git a/packages/ai/src/scripts/createSquiggle.ts b/packages/ai/src/scripts/createSquiggle.ts index 1ada64fe2e..6a4f545b1c 100644 --- a/packages/ai/src/scripts/createSquiggle.ts +++ b/packages/ai/src/scripts/createSquiggle.ts @@ -1,5 +1,9 @@ +import { config } from "dotenv"; + import { SquiggleWorkflow } from "../workflows/SquiggleWorkflow.js"; +config(); + async function main() { const prompt = "Generate a function that takes a list of numbers and returns the sum of the numbers"; diff --git a/packages/ai/src/scripts/editSquiggle.ts b/packages/ai/src/scripts/editSquiggle.ts index 6fec242e57..970cca0b2b 100644 --- a/packages/ai/src/scripts/editSquiggle.ts +++ b/packages/ai/src/scripts/editSquiggle.ts @@ -1,9 +1,15 @@ +import { config } from "dotenv"; + import { SquiggleWorkflow } from "../workflows/SquiggleWorkflow.js"; +config(); + async function main() { const initialCode = ` foo = 0 to 100 bar = 30 + foo + bar `; const { totalPrice, runTimeMs, llmRunCount, code, isValid, logSummary } = diff --git a/packages/ai/src/workflows/SquiggleWorkflow.ts b/packages/ai/src/workflows/SquiggleWorkflow.ts index fa2ba81418..367484ca97 100644 --- a/packages/ai/src/workflows/SquiggleWorkflow.ts +++ b/packages/ai/src/workflows/SquiggleWorkflow.ts @@ -38,17 +38,16 @@ export class SquiggleWorkflow extends ControlledWorkflow { protected configureControllerLoop(): void { this.workflow.addEventListener("stepFinished", ({ data: { step } }) => { - if (step.getState().kind !== "DONE") { - return; - } + const code = step.getOutputs()["code"]; + const hasCompleted = step.getState().kind === "DONE"; // output name is hardcoded, should we scan all outputs? - const code = step.getOutputs()["code"]; - if (code?.kind !== "code") { - return; - } + const hasCode = code?.kind === "code"; - if (code.value.type === "success") { + if (!hasCompleted || !hasCode) { + // find previous good step and run that one. + this.workflow.repeatPreviousStep(); + } else if (code.value.type === "success") { this.workflow.addStep(adjustToFeedbackStep, { prompt: this.prompt, code, diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts index 02444ea4db..e63c3130f4 100644 --- a/packages/ai/src/workflows/Workflow.ts +++ b/packages/ai/src/workflows/Workflow.ts @@ -133,6 +133,14 @@ export class Workflow { return step; } + repeatPreviousStep() { + const previousStep = this.steps.at(-1); + if (!previousStep) { + return; + } + this.addStep(previousStep.template, previousStep.inputs); + } + private async runNextStep(): Promise { const step = this.getCurrentStep(); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5b5c46c4e7..e7d1069b60 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -50,6 +50,9 @@ importers: clsx: specifier: ^2.1.1 version: 2.1.1 + dotenv: + specifier: ^16.4.5 + version: 16.4.5 next: specifier: 14.2.4 version: 14.2.4(@babel/core@7.25.2)(react-dom@18.3.1)(react@18.3.1) @@ -12718,7 +12721,6 @@ packages: /dotenv@16.4.5: resolution: {integrity: sha512-ZmdL2rui+eB2YwhsWzjInR8LldtZHGDoQ1ugH85ppHKwpUHL7j7rN0Ti9NCnGiQbhaZ11FpR+7ao1dNsmduNUg==} engines: {node: '>=12'} - dev: true /dset@3.1.2: resolution: {integrity: sha512-g/M9sqy3oHe477Ar4voQxWtaPIFw1jTdKZuomOjhCcBx9nHUNn0pu6NopuFFrTh/TRZIKEj+76vLWFu9BNKk+Q==} From d6f5a6c7781749153fe04e961215415f00501818 Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Mon, 23 Sep 2024 14:53:37 -0700 Subject: [PATCH 2/8] Minor cleanup --- packages/ai/src/workflows/SquiggleWorkflow.ts | 6 ++++-- packages/ai/src/workflows/Workflow.ts | 2 +- packages/hub/src/app/ai/WorkflowViewer/index.tsx | 12 +++++++++++- packages/hub/src/app/ai/useSquiggleWorkflows.tsx | 2 +- 4 files changed, 17 insertions(+), 5 deletions(-) diff --git a/packages/ai/src/workflows/SquiggleWorkflow.ts b/packages/ai/src/workflows/SquiggleWorkflow.ts index 367484ca97..c23d1e029c 100644 --- a/packages/ai/src/workflows/SquiggleWorkflow.ts +++ b/packages/ai/src/workflows/SquiggleWorkflow.ts @@ -44,9 +44,11 @@ export class SquiggleWorkflow extends ControlledWorkflow { // output name is hardcoded, should we scan all outputs? const hasCode = code?.kind === "code"; - if (!hasCompleted || !hasCode) { + if (!hasCompleted) { // find previous good step and run that one. - this.workflow.repeatPreviousStep(); + this.workflow.addDuplicateOfPreviousStep(); + } else if (!hasCode) { + return; } else if (code.value.type === "success") { this.workflow.addStep(adjustToFeedbackStep, { prompt: this.prompt, diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts index e63c3130f4..7d24d57e9b 100644 --- a/packages/ai/src/workflows/Workflow.ts +++ b/packages/ai/src/workflows/Workflow.ts @@ -133,7 +133,7 @@ export class Workflow { return step; } - repeatPreviousStep() { + addDuplicateOfPreviousStep() { const previousStep = this.steps.at(-1); if (!previousStep) { return; diff --git a/packages/hub/src/app/ai/WorkflowViewer/index.tsx b/packages/hub/src/app/ai/WorkflowViewer/index.tsx index 5717dba2e9..cc195c3479 100644 --- a/packages/hub/src/app/ai/WorkflowViewer/index.tsx +++ b/packages/hub/src/app/ai/WorkflowViewer/index.tsx @@ -121,7 +121,17 @@ export const WorkflowViewer: FC = ({ case "loading": return ; case "error": - return
{workflow.result}
; + return ( +
+

+ Server Error +

+

{workflow.result}

+

+ Please try refreshing the page or attempt your action again. +

+
+ ); default: throw workflow satisfies never; } diff --git a/packages/hub/src/app/ai/useSquiggleWorkflows.tsx b/packages/hub/src/app/ai/useSquiggleWorkflows.tsx index 0d48eb8343..d337ef59c4 100644 --- a/packages/hub/src/app/ai/useSquiggleWorkflows.tsx +++ b/packages/hub/src/app/ai/useSquiggleWorkflows.tsx @@ -83,7 +83,7 @@ export function useSquiggleWorkflows(initialWorkflows: SerializedWorkflow[]) { updateWorkflow(id, (workflow) => ({ ...workflow, status: "error", - result: `Error: ${error instanceof Error ? error.toString() : "Unknown error"}`, + result: `Server error: ${error instanceof Error ? error.toString() : "Unknown error"}.`, })); } }, From 477cb50ae9b681d3a3c30927bd0d394aa0a1a396 Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Mon, 23 Sep 2024 15:00:30 -0700 Subject: [PATCH 3/8] Inverse order of workflows --- packages/hub/src/app/ai/WorkflowSummaryList.tsx | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/hub/src/app/ai/WorkflowSummaryList.tsx b/packages/hub/src/app/ai/WorkflowSummaryList.tsx index d39e794513..8fd2c68f86 100644 --- a/packages/hub/src/app/ai/WorkflowSummaryList.tsx +++ b/packages/hub/src/app/ai/WorkflowSummaryList.tsx @@ -1,3 +1,4 @@ +import { orderBy } from "lodash"; import { FC } from "react"; import { SerializedWorkflow } from "@quri/squiggle-ai"; @@ -9,9 +10,11 @@ export const WorkflowSummaryList: FC<{ selectedWorkflow: SerializedWorkflow | undefined; selectWorkflow: (id: string) => void; }> = ({ workflows, selectedWorkflow, selectWorkflow }) => { + const sortedWorkflows = orderBy(workflows, ["timestamp"], ["asc"]); + return (
- {workflows.map((workflow) => ( + {sortedWorkflows.map((workflow) => ( Date: Mon, 23 Sep 2024 15:07:56 -0700 Subject: [PATCH 4/8] Begin Actions view, with last action delected --- packages/hub/src/app/ai/WorkflowSummaryList.tsx | 2 +- packages/hub/src/app/ai/WorkflowViewer/WorkflowActions.tsx | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/hub/src/app/ai/WorkflowSummaryList.tsx b/packages/hub/src/app/ai/WorkflowSummaryList.tsx index 8fd2c68f86..f41be60b59 100644 --- a/packages/hub/src/app/ai/WorkflowSummaryList.tsx +++ b/packages/hub/src/app/ai/WorkflowSummaryList.tsx @@ -13,7 +13,7 @@ export const WorkflowSummaryList: FC<{ const sortedWorkflows = orderBy(workflows, ["timestamp"], ["asc"]); return ( -
+
{sortedWorkflows.map((workflow) => ( void; }> = ({ workflow, height, onNodeClick }) => { - const [selectedNodeIndex, setSelectedNodeIndex] = useState(0); + const [selectedNodeIndex, setSelectedNodeIndex] = useState( + workflow.steps.length - 1 + ); const prevStepsLengthRef = useRef(workflow.steps.length); useEffect(() => { From f1cafc218d07f8b19fcc13c3e54004bffc78abb8 Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Mon, 23 Sep 2024 17:34:20 -0700 Subject: [PATCH 5/8] Added retry policy for minor errors in AI --- packages/ai/src/LLMStep.ts | 50 +++++++++------ .../ai/src/steps/fixCodeUntilItRunsStep.ts | 12 ++-- packages/ai/src/workflows/SquiggleWorkflow.ts | 63 +++++++++++++------ packages/ai/src/workflows/Workflow.ts | 15 +++-- .../hub/src/app/ai/WorkflowSummaryList.tsx | 2 +- 5 files changed, 92 insertions(+), 50 deletions(-) diff --git a/packages/ai/src/LLMStep.ts b/packages/ai/src/LLMStep.ts index 02db56a1ad..58363c7893 100644 --- a/packages/ai/src/LLMStep.ts +++ b/packages/ai/src/LLMStep.ts @@ -14,6 +14,8 @@ import { LlmId } from "./modelConfigs.js"; import { PromptPair } from "./prompts.js"; import { Workflow } from "./workflows/Workflow.js"; +export type ErrorType = "CRITICAL" | "MINOR"; + export type StepState = | { kind: "PENDING"; @@ -24,8 +26,9 @@ export type StepState = } | { kind: "FAILED"; + errorType: ErrorType; durationMs: number; - error: string; + message: string; }; export type StepShape< @@ -43,6 +46,7 @@ type ExecuteContext = { ): void; queryLLM(promptPair: PromptPair): Promise; log(log: LogEntry): void; + fail(errorType: ErrorType, message: string): void; // workflow: Workflow; // intentionally not exposed, but if you need it, add it here }; @@ -69,9 +73,10 @@ export class LLMStepTemplate { instantiate( workflow: Workflow, - inputs: Inputs + inputs: Inputs, + retryingStep?: LLMStepInstance | undefined ): LLMStepInstance { - return new LLMStepInstance(this, workflow, inputs); + return new LLMStepInstance(this, workflow, inputs, retryingStep); } } @@ -88,7 +93,8 @@ export class LLMStepInstance { constructor( public readonly template: LLMStepTemplate, public readonly workflow: Workflow, - inputs: Inputs + inputs: Inputs, + public retryingStep?: LLMStepInstance | undefined ) { this.startTime = Date.now(); this.id = crypto.randomUUID(); @@ -100,6 +106,10 @@ export class LLMStepInstance { return this.logger.logs; } + isRetrying(): boolean { + return !!this.retryingStep; + } + getConversationMessages(): Message[] { return this.conversationMessages; } @@ -111,7 +121,7 @@ export class LLMStepInstance { const limits = this.workflow.checkResourceLimits(); if (limits) { - this.criticalError(limits); + this.fail("CRITICAL", limits); return; } @@ -119,17 +129,24 @@ export class LLMStepInstance { setOutput: (key, value) => this.setOutput(key, value), log: (log) => this.log(log), queryLLM: (promptPair) => this.queryLLM(promptPair), + fail: (errorType, message) => this.fail(errorType, message), }; try { await this.template.execute(executeContext, this.inputs); } catch (error) { - this.criticalError( + this.fail( + "MINOR", error instanceof Error ? error.message : String(error) ); return; } - this.complete(); + + const hasFailed = (this.state as StepState).kind === "FAILED"; + + if (!hasFailed) { + this.state = { kind: "DONE", durationMs: this.calculateDuration() }; + } } getState() { @@ -169,7 +186,7 @@ export class LLMStepInstance { value: Outputs[K] | Outputs[K]["value"] ): void { if (key in this.outputs) { - this.criticalError(`Output ${key} is already set`); + this.fail("MINOR", `Output ${key} is already set`); return; } @@ -188,12 +205,13 @@ export class LLMStepInstance { this.logger.log(log); } - private criticalError(error: string) { - this.log({ type: "error", message: error }); + private fail(errorType: ErrorType, message: string) { + this.log({ type: "error", message }); this.state = { kind: "FAILED", durationMs: this.calculateDuration(), - error, + errorType, + message, }; } @@ -201,13 +219,6 @@ export class LLMStepInstance { return Date.now() - this.startTime; } - private complete() { - if (this.state.kind === "FAILED") { - return; - } - this.state = { kind: "DONE", durationMs: this.calculateDuration() }; - } - private addConversationMessage(message: Message): void { this.conversationMessages.push(message); } @@ -261,7 +272,8 @@ export class LLMStepInstance { return completion.content; } catch (error) { - this.criticalError( + this.fail( + "MINOR", `Error in queryLLM: ${error instanceof Error ? error.message : error}` ); return null; diff --git a/packages/ai/src/steps/fixCodeUntilItRunsStep.ts b/packages/ai/src/steps/fixCodeUntilItRunsStep.ts index d9845b92cf..5190932140 100644 --- a/packages/ai/src/steps/fixCodeUntilItRunsStep.ts +++ b/packages/ai/src/steps/fixCodeUntilItRunsStep.ts @@ -56,16 +56,14 @@ export const fixCodeUntilItRunsStep = new LLMStepTemplate( const completion = await context.queryLLM(promptPair); if (completion) { - const nextState = await diffCompletionContentToCode( - completion, - code.value - ); - if (nextState.ok) { - context.setOutput("code", nextState.value); + const newCode = await diffCompletionContentToCode(completion, code.value); + if (newCode.ok) { + context.setOutput("code", newCode.value); } else { + context.fail("MINOR", newCode.value); context.log({ type: "codeRunError", - error: nextState.value, + error: newCode.value, }); } } diff --git a/packages/ai/src/workflows/SquiggleWorkflow.ts b/packages/ai/src/workflows/SquiggleWorkflow.ts index c23d1e029c..7150fc60c8 100644 --- a/packages/ai/src/workflows/SquiggleWorkflow.ts +++ b/packages/ai/src/workflows/SquiggleWorkflow.ts @@ -11,6 +11,7 @@ import { LlmConfig } from "./Workflow.js"; export type SquiggleWorkflowInput = z.infer; +const MAX_RETRIES = 5; /** * This is a basic workflow for generating Squiggle code. * @@ -39,29 +40,55 @@ export class SquiggleWorkflow extends ControlledWorkflow { protected configureControllerLoop(): void { this.workflow.addEventListener("stepFinished", ({ data: { step } }) => { const code = step.getOutputs()["code"]; + const state = step.getState(); - const hasCompleted = step.getState().kind === "DONE"; - // output name is hardcoded, should we scan all outputs? - const hasCode = code?.kind === "code"; + if (this.handleFailedState(state)) return; + if (!this.isValidCodeOutput(code)) return; - if (!hasCompleted) { - // find previous good step and run that one. - this.workflow.addDuplicateOfPreviousStep(); - } else if (!hasCode) { - return; - } else if (code.value.type === "success") { - this.workflow.addStep(adjustToFeedbackStep, { - prompt: this.prompt, - code, - }); - } else { - this.workflow.addStep(fixCodeUntilItRunsStep, { - code, - }); - } + this.addNextStep(code); }); } + private handleFailedState(state: any): boolean { + if (state.kind === "FAILED") { + if (state.errorType === "MINOR") { + if (this.retryCount() < MAX_RETRIES) { + this.workflow.addRetryOfPreviousStep(); + } + } + return true; + } + return false; + } + + private retryCount(): number { + const steps = this.workflow.getSteps(); + const currentRetryingStepId = steps.at(-1)?.retryingStep?.id; + + if (!currentRetryingStepId) return 0; + + return steps.filter( + (step) => step.retryingStep?.id === currentRetryingStepId + ).length; + } + + private isValidCodeOutput(code: any): boolean { + return code?.kind === "code"; + } + + private addNextStep(code: any): void { + if (code.value.type === "success") { + this.workflow.addStep(adjustToFeedbackStep, { + prompt: this.prompt, + code, + }); + } else { + this.workflow.addStep(fixCodeUntilItRunsStep, { + code, + }); + } + } + protected configureInitialSteps(): void { if (this.input.type === "Create") { this.workflow.addStep(generateCodeStep, { prompt: this.prompt }); diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts index 7d24d57e9b..7192b2b2f6 100644 --- a/packages/ai/src/workflows/Workflow.ts +++ b/packages/ai/src/workflows/Workflow.ts @@ -121,10 +121,15 @@ export class Workflow { addStep( template: LLMStepTemplate, - inputs: Inputs + inputs: Inputs, + retryingStep?: LLMStepInstance ): LLMStepInstance { // sorry for "any"; countervariance issues - const step: LLMStepInstance = template.instantiate(this, inputs); + const step: LLMStepInstance = template.instantiate( + this, + inputs, + retryingStep + ); this.steps.push(step); this.dispatchEvent({ type: "stepAdded", @@ -133,12 +138,12 @@ export class Workflow { return step; } - addDuplicateOfPreviousStep() { - const previousStep = this.steps.at(-1); + addRetryOfPreviousStep() { + const previousStep = this.steps.filter((step) => !step.isRetrying()).at(-1); if (!previousStep) { return; } - this.addStep(previousStep.template, previousStep.inputs); + this.addStep(previousStep.template, previousStep.inputs, previousStep); } private async runNextStep(): Promise { diff --git a/packages/hub/src/app/ai/WorkflowSummaryList.tsx b/packages/hub/src/app/ai/WorkflowSummaryList.tsx index f41be60b59..d5cc4ead05 100644 --- a/packages/hub/src/app/ai/WorkflowSummaryList.tsx +++ b/packages/hub/src/app/ai/WorkflowSummaryList.tsx @@ -10,7 +10,7 @@ export const WorkflowSummaryList: FC<{ selectedWorkflow: SerializedWorkflow | undefined; selectWorkflow: (id: string) => void; }> = ({ workflows, selectedWorkflow, selectWorkflow }) => { - const sortedWorkflows = orderBy(workflows, ["timestamp"], ["asc"]); + const sortedWorkflows = orderBy(workflows, ["timestamp"], ["desc"]); return (
From 1f31ee4451b84c4ac61619ba7a7f0f363e3de4e7 Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Mon, 23 Sep 2024 17:50:40 -0700 Subject: [PATCH 6/8] Minor refactors --- packages/ai/src/workflows/SquiggleWorkflow.ts | 31 ++++--------------- packages/ai/src/workflows/Workflow.ts | 23 ++++++++++++-- 2 files changed, 26 insertions(+), 28 deletions(-) diff --git a/packages/ai/src/workflows/SquiggleWorkflow.ts b/packages/ai/src/workflows/SquiggleWorkflow.ts index 7150fc60c8..e6c6b53845 100644 --- a/packages/ai/src/workflows/SquiggleWorkflow.ts +++ b/packages/ai/src/workflows/SquiggleWorkflow.ts @@ -11,7 +11,6 @@ import { LlmConfig } from "./Workflow.js"; export type SquiggleWorkflowInput = z.infer; -const MAX_RETRIES = 5; /** * This is a basic workflow for generating Squiggle code. * @@ -42,34 +41,16 @@ export class SquiggleWorkflow extends ControlledWorkflow { const code = step.getOutputs()["code"]; const state = step.getState(); - if (this.handleFailedState(state)) return; - if (!this.isValidCodeOutput(code)) return; - - this.addNextStep(code); - }); - } - - private handleFailedState(state: any): boolean { - if (state.kind === "FAILED") { - if (state.errorType === "MINOR") { - if (this.retryCount() < MAX_RETRIES) { + if (state.kind === "FAILED") { + if (state.errorType === "MINOR") { this.workflow.addRetryOfPreviousStep(); } + return true; } - return true; - } - return false; - } - - private retryCount(): number { - const steps = this.workflow.getSteps(); - const currentRetryingStepId = steps.at(-1)?.retryingStep?.id; - - if (!currentRetryingStepId) return 0; + if (!this.isValidCodeOutput(code)) return; - return steps.filter( - (step) => step.retryingStep?.id === currentRetryingStepId - ).length; + this.addNextStep(code); + }); } private isValidCodeOutput(code: any): boolean { diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts index 7192b2b2f6..ad26920045 100644 --- a/packages/ai/src/workflows/Workflow.ts +++ b/packages/ai/src/workflows/Workflow.ts @@ -84,6 +84,8 @@ export type WorkflowEventListener = ( * See `ControlledWorkflow` for a common base class that controls the workflow * by injecting new steps based on events. */ + +const MAX_RETRIES = 5; export class Workflow { private steps: LLMStepInstance[] = []; private priceLimit: number; @@ -139,11 +141,26 @@ export class Workflow { } addRetryOfPreviousStep() { - const previousStep = this.steps.filter((step) => !step.isRetrying()).at(-1); - if (!previousStep) { + // If the last step is retrying, we want to retry the step it's retrying. If it's not retrying, we want to retry that last step. + const retryingStep = this.steps.at(-1)?.retryingStep || this.steps.at(-1); + if (!retryingStep) { + return; + } + if (this.getCurrentRetryAttempts() >= MAX_RETRIES) { return; } - this.addStep(previousStep.template, previousStep.inputs, previousStep); + this.addStep(retryingStep.template, retryingStep.inputs, retryingStep); + } + + public getCurrentRetryAttempts(): number { + const steps = this.getSteps(); + const currentRetryingStepId = steps.at(-1)?.retryingStep?.id; + + if (!currentRetryingStepId) return 0; + + return steps.filter( + (step) => step.retryingStep?.id === currentRetryingStepId + ).length; } private async runNextStep(): Promise { From d11b23f0e16c426614fa0a263bc4bdb48049a8c7 Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Mon, 23 Sep 2024 18:47:28 -0700 Subject: [PATCH 7/8] Minor refactor --- packages/ai/src/workflows/Workflow.ts | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts index ad26920045..7a7b38ca8d 100644 --- a/packages/ai/src/workflows/Workflow.ts +++ b/packages/ai/src/workflows/Workflow.ts @@ -141,26 +141,21 @@ export class Workflow { } addRetryOfPreviousStep() { - // If the last step is retrying, we want to retry the step it's retrying. If it's not retrying, we want to retry that last step. - const retryingStep = this.steps.at(-1)?.retryingStep || this.steps.at(-1); - if (!retryingStep) { - return; - } - if (this.getCurrentRetryAttempts() >= MAX_RETRIES) { + const lastStep = this.steps.at(-1); + if (!lastStep) return; + + const retryingStep = lastStep.retryingStep || lastStep; + const retryAttempts = this.getCurrentRetryAttempts(retryingStep.id); + + if (retryAttempts >= MAX_RETRIES) { return; } + this.addStep(retryingStep.template, retryingStep.inputs, retryingStep); } - public getCurrentRetryAttempts(): number { - const steps = this.getSteps(); - const currentRetryingStepId = steps.at(-1)?.retryingStep?.id; - - if (!currentRetryingStepId) return 0; - - return steps.filter( - (step) => step.retryingStep?.id === currentRetryingStepId - ).length; + public getCurrentRetryAttempts(stepId: string): number { + return this.steps.filter((step) => step.retryingStep?.id === stepId).length; } private async runNextStep(): Promise { From 4f1df55fe48d0656c18a8c7a90c78c3dc669b76d Mon Sep 17 00:00:00 2001 From: Ozzie Gooen Date: Tue, 24 Sep 2024 17:35:05 -0700 Subject: [PATCH 8/8] Implemented most CR changes --- packages/ai/src/LLMStep.ts | 8 +++-- .../ai/src/steps/fixCodeUntilItRunsStep.ts | 13 +++++--- packages/ai/src/workflows/SquiggleWorkflow.ts | 31 +++++++------------ packages/ai/src/workflows/Workflow.ts | 6 ++-- 4 files changed, 28 insertions(+), 30 deletions(-) diff --git a/packages/ai/src/LLMStep.ts b/packages/ai/src/LLMStep.ts index 58363c7893..9724324084 100644 --- a/packages/ai/src/LLMStep.ts +++ b/packages/ai/src/LLMStep.ts @@ -82,7 +82,6 @@ export class LLMStepTemplate { export class LLMStepInstance { public id: string; - public inputs: Inputs; private logger: Logger; private conversationMessages: Message[] = []; public llmMetricsList: LlmMetrics[] = []; @@ -93,7 +92,7 @@ export class LLMStepInstance { constructor( public readonly template: LLMStepTemplate, public readonly workflow: Workflow, - inputs: Inputs, + public readonly inputs: Inputs, public retryingStep?: LLMStepInstance | undefined ) { this.startTime = Date.now(); @@ -186,7 +185,10 @@ export class LLMStepInstance { value: Outputs[K] | Outputs[K]["value"] ): void { if (key in this.outputs) { - this.fail("MINOR", `Output ${key} is already set`); + this.fail( + "CRITICAL", + `Output ${key} is already set. This is a bug with the workflow code.` + ); return; } diff --git a/packages/ai/src/steps/fixCodeUntilItRunsStep.ts b/packages/ai/src/steps/fixCodeUntilItRunsStep.ts index 5190932140..4d982f7551 100644 --- a/packages/ai/src/steps/fixCodeUntilItRunsStep.ts +++ b/packages/ai/src/steps/fixCodeUntilItRunsStep.ts @@ -56,14 +56,17 @@ export const fixCodeUntilItRunsStep = new LLMStepTemplate( const completion = await context.queryLLM(promptPair); if (completion) { - const newCode = await diffCompletionContentToCode(completion, code.value); - if (newCode.ok) { - context.setOutput("code", newCode.value); + const newCodeResult = await diffCompletionContentToCode( + completion, + code.value + ); + if (newCodeResult.ok) { + context.setOutput("code", newCodeResult.value); } else { - context.fail("MINOR", newCode.value); + context.fail("MINOR", newCodeResult.value); context.log({ type: "codeRunError", - error: newCode.value, + error: newCodeResult.value, }); } } diff --git a/packages/ai/src/workflows/SquiggleWorkflow.ts b/packages/ai/src/workflows/SquiggleWorkflow.ts index e6c6b53845..e46572ecc0 100644 --- a/packages/ai/src/workflows/SquiggleWorkflow.ts +++ b/packages/ai/src/workflows/SquiggleWorkflow.ts @@ -47,27 +47,20 @@ export class SquiggleWorkflow extends ControlledWorkflow { } return true; } - if (!this.isValidCodeOutput(code)) return; - this.addNextStep(code); - }); - } + if (code === undefined || code.kind !== "code") return; - private isValidCodeOutput(code: any): boolean { - return code?.kind === "code"; - } - - private addNextStep(code: any): void { - if (code.value.type === "success") { - this.workflow.addStep(adjustToFeedbackStep, { - prompt: this.prompt, - code, - }); - } else { - this.workflow.addStep(fixCodeUntilItRunsStep, { - code, - }); - } + if (code.value.type === "success") { + this.workflow.addStep(adjustToFeedbackStep, { + prompt: this.prompt, + code, + }); + } else { + this.workflow.addStep(fixCodeUntilItRunsStep, { + code, + }); + } + }); } protected configureInitialSteps(): void { diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts index 7a7b38ca8d..72010a9afd 100644 --- a/packages/ai/src/workflows/Workflow.ts +++ b/packages/ai/src/workflows/Workflow.ts @@ -124,13 +124,13 @@ export class Workflow { addStep( template: LLMStepTemplate, inputs: Inputs, - retryingStep?: LLMStepInstance + options?: { retryingStep?: LLMStepInstance } ): LLMStepInstance { // sorry for "any"; countervariance issues const step: LLMStepInstance = template.instantiate( this, inputs, - retryingStep + options?.retryingStep ); this.steps.push(step); this.dispatchEvent({ @@ -151,7 +151,7 @@ export class Workflow { return; } - this.addStep(retryingStep.template, retryingStep.inputs, retryingStep); + this.addStep(retryingStep.template, retryingStep.inputs, { retryingStep }); } public getCurrentRetryAttempts(stepId: string): number {