修改并适配硅基流动模型

src/config/__init__.py

@@ -59,10 +59,10 @@ class Config(SimpleConfig):
         # 模型配置
         ## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
         ## 如果需要自定义本地模型路径，则在 src/.env 中配置 MODEL_DIR
-        self.add_item("model_provider", default="zhipu", des="模型提供商", choices=list(MODEL_NAMES.keys()))
+        self.add_item("model_provider", default="siliconflow", des="模型提供商", choices=list(MODEL_NAMES.keys()))
-        self.add_item("model_name", default=None, des="模型名称")
+        self.add_item("model_name", default="Qwen/Qwen2.5-7B-Instruct", des="模型名称")
-        self.add_item("embed_model", default="zhipu-embedding-3", des="Embedding 模型", choices=list(EMBED_MODEL_INFO.keys()))
+        self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(EMBED_MODEL_INFO.keys()))
-        self.add_item("reranker", default="bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(RERANKER_LIST.keys()))
+        self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(RERANKER_LIST.keys()))
         self.add_item("model_local_paths", default={}, des="本地模型路径")
         self.add_item("use_rewrite_query", default="off", des="重写查询", choices=["off", "on", "hyde"])
         ### <<< 默认配置结束

src/core/knowledgebase.py

@@ -1,6 +1,5 @@
 import os
 
-from src.models.embedding import EmbeddingModel
 from pymilvus import MilvusClient, MilvusException
 from src.utils import setup_logger, hashstr
 logger = setup_logger("KnowledgeBase")
@@ -83,7 +82,7 @@ class KnowledgeBase:
 
     def search(self, query, collection_name, limit=3):
 
-        query_vectors = self.embed_model.encode_queries([query])
+        query_vectors = self.embed_model.batch_encode([query])
         return self.search_by_vector(query_vectors[0], collection_name, limit)
 
     def search_by_vector(self, vector, collection_name, limit=3):

src/core/retriever.py

@@ -1,4 +1,4 @@
-from src.models.embedding import Reranker
+from src.models.rerank_model import get_reranker
 from src.utils.logging_config import setup_logger
 
 logger = setup_logger("server-common")
@@ -12,7 +12,7 @@ class Retriever:
         self.model = model
 
         if self.config.enable_reranker:
-            self.reranker = Reranker(config)
+            self.reranker = get_reranker(config)
 
         if self.config.enable_web_search:
             from src.utils.web_search import WebSearcher
@@ -110,15 +110,13 @@ class Retriever:
         for r in all_kb_res:
             r["file"] = kb.id2file(r["entity"]["file_id"])
 
-        # use distance threshold to filter results
-        if meta.get("mode") == "search":
-            kb_res = all_kb_res
-        else:
         kb_res = [r for r in all_kb_res if r["distance"] > distance_threshold]
 
-        if self.config.enable_reranker:
+        if self.config.enable_reranker and len(kb_res) > 0:
-            for r in kb_res:
+            texts = [r["entity"]["text"] for r in kb_res]
-                r["rerank_score"] = self.reranker.compute_score([rw_query, r["entity"]["text"]], normalize=True)[0]
+            rerank_scores = self.reranker.compute_score([rw_query, texts], normalize=True)
+            for i, r in enumerate(kb_res):
+                r["rerank_score"] = rerank_scores[i]
             kb_res.sort(key=lambda x: x["rerank_score"], reverse=True)
             kb_res = [_res for _res in kb_res if _res["rerank_score"] > rerank_threshold]
 

src/models/embedding.py

@@ -1,105 +1,125 @@
 import os
-from FlagEmbedding import FlagModel, FlagReranker
+import json
+import requests
+from FlagEmbedding import FlagModel
 
-from src.config import EMBED_MODEL_INFO, RERANKER_LIST
+from src.config import EMBED_MODEL_INFO
 from src.utils.logging_config import setup_logger
 from src.utils import hashstr
 
 
 logger = setup_logger("EmbeddingModel")
 
-GLOBAL_EMBED_STATE = {}
+class LocalEmbeddingModel(FlagModel):
-
+    def __init__(self, config, **kwargs):
-
+        info = EMBED_MODEL_INFO[config.embed_model]
-class EmbeddingModel(FlagModel):
+        model_name_or_path = config.model_local_paths.get(info["name"], info.get("default_path"))
-    def __init__(self, model_info, config, **kwargs):
+        logger.info(f"Loading embedding model {info['name']} from {model_name_or_path}")
-        self.info = model_info
-        model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("default_path"))
-        logger.info(f"Loading embedding model {model_info['name']} from {model_name_or_path}")
 
         super().__init__(model_name_or_path,
-                query_instruction_for_retrieval=model_info.get("query_instruction", None),
+                query_instruction_for_retrieval=info.get("query_instruction", None),
                 use_fp16=False, **kwargs)
 
-        logger.info(f"Embedding model {model_info['name']} loaded")
+        logger.info(f"Embedding model {info['name']} loaded")
 
 
-class Reranker(FlagReranker):
-    def __init__(self, config, **kwargs):
-
-        assert config.reranker in RERANKER_LIST.keys(), f"Unsupported Reranker: {config.reranker}, only support {RERANKER_LIST.keys()}"
-
-        model_info = RERANKER_LIST[config.reranker]
-        model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("default_path"))
-        logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}")
-
-        super().__init__(model_name_or_path, use_fp16=True, **kwargs)
-        logger.info(f"Reranker model {config.reranker} loaded")
-
 
 from zhipuai import ZhipuAI
 
-class ZhipuEmbedding:
 
-    def __init__(self, model_info, config) -> None:
+class RemoteEmbeddingModel:
-        self.config = config
+    embed_state = {}
-        self.model_info = model_info
-        self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))
-        logger.info("Zhipu Embedding model loaded")
-        self.query_instruction_for_retrieval = "为这个句子生成表示以用于检索相关文章："
 
-    def predict(self, message):
+    def batch_encode(self, messages, batch_size=20):
         data = []
-        batch_size = 20
 
-        if len(message) > batch_size:
+        if len(messages) > batch_size:
-            global GLOBAL_EMBED_STATE
+            task_id = hashstr(messages)
-            task_id = hashstr(message)
+            self.embed_state[task_id] = {
-            logger.info(f"Creating new state for process {task_id}")
-            GLOBAL_EMBED_STATE[task_id] = {
                 'status': 'in-progress',
-                'total': len(message),
+                'total': len(messages),
                 'progress': 0
             }
 
-        for i in range(0, len(message), batch_size):
+        for i in range(0, len(messages), batch_size):
-            if len(message) > batch_size:
+            group_msg = messages[i:i+batch_size]
-                logger.info(f"Encoding {i} to {i+batch_size} with {len(message)} messages")
+            logger.info(f"Encoding {i} to {i+batch_size} with {len(messages)} messages")
-                GLOBAL_EMBED_STATE[task_id]['progress'] = i
+            response = self.encode_queries(group_msg)
+            data.extend(response)
 
-            group_msg = message[i:i+batch_size]
+        if len(messages) > batch_size:
+            self.embed_state[task_id]['progress'] = len(messages)
+            self.embed_state[task_id]['status'] = 'completed'
+
+        return data
+
+class ZhipuEmbedding(RemoteEmbeddingModel):
+
+    def __init__(self, config) -> None:
+        self.config = config
+        self.model = EMBED_MODEL_INFO[config.embed_model]["name"]
+        self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))
+
+    def predict(self, message):
         response = self.client.embeddings.create(
-                model=self.model_info.get("default_path", None),
+            model=self.model,
-                input=group_msg,
+            input=message,
         )
-
+        data = [a.embedding for a in response.data]
-            data.extend([a.embedding for a in response.data])
-
-        if len(message) > batch_size:
-            GLOBAL_EMBED_STATE[task_id]['progress'] = len(message)
-            GLOBAL_EMBED_STATE[task_id]['status'] = 'completed'
-
         return data
 
     def encode(self, message):
         return self.predict(message)
 
     def encode_queries(self, queries):
-        # queries = [self.query_instruction_for_retrieval + query for query in queries]
         return self.predict(queries)
 
 
+class SiliconFlowEmbedding(RemoteEmbeddingModel):
+
+    def __init__(self, config) -> None:
+        self.url = "https://api.siliconflow.cn/v1/embeddings"
+        self.model = EMBED_MODEL_INFO[config.embed_model]["name"]
+        api_key = os.getenv("SILICONFLOW_API_KEY")
+        assert api_key, "SILICONFLOW_API_KEY is required"
+        self.headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json"
+        }
+
+    def encode(self, message):
+        payload = self.build_payload(message)
+        response = requests.request("POST", self.url, json=payload, headers=self.headers)
+        response = json.loads(response.text)
+        # logger.debug(f"SiliconFlow Embedding response: {response}")
+        assert response["data"], f"SiliconFlow Embedding failed: {response}"
+        data = [a["embedding"] for a in response["data"]]
+        return data
+
+    def encode_queries(self, queries):
+        return self.encode(queries)
+
+    def build_payload(self, message):
+        return {
+            "model": self.model,
+            "input": message,
+        }
+
 def get_embedding_model(config):
     if not config.enable_knowledge_base:
         return None
 
+    provider, model_name = config.embed_model.split('/', 1)
     assert config.embed_model in EMBED_MODEL_INFO.keys(), f"Unsupported embed model: {config.embed_model}, only support {EMBED_MODEL_INFO.keys()}"
+    logger.debug(f"Loading embedding model {config.embed_model}")
+    if provider == "local":
+        model = LocalEmbeddingModel(config)
 
-    if config.embed_model in ["bge-large-zh-v1.5"]:
+    if provider == "zhipu":
-        model = EmbeddingModel(EMBED_MODEL_INFO[config.embed_model], config)
+        model = ZhipuEmbedding(config)
 
-    if config.embed_model in ["zhipu-embedding-2", "zhipu-embedding-3"]:
+    if provider == "siliconflow":
-        model = ZhipuEmbedding(EMBED_MODEL_INFO[config.embed_model], config)
+        model = SiliconFlowEmbedding(config)
 
     return model
 

src/models/rerank_model.py

@@ -0,0 +1,70 @@
+import os
+import json
+import requests
+import numpy as np
+from FlagEmbedding import FlagReranker
+
+from src.config import RERANKER_LIST
+from src.utils.logging_config import setup_logger
+
+
+logger = setup_logger("RerankModel")
+
+
+class LocalReranker(FlagReranker):
+    def __init__(self, config, **kwargs):
+        model_info = RERANKER_LIST[config.reranker]
+        model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("default_path"))
+        logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}")
+
+        super().__init__(model_name_or_path, use_fp16=True, **kwargs)
+        logger.info(f"Reranker model {config.reranker} loaded")
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+class SilconFlowReranker():
+    def __init__(self, config, **kwargs):
+        self.url = "https://api.siliconflow.cn/v1/rerank"
+        self.model = RERANKER_LIST[config.reranker]["name"]
+
+        api_key = os.getenv("SILICONFLOW_API_KEY")
+        assert api_key, "SILICONFLOW_API_KEY is required"
+        self.headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json"
+        }
+
+    def compute_score(self, sentence_pairs, batch_size = 256, max_length = 512, normalize = False):
+        # TODO 还没实现 batch_size
+        query, sentences = sentence_pairs[0], sentence_pairs[1]
+        payload = self.build_payload(query, sentences, max_length)
+        response = requests.request("POST", self.url, json=payload, headers=self.headers)
+        response = json.loads(response.text)
+        logger.debug(f"SiliconFlow Reranker response: {response}")
+
+        results = sorted(response["results"], key=lambda x: x["index"])
+        all_scores = [result["relevance_score"] for result in results]
+
+        if normalize:
+            all_scores = [sigmoid(score) for score in all_scores]
+
+        return all_scores
+
+    def build_payload(self, query, sentences, max_length = 512):
+        return {
+            "model": self.model,
+            "query": query,
+            "documents": sentences,
+            "max_chunks_per_doc": max_length,
+        }
+
+def get_reranker(config):
+    assert config.reranker in RERANKER_LIST.keys(), f"Unsupported Reranker: {config.reranker}, only support {RERANKER_LIST.keys()}"
+    provider, model_name = config.reranker.split('/', 1)
+    if provider == "local":
+        return LocalReranker(config)
+    elif provider == "siliconflow":
+        return SilconFlowReranker(config)
+

src/static/models.yaml

@@ -68,30 +68,33 @@ MODEL_NAMES:
   siliconflow:
     name: SiliconFlow
     url: https://cloud.siliconflow.cn/models
-    default: meta-llama/Meta-Llama-3.1-8B-Instruct
+    default: Qwen/Qwen2.5-7B-Instruct
     env:
       - SILICONFLOW_API_KEY
     models:
       - meta-llama/Meta-Llama-3.1-8B-Instruct
-      - meta-llama/Meta-Llama-3.1-70B-Instruct
+      - Qwen/Qwen2.5-7B-Instruct
-      - meta-llama/Meta-Llama-3.1-405B-Instruct
       - deepseek-ai/DeepSeek-R1
+      - deepseek-ai/DeepSeek-V3
 
 EMBED_MODEL_INFO:
-  bge-m3:
+  local/BAAI/bge-m3:
     name: BAAI/bge-m3
     default_path: BAAI/bge-m3
     dimension: 1024
-  zhipu-embedding-2:
+  zhipu/zhipu-embedding-2:
-    name: zhipu-embedding-2
+    name: embedding-2
-    default_path: embedding-2
     dimension: 1024
-  zhipu-embedding-3:
+  zhipu/zhipu-embedding-3:
-    name: zhipu-embedding-3
+    name: embedding-3
-    default_path: embedding-3
     dimension: 2048
+  siliconflow/BAAI/bge-m3:
+    name: BAAI/bge-m3
+    dimension: 1024
 
 RERANKER_LIST:
-  bge-reranker-v2-m3:
+  local/BAAI/bge-reranker-v2-m3:
     name: BAAI/bge-reranker-v2-m3
     default_path: BAAI/bge-reranker-v2-m3
+  siliconflow/BAAI/bge-reranker-v2-m3:
+    name: BAAI/bge-reranker-v2-m3

web/src/components/ChatComponent.vue

@@ -46,7 +46,7 @@
           <div class="flex-center" @click="meta.summary_title = !meta.summary_title">
             总结对话标题 <div @click.stop><a-switch v-model:checked="meta.summary_title" /></div>
           </div>
-          <div class="flex-center" @click="meta.enable_retrieval = !meta.enable_retrieval">
+          <div class="flex-center" @click="meta.enable_retrieval = !meta.enable_retrieval" v-if="configStore.config.enable_knowledge_base">
             启用检索 <div @click.stop><a-switch v-model:checked="meta.enable_retrieval" /></div>
           </div>
           <div class="flex-center">
@@ -117,13 +117,13 @@
           <div></div>
         </div>
         <div v-else-if="message.status == 'searching' && isStreaming" class="searching-msg"><i>正在检索……</i></div>
-        <div v-else-if="message.status == 'reasoning' && isStreaming" class="searching-msg"><i>正在思考…… {{ message.reasoning }}</i></div>
+        <div v-else-if="message.status == 'reasoning' && isStreaming" class="searching-msg"><i>正在思考…… {{ message.reasoning_content }}</i></div>
         <div
           v-else-if="message.text.length == 0 || message.status == 'error' || (message.status != 'finished' && !isStreaming)"
           class="err-msg"
           @click="retryMessage(message.id)"
         >
-          请求错误，请重试
+          请求错误，请重试。{{ message.message }}
         </div>
         <div v-else
           v-html="renderMarkdown(message)"
@@ -234,15 +234,15 @@ const marked = new Marked(
   })
 );
 
-const consoleMsg = (message) => console.log(message)
+const consoleMsg = (msg) => console.log(msg)
 onClickOutside(panel, () => setTimeout(() => opts.showPanel = false, 30))
 onClickOutside(modelCard, () => setTimeout(() => opts.showModelCard = false, 30))
 
-const renderMarkdown = (message) => {
+const renderMarkdown = (msg) => {
-  if (message.status === 'loading') {
+  if (msg.status === 'loading') {
-    return marked.parse(message.text + '🟢')
+    return marked.parse(msg.text + '🟢')
   } else {
-    return marked.parse(message.text)
+    return marked.parse(msg.text)
   }
 }
 
@@ -306,11 +306,11 @@ const generateRandomHash = (length) => {
     return hash;
 }
 
-const appendUserMessage = (message) => {
+const appendUserMessage = (msg) => {
   conv.value.messages.push({
     id: generateRandomHash(16),
     role: 'sent',
-    text: message
+    text: msg
   })
   scrollToBottom()
 }
@@ -320,7 +320,7 @@ const appendAiMessage = (text, refs=null) => {
     id: generateRandomHash(16),
     role: 'received',
     text: text,
-    reasoning: '',
+    reasoning_content: '',
     refs,
     status: "init",
     meta: {},
@@ -329,40 +329,44 @@ const appendAiMessage = (text, refs=null) => {
 }
 
 const updateMessage = (info) => {
-  const message = conv.value.messages.find((message) => message.id === info.id);
+  const msg = conv.value.messages.find((msg) => msg.id === info.id);
-  if (message) {
+  if (msg) {
     try {
       // 只有在 text 不为空时更新
       if (info.text !== null && info.text !== undefined && info.text !== '') {
-        message.text = info.text;
+        msg.text = info.text;
       }
 
-      if (info.reasoning !== null && info.reasoning !== undefined && info.reasoning !== '') {
+      if (info.reasoning_content !== null && info.reasoning_content !== undefined && info.reasoning_content !== '') {
-        message.reasoning = info.reasoning;
+        msg.reasoning_content = info.reasoning_content;
       }
 
       // 只有在 refs 不为空时更新
       if (info.refs !== null && info.refs !== undefined) {
-        message.refs = info.refs;
+        msg.refs = info.refs;
       }
 
       if (info.model_name !== null && info.model_name !== undefined && info.model_name !== '') {
-        message.model_name = info.model_name;
+        msg.model_name = info.model_name;
       }
 
       // 只有在 status 不为空时更新
       if (info.status !== null && info.status !== undefined && info.status !== '') {
-        message.status = info.status;
+        msg.status = info.status;
       }
 
       if (info.meta !== null && info.meta !== undefined) {
-        message.meta = info.meta;
+        msg.meta = info.meta;
+      }
+
+      if (info.message !== null && info.message !== undefined) {
+        msg.message = info.message;
       }
       scrollToBottom();
     } catch (error) {
       console.error('Error updating message:', error);
-      message.status = 'error';
+      msg.status = 'error';
-      message.text = '消息更新失败';
+      msg.text = '消息更新失败';
     }
   } else {
     console.error('Message not found:', info.id);
@@ -371,9 +375,9 @@ const updateMessage = (info) => {
 
 
 const groupRefs = (id) => {
-  const message = conv.value.messages.find((message) => message.id === id)
+  const msg = conv.value.messages.find((msg) => msg.id === id)
-  if (message.refs && message.refs.knowledge_base.results.length > 0) {
+  if (msg.refs && msg.refs.knowledge_base.results.length > 0) {
-    message.groupedResults = message.refs.knowledge_base.results
+    msg.groupedResults = msg.refs.knowledge_base.results
     .filter(result => result.file && result.file.filename)
     .reduce((acc, result) => {
       const { filename } = result.file;
@@ -387,11 +391,11 @@ const groupRefs = (id) => {
   scrollToBottom()
 }
 
-const simpleCall = (message) => {
+const simpleCall = (msg) => {
   return new Promise((resolve, reject) => {
     fetch('/api/chat/call_lite', {
       method: 'POST',
-      body: JSON.stringify({ query: message, }),
+      body: JSON.stringify({ query: msg, }),
       headers: { 'Content-Type': 'application/json' }
     })
     .then((response) => response.json())
@@ -432,24 +436,11 @@ const fetchChatResponse = (user_input, cur_res_id) => {
     const readChunk = () => {
       return reader.read().then(({ done, value }) => {
         if (done) {
-          const message = conv.value.messages.find((message) => message.id === cur_res_id)
+          const msg = conv.value.messages.find((msg) => msg.id === cur_res_id)
-          console.log(message)
+          console.log(msg)
-          if (message.meta.enable_retrieval) {
+          if (msg.meta.enable_retrieval) {
             console.log("fetching refs")
-            fetchRefs(cur_res_id).then((data) => {
-              console.log(data)
-              updateMessage({
-                id: cur_res_id,
-                refs: data,
-                status: "finished",
-              });
             groupRefs(cur_res_id);
-            })
-          } else {
-            updateMessage({
-              id: cur_res_id,
-              status: "finished",
-            });
           }
           isStreaming.value = false;
           if (conv.value.messages.length === 2) { renameTitle(); }
@@ -468,12 +459,11 @@ const fetchChatResponse = (user_input, cur_res_id) => {
               updateMessage({
                 id: cur_res_id,
                 text: data.response,
-                reasoning: data.reasoning_response,
+                reasoning_content: data.reasoning_content,
-                model_name: data.model_name,
                 status: data.status,
                 meta: data.meta,
+                ...data,
               });
-              // console.log(data)
               // console.log("Last message", conv.value.messages[conv.value.messages.length - 1].text)
               // console.log("Last message", conv.value.messages[conv.value.messages.length - 1].status)
 
@@ -541,7 +531,7 @@ const sendMessage = () => {
 
 const retryMessage = (id) => {
   // 找到 id 对应的 message，然后删除包含 message 在内以及后面所有的 message
-  const index = conv.value.messages.findIndex(message => message.id === id);
+  const index = conv.value.messages.findIndex(msg => msg.id === id);
   const pastMessage = conv.value.messages[index-1]
   console.log("retryMessage", id, pastMessage)
   conv.value.inputText = pastMessage.text
@@ -552,8 +542,8 @@ const retryMessage = (id) => {
   sendMessage();
 }
 
-const autoSend = (message) => {
+const autoSend = (msg) => {
-  conv.value.inputText = message
+  conv.value.inputText = msg
   sendMessage()
 }
 
@@ -750,12 +740,12 @@ watch(
     /* animation: slideInUp 0.1s ease-in; */
 
     .err-msg {
-      color: #FF6B6B;
+      color: #eb8080;
-      border: 1px solid #FF6B6B;
+      border: 1px solid #eb8080;
-      padding: 0.2rem 1rem;
+      padding: 0.5rem 1rem;
       border-radius: 8px;
-      text-align: center;
+      text-align: left;
-      background: #FFF0F0;
+      background: #FFF5F5;
       margin-bottom: 10px;
       cursor: pointer;
     }

web/src/components/RefsComponent.vue

@@ -1,10 +1,10 @@
 <template>
   <div class="refs" v-if="showRefs">
     <div class="tags">
-      <span class="item btn" @click="copyText(msg.text)"><CopyOutlined /></span>
       <!-- <span class="item btn" @click="likeThisResponse(msg)"><LikeOutlined /></span> -->
       <!-- <span class="item btn" @click="dislikeThisResponse(msg)"><DislikeOutlined /></span> -->
       <span class="item"><GlobalOutlined /> {{ msg.model_name }}</span>
+      <span class="item btn" @click="copyText(msg.text)"><CopyOutlined /></span>
       <span
         class="item btn"
         @click="openSubGraph(msg)"
@@ -165,15 +165,15 @@ const getPercent = (value) => {
   display: flex;
   margin-bottom: 20px;
   color: var(--gray-500);
-  font-size: 14px;
+  font-size: 13px;
   gap: 10px;
 
   .item {
     background: var(--gray-100);
-    color: var(--gray-800);
+    color: var(--gray-700);
     padding: 2px 8px;
     border-radius: 8px;
-    font-size: 14px;
+    font-size: 13px;
     user-select: none;
 
     &.btn {