api/schema: Make AI schema automatic

.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.