model_garden.ts•3.85 kB
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { z, type Genkit, type ModelReference } from 'genkit';
import { modelRef, type GenerateRequest, type ModelAction } from 'genkit/model';
import type { GoogleAuth } from 'google-auth-library';
import OpenAI from 'openai';
import { getGenkitClientHeader } from '../../common/index.js';
import {
OpenAIConfigSchema,
openaiCompatibleModel,
} from './openai_compatibility.js';
/** @deprecated */
export const ModelGardenModelConfigSchema = OpenAIConfigSchema.extend({
location: z.string().optional(),
});
/** @deprecated */
export const llama31 = modelRef({
name: 'vertexai/llama-3.1',
info: {
label: 'Llama 3.1',
supports: {
multiturn: true,
tools: true,
media: false,
systemRole: true,
output: ['text', 'json'],
},
versions: [
'meta/llama3-405b-instruct-maas',
// 8b and 70b versions are coming soon
],
},
configSchema: ModelGardenModelConfigSchema,
version: 'meta/llama3-405b-instruct-maas',
}) as ModelReference<typeof ModelGardenModelConfigSchema>;
/** @deprecated */
export const llama32 = modelRef({
name: 'vertexai/llama-3.2',
info: {
label: 'Llama 3.2',
supports: {
multiturn: true,
tools: true,
media: true,
systemRole: true,
output: ['text', 'json'],
},
versions: ['meta/llama-3.2-90b-vision-instruct-maas'],
},
configSchema: ModelGardenModelConfigSchema,
version: 'meta/llama-3.2-90b-vision-instruct-maas',
}) as ModelReference<typeof ModelGardenModelConfigSchema>;
/**
* @deprecated use `llama31` instead
*/
export const llama3 = modelRef({
name: 'vertexai/llama3-405b',
info: {
label: 'Llama 3.1 405b',
supports: {
multiturn: true,
tools: true,
media: false,
systemRole: true,
output: ['text'],
},
versions: ['meta/llama3-405b-instruct-maas'],
},
configSchema: ModelGardenModelConfigSchema,
version: 'meta/llama3-405b-instruct-maas',
}) as ModelReference<typeof ModelGardenModelConfigSchema>;
/** @deprecated */
export const SUPPORTED_OPENAI_FORMAT_MODELS = {
'llama3-405b': llama3,
'llama-3.1': llama31,
'llama-3.2': llama32,
};
/** @deprecated */
export function modelGardenOpenaiCompatibleModel(
ai: Genkit,
name: string,
projectId: string,
location: string,
googleAuth: GoogleAuth,
baseUrlTemplate: string | undefined
): ModelAction<typeof ModelGardenModelConfigSchema> {
const model = SUPPORTED_OPENAI_FORMAT_MODELS[name];
if (!model) throw new Error(`Unsupported model: ${name}`);
if (!baseUrlTemplate) {
baseUrlTemplate =
'https://{location}-aiplatform.googleapis.com/v1beta1/projects/{projectId}/locations/{location}/endpoints/openapi';
}
const clientFactory = async (
request: GenerateRequest<typeof ModelGardenModelConfigSchema>
): Promise<OpenAI> => {
const requestLocation = request.config?.location || location;
return new OpenAI({
baseURL: baseUrlTemplate!
.replace(/{location}/g, requestLocation)
.replace(/{projectId}/g, projectId),
apiKey: (await googleAuth.getAccessToken())!,
defaultHeaders: {
'X-Goog-Api-Client': getGenkitClientHeader(),
},
});
};
return openaiCompatibleModel(ai, model, clientFactory);
}