graph support

src/core/graphbase.py

@@ -268,18 +268,18 @@ class GraphDatabase:
             return self.query_by_vector(entity_name=entity_name, **kwargs)
 
     def query_by_vector(self, entity_name, threshold=0.9, kgdb_name='neo4j', hops=2, num_of_res=5):
-        result = self.query_by_vector_tep(entity_name=entity_name)
+        results = self.query_by_vector_tep(entity_name=entity_name)
-        querys = []
+
-        for i in range(num_of_res):
+        # 筛选出分数高于阈值的实体
-            if result[i][1] > threshold:
+        qualified_entities = [result[0] for result in results[:num_of_res] if result[1] > threshold]
-                querys.append(result[i][0])
+
-            else:
+        # 对每个合格的实体进行查询
-                break
+        all_query_results = []
-        ans = []
+        for entity in qualified_entities:
-        for query in querys:
+            query_result = self.query_specific_entity(entity_name=entity, hops=hops, kgdb_name=kgdb_name)
-            tep = self.query_specific_entity(entity_name=query, hops=hops) # 这里是只获取第一个 TODO: 优化
+            all_query_results.extend(query_result)
-            ans.extend(tep)
+
-        return ans
+        return all_query_results
 
     def query_by_vector_tep(self, entity_name, kgdb_name='neo4j'):
         """向量查询"""

src/core/retriever.py

@@ -25,27 +25,28 @@ class Retriever:
         return refs
 
     def construct_query(self, query, refs, meta):
-        if len(refs) == 0:
+        if not refs or len(refs) == 0:
             return query
 
-        external = ""
+        external_parts = []
 
         # 解析知识库的结果
-        kb_res = refs.get("knowledge_base").get("results", [])
+        kb_res = refs.get("knowledge_base", {}).get("results", [])
-        if len(kb_res) > 0:
+        if kb_res:
-            kb_text = "\n".join([f"{r['id']}: {r['entity']['text']}" for r in kb_res])
+            kb_text = "\n".join(f"{r['id']}: {r['entity']['text']}" for r in kb_res)
-            external += f"知识库信息: \n\n{kb_text}"
+            external_parts.extend(["知识库信息:", kb_text])
 
         # 解析图数据库的结果
-        db_res = refs.get("graph_base").get("results", [])
+        db_res = refs.get("graph_base", {}).get("results", {})
-        if len(db_res["nodes"]) > 0:
+        if db_res.get("nodes") and len(db_res["nodes"]) > 0:
             db_text = "\n".join(
-                [f"{edge['source_name']}和{edge['target_name']}的关系是{edge['type']}" for edge in db_res["edges"]]
+                [f"{edge['source_name']}和{edge['target_name']}的关系是{edge['type']}" for edge in db_res.get("edges", [])]
             )
-            external += f"图数据库信息: \n\n{db_text}"
+            external_parts.extend(["图数据库信息:", db_text])
 
         # 构造查询
-        if len(external) > 0:
+        if external_parts and len(external_parts) > 0:
+            external = "\n\n".join(external_parts)
             query = f"参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答问题。\n\n问题：{query}\n\n回答："
 
         return query
@@ -73,7 +74,9 @@ class Retriever:
 
         kb_res = []
         final_res = []
-        if not refs["meta"].get("db_name") or not self.config.enable_knowledge_base:
+
+        db_name = refs["meta"].get("db_name")
+        if not db_name or not self.config.enable_knowledge_base:
             return {
                 "results": final_res,
                 "all_results": kb_res,
@@ -83,21 +86,21 @@ class Retriever:
 
         rw_query = self.rewrite_query(query, history, refs)
 
-        db_name = refs["meta"]["db_name"]
         kb = self.dbm.metaname2db[db_name]
         logger.debug(f"{refs['meta']=}")
 
-        max_query_count = refs["meta"].get("maxQueryCount", 10)
+        meta = refs["meta"]
-        rerank_threshold = refs["meta"].get("rerankThreshold", 0.1)
+        max_query_count = meta.get("maxQueryCount", 10)
-        distance_threshold = refs["meta"].get("distanceThreshold", 0)
+        rerank_threshold = meta.get("rerankThreshold", 0.1)
-        top_k = refs["meta"].get("topK", 5)
+        distance_threshold = meta.get("distanceThreshold", 0)
+        top_k = meta.get("topK", 5)
 
         all_kb_res = self.dbm.knowledge_base.search(rw_query, db_name, limit=max_query_count)
         for r in all_kb_res:
             r["file"] = kb.id2file(r["entity"]["file_id"])
 
         # use distance threshold to filter results
-        if refs["meta"].get("mode") == "search":
+        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]
@@ -136,11 +139,12 @@ class Retriever:
 
         entities = []
         if refs["meta"].get("use_graph"):
-            from src.utils.prompts import entity_extraction_prompt_template
+            from src.utils.prompts import entity_extraction_prompt_template as entity_template
+            from src.utils.prompts import keywords_prompt_template as entity_template
 
-            entity_extraction_prompt = entity_extraction_prompt_template.format(text=query)
+            entity_extraction_prompt = entity_template.format(text=query)
-            entities = self.model.predict(entity_extraction_prompt).content.split(",")
+            entities = self.model.predict(entity_extraction_prompt).content.split("<->")
-            entities = [entity for entity in entities if all(char.isalnum() or char in "汉字" for char in entity)]
+            # entities = [entity for entity in entities if all(char.isalnum() or char in "汉字" for char in entity)]
 
         return entities
 
@@ -202,7 +206,7 @@ class Retriever:
         node_dict = {}
 
         for item in results:
-            if not isinstance(item[1], list) or len(item[1]) == 0:
+            if not isinstance(item[1], list) or not item[1]:
                 continue
 
             relationship = item[1][0]
@@ -213,10 +217,7 @@ class Retriever:
             if node_info is None or edge_info is None:
                 continue
 
-            for node in node_info:
+            node_dict.update({node["id"]: node for node in node_info})
-                if node["id"] not in node_dict:
-                    node_dict[node["id"]] = node
-
             formatted_results["edges"].append(edge_info)
 
         formatted_results["nodes"] = list(node_dict.values())

src/utils/prompts.py

@@ -31,3 +31,11 @@ entity_extraction_prompt_template = """
 4.返回的实体用逗号隔开<内容要求>
 <文本>{text}</文本>
 """
+
+keywords_prompt_template = """
+你是用来辅助查询的助手，请对以下文本进行关键词提取，返回提取出的关键词。
+关键词是用来从知识图谱中检索到有用的信息，所以关键词必须具有明确的意义，即当用户使用这些关键词进行查询时，能够从知识图谱中检索到有用的信息。
+返回的实体使用<->隔开。如：关键词1<->关键词<->关键词3
+不要改变关键词的语言
+<文本>{text}</文本>
+"""

web/src/components/ChatComponent.vue

@@ -572,19 +572,19 @@ watch(
     gap: 10px;
 
     .opt__button {
-      background-color: white;
+      background-color: #f2f5f5;
-      color: #222;
+      color: #333;
-      padding: 6px 1rem;
+      padding: .5rem 1.5rem;
-      border-radius: 1rem;
+      border-radius: 2rem;
       cursor: pointer;
       // border: 2px solid var(--main-light-4);
       transition: background-color 0.3s;
-      box-shadow: 0px 0px 10px 4px var(--main-light-4);
+      // box-shadow: 0px 0px 10px 2px var(--main-light-4);
 
 
       &:hover {
-        background-color: #fcfcfc;
+        background-color: #f0f1f1;
-        box-shadow: 0px 0px 10px 1px rgba(0, 0, 0, 0.1);
+        // box-shadow: 0px 0px 10px 1px rgba(0, 0, 0, 0.1);
       }
     }
   }
@@ -672,19 +672,19 @@ watch(
     max-width: 900px;
     margin: 0 auto;
     align-items: flex-end;
-    padding: 0.5rem;
+    padding: 0.25rem 0.5rem;
-    box-shadow: rgba(42, 60, 79, 0.1) 0px 6px 10px 0px;
+    // box-shadow: rgba(42, 60, 79, 0.1) 0px 6px 10px 0px;
-    border: 1px solid #E5E5E5;
+    border: 2px solid #E5E5E5;
-    border-radius: 2rem;
+    border-radius: 1rem;
-    background: #fafafa;
+    background: #fcfdfd;
     transition: background 0.3s, box-shadow 0.3s;
     &:focus-within {
-      border: 1px solid #BABABA;
+      border: 2px solid var(--main-500);
       background: white;
-      box-shadow: rgba(42, 60, 79, 0.1) 0px 6px 10px 0px;
+      // box-shadow: rgb(42 60 79 / 5%) 0px 4px 10px 0px;
     }
 
-    .user-input {
+    textarea.user-input {
       flex: 1;
       height: 40px;
       padding: 0.5rem 0.5rem;
@@ -696,7 +696,7 @@ watch(
       font-size: 16px;
       font-variation-settings: 'wght' 400, 'opsz' 10.5;
       outline: none;
-
+      resize: none;
       &:focus {
         outline: none;
         box-shadow: none;
@@ -706,27 +706,27 @@ watch(
         outline: none;
       }
     }
-
-    .send-btn {
-      border: none;
-      background: transparent;
-      cursor: pointer;
-      font-weight: 500;
-      padding: 0.5rem 1rem;
-      border-radius: 1rem;
-      transition: background-color 0.3s;
-
-      &:hover {
-        background-color: var(--main-light-3);
-      }
-
-      &:disabled {
-        cursor: not-allowed;
-        background-color: #DCDCDC;
-      }
-    }
   }
 
+  button.ant-btn-icon-only {
+    font-size: 1.25rem;
+    cursor: pointer;
+    background-color: transparent;
+    border: none;
+    transition: color 0.3s;
+    box-shadow: none;
+    color: var(--main-700);;
+    padding: 0;
+
+    &:hover {
+      color: var(--c-text-dark-1);
+    }
+
+    &:disabled {
+      color: #ccc;
+      cursor: not-allowed;
+    }
+  }
   .note {
     width: 100%;
     font-size: small;
@@ -745,27 +745,6 @@ watch(
   cursor: pointer;
 }
 
-.ant-btn-icon-only {
-  font-size: 16px;
-  cursor: pointer;
-  background-color: transparent;
-  border: none;
-  height: 2.5rem;
-  background-color: var(--main-color);
-  border-radius: 3rem;
-  color: white;
-  transition: background-color 0.3s;
-
-  &:hover {
-    color: white;
-  }
-}
-
-button:disabled {
-  background: #E0E0E0;
-  cursor: not-allowed;
-}
-
 
 
 .chat::-webkit-scrollbar {

web/src/components/GraphContainer.vue

@@ -0,0 +1,95 @@
+<template>
+    <div class="graph-container" ref="container"></div>
+  </template>
+
+  <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;
+  }
+  </style>

web/src/components/RefsComponent.vue

@@ -1,17 +1,26 @@
-<!-- RefsComponent.vue -->
 <template>
   <div class="refs" v-if="showRefs">
     <div class="tags">
-      <span class="item"><GlobalOutlined /> {{ message.model_name }}</span>
+      <span class="item btn" @click="copyText(msg.text)"><CopyOutlined /></span>
-      <span class="filetag item"
+      <span class="item btn" @click="likeThisResponse(msg)"><LikeOutlined /></span>
-        v-for="(results, filename) in message.groupedResults"
+      <span class="item btn" @click="dislikeThisResponse(msg)"><DislikeOutlined /></span>
+      <span class="item"><GlobalOutlined /> {{ msg.model_name }}</span>
+      <span
+        class="item btn"
+        @click="openSubGraph(msg)"
+        v-if="hasSubGraphData(msg)"
+      >
+        <GlobalOutlined /> 关系图
+      </span>
+      <span class="filetag item btn"
+        v-for="(results, filename) in msg.groupedResults"
         :key="filename"
         @click="toggleDrawer(filename)"
       >
         <FileTextOutlined /> {{ filename }}
         <a-drawer
           v-model:open="openDetail[filename]"
-          title="检索详情"
+          :title="filename"
           width="800"
           :contentWrapperStyle="{ maxWidth: '100%'}"
           placement="right"
@@ -19,51 +28,95 @@
           rootClassName="root"
         >
           <div class="fileinfo">
-            <p><strong>文件名:</strong> {{ filename }}</p>
+            <p><FileOutlined /> {{ results[0].file.type }}</p>
-            <p><strong>文件类型:</strong> {{ results[0].file.type }}</p>
+            <p><ClockCircleOutlined /> {{ formatDate(results[0].file.created_at) }}</p>
-            <p><strong>创建时间:</strong> {{ new Date(results[0].file.created_at * 1000).toLocaleString() }}</p>
           </div>
-          <div v-for="(res, idx) in results" :key="idx" class="result-item">
+          <div class="results-list">
-            <p class="result-id"><strong>ID:</strong> #{{ res.id }}</p>
+            <div v-for="res in results" :key="res.id" class="result-item">
-            <p class="result-distance">
+              <div class="result-meta">
-              <strong>相似度距离:</strong>
+                <div class="score-info">
-              <div class="scorebar">
+                  <span>
-                <a-progress :percent="parseFloat((res.distance * 100).toFixed(2))" stroke-color="#1677FF" :size="[200, 10]"/>
+                    <strong>相似度：</strong>
+                    <a-progress :percent="getPercent(res.distance)"/>
+                  </span>
+                  <span v-if="res.rerank_score">
+                    <strong>重排序：</strong>
+                    <a-progress :percent="getPercent(res.rerank_score)"/>
+                  </span>
+                </div>
+                <div class="result-id">ID: #{{ res.id }}</div>
               </div>
-            </p>
+              <div class="result-text">{{ res.entity.text }}</div>
-            <p class="result-rerank-score" v-if="res.rerank_score">
+            </div>
-              <strong>重排序分数:</strong>
-              <div class="scorebar">
-                <a-progress :percent="parseFloat((res.rerank_score * 100).toFixed(2))" stroke-color="#1677FF" :size="[200, 10]"/>
-              </div>
-            </p>
-            <a-divider />
-            <p class="result-text">{{ res.entity.text }}</p>
           </div>
         </a-drawer>
       </span>
     </div>
+    <a-modal
+      v-model:open="subGraphVisible"
+      title="相关实体与关系"
+      :width="800"
+      :footer="null"
+    >
+      <GraphContainer :graphData="subGraphData" />
+    </a-modal>
   </div>
 </template>
 
 <script setup>
 import { ref, computed, reactive } from 'vue'
+import { useClipboard } from '@vueuse/core'
+import { message } from 'ant-design-vue'
 import {
   GlobalOutlined,
   FileTextOutlined,
+  CopyOutlined,
+  LikeOutlined,
+  DislikeOutlined,
+  FileOutlined,
+  ClockCircleOutlined
 } from '@ant-design/icons-vue'
+import GraphContainer from './GraphContainer.vue'  // 导入 GraphContainer 组件
 
 const props = defineProps({
   message: Object,
 })
 
-const message = ref(props.message)
+const msg = ref(props.message)
+
+// 使用 useClipboard 实现复制功能
+const { copy, isSupported } = useClipboard()
+
+// 定义 copy 方法
+const copyText = async (text) => {
+  if (isSupported) {
+    try {
+      await copy(text)
+      message.success('文本已复制到剪贴板')
+    } catch (error) {
+      console.error('复制失败:', error)
+      message.error('复制失败，请手动复制')
+    }
+  } else {
+    console.warn('浏览器不支持自动复制')
+    message.warning('浏览器不支持自动复制，请手动复制')
+  }
+}
+
+// 其他方法
+const likeThisResponse = (msg) => {
+  console.log('Like this response:', msg)
+}
+
+const dislikeThisResponse = (msg) => {
+  console.log('Dislike this response:', msg)
+}
 
 // 使用 reactive 创建一个响应式对象来存储每个文件的抽屉状态
 const openDetail = reactive({})
 
 // 初始化 openDetail 对象
-for (const filename in message.value.groupedResults) {
+for (const filename in msg.value.groupedResults) {
   openDetail[filename] = false
 }
 
@@ -71,7 +124,40 @@ const toggleDrawer = (filename) => {
   openDetail[filename] = !openDetail[filename]
 }
 
-const showRefs = computed(() => message.value.role=='received' && message.value.status=='finished')
+const showRefs = computed(() => msg.value.role=='received' && msg.value.status=='finished')
+
+const subGraphVisible = ref(false)
+const subGraphData = ref(null)
+
+
+const openSubGraph = (msg) => {
+  if (hasSubGraphData(msg)) {
+    subGraphData.value = msg.refs.graph_base.results
+    subGraphVisible.value = true
+  } else {
+    console.error('无法获取子图数据')
+  }
+}
+
+const closeSubGraph = () => {
+  subGraphVisible.value = false
+}
+
+const hasSubGraphData = (msg) => {
+  return msg.refs &&
+         msg.refs.graph_base &&
+         msg.refs.graph_base.results.nodes.length > 0;
+}
+
+// 添加日期格式化函数
+const formatDate = (timestamp) => {
+  return new Date(timestamp * 1000).toLocaleString()
+}
+
+// 添加百分比计算函数
+const getPercent = (value) => {
+  return parseFloat((value * 100).toFixed(2))
+}
 </script>
 
 <style lang="less" scoped>
@@ -83,11 +169,21 @@ const showRefs = computed(() => message.value.role=='received' && message.value.
   gap: 10px;
 
   .item {
-    background: var(--main-10);
+    background: var(--main-25);
-    color: var(--main-600);
+    color: var(--main-500);
     padding: 2px 8px;
     border-radius: 8px;
     font-size: 14px;
+
+    &.btn {
+      cursor: pointer;
+      &:hover {
+        background: var(--main-25);
+      }
+      &:active {
+        background: var(--main-50);
+      }
+    }
   }
 
   .tags {
@@ -99,56 +195,83 @@ const showRefs = computed(() => message.value.role=='received' && message.value.
       display: flex;
       align-items: center;
       gap: 5px;
-      cursor: pointer;
-
-      &:hover {
-        background: var(--main-25);
-      }
     }
   }
 }
 
-
 .retrieval-detail {
   .fileinfo {
-    margin-bottom: 20px;
+    display: flex;
-    padding: 1rem;
+    justify-content: space-between;
-    background: var(--main-25);
+    padding: 12px 16px;
-    color: var(--main-800);
+    background-color: #f5f5f5;
-    border-radius: 8px;
+    border-radius: 4px;
-    border: 1px solid var(--main-100);
+    margin-bottom: 16px;
+
     p {
-      margin: 10px;
+      margin: 0;
-      line-height: 1.5;
+      color: #666;
     }
   }
 
-  .result-item {
+  .score-info {
-    margin-bottom: 20px;
+    display: flex;
-    padding: 24px 16px 10px 16px;
+    flex-wrap: wrap;
-    border: 1px solid #e8e8e8;
+    gap: 2rem;
-    border-radius: 8px;
+    margin-bottom: 8px;
-    background: var(--main-light-6);
 
-    .result-id,
+    span {
-    .result-distance,
+      display: flex;
-    .result-rerank-score,
+      align-items: center;
-    .result-text-label,
-    .result-text {
-      margin: 5px 0;
-    }
 
-    .scorebar {
+      strong {
-      margin-left: 10px;
+        margin-right: 8px;
-      display: inline-block;
+        white-space: nowrap;
-      width: 200px;
+        color: #666;
-      padding-bottom: 2px;
+      }
 
+      .ant-progress {
+        width: 170px;
+        margin-bottom: 0;
+        margin-inline: 10px;
 
-      & > * {
+        .ant-progress-bg {
-        margin: 0;
+          background-color: #666;
+        }
       }
     }
   }
+
+  .result-id {
+    font-size: 12px;
+    color: #999;
+    margin-bottom: 8px;
+  }
+
+  .result-text {
+    font-size: 14px;
+    line-height: 1.6;
+    white-space: pre-wrap;
+    word-break: break-word;
+    background-color: #f9f9f9;
+    padding: 12px;
+    border-radius: 4px;
+    border: 1px solid #e8e8e8;
+  }
+}
+
+.results-list {
+  .result-item {
+    border-bottom: 1px solid #f0f0f0;
+    padding: 16px 0;
+
+    &:last-child {
+      border-bottom: none;
+    }
+  }
+
+  .result-meta {
+    margin-bottom: 12px;
+  }
 }
 </style>

web/src/layouts/AppLayout.vue

@@ -58,8 +58,18 @@ console.log(route)
 
 <template>
   <div class="app-layout">
-    <div class="debug-panel" @click="showDebug=!showDebug">
+    <div class="debug-panel" >
-      <div class="shown-btn"><BugOutlined /></div>
+      <a-float-button 
+        @click="showDebug=!showDebug"
+        tooltip="调试面板"
+        :style="{
+          right: '12px',
+        }"
+      >
+        <template #icon>
+          <BugOutlined />
+        </template>
+      </a-float-button>
       <a-drawer
         v-model:open="showDebug"
         title="调试面板"
@@ -138,12 +148,6 @@ console.log(route)
     right: 0;
     bottom: 50px;
     border-radius: 20px 0 0 20px;
-    background-color: var(--main-light-4);
-    padding: 6px;
-    padding-left: 12px;
-    box-shadow: 0 0 10px 5px rgba(0, 0, 0, 0.05);
-    border: 1px solid var(--c-black-soft);
-    transition: right 0.3s ease-in-out;
     cursor: pointer;
   }
 }

web/src/views/DataBaseInfoView.vue

@@ -868,29 +868,36 @@ onMounted(() => {
 </style>
 
 <style lang="less">
-.atab-container {
+.db-main-container {
-  padding: 0;
+  .atab-container {
-  width: 100%;
+    padding: 0;
-  max-height: 100%;
+    width: 100%;
-  overflow: auto;
+    max-height: 100%;
+    overflow: auto;
 
-  div.ant-tabs-nav {
+    div.ant-tabs-nav {
-    background: var(--main-light-5);
+      background: var(--main-light-5);
-    padding: 8px 20px;
+      padding: 8px 20px;
-    padding-bottom: 0;
+      padding-bottom: 0;
+    }
+
+    .ant-tabs-content-holder {
+      padding: 0 20px;
+    }
   }
 
-  .ant-tabs-content-holder {
+  .params-item.col .ant-segmented {
-    padding: 0 20px;
+    width: 100%;
+    div.ant-segmented-group {
+      display: flex;
+      justify-content: space-around;
+    }
+    label.ant-segmented-item {
+      flex: 1;
+      text-align: center;
+    }
   }
 }
 
-.params-item.col .ant-segmented {
-  width: 100%;
 
-  div.ant-segmented-group {
-    display: flex;
-    justify-content: space-around;
-  }
-}
 </style>