generate-content.ts

Author: russellwheatley

packages/ai/lib/methods/generate-content.ts

@@ -26,13 +26,18 @@ import { Task, makeRequest } from '../requests/request';
 import { createEnhancedContentResponse } from '../requests/response-helpers';
 import { processStream } from '../requests/stream-reader';
 import { ApiSettings } from '../types/internal';
+import { BackendType } from '../public-types';
+import * as GoogleAIMapper from '../googleai-mappers';
 
 export async function generateContentStream(
   apiSettings: ApiSettings,
   model: string,
   params: GenerateContentRequest,
   requestOptions?: RequestOptions,
 ): Promise<GenerateContentStreamResult> {
+  if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
+    params = GoogleAIMapper.mapGenerateContentRequest(params);
+  }
   const response = await makeRequest(
     model,
     Task.STREAM_GENERATE_CONTENT,
@@ -41,7 +46,7 @@ export async function generateContentStream(
     JSON.stringify(params),
     requestOptions,
   );
-  return processStream(response);
+  return processStream(response, apiSettings);
 }
 
 export async function generateContent(
@@ -50,6 +55,9 @@ export async function generateContent(
   params: GenerateContentRequest,
   requestOptions?: RequestOptions,
 ): Promise<GenerateContentResult> {
+  if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
+    params = GoogleAIMapper.mapGenerateContentRequest(params);
+  }
   const response = await makeRequest(
     model,
     Task.GENERATE_CONTENT,
@@ -58,9 +66,21 @@ export async function generateContent(
     JSON.stringify(params),
     requestOptions,
   );
-  const responseJson: GenerateContentResponse = await response.json();
-  const enhancedResponse = createEnhancedContentResponse(responseJson);
+  const generateContentResponse = await processGenerateContentResponse(response, apiSettings);
+  const enhancedResponse = createEnhancedContentResponse(generateContentResponse);
   return {
     response: enhancedResponse,
   };
 }
+
+async function processGenerateContentResponse(
+  response: Response,
+  apiSettings: ApiSettings,
+): Promise<GenerateContentResponse> {
+  const responseJson = await response.json();
+  if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
+    return GoogleAIMapper.mapGenerateContentResponse(responseJson);
+  } else {
+    return responseJson;
+  }
+}