Add a request option for custom headers

google-gemini · May 1, 2024 · 87cb3f5 · 87cb3f5
1 parent d785a79
commit 87cb3f5
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 0 deletions.
diff --git a/packages/main/src/requests/request.test.ts b/packages/main/src/requests/request.test.ts
@@ -155,6 +155,52 @@ 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 () => {
+      const request = await constructRequest(
+        "model-name",
+        Task.GENERATE_CONTENT,
+        "key",
+        true,
+        "",
+        {
+          customHeaders: new Headers({ "x-goog-api-client": "client/version" }),
+        },
+      );
+      expect(
+        (request.fetchOptions.headers as Headers).get("x-goog-api-client"),
+      ).to.equal("client/version");
+    });
+    it("passes apiClient and custom x-goog-api-client header", async () => {
+      const request = await constructRequest(
+        "model-name",
+        Task.GENERATE_CONTENT,
+        "key",
+        true,
+        "",
+        {
+          apiClient: "client/version",
+          customHeaders: new Headers({ "x-goog-api-client": "client/version2" }),
+        },
+      );
+      expect(
+        (request.fetchOptions.headers as Headers).get("x-goog-api-client"),
+      ).to.equal("client/version genai-js/__PACKAGE_VERSION__, client/version2");
+    });
   });
   describe("_makeRequestInternal", () => {
     it("no error", async () => {

diff --git a/packages/main/src/requests/request.ts b/packages/main/src/requests/request.ts
@@ -63,6 +63,7 @@ export class RequestUrl {
  * Simple, but may become more complex if we add more versions to log.
  */
 export function getClientHeaders(requestOptions: RequestOptions): string {
+  // FIXME: Should we do this for an apiClient entry in customHeaders?
   const clientHeaders = [];
   if (requestOptions?.apiClient) {
     clientHeaders.push(requestOptions.apiClient);
@@ -76,6 +77,16 @@ 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);
+
+  if (url.requestOptions.customHeaders) {
+    for (const [
+      headerName,
+      headerValue,
+    ] of url.requestOptions.customHeaders.entries()) {
+      headers.append(headerName, headerValue);
+    }
+  }
+
   return headers;
 }
 

diff --git a/packages/main/types/requests.ts b/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;
 }
 
 /**