Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add AI LLM endpoint #2319

Merged
merged 4 commits into from
Oct 7, 2024
Merged

Add AI LLM endpoint #2319

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f4862ae
Add AI LLM endpoint
mjh1 Oct 2, 2024
110d060
Fix type generation
mjh1 Oct 2, 2024
22345df
fix validator lookup and test
mjh1 Oct 2, 2024
24ed760
Merge branch 'master' into mh/llm
mjh1 Oct 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 22 additions & 3 deletions packages/api/src/controllers/generate.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ describe("controllers/generate", () => {
"image-to-video",
"upscale",
"segment-anything-2",
"llm",
];
for (const api of apis) {
aiGatewayServer.app.post(`/${api}`, async (req, res) => {
Expand Down Expand Up @@ -145,13 +146,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;
};

Expand Down Expand Up @@ -237,6 +243,19 @@ describe("controllers/generate", () => {
});
expect(aiGatewayCalls).toEqual({ "segment-anything-2": 1 });
});

it("should call the AI Gateway for generate API /llm", async () => {
const res = await client.fetch("/beta/generate/llm", {
method: "POST",
body: buildForm({ prompt: "foo" }),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol so the LLM API is a multipart form even tho there's no files in the payload? or is it multi-modal and could receive files? if it could, I'd say we test with a dummy file here as well just in case

Copy link
Contributor Author

@mjh1 mjh1 Oct 7, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah.... it did strike me as odd too but wasn't sure. there aren't any file fields in the spec 😬 Maybe we're futureproofing..?

});
expect(res.status).toBe(200);
expect(await res.json()).toEqual({
message: "success",
reqContentType: expect.stringMatching("^multipart/form-data"),
});
expect(aiGatewayCalls).toEqual({ llm: 1 });
});
});

describe("validates multipart schema", () => {
Expand Down
5 changes: 5 additions & 0 deletions packages/api/src/controllers/generate.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -179,6 +180,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}`),
Expand Down Expand Up @@ -252,5 +256,6 @@ registerGenerateHandler("image-to-video");
registerGenerateHandler("upscale");
registerGenerateHandler("audio-to-text");
registerGenerateHandler("segment-anything-2");
registerGenerateHandler("llm");

export default app;
3 changes: 1 addition & 2 deletions packages/api/src/schema/ai-api-schema.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -608,7 +608,7 @@ components:
model_id:
type: string
title: Model Id
default: ''
default: meta-llama/Meta-Llama-3.1-8B-Instruct
system_msg:
type: string
title: System Msg
Expand All @@ -632,7 +632,6 @@ components:
type: object
required:
- prompt
- model_id
title: Body_genLLM
additionalProperties: false
Body_genSegmentAnything2:
Expand Down
4 changes: 4 additions & 0 deletions packages/api/src/schema/compile-schemas.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ function removeAllTitles(schema) {
}
}

if (schema.oneOf && Array.isArray(schema.oneOf)) {
schema.oneOf = schema.oneOf.map((item) => removeAllTitles(item));
}

return schema;
}

Expand Down
2 changes: 2 additions & 0 deletions packages/api/src/schema/db-schema.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1454,13 +1454,15 @@ components:
- image-to-video
- upscale
- segment-anything-2
- llm
request:
oneOf:
- $ref: "./ai-api-schema.yaml#/components/schemas/TextToImageParams"
- $ref: "./ai-api-schema.yaml#/components/schemas/Body_genImageToImage"
- $ref: "./ai-api-schema.yaml#/components/schemas/Body_genImageToVideo"
- $ref: "./ai-api-schema.yaml#/components/schemas/Body_genUpscale"
- $ref: "./ai-api-schema.yaml#/components/schemas/Body_genSegmentAnything2"
- $ref: "./ai-api-schema.yaml#/components/schemas/Body_genLLM"
statusCode:
type: integer
description: HTTP status code received from the AI gateway
Expand Down
1 change: 1 addition & 0 deletions packages/api/src/schema/pull-ai-schema.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ export const defaultModels = {
upscale: "stabilityai/stable-diffusion-x4-upscaler",
"audio-to-text": "openai/whisper-large-v3",
"segment-anything-2": "facebook/sam2-hiera-large",
llm: "meta-llama/Meta-Llama-3.1-8B-Instruct",
};
const schemaDir = path.resolve(__dirname, ".");
const aiSchemaUrl =
Expand Down
Loading