feat(google-genai): Context Caching (#7169)

Co-authored-by: Chau Nguyen <[email protected]>
Co-authored-by: jacoblee93 <[email protected]>

docs/core_docs/docs/integrations/chat/google_generativeai.ipynb

@@ -796,6 +796,80 @@
         "console.log(codeExecutionExplanation.content);"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "id": "a464c1a9",
+      "metadata": {},
+      "source": [
+        "### Context Caching\n",
+        "\n",
+        "Context caching allows you to pass some content to the model once, cache the input tokens, and then refer to the cached tokens for subsequent requests to reduce cost. You can create a `CachedContent` object using `GoogleAICacheManager` class and then pass the `CachedContent` object to your `ChatGoogleGenerativeAIModel` with `enableCachedContent()` method."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "9a649be0",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import { ChatGoogleGenerativeAI } from \"@langchain/google-genai\";\n",
+        "import {\n",
+        "  GoogleAICacheManager,\n",
+        "  GoogleAIFileManager,\n",
+        "} from \"@google/generative-ai/server\";\n",
+        "\n",
+        "const fileManager = new GoogleAIFileManager(process.env.GOOGLE_API_KEY);\n",
+        "const cacheManager = new GoogleAICacheManager(process.env.GOOGLE_API_KEY);\n",
+        "\n",
+        "// uploads file for caching\n",
+        "const pathToVideoFile = \"/path/to/video/file\";\n",
+        "const displayName = \"example-video\";\n",
+        "const fileResult = await fileManager.uploadFile(pathToVideoFile, {\n",
+        "    displayName,\n",
+        "    mimeType: \"video/mp4\",\n",
+        "});\n",
+        "\n",
+        "// creates cached content AFTER uploading is finished\n",
+        "const cachedContent = await cacheManager.create({\n",
+        "    model: \"models/gemini-1.5-flash-001\",\n",
+        "    displayName: displayName,\n",
+        "    systemInstruction: \"You are an expert video analyzer, and your job is to answer \" +\n",
+        "      \"the user's query based on the video file you have access to.\",\n",
+        "    contents: [\n",
+        "        {\n",
+        "            role: \"user\",\n",
+        "            parts: [\n",
+        "                {\n",
+        "                    fileData: {\n",
+        "                        mimeType: fileResult.file.mimeType,\n",
+        "                        fileUri: fileResult.file.uri,\n",
+        "                    },\n",
+        "                },\n",
+        "            ],\n",
+        "        },\n",
+        "    ],\n",
+        "    ttlSeconds: 300,\n",
+        "});\n",
+        "\n",
+        "// passes cached video to model\n",
+        "const model = new ChatGoogleGenerativeAI({});\n",
+        "model.useCachedContent(cachedContent);\n",
+        "\n",
+        "// invokes model with cached video\n",
+        "await model.invoke(\"Summarize the video\");"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "12e978ff",
+      "metadata": {},
+      "source": [
+        "**Note**\n",
+        "- Context caching supports both Gemini 1.5 Pro and Gemini 1.5 Flash. Context caching is only available for stable models with fixed versions (for example, gemini-1.5-pro-001). You must include the version postfix (for example, the -001 in gemini-1.5-pro-001).\n",
+        "- The minimum input token count for context caching is 32,768, and the maximum is the same as the maximum for the given model."
+      ]
+    },
     {
       "cell_type": "markdown",
       "id": "0c6a950f",

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

@@ -7,6 +7,9 @@ import {
   GenerateContentRequest,
   SafetySetting,
   Part as GenerativeAIPart,
+  ModelParams,
+  RequestOptions,
+  type CachedContent,
 } from "@google/generative-ai";
 import { CallbackManagerForLLMRun } from "@langchain/core/callbacks/manager";
 import {
@@ -662,6 +665,21 @@ export class ChatGoogleGenerativeAI
     this.streamUsage = fields?.streamUsage ?? this.streamUsage;
   }
 
+  useCachedContent(
+    cachedContent: CachedContent,
+    modelParams?: ModelParams,
+    requestOptions?: RequestOptions
+  ): void {
+    if (!this.apiKey) return;
+    this.client = new GenerativeAI(
+      this.apiKey
+    ).getGenerativeModelFromCachedContent(
+      cachedContent,
+      modelParams,
+      requestOptions
+    );
+  }
+
   get useSystemInstruction(): boolean {
     return typeof this.convertSystemMessageToHumanContent === "boolean"
       ? !this.convertSystemMessageToHumanContent

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

@@ -0,0 +1,84 @@
+/* eslint-disable no-process-env */
+
+import { test } from "@jest/globals";
+
+import { fileURLToPath } from "node:url";
+import * as path from "node:path";
+
+import {
+  FileState,
+  UploadFileResponse,
+  GoogleAIFileManager,
+  GoogleAICacheManager,
+} from "@google/generative-ai/server";
+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 GoogleAICacheManager(
+    process.env.GOOGLE_API_KEY || ""
+  );
+  const fileCache = new GoogleAIFileManager(process.env.GOOGLE_API_KEY || "");
+  fileResult = await fileCache.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 fileCache.getFile(name);
+  while (file.state === FileState.PROCESSING) {
+    // Sleep for 2 seconds
+    await new Promise((resolve) => {
+      setTimeout(resolve, 2_000);
+    });
+    file = await fileCache.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.create({
+    model: "models/gemini-1.5-flash-001",
+    displayName: "gettysburg audio",
+    systemInstruction,
+    contents: [
+      {
+        role: "user",
+        parts: [
+          {
+            fileData: {
+              mimeType: fileResult.file.mimeType,
+              fileUri: fileResult.file.uri,
+            },
+          },
+        ],
+      },
+    ],
+    ttlSeconds: 300,
+  });
+
+  model.useCachedContent(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();
+});