Move GoogleAI types to single file

Author: dlarocque

common/api-review/vertexai.api.md

@@ -125,11 +125,9 @@ export interface Citation {
     endIndex?: number;
     // (undocumented)
     license?: string;
-    // (undocumented)
     publicationDate?: Date_2;
     // (undocumented)
     startIndex?: number;
-    // (undocumented)
     title?: string;
     // (undocumented)
     uri?: string;
@@ -420,6 +418,60 @@ export function getImagenModel(ai: AI, modelParams: ImagenModelParams, requestOp
 // @public
 export function getVertexAI(app?: FirebaseApp, options?: VertexAIOptions): VertexAI;
 
+// Warning: (ae-internal-missing-underscore) The name "GoogleAICitationMetadata" should be prefixed with an underscore because the declaration is marked as @internal
+//
+// @internal (undocumented)
+export interface GoogleAICitationMetadata {
+    // (undocumented)
+    citationSources: Citation[];
+}
+
+// Warning: (ae-internal-missing-underscore) The name "GoogleAICountTokensRequest" should be prefixed with an underscore because the declaration is marked as @internal
+//
+// @internal (undocumented)
+export interface GoogleAICountTokensRequest {
+    // (undocumented)
+    generateContentRequest: {
+        model: string;
+        contents: Content[];
+        systemInstruction?: string | Part | Content;
+        tools?: Tool[];
+        generationConfig?: GenerationConfig;
+    };
+}
+
+// Warning: (ae-internal-missing-underscore) The name "GoogleAIGenerateContentCandidate" should be prefixed with an underscore because the declaration is marked as @internal
+//
+// @internal (undocumented)
+export interface GoogleAIGenerateContentCandidate {
+    // (undocumented)
+    citationMetadata?: GoogleAICitationMetadata;
+    // (undocumented)
+    content: Content;
+    // (undocumented)
+    finishMessage?: string;
+    // (undocumented)
+    finishReason?: FinishReason;
+    // (undocumented)
+    groundingMetadata?: GroundingMetadata;
+    // (undocumented)
+    index: number;
+    // (undocumented)
+    safetyRatings?: SafetyRating[];
+}
+
+// Warning: (ae-internal-missing-underscore) The name "GoogleAIGenerateContentResponse" should be prefixed with an underscore because the declaration is marked as @internal
+//
+// @internal (undocumented)
+export interface GoogleAIGenerateContentResponse {
+    // (undocumented)
+    candidates?: GoogleAIGenerateContentCandidate[];
+    // (undocumented)
+    promptFeedback?: PromptFeedback;
+    // (undocumented)
+    usageMetadata?: UsageMetadata;
+}
+
 // @public @deprecated (undocumented)
 export interface GroundingAttribution {
     // (undocumented)
@@ -653,7 +705,6 @@ export const POSSIBLE_ROLES: readonly ["user", "model", "function", "system"];
 export interface PromptFeedback {
     // (undocumented)
     blockReason?: BlockReason;
-    // (undocumented)
     blockReasonMessage?: string;
     // (undocumented)
     safetyRatings: SafetyRating[];
@@ -684,11 +735,8 @@ export interface SafetyRating {
     category: HarmCategory;
     // (undocumented)
     probability: HarmProbability;
-    // (undocumented)
     probabilityScore: number;
-    // (undocumented)
     severity: HarmSeverity;
-    // (undocumented)
     severityScore: number;
 }
 

packages/vertexai/src/googleAIMappers.test.ts renamed to packages/vertexai/src/googleai-mappers.test.ts

@@ -24,7 +24,7 @@ import {
   mapGenerateContentRequest,
   mapGenerateContentResponse,
   mapPromptFeedback
-} from './googleAIMappers';
+} from './googleai-mappers';
 import {
   BlockReason,
   Content,
@@ -44,7 +44,7 @@ import {
   GoogleAIGenerateContentResponse,
   GoogleAIGenerateContentCandidate,
   GoogleAICountTokensRequest
-} from './types/googleAI';
+} from './types/googleai';
 import { logger } from './logger';
 import { AIError } from './errors';
 import { getMockResponse } from '../test-utils/mock-response';
@@ -281,7 +281,7 @@ describe('Google AI Mappers', () => {
         .undefined; // Not in Google AI
     });
 
-    it('should add default safety rating properties and warn', () => {
+    it('should add default safety rating properties', () => {
       const candidates: GoogleAIGenerateContentCandidate[] = [
         {
           index: 0,
@@ -297,9 +297,6 @@ describe('Google AI Mappers', () => {
         }
       ];
       const mapped = mapGenerateContentCandidates(candidates);
-      expect(loggerWarnStub).to.have.been.calledOnceWith(
-        "Candidate safety rating properties 'severity', 'severityScore', and 'probabilityScore' are not included in responses from Google AI. Properties have been assigned to default values."
-      );
       expect(mapped[0].safetyRatings).to.exist;
       const safetyRating = mapped[0].safetyRatings?.[0] as SafetyRating; // Type assertion
       expect(safetyRating.severity).to.equal(
@@ -359,7 +356,7 @@ describe('Google AI Mappers', () => {
   });
 
   describe('mapPromptFeedback', () => {
-    it('should add default safety rating properties and warn', () => {
+    it('should add default safety rating properties', () => {
       const feedback: PromptFeedback = {
         blockReason: BlockReason.OTHER,
         safetyRatings: [
@@ -373,9 +370,6 @@ describe('Google AI Mappers', () => {
         // Missing blockReasonMessage
       };
       const mapped = mapPromptFeedback(feedback);
-      expect(loggerWarnStub).to.have.been.calledOnceWith(
-        "PromptFeedback safety ratings' properties severity, severityScore, and probabilityScore are not included in responses from Google AI. Properties have been assigned to default values."
-      );
       expect(mapped.safetyRatings).to.exist;
       const safetyRating = mapped.safetyRatings[0] as SafetyRating; // Type assertion
       expect(safetyRating.severity).to.equal(

packages/vertexai/src/googleAIMappers.ts renamed to packages/vertexai/src/googleai-mappers.ts

@@ -33,7 +33,7 @@ import {
   GoogleAIGenerateContentResponse,
   GoogleAIGenerateContentCandidate,
   GoogleAICountTokensRequest
-} from './types/googleAI';
+} from './types/googleai';
 
 /**
  * This SDK supports both Vertex AI and Google AI APIs.
@@ -159,17 +159,14 @@ export function mapGenerateContentCandidates(
         };
       }
 
-      // Assign missing candidate SafetyRatings properties to their defaults.
+      // Assign missing candidate SafetyRatings properties to their defaults if undefined.
       if (candidate.safetyRatings) {
-        logger.warn(
-          "Candidate safety rating properties 'severity', 'severityScore', and 'probabilityScore' are not included in responses from Google AI. Properties have been assigned to default values."
-        );
         mappedSafetyRatings = candidate.safetyRatings.map(safetyRating => {
           return {
             ...safetyRating,
-            severity: HarmSeverity.HARM_SEVERITY_UNSUPPORTED,
-            probabilityScore: 0,
-            severityScore: 0
+            severity: safetyRating.severity ?? HarmSeverity.HARM_SEVERITY_UNSUPPORTED,
+            probabilityScore: safetyRating.probabilityScore ?? 0,
+            severityScore: safetyRating.severityScore ?? 0
           };
         });
       }
@@ -207,21 +204,18 @@ export function mapGenerateContentCandidates(
 export function mapPromptFeedback(
   promptFeedback: PromptFeedback
 ): PromptFeedback {
-  // Assign missing PromptFeedback SafetyRatings properties to their defaults.
+  // Assign missing SafetyRating properties to their defaults if undefined.
   const mappedSafetyRatings: SafetyRating[] = [];
   promptFeedback.safetyRatings.forEach(safetyRating => {
     mappedSafetyRatings.push({
       category: safetyRating.category,
       probability: safetyRating.probability,
-      severity: HarmSeverity.HARM_SEVERITY_UNSUPPORTED,
-      probabilityScore: 0,
-      severityScore: 0,
+      severity: safetyRating.severity ?? HarmSeverity.HARM_SEVERITY_UNSUPPORTED,
+      probabilityScore: safetyRating.probabilityScore ?? 0,
+      severityScore: safetyRating.severityScore ?? 0,
       blocked: safetyRating.blocked
     });
   });
-  logger.warn(
-    "PromptFeedback safety ratings' properties severity, severityScore, and probabilityScore are not included in responses from Google AI. Properties have been assigned to default values."
-  );
 
   const mappedPromptFeedback: PromptFeedback = {
     blockReason: promptFeedback.blockReason,

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

@@ -25,7 +25,7 @@ import { countTokens } from './count-tokens';
 import { CountTokensRequest } from '../types';
 import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
-import { mapCountTokensRequest } from '../googleAIMappers';
+import { mapCountTokensRequest } from '../googleai-mappers';
 import { GoogleAIBackend, VertexAIBackend } from '../backend';
 
 use(sinonChai);

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

@@ -22,7 +22,7 @@ import {
 } from '../types';
 import { Task, makeRequest } from '../requests/request';
 import { ApiSettings } from '../types/internal';
-import * as GoogleAIMapper from '../googleAIMappers';
+import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
 
 export async function countTokens(

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

@@ -32,7 +32,7 @@ import {
 import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
 import { AIError } from '../api';
-import { mapGenerateContentRequest } from '../googleAIMappers';
+import { mapGenerateContentRequest } from '../googleai-mappers';
 import { GoogleAIBackend, VertexAIBackend } from '../backend';
 
 use(sinonChai);

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

@@ -26,7 +26,7 @@ import { Task, makeRequest } from '../requests/request';
 import { createEnhancedContentResponse } from '../requests/response-helpers';
 import { processStream } from '../requests/stream-reader';
 import { ApiSettings } from '../types/internal';
-import * as GoogleAIMapper from '../googleAIMappers';
+import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
 
 export async function generateContentStream(
@@ -66,7 +66,7 @@ export async function generateContent(
     JSON.stringify(params),
     requestOptions
   );
-  const generateContentResponse = await handleGenerateContentResponse(
+  const generateContentResponse = await processGenerateContentResponse(
     response,
     apiSettings
   );
@@ -78,7 +78,7 @@ export async function generateContent(
   };
 }
 
-async function handleGenerateContentResponse(
+async function processGenerateContentResponse(
   response: Response,
   apiSettings: ApiSettings
 ): Promise<GenerateContentResponse> {

packages/vertexai/src/public-types.ts

@@ -23,7 +23,7 @@ export * from './types';
 /**
  * An instance of the Firebase AI SDK.
  *
- * For more information, refer to the documentation for the new {@link AI}.
+ * For more information, refer to the documentation for the new {@link AI} interface.
  *
  * @public
  */

packages/vertexai/src/requests/stream-reader.ts

@@ -25,8 +25,8 @@ import {
 } from '../types';
 import { AIError } from '../errors';
 import { createEnhancedContentResponse } from './response-helpers';
-import * as GoogleAIMapper from '../googleAIMappers';
-import { GoogleAIGenerateContentResponse } from '../types/googleAI';
+import * as GoogleAIMapper from '../googleai-mappers';
+import { GoogleAIGenerateContentResponse } from '../types/googleai';
 import { ApiSettings } from '../types/internal';
 import { BackendType } from '../public-types';