Skip to content

Commit 648de84

Browse files
dlarocquedlarocque
authored
Add X-Firebase-AppId header to VertexAI requests (#8809)
1 parent edb4001 commit 648de84

15 files changed

+137
-8
lines changed

‎.changeset/red-hornets-peel.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@firebase/vertexai': patch
3+
---
4+
5+
Throw an error when initializing models if `appId` is not defined in the given `VertexAI` instance.

‎common/api-review/vertexai.api.md

+1
Original file line numberDiff line numberDiff line change
@@ -802,6 +802,7 @@ export const enum VertexAIErrorCode {
802802
INVALID_CONTENT = "invalid-content",
803803
INVALID_SCHEMA = "invalid-schema",
804804
NO_API_KEY = "no-api-key",
805+
NO_APP_ID = "no-app-id",
805806
NO_MODEL = "no-model",
806807
NO_PROJECT_ID = "no-project-id",
807808
PARSE_FAILED = "parse-failed",

‎docs-devsite/vertexai.md

+1
Original file line numberDiff line numberDiff line change
@@ -551,6 +551,7 @@ export declare const enum VertexAIErrorCode
551551
| INVALID\_CONTENT | <code>&quot;invalid-content&quot;</code> | An error associated with a Content object. |
552552
| INVALID\_SCHEMA | <code>&quot;invalid-schema&quot;</code> | An error due to invalid Schema input. |
553553
| NO\_API\_KEY | <code>&quot;no-api-key&quot;</code> | An error occurred due to a missing Firebase API key. |
554+
| NO\_APP\_ID | <code>&quot;no-app-id&quot;</code> | An error occured due to a missing Firebase app ID. |
554555
| NO\_MODEL | <code>&quot;no-model&quot;</code> | An error occurred due to a model name not being specified during initialization. |
555556
| NO\_PROJECT\_ID | <code>&quot;no-project-id&quot;</code> | An error occurred due to a missing project ID. |
556557
| PARSE\_FAILED | <code>&quot;parse-failed&quot;</code> | An error occurred while parsing. |

‎packages/vertexai/src/api.test.ts

+38-5
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ const fakeVertexAI: VertexAI = {
2727
automaticDataCollectionEnabled: true,
2828
options: {
2929
apiKey: 'key',
30-
projectId: 'my-project'
30+
projectId: 'my-project',
31+
appId: 'my-appid'
3132
}
3233
},
3334
location: 'us-central1'
@@ -48,7 +49,7 @@ describe('Top level API', () => {
4849
it('getGenerativeModel throws if no apiKey is provided', () => {
4950
const fakeVertexNoApiKey = {
5051
...fakeVertexAI,
51-
app: { options: { projectId: 'my-project' } }
52+
app: { options: { projectId: 'my-project', appId: 'my-appid' } }
5253
} as VertexAI;
5354
try {
5455
getGenerativeModel(fakeVertexNoApiKey, { model: 'my-model' });
@@ -64,7 +65,7 @@ describe('Top level API', () => {
6465
it('getGenerativeModel throws if no projectId is provided', () => {
6566
const fakeVertexNoProject = {
6667
...fakeVertexAI,
67-
app: { options: { apiKey: 'my-key' } }
68+
app: { options: { apiKey: 'my-key', appId: 'my-appid' } }
6869
} as VertexAI;
6970
try {
7071
getGenerativeModel(fakeVertexNoProject, { model: 'my-model' });
@@ -79,6 +80,22 @@ describe('Top level API', () => {
7980
);
8081
}
8182
});
83+
it('getGenerativeModel throws if no appId is provided', () => {
84+
const fakeVertexNoProject = {
85+
...fakeVertexAI,
86+
app: { options: { apiKey: 'my-key', projectId: 'my-projectid' } }
87+
} as VertexAI;
88+
try {
89+
getGenerativeModel(fakeVertexNoProject, { model: 'my-model' });
90+
} catch (e) {
91+
expect((e as VertexAIError).code).includes(VertexAIErrorCode.NO_APP_ID);
92+
expect((e as VertexAIError).message).equals(
93+
`VertexAI: The "appId" field is empty in the local` +
94+
` Firebase config. Firebase VertexAI requires this field ` +
95+
`to contain a valid app ID. (vertexAI/${VertexAIErrorCode.NO_APP_ID})`
96+
);
97+
}
98+
});
8299
it('getGenerativeModel gets a GenerativeModel', () => {
83100
const genModel = getGenerativeModel(fakeVertexAI, { model: 'my-model' });
84101
expect(genModel).to.be.an.instanceOf(GenerativeModel);
@@ -98,7 +115,7 @@ describe('Top level API', () => {
98115
it('getImagenModel throws if no apiKey is provided', () => {
99116
const fakeVertexNoApiKey = {
100117
...fakeVertexAI,
101-
app: { options: { projectId: 'my-project' } }
118+
app: { options: { projectId: 'my-project', appId: 'my-appid' } }
102119
} as VertexAI;
103120
try {
104121
getImagenModel(fakeVertexNoApiKey, { model: 'my-model' });
@@ -114,7 +131,7 @@ describe('Top level API', () => {
114131
it('getImagenModel throws if no projectId is provided', () => {
115132
const fakeVertexNoProject = {
116133
...fakeVertexAI,
117-
app: { options: { apiKey: 'my-key' } }
134+
app: { options: { apiKey: 'my-key', appId: 'my-appid' } }
118135
} as VertexAI;
119136
try {
120137
getImagenModel(fakeVertexNoProject, { model: 'my-model' });
@@ -129,6 +146,22 @@ describe('Top level API', () => {
129146
);
130147
}
131148
});
149+
it('getImagenModel throws if no appId is provided', () => {
150+
const fakeVertexNoProject = {
151+
...fakeVertexAI,
152+
app: { options: { apiKey: 'my-key', projectId: 'my-project' } }
153+
} as VertexAI;
154+
try {
155+
getImagenModel(fakeVertexNoProject, { model: 'my-model' });
156+
} catch (e) {
157+
expect((e as VertexAIError).code).includes(VertexAIErrorCode.NO_APP_ID);
158+
expect((e as VertexAIError).message).equals(
159+
`VertexAI: The "appId" field is empty in the local` +
160+
` Firebase config. Firebase VertexAI requires this field ` +
161+
`to contain a valid app ID. (vertexAI/${VertexAIErrorCode.NO_APP_ID})`
162+
);
163+
}
164+
});
132165
it('getImagenModel gets an ImagenModel', () => {
133166
const genModel = getImagenModel(fakeVertexAI, { model: 'my-model' });
134167
expect(genModel).to.be.an.instanceOf(ImagenModel);

‎packages/vertexai/src/methods/chat-session.test.ts

+1
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ use(chaiAsPromised);
3030
const fakeApiSettings: ApiSettings = {
3131
apiKey: 'key',
3232
project: 'my-project',
33+
appId: 'my-appid',
3334
location: 'us-central1'
3435
};
3536

‎packages/vertexai/src/methods/count-tokens.test.ts

+1
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ use(chaiAsPromised);
3232
const fakeApiSettings: ApiSettings = {
3333
apiKey: 'key',
3434
project: 'my-project',
35+
appId: 'my-appid',
3536
location: 'us-central1'
3637
};
3738

‎packages/vertexai/src/methods/generate-content.test.ts

+1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ use(chaiAsPromised);
3737
const fakeApiSettings: ApiSettings = {
3838
apiKey: 'key',
3939
project: 'my-project',
40+
appId: 'my-appid',
4041
location: 'us-central1'
4142
};
4243

‎packages/vertexai/src/models/generative-model.test.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ const fakeVertexAI: VertexAI = {
3030
automaticDataCollectionEnabled: true,
3131
options: {
3232
apiKey: 'key',
33-
projectId: 'my-project'
33+
projectId: 'my-project',
34+
appId: 'my-appid'
3435
}
3536
},
3637
location: 'us-central1'

‎packages/vertexai/src/models/imagen-model.test.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,8 @@ const fakeVertexAI: VertexAI = {
3737
automaticDataCollectionEnabled: true,
3838
options: {
3939
apiKey: 'key',
40-
projectId: 'my-project'
40+
projectId: 'my-project',
41+
appId: 'my-appid'
4142
}
4243
},
4344
location: 'us-central1'

‎packages/vertexai/src/models/vertexai-model.test.ts

+20-1
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@ const fakeVertexAI: VertexAI = {
3838
automaticDataCollectionEnabled: true,
3939
options: {
4040
apiKey: 'key',
41-
projectId: 'my-project'
41+
projectId: 'my-project',
42+
appId: 'my-appid'
4243
}
4344
},
4445
location: 'us-central1'
@@ -100,4 +101,22 @@ describe('VertexAIModel', () => {
100101
);
101102
}
102103
});
104+
it('throws if not passed an app ID', () => {
105+
const fakeVertexAI: VertexAI = {
106+
app: {
107+
name: 'DEFAULT',
108+
automaticDataCollectionEnabled: true,
109+
options: {
110+
apiKey: 'key',
111+
projectId: 'my-project'
112+
}
113+
},
114+
location: 'us-central1'
115+
};
116+
try {
117+
new TestModel(fakeVertexAI, 'my-model');
118+
} catch (e) {
119+
expect((e as VertexAIError).code).to.equal(VertexAIErrorCode.NO_APP_ID);
120+
}
121+
});
103122
});

‎packages/vertexai/src/models/vertexai-model.ts

+8
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,18 @@ export abstract class VertexAIModel {
6868
VertexAIErrorCode.NO_PROJECT_ID,
6969
`The "projectId" field is empty in the local Firebase config. Firebase VertexAI requires this field to contain a valid project ID.`
7070
);
71+
} else if (!vertexAI.app?.options?.appId) {
72+
throw new VertexAIError(
73+
VertexAIErrorCode.NO_APP_ID,
74+
`The "appId" field is empty in the local Firebase config. Firebase VertexAI requires this field to contain a valid app ID.`
75+
);
7176
} else {
7277
this._apiSettings = {
7378
apiKey: vertexAI.app.options.apiKey,
7479
project: vertexAI.app.options.projectId,
80+
appId: vertexAI.app.options.appId,
81+
automaticDataCollectionEnabled:
82+
vertexAI.app.automaticDataCollectionEnabled,
7583
location: vertexAI.location
7684
};
7785

‎packages/vertexai/src/requests/request.test.ts

+49
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ use(chaiAsPromised);
3232
const fakeApiSettings: ApiSettings = {
3333
apiKey: 'key',
3434
project: 'my-project',
35+
appId: 'my-appid',
3536
location: 'us-central1'
3637
};
3738

@@ -103,6 +104,7 @@ describe('request methods', () => {
103104
const fakeApiSettings: ApiSettings = {
104105
apiKey: 'key',
105106
project: 'myproject',
107+
appId: 'my-appid',
106108
location: 'moon',
107109
getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
108110
getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
@@ -124,6 +126,50 @@ describe('request methods', () => {
124126
const headers = await getHeaders(fakeUrl);
125127
expect(headers.get('x-goog-api-key')).to.equal('key');
126128
});
129+
it('adds app id if automatedDataCollectionEnabled is true', async () => {
130+
const fakeApiSettings: ApiSettings = {
131+
apiKey: 'key',
132+
project: 'myproject',
133+
appId: 'my-appid',
134+
location: 'moon',
135+
automaticDataCollectionEnabled: true,
136+
getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
137+
getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
138+
};
139+
const fakeUrl = new RequestUrl(
140+
'models/model-name',
141+
Task.GENERATE_CONTENT,
142+
fakeApiSettings,
143+
true,
144+
{}
145+
);
146+
const headers = await getHeaders(fakeUrl);
147+
expect(headers.get('X-Firebase-Appid')).to.equal('my-appid');
148+
});
149+
it('does not add app id if automatedDataCollectionEnabled is undefined', async () => {
150+
const headers = await getHeaders(fakeUrl);
151+
expect(headers.get('X-Firebase-Appid')).to.be.null;
152+
});
153+
it('does not add app id if automatedDataCollectionEnabled is false', async () => {
154+
const fakeApiSettings: ApiSettings = {
155+
apiKey: 'key',
156+
project: 'myproject',
157+
appId: 'my-appid',
158+
location: 'moon',
159+
automaticDataCollectionEnabled: false,
160+
getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
161+
getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
162+
};
163+
const fakeUrl = new RequestUrl(
164+
'models/model-name',
165+
Task.GENERATE_CONTENT,
166+
fakeApiSettings,
167+
true,
168+
{}
169+
);
170+
const headers = await getHeaders(fakeUrl);
171+
expect(headers.get('X-Firebase-Appid')).to.be.null;
172+
});
127173
it('adds app check token if it exists', async () => {
128174
const headers = await getHeaders(fakeUrl);
129175
expect(headers.get('X-Firebase-AppCheck')).to.equal('appchecktoken');
@@ -135,6 +181,7 @@ describe('request methods', () => {
135181
{
136182
apiKey: 'key',
137183
project: 'myproject',
184+
appId: 'my-appid',
138185
location: 'moon'
139186
},
140187
true,
@@ -167,6 +214,7 @@ describe('request methods', () => {
167214
{
168215
apiKey: 'key',
169216
project: 'myproject',
217+
appId: 'my-appid',
170218
location: 'moon',
171219
getAppCheckToken: () =>
172220
Promise.resolve({ token: 'dummytoken', error: Error('oops') })
@@ -193,6 +241,7 @@ describe('request methods', () => {
193241
{
194242
apiKey: 'key',
195243
project: 'myproject',
244+
appId: 'my-appid',
196245
location: 'moon'
197246
},
198247
true,

‎packages/vertexai/src/requests/request.ts

+3
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,9 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
8484
headers.append('Content-Type', 'application/json');
8585
headers.append('x-goog-api-client', getClientHeaders());
8686
headers.append('x-goog-api-key', url.apiSettings.apiKey);
87+
if (url.apiSettings.automaticDataCollectionEnabled) {
88+
headers.append('X-Firebase-Appid', url.apiSettings.appId);
89+
}
8790
if (url.apiSettings.getAppCheckToken) {
8891
const appCheckToken = await url.apiSettings.getAppCheckToken();
8992
if (appCheckToken) {

‎packages/vertexai/src/types/error.ts

+3
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,9 @@ export const enum VertexAIErrorCode {
8787
/** An error occurred due to a missing Firebase API key. */
8888
NO_API_KEY = 'no-api-key',
8989

90+
/** An error occured due to a missing Firebase app ID. */
91+
NO_APP_ID = 'no-app-id',
92+
9093
/** An error occurred due to a model name not being specified during initialization. */
9194
NO_MODEL = 'no-model',
9295

‎packages/vertexai/src/types/internal.ts

+2
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ export * from './imagen/internal';
2323
export interface ApiSettings {
2424
apiKey: string;
2525
project: string;
26+
appId: string;
2627
location: string;
28+
automaticDataCollectionEnabled?: boolean;
2729
getAuthToken?: () => Promise<FirebaseAuthTokenData | null>;
2830
getAppCheckToken?: () => Promise<AppCheckTokenResult>;
2931
}

0 commit comments

Comments
 (0)