fix validator lookup and test

livepeer · Oct 2, 2024 · 22345df · 22345df
1 parent 110d060
commit 22345df
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/packages/api/src/controllers/generate.test.ts b/packages/api/src/controllers/generate.test.ts
@@ -91,6 +91,7 @@ describe("controllers/generate", () => {
       "image-to-video",
       "upscale",
       "segment-anything-2",
+      "llm",
     ];
     for (const api of apis) {
       aiGatewayServer.app.post(`/${api}`, async (req, res) => {
@@ -135,13 +136,18 @@ describe("controllers/generate", () => {
     textFields: Record<string, any>,
     multipartField = { name: "image", contentType: "image/png" },
   ) => {
+    const form = buildForm(textFields);
+    form.append(multipartField.name, "dummy", {
+      contentType: multipartField.contentType,
+    });
+    return form;
+  };
+
+  const buildForm = (textFields: Record<string, any>) => {
     const form = new FormData();
     for (const [k, v] of Object.entries(textFields)) {
       form.append(k, v);
     }
-    form.append(multipartField.name, "dummy", {
-      contentType: multipartField.contentType,
-    });
     return form;
   };
 
@@ -231,7 +237,7 @@ describe("controllers/generate", () => {
     it("should call the AI Gateway for generate API /llm", async () => {
       const res = await client.fetch("/beta/generate/llm", {
         method: "POST",
-        body: buildMultipartBody({}),
+        body: buildForm({ prompt: "foo" }),
       });
       expect(res.status).toBe(200);
       expect(await res.json()).toEqual({

diff --git a/packages/api/src/controllers/generate.ts b/packages/api/src/controllers/generate.ts
@@ -15,6 +15,7 @@ import { BadRequestError } from "../store/errors";
 import { fetchWithTimeout, kebabToCamel } from "../util";
 import { experimentSubjectsOnly } from "./experiment";
 import { pathJoin2 } from "./helpers";
+import validators from "../schema/validators";
 
 const AI_GATEWAY_TIMEOUT = 10 * 60 * 1000; // 10 minutes
 const RATE_LIMIT_WINDOW = 60 * 1000; // 1 minute
@@ -181,6 +182,9 @@ function registerGenerateHandler(
   if (isJSONReq) {
     payloadParsers = [validatePost(`${camelType}Params`)];
   } else {
+    if (!validators[`Body_gen${camelType}`]) {
+      camelType = type.toUpperCase();
+    }
     payloadParsers = [
       multipart.any(),
       validateFormData(`Body_gen${camelType}`),