细节优化

2025-02-23 16:39:52 +08:00 · 2025-02-23 16:39:52 +08:00 · 19a8e4714e
commit 19a8e4714e
parent b2dbc17fec
7 changed files with 158 additions and 150 deletions
--- a/src/core/graphbase.py
+++ b/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):
--- a/src/models/ollama_embedding.py
+++ b/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 进行文本嵌入的类
    """
--- a/src/routers/chat_router.py
+++ b/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)
-            refs_pool[cur_res_id] = refs
+            try:
+                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,
-                               reasoning_content=reasoning_content,
-                               status="loading",
-                               history=history_manager.update_ai(content))
+            chunk = make_chunk(content=content, status="loading")
            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')

--- a/web/src/components/GraphContainer.vue
+++ b/web/src/components/GraphContainer.vue
@ -1,95 +1,95 @@
 <template>
-    <div class="graph-container" ref="container"></div>
-  </template>
+  <div class="graph-container" ref="container"></div>
+</template>

-  <script setup>
-  import { Graph } from "@antv/g6";
-  import { onMounted, watch, ref } from 'vue';
+<script setup>
+import { Graph } from "@antv/g6";
+import { onMounted, watch, ref } from 'vue';

-  const props = defineProps({
-    graphData: {
-      type: Object,
-      required: true,
-      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;
+const props = defineProps({
+  graphData: {
+    type: Object,
+    required: true,
+    default: () => ({ nodes: [], edges: [] })
  }
-  </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>
--- a/web/src/views/GraphView.vue
+++ b/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) => {
--- a/web/src/views/SettingView.vue
+++ b/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"