Custom Headers (#158)

* Allow users to pass custom headers Co-authored-by: Naveen <[email protected]> Co-authored-by: Bruce MacDonald <[email protected]>
ollama · Nov 7, 2024 · 6e7e496 · 6e7e496
1 parent c97f231
commit 6e7e496
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 18 deletions.
diff --git a/src/browser.ts b/src/browser.ts
@@ -1,5 +1,5 @@
 import * as utils from './utils.js'
-import { AbortableAsyncIterator, parseJSON, post } from './utils.js'
+import { AbortableAsyncIterator, parseJSON } from './utils.js'
 import 'whatwg-fetch'
 
 import type {
@@ -34,15 +34,14 @@ export class Ollama {
   constructor(config?: Partial<Config>) {
     this.config = {
       host: '',
+      headers: config?.headers
     }
+
     if (!config?.proxy) {
       this.config.host = utils.formatHost(config?.host ?? 'http://127.0.0.1:11434')
     }
 
-    this.fetch = fetch
-    if (config?.fetch != null) {
-      this.fetch = config.fetch
-    }
+    this.fetch = config?.fetch ?? fetch
   }
 
   // Abort any ongoing streamed requests to Ollama
@@ -72,7 +71,7 @@ export class Ollama {
     const host = `${this.config.host}/api/${endpoint}`
     if (request.stream) {
       const abortController = new AbortController()
-      const response = await post(this.fetch, host, request, {
+      const response = await utils.post(this.fetch, host, request, {
         signal: abortController.signal,
         headers: this.config.headers
       })
@@ -236,9 +235,12 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
    * @returns {Promise<StatusResponse>} - The response object.
    */
   async delete(request: DeleteRequest): Promise<StatusResponse> {
-    await utils.del(this.fetch, `${this.config.host}/api/delete`, {
-      name: request.model,
-    })
+    await utils.del(
+      this.fetch,
+      `${this.config.host}/api/delete`,
+      { name: request.model },
+      { headers: this.config.headers }
+    )
     return { status: 'success' }
   }
 
@@ -249,7 +251,9 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
    * @returns {Promise<StatusResponse>} - The response object.
    */
   async copy(request: CopyRequest): Promise<StatusResponse> {
-    await utils.post(this.fetch, `${this.config.host}/api/copy`, { ...request })
+    await utils.post(this.fetch, `${this.config.host}/api/copy`, { ...request }, {
+      headers: this.config.headers
+    })
     return { status: 'success' }
   }
 
@@ -259,7 +263,9 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
    * @throws {Error} - If the response body is missing.
    */
   async list(): Promise<ListResponse> {
-    const response = await utils.get(this.fetch, `${this.config.host}/api/tags`)
+    const response = await utils.get(this.fetch, `${this.config.host}/api/tags`, {
+      headers: this.config.headers
+    })
     return (await response.json()) as ListResponse
   }
 
@@ -271,6 +277,8 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
   async show(request: ShowRequest): Promise<ShowResponse> {
     const response = await utils.post(this.fetch, `${this.config.host}/api/show`, {
       ...request,
+    }, {
+      headers: this.config.headers
     })
     return (await response.json()) as ShowResponse
   }
@@ -283,6 +291,8 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
     async embed(request: EmbedRequest): Promise<EmbedResponse> {
       const response = await utils.post(this.fetch, `${this.config.host}/api/embed`, {
         ...request,
+      }, {
+        headers: this.config.headers
       })
       return (await response.json()) as EmbedResponse
     }
@@ -295,6 +305,8 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
   async embeddings(request: EmbeddingsRequest): Promise<EmbeddingsResponse> {
     const response = await utils.post(this.fetch, `${this.config.host}/api/embeddings`, {
       ...request,
+    }, {
+      headers: this.config.headers
     })
     return (await response.json()) as EmbeddingsResponse
   }
@@ -305,7 +317,9 @@ async encodeImage(image: Uint8Array | string): Promise<string> {
    * @throws {Error} - If the response body is missing.
    */
   async ps(): Promise<ListResponse> {
-    const response = await utils.get(this.fetch, `${this.config.host}/api/ps`)
+    const response = await utils.get(this.fetch, `${this.config.host}/api/ps`, {
+      headers: this.config.headers
+    })
     return (await response.json()) as ListResponse
   }
 }

diff --git a/src/interfaces.ts b/src/interfaces.ts
@@ -4,7 +4,7 @@ export interface Config {
   host: string
   fetch?: Fetch
   proxy?: boolean
-  headers?: Headers
+  headers?: HeadersInit
 }
 
 // request types

diff --git a/src/utils.ts b/src/utils.ts
@@ -115,15 +115,20 @@ const fetchWithHeaders = async (
     'Content-Type': 'application/json',
     Accept: 'application/json',
     'User-Agent': `ollama-js/${version} (${getPlatform()})`,
-  }
+  } as HeadersInit
 
   if (!options.headers) {
     options.headers = {}
   }
 
+  // Filter out default headers from custom headers
+  const customHeaders = Object.fromEntries(
+    Object.entries(options.headers).filter(([key]) => !Object.keys(defaultHeaders).some(defaultKey => defaultKey.toLowerCase() === key.toLowerCase()))
+  )
+
   options.headers = {
     ...defaultHeaders,
-    ...options.headers,
+    ...customHeaders
   }
 
   return fetch(url, options)
@@ -135,8 +140,10 @@ const fetchWithHeaders = async (
  * @param host {string} - The host to fetch
  * @returns {Promise<Response>} - The fetch response
  */
-export const get = async (fetch: Fetch, host: string): Promise<Response> => {
-  const response = await fetchWithHeaders(fetch, host)
+export const get = async (fetch: Fetch, host: string, options?: { headers?: HeadersInit }): Promise<Response> => {
+  const response = await fetchWithHeaders(fetch, host, {
+    headers: options?.headers
+  })
 
   await checkOk(response)
 
@@ -169,7 +176,7 @@ export const post = async (
   fetch: Fetch,
   host: string,
   data?: Record<string, unknown> | BodyInit,
-  options?: { signal?: AbortSignal, headers?: Headers },
+  options?: { signal?: AbortSignal, headers?: HeadersInit },
 ): Promise<Response> => {
   const isRecord = (input: any): input is Record<string, unknown> => {
     return input !== null && typeof input === 'object' && !Array.isArray(input)
@@ -199,10 +206,12 @@ export const del = async (
   fetch: Fetch,
   host: string,
   data?: Record<string, unknown>,
+  options?: { headers?: HeadersInit },
 ): Promise<Response> => {
   const response = await fetchWithHeaders(fetch, host, {
     method: 'DELETE',
     body: JSON.stringify(data),
+    headers: options?.headers
   })
 
   await checkOk(response)