Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions packages/vertexai/src/models/vertexai-model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ export abstract class VertexAIModel {
apiKey: vertexAI.app.options.apiKey,
project: vertexAI.app.options.projectId,
appId: vertexAI.app.options.appId,
automaticDataCollectionEnabled:
vertexAI.app.automaticDataCollectionEnabled,
location: vertexAI.location
};

Expand Down
44 changes: 44 additions & 0 deletions packages/vertexai/src/requests/request.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,50 @@ describe('request methods', () => {
const headers = await getHeaders(fakeUrl);
expect(headers.get('x-goog-api-key')).to.equal('key');
});
it('adds app id if automatedDataCollectionEnabled is undefined', async () => {
const headers = await getHeaders(fakeUrl);
expect(headers.get('X-Firebase-AppId')).to.equal('my-appid');
});
it('adds app id if automatedDataCollectionEnabled is true', async () => {
const fakeApiSettings: ApiSettings = {
apiKey: 'key',
project: 'myproject',
appId: 'my-appid',
location: 'moon',
automaticDataCollectionEnabled: true,
getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
};
const fakeUrl = new RequestUrl(
'models/model-name',
Task.GENERATE_CONTENT,
fakeApiSettings,
true,
{}
);
const headers = await getHeaders(fakeUrl);
expect(headers.get('X-Firebase-AppId')).to.equal('my-appid');
});
it('does not add app id if automatedDataCollectionEnabled is false', async () => {
const fakeApiSettings: ApiSettings = {
apiKey: 'key',
project: 'myproject',
appId: 'my-appid',
location: 'moon',
automaticDataCollectionEnabled: false,
getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
};
const fakeUrl = new RequestUrl(
'models/model-name',
Task.GENERATE_CONTENT,
fakeApiSettings,
true,
{}
);
const headers = await getHeaders(fakeUrl);
expect(headers.get('X-Firebase-AppId')).to.be.null;
});
it('adds app check token if it exists', async () => {
const headers = await getHeaders(fakeUrl);
expect(headers.get('X-Firebase-AppCheck')).to.equal('appchecktoken');
Expand Down
4 changes: 3 additions & 1 deletion packages/vertexai/src/requests/request.ts
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,9 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
headers.append('Content-Type', 'application/json');
headers.append('x-goog-api-client', getClientHeaders());
headers.append('x-goog-api-key', url.apiSettings.apiKey);
headers.append('X-Firebase-AppId', url.apiSettings.appId); // Will be converted to 'X-Firebase-Appid' before it's sent in the browser.
if (url.apiSettings.automaticDataCollectionEnabled !== false) {
headers.append('X-Firebase-AppId', url.apiSettings.appId);
}
if (url.apiSettings.getAppCheckToken) {
const appCheckToken = await url.apiSettings.getAppCheckToken();
if (appCheckToken) {
Expand Down
1 change: 1 addition & 0 deletions packages/vertexai/src/types/internal.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ export interface ApiSettings {
project: string;
appId: string;
location: string;
automaticDataCollectionEnabled?: boolean;
getAuthToken?: () => Promise<FirebaseAuthTokenData | null>;
getAppCheckToken?: () => Promise<AppCheckTokenResult>;
}
Loading