From aed1bbbcaa2edb224dbf9506add0b64ec5c75db8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E8=85=BE?=
 <101850389+hangters@users.noreply.github.com>
Date: Thu, 8 Aug 2024 12:09:50 +0800
Subject: [PATCH] add supprot for lepton (#1866)

### What problem does this PR solve?

add supprot for lepton
#1853

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
---
 conf/llm_factories.json                       | 98 +++++++++++++++++++
 rag/llm/__init__.py                           |  3 +-
 rag/llm/chat_model.py                         |  9 +-
 web/src/assets/svg/llm/lepton.svg             | 44 +++++++++
 .../user-setting/setting-model/constant.ts    |  3 +-
 5 files changed, 154 insertions(+), 3 deletions(-)
 create mode 100644 web/src/assets/svg/llm/lepton.svg

diff --git a/conf/llm_factories.json b/conf/llm_factories.json
index 57b26ffa7..a4b99eb86 100644
--- a/conf/llm_factories.json
+++ b/conf/llm_factories.json
@@ -2326,6 +2326,104 @@
                     "model_type": "rerank"
                 }
             ]
+        },
+        {
+            "name": "Lepton",
+            "logo": "",
+            "tags": "LLM",
+            "status": "1",
+            "llm": [
+                {
+                    "llm_name": "dolphin-mixtral-8x7b",
+                    "tags": "LLM,CHAT,32k",
+                    "max_tokens": 32768,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "gemma-7b",
+                    "tags": "LLM,CHAT,8k",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "llama3-1-8b",
+                    "tags": "LLM,CHAT,4k",
+                    "max_tokens": 4096,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "llama3-8b",
+                    "tags": "LLM,CHAT,8K",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "llama2-13b",
+                    "tags": "LLM,CHAT,4K",
+                    "max_tokens": 4096,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "llama3-1-70b",
+                    "tags": "LLM,CHAT,8k",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "llama3-70b",
+                    "tags": "LLM,CHAT,8k",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "llama3-1-405b",
+                    "tags": "LLM,CHAT,8k",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "mistral-7b",
+                    "tags": "LLM,CHAT,8K",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "mistral-8x7b",
+                    "tags": "LLM,CHAT,8K",
+                    "max_tokens": 8192,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "nous-hermes-llama2",
+                    "tags": "LLM,CHAT,4k",
+                    "max_tokens": 4096,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "openchat-3-5",
+                    "tags": "LLM,CHAT,4k",
+                    "max_tokens": 4096,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "toppy-m-7b",
+                    "tags": "LLM,CHAT,4k",
+                    "max_tokens": 4096,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "wizardlm-2-7b",
+                    "tags": "LLM,CHAT,32k",
+                    "max_tokens": 32768,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "wizardlm-2-8x22b",
+                    "tags": "LLM,CHAT,64K",
+                    "max_tokens": 65536,
+                    "model_type": "chat"
+                }
+            ]
         }
     ]
 }
\ No newline at end of file
diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
index f652b6336..eebe8bfa9 100644
--- a/rag/llm/__init__.py
+++ b/rag/llm/__init__.py
@@ -83,7 +83,8 @@ ChatModel = {
     "NVIDIA": NvidiaChat,
     "LM-Studio": LmStudioChat,
     "OpenAI-API-Compatible": OpenAI_APIChat,
-    "cohere": CoHereChat
+    "cohere": CoHereChat,
+    "LeptonAI": LeptonAIChat
 }
 
 
diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
index ee3acf1d6..f59c4785f 100644
--- a/rag/llm/chat_model.py
+++ b/rag/llm/chat_model.py
@@ -71,7 +71,7 @@ class Base(ABC):
                         total_tokens
                         + num_tokens_from_string(resp.choices[0].delta.content)
                     )
-                    if not hasattr(resp, "usage")
+                    if not hasattr(resp, "usage") or not resp.usage
                     else resp.usage["total_tokens"]
                 )
                 if resp.choices[0].finish_reason == "length":
@@ -981,3 +981,10 @@ class CoHereChat(Base):
             yield ans + "\n**ERROR**: " + str(e)
 
         yield total_tokens
+
+
+class LeptonAIChat(Base):
+    def __init__(self, key, model_name, base_url=None):
+        if not base_url:
+            base_url = os.path.join("https://"+model_name+".lepton.run","api","v1")
+        super().__init__(key, model_name, base_url)
\ No newline at end of file
diff --git a/web/src/assets/svg/llm/lepton.svg b/web/src/assets/svg/llm/lepton.svg
new file mode 100644
index 000000000..b7ccd3d26
--- /dev/null
+++ b/web/src/assets/svg/llm/lepton.svg
@@ -0,0 +1,44 @@
+<?xml version="1.0" encoding="utf-8"?>
+<svg version="1.1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 85 85" style="enable-background:new 0 0 85 85;" xml:space="preserve">
+<style type="text/css">
+	.st0{fill-rule:evenodd;clip-rule:evenodd;fill:#2D9CDB;}
+	.st1{fill-rule:evenodd;clip-rule:evenodd;fill:#2F80ED;}
+	.st2{fill:#2F80ED;}
+	.st3{fill:#2D9CDB;}
+</style>
+<path class="st0" d="M75.9,48.1V36.9c0-2,0-3.1-0.1-3.9c0-0.4-0.1-0.6-0.1-0.7c-0.1-0.3-0.2-0.5-0.4-0.7c-0.1,0-0.2-0.2-0.6-0.4
+	c-0.7-0.5-1.6-1-3.3-2l-9.7-5.6c-1.7-1-2.7-1.5-3.4-1.9c-0.4-0.2-0.6-0.3-0.6-0.3c-0.3-0.1-0.6-0.1-0.9,0c-0.1,0-0.3,0.1-0.6,0.3
+	c-0.7,0.4-1.7,0.9-3.4,1.9l-9.7,5.6c-1.7,1-2.7,1.5-3.3,2c-0.3,0.2-0.5,0.4-0.6,0.4c-0.2,0.2-0.3,0.5-0.4,0.7c0,0.1,0,0.3-0.1,0.7
+	c0,0.8-0.1,1.9-0.1,3.9v11.2c0,2,0,3.1,0.1,3.9c0,0.4,0.1,0.6,0.1,0.7c0.1,0.3,0.2,0.5,0.4,0.7c0.1,0,0.2,0.2,0.6,0.4
+	c0.7,0.5,1.6,1,3.3,2l9.7,5.6c1.7,1,2.7,1.5,3.4,1.9c0.4,0.2,0.6,0.3,0.6,0.3c0.3,0.1,0.6,0.1,0.9,0c0.1,0,0.3-0.1,0.6-0.3
+	c0.7-0.4,1.7-0.9,3.4-1.9l9.7-5.6c1.7-1,2.7-1.5,3.3-2c0.3-0.2,0.5-0.4,0.6-0.4c0.2-0.2,0.3-0.5,0.4-0.7c0-0.1,0-0.3,0.1-0.7
+	C75.9,51.2,75.9,50.1,75.9,48.1z M75.7,52.7C75.7,52.7,75.7,52.7,75.7,52.7C75.7,52.7,75.7,52.7,75.7,52.7z M75.3,53.4
+	C75.3,53.4,75.3,53.4,75.3,53.4C75.3,53.4,75.3,53.4,75.3,53.4z M57.7,63.7C57.7,63.7,57.7,63.7,57.7,63.7
+	C57.7,63.7,57.7,63.7,57.7,63.7z M56.9,63.7C56.9,63.7,56.9,63.7,56.9,63.7C56.9,63.7,56.9,63.7,56.9,63.7z M39.3,53.4
+	C39.3,53.4,39.3,53.4,39.3,53.4C39.3,53.4,39.3,53.4,39.3,53.4z M38.9,52.7C38.9,52.7,38.9,52.7,38.9,52.7
+	C38.9,52.7,38.9,52.7,38.9,52.7z M38.9,32.3C38.9,32.3,38.9,32.3,38.9,32.3C38.9,32.3,38.9,32.3,38.9,32.3z M39.3,31.6
+	C39.3,31.6,39.3,31.6,39.3,31.6C39.3,31.6,39.3,31.6,39.3,31.6z M56.9,21.4C56.9,21.4,56.9,21.4,56.9,21.4
+	C56.9,21.4,56.9,21.4,56.9,21.4z M57.7,21.4C57.7,21.4,57.7,21.4,57.7,21.4C57.7,21.4,57.7,21.4,57.7,21.4z M75.3,31.6
+	C75.3,31.6,75.3,31.6,75.3,31.6C75.3,31.6,75.3,31.6,75.3,31.6z M75.7,32.3C75.7,32.3,75.7,32.3,75.7,32.3
+	C75.7,32.3,75.7,32.3,75.7,32.3z M81.9,25.6c-1.2-1.3-2.8-2.3-6-4.1l-9.7-5.6C63,14,61.3,13,59.6,12.7c-1.5-0.3-3.1-0.3-4.6,0
+	c-1.7,0.4-3.3,1.3-6.6,3.2l-9.7,5.6c-3.2,1.9-4.9,2.8-6,4.1c-1,1.2-1.8,2.5-2.3,4c-0.5,1.7-0.5,3.6-0.5,7.3v11.2
+	c0,3.8,0,5.6,0.5,7.3c0.5,1.5,1.3,2.9,2.3,4c1.2,1.3,2.8,2.3,6,4.1l9.7,5.6c3.2,1.9,4.9,2.8,6.6,3.2c1.5,0.3,3.1,0.3,4.6,0
+	c1.7-0.4,3.3-1.3,6.6-3.2l9.7-5.6c3.2-1.9,4.9-2.8,6-4.1c1-1.2,1.8-2.5,2.3-4c0.5-1.7,0.5-3.6,0.5-7.3V36.9c0-3.8,0-5.6-0.5-7.3
+	C83.7,28.1,82.9,26.7,81.9,25.6z"/>
+<path class="st1" d="M46.3,48.1V36.9c0-2,0-3.1-0.1-3.9c0-0.4-0.1-0.6-0.1-0.7c-0.1-0.3-0.2-0.5-0.4-0.7c-0.1,0-0.2-0.2-0.6-0.4
+	c-0.7-0.5-1.6-1-3.3-2l-9.7-5.6c-1.7-1-2.7-1.5-3.4-1.9c-0.4-0.2-0.6-0.3-0.6-0.3c-0.3-0.1-0.6-0.1-0.9,0c-0.1,0-0.3,0.1-0.6,0.3
+	c-0.7,0.4-1.7,0.9-3.4,1.9l-9.7,5.6c-1.7,1-2.7,1.5-3.3,2c-0.3,0.2-0.5,0.4-0.6,0.4c-0.2,0.2-0.3,0.5-0.4,0.7c0,0.1,0,0.3-0.1,0.7
+	c0,0.8-0.1,1.9-0.1,3.9v11.2c0,2,0,3.1,0.1,3.9c0,0.4,0.1,0.6,0.1,0.7c0.1,0.3,0.2,0.5,0.4,0.7c0.1,0,0.2,0.2,0.6,0.4
+	c0.7,0.5,1.6,1,3.3,2l9.7,5.6c1.7,1,2.7,1.5,3.4,1.9c0.4,0.2,0.6,0.3,0.6,0.3c0.3,0.1,0.6,0.1,0.9,0c0.1,0,0.3-0.1,0.6-0.3
+	c0.7-0.4,1.7-0.9,3.4-1.9l9.7-5.6c1.7-1,2.7-1.5,3.3-2c0.3-0.2,0.5-0.4,0.6-0.4c0.2-0.2,0.3-0.5,0.4-0.7c0-0.1,0-0.3,0.1-0.7
+	C46.3,51.2,46.3,50.1,46.3,48.1z M52.3,25.6c-1.2-1.3-2.8-2.3-6-4.1l-9.7-5.6C33.4,14,31.8,13,30,12.7c-1.5-0.3-3.1-0.3-4.6,0
+	c-1.7,0.4-3.3,1.3-6.6,3.2l-9.7,5.6c-3.2,1.9-4.9,2.8-6,4.1c-1,1.2-1.8,2.5-2.3,4c-0.5,1.7-0.5,3.6-0.5,7.3v11.2
+	c0,3.8,0,5.6,0.5,7.3c0.5,1.5,1.3,2.9,2.3,4c1.2,1.3,2.8,2.3,6,4.1l9.7,5.6c3.2,1.9,4.9,2.8,6.6,3.2c1.5,0.3,3.1,0.3,4.6,0
+	c1.7-0.4,3.3-1.3,6.6-3.2l9.7-5.6c3.2-1.9,4.9-2.8,6-4.1c1-1.2,1.8-2.5,2.3-4c0.5-1.7,0.5-3.6,0.5-7.3V36.9c0-3.8,0-5.6-0.5-7.3
+	C54.2,28.1,53.4,26.7,52.3,25.6z"/>
+<path class="st2" d="M42.5,55.5c0.2,0.1,0.4,0.3,0.7,0.4l8,4.6c-1.1,0.9-2.6,1.7-4.9,3.1l-3.8,2.2l-3.8-2.2
+	c-2.3-1.4-3.8-2.2-4.9-3.1l8-4.6C42.1,55.7,42.3,55.6,42.5,55.5z"/>
+<path class="st3" d="M51.2,24.5c-1.1-0.9-2.6-1.7-4.9-3.1l-3.8-2.2l-3.8,2.2c-2.3,1.4-3.8,2.2-4.9,3.1l8,4.6
+	c0.2,0.1,0.5,0.3,0.7,0.4c0.2-0.1,0.4-0.3,0.7-0.4L51.2,24.5z"/>
+</svg>
diff --git a/web/src/pages/user-setting/setting-model/constant.ts b/web/src/pages/user-setting/setting-model/constant.ts
index 0ac73d6de..b4d726484 100644
--- a/web/src/pages/user-setting/setting-model/constant.ts
+++ b/web/src/pages/user-setting/setting-model/constant.ts
@@ -23,7 +23,8 @@ export const IconMap = {
   NVIDIA:'nvidia',
   'LM-Studio':'lm-studio',
   'OpenAI-API-Compatible':'openai-api',
-  'cohere':'cohere'
+  'cohere':'cohere',
+  'LeptonAI':'lepton'
 };
 
 export const BedrockRegionList = [