细节优化

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
 
-            modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta)
+            try:
-            refs_pool[cur_res_id] = refs
+                modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta)
+            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/components/GraphContainer.vue

@@ -1,95 +1,95 @@
 <template>
-    <div class="graph-container" ref="container"></div>
+  <div class="graph-container" ref="container"></div>
-  </template>
+</template>
 
-  <script setup>
+<script setup>
-  import { Graph } from "@antv/g6";
+import { Graph } from "@antv/g6";
-  import { onMounted, watch, ref } from 'vue';
+import { onMounted, watch, ref } from 'vue';
 
-  const props = defineProps({
+const props = defineProps({
-    graphData: {
+  graphData: {
-      type: Object,
+    type: Object,
-      required: true,
+    required: true,
-      default: () => ({ nodes: [], edges: [] })
+    default: () => ({ nodes: [], edges: [] })
-    }
-  });
-
-  const container = ref(null);
-  let graphInstance = null;
-
-  const initGraph = () => {
-    graphInstance = new Graph({
-      container: container.value,
-      width: container.value.offsetWidth,
-      height: container.value.offsetHeight,
-      autoFit: true,
-      autoResize: true,
-      layout: {
-        type: 'd3-force',
-        preventOverlap: true,
-        kr: 20,
-        collide: {
-          strength: 1.0,
-        },
-      },
-      node: {
-        type: 'circle',
-        style: {
-          labelText: (d) => d.data.label,
-          size: 70,
-        },
-        palette: {
-          field: 'label',
-          color: 'tableau',
-        },
-      },
-      edge: {
-        type: 'line',
-        style: {
-          labelText: (d) => d.data.label,
-          labelBackground: '#fff',
-          endArrow: true,
-        },
-      },
-      behaviors: ['drag-element', 'zoom-canvas', 'drag-canvas'],
-    });
-  };
-
-  const renderGraph = () => {
-    if (!graphInstance) {
-      initGraph();
-    }
-
-    const formattedData = {
-      nodes: props.graphData.nodes.map(node => ({
-        id: node.id,
-        data: { label: node.name }
-      })),
-      edges: props.graphData.edges.map(edge => ({
-        source: edge.source_id,
-        target: edge.target_id,
-        data: { label: edge.type }
-      }))
-    };
-
-    graphInstance.setData(formattedData);
-    graphInstance.render();
-  };
-
-  onMounted(() => {
-    renderGraph();
-    window.addEventListener('resize', renderGraph);
-  });
-
-  watch(() => props.graphData, renderGraph, { deep: true });
-  </script>
-
-  <style scoped>
-  .graph-container {
-    background: #F7F7F7;
-    border-radius: 16px;
-    width: 100%;
-    height: 600px;
-    overflow: hidden;
   }
-  </style>
+});
+
+const container = ref(null);
+let graphInstance = null;
+
+const initGraph = () => {
+  graphInstance = new Graph({
+    container: container.value,
+    width: container.value.offsetWidth,
+    height: container.value.offsetHeight,
+    autoFit: true,
+    autoResize: true,
+    layout: {
+      type: 'd3-force',
+      preventOverlap: true,
+      kr: 20,
+      collide: {
+        strength: 1.0,
+      },
+    },
+    node: {
+      type: 'circle',
+      style: {
+        labelText: (d) => d.data.label,
+        size: 70,
+      },
+      palette: {
+        field: 'label',
+        color: 'tableau',
+      },
+    },
+    edge: {
+      type: 'line',
+      style: {
+        labelText: (d) => d.data.label,
+        labelBackground: '#fff',
+        endArrow: true,
+      },
+    },
+    behaviors: ['drag-element', 'zoom-canvas', 'drag-canvas'],
+  });
+};
+
+const renderGraph = () => {
+  if (!graphInstance) {
+    initGraph();
+  }
+
+  const formattedData = {
+    nodes: props.graphData.nodes.map(node => ({
+      id: node.id,
+      data: { label: node.name }
+    })),
+    edges: props.graphData.edges.map(edge => ({
+      source: edge.source_id,
+      target: edge.target_id,
+      data: { label: edge.type }
+    }))
+  };
+
+  graphInstance.setData(formattedData);
+  graphInstance.render();
+};
+
+onMounted(() => {
+  renderGraph();
+  window.addEventListener('resize', renderGraph);
+});
+
+watch(() => props.graphData, renderGraph, { deep: true });
+</script>
+
+<style scoped>
+.graph-container {
+  background: #F7F7F7;
+  border-radius: 16px;
+  width: 100%;
+  height: 600px;
+  overflow: hidden;
+}
+</style>

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">
@@ -43,7 +43,7 @@
       </div>
     </div>
     <div class="main" id="container" ref="container" v-show="graphData.nodes.length > 0"></div>
-    <a-empty v-show="graphData.nodes.length === 0"  style="padding: 4rem 0;"/>
+    <a-empty v-show="graphData.nodes.length === 0" style="padding: 4rem 0;"/>
 
     <a-modal
       :open="state.showModal" title="上传文件"
@@ -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"