stream-reader.ts

russellwheatley · mikehardy · commit 9a6b7958f2f9 · 2025-08-05T17:12:16.000-05:00
diff --git a/packages/ai/lib/googleai-mappers.ts b/packages/ai/lib/googleai-mappers.ts
@@ -0,0 +1,218 @@
+/**
+ * @license
+ * Copyright 2025 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 { AIError } from './errors';
+import { logger } from './logger';
+import {
+  CitationMetadata,
+  CountTokensRequest,
+  GenerateContentCandidate,
+  GenerateContentRequest,
+  GenerateContentResponse,
+  HarmSeverity,
+  InlineDataPart,
+  PromptFeedback,
+  SafetyRating,
+  AIErrorCode,
+} from './types';
+import {
+  GoogleAIGenerateContentResponse,
+  GoogleAIGenerateContentCandidate,
+  GoogleAICountTokensRequest,
+} from './types/googleai';
+
+/**
+ * This SDK supports both the Vertex AI Gemini API and the Gemini Developer API (using Google AI).
+ * The public API prioritizes the format used by the Vertex AI Gemini API.
+ * We avoid having two sets of types by translating requests and responses between the two API formats.
+ * This translation allows developers to switch between the Vertex AI Gemini API and the Gemini Developer API
+ * with minimal code changes.
+ *
+ * In here are functions that map requests and responses between the two API formats.
+ * Requests in the Vertex AI format are mapped to the Google AI format before being sent.
+ * Responses from the Google AI backend are mapped back to the Vertex AI format before being returned to the user.
+ */
+
+/**
+ * Maps a Vertex AI {@link GenerateContentRequest} to a format that can be sent to Google AI.
+ *
+ * @param generateContentRequest The {@link GenerateContentRequest} to map.
+ * @returns A {@link GenerateContentResponse} that conforms to the Google AI format.
+ *
+ * @throws If the request contains properties that are unsupported by Google AI.
+ *
+ * @internal
+ */
+export function mapGenerateContentRequest(
+  generateContentRequest: GenerateContentRequest,
+): GenerateContentRequest {
+  generateContentRequest.safetySettings?.forEach(safetySetting => {
+    if (safetySetting.method) {
+      throw new AIError(
+        AIErrorCode.UNSUPPORTED,
+        'SafetySetting.method is not supported in the the Gemini Developer API. Please remove this property.',
+      );
+    }
+  });
+
+  if (generateContentRequest.generationConfig?.topK) {
+    const roundedTopK = Math.round(generateContentRequest.generationConfig.topK);
+
+    if (roundedTopK !== generateContentRequest.generationConfig.topK) {
+      logger.warn(
+        'topK in GenerationConfig has been rounded to the nearest integer to match the format for requests to the Gemini Developer API.',
+      );
+      generateContentRequest.generationConfig.topK = roundedTopK;
+    }
+  }
+
+  return generateContentRequest;
+}
+
+/**
+ * Maps a {@link GenerateContentResponse} from Google AI to the format of the
+ * {@link GenerateContentResponse} that we get from VertexAI that is exposed in the public API.
+ *
+ * @param googleAIResponse The {@link GenerateContentResponse} from Google AI.
+ * @returns A {@link GenerateContentResponse} that conforms to the public API's format.
+ *
+ * @internal
+ */
+export function mapGenerateContentResponse(
+  googleAIResponse: GoogleAIGenerateContentResponse,
+): GenerateContentResponse {
+  const generateContentResponse = {
+    candidates: googleAIResponse.candidates
+      ? mapGenerateContentCandidates(googleAIResponse.candidates)
+      : undefined,
+    prompt: googleAIResponse.promptFeedback
+      ? mapPromptFeedback(googleAIResponse.promptFeedback)
+      : undefined,
+    usageMetadata: googleAIResponse.usageMetadata,
+  };
+
+  return generateContentResponse;
+}
+
+/**
+ * Maps a Vertex AI {@link CountTokensRequest} to a format that can be sent to Google AI.
+ *
+ * @param countTokensRequest The {@link CountTokensRequest} to map.
+ * @param model The model to count tokens with.
+ * @returns A {@link CountTokensRequest} that conforms to the Google AI format.
+ *
+ * @internal
+ */
+export function mapCountTokensRequest(
+  countTokensRequest: CountTokensRequest,
+  model: string,
+): GoogleAICountTokensRequest {
+  const mappedCountTokensRequest: GoogleAICountTokensRequest = {
+    generateContentRequest: {
+      model,
+      ...countTokensRequest,
+    },
+  };
+
+  return mappedCountTokensRequest;
+}
+
+/**
+ * Maps a Google AI {@link GoogleAIGenerateContentCandidate} to a format that conforms
+ * to the Vertex AI API format.
+ *
+ * @param candidates The {@link GoogleAIGenerateContentCandidate} to map.
+ * @returns A {@link GenerateContentCandidate} that conforms to the Vertex AI format.
+ *
+ * @throws If any {@link Part} in the candidates has a `videoMetadata` property.
+ *
+ * @internal
+ */
+export function mapGenerateContentCandidates(
+  candidates: GoogleAIGenerateContentCandidate[],
+): GenerateContentCandidate[] {
+  const mappedCandidates: GenerateContentCandidate[] = [];
+  let mappedSafetyRatings: SafetyRating[];
+  if (mappedCandidates) {
+    candidates.forEach(candidate => {
+      // Map citationSources to citations.
+      let citationMetadata: CitationMetadata | undefined;
+      if (candidate.citationMetadata) {
+        citationMetadata = {
+          citations: candidate.citationMetadata.citationSources,
+        };
+      }
+
+      // Assign missing candidate SafetyRatings properties to their defaults if undefined.
+      if (candidate.safetyRatings) {
+        mappedSafetyRatings = candidate.safetyRatings.map(safetyRating => {
+          return {
+            ...safetyRating,
+            severity: safetyRating.severity ?? HarmSeverity.HARM_SEVERITY_UNSUPPORTED,
+            probabilityScore: safetyRating.probabilityScore ?? 0,
+            severityScore: safetyRating.severityScore ?? 0,
+          };
+        });
+      }
+
+      // videoMetadata is not supported.
+      // Throw early since developers may send a long video as input and only expect to pay
+      // for inference on a small portion of the video.
+      if (candidate.content?.parts.some(part => (part as InlineDataPart)?.videoMetadata)) {
+        throw new AIError(
+          AIErrorCode.UNSUPPORTED,
+          'Part.videoMetadata is not supported in the Gemini Developer API. Please remove this property.',
+        );
+      }
+
+      const mappedCandidate = {
+        index: candidate.index,
+        content: candidate.content,
+        finishReason: candidate.finishReason,
+        finishMessage: candidate.finishMessage,
+        safetyRatings: mappedSafetyRatings,
+        citationMetadata,
+        groundingMetadata: candidate.groundingMetadata,
+      };
+      mappedCandidates.push(mappedCandidate);
+    });
+  }
+
+  return mappedCandidates;
+}
+
+export function mapPromptFeedback(promptFeedback: PromptFeedback): PromptFeedback {
+  // 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: safetyRating.severity ?? HarmSeverity.HARM_SEVERITY_UNSUPPORTED,
+      probabilityScore: safetyRating.probabilityScore ?? 0,
+      severityScore: safetyRating.severityScore ?? 0,
+      blocked: safetyRating.blocked,
+    });
+  });
+
+  const mappedPromptFeedback: PromptFeedback = {
+    blockReason: promptFeedback.blockReason,
+    safetyRatings: mappedSafetyRatings,
+    blockReasonMessage: promptFeedback.blockReasonMessage,
+  };
+  return mappedPromptFeedback;
+}
diff --git a/packages/ai/lib/requests/stream-reader.ts b/packages/ai/lib/requests/stream-reader.ts
@@ -22,10 +22,14 @@ import {
   GenerateContentResponse,
   GenerateContentStreamResult,
   Part,
-  VertexAIErrorCode,
+  AIErrorCode,
 } from '../types';
-import { VertexAIError } from '../errors';
+import { AIError } from '../errors';
 import { createEnhancedContentResponse } from './response-helpers';
+import { ApiSettings } from '../types/internal';
+import { BackendType } from '../public-types';
+import * as GoogleAIMapper from '../googleai-mappers';
+import { GoogleAIGenerateContentResponse } from '../types/googleai';
 
 const responseLineRE = /^data\: (.*)(?:\n\n|\r\r|\r\n\r\n)/;
 
@@ -37,7 +41,10 @@ const responseLineRE = /^data\: (.*)(?:\n\n|\r\r|\r\n\r\n)/;
  *
  * @param response - Response from a fetch call
  */
-export function processStream(response: Response): GenerateContentStreamResult {
+export function processStream(
+  response: Response,
+  apiSettings: ApiSettings,
+): GenerateContentStreamResult {
   const inputStream = new ReadableStream<string>({
     async start(controller) {
       const reader = response.body!.getReader();
@@ -56,28 +63,36 @@ export function processStream(response: Response): GenerateContentStreamResult {
   const responseStream = getResponseStream<GenerateContentResponse>(inputStream);
   const [stream1, stream2] = responseStream.tee();
   return {
-    stream: generateResponseSequence(stream1),
-    response: getResponsePromise(stream2),
+    stream: generateResponseSequence(stream1, apiSettings),
+    response: getResponsePromise(stream2, apiSettings),
   };
 }
 
 async function getResponsePromise(
   stream: ReadableStream<GenerateContentResponse>,
+  apiSettings: ApiSettings,
 ): Promise<EnhancedGenerateContentResponse> {
   const allResponses: GenerateContentResponse[] = [];
   const reader = stream.getReader();
   while (true) {
     const { done, value } = await reader.read();
     if (done) {
-      const enhancedResponse = createEnhancedContentResponse(aggregateResponses(allResponses));
-      return enhancedResponse;
+      let generateContentResponse = aggregateResponses(allResponses);
+      if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
+        generateContentResponse = GoogleAIMapper.mapGenerateContentResponse(
+          generateContentResponse as GoogleAIGenerateContentResponse,
+        );
+      }
+      return createEnhancedContentResponse(generateContentResponse);
     }
+
     allResponses.push(value);
   }
 }
 
 async function* generateResponseSequence(
   stream: ReadableStream<GenerateContentResponse>,
+  apiSettings: ApiSettings,
 ): AsyncGenerator<EnhancedGenerateContentResponse> {
   const reader = stream.getReader();
   while (true) {
@@ -86,7 +101,15 @@ async function* generateResponseSequence(
       break;
     }
 
-    const enhancedResponse = createEnhancedContentResponse(value);
+    let enhancedResponse: EnhancedGenerateContentResponse;
+    if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
+      enhancedResponse = createEnhancedContentResponse(
+        GoogleAIMapper.mapGenerateContentResponse(value as GoogleAIGenerateContentResponse),
+      );
+    } else {
+      enhancedResponse = createEnhancedContentResponse(value);
+    }
+
     yield enhancedResponse;
   }
 }
@@ -106,9 +129,7 @@ export function getResponseStream<T>(inputStream: ReadableStream<string>): Reada
         return reader.read().then(({ value, done }) => {
           if (done) {
             if (currentText.trim()) {
-              controller.error(
-                new VertexAIError(VertexAIErrorCode.PARSE_FAILED, 'Failed to parse stream'),
-              );
+              controller.error(new AIError(AIErrorCode.PARSE_FAILED, 'Failed to parse stream'));
               return;
             }
             controller.close();
@@ -123,10 +144,7 @@ export function getResponseStream<T>(inputStream: ReadableStream<string>): Reada
               parsedResponse = JSON.parse(match[1]!);
             } catch (_) {
               controller.error(
-                new VertexAIError(
-                  VertexAIErrorCode.PARSE_FAILED,
-                  `Error parsing JSON response: "${match[1]}`,
-                ),
+                new AIError(AIErrorCode.PARSE_FAILED, `Error parsing JSON response: "${match[1]}`),
               );
               return;
             }
@@ -197,8 +215,8 @@ export function aggregateResponses(responses: GenerateContentResponse[]): Genera
               newPart.functionCall = part.functionCall;
             }
             if (Object.keys(newPart).length === 0) {
-              throw new VertexAIError(
-                VertexAIErrorCode.INVALID_CONTENT,
+              throw new AIError(
+                AIErrorCode.INVALID_CONTENT,
                 'Part should have at least one property, but there are none. This is likely caused ' +
                   'by a malformed response from the backend.',
               );
diff --git a/packages/ai/lib/types/enums.ts b/packages/ai/lib/types/enums.ts
@@ -91,6 +91,13 @@ export enum HarmSeverity {
   HARM_SEVERITY_MEDIUM = 'HARM_SEVERITY_MEDIUM',
   // High level of harm severity.
   HARM_SEVERITY_HIGH = 'HARM_SEVERITY_HIGH',
+  /**
+   * Harm severity is not supported.
+   *
+   * @remarks
+   * The GoogleAI backend does not support `HarmSeverity`, so this value is used as a fallback.
+   */
+  HARM_SEVERITY_UNSUPPORTED = 'HARM_SEVERITY_UNSUPPORTED',
 }
 
 /**
diff --git a/packages/ai/lib/types/googleai.ts b/packages/ai/lib/types/googleai.ts
