From 96c65a2aa9ee8f9b1b2059f3063f922aea251ec5 Mon Sep 17 00:00:00 2001 From: Christina Holland Date: Wed, 23 Oct 2024 10:13:01 -0700 Subject: [PATCH 1/2] Use logger class and send App Check dummy token in header. --- .changeset/tall-peas-tell.md | 5 +++++ packages/vertexai/src/logger.ts | 20 +++++++++++++++++++ .../vertexai/src/methods/chat-session.test.ts | 2 +- packages/vertexai/src/methods/chat-session.ts | 7 ++++--- .../vertexai/src/requests/request.test.ts | 12 ++++++++--- packages/vertexai/src/requests/request.ts | 8 +++++++- .../vertexai/src/requests/response-helpers.ts | 5 +++-- 7 files changed, 49 insertions(+), 10 deletions(-) create mode 100644 .changeset/tall-peas-tell.md create mode 100644 packages/vertexai/src/logger.ts diff --git a/.changeset/tall-peas-tell.md b/.changeset/tall-peas-tell.md new file mode 100644 index 00000000000..59af1cb8ab8 --- /dev/null +++ b/.changeset/tall-peas-tell.md @@ -0,0 +1,5 @@ +--- +'@firebase/vertexai': patch +--- + +Send App Check dummy token in header if there is an App Check getToken error. diff --git a/packages/vertexai/src/logger.ts b/packages/vertexai/src/logger.ts new file mode 100644 index 00000000000..7664d6b4e31 --- /dev/null +++ b/packages/vertexai/src/logger.ts @@ -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'); diff --git a/packages/vertexai/src/methods/chat-session.test.ts b/packages/vertexai/src/methods/chat-session.test.ts index c8c92e046b9..7741c33ea0b 100644 --- a/packages/vertexai/src/methods/chat-session.test.ts +++ b/packages/vertexai/src/methods/chat-session.test.ts @@ -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' ); diff --git a/packages/vertexai/src/methods/chat-session.ts b/packages/vertexai/src/methods/chat-session.ts index bfa3f2a660f..dd22b29a7c8 100644 --- a/packages/vertexai/src/methods/chat-session.ts +++ b/packages/vertexai/src/methods/chat-session.ts @@ -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. @@ -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.` ); } @@ -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.` ); } @@ -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; diff --git a/packages/vertexai/src/requests/request.test.ts b/packages/vertexai/src/requests/request.test.ts index da32ec7d59c..957a3321125 100644 --- a/packages/vertexai/src/requests/request.test.ts +++ b/packages/vertexai/src/requests/request.test.ts @@ -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'; @@ -25,6 +25,7 @@ import { DEFAULT_API_VERSION } from '../constants'; import { VertexAIErrorCode } from '../types'; import { VertexAIError } from '../errors'; import { getMockResponse } from '../../test-utils/mock-response'; +import { base64 } from '@firebase/util'; use(sinonChai); use(chaiAsPromised); @@ -169,13 +170,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); diff --git a/packages/vertexai/src/requests/request.ts b/packages/vertexai/src/requests/request.ts index edb894d7f76..26a5ac523f3 100644 --- a/packages/vertexai/src/requests/request.ts +++ b/packages/vertexai/src/requests/request.ts @@ -24,6 +24,7 @@ import { LANGUAGE_TAG, PACKAGE_VERSION } from '../constants'; +import { logger } from '../logger'; export enum Task { GENERATE_CONTENT = 'generateContent', @@ -83,8 +84,13 @@ export async function getHeaders(url: RequestUrl): Promise { 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}` + ); + } } } diff --git a/packages/vertexai/src/requests/response-helpers.ts b/packages/vertexai/src/requests/response-helpers.ts index 0435f1711ce..27347d10f0d 100644 --- a/packages/vertexai/src/requests/response-helpers.ts +++ b/packages/vertexai/src/requests/response-helpers.ts @@ -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 @@ -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.` @@ -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.` From 9e8ef567a0c801988ecf4bf24aa1abb68230b5dc Mon Sep 17 00:00:00 2001 From: Christina Holland Date: Wed, 23 Oct 2024 13:16:36 -0700 Subject: [PATCH 2/2] Remove extra import --- packages/vertexai/src/requests/request.test.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/vertexai/src/requests/request.test.ts b/packages/vertexai/src/requests/request.test.ts index 957a3321125..b6d0ecb9b71 100644 --- a/packages/vertexai/src/requests/request.test.ts +++ b/packages/vertexai/src/requests/request.test.ts @@ -25,7 +25,6 @@ import { DEFAULT_API_VERSION } from '../constants'; import { VertexAIErrorCode } from '../types'; import { VertexAIError } from '../errors'; import { getMockResponse } from '../../test-utils/mock-response'; -import { base64 } from '@firebase/util'; use(sinonChai); use(chaiAsPromised);