feat(google-genai): Context Caching

libs/langchain-google-genai/src/chat_models.ts

@@ -7,6 +7,9 @@ import {
   GenerateContentRequest,
   SafetySetting,
   Part as GenerativeAIPart,
+  ModelParams,
+  RequestOptions,
+  CachedContent,
 } from "@google/generative-ai";
 import { CallbackManagerForLLMRun } from "@langchain/core/callbacks/manager";
 import {
@@ -651,6 +654,21 @@ export class ChatGoogleGenerativeAI
     this.streamUsage = fields?.streamUsage ?? this.streamUsage;
   }
 
+  enableCachedContent(
+    cachedContent: CachedContent,
+    modelParams?: ModelParams,
+    requestOptions?: RequestOptions
+  ): void {
+    if (!this.apiKey) return;
+    this.client = new GenerativeAI(
+      this.apiKey
+    ).getGenerativeModelFromCachedContent(
+      cachedContent,
+      modelParams,
+      requestOptions
+    );
+  }
+
   getLsParams(options: this["ParsedCallOptions"]): LangSmithParams {
     return {
       ls_provider: "google_genai",

libs/langchain-google-genai/src/context_caching.ts

@@ -0,0 +1,88 @@
+import {
+  CachedContentCreateParams,
+  CachedContentUpdateParams,
+  FileMetadata,
+  FileMetadataResponse,
+  GoogleAICacheManager,
+  ListCacheResponse,
+  ListFilesResponse,
+  ListParams,
+  UploadFileResponse,
+  GoogleAIFileManager,
+} from "@google/generative-ai/server";
+import {
+  CachedContent,
+  RequestOptions,
+  SingleRequestOptions,
+} from "@google/generative-ai";
+
+export class GoogleGenerativeAIContextCache {
+  private fileManager: GoogleAIFileManager;
+
+  private cacheManager: GoogleAICacheManager;
+
+  constructor(
+    apiKey: string,
+    fileManagerRequestOptions?: RequestOptions,
+    cacheManagerRequestOptions?: RequestOptions
+  ) {
+    this.fileManager = new GoogleAIFileManager(
+      apiKey,
+      fileManagerRequestOptions
+    );
+    this.cacheManager = new GoogleAICacheManager(
+      apiKey,
+      cacheManagerRequestOptions
+    );
+  }
+
+  uploadFile(
+    filePath: string,
+    fileMetadata: FileMetadata
+  ): Promise<UploadFileResponse> {
+    return this.fileManager.uploadFile(filePath, fileMetadata);
+  }
+
+  listFiles(
+    listParams?: ListParams,
+    requestOptions?: SingleRequestOptions
+  ): Promise<ListFilesResponse> {
+    return this.fileManager.listFiles(listParams, requestOptions);
+  }
+
+  getFile(
+    fileId: string,
+    requestOptions?: SingleRequestOptions
+  ): Promise<FileMetadataResponse> {
+    return this.fileManager.getFile(fileId, requestOptions);
+  }
+
+  deleteFile(fileId: string): Promise<void> {
+    return this.fileManager.deleteFile(fileId);
+  }
+
+  createCache(
+    createOptions: CachedContentCreateParams
+  ): Promise<CachedContent> {
+    return this.cacheManager.create(createOptions);
+  }
+
+  listCaches(listParams?: ListParams): Promise<ListCacheResponse> {
+    return this.cacheManager.list(listParams);
+  }
+
+  getCache(name: string): Promise<CachedContent> {
+    return this.cacheManager.get(name);
+  }
+
+  updateCache(
+    name: string,
+    updateParams: CachedContentUpdateParams
+  ): Promise<CachedContent> {
+    return this.cacheManager.update(name, updateParams);
+  }
+
+  deleteCache(name: string): Promise<void> {
+    return this.cacheManager.delete(name);
+  }
+}

libs/langchain-google-genai/src/tests/context_caching.int.test.ts

@@ -0,0 +1,79 @@
+/* eslint-disable no-process-env */
+
+import { test } from "@jest/globals";
+
+import { fileURLToPath } from "node:url";
+import * as path from "node:path";
+
+import { FileState, UploadFileResponse } from "@google/generative-ai/server";
+import { GoogleGenerativeAIContextCache } from "../context_caching.js";
+import { ChatGoogleGenerativeAI } from "../chat_models.js";
+
+const model = new ChatGoogleGenerativeAI({});
+let fileResult: UploadFileResponse;
+
+beforeAll(async () => {
+  // Download video file and save in src/tests/data
+  // curl -O https://storage.googleapis.com/generativeai-downloads/data/Sherlock_Jr_FullMovie.mp4
+  const displayName = "Sherlock Jr. video";
+
+  const filename = fileURLToPath(import.meta.url);
+  const dirname = path.dirname(filename);
+  const pathToVideoFile = path.join(dirname, "/data/Sherlock_Jr_FullMovie.mp4");
+
+  const contextCache = new GoogleGenerativeAIContextCache(
+    process.env.GOOGLE_API_KEY || ""
+  );
+  fileResult = await contextCache.uploadFile(pathToVideoFile, {
+    displayName,
+    mimeType: "video/mp4",
+  });
+
+  const { name } = fileResult.file;
+
+  // Poll getFile() on a set interval (2 seconds here) to check file state.
+  let file = await contextCache.getFile(name);
+  while (file.state === FileState.PROCESSING) {
+    // Sleep for 2 seconds
+    await new Promise((resolve) => {
+      setTimeout(resolve, 2_000);
+    });
+    file = await contextCache.getFile(name);
+  }
+
+  const systemInstruction =
+    "You are an expert video analyzer, and your job is to answer " +
+    "the user's query based on the video file you have access to.";
+  const cachedContent = await contextCache.createCache({
+    model: "models/gemini-1.5-flash-001",
+    displayName: "sherlock jr movie",
+    systemInstruction,
+    contents: [
+      {
+        role: "user",
+        parts: [
+          {
+            fileData: {
+              mimeType: fileResult.file.mimeType,
+              fileUri: fileResult.file.uri,
+            },
+          },
+        ],
+      },
+    ],
+    ttlSeconds: 300,
+  });
+
+  model.enableCachedContent(cachedContent);
+}, 10 * 60 * 1000); // Set timeout to 10 minutes to upload file
+
+test("Test Google AI", async () => {
+  const res = await model.invoke(
+    "Introduce different characters in the movie by describing " +
+      "their personality, looks, and names. Also list the " +
+      "timestamps they were introduced for the first time."
+  );
+
+  console.log(res);
+  expect(res).toBeTruthy();
+});