细节优化

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,18 +2,20 @@ 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 进行文本嵌入的类
     """
     def __init__(self, model_info: Dict, config) -> None:
         """
         初始化 Ollama Embedding 模型
-        
+
         Args:
             model_info: 模型信息字典
             config: 配置对象
@@ -28,10 +30,10 @@ class OllamaEmbedding:
     def _get_embedding(self, text: str) -> List[float]:
         """
         获取单个文本的嵌入向量
-        
+
         Args:
             text: 输入文本
-            
+
         Returns:
             嵌入向量
         """
@@ -50,10 +52,10 @@ class OllamaEmbedding:
     def predict(self, messages: List[str]) -> List[List[float]]:
         """
         批量获取文本嵌入向量
-        
+
         Args:
             messages: 文本列表
-            
+
         Returns:
             嵌入向量列表
         """
@@ -63,23 +65,23 @@ class OllamaEmbedding:
         for i in range(0, len(messages), batch_size):
             batch = messages[i:i + batch_size]
             logger.info(f"Processing batch {i//batch_size + 1}, size: {len(batch)}")
-            
+
             batch_embeddings = []
             for text in batch:
                 embedding = self._get_embedding(text)
                 batch_embeddings.append(embedding)
-            
+
             embeddings.extend(batch_embeddings)
-            
+
         return embeddings
 
     def encode(self, messages: Union[str, List[str]]) -> List[List[float]]:
         """
         编码文本
-        
+
         Args:
             messages: 单个文本或文本列表
-            
+
         Returns:
             嵌入向量列表
         """
@@ -90,10 +92,10 @@ class OllamaEmbedding:
     def encode_queries(self, queries: List[str]) -> List[List[float]]:
         """
         编码查询文本
-        
+
         Args:
             queries: 查询文本列表
-            
+
         Returns:
             查询文本的嵌入向量列表
         """
@@ -107,7 +109,7 @@ class OllamaReranker:
     def __init__(self, config) -> None:
         """
         初始化 Ollama Reranker
-        
+
         Args:
             config: 配置对象
         """
@@ -119,16 +121,16 @@ class OllamaReranker:
     def compute_score(self, query: str, passage: str) -> float:
         """
         计算查询和文本段落之间的相关性分数
-        
+
         Args:
             query: 查询文本
             passage: 段落文本
-            
+
         Returns:
             相关性分数
         """
         prompt = f"Query: {query}\nPassage: {passage}\nRate the relevance of the passage to the query on a scale of 0 to 1:"
-        
+
         try:
             response = requests.post(
                 f"{self.base_url}/api/generate",
@@ -139,7 +141,7 @@ class OllamaReranker:
                 }
             )
             response.raise_for_status()
-            
+
             # 提取生成的数字作为分数
             result = response.json()["response"].strip()
             try:
@@ -148,7 +150,7 @@ class OllamaReranker:
             except ValueError:
                 logger.warning(f"Could not parse score from response: {result}")
                 return 0.0
-                
+
         except Exception as e:
             logger.error(f"Error computing rerank score: {str(e)}")
             return 0.0
@@ -156,12 +158,12 @@ class OllamaReranker:
     def rerank(self, query: str, passages: List[str], top_n: int = None) -> List[Dict]:
         """
         重新排序文本段落
-        
+
         Args:
             query: 查询文本
             passages: 段落文本列表
             top_n: 返回前 n 个结果
-            
+
         Returns:
             排序后的结果列表，每个元素包含索引和分数
         """
@@ -169,11 +171,11 @@ class OllamaReranker:
         for i, passage in enumerate(passages):
             score = self.compute_score(query, passage)
             scores.append({"index": i, "score": score})
-        
+
         # 按分数降序排序
         sorted_results = sorted(scores, key=lambda x: x["score"], reverse=True)
-        
+
         if top_n:
             sorted_results = sorted_results[:top_n]
-            
+
-        return sorted_results 
+        return sorted_results

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')
 

src/utils/web_search.py

@@ -16,11 +16,11 @@ class WebSearcher:
     def search(self, query: str, max_results: int = 1) -> List[Dict]:
         """
         使用 Tavily 搜索相关内容
-        
+
         Args:
             query: 搜索查询
             max_results: 最大返回结果数
-            
+
         Returns:
             搜索结果列表
         """
@@ -30,7 +30,7 @@ class WebSearcher:
                 search_depth="basic",
                 max_results=max_results
             )
-            
+
             # 提取需要的信息
             formatted_results = []
             for result in search_results['results'][:max_results]:
@@ -40,9 +40,9 @@ class WebSearcher:
                     'url': result.get('url', ''),
                     'score': result.get('score', 0)
                 })
-                
+
             return formatted_results
-            
+
         except Exception as e:
             logger.error(f"Error during web search: {str(e)}")
             return []
@@ -50,20 +50,20 @@ class WebSearcher:
     def format_search_results(self, results: List[Dict]) -> str:
         """
         将搜索结果格式化为文本
-        
+
         Args:
             results: 搜索结果列表
-            
+
         Returns:
             格式化后的文本
         """
         if not results:
             return "没有找到相关的网络搜索结果。"
-            
+
         formatted_text = "以下是相关的网络搜索结果：\n\n"
         for i, result in enumerate(results, 1):
             formatted_text += f"{i}. {result['title']}\n"
             formatted_text += f"   {result['content']}\n"
             formatted_text += f"   来源: {result['url']}\n\n"
-            
+
-        return formatted_text 
+        return formatted_text

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"