Add a request option for custom headers (#118)

* Add a request option for custom headers

* Allow non-Header object as customHeaders and try to convert to Headers

* Reject customHeaders for reserved header names with request input error

* Use distinct error messages for customHeaders with api-key and api-client

* Add changeset

Author: dlarocque

.changeset/strange-papayas-accept.md

@@ -0,0 +1,5 @@
+---
+"@google/generative-ai": minor
+---
+
+Add a request option for custom headers

packages/main/src/errors.ts

@@ -56,6 +56,10 @@ export class GoogleGenerativeAIFetchError extends GoogleGenerativeAIError {
   }
 }
 
+/**
+ */
+export class GoogleGenerativeAIRequestInputError extends GoogleGenerativeAIError {}
+
 /**
  * Details object that may be included in an error response.
  * @public

packages/main/src/requests/request.test.ts

@@ -27,7 +27,10 @@ import {
   _makeRequestInternal,
   constructRequest,
 } from "./request";
-import { GoogleGenerativeAIFetchError } from "../errors";
+import {
+  GoogleGenerativeAIFetchError,
+  GoogleGenerativeAIRequestInputError,
+} from "../errors";
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -155,6 +158,40 @@ describe("request methods", () => {
       );
       expect(request.fetchOptions.signal).to.be.instanceOf(AbortSignal);
     });
+    it("passes custom headers", async () => {
+      const request = await constructRequest(
+        "model-name",
+        Task.GENERATE_CONTENT,
+        "key",
+        true,
+        "",
+        {
+          customHeaders: new Headers({ customerHeader: "customerHeaderValue" }),
+        },
+      );
+      expect(
+        (request.fetchOptions.headers as Headers).get("customerHeader"),
+      ).to.equal("customerHeaderValue");
+    });
+    it("passes custom x-goog-api-client header", async () => {
+      await expect(
+        constructRequest("model-name", Task.GENERATE_CONTENT, "key", true, "", {
+          customHeaders: new Headers({
+            "x-goog-api-client": "client/version",
+          }),
+        }),
+      ).to.be.rejectedWith(GoogleGenerativeAIRequestInputError);
+    });
+    it("passes apiClient and custom x-goog-api-client header", async () => {
+      await expect(
+        constructRequest("model-name", Task.GENERATE_CONTENT, "key", true, "", {
+          apiClient: "client/version",
+          customHeaders: new Headers({
+            "x-goog-api-client": "client/version2",
+          }),
+        }),
+      ).to.be.rejectedWith(GoogleGenerativeAIRequestInputError);
+    });
   });
   describe("_makeRequestInternal", () => {
     it("no error", async () => {
@@ -309,5 +346,23 @@ describe("request methods", () => {
       }
       expect(fetchStub).to.be.calledOnce;
     });
+    it("has invalid custom header", async () => {
+      const fetchStub = stub();
+      await expect(
+        _makeRequestInternal(
+          "model-name",
+          Task.GENERATE_CONTENT,
+          "key",
+          true,
+          "",
+          {
+            customHeaders: new Headers({
+              "x-goog-api-client": "client/version",
+            }),
+          },
+          fetchStub as typeof fetch,
+        ),
+      ).to.be.rejectedWith(GoogleGenerativeAIRequestInputError);
+    });
   });
 });

packages/main/src/requests/request.ts

@@ -19,6 +19,7 @@ import { RequestOptions } from "../../types";
 import {
   GoogleGenerativeAIError,
   GoogleGenerativeAIFetchError,
+  GoogleGenerativeAIRequestInputError,
 } from "../errors";
 
 export const DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com";
@@ -76,6 +77,36 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
   headers.append("Content-Type", "application/json");
   headers.append("x-goog-api-client", getClientHeaders(url.requestOptions));
   headers.append("x-goog-api-key", url.apiKey);
+
+  let customHeaders = url.requestOptions.customHeaders;
+  if (customHeaders) {
+    if (!(customHeaders instanceof Headers)) {
+      try {
+        customHeaders = new Headers(customHeaders);
+      } catch (e) {
+        throw new GoogleGenerativeAIRequestInputError(
+          `unable to convert customHeaders value ${JSON.stringify(
+            customHeaders,
+          )} to Headers: ${e.message}`,
+        );
+      }
+    }
+
+    for (const [headerName, headerValue] of customHeaders.entries()) {
+      if (headerName === "x-goog-api-key") {
+        throw new GoogleGenerativeAIRequestInputError(
+          `Cannot set reserved header name ${headerName}`,
+        );
+      } else if (headerName === "x-goog-api-client") {
+        throw new GoogleGenerativeAIRequestInputError(
+          `Header name ${headerName} can only be set using the apiClient field`,
+        );
+      }
+
+      headers.append(headerName, headerValue);
+    }
+  }
+
   return headers;
 }
 
@@ -168,7 +199,12 @@ export async function _makeRequestInternal(
     }
   } catch (e) {
     let err = e;
-    if (!(e instanceof GoogleGenerativeAIFetchError)) {
+    if (
+      !(
+        e instanceof GoogleGenerativeAIFetchError ||
+        e instanceof GoogleGenerativeAIRequestInputError
+      )
+    ) {
       err = new GoogleGenerativeAIError(
         `Error fetching from ${url.toString()}: ${e.message}`,
       );

packages/main/types/requests.ts

@@ -136,6 +136,10 @@ export interface RequestOptions {
    * Base endpoint url. Defaults to "https://generativelanguage.googleapis.com"
    */
   baseUrl?: string;
+  /**
+   * Custom HTTP request headers.
+   */
+  customHeaders?: Headers | Record<string, string>;
 }
 
 /**