add support for 01.AI (#1951)

### What problem does this PR solve? #1853 add support for 01.AI ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
infiniflow · Aug 15, 2024 · c59c1b6 · c59c1b6
1 parent c9caccf
commit c59c1b6
Show file tree

Hide file tree

Showing 6 changed files with 95 additions and 7 deletions.
diff --git a/conf/llm_factories.json b/conf/llm_factories.json
@@ -3051,6 +3051,68 @@
                     "model_type": "rerank"
                 }
             ]
-        }        
+        },
+        {
+            "name": "01.AI",
+            "logo": "",
+            "tags": "LLM,IMAGE2TEXT",
+            "status": "1",
+            "llm": [
+                {
+                    "llm_name": "yi-large",
+                    "tags": "LLM,CHAT,32k",
+                    "max_tokens": 32768,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-medium",
+                    "tags": "LLM,CHAT,16k",
+                    "max_tokens": 16384,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-medium-200k",
+                    "tags": "LLM,CHAT,200k",
+                    "max_tokens": 204800,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-spark",
+                    "tags": "LLM,CHAT,16k",
+                    "max_tokens": 16384,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-large-rag",
+                    "tags": "LLM,CHAT,16k",
+                    "max_tokens": 16384,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-large-fc",
+                    "tags": "LLM,CHAT,32k",
+                    "max_tokens": 32768,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-large-turbo",
+                    "tags": "LLM,CHAT,16k",
+                    "max_tokens": 16384,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-large-preview",
+                    "tags": "LLM,CHAT,16k",
+                    "max_tokens": 16384,
+                    "model_type": "chat"
+                },
+                {
+                    "llm_name": "yi-vision",
+                    "tags": "LLM,CHAT,IMAGE2TEXT,16k",
+                    "max_tokens": 16384,
+                    "model_type": "image2text"
+                }
+            ]
+        }                
     ]
 }
diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
@@ -61,7 +61,8 @@
     "LM-Studio": LmStudioCV,
     "StepFun":StepFunCV,
     "OpenAI-API-Compatible": OpenAI_APICV,
-    "TogetherAI": TogetherAICV
+    "TogetherAI": TogetherAICV,
+    "01.AI": YiCV
 }
 
 
@@ -94,7 +95,8 @@
     "PerfXCloud": PerfXCloudChat,
     "Upstage":UpstageChat,
     "novita.ai": NovitaAIChat,
-    "SILICONFLOW": SILICONFLOWChat
+    "SILICONFLOW": SILICONFLOWChat,
+    "01.AI": YiChat
 }
 
 

diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
@@ -1022,4 +1022,11 @@ class SILICONFLOWChat(Base):
     def __init__(self, key, model_name, base_url="https://api.siliconflow.cn/v1"):
         if not base_url:
             base_url = "https://api.siliconflow.cn/v1"
+        super().__init__(key, model_name, base_url)
+
+
+class YiChat(Base):
+    def __init__(self, key, model_name, base_url="https://api.lingyiwanwu.com/v1"):
+        if not base_url:
+            base_url = "https://api.lingyiwanwu.com/v1"
         super().__init__(key, model_name, base_url)
diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py
@@ -622,15 +622,17 @@ def chat_prompt(self, text, b64):
             }
         ]
 
+
 class StepFunCV(GptV4):
     def __init__(self, key, model_name="step-1v-8k", lang="Chinese", base_url="https://api.stepfun.com/v1"):
         if not base_url: base_url="https://api.stepfun.com/v1"
         self.client = OpenAI(api_key=key, base_url=base_url)
         self.model_name = model_name
         self.lang = lang
 
+
 class LmStudioCV(GptV4):
-    def __init__(self, key, model_name, base_url, lang="Chinese"):
+    def __init__(self, key, model_name, lang="Chinese", base_url=""):
         if not base_url:
             raise ValueError("Local llm url cannot be None")
         if base_url.split("/")[-1] != "v1":
@@ -641,7 +643,7 @@ def __init__(self, key, model_name, base_url, lang="Chinese"):
 
 
 class OpenAI_APICV(GptV4):
-    def __init__(self, key, model_name, base_url, lang="Chinese"):
+    def __init__(self, key, model_name, lang="Chinese", base_url=""):
         if not base_url:
             raise ValueError("url cannot be None")
         if base_url.split("/")[-1] != "v1":
@@ -652,7 +654,14 @@ def __init__(self, key, model_name, base_url, lang="Chinese"):
 
 
 class TogetherAICV(GptV4):
-    def __init__(self, key, model_name, base_url="https://api.together.xyz/v1"):
+    def __init__(self, key, model_name, lang="Chinese", base_url="https://api.together.xyz/v1"):
         if not base_url:
             base_url = "https://api.together.xyz/v1"
-        super().__init__(key, model_name, base_url)
+        super().__init__(key, model_name,lang,base_url)
+
+
+class YiCV(GptV4):
+    def __init__(self, key, model_name, lang="Chinese",base_url="https://api.lingyiwanwu.com/v1",):
+        if not base_url:
+            base_url = "https://api.lingyiwanwu.com/v1"
+        super().__init__(key, model_name,lang,base_url)
diff --git a/web/src/assets/svg/llm/yi.svg b/web/src/assets/svg/llm/yi.svg
diff --git a/web/src/pages/user-setting/setting-model/constant.ts b/web/src/pages/user-setting/setting-model/constant.ts
@@ -30,6 +30,7 @@ export const IconMap = {
   Upstage: 'upstage',
   'novita.ai': 'novita-ai',
   SILICONFLOW: 'siliconflow',
+  "01.AI": 'yi'
 };
 
 export const BedrockRegionList = [