细节优化

src/core/graphbase.py

@@ -27,6 +27,7 @@ class GraphDatabase:
         self.kgdb_name = kgdb_name
         assert embed_model, "embed_model=None"
         self.embed_model = embed_model
+        self.embed_model_name = None
 
     def start(self):
         uri = os.environ.get("NEO4J_URI", "bolt://localhost:7687")
@@ -87,7 +88,8 @@ class GraphDatabase:
                 "relationship_count": relationship_count,
                 "triples_count": triples_count,
                 "labels": labels,
-                "status": self.status
+                "status": self.status,
+                "embed_model_name": self.embed_model_name
             }
 
         with self.driver.session() as session:
@@ -171,6 +173,7 @@ class GraphDatabase:
     def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'):
         self.status = "processing"
         kgdb_name = kgdb_name or 'neo4j'
+        self.embed_model_name = self.embed_model_name or self.config.embed_model
         self.use_database(kgdb_name)  # 切换到指定数据库
 
         def read_triples(file_path):

src/models/ollama_embedding.py

@@ -2,11 +2,13 @@ import os
 import requests
 import numpy as np
 from typing import List, Union, Dict
+
+from src.models.embedding import RemoteEmbeddingModel
 from src.utils.logging_config import setup_logger
 
 logger = setup_logger("OllamaEmbedding")
 
-class OllamaEmbedding:
+class OllamaEmbedding(RemoteEmbeddingModel):
     """
     使用 Ollama API 进行文本嵌入的类
     """

src/routers/chat_router.py

@@ -27,26 +27,29 @@ def chat_post(
 
     history_manager = HistoryManager(history)
 
-    def make_chunk(content=None, status=None, history=None, reasoning_content=None):
+    def make_chunk(content=None, **kwargs):
         return json.dumps({
             "response": content,
-            "reasoning_response": reasoning_content,
-            "history": history,
             "model_name": startup.config.model_name,
-            "status": status,
             "meta": meta,
+            **kwargs
         }, ensure_ascii=False).encode('utf-8') + b"\n"
 
     def generate_response():
         modified_query = query
+        refs = None
 
         # 处理知识库检索
         if meta and meta.get("enable_retrieval"):
             chunk = make_chunk(status="searching")
             yield chunk
 
+            try:
                 modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta)
-            refs_pool[cur_res_id] = refs
+            except Exception as e:
+                logger.error(f"Retriever error: {e}")
+                yield make_chunk(message=f"Retriever error: {e}", status="error")
+                return
 
         messages = history_manager.get_history_with_msg(modified_query, max_rounds=meta.get('history_round'))
         history_manager.add_user(query)  # 注意这里使用原始查询
@@ -56,7 +59,7 @@ def chat_post(
         reasoning_content = ""
         for delta in startup.model.predict(messages, stream=True):
             if not delta.content and hasattr(delta, 'reasoning_content'):
-                reasoning_content += delta.reasoning_content
+                reasoning_content += delta.reasoning_content or ""
                 chunk = make_chunk(reasoning_content=reasoning_content, status="reasoning")
                 yield chunk
                 continue
@@ -67,14 +70,15 @@ def chat_post(
             else:
                 content += delta.content or ""
 
-            chunk = make_chunk(content=content,
+            chunk = make_chunk(content=content, status="loading")
-                               reasoning_content=reasoning_content,
-                               status="loading",
-                               history=history_manager.update_ai(content))
             yield chunk
 
         logger.debug(f"Final response: {content}")
         logger.debug(f"Final reasoning response: {reasoning_content}")
+        yield make_chunk(content=content,
+                         status="finished",
+                         history=history_manager.update_ai(content),
+                         refs=refs)
 
     return StreamingResponse(generate_response(), media_type='application/json')
 

web/src/views/GraphView.vue

@@ -11,7 +11,7 @@
   <div class="graph-container layout-container" v-else>
     <HeaderComponent
       title="图数据库"
-      :description="`${graphInfo?.database_name || ''} - 共 ${graphInfo?.entity_count || 0} 实体，${graphInfo?.relationship_count || 0} 个关系`"
+      :description="`${graphInfo?.database_name || ''} - 共 ${graphInfo?.entity_count || 0} 实体，${graphInfo?.relationship_count || 0} 个关系。向量模型：${graphInfo?.embed_model_name || '未上传文件'}`"
     >
       <template #actions>
         <div class="status-wrapper">
@@ -51,6 +51,11 @@
       @cancel="() => state.showModal = false"
       ok-text="添加到图数据库" cancel-text="取消"
       :confirm-loading="state.precessing">
+      <div v-if="graphInfo?.embed_model_name">
+        <p>当前图数据库向量模型：{{ graphInfo?.embed_model_name }}</p>
+        <p>当前所选择的向量模型是 {{ configStore.config.embed_model }}</p>
+      </div>
+      <p v-else>第一次创建之后将无法修改向量模型，当前向量模型 {{ configStore.config.embed_model }}</p>
       <div class="upload">
         <a-upload-dragger
           class="upload-dragger"
@@ -58,7 +63,7 @@
           name="file"
           :fileList="fileList"
           :max-count="1"
-          :disabled="state.precessing"
+          :disabled="state.precessing || (graphInfo?.embed_model_name && graphInfo?.embed_model_name !== configStore.config.embed_model)"
           action="/api/data/upload"
           @change="handleFileUpload"
           @drop="handleDrop"
@@ -181,7 +186,7 @@ const loadSampleNodes = () => {
       graphData.nodes = data.result.nodes
       graphData.edges = data.result.edges
       console.log(graphData)
-      randerGraph()
+      setTimeout(() => randerGraph(), 500)
     })
     .catch((error) => {
       message.error(error.message);
@@ -195,12 +200,6 @@ const onSearch = () => {
     return
   }
 
-  const cur_embed_model = configStore.config.embed_model
-  if (cur_embed_model !== 'zhipu-embedding-3') {
-    message.error('当前不支持实体检索，请在设置中选择向量模型为 zhipu-embedding-3')
-    return
-  }
-
   state.searchLoading = true
   fetch(`/api/data/graph/node?entity_name=${state.searchInput}`)
     .then((res) => {

web/src/views/SettingView.vue

@@ -24,7 +24,7 @@
         <div class="section">
           <div class="card">
             <span class="label">{{ items?.embed_model.des }}</span>
-            <a-select style="width: 200px"
+            <a-select style="width: 300px"
               :value="configStore.config?.embed_model"
               @change="handleChange('embed_model', $event)"
             >
@@ -36,7 +36,7 @@
           </div>
           <div class="card">
             <span class="label">{{ items?.reranker.des }}</span>
-            <a-select style="width: 200px"
+            <a-select style="width: 300px"
               :value="configStore.config?.reranker"
               @change="handleChange('reranker', $event)"
               :disabled="!configStore.config.enable_reranker"