From 4edef30e0e5338ad359fd07efc6fa3c43f6dfdc0 Mon Sep 17 00:00:00 2001
From: NamH <NamNh0122@gmail.com>
Date: Fri, 31 May 2024 13:15:06 +0700
Subject: [PATCH] feat: allow user to register their access token (#2974)

Signed-off-by: James <james@jan.ai>
Co-authored-by: James <james@jan.ai>
---
 .../model-extension/resources/settings.json   | 12 ++++++++++
 extensions/model-extension/rollup.config.ts   |  2 ++
 extensions/model-extension/src/index.ts       | 22 ++++++++++++++++++-
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 extensions/model-extension/resources/settings.json

diff --git a/extensions/model-extension/resources/settings.json b/extensions/model-extension/resources/settings.json
new file mode 100644
index 000000000..5cdeeeba0
--- /dev/null
+++ b/extensions/model-extension/resources/settings.json
@@ -0,0 +1,12 @@
+[
+  {
+    "key": "hugging-face-access-token",
+    "title": "Hugging Face Access Token",
+    "description": "Access tokens programmatically authenticate your identity to the Hugging Face Hub, allowing applications to perform specific actions specified by the scope of permissions granted.",
+    "controllerType": "input",
+    "controllerProps": {
+      "value": "",
+      "placeholder": "hf_**********************************"
+    }
+  }
+]
diff --git a/extensions/model-extension/rollup.config.ts b/extensions/model-extension/rollup.config.ts
index abd12890e..aa22bd1f6 100644
--- a/extensions/model-extension/rollup.config.ts
+++ b/extensions/model-extension/rollup.config.ts
@@ -4,6 +4,7 @@ import typescript from 'rollup-plugin-typescript2'
 import json from '@rollup/plugin-json'
 import replace from '@rollup/plugin-replace'
 
+const settingJson = require('./resources/settings.json')
 const packageJson = require('./package.json')
 const defaultModelJson = require('./resources/default-model.json')
 
@@ -20,6 +21,7 @@ export default [
       replace({
         preventAssignment: true,
         DEFAULT_MODEL: JSON.stringify(defaultModelJson),
+        SETTINGS: JSON.stringify(settingJson),
         NODE: JSON.stringify(`${packageJson.name}/${packageJson.node}`),
       }),
       // Allow json resolution
diff --git a/extensions/model-extension/src/index.ts b/extensions/model-extension/src/index.ts
index dbe7605ea..887ce7474 100644
--- a/extensions/model-extension/src/index.ts
+++ b/extensions/model-extension/src/index.ts
@@ -31,6 +31,11 @@ import { GGUFMetadata, gguf } from '@huggingface/gguf'
 import { NotSupportedModelError } from './@types/NotSupportModelError'
 import { InvalidHostError } from './@types/InvalidHostError'
 
+declare const SETTINGS: Array<any>
+enum Settings {
+  huggingFaceAccessToken = 'hugging-face-access-token',
+}
+
 /**
  * A extension for models
  */
@@ -63,6 +68,7 @@ export default class JanModelExtension extends ModelExtension {
    */
   async onLoad() {
     // Handle Desktop Events
+    this.registerSettings(SETTINGS)
     this.handleDesktopEvents()
   }
 
@@ -195,7 +201,21 @@ export default class JanModelExtension extends ModelExtension {
     const sanitizedUrl = this.toHuggingFaceUrl(repoId)
     console.debug('sanitizedUrl', sanitizedUrl)
 
-    const res = await fetch(sanitizedUrl)
+    const huggingFaceAccessToken = (
+      await this.getSetting<string>(Settings.huggingFaceAccessToken, '')
+    ).trim()
+
+    const headers = {
+      Accept: 'application/json',
+    }
+
+    if (huggingFaceAccessToken.length > 0) {
+      headers['Authorization'] = `Bearer ${huggingFaceAccessToken}`
+    }
+
+    const res = await fetch(sanitizedUrl, {
+      headers: headers,
+    })
     const response = await res.json()
     if (response['error'] != null) {
       throw new Error(response['error'])