Cleanup

Author: dlarocque

packages/vertexai/src/googleAIMappers.test.ts

@@ -92,7 +92,7 @@ describe('Google AI Mappers', () => {
       };
       const mappedRequest = mapGenerateContentRequest(request);
       expect(loggerWarnStub).to.have.been.calledOnceWith(
-        'topK in GenerationConfig has been rounded to the nearest integer.'
+        'topK in GenerationConfig has been rounded to the nearest integer to match the format for Google AI requests.'
       );
       expect(mappedRequest.generationConfig?.topK).to.equal(16);
     });
@@ -133,12 +133,8 @@ describe('Google AI Mappers', () => {
       const googleAIMockResponse: GoogleAIGenerateContentResponse = await (
         getMockResponse('googleAI', 'unary-success-citations.txt') as Response
       ).json();
-      console.log(JSON.stringify(googleAIMockResponse));
       const mappedResponse = mapGenerateContentResponse(googleAIMockResponse);
 
-      console.log(JSON.stringify(googleAIMockResponse));
-      console.log(JSON.stringify(mappedResponse));
-
       expect(mappedResponse.candidates).to.exist;
       expect(mappedResponse.candidates?.[0].content.parts[0].text).to.contain(
         'quantum mechanics'
@@ -242,16 +238,17 @@ describe('Google AI Mappers', () => {
 
     it('should map a minimal Vertex AI CountTokensRequest', () => {
       const vertexRequest: CountTokensRequest = {
-        contents: fakeContents
+        contents: fakeContents,
+        systemInstruction: { role: 'system', parts: [{ text: 'Be nice' }] },
+        generationConfig: { temperature: 0.8 }
       };
 
       const expectedGoogleAIRequest: GoogleAICountTokensRequest = {
         generateContentRequest: {
           model: fakeModel,
           contents: vertexRequest.contents,
-          systemInstruction: undefined,
-          tools: undefined,
-          generationConfig: undefined
+          systemInstruction: { role: 'system', parts: [{ text: 'Be nice' }] },
+          generationConfig: { temperature: 0.8 }
         }
       };
 

packages/vertexai/src/googleAIMappers.ts

@@ -76,7 +76,7 @@ export function mapGenerateContentRequest(
 
     if (roundedTopK !== generateContentRequest.generationConfig.topK) {
       logger.warn(
-        'topK in GenerationConfig has been rounded to the nearest integer.'
+        'topK in GenerationConfig has been rounded to the nearest integer to match the format for Google AI requests.'
       );
       generateContentRequest.generationConfig.topK = roundedTopK;
     }
@@ -126,10 +126,7 @@ export function mapCountTokensRequest(
   const mappedCountTokensRequest: GoogleAICountTokensRequest = {
     generateContentRequest: {
       model,
-      contents: countTokensRequest.contents,
-      systemInstruction: countTokensRequest.systemInstruction,
-      tools: countTokensRequest.tools,
-      generationConfig: countTokensRequest.generationConfig
+      ...countTokensRequest
     }
   };
 

packages/vertexai/src/helpers.test.ts

@@ -29,7 +29,6 @@ describe('Identifier Encoding/Decoding', () => {
         backendType: BackendType.VERTEX_AI,
         location: 'us-central1'
       };
-      console.log(identifier);
       const expected = `${AI_TYPE}/vertexai/us-central1`;
       expect(encodeInstanceIdentifier(identifier)).to.equal(expected);
     });

packages/vertexai/test-utils/mock-response.ts

@@ -72,9 +72,6 @@ export function getMockResponse(
   filename: string
 ): Partial<Response> {
   const mocksLookup = mockSetMaps[backendName];
-  if (backendName === 'googleAI') {
-    console.log(Object.keys(mocksLookup));
-  }
   if (!(filename in mocksLookup)) {
     throw Error(`${backendName} mock response file '${filename}' not found.`);
   }