fix: format

This commit is contained in:
BennyKok 2024-01-22 18:31:56 +08:00
parent 47168930dc
commit c0450b58d5

View File

@ -2,9 +2,9 @@
import { db } from "@/db/db"; import { db } from "@/db/db";
import type { import type {
MachineType, MachineType,
WorkflowRunOriginType, WorkflowRunOriginType,
WorkflowVersionType, WorkflowVersionType,
} from "@/db/schema"; } from "@/db/schema";
import { machinesTable, workflowRunsTable } from "@/db/schema"; import { machinesTable, workflowRunsTable } from "@/db/schema";
import type { APIKeyUserType } from "@/server/APIKeyBodyRequest"; import type { APIKeyUserType } from "@/server/APIKeyBodyRequest";
@ -18,226 +18,226 @@ import { v4 } from "uuid";
import { withServerPromise } from "./withServerPromise"; import { withServerPromise } from "./withServerPromise";
export const createRun = withServerPromise( export const createRun = withServerPromise(
async ({ async ({
origin, origin,
workflow_version_id, workflow_version_id,
machine_id, machine_id,
inputs, inputs,
runOrigin, runOrigin,
apiUser, apiUser,
}: { }: {
origin: string; origin: string;
workflow_version_id: string | WorkflowVersionType; workflow_version_id: string | WorkflowVersionType;
machine_id: string | MachineType; machine_id: string | MachineType;
inputs?: Record<string, string | number>; inputs?: Record<string, string | number>;
runOrigin?: WorkflowRunOriginType; runOrigin?: WorkflowRunOriginType;
apiUser?: APIKeyUserType; apiUser?: APIKeyUserType;
}) => { }) => {
const machine = const machine =
typeof machine_id === "string" typeof machine_id === "string"
? await db.query.machinesTable.findFirst({ ? await db.query.machinesTable.findFirst({
where: and( where: and(
eq(machinesTable.id, machine_id), eq(machinesTable.id, machine_id),
eq(machinesTable.disabled, false), eq(machinesTable.disabled, false),
), ),
}) })
: machine_id; : machine_id;
if (!machine) { if (!machine) {
throw new Error("Machine not found"); throw new Error("Machine not found");
} }
const workflow_version_data = const workflow_version_data =
typeof workflow_version_id === "string" typeof workflow_version_id === "string"
? await db.query.workflowVersionTable.findFirst({ ? await db.query.workflowVersionTable.findFirst({
where: eq(workflowRunsTable.id, workflow_version_id), where: eq(workflowRunsTable.id, workflow_version_id),
with: { with: {
workflow: { workflow: {
columns: { columns: {
org_id: true, org_id: true,
user_id: true, user_id: true,
}, },
}, },
}, },
}) })
: workflow_version_id; : workflow_version_id;
if (!workflow_version_data) { if (!workflow_version_data) {
throw new Error("Workflow version not found"); throw new Error("Workflow version not found");
} }
if (apiUser) if (apiUser)
if (apiUser.org_id) { if (apiUser.org_id) {
// is org api call, check org only // is org api call, check org only
if (apiUser.org_id != workflow_version_data.workflow.org_id) { if (apiUser.org_id != workflow_version_data.workflow.org_id) {
throw new Error("Workflow not found"); throw new Error("Workflow not found");
} }
} else { } else {
// is user api call, check user only // is user api call, check user only
if ( if (
apiUser.user_id != workflow_version_data.workflow.user_id && apiUser.user_id != workflow_version_data.workflow.user_id &&
workflow_version_data.workflow.org_id == null workflow_version_data.workflow.org_id == null
) { ) {
throw new Error("Workflow not found"); throw new Error("Workflow not found");
} }
} }
const workflow_api = workflow_version_data.workflow_api; const workflow_api = workflow_version_data.workflow_api;
// Replace the inputs // Replace the inputs
if (inputs && workflow_api) { if (inputs && workflow_api) {
for (const key in inputs) { for (const key in inputs) {
Object.entries(workflow_api).forEach(([_, node]) => { Object.entries(workflow_api).forEach(([_, node]) => {
if (node.inputs["input_id"] === key) { if (node.inputs["input_id"] === key) {
node.inputs["input_id"] = inputs[key]; node.inputs["input_id"] = inputs[key];
} }
}); });
} }
} }
let prompt_id: string | undefined = undefined; let prompt_id: string | undefined = undefined;
const shareData = { const shareData = {
workflow_api: workflow_api, workflow_api: workflow_api,
status_endpoint: `${origin}/api/update-run`, status_endpoint: `${origin}/api/update-run`,
file_upload_endpoint: `${origin}/api/file-upload`, file_upload_endpoint: `${origin}/api/file-upload`,
}; };
prompt_id = v4(); prompt_id = v4();
// Add to our db // Add to our db
const workflow_run = await db const workflow_run = await db
.insert(workflowRunsTable) .insert(workflowRunsTable)
.values({ .values({
id: prompt_id, id: prompt_id,
workflow_id: workflow_version_data.workflow_id, workflow_id: workflow_version_data.workflow_id,
workflow_version_id: workflow_version_data.id, workflow_version_id: workflow_version_data.id,
workflow_inputs: inputs, workflow_inputs: inputs,
machine_id: machine.id, machine_id: machine.id,
origin: runOrigin, origin: runOrigin,
}) })
.returning(); .returning();
revalidatePath(`/${workflow_version_data.workflow_id}`); revalidatePath(`/${workflow_version_data.workflow_id}`);
try { try {
switch (machine.type) { switch (machine.type) {
case "comfy-deploy-serverless": case "comfy-deploy-serverless":
case "modal-serverless": case "modal-serverless":
const _data = { const _data = {
input: { input: {
...shareData, ...shareData,
prompt_id: prompt_id, prompt_id: prompt_id,
}, },
}; };
const ___result = await fetch(`${machine.endpoint}/run`, { const ___result = await fetch(`${machine.endpoint}/run`, {
method: "POST", method: "POST",
headers: { headers: {
"Content-Type": "application/json", "Content-Type": "application/json",
}, },
body: JSON.stringify(_data), body: JSON.stringify(_data),
cache: "no-store", cache: "no-store",
}); });
console.log(___result); console.log(___result);
if (!___result.ok) if (!___result.ok)
throw new Error( throw new Error(
`Error creating run, ${ `Error creating run, ${
___result.statusText ___result.statusText
} ${await ___result.text()}`, } ${await ___result.text()}`,
); );
console.log(_data, ___result); console.log(_data, ___result);
break; break;
case "runpod-serverless": case "runpod-serverless":
const data = { const data = {
input: { input: {
...shareData, ...shareData,
prompt_id: prompt_id, prompt_id: prompt_id,
}, },
}; };
if ( if (
!machine.auth_token && !machine.auth_token &&
!machine.endpoint.includes("localhost") && !machine.endpoint.includes("localhost") &&
!machine.endpoint.includes("127.0.0.1") !machine.endpoint.includes("127.0.0.1")
) { ) {
throw new Error("Machine auth token not found"); throw new Error("Machine auth token not found");
} }
const __result = await fetch(`${machine.endpoint}/run`, { const __result = await fetch(`${machine.endpoint}/run`, {
method: "POST", method: "POST",
headers: { headers: {
"Content-Type": "application/json", "Content-Type": "application/json",
Authorization: `Bearer ${machine.auth_token}`, Authorization: `Bearer ${machine.auth_token}`,
}, },
body: JSON.stringify(data), body: JSON.stringify(data),
cache: "no-store", cache: "no-store",
}); });
console.log(__result); console.log(__result);
if (!__result.ok) if (!__result.ok)
throw new Error( throw new Error(
`Error creating run, ${ `Error creating run, ${
__result.statusText __result.statusText
} ${await __result.text()}`, } ${await __result.text()}`,
); );
console.log(data, __result); console.log(data, __result);
break; break;
case "classic": case "classic":
const body = { const body = {
...shareData, ...shareData,
prompt_id: prompt_id, prompt_id: prompt_id,
}; };
// console.log(body); // console.log(body);
const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`; const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
const _result = await fetch(comfyui_endpoint, { const _result = await fetch(comfyui_endpoint, {
method: "POST", method: "POST",
body: JSON.stringify(body), body: JSON.stringify(body),
cache: "no-store", cache: "no-store",
}); });
// console.log(_result); // console.log(_result);
if (!_result.ok) { if (!_result.ok) {
let message = `Error creating run, ${_result.statusText}`; let message = `Error creating run, ${_result.statusText}`;
try { try {
const result = await ComfyAPI_Run.parseAsync( const result = await ComfyAPI_Run.parseAsync(
await _result.json(), await _result.json(),
); );
message += ` ${result.node_errors}`; message += ` ${result.node_errors}`;
} catch (error) {} } catch (error) {}
throw new Error(message); throw new Error(message);
} }
// prompt_id = result.prompt_id; // prompt_id = result.prompt_id;
break; break;
} }
} catch (e) { } catch (e) {
console.error(e); console.error(e);
await db await db
.update(workflowRunsTable) .update(workflowRunsTable)
.set({ .set({
status: "failed", status: "failed",
}) })
.where(eq(workflowRunsTable.id, workflow_run[0].id)); .where(eq(workflowRunsTable.id, workflow_run[0].id));
throw e; throw e;
} }
// It successfully started, update the started_at time // It successfully started, update the started_at time
await db await db
.update(workflowRunsTable) .update(workflowRunsTable)
.set({ .set({
started_at: new Date(), started_at: new Date(),
}) })
.where(eq(workflowRunsTable.id, workflow_run[0].id)); .where(eq(workflowRunsTable.id, workflow_run[0].id));
return { return {
workflow_run_id: workflow_run[0].id, workflow_run_id: workflow_run[0].id,
message: "Successful workflow run", message: "Successful workflow run",
}; };
}, },
); );
export async function checkStatus(run_id: string) { export async function checkStatus(run_id: string) {
const { userId } = auth(); const { userId } = auth();
if (!userId) throw new Error("User not found"); if (!userId) throw new Error("User not found");
return await getRunsData(run_id); return await getRunsData(run_id);
} }