Skip to content

Commit

Permalink
better handling of workflow errors
Browse files Browse the repository at this point in the history
  • Loading branch information
berekuk committed Dec 27, 2024
1 parent 5a1bc05 commit 3cd6711
Show file tree
Hide file tree
Showing 10 changed files with 106 additions and 57 deletions.
12 changes: 10 additions & 2 deletions apps/hub/src/ai/data/storage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,18 @@ export function decodeDbWorkflowToClientWorkflow(
return {
id: row.id,
timestamp: row.createdAt.getTime(),
status: "error",
status: "finished",
inputs: {},
steps: [],
result: `Invalid workflow format in the database: ${e}`,
result: {
code: "",
isValid: false,
totalPrice: 0,
runTimeMs: 0,
llmRunCount: 0,
logSummary: "",
error: `Invalid workflow format in the database: ${e}`,
},
} satisfies ClientWorkflow;
}
}
64 changes: 44 additions & 20 deletions apps/hub/src/ai/data/v1_0.ts
Original file line number Diff line number Diff line change
Expand Up @@ -94,16 +94,15 @@ export const v1WorkflowSchema = z.discriminatedUnion("status", [
}),
]);

// upgrading legacy workflow to new client workflow shape
export function decodeV1_0JsonToClientWorkflow(
json: Prisma.JsonValue
): ClientWorkflow {
// `input` doesn't exist in the new ClientWorkflow shape, so we need to pull it out
const { input, ...v1Workflow } = v1WorkflowSchema.parse(json);

// upgrading legacy workflow to new client workflow shape
return {
...v1Workflow,
steps: v1Workflow.steps.map(({ outputs, ...step }) => ({
const steps: ClientWorkflow["steps"] = v1Workflow.steps.map(
({ outputs, ...step }) => ({
...step,
// modern steps in ClientWorkflow store state as an object
state:
Expand All @@ -122,22 +121,47 @@ export function decodeV1_0JsonToClientWorkflow(
}
: { kind: "PENDING" },
startTime: v1Workflow.timestamp, // old workflow steps don't have start times
})),
inputs:
input.type === "Create"
? {
prompt: {
value: input.prompt,
kind: "prompt",
id: `${v1Workflow.id}-prompt`,
},
}
: {
source: {
value: input.source,
kind: "source",
id: `${v1Workflow.id}-source`,
},
})
);

const inputs: ClientWorkflow["inputs"] =
input.type === "Create"
? {
prompt: {
value: input.prompt,
kind: "prompt",
id: `${v1Workflow.id}-prompt`,
},
}
: {
source: {
value: input.source,
kind: "source",
id: `${v1Workflow.id}-source`,
},
};

if (v1Workflow.status === "error") {
return {
...v1Workflow,
status: "finished",
result: {
code: "",
isValid: false,
totalPrice: 0,
runTimeMs: 0,
llmRunCount: 0,
logSummary: "",
error: v1Workflow.result,
},
steps,
inputs,
};
}

return {
...v1Workflow,
steps,
inputs,
};
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ export function isWorkflowOutdated(workflow: ClientWorkflow): boolean {

function getWorkflowStatusForIcon(
workflow: ClientWorkflow
): ClientWorkflow["status"] {
): "loading" | "finished" | "error" {
if (workflow.status === "loading") {
return isWorkflowOutdated(workflow) ? "error" : "loading";
}
Expand All @@ -36,5 +36,7 @@ export const WorkflowStatusIcon: FC<{ workflow: ClientWorkflow }> = ({
return <CheckCircleIcon className="text-emerald-400" size={16} />;
case "error":
return <ErrorIcon className="text-red-400" size={16} />;
default:
throw new Error(status satisfies never);
}
};
40 changes: 21 additions & 19 deletions apps/hub/src/app/ai/WorkflowViewer/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { format } from "date-fns";
import { Children, FC } from "react";

import { ClientWorkflow } from "@quri/squiggle-ai";
import { ErrorIcon, StyledTab } from "@quri/ui";
import { ErrorIcon, StyledTab, TextTooltip } from "@quri/ui";

import { commonDateFormat } from "@/lib/constants";
import { useAvailableHeight } from "@/lib/hooks/useAvailableHeight";
Expand Down Expand Up @@ -53,12 +53,26 @@ const FinishedWorkflowViewer: FC<WorkflowViewerProps<"finished">> = ({
<Header
workflow={workflow}
renderLeft={() => (
<LineSeparatedList>
<span>{(workflow.result.runTimeMs / 1000).toFixed(2)}s</span>
<span>${workflow.result.totalPrice.toFixed(2)}</span>
<span>{workflow.result.llmRunCount} LLM runs</span>
<WorkflowDate workflow={workflow} />
</LineSeparatedList>
<div className="flex items-center gap-2">
<LineSeparatedList>
<span>{(workflow.result.runTimeMs / 1000).toFixed(2)}s</span>
<span>${workflow.result.totalPrice.toFixed(2)}</span>
<span>{workflow.result.llmRunCount} LLM runs</span>
<WorkflowDate workflow={workflow} />
</LineSeparatedList>
{workflow.result.error ? (
<div className="flex items-center gap-1">
<TextTooltip text={workflow.result.error}>
<span>
<ErrorIcon
className="cursor-pointer text-red-400"
size={16}
/>
</span>
</TextTooltip>
</div>
) : null}
</div>
)}
renderRight={() => (
<div className="flex items-center gap-2">
Expand Down Expand Up @@ -137,18 +151,6 @@ export const WorkflowViewer: FC<WorkflowViewerProps> = ({
return <FinishedWorkflowViewer {...props} workflow={workflow} />;
case "loading":
return <LoadingWorkflowViewer {...props} workflow={workflow} />;
case "error":
return (
<div className="mt-2 rounded-md border border-red-300 bg-red-50 p-4">
<h3 className="mb-2 text-lg font-semibold text-red-800">
Server Error
</h3>
<p className="mb-4 text-red-700">{workflow.result}</p>
<p className="text-sm text-red-600">
Please try refreshing the page or attempt your action again.
</p>
</div>
);
default:
throw workflow satisfies never;
}
Expand Down
2 changes: 1 addition & 1 deletion apps/hub/src/app/ai/api/create/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ function aiRequestToWorkflow(request: AiRequestBody) {
const llmConfig: LlmConfig = {
llmId: request.model ?? "Claude-Sonnet",
priceLimit: 0.3,
durationLimitMinutes: 2,
durationLimitMinutes: 0.01,
messagesInHistoryToKeep: 4,
numericSteps: request.numericSteps,
styleGuideSteps: request.styleGuideSteps,
Expand Down
12 changes: 10 additions & 2 deletions apps/hub/src/app/ai/useSquiggleWorkflows.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,16 @@ export function useSquiggleWorkflows(preloadedWorkflows: AiWorkflow[]) {
} catch (error) {
updateWorkflow(id, (workflow) => ({
...workflow,
status: "error",
result: `Server error: ${error instanceof Error ? error.toString() : "Unknown error"}.`,
status: "finished",
result: {
code: "",
isValid: false,
totalPrice: 0,
runTimeMs: 0,
llmRunCount: 0,
logSummary: "",
error: `Server error: ${error instanceof Error ? error.toString() : "Unknown error"}.`,
},
}));
}
},
Expand Down
6 changes: 1 addition & 5 deletions internal-packages/ai/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ const stepUpdatedSchema = stepSchema.partial().required({
export const clientWorkflowResultSchema = z.object({
code: z.string().describe("Squiggle code snippet"),
isValid: z.boolean(),
error: z.string().optional(),
totalPrice: z.number(),
runTimeMs: z.number(),
llmRunCount: z.number(),
Expand Down Expand Up @@ -142,11 +143,6 @@ export const clientWorkflowSchema = z.discriminatedUnion("status", [
status: z.literal("finished"),
result: clientWorkflowResultSchema,
}),
z.object({
...commonClientWorkflowFields,
status: z.literal("error"),
result: z.string(),
}),
]);

export type ClientWorkflow = z.infer<typeof clientWorkflowSchema>;
20 changes: 14 additions & 6 deletions internal-packages/ai/src/workflows/Workflow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -111,10 +111,13 @@ export class Workflow<Shape extends IOShape = IOShape> {

public llmClient: LLMClient;

public error?: string;

constructor(params: WorkflowInstanceParams<Shape>) {
this.id = params.id ?? crypto.randomUUID();
this.template = params.template;
this.inputs = params.inputs;
this.error = params.error ?? undefined;

this.llmConfig = params.llmConfig ?? llmConfigDefault;
this.startTime = params.steps.at(0)?.startTime ?? Date.now();
Expand Down Expand Up @@ -201,9 +204,9 @@ export class Workflow<Shape extends IOShape = IOShape> {
payload: { step },
});

// apply the transition rule, produce the next step in PENDING state
// this code is inlined in this method, which guarantees that we always have one pending step
// (until the transition rule decides to finish)
// Apply the transition rule, produce the next step in PENDING state or stop the workflow.
// This code is inlined in this method, which guarantees that we always have one pending step
// (until the transition rule decides to finish).
const result = this.transitionRule(step, new WorkflowGuardHelpers(step));
switch (result.kind) {
case "repeat":
Expand All @@ -213,10 +216,11 @@ export class Workflow<Shape extends IOShape = IOShape> {
this.addStep(result.step.prepare(result.inputs));
break;
case "finish":
// no new steps to add
// no new steps to add, we're done
break;
case "fatal":
throw new Error(result.message);
this.error = result.message;
break;
}
}

Expand Down Expand Up @@ -299,7 +303,7 @@ export class Workflow<Shape extends IOShape = IOShape> {

getFinalResult(): ClientWorkflowResult {
const finalStep = this.getRecentStepWithCode();
const isValid = finalStep?.step.getState().kind === "DONE";
const isValid = finalStep?.step.getState().kind === "DONE" && !this.error;

// compute run time
let runTimeMs: number;
Expand All @@ -326,6 +330,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
return {
code: finalStep?.code ?? "",
isValid,
error: this.error,
totalPrice,
runTimeMs,
llmRunCount,
Expand Down Expand Up @@ -427,6 +432,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
return {
id: this.id,
templateName: this.template.name,
error: this.error ?? null,
inputIds: Object.fromEntries(
Object.entries(this.inputs).map(([key, input]) => [
key,
Expand Down Expand Up @@ -454,6 +460,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
}): Workflow<IOShape> {
const workflow = new Workflow({
id: node.id,
error: node.error ?? undefined,
template: getWorkflowTemplateByName(node.templateName),
inputs: Object.fromEntries(
Object.entries(node.inputIds).map(([key, id]) => [
Expand Down Expand Up @@ -509,4 +516,5 @@ export type SerializedWorkflow = {
inputIds: Record<string, number>;
llmConfig: LlmConfig;
stepIds: number[];
error: string | null;
};
1 change: 1 addition & 0 deletions internal-packages/ai/src/workflows/WorkflowTemplate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ export type WorkflowInstanceParams<Shape extends IOShape> = {
llmConfig?: LlmConfig;
openaiApiKey?: string;
anthropicApiKey?: string;
error?: string;
};

/**
Expand Down
2 changes: 1 addition & 1 deletion internal-packages/ai/src/workflows/controllers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,6 @@ export function getDefaultTransitionRule<Shape extends IOShape>(
}
}

return h.fatal("Unknown step");
return h.fatal(ERROR_MESSAGES.UNKNOWN_STEP);
};
}

0 comments on commit 3cd6711

Please sign in to comment.