@ai-sdk/google-vertex 4.0.22 → 4.0.24
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/CHANGELOG.md +16 -0
- package/dist/anthropic/edge/index.js +1 -1
- package/dist/anthropic/edge/index.mjs +1 -1
- package/dist/edge/index.js +1 -1
- package/dist/edge/index.mjs +1 -1
- package/dist/index.js +1 -1
- package/dist/index.mjs +1 -1
- package/package.json +7 -6
- package/src/__snapshots__/google-vertex-embedding-model.test.ts.snap +39 -0
- package/src/anthropic/edge/google-vertex-anthropic-provider-edge.test.ts +87 -0
- package/src/anthropic/edge/google-vertex-anthropic-provider-edge.ts +41 -0
- package/src/anthropic/edge/index.ts +8 -0
- package/src/anthropic/google-vertex-anthropic-messages-options.ts +15 -0
- package/src/anthropic/google-vertex-anthropic-provider-node.test.ts +73 -0
- package/src/anthropic/google-vertex-anthropic-provider-node.ts +40 -0
- package/src/anthropic/google-vertex-anthropic-provider.test.ts +208 -0
- package/src/anthropic/google-vertex-anthropic-provider.ts +210 -0
- package/src/anthropic/index.ts +8 -0
- package/src/edge/google-vertex-auth-edge.test.ts +308 -0
- package/src/edge/google-vertex-auth-edge.ts +161 -0
- package/src/edge/google-vertex-provider-edge.test.ts +105 -0
- package/src/edge/google-vertex-provider-edge.ts +50 -0
- package/src/edge/index.ts +5 -0
- package/src/google-vertex-auth-google-auth-library.test.ts +59 -0
- package/src/google-vertex-auth-google-auth-library.ts +27 -0
- package/src/google-vertex-config.ts +8 -0
- package/src/google-vertex-embedding-model.test.ts +315 -0
- package/src/google-vertex-embedding-model.ts +135 -0
- package/src/google-vertex-embedding-options.ts +63 -0
- package/src/google-vertex-error.ts +19 -0
- package/src/google-vertex-image-model.test.ts +926 -0
- package/src/google-vertex-image-model.ts +288 -0
- package/src/google-vertex-image-settings.ts +8 -0
- package/src/google-vertex-options.ts +32 -0
- package/src/google-vertex-provider-node.test.ts +88 -0
- package/src/google-vertex-provider-node.ts +49 -0
- package/src/google-vertex-provider.test.ts +318 -0
- package/src/google-vertex-provider.ts +217 -0
- package/src/google-vertex-tools.ts +11 -0
- package/src/index.ts +7 -0
- package/src/version.ts +6 -0
|
@@ -0,0 +1,210 @@
|
|
|
1
|
+
import {
|
|
2
|
+
LanguageModelV3,
|
|
3
|
+
NoSuchModelError,
|
|
4
|
+
ProviderV3,
|
|
5
|
+
} from '@ai-sdk/provider';
|
|
6
|
+
import {
|
|
7
|
+
FetchFunction,
|
|
8
|
+
Resolvable,
|
|
9
|
+
loadOptionalSetting,
|
|
10
|
+
withoutTrailingSlash,
|
|
11
|
+
} from '@ai-sdk/provider-utils';
|
|
12
|
+
import {
|
|
13
|
+
anthropicTools,
|
|
14
|
+
AnthropicMessagesLanguageModel,
|
|
15
|
+
} from '@ai-sdk/anthropic/internal';
|
|
16
|
+
import { GoogleVertexAnthropicMessagesModelId } from './google-vertex-anthropic-messages-options';
|
|
17
|
+
|
|
18
|
+
/**
|
|
19
|
+
* Tools supported by Google Vertex Anthropic.
|
|
20
|
+
* This is a subset of the full Anthropic tools - only these are recognized by the Vertex API.
|
|
21
|
+
*/
|
|
22
|
+
export const vertexAnthropicTools = {
|
|
23
|
+
/**
|
|
24
|
+
* The bash tool enables Claude to execute shell commands in a persistent bash session,
|
|
25
|
+
* allowing system operations, script execution, and command-line automation.
|
|
26
|
+
*
|
|
27
|
+
* Image results are supported.
|
|
28
|
+
*/
|
|
29
|
+
bash_20241022: anthropicTools.bash_20241022,
|
|
30
|
+
|
|
31
|
+
/**
|
|
32
|
+
* The bash tool enables Claude to execute shell commands in a persistent bash session,
|
|
33
|
+
* allowing system operations, script execution, and command-line automation.
|
|
34
|
+
*
|
|
35
|
+
* Image results are supported.
|
|
36
|
+
*/
|
|
37
|
+
bash_20250124: anthropicTools.bash_20250124,
|
|
38
|
+
|
|
39
|
+
/**
|
|
40
|
+
* Claude can use an Anthropic-defined text editor tool to view and modify text files,
|
|
41
|
+
* helping you debug, fix, and improve your code or other text documents.
|
|
42
|
+
*
|
|
43
|
+
* Supported models: Claude Sonnet 3.5
|
|
44
|
+
*/
|
|
45
|
+
textEditor_20241022: anthropicTools.textEditor_20241022,
|
|
46
|
+
|
|
47
|
+
/**
|
|
48
|
+
* Claude can use an Anthropic-defined text editor tool to view and modify text files,
|
|
49
|
+
* helping you debug, fix, and improve your code or other text documents.
|
|
50
|
+
*
|
|
51
|
+
* Supported models: Claude Sonnet 3.7
|
|
52
|
+
*/
|
|
53
|
+
textEditor_20250124: anthropicTools.textEditor_20250124,
|
|
54
|
+
|
|
55
|
+
/**
|
|
56
|
+
* Claude can use an Anthropic-defined text editor tool to view and modify text files.
|
|
57
|
+
* Note: This version does not support the "undo_edit" command.
|
|
58
|
+
* @deprecated Use textEditor_20250728 instead
|
|
59
|
+
*/
|
|
60
|
+
textEditor_20250429: anthropicTools.textEditor_20250429,
|
|
61
|
+
|
|
62
|
+
/**
|
|
63
|
+
* Claude can use an Anthropic-defined text editor tool to view and modify text files.
|
|
64
|
+
* Note: This version does not support the "undo_edit" command and adds optional max_characters parameter.
|
|
65
|
+
* Supported models: Claude Sonnet 4, Opus 4, and Opus 4.1
|
|
66
|
+
*/
|
|
67
|
+
textEditor_20250728: anthropicTools.textEditor_20250728,
|
|
68
|
+
|
|
69
|
+
/**
|
|
70
|
+
* Claude can interact with computer environments through the computer use tool, which
|
|
71
|
+
* provides screenshot capabilities and mouse/keyboard control for autonomous desktop interaction.
|
|
72
|
+
*
|
|
73
|
+
* Image results are supported.
|
|
74
|
+
*/
|
|
75
|
+
computer_20241022: anthropicTools.computer_20241022,
|
|
76
|
+
|
|
77
|
+
/**
|
|
78
|
+
* Creates a web search tool that gives Claude direct access to real-time web content.
|
|
79
|
+
*/
|
|
80
|
+
webSearch_20250305: anthropicTools.webSearch_20250305,
|
|
81
|
+
};
|
|
82
|
+
export interface GoogleVertexAnthropicProvider extends ProviderV3 {
|
|
83
|
+
/**
|
|
84
|
+
Creates a model for text generation.
|
|
85
|
+
*/
|
|
86
|
+
(modelId: GoogleVertexAnthropicMessagesModelId): LanguageModelV3;
|
|
87
|
+
|
|
88
|
+
/**
|
|
89
|
+
Creates a model for text generation.
|
|
90
|
+
*/
|
|
91
|
+
languageModel(modelId: GoogleVertexAnthropicMessagesModelId): LanguageModelV3;
|
|
92
|
+
|
|
93
|
+
/**
|
|
94
|
+
* Anthropic tools supported by Google Vertex.
|
|
95
|
+
* Note: Only a subset of Anthropic tools are available on Vertex.
|
|
96
|
+
* Supported tools: bash_20241022, bash_20250124, textEditor_20241022,
|
|
97
|
+
* textEditor_20250124, textEditor_20250429, textEditor_20250728,
|
|
98
|
+
* computer_20241022, webSearch_20250305
|
|
99
|
+
*/
|
|
100
|
+
tools: typeof vertexAnthropicTools;
|
|
101
|
+
|
|
102
|
+
/**
|
|
103
|
+
* @deprecated Use `embeddingModel` instead.
|
|
104
|
+
*/
|
|
105
|
+
textEmbeddingModel(modelId: string): never;
|
|
106
|
+
}
|
|
107
|
+
|
|
108
|
+
export interface GoogleVertexAnthropicProviderSettings {
|
|
109
|
+
/**
|
|
110
|
+
* Google Cloud project ID. Defaults to the value of the `GOOGLE_VERTEX_PROJECT` environment variable.
|
|
111
|
+
*/
|
|
112
|
+
project?: string;
|
|
113
|
+
|
|
114
|
+
/**
|
|
115
|
+
* Google Cloud region. Defaults to the value of the `GOOGLE_VERTEX_LOCATION` environment variable.
|
|
116
|
+
*/
|
|
117
|
+
location?: string;
|
|
118
|
+
|
|
119
|
+
/**
|
|
120
|
+
Use a different URL prefix for API calls, e.g. to use proxy servers.
|
|
121
|
+
The default prefix is `https://api.anthropic.com/v1`.
|
|
122
|
+
*/
|
|
123
|
+
baseURL?: string;
|
|
124
|
+
|
|
125
|
+
/**
|
|
126
|
+
Custom headers to include in the requests.
|
|
127
|
+
*/
|
|
128
|
+
headers?: Resolvable<Record<string, string | undefined>>;
|
|
129
|
+
|
|
130
|
+
/**
|
|
131
|
+
Custom fetch implementation. You can use it as a middleware to intercept requests,
|
|
132
|
+
or to provide a custom fetch implementation for e.g. testing.
|
|
133
|
+
*/
|
|
134
|
+
fetch?: FetchFunction;
|
|
135
|
+
}
|
|
136
|
+
|
|
137
|
+
/**
|
|
138
|
+
Create a Google Vertex Anthropic provider instance.
|
|
139
|
+
*/
|
|
140
|
+
export function createVertexAnthropic(
|
|
141
|
+
options: GoogleVertexAnthropicProviderSettings = {},
|
|
142
|
+
): GoogleVertexAnthropicProvider {
|
|
143
|
+
const getBaseURL = () => {
|
|
144
|
+
const location = loadOptionalSetting({
|
|
145
|
+
settingValue: options.location,
|
|
146
|
+
environmentVariableName: 'GOOGLE_VERTEX_LOCATION',
|
|
147
|
+
});
|
|
148
|
+
const project = loadOptionalSetting({
|
|
149
|
+
settingValue: options.project,
|
|
150
|
+
environmentVariableName: 'GOOGLE_VERTEX_PROJECT',
|
|
151
|
+
});
|
|
152
|
+
|
|
153
|
+
return (
|
|
154
|
+
withoutTrailingSlash(options.baseURL) ??
|
|
155
|
+
`https://${location === 'global' ? '' : location + '-'}aiplatform.googleapis.com/v1/projects/${project}/locations/${location}/publishers/anthropic/models`
|
|
156
|
+
);
|
|
157
|
+
};
|
|
158
|
+
|
|
159
|
+
const createChatModel = (modelId: GoogleVertexAnthropicMessagesModelId) =>
|
|
160
|
+
new AnthropicMessagesLanguageModel(modelId, {
|
|
161
|
+
provider: 'vertex.anthropic.messages',
|
|
162
|
+
baseURL: getBaseURL(),
|
|
163
|
+
headers: options.headers ?? {},
|
|
164
|
+
fetch: options.fetch,
|
|
165
|
+
|
|
166
|
+
buildRequestUrl: (baseURL, isStreaming) =>
|
|
167
|
+
`${baseURL}/${modelId}:${
|
|
168
|
+
isStreaming ? 'streamRawPredict' : 'rawPredict'
|
|
169
|
+
}`,
|
|
170
|
+
transformRequestBody: args => {
|
|
171
|
+
// Remove model from args and add anthropic version
|
|
172
|
+
const { model, ...rest } = args;
|
|
173
|
+
return {
|
|
174
|
+
...rest,
|
|
175
|
+
anthropic_version: 'vertex-2023-10-16',
|
|
176
|
+
};
|
|
177
|
+
},
|
|
178
|
+
// Google Vertex Anthropic doesn't support URL sources, force download and base64 conversion
|
|
179
|
+
supportedUrls: () => ({}),
|
|
180
|
+
// force the use of JSON tool fallback for structured outputs since beta header isn't supported
|
|
181
|
+
supportsNativeStructuredOutput: false,
|
|
182
|
+
});
|
|
183
|
+
|
|
184
|
+
const provider = function (modelId: GoogleVertexAnthropicMessagesModelId) {
|
|
185
|
+
if (new.target) {
|
|
186
|
+
throw new Error(
|
|
187
|
+
'The Anthropic model function cannot be called with the new keyword.',
|
|
188
|
+
);
|
|
189
|
+
}
|
|
190
|
+
|
|
191
|
+
return createChatModel(modelId);
|
|
192
|
+
};
|
|
193
|
+
|
|
194
|
+
provider.specificationVersion = 'v3' as const;
|
|
195
|
+
provider.languageModel = createChatModel;
|
|
196
|
+
provider.chat = createChatModel;
|
|
197
|
+
provider.messages = createChatModel;
|
|
198
|
+
|
|
199
|
+
provider.embeddingModel = (modelId: string) => {
|
|
200
|
+
throw new NoSuchModelError({ modelId, modelType: 'embeddingModel' });
|
|
201
|
+
};
|
|
202
|
+
provider.textEmbeddingModel = provider.embeddingModel;
|
|
203
|
+
provider.imageModel = (modelId: string) => {
|
|
204
|
+
throw new NoSuchModelError({ modelId, modelType: 'imageModel' });
|
|
205
|
+
};
|
|
206
|
+
|
|
207
|
+
provider.tools = vertexAnthropicTools;
|
|
208
|
+
|
|
209
|
+
return provider;
|
|
210
|
+
}
|
|
@@ -0,0 +1,308 @@
|
|
|
1
|
+
import {
|
|
2
|
+
generateAuthToken,
|
|
3
|
+
GoogleCredentials,
|
|
4
|
+
} from './google-vertex-auth-edge';
|
|
5
|
+
import { beforeEach, afterEach, describe, expect, it, vi } from 'vitest';
|
|
6
|
+
|
|
7
|
+
// Mock provider-utils to control runtime environment detection
|
|
8
|
+
vi.mock('@ai-sdk/provider-utils', async () => {
|
|
9
|
+
const actual = await vi.importActual('@ai-sdk/provider-utils');
|
|
10
|
+
return {
|
|
11
|
+
...actual,
|
|
12
|
+
getRuntimeEnvironmentUserAgent: vi.fn(() => 'runtime/testenv'),
|
|
13
|
+
withUserAgentSuffix: actual.withUserAgentSuffix,
|
|
14
|
+
};
|
|
15
|
+
});
|
|
16
|
+
|
|
17
|
+
vi.mock('../version', () => ({
|
|
18
|
+
VERSION: '0.0.0-test',
|
|
19
|
+
}));
|
|
20
|
+
|
|
21
|
+
describe('Google Vertex Edge Auth', () => {
|
|
22
|
+
const mockCredentials: GoogleCredentials = {
|
|
23
|
+
clientEmail: 'test@test.iam.gserviceaccount.com',
|
|
24
|
+
privateKey: 'mock-private-key',
|
|
25
|
+
privateKeyId: 'test-key-id',
|
|
26
|
+
};
|
|
27
|
+
|
|
28
|
+
const setupAtobStub = (credentials: typeof mockCredentials) => {
|
|
29
|
+
vi.stubGlobal(
|
|
30
|
+
'atob',
|
|
31
|
+
vi.fn().mockImplementation(str => {
|
|
32
|
+
const payload = {
|
|
33
|
+
alg: 'RS256',
|
|
34
|
+
typ: 'JWT',
|
|
35
|
+
iss: credentials.clientEmail,
|
|
36
|
+
scope: 'https://www.googleapis.com/auth/cloud-platform',
|
|
37
|
+
aud: 'https://oauth2.googleapis.com/token',
|
|
38
|
+
iat: 1616161616,
|
|
39
|
+
exp: 1616165216,
|
|
40
|
+
};
|
|
41
|
+
|
|
42
|
+
if (credentials.privateKeyId) {
|
|
43
|
+
Object.assign(payload, { kid: credentials.privateKeyId });
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
return JSON.stringify(payload);
|
|
47
|
+
}),
|
|
48
|
+
);
|
|
49
|
+
};
|
|
50
|
+
|
|
51
|
+
beforeEach(() => {
|
|
52
|
+
// Mock WebCrypto
|
|
53
|
+
const mockSubtleCrypto = {
|
|
54
|
+
importKey: vi.fn().mockResolvedValue('mock-crypto-key'),
|
|
55
|
+
sign: vi.fn().mockResolvedValue(new Uint8Array([1, 2, 3])),
|
|
56
|
+
};
|
|
57
|
+
|
|
58
|
+
const mockCrypto = {
|
|
59
|
+
subtle: mockSubtleCrypto,
|
|
60
|
+
};
|
|
61
|
+
|
|
62
|
+
// Use vi.stubGlobal instead of direct assignment
|
|
63
|
+
vi.stubGlobal('crypto', mockCrypto);
|
|
64
|
+
|
|
65
|
+
// Mock fetch
|
|
66
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
67
|
+
ok: true,
|
|
68
|
+
json: () => Promise.resolve({ access_token: 'mock.jwt.token' }),
|
|
69
|
+
});
|
|
70
|
+
|
|
71
|
+
setupAtobStub(mockCredentials);
|
|
72
|
+
});
|
|
73
|
+
|
|
74
|
+
afterEach(() => {
|
|
75
|
+
vi.restoreAllMocks();
|
|
76
|
+
vi.unstubAllGlobals();
|
|
77
|
+
});
|
|
78
|
+
|
|
79
|
+
it('should generate a valid JWT token', async () => {
|
|
80
|
+
// Mock successful token exchange
|
|
81
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
82
|
+
ok: true,
|
|
83
|
+
json: () =>
|
|
84
|
+
Promise.resolve({
|
|
85
|
+
access_token:
|
|
86
|
+
'eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6InRlc3Qta2V5LWlkIn0.eyJpc3MiOiJ0ZXN0QHRlc3QuaWFtLmdzZXJ2aWNlYWNjb3VudC5jb20iLCJzY29wZSI6Imh0dHBzOi8vd3d3Lmdvb2dsZWFwaXMuY29tL2F1dGgvY2xvdWQtcGxhdGZvcm0iLCJhdWQiOiJodHRwczovL29hdXRoMi5nb29nbGVhcGlzLmNvbS90b2tlbiIsImlhdCI6MTYxNjE2MTYxNiwiZXhwIjoxNjE2MTY1MjE2fQ.signature',
|
|
87
|
+
}),
|
|
88
|
+
});
|
|
89
|
+
|
|
90
|
+
const token = await generateAuthToken(mockCredentials);
|
|
91
|
+
|
|
92
|
+
// JWT structure validation
|
|
93
|
+
const parts = token.split('.');
|
|
94
|
+
expect(parts).toHaveLength(3);
|
|
95
|
+
|
|
96
|
+
// Header validation
|
|
97
|
+
const header = JSON.parse(atob(parts[0]));
|
|
98
|
+
expect(header).toMatchObject({
|
|
99
|
+
alg: 'RS256',
|
|
100
|
+
typ: 'JWT',
|
|
101
|
+
kid: mockCredentials.privateKeyId,
|
|
102
|
+
iss: mockCredentials.clientEmail,
|
|
103
|
+
});
|
|
104
|
+
|
|
105
|
+
// Payload validation
|
|
106
|
+
const payload = JSON.parse(atob(parts[1]));
|
|
107
|
+
expect(payload).toHaveProperty('iss', mockCredentials.clientEmail);
|
|
108
|
+
expect(payload).toHaveProperty(
|
|
109
|
+
'scope',
|
|
110
|
+
'https://www.googleapis.com/auth/cloud-platform',
|
|
111
|
+
);
|
|
112
|
+
expect(payload).toHaveProperty(
|
|
113
|
+
'aud',
|
|
114
|
+
'https://oauth2.googleapis.com/token',
|
|
115
|
+
);
|
|
116
|
+
expect(payload).toHaveProperty('iat');
|
|
117
|
+
expect(payload).toHaveProperty('exp');
|
|
118
|
+
|
|
119
|
+
// Verify exp is ~1 hour after iat
|
|
120
|
+
expect(payload.exp - payload.iat).toBe(3600);
|
|
121
|
+
});
|
|
122
|
+
|
|
123
|
+
it('should throw error with invalid credentials', async () => {
|
|
124
|
+
// Mock failed token exchange
|
|
125
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
126
|
+
ok: false,
|
|
127
|
+
status: 400,
|
|
128
|
+
statusText: 'Bad Request',
|
|
129
|
+
json: () => Promise.resolve({ error: 'invalid_grant' }),
|
|
130
|
+
});
|
|
131
|
+
|
|
132
|
+
const invalidCredentials = {
|
|
133
|
+
...mockCredentials,
|
|
134
|
+
private_key: 'invalid-key',
|
|
135
|
+
};
|
|
136
|
+
|
|
137
|
+
await expect(generateAuthToken(invalidCredentials)).rejects.toThrow(
|
|
138
|
+
'Token request failed: Bad Request',
|
|
139
|
+
);
|
|
140
|
+
});
|
|
141
|
+
|
|
142
|
+
it('should load credentials from environment variables', async () => {
|
|
143
|
+
process.env.GOOGLE_CLIENT_EMAIL = mockCredentials.clientEmail;
|
|
144
|
+
process.env.GOOGLE_PRIVATE_KEY = mockCredentials.privateKey;
|
|
145
|
+
process.env.GOOGLE_PRIVATE_KEY_ID = mockCredentials.privateKeyId;
|
|
146
|
+
|
|
147
|
+
// Mock successful token exchange
|
|
148
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
149
|
+
ok: true,
|
|
150
|
+
json: () => Promise.resolve({ access_token: 'mock.jwt.token' }),
|
|
151
|
+
});
|
|
152
|
+
|
|
153
|
+
const token = await generateAuthToken();
|
|
154
|
+
expect(token).toBeTruthy();
|
|
155
|
+
expect(token.split('.')).toHaveLength(3);
|
|
156
|
+
|
|
157
|
+
// Clean up
|
|
158
|
+
delete process.env.GOOGLE_CLIENT_EMAIL;
|
|
159
|
+
delete process.env.GOOGLE_PRIVATE_KEY;
|
|
160
|
+
delete process.env.GOOGLE_PRIVATE_KEY_ID;
|
|
161
|
+
});
|
|
162
|
+
|
|
163
|
+
it('should throw error when client email is missing', async () => {
|
|
164
|
+
delete process.env.GOOGLE_CLIENT_EMAIL;
|
|
165
|
+
process.env.GOOGLE_PRIVATE_KEY = mockCredentials.privateKey;
|
|
166
|
+
process.env.GOOGLE_PRIVATE_KEY_ID = mockCredentials.privateKeyId;
|
|
167
|
+
|
|
168
|
+
await expect(generateAuthToken()).rejects.toThrow(
|
|
169
|
+
"Google client email setting is missing. Pass it using the 'clientEmail' parameter or the GOOGLE_CLIENT_EMAIL environment variable.",
|
|
170
|
+
);
|
|
171
|
+
|
|
172
|
+
// Clean up
|
|
173
|
+
delete process.env.GOOGLE_PRIVATE_KEY;
|
|
174
|
+
delete process.env.GOOGLE_PRIVATE_KEY_ID;
|
|
175
|
+
});
|
|
176
|
+
|
|
177
|
+
it('should throw error when private key is missing', async () => {
|
|
178
|
+
process.env.GOOGLE_CLIENT_EMAIL = mockCredentials.clientEmail;
|
|
179
|
+
delete process.env.GOOGLE_PRIVATE_KEY;
|
|
180
|
+
process.env.GOOGLE_PRIVATE_KEY_ID = mockCredentials.privateKeyId;
|
|
181
|
+
|
|
182
|
+
await expect(generateAuthToken()).rejects.toThrow(
|
|
183
|
+
"Google private key setting is missing. Pass it using the 'privateKey' parameter or the GOOGLE_PRIVATE_KEY environment variable.",
|
|
184
|
+
);
|
|
185
|
+
|
|
186
|
+
// Clean up
|
|
187
|
+
delete process.env.GOOGLE_CLIENT_EMAIL;
|
|
188
|
+
delete process.env.GOOGLE_PRIVATE_KEY_ID;
|
|
189
|
+
});
|
|
190
|
+
|
|
191
|
+
it('should work with or without private key ID', async () => {
|
|
192
|
+
// Test with private key ID
|
|
193
|
+
process.env.GOOGLE_CLIENT_EMAIL = mockCredentials.clientEmail;
|
|
194
|
+
process.env.GOOGLE_PRIVATE_KEY = mockCredentials.privateKey;
|
|
195
|
+
process.env.GOOGLE_PRIVATE_KEY_ID = mockCredentials.privateKeyId;
|
|
196
|
+
|
|
197
|
+
// Mock successful token exchange
|
|
198
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
199
|
+
ok: true,
|
|
200
|
+
json: () => Promise.resolve({ access_token: 'mock.jwt.token' }),
|
|
201
|
+
});
|
|
202
|
+
|
|
203
|
+
const tokenWithKeyId = await generateAuthToken();
|
|
204
|
+
expect(tokenWithKeyId).toBeTruthy();
|
|
205
|
+
|
|
206
|
+
// Test without private key ID
|
|
207
|
+
delete process.env.GOOGLE_PRIVATE_KEY_ID;
|
|
208
|
+
|
|
209
|
+
const tokenWithoutKeyId = await generateAuthToken();
|
|
210
|
+
expect(tokenWithoutKeyId).toBeTruthy();
|
|
211
|
+
|
|
212
|
+
// Clean up
|
|
213
|
+
delete process.env.GOOGLE_CLIENT_EMAIL;
|
|
214
|
+
delete process.env.GOOGLE_PRIVATE_KEY;
|
|
215
|
+
});
|
|
216
|
+
|
|
217
|
+
it('should handle newlines in private key from env vars', async () => {
|
|
218
|
+
process.env.GOOGLE_CLIENT_EMAIL = mockCredentials.clientEmail;
|
|
219
|
+
process.env.GOOGLE_PRIVATE_KEY = mockCredentials.privateKey.replace(
|
|
220
|
+
/\n/g,
|
|
221
|
+
'\\n',
|
|
222
|
+
);
|
|
223
|
+
process.env.GOOGLE_PRIVATE_KEY_ID = mockCredentials.privateKeyId;
|
|
224
|
+
|
|
225
|
+
// Mock successful token exchange
|
|
226
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
227
|
+
ok: true,
|
|
228
|
+
json: () => Promise.resolve({ access_token: 'mock.jwt.token' }),
|
|
229
|
+
});
|
|
230
|
+
|
|
231
|
+
const token = await generateAuthToken();
|
|
232
|
+
expect(token).toBeTruthy();
|
|
233
|
+
|
|
234
|
+
// Clean up
|
|
235
|
+
delete process.env.GOOGLE_CLIENT_EMAIL;
|
|
236
|
+
delete process.env.GOOGLE_PRIVATE_KEY;
|
|
237
|
+
delete process.env.GOOGLE_PRIVATE_KEY_ID;
|
|
238
|
+
});
|
|
239
|
+
|
|
240
|
+
it('should throw error on fetch failure', async () => {
|
|
241
|
+
const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
|
|
242
|
+
global.fetch = vi.fn().mockRejectedValue(new Error('Network error'));
|
|
243
|
+
|
|
244
|
+
await expect(generateAuthToken(mockCredentials)).rejects.toThrow(
|
|
245
|
+
'Network error',
|
|
246
|
+
);
|
|
247
|
+
|
|
248
|
+
consoleSpy.mockRestore();
|
|
249
|
+
});
|
|
250
|
+
|
|
251
|
+
it('should throw error when token request fails', async () => {
|
|
252
|
+
// Mock a failed response from the token endpoint
|
|
253
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
254
|
+
ok: false,
|
|
255
|
+
statusText: 'Unauthorized',
|
|
256
|
+
status: 401,
|
|
257
|
+
json: () => Promise.resolve({ error: 'unauthorized' }),
|
|
258
|
+
});
|
|
259
|
+
|
|
260
|
+
await expect(generateAuthToken(mockCredentials)).rejects.toThrow(
|
|
261
|
+
'Token request failed: Unauthorized',
|
|
262
|
+
);
|
|
263
|
+
});
|
|
264
|
+
|
|
265
|
+
it('should work without privateKeyId', async () => {
|
|
266
|
+
const credentialsWithoutKeyId = {
|
|
267
|
+
clientEmail: mockCredentials.clientEmail,
|
|
268
|
+
privateKey: mockCredentials.privateKey,
|
|
269
|
+
};
|
|
270
|
+
setupAtobStub(credentialsWithoutKeyId);
|
|
271
|
+
|
|
272
|
+
// Mock successful token exchange
|
|
273
|
+
global.fetch = vi.fn().mockResolvedValue({
|
|
274
|
+
ok: true,
|
|
275
|
+
json: () => Promise.resolve({ access_token: 'mock.jwt.token' }),
|
|
276
|
+
});
|
|
277
|
+
|
|
278
|
+
const token = await generateAuthToken(credentialsWithoutKeyId);
|
|
279
|
+
expect(token).toBeTruthy();
|
|
280
|
+
|
|
281
|
+
// Verify the JWT structure
|
|
282
|
+
const parts = token.split('.');
|
|
283
|
+
expect(parts).toHaveLength(3);
|
|
284
|
+
|
|
285
|
+
// Verify header doesn't include kid when privateKeyId is not provided
|
|
286
|
+
const header = JSON.parse(atob(parts[0]));
|
|
287
|
+
expect(header).not.toHaveProperty('kid');
|
|
288
|
+
});
|
|
289
|
+
|
|
290
|
+
it('should include correct user-agent header', async () => {
|
|
291
|
+
const mockFetch = vi.fn().mockResolvedValue({
|
|
292
|
+
ok: true,
|
|
293
|
+
json: () => Promise.resolve({ access_token: 'mock.jwt.token' }),
|
|
294
|
+
});
|
|
295
|
+
global.fetch = mockFetch;
|
|
296
|
+
|
|
297
|
+
await generateAuthToken(mockCredentials);
|
|
298
|
+
|
|
299
|
+
expect(mockFetch).toHaveBeenCalledWith(
|
|
300
|
+
'https://oauth2.googleapis.com/token',
|
|
301
|
+
expect.objectContaining({
|
|
302
|
+
headers: expect.objectContaining({
|
|
303
|
+
'user-agent': 'ai-sdk/google-vertex/0.0.0-test runtime/testenv',
|
|
304
|
+
}),
|
|
305
|
+
}),
|
|
306
|
+
);
|
|
307
|
+
});
|
|
308
|
+
});
|
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
import {
|
|
2
|
+
loadOptionalSetting,
|
|
3
|
+
loadSetting,
|
|
4
|
+
withUserAgentSuffix,
|
|
5
|
+
getRuntimeEnvironmentUserAgent,
|
|
6
|
+
} from '@ai-sdk/provider-utils';
|
|
7
|
+
import { VERSION } from '../version';
|
|
8
|
+
|
|
9
|
+
export interface GoogleCredentials {
|
|
10
|
+
/**
|
|
11
|
+
* The client email for the Google Cloud service account. Defaults to the
|
|
12
|
+
* value of the `GOOGLE_CLIENT_EMAIL` environment variable.
|
|
13
|
+
*/
|
|
14
|
+
clientEmail: string;
|
|
15
|
+
|
|
16
|
+
/**
|
|
17
|
+
* The private key for the Google Cloud service account. Defaults to the
|
|
18
|
+
* value of the `GOOGLE_PRIVATE_KEY` environment variable.
|
|
19
|
+
*/
|
|
20
|
+
privateKey: string;
|
|
21
|
+
|
|
22
|
+
/**
|
|
23
|
+
* Optional. The private key ID for the Google Cloud service account. Defaults
|
|
24
|
+
* to the value of the `GOOGLE_PRIVATE_KEY_ID` environment variable.
|
|
25
|
+
*/
|
|
26
|
+
privateKeyId?: string;
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
const loadCredentials = async (): Promise<GoogleCredentials> => {
|
|
30
|
+
try {
|
|
31
|
+
return {
|
|
32
|
+
clientEmail: loadSetting({
|
|
33
|
+
settingValue: undefined,
|
|
34
|
+
settingName: 'clientEmail',
|
|
35
|
+
environmentVariableName: 'GOOGLE_CLIENT_EMAIL',
|
|
36
|
+
description: 'Google client email',
|
|
37
|
+
}),
|
|
38
|
+
privateKey: loadSetting({
|
|
39
|
+
settingValue: undefined,
|
|
40
|
+
settingName: 'privateKey',
|
|
41
|
+
environmentVariableName: 'GOOGLE_PRIVATE_KEY',
|
|
42
|
+
description: 'Google private key',
|
|
43
|
+
}),
|
|
44
|
+
privateKeyId: loadOptionalSetting({
|
|
45
|
+
settingValue: undefined,
|
|
46
|
+
environmentVariableName: 'GOOGLE_PRIVATE_KEY_ID',
|
|
47
|
+
}),
|
|
48
|
+
};
|
|
49
|
+
} catch (error: any) {
|
|
50
|
+
throw new Error(`Failed to load Google credentials: ${error.message}`);
|
|
51
|
+
}
|
|
52
|
+
};
|
|
53
|
+
|
|
54
|
+
// Convert a string to base64url
|
|
55
|
+
const base64url = (str: string) => {
|
|
56
|
+
return btoa(str).replace(/\+/g, '-').replace(/\//g, '_').replace(/=/g, '');
|
|
57
|
+
};
|
|
58
|
+
const importPrivateKey = async (pemKey: string) => {
|
|
59
|
+
const pemHeader = '-----BEGIN PRIVATE KEY-----';
|
|
60
|
+
const pemFooter = '-----END PRIVATE KEY-----';
|
|
61
|
+
|
|
62
|
+
// Remove header, footer, and any whitespace/newlines
|
|
63
|
+
const pemContents = pemKey
|
|
64
|
+
.replace(pemHeader, '')
|
|
65
|
+
.replace(pemFooter, '')
|
|
66
|
+
.replace(/\s/g, '');
|
|
67
|
+
|
|
68
|
+
// Decode base64 to binary
|
|
69
|
+
const binaryString = atob(pemContents);
|
|
70
|
+
|
|
71
|
+
// Convert binary string to Uint8Array
|
|
72
|
+
const binaryData = new Uint8Array(binaryString.length);
|
|
73
|
+
for (let i = 0; i < binaryString.length; i++) {
|
|
74
|
+
binaryData[i] = binaryString.charCodeAt(i);
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
return await crypto.subtle.importKey(
|
|
78
|
+
'pkcs8',
|
|
79
|
+
binaryData,
|
|
80
|
+
{ name: 'RSASSA-PKCS1-v1_5', hash: 'SHA-256' },
|
|
81
|
+
true,
|
|
82
|
+
['sign'],
|
|
83
|
+
);
|
|
84
|
+
};
|
|
85
|
+
|
|
86
|
+
const buildJwt = async (credentials: GoogleCredentials) => {
|
|
87
|
+
const now = Math.floor(Date.now() / 1000);
|
|
88
|
+
|
|
89
|
+
// Only include kid in header if privateKeyId is provided
|
|
90
|
+
const header: { alg: string; typ: string; kid?: string } = {
|
|
91
|
+
alg: 'RS256',
|
|
92
|
+
typ: 'JWT',
|
|
93
|
+
};
|
|
94
|
+
|
|
95
|
+
if (credentials.privateKeyId) {
|
|
96
|
+
header.kid = credentials.privateKeyId;
|
|
97
|
+
}
|
|
98
|
+
|
|
99
|
+
const payload = {
|
|
100
|
+
iss: credentials.clientEmail,
|
|
101
|
+
scope: 'https://www.googleapis.com/auth/cloud-platform',
|
|
102
|
+
aud: 'https://oauth2.googleapis.com/token',
|
|
103
|
+
exp: now + 3600,
|
|
104
|
+
iat: now,
|
|
105
|
+
};
|
|
106
|
+
|
|
107
|
+
const privateKey = await importPrivateKey(credentials.privateKey);
|
|
108
|
+
|
|
109
|
+
const signingInput = `${base64url(JSON.stringify(header))}.${base64url(
|
|
110
|
+
JSON.stringify(payload),
|
|
111
|
+
)}`;
|
|
112
|
+
const encoder = new TextEncoder();
|
|
113
|
+
const data = encoder.encode(signingInput);
|
|
114
|
+
|
|
115
|
+
const signature = await crypto.subtle.sign(
|
|
116
|
+
'RSASSA-PKCS1-v1_5',
|
|
117
|
+
privateKey,
|
|
118
|
+
data,
|
|
119
|
+
);
|
|
120
|
+
|
|
121
|
+
const signatureBase64 = base64url(
|
|
122
|
+
String.fromCharCode(...new Uint8Array(signature)),
|
|
123
|
+
);
|
|
124
|
+
|
|
125
|
+
return `${base64url(JSON.stringify(header))}.${base64url(
|
|
126
|
+
JSON.stringify(payload),
|
|
127
|
+
)}.${signatureBase64}`;
|
|
128
|
+
};
|
|
129
|
+
|
|
130
|
+
/**
|
|
131
|
+
* Generate an authentication token for Google Vertex AI in a manner compatible
|
|
132
|
+
* with the Edge runtime.
|
|
133
|
+
*/
|
|
134
|
+
export async function generateAuthToken(credentials?: GoogleCredentials) {
|
|
135
|
+
try {
|
|
136
|
+
const creds = credentials || (await loadCredentials());
|
|
137
|
+
const jwt = await buildJwt(creds);
|
|
138
|
+
|
|
139
|
+
const response = await fetch('https://oauth2.googleapis.com/token', {
|
|
140
|
+
method: 'POST',
|
|
141
|
+
headers: withUserAgentSuffix(
|
|
142
|
+
{ 'Content-Type': 'application/x-www-form-urlencoded' },
|
|
143
|
+
`ai-sdk/google-vertex/${VERSION}`,
|
|
144
|
+
getRuntimeEnvironmentUserAgent(),
|
|
145
|
+
),
|
|
146
|
+
body: new URLSearchParams({
|
|
147
|
+
grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer',
|
|
148
|
+
assertion: jwt,
|
|
149
|
+
}),
|
|
150
|
+
});
|
|
151
|
+
|
|
152
|
+
if (!response.ok) {
|
|
153
|
+
throw new Error(`Token request failed: ${response.statusText}`);
|
|
154
|
+
}
|
|
155
|
+
|
|
156
|
+
const data = await response.json();
|
|
157
|
+
return data.access_token;
|
|
158
|
+
} catch (error) {
|
|
159
|
+
throw error;
|
|
160
|
+
}
|
|
161
|
+
}
|