api/schema: Make AI schema automatic (#2295)

* api/schema: Add AI generate APIs to schema

Mostly translated from the AI API schema at https://github.com/livepeer/livepeer-ai-sdks/blob/main/openapi_ai_gateway.json

Converted to JSON then made several adjustments.

* api/schema: Add tag definition

* api: Move compile-schemas to schema folde

* api: Reference AI Runner schema instead of duplicating

* api: Adjust API handlers to new ai schema

* api/schema: Avoid circular references between ai-api schemas

* api/schema: Make the openapi ref indirect

The global $ref is not supported so we need to configure
them as 2 inputs.

* pkg: Allow running schema scripts from root

* api/schema: Fix TypeScript generation for AI types

Gotta remove the title or it will create duplicate identifiers
for each field.

* api/schema: Update schema README

* api/schema: Fix patching of AI input payloads

* api/schema: Stop prettying AI schema file

* api: Add test for segment-anything-2

.prettierignore

@@ -17,3 +17,5 @@ packages/api/src/schema/schema.yaml
 packages/www/static-build
 packages/www/static-build-app
 packages/api/dist-esbuild/api.js
+# auto-generated from AI Gateway schema
+packages/api/src/schema/ai-api-schema.yaml

package.json

@@ -26,6 +26,8 @@
     "release:dry-run": "lerna publish --exact --skip-git --skip-npm --cd-version prerelease --conventional-commits --yes",
     "release:alpha": "lerna publish --exact --cd-version prerelease --conventional-commits",
     "test": "cd packages/api && yarn run test",
+    "compile-schemas": "cd packages/api && yarn run compile-schemas",
+    "pull-ai-schema": "cd packages/api && yarn run pull-ai-schema",
     "dev": "touch .env.local && cp .env.local packages/www && lerna run --stream --no-sort --concurrency=999 dev",
     "updated": "lerna updated --json",
     "prettier:base": "prettier '**/*.{ts,js,css,html,md,tsx,mdx,yaml,yml}'",

packages/api/package.json

@@ -21,7 +21,8 @@
     "prepare:redoc": "redoc-cli bundle --cdn -o docs/index.html src/schema/schema.yaml",
     "prepare:type-check": "tsc --pretty --noEmit",
     "prepare": "run-s compile-schemas && run-p \"prepare:**\"",
-    "compile-schemas": "node -r esm src/compile-schemas.js",
+    "compile-schemas": "node -r esm src/schema/compile-schemas.js",
+    "pull-ai-schema": "node -r esm src/schema/pull-ai-schema.js",
     "dev-server": "run-s compile-schemas && node dist/cli.js",
     "redoc": "nodemon -w src/schema/schema.yaml -x npm run prepare:redoc",
     "siserver": "nodemon -w dist -x node -r esm dist/stream-info-service.js -e js,yaml",

packages/api/src/controllers/generate.test.ts

@@ -90,6 +90,7 @@ describe("controllers/generate", () => {
       "image-to-image",
       "image-to-video",
       "upscale",
+      "segment-anything-2",
     ];
     for (const api of apis) {
       aiGatewayServer.app.post(`/${api}`, async (req, res) => {
@@ -213,6 +214,19 @@ describe("controllers/generate", () => {
       });
       expect(aiGatewayCalls).toEqual({ upscale: 1 });
     });
+
+    it("should call the AI Gateway for generate API /segment-anything-2", async () => {
+      const res = await client.fetch("/beta/generate/segment-anything-2", {
+        method: "POST",
+        body: buildMultipartBody({}),
+      });
+      expect(res.status).toBe(200);
+      expect(await res.json()).toEqual({
+        message: "success",
+        reqContentType: expect.stringMatching("^multipart/form-data"),
+      });
+      expect(aiGatewayCalls).toEqual({ "segment-anything-2": 1 });
+    });
   });
 
   describe("validates multipart schema", () => {

packages/api/src/controllers/generate.ts

@@ -8,10 +8,11 @@ import sql from "sql-template-strings";
 import { v4 as uuid } from "uuid";
 import logger from "../logger";
 import { authorizer, validateFormData, validatePost } from "../middleware";
+import { defaultModels } from "../schema/pull-ai-schema";
 import { AiGenerateLog } from "../schema/types";
 import { db } from "../store";
 import { BadRequestError } from "../store/errors";
-import { fetchWithTimeout } from "../util";
+import { fetchWithTimeout, kebabToCamel } from "../util";
 import { experimentSubjectsOnly } from "./experiment";
 import { pathJoin2 } from "./helpers";
 
@@ -170,13 +171,24 @@ function logAiGenerateRequest(
 
 function registerGenerateHandler(
   type: AiGenerateType,
-  defaultModel: string,
   isJSONReq = false, // multipart by default
 ): RequestHandler {
   const path = `/${type}`;
-  const payloadParsers = isJSONReq
-    ? [validatePost(`${type}-payload`)]
-    : [multipart.any(), validateFormData(`${type}-payload`)];
+
+  let payloadParsers: RequestHandler[];
+  let camelType = kebabToCamel(type);
+  camelType = camelType[0].toUpperCase() + camelType.slice(1);
+  if (isJSONReq) {
+    payloadParsers = [validatePost(`${camelType}Params`)];
+  } else {
+    payloadParsers = [
+      multipart.any(),
+      validateFormData(`Body_gen${camelType}`),
+    ];
+  }
+
+  const defaultModel = defaultModels[type];
+
   return app.post(
     path,
     authorizer({}),
@@ -236,17 +248,11 @@ function registerGenerateHandler(
   );
 }
 
-registerGenerateHandler(
-  "text-to-image",
-  "SG161222/RealVisXL_V4.0_Lightning",
-  true,
-);
-registerGenerateHandler("image-to-image", "timbrooks/instruct-pix2pix");
-registerGenerateHandler(
-  "image-to-video",
-  "stabilityai/stable-video-diffusion-img2vid-xt-1-1",
-);
-registerGenerateHandler("upscale", "stabilityai/stable-diffusion-x4-upscaler");
-registerGenerateHandler("audio-to-text", "openai/whisper-large-v3");
+registerGenerateHandler("text-to-image", true);
+registerGenerateHandler("image-to-image");
+registerGenerateHandler("image-to-video");
+registerGenerateHandler("upscale");
+registerGenerateHandler("audio-to-text");
+registerGenerateHandler("segment-anything-2");
 
 export default app;

packages/api/src/schema/README.md

@@ -1,14 +1,21 @@
 # Studio API Schema
 
-Our API schema is generated from the 2 YAML files in this repository:
-
-- `api-schema.yaml` - the schema file for our public API
-- `db-schema.yaml` - the schema file for internal fields we use in our code
-
-These 2 files are deep merged on a key-by-key basis to generate the final schema
-file, with `api-schema.yaml` going first (so `db-schema` can override values).
-It is recursive, so if you want to set only 1 key in a nested object you can set
-only that and all the other fields in the objects will be left intact.
+Our API schema is generated from the 3 YAML files in this repository:
+
+- `api-schema.yaml` - The schema file for our public API. This is the base and
+  used by the code and docs/SDKs
+- `ai-api-schema.yaml` - The schema for the AI Gateway APIs. This is also used
+  by code and docs/SDKs, but is kept separate since it's pulled from
+  `livepeer/ai-worker`.
+- `db-schema.yaml` - The schema file for internal fields we use in our code.
+  This is used by the code, but not for docs/SDK generation since it contains
+  internal abstractions.
+
+These files are deep merged on a key-by-key basis to generate the final schema
+file, in the order specified above (the later can override the previous ones).
+It is recursive, so if you want to set only 1 key in a nested object you can
+specify only the nested field and all the other fields in the objects path will
+be left intact.
 
 e.g. `{a:{b:{c:d:"hello"}}}` will set only the `d` field in the `c` nested obj.
 
@@ -26,6 +33,10 @@ possible reasons to use `db-schema` instead:
   returned objects in our code (e.g. `password`, `createdByTokenId`)
 - Deprecated fields we don't want anyone using (e.g. `wowza`, `detection`)
 
+The `ai-api-schema.yaml` file should never be edited manually. Instead, run
+`yarn pull-ai-schema` to update it from the source of truth
+(`livepeer/ai-worker`).
+
 ## Outputs
 
 The schema files are used to generate the following files:
@@ -39,3 +50,14 @@ The schema files are used to generate the following files:
   our API code to validate request payloads (`middleware/validators.js`)
 
 Check `compile-schemas.js` for more details on the whole process.
+
+## AI APIs
+
+The flow for the AI Gateway schemas is:
+
+- When there are changes to the upstream AI Gateway schema, a developer can run
+  `yarn pull-ai-schema` to update the version in the repository with it.
+- The `ai-api-schema.yaml` file is merged into the code abstractions in the
+  `compile-schemas.js` script above.
+- The `ai-api-schema.yaml` file is also used on the automatic SDK and docs
+  generation to include the AI APIs.