utils.ts•11.1 kB
/**
* Copyright 2025 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 { GenkitError } from 'genkit';
import { GoogleAuth } from 'google-auth-library';
import type {
ClientOptions,
ExpressClientOptions,
GlobalClientOptions,
RegionalClientOptions,
VertexPluginOptions,
} from './types.js';
export {
checkModelName,
checkSupportedMimeType,
cleanSchema,
extractMedia,
extractMimeType,
extractText,
extractVersion,
modelName,
} from '../common/utils.js';
const CLOUD_PLATFORM_OAUTH_SCOPE =
'https://www.googleapis.com/auth/cloud-platform';
function parseFirebaseProjectId(): string | undefined {
if (!process.env.FIREBASE_CONFIG) return undefined;
try {
return JSON.parse(process.env.FIREBASE_CONFIG).projectId as string;
} catch {
return undefined;
}
}
let __mockDerivedOptions: ClientOptions | undefined = undefined;
function setMockDerivedOptions(options: ClientOptions | undefined): void {
__mockDerivedOptions = options;
}
export const TEST_ONLY = { setMockDerivedOptions };
export async function getDerivedOptions(
options?: VertexPluginOptions,
AuthClass: typeof GoogleAuth = GoogleAuth // Injectable testing
): Promise<ClientOptions> {
if (__mockDerivedOptions) {
return Promise.resolve(__mockDerivedOptions);
}
// Figure out the type of preferred options if possible
// The order of the if statements is important.
if (options?.location == 'global') {
return await getGlobalDerivedOptions(AuthClass, options);
} else if (options?.location) {
return await getRegionalDerivedOptions(AuthClass, options);
} else if (options?.apiKey !== undefined) {
// apiKey = false still indicates apiKey expectation
return getExpressDerivedOptions(options);
}
// If we got here then we're relying on environment variables.
// Try regional first, it's the most common usage.
try {
const regionalOptions = await getRegionalDerivedOptions(AuthClass, options);
return regionalOptions;
} catch (e: unknown) {
/* no-op - try global next */
}
try {
const globalOptions = await getGlobalDerivedOptions(AuthClass, options);
return globalOptions;
} catch (e: unknown) {
/* no-op - try express last */
}
try {
const expressOptions = getExpressDerivedOptions(options);
return expressOptions;
} catch (e: unknown) {
/* no-op */
}
// We did not have enough information in the options or in environment variables
// to properly determine client options.
throw new GenkitError({
status: 'INVALID_ARGUMENT',
message:
'Unable to determine client options. Please set either apiKey or projectId and location',
});
}
async function getGlobalDerivedOptions(
AuthClass: typeof GoogleAuth,
options?: VertexPluginOptions
): Promise<GlobalClientOptions> {
let authOptions = options?.googleAuth;
let authClient: GoogleAuth;
const providedProjectId =
options?.projectId ||
process.env.GCLOUD_PROJECT ||
parseFirebaseProjectId();
if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
const serviceAccountCreds = JSON.parse(
process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
);
authOptions = {
credentials: serviceAccountCreds,
scopes: [CLOUD_PLATFORM_OAUTH_SCOPE],
projectId: providedProjectId,
};
authClient = new AuthClass(authOptions);
} else {
authClient = new AuthClass(
authOptions ?? {
scopes: [CLOUD_PLATFORM_OAUTH_SCOPE],
projectId: providedProjectId,
}
);
}
const projectId =
options?.projectId ||
process.env.GCLOUD_PROJECT ||
(await authClient.getProjectId());
if (!projectId) {
throw new Error(
`VertexAI Plugin is missing the 'project' configuration. Please set the 'GCLOUD_PROJECT' environment variable or explicitly pass 'project' into genkit config.`
);
}
const clientOpt: GlobalClientOptions = {
kind: 'global',
location: 'global',
projectId,
authClient,
};
if (options?.apiKey) {
clientOpt.apiKey = options.apiKey;
}
return clientOpt;
}
function getExpressDerivedOptions(
options?: VertexPluginOptions
): ExpressClientOptions {
const apiKey = checkApiKey(options?.apiKey);
return {
kind: 'express',
apiKey,
};
}
async function getRegionalDerivedOptions(
AuthClass: typeof GoogleAuth,
options?: VertexPluginOptions
): Promise<RegionalClientOptions> {
let authOptions = options?.googleAuth;
let authClient: GoogleAuth;
const providedProjectId =
options?.projectId ||
process.env.GCLOUD_PROJECT ||
parseFirebaseProjectId();
if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
const serviceAccountCreds = JSON.parse(
process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
);
authOptions = {
credentials: serviceAccountCreds,
scopes: [CLOUD_PLATFORM_OAUTH_SCOPE],
projectId: providedProjectId,
};
authClient = new AuthClass(authOptions);
} else {
authClient = new AuthClass(
authOptions ?? {
scopes: [CLOUD_PLATFORM_OAUTH_SCOPE],
projectId: providedProjectId,
}
);
}
const projectId =
options?.projectId ||
process.env.GCLOUD_PROJECT ||
(await authClient.getProjectId());
const location =
options?.location || process.env.GCLOUD_LOCATION || 'us-central1';
if (!location) {
throw new Error(
`VertexAI Plugin is missing the 'location' configuration. Please set the 'GCLOUD_LOCATION' environment variable or explicitly pass 'location' into genkit config.`
);
}
if (!projectId) {
throw new Error(
`VertexAI Plugin is missing the 'project' configuration. Please set the 'GCLOUD_PROJECT' environment variable or explicitly pass 'project' into genkit config.`
);
}
const clientOpt: RegionalClientOptions = {
kind: 'regional',
location,
projectId,
authClient,
};
if (options?.apiKey) {
clientOpt.apiKey = options.apiKey;
}
return clientOpt;
}
/**
* Retrieves an API key from environment variables.
*
* @returns The API key as a string, or `undefined` if none of the specified
* environment variables are set.
*/
export function getApiKeyFromEnvVar(): string | undefined {
return (
process.env.VERTEX_API_KEY ||
process.env.GOOGLE_API_KEY ||
process.env.GOOGLE_GENAI_API_KEY
);
}
export const MISSING_API_KEY_ERROR = new GenkitError({
status: 'FAILED_PRECONDITION',
message:
'Please pass in the API key or set the VERTEX_API_KEY or GOOGLE_API_KEY environment variable.\n' +
'For more details see https://firebase.google.com/docs/genkit/plugins/google-genai',
});
export const API_KEY_FALSE_ERROR = new GenkitError({
status: 'INVALID_ARGUMENT',
message:
'VertexAI plugin was initialized with {apiKey: false} but no apiKey configuration was passed at call time.',
});
export const NOT_SUPPORTED_IN_EXPRESS_ERROR = new GenkitError({
status: 'PERMISSION_DENIED',
message:
'This method or model is not supported in Vertex AI Express Mode.\n' +
'For more details see https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/vertex-ai-express-mode-api-reference',
});
/**
* Checks and retrieves an API key based on the provided argument and environment variables.
*
* - If `pluginApiKey` is a non-empty string, it's used as the API key.
* - If `pluginApiKey` is `undefined` or an empty string, it attempts to fetch the API key from environment
* - If `pluginApiKey` is `false`, key retrieval from the environment is skipped, and the function
* will return `undefined`. This mode indicates that the API key is expected to be provided
* at a later stage or in a different context.
*
* @param pluginApiKey - An optional API key string, `undefined` to check the environment, or `false` to bypass all checks in this function.
* @returns The resolved API key as a string, or `undefined` if `pluginApiKey` is `false`.
* @throws {Error} MISSING_API_KEY_ERROR - Thrown if `pluginApiKey` is not `false` and no API key
* can be found either in the `pluginApiKey` argument or from the environment.
*/
export function checkApiKey(
pluginApiKey: string | false | undefined
): string | undefined {
let apiKey: string | undefined;
// Don't get the key from the environment if pluginApiKey is false
if (pluginApiKey !== false) {
apiKey = pluginApiKey || getApiKeyFromEnvVar();
}
// If pluginApiKey is false, then we don't throw because we are waiting for
// the apiKey passed into the individual call
if (pluginApiKey !== false && !apiKey) {
throw MISSING_API_KEY_ERROR;
}
return apiKey;
}
/**
* Calculates and returns the effective API key based on multiple potential sources.
* The order of precedence for determining the API key is:
* 1. `requestApiKey` (if provided)
* 2. `pluginApiKey` (if provided and not `false`)
* 3. Environment variable (if `pluginApiKey` is not `false` and `pluginApiKey` is not provided)
*
* @param pluginApiKey - The apiKey value provided during plugin initialization.
* @param requestApiKey - The apiKey provided to an individual generate call.
* @returns The resolved API key as a string.
* @throws {Error} API_KEY_FALSE_ERROR - Thrown if `pluginApiKey` is `false` and `requestApiKey` is not provided
* @throws {Error} MISSING_API_KEY_ERROR - Thrown if no API key can be resolved from any source
*/
export function calculateApiKey(
pluginApiKey: string | false | undefined,
requestApiKey: string | undefined
): string {
let apiKey: string | undefined;
// Don't get the key from the environment if pluginApiKey is false
if (pluginApiKey !== false) {
apiKey = pluginApiKey || getApiKeyFromEnvVar();
}
apiKey = requestApiKey || apiKey;
if (pluginApiKey === false && !requestApiKey) {
throw API_KEY_FALSE_ERROR;
}
if (!apiKey) {
throw MISSING_API_KEY_ERROR;
}
return apiKey;
}
/** Vertex Express Mode lets you try a *subset* of Vertex AI features */
export function checkSupportedResourceMethod(params: {
clientOptions: ClientOptions;
resourcePath?: string;
resourceMethod?: string;
}) {
if (params.resourcePath == '') {
// This is how we get a base url for metadata
return;
}
const supportedExpressMethods = [
'countTokens',
'generateContent',
'streamGenerateContent',
];
if (
params.clientOptions.kind === 'express' &&
(!supportedExpressMethods.includes(params.resourceMethod ?? '') ||
params.resourcePath?.includes('endpoints/'))
) {
throw NOT_SUPPORTED_IN_EXPRESS_ERROR;
}
}