feat: Implement cached_content with generateContent methods

PiperOrigin-RevId: 675825556

src/functions/generate_content.ts

@@ -72,6 +72,7 @@ export async function generateContent(
   const generateContentRequest: GenerateContentRequest = {
     contents: request.contents,
     systemInstruction: request.systemInstruction,
+    cachedContent: request.cachedContent,
     generationConfig: request.generationConfig ?? generationConfig,
     safetySettings: request.safetySettings ?? safetySettings,
     tools: request.tools ?? tools,
@@ -126,6 +127,7 @@ export async function generateContentStream(
   const generateContentRequest: GenerateContentRequest = {
     contents: request.contents,
     systemInstruction: request.systemInstruction,
+    cachedContent: request.cachedContent,
     generationConfig: request.generationConfig ?? generationConfig,
     safetySettings: request.safetySettings ?? safetySettings,
     tools: request.tools ?? tools,

src/functions/pre_fetch_processing.ts

@@ -80,7 +80,9 @@ export function validateGenerationConfig(
 export function getApiVersion(
   request: GenerateContentRequest
 ): 'v1' | 'v1beta1' {
-  return hasVertexRagStore(request) ? 'v1beta1' : 'v1';
+  return hasVertexRagStore(request) || hasCachedContent(request)
+    ? 'v1beta1'
+    : 'v1';
 }
 
 export function hasVertexRagStore(request: GenerateContentRequest): boolean {
@@ -94,6 +96,10 @@ export function hasVertexRagStore(request: GenerateContentRequest): boolean {
   return false;
 }
 
+function hasCachedContent(request: GenerateContentRequest): boolean {
+  return !!request.cachedContent;
+}
+
 export function hasVertexAISearch(request: GenerateContentRequest): boolean {
   for (const tool of request?.tools ?? []) {
     const retrieval = (tool as RetrievalTool).retrieval;

src/functions/util.ts

@@ -0,0 +1,32 @@
+/**
+ * @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
+ *
+ *     https://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 {Content} from '../types/content';
+import {constants} from '../util';
+
+export function formulateSystemInstructionIntoContent(
+  systemInstruction: string | Content
+): Content {
+  if (typeof systemInstruction === 'string') {
+    return {
+      role: constants.SYSTEM_ROLE,
+      parts: [{text: systemInstruction}],
+    } as Content;
+  }
+  systemInstruction.role = constants.SYSTEM_ROLE;
+  return systemInstruction;
+}

src/models/chat_session.ts

@@ -18,7 +18,7 @@
 /* tslint:disable */
 import {GoogleAuth} from 'google-auth-library';
 
-import {formulateSystemInstructionIntoContent} from './util';
+import {formulateSystemInstructionIntoContent} from '../functions/util';
 import {
   generateContent,
   generateContentStream,

src/models/generative_models.ts

@@ -18,13 +18,14 @@
 /* tslint:disable */
 import {GoogleAuth} from 'google-auth-library';
 
-import {formulateSystemInstructionIntoContent} from './util';
+import {formulateSystemInstructionIntoContent} from '../functions/util';
 import {countTokens} from '../functions/count_tokens';
 import {
   generateContent,
   generateContentStream,
 } from '../functions/generate_content';
 import {
+  CachedContent,
   Content,
   CountTokensRequest,
   CountTokensResponse,
@@ -295,6 +296,7 @@ export class GenerativeModelPreview {
   private readonly publisherModelEndpoint: string;
   private readonly resourcePath: string;
   private readonly apiEndpoint?: string;
+  private readonly cachedContent?: CachedContent;
 
   /**
    * @constructor
@@ -310,6 +312,7 @@ export class GenerativeModelPreview {
     this.safetySettings = getGenerativeModelParams.safetySettings;
     this.tools = getGenerativeModelParams.tools;
     this.toolConfig = getGenerativeModelParams.toolConfig;
+    this.cachedContent = getGenerativeModelParams.cachedContent;
     this.requestOptions = getGenerativeModelParams.requestOptions ?? {};
     if (getGenerativeModelParams.systemInstruction) {
       this.systemInstruction = formulateSystemInstructionIntoContent(
@@ -358,11 +361,13 @@ export class GenerativeModelPreview {
     request: GenerateContentRequest | string
   ): Promise<GenerateContentResult> {
     request = formulateRequestToGenerateContentRequest(request);
-    const formulatedRequest =
-      formulateSystemInstructionIntoGenerateContentRequest(
+    const formulatedRequest = {
+      ...formulateSystemInstructionIntoGenerateContentRequest(
         request,
         this.systemInstruction
-      );
+      ),
+      cachedContent: this.cachedContent?.name,
+    };
     return generateContent(
       this.location,
       this.resourcePath,
@@ -405,11 +410,13 @@ export class GenerativeModelPreview {
     request: GenerateContentRequest | string
   ): Promise<StreamGenerateContentResult> {
     request = formulateRequestToGenerateContentRequest(request);
-    const formulatedRequest =
-      formulateSystemInstructionIntoGenerateContentRequest(
+    const formulatedRequest = {
+      ...formulateSystemInstructionIntoGenerateContentRequest(
         request,
         this.systemInstruction
-      );
+      ),
+      cachedContent: this.cachedContent?.name,
+    };
     return generateContentStream(
       this.location,
       this.resourcePath,
@@ -486,6 +493,7 @@ export class GenerativeModelPreview {
       resourcePath: this.resourcePath,
       tools: this.tools,
       systemInstruction: this.systemInstruction,
+      cachedContent: this.cachedContent?.name,
     };
 
     if (request) {
@@ -497,9 +505,23 @@ export class GenerativeModelPreview {
       startChatRequest.tools = request.tools ?? this.tools;
       startChatRequest.systemInstruction =
         request.systemInstruction ?? this.systemInstruction;
+      startChatRequest.cachedContent =
+        request.cachedContent ?? this.cachedContent?.name;
     }
     return new ChatSessionPreview(startChatRequest, this.requestOptions);
   }
+
+  getModelName(): string {
+    return this.model;
+  }
+
+  getCachedContent(): CachedContent | undefined {
+    return this.cachedContent;
+  }
+
+  getSystemInstruction(): Content | undefined {
+    return this.systemInstruction;
+  }
 }
 
 function formulateResourcePathFromModel(