feat: Allow labels in GenerateContentRequest (#466) (#501)

* feat: Allow labels in GenerateContentRequest

* Add optional `labels` record to GenerateContentRequest

* include labels in sent `generateContentRequest`

* Add unit tests

* Test that `generateContent` provides labels to the Vertex endpoint
* Test that `generateContentStream` provides labels to the Vertex endpoint

* cleanup log and spacing

* fix style

---------

Co-authored-by: Yvonne Yu <[email protected]>

Author: julian-computes

src/functions/generate_content.ts

@@ -77,6 +77,7 @@ export async function generateContent(
     safetySettings: request.safetySettings ?? safetySettings,
     tools: request.tools ?? tools,
     toolConfig: request.toolConfig ?? toolConfig,
+    labels: request.labels,
   };
   const response: Response | undefined = await postRequest({
     region: location,
@@ -132,6 +133,7 @@ export async function generateContentStream(
     safetySettings: request.safetySettings ?? safetySettings,
     tools: request.tools ?? tools,
     toolConfig: request.toolConfig ?? toolConfig,
+    labels: request.labels,
   };
   const response = await postRequest({
     region: location,

src/functions/test/functions_test.ts

@@ -226,6 +226,10 @@ const TEST_TOOLS_WITH_RAG: Tool[] = [
   },
 ];
 
+const TEST_LABELS: Record<string, string> = {
+  test_key: 'test_value',
+};
+
 const fetchResponseObj = {
   status: 200,
   statusText: 'OK',
@@ -648,6 +652,26 @@ describe('generateContent', () => {
     const vertexEndpoint = fetchSpy.calls.allArgs()[0][0];
     expect(vertexEndpoint).toContain('/v1beta1/');
   });
+
+  it('provides labels to the vertex endpoint', async () => {
+    const request: GenerateContentRequest = {
+      contents: CONTENTS,
+      labels: TEST_LABELS,
+    };
+    fetchSpy.and.resolveTo(buildFetchResponse(TEST_MODEL_RESPONSE));
+
+    await generateContent(
+      TEST_LOCATION,
+      TEST_RESOURCE_PATH,
+      TEST_TOKEN_PROMISE,
+      request,
+      TEST_API_ENDPOINT
+    );
+
+    const httpRequest = fetchSpy.calls.allArgs()[0][1];
+    const body = JSON.parse(httpRequest.body);
+    expect(body.labels).toEqual(TEST_LABELS);
+  });
 });
 
 describe('generateContentStream', () => {
@@ -830,4 +854,29 @@ describe('generateContentStream', () => {
       )
     ).toHaveSize(0);
   });
+
+  it('provides labels to the vertex endpoint', async () => {
+    const request: GenerateContentRequest = {
+      contents: CONTENTS,
+      labels: TEST_LABELS,
+    };
+    const expectedStreamResult: StreamGenerateContentResult = {
+      response: Promise.resolve(TEST_MODEL_RESPONSE_WITH_INVALID_DATA),
+      stream: testGenerator(),
+    };
+    fetchSpy = spyOn(global, 'fetch').and.resolveTo(fetchResult);
+    spyOn(StreamFunctions, 'processStream').and.resolveTo(expectedStreamResult);
+
+    await generateContentStream(
+      TEST_LOCATION,
+      TEST_RESOURCE_PATH,
+      TEST_TOKEN_PROMISE,
+      request,
+      TEST_API_ENDPOINT
+    );
+
+    const httpRequest = fetchSpy.calls.allArgs()[0][1];
+    const body = JSON.parse(httpRequest.body);
+    expect(body.labels).toEqual(TEST_LABELS);
+  });
 });

src/types/content.ts

@@ -68,6 +68,12 @@ export declare interface GenerateContentRequest extends BaseModelParams {
    * This is the name of a `CachedContent` and not the cache object itself.
    */
   cachedContent?: string;
+
+  /**
+   * Optional. Custom metadata labels for organizing API calls and managing costs at scale. See
+   * https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls
+   */
+  labels?: Record<string, string>;
 }
 
 /**