UNPKG

@ai-sdk/google-vertex

Version:
149 lines (148 loc) • 4.77 kB
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