feat: Support Custom Metadata Labels on Vertex AI

Signed-off-by: Yu Ishikawa <[email protected]>

Author: yu-iskw

libs/langchain-google-common/src/types.ts

@@ -308,6 +308,11 @@ export interface GoogleAIModelRequestParams extends GoogleAIModelParams {
    * https://cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-use
    */
   cachedContent?: string;
+
+  /**
+   * Custom metadata labels to associate with the API call.
+   */
+  labels?: Record<string, string>;
 }
 
 export interface GoogleAIBaseLLMInput<AuthOptions>
@@ -585,6 +590,11 @@ export interface GeminiRequest {
   safetySettings?: GeminiSafetySetting[];
   generationConfig?: GeminiGenerationConfig;
   cachedContent?: string;
+
+  /**
+   * Custom metadata labels to associate with the API call.
+   */
+  labels?: Record<string, string>;
 }
 
 export interface GeminiResponseCandidate {

libs/langchain-google-common/src/utils/gemini.ts

@@ -1594,6 +1594,7 @@ export function getGeminiAPI(config?: GeminiAPIConfig): GoogleAIAPI {
     const ret: GeminiRequest = {
       contents,
       generationConfig,
+      labels: parameters.labels,
     };
     if (tools && tools.length) {
       ret.tools = tools;

libs/langchain-google-vertexai/src/tests/labels.test.ts

@@ -0,0 +1,29 @@
+import { test, expect } from "@jest/globals";
+import { ChatVertexAI } from "../chat_models.js";
+import { MockClient } from "@langchain/google-common/src/tests/mock.ts";
+
+test("ChatVertexAI should pass labels to the API call", async () => {
+  const mockClient = new MockClient();
+  const chatModel = new ChatVertexAI({
+    model: "gemini-pro",
+    // We need to force the client to be the mock client for testing purposes.
+    // This is a bit of a hack, ideally there would be a cleaner way to inject mocks.
+    // For now, we'll cast the chatModel to access the protected connection property.
+    // In a real scenario, testing lower-level components or refactoring for dependency injection would be better.
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    client: mockClient as any,
+  });
+
+  const labels = {
+    "user-id": "123",
+    "session-id": "abc",
+  };
+
+  await chatModel.invoke("hello", {
+    labels,
+  });
+
+  // Assert that the mock client's request record contains the labels
+  // The request format is based on the Gemini API request structure
+  expect(mockClient.record.opts.data.labels).toEqual(labels);
+});