Skip to content

Commit

Permalink
Merge 9e8ef56 into 052e438
Browse files Browse the repository at this point in the history
  • Loading branch information
hsubox76 authored Oct 23, 2024
2 parents 052e438 + 9e8ef56 commit 109f9be
Show file tree
Hide file tree
Showing 7 changed files with 48 additions and 10 deletions.
5 changes: 5 additions & 0 deletions .changeset/tall-peas-tell.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@firebase/vertexai': patch
---

Send App Check dummy token in header if there is an App Check getToken error.
20 changes: 20 additions & 0 deletions packages/vertexai/src/logger.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
/**
* @license
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

import { Logger } from '@firebase/logger';

export const logger = new Logger('@firebase/vertexai');
2 changes: 1 addition & 1 deletion packages/vertexai/src/methods/chat-session.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ describe('ChatSession', () => {
match.any
);
await clock.runAllAsync();
expect(consoleStub.args[0][0].toString()).to.include(
expect(consoleStub.args[0][1].toString()).to.include(
// Firefox has different wording when a property is undefined
'undefined'
);
Expand Down
7 changes: 4 additions & 3 deletions packages/vertexai/src/methods/chat-session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import { formatBlockErrorMessage } from '../requests/response-helpers';
import { validateChatHistory } from './chat-session-helpers';
import { generateContent, generateContentStream } from './generate-content';
import { ApiSettings } from '../types/internal';
import { logger } from '../logger';

/**
* Do not log a message for this error.
Expand Down Expand Up @@ -112,7 +113,7 @@ export class ChatSession {
} else {
const blockErrorMessage = formatBlockErrorMessage(result.response);
if (blockErrorMessage) {
console.warn(
logger.warn(
`sendMessage() was unsuccessful. ${blockErrorMessage}. Inspect response object for details.`
);
}
Expand Down Expand Up @@ -169,7 +170,7 @@ export class ChatSession {
} else {
const blockErrorMessage = formatBlockErrorMessage(response);
if (blockErrorMessage) {
console.warn(
logger.warn(
`sendMessageStream() was unsuccessful. ${blockErrorMessage}. Inspect response object for details.`
);
}
Expand All @@ -182,7 +183,7 @@ export class ChatSession {
if (e.message !== SILENT_ERROR) {
// Users do not have access to _sendPromise to catch errors
// downstream from streamPromise, so they should not throw.
console.error(e);
logger.error(e);
}
});
return streamPromise;
Expand Down
11 changes: 8 additions & 3 deletions packages/vertexai/src/requests/request.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
*/

import { expect, use } from 'chai';
import { restore, stub } from 'sinon';
import { match, restore, stub } from 'sinon';
import sinonChai from 'sinon-chai';
import chaiAsPromised from 'chai-as-promised';
import { RequestUrl, Task, getHeaders, makeRequest } from './request';
Expand Down Expand Up @@ -169,13 +169,18 @@ describe('request methods', () => {
project: 'myproject',
location: 'moon',
getAppCheckToken: () =>
Promise.resolve({ token: 'token', error: Error('oops') })
Promise.resolve({ token: 'dummytoken', error: Error('oops') })
},
true,
{}
);
const warnStub = stub(console, 'warn');
const headers = await getHeaders(fakeUrl);
expect(headers.has('X-Firebase-AppCheck')).to.be.false;
expect(headers.get('X-Firebase-AppCheck')).to.equal('dummytoken');
expect(warnStub).to.be.calledWith(
match(/vertexai/),
match(/App Check.*oops/)
);
});
it('adds auth token if it exists', async () => {
const headers = await getHeaders(fakeUrl);
Expand Down
8 changes: 7 additions & 1 deletion packages/vertexai/src/requests/request.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
LANGUAGE_TAG,
PACKAGE_VERSION
} from '../constants';
import { logger } from '../logger';

export enum Task {
GENERATE_CONTENT = 'generateContent',
Expand Down Expand Up @@ -83,8 +84,13 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
headers.append('x-goog-api-key', url.apiSettings.apiKey);
if (url.apiSettings.getAppCheckToken) {
const appCheckToken = await url.apiSettings.getAppCheckToken();
if (appCheckToken && !appCheckToken.error) {
if (appCheckToken) {
headers.append('X-Firebase-AppCheck', appCheckToken.token);
if (appCheckToken.error) {
logger.warn(
`Unable to obtain a valid App Check token: ${appCheckToken.error.message}`
);
}
}
}

Expand Down
5 changes: 3 additions & 2 deletions packages/vertexai/src/requests/response-helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
VertexAIErrorCode
} from '../types';
import { VertexAIError } from '../errors';
import { logger } from '../logger';

/**
* Creates an EnhancedGenerateContentResponse object that has helper functions and
Expand Down Expand Up @@ -56,7 +57,7 @@ export function addHelpers(
(response as EnhancedGenerateContentResponse).text = () => {
if (response.candidates && response.candidates.length > 0) {
if (response.candidates.length > 1) {
console.warn(
logger.warn(
`This response had ${response.candidates.length} ` +
`candidates. Returning text from the first candidate only. ` +
`Access response.candidates directly to use the other candidates.`
Expand Down Expand Up @@ -88,7 +89,7 @@ export function addHelpers(
(response as EnhancedGenerateContentResponse).functionCalls = () => {
if (response.candidates && response.candidates.length > 0) {
if (response.candidates.length > 1) {
console.warn(
logger.warn(
`This response had ${response.candidates.length} ` +
`candidates. Returning function calls from the first candidate only. ` +
`Access response.candidates directly to use the other candidates.`
Expand Down

0 comments on commit 109f9be

Please sign in to comment.