@ai-sdk/google-vertex
Version:
149 lines (148 loc) • 4.77 kB
JavaScript
import { loadOptionalSetting, loadSetting, resolve, withoutTrailingSlash } from "@ai-sdk/provider-utils";
import { GoogleAuth } from "google-auth-library";
import { NoSuchModelError } from "@ai-sdk/provider";
import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
//#region src/google-vertex-auth-google-auth-library.ts
function createAuthTokenGenerator(options) {
const auth = new GoogleAuth({
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
...options
});
return async function generateAuthToken() {
return (await (await auth.getClient()).getAccessToken())?.token ?? null;
};
}
//#endregion
//#region src/xai/google-vertex-xai-provider.ts
function convertGoogleVertexXaiUsage(usage) {
if (usage == null) return {
inputTokens: {
total: void 0,
noCache: void 0,
cacheRead: void 0,
cacheWrite: void 0
},
outputTokens: {
total: void 0,
text: void 0,
reasoning: void 0
},
raw: void 0
};
const promptTokens = usage.prompt_tokens ?? 0;
const completionTokens = usage.completion_tokens ?? 0;
const cacheReadTokens = usage.prompt_tokens_details?.cached_tokens ?? 0;
const reasoningTokens = usage.completion_tokens_details?.reasoning_tokens ?? 0;
return {
inputTokens: {
total: promptTokens,
noCache: promptTokens - cacheReadTokens,
cacheRead: cacheReadTokens,
cacheWrite: void 0
},
outputTokens: {
total: completionTokens + reasoningTokens,
text: completionTokens,
reasoning: reasoningTokens
},
raw: usage
};
}
function transformGoogleVertexXaiRequestBody(args) {
const { reasoning_effort: _reasoningEffort, ...rest } = args;
return rest;
}
/**
* Create a Google Vertex AI xAI provider instance.
* Uses the OpenAI-compatible Chat Completions API for Grok partner models.
*
* @see https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/grok
*/
function createGoogleVertexXai$1(options = {}) {
const loadLocation = () => loadOptionalSetting({
settingValue: options.location,
environmentVariableName: "GOOGLE_VERTEX_LOCATION"
});
const loadProject = () => loadSetting({
settingValue: options.project,
settingName: "project",
environmentVariableName: "GOOGLE_VERTEX_PROJECT",
description: "Google Vertex project"
});
const constructBaseURL = () => {
return `https://aiplatform.googleapis.com/v1/projects/${loadProject()}/locations/${loadLocation() ?? "global"}/endpoints/openapi`;
};
const loadBaseURL = () => withoutTrailingSlash(options.baseURL ?? "") || constructBaseURL();
let cachedProvider;
const getProvider = () => cachedProvider ??= createOpenAICompatible({
name: "googleVertex.xai",
baseURL: loadBaseURL(),
fetch: options.fetch,
includeUsage: true,
supportsStructuredOutputs: true,
supportedUrls: () => ({ "image/*": [/^https?:\/\/.*$/] }),
transformRequestBody: transformGoogleVertexXaiRequestBody,
convertUsage: convertGoogleVertexXaiUsage
});
const createChatModel = (modelId) => getProvider().languageModel(modelId);
const provider = function(modelId) {
if (new.target) throw new Error("The Google Vertex xAI model function cannot be called with the new keyword.");
return createChatModel(modelId);
};
provider.specificationVersion = "v4";
provider.languageModel = createChatModel;
provider.chatModel = (modelId) => getProvider().chatModel(modelId);
provider.embeddingModel = (modelId) => {
throw new NoSuchModelError({
modelId,
modelType: "embeddingModel"
});
};
provider.textEmbeddingModel = provider.embeddingModel;
provider.imageModel = (modelId) => {
throw new NoSuchModelError({
modelId,
modelType: "imageModel"
});
};
return provider;
}
//#endregion
//#region src/xai/google-vertex-xai-provider-node.ts
/**
* Create a Google Vertex AI xAI provider instance for Node.js.
* Uses the OpenAI-compatible Chat Completions API for Grok partner models.
* Automatically handles Google Cloud authentication.
*
* @see https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/grok
*/
function createGoogleVertexXai(options = {}) {
const generateAuthToken = createAuthTokenGenerator(options.googleAuthOptions);
const customFetch = async (url, init) => {
const token = await generateAuthToken();
const authHeaders = {
...await resolve(options.headers),
Authorization: `Bearer ${token}`
};
const fetchInit = {
...init,
headers: {
...init?.headers,
...authHeaders
}
};
return (options.fetch ?? fetch)(url, fetchInit);
};
return createGoogleVertexXai$1({
...options,
fetch: customFetch,
headers: void 0
});
}
/**
* Default Google Vertex AI xAI provider instance for Node.js.
*/
const googleVertexXai = createGoogleVertexXai();
//#endregion
export { createGoogleVertexXai, googleVertexXai };
//# sourceMappingURL=index.js.map