update graph page

src/core/database.py

@@ -51,10 +51,10 @@ class DataBaseManager:
         self.database_path = "data/databases.json"
         self.embed_model = get_embedding_model(config)
         self.knowledge_base = KnowledgeBase(config, self.embed_model)
+        self.graph_base = GraphDatabase(self.config, self.embed_model)
         self.data = {"databases": [], "graph": {}}
 
         if self.config.enable_knowledge_graph:
-            self.graph_base = GraphDatabase(self.config, self.embed_model)
             self.graph_base.start()
 
         self._load_databases()
@@ -130,10 +130,9 @@ class DataBaseManager:
         for idx, file in new_files:
             db.files[idx]["status"] = "processing"
 
-            text = self.read_text(file)
-            chunks = self.chunking(text)
-
             try:
+                text = self.read_text(file)
+                chunks = self.chunking(text)
                 self.knowledge_base.add_documents(
                     docs=chunks,
                     collection_name=db.metaname,
@@ -156,7 +155,7 @@ class DataBaseManager:
             return db.to_dict()
 
     def read_text(self, file):
-        support_format = [".pdf", ".txt", "*.md"]
+        support_format = [".pdf", ".txt", ".md"]
         assert os.path.exists(file), "File not found"
         logger.info(f"Try to read file {file}")
 

src/core/graphbase.py

@@ -6,7 +6,10 @@ from neo4j import GraphDatabase as GD
 from transformers import AutoTokenizer, AutoModel
 from FlagEmbedding import FlagModel, FlagReranker
 
+from plugins import pdf2txt
+from plugins.oneke import OneKE
 
+UIE_MODEL = None
 
 class GraphDatabase:
     def __init__(self, config, embed_model=None):
@@ -85,8 +88,10 @@ class GraphDatabase:
     def pdf_file_add_entity(self, file_path, output_path, kgdb_name='neo4j'):
         self.use_database(kgdb_name)  # 切换到指定数据库
         text_path = pdf2txt(file_path)
-        oneke = OneKE()
+        global UIE_MODEL
-        triples_path = oneke.processing_text_to_kg(text_path, output_path)
+        if UIE_MODEL is None:
+            UIE_MODEL = OneKE()
+        triples_path = UIE_MODEL.processing_text_to_kg(text_path, output_path)
         def read_triples(file_path):
             with open(file_path, 'r', encoding='utf-8') as file:
                 for line in file:
@@ -95,7 +100,7 @@ class GraphDatabase:
         for trio in read_triples(triples_path):
             self.txt_add_entity(trio, kgdb_name)
         return kgdb_name
-    
+
     def txt_add_vector_entity(self, triples, kgdb_name='neo4j'):
         """添加实体三元组"""
         self.use_database(kgdb_name)
@@ -134,15 +139,20 @@ class GraphDatabase:
                 session.execute_write(self.set_embedding, entry['t'], embedding_t)
 
     def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'):
+        self.status = "processing"
         self.use_database(kgdb_name)  # 切换到指定数据库
         triples_path = file_path
+
         def read_triples(file_path):
             with open(file_path, 'r', encoding='utf-8') as file:
                 for line in file:
                     item = json.loads(line.strip())
                     yield [item]
+
         for trio in read_triples(triples_path):
             self.txt_add_entity(trio, kgdb_name)
+
+        self.status = "open"
         return kgdb_name
 
     def delete_entity(self, entity_name=None, kgdb_name="neo4j"):
@@ -221,7 +231,7 @@ class GraphDatabase:
 
         with self.driver.session() as session:
             return session.execute_read(query, keyword, hops)
-        
+
     def query_by_vector_tep(self, keyword, kgdb_name='neo4j'):
         """向量查询"""
         self.use_database(kgdb_name)
@@ -238,11 +248,11 @@ class GraphDatabase:
 
         with self.driver.session() as session:
             return session.execute_read(query, keyword)
-    
+
-    def query_by_vector(self, entity_name, kgdb_name='neo4j'):
+    def query_by_vector(self, entity_name, kgdb_name='neo4j', hops=2):
         self.use_database(kgdb_name)
         result = self.query_by_vector_tep(entity_name)
-        ans = self.query_specific_entity(result[0][0])
+        ans = self.query_specific_entity(result[0][0], hops) # 这里是只获取第一个 TODO: 优化
         return ans
 
     def query_node_info(self, node_name, kgdb_name='neo4j', hops = 2):
@@ -258,20 +268,20 @@ class GraphDatabase:
 
         with self.driver.session() as session:
             return session.execute_read(query, node_name, hops)
-        
+
     def get_embedding(self, text):
         inputs = [text]
         with torch.no_grad():
             outputs = self.embed_model.encode(inputs)
         embeddings = outputs[0] # 假设取平均作为文本的嵌入向量
         return embeddings
-    
+
     def set_embedding(self, tx, entity_name, embedding):
         tx.run("""
         MATCH (e:Entity {name: $name})
         CALL db.create.setNodeVectorProperty(e, 'embedding', $embedding)
         """, name=entity_name, embedding=embedding)
-        
+
     # def format_query_results(self, results):
     #     formatted_results = []
     #     for row in results:
@@ -301,7 +311,7 @@ if __name__ == "__main__":
 
             super().__init__(model_name_or_path, use_fp16=False, **kwargs)
 
-   
+
     model = EmbeddingModel(config)
     # 初始化知识图谱数据库
     kgdb = GraphDatabase(config, model)
@@ -311,7 +321,7 @@ if __name__ == "__main__":
     # 返回指定数据库信息
     # info = kgdb.get_database_info(kgdb_name)
     # print(info)
-        
+
     # triples = [
     #         {
     #             "h": "CCC",
@@ -324,14 +334,14 @@ if __name__ == "__main__":
     #             "r": "同事"
     #         }
     #         ]
-    
+
     # print(kgdb.query_by_vector("维C"))
     def format_query_results(results):
         formatted_results = {"nodes": [], "edges": []}
-        
+
         # 用于存储所有唯一的节点信息
         node_dict = {}
-        
+
         for item in results:
             # 确保item[1]是一个非空的列表
             if isinstance(item[1], list) and len(item[1]) > 0:
@@ -340,25 +350,25 @@ if __name__ == "__main__":
                 nodes = relationship.nodes
                 if len(nodes) == 2:
                     node1, node2 = nodes
-                    
+
                     # 提取源节点和目标节点信息
                     node1_id = node1.element_id
                     node2_id = node2.element_id
                     node1_name = item[0]  # 假设节点名称和列表中的第一个元素相同
                     node2_name = item[2] if len(item) > 2 else 'unknown'
-                    
+
                     # 记录节点信息
                     if node1_id not in node_dict:
                         node_dict[node1_id] = {"id": node1_id, "name": node1_name}
                     if node2_id not in node_dict:
                         node_dict[node2_id] = {"id": node2_id, "name": node2_name}
-                    
+
                     # 确定关系类型
-                    
+
                     relationship_type = relationship._properties.get('type', 'unknown')
                     if relationship_type == 'unknown':
                         relationship_type = relationship.type
-                    
+
                     # 记录边的信息
                     formatted_results["edges"].append({
                         "id": rel_id,
@@ -368,10 +378,10 @@ if __name__ == "__main__":
                         "source_name": node1_name,
                         "target_name": node2_name
                     })
-        
+
         # 将唯一的节点信息添加到结果中
         formatted_results["nodes"] = list(node_dict.values())
-        
+
         return formatted_results
 
     entities = ['jqy', '维c']
@@ -379,7 +389,7 @@ if __name__ == "__main__":
     for entitie in entities:
         result = kgdb.query_by_vector(entitie)
         if result != []:
-            results.extend(result) 
+            results.extend(result)
     print(format_query_results(results))
 
 

src/views/database_view.py

@@ -86,9 +86,10 @@ def upload_file():
     file = request.files['file']
     if file.filename == '':
         return jsonify({'message': 'No selected file'}), 400
-    elif file.filename.split('.')[-1] not in ['pdf', 'txt', 'md']:
+    # elif file.filename.split('.')[-1] not in ['pdf', 'txt', 'md']:
-        return jsonify({'message': 'Unsupported file type'}), 400
+    #     return jsonify({'message': 'Unsupported file type'}), 400
     if file:
+        os.makedirs("data/uploads", exist_ok=True)
         filename = file.filename
         file_path = os.path.join("data/uploads", filename)
         file.save(file_path)
@@ -98,3 +99,28 @@ def upload_file():
 def get_graph_info():
     graph_info = startup.dbm.get_graph()
     return jsonify(graph_info)
+
+@db.route('/graph/node', methods=['GET'])
+def get_graph_node():
+    entity_name = request.args.get('entity_name')
+    kgdb_name = request.args.get('kgdb_name')
+    hops = request.args.get('hops')
+    if not entity_name:
+        return jsonify({'message': 'entity_name and kgdb_name are required'}), 400
+
+    logger.debug(f"Get graph node {entity_name} in {kgdb_name} with {hops} hops")
+    result = startup.dbm.graph_base.query_by_vector(entity_name, kgdb_name, hops)
+    return jsonify({'result': startup.retriever.format_query_results(result), 'message': 'success'}), 200
+
+@db.route('/graph/add', methods=['POST'])
+def add_graph_entity():
+    data = json.loads(request.data)
+    kgdb_name = data.get('kgdb_name')
+    file_path = data.get('file_path')
+
+    if file_path.endswith('.jsonl'):
+        startup.dbm.graph_base.jsonl_file_add_entity(file_path, kgdb_name)
+    else:
+        return jsonify({'message': 'Unsupported file type'}), 400
+
+    return jsonify({'message': 'Entity successfully added'}), 200

web/src/components/ChatComponent.vue

@@ -285,8 +285,9 @@ const sendMessage = () => {
   if (conv.value.inputText.trim()) {
     isStreaming.value = true
     appendUserMessage(conv.value.inputText)
+    appendAiMessage("检索中……", null)
+    const cur_res_id = conv.value.messages[conv.value.messages.length - 1].id
     const user_input = conv.value.inputText
-    var cur_res_id = null
     conv.value.inputText = ''
     fetch('/api/chat', {
       method: 'POST',
@@ -319,12 +320,7 @@ const sendMessage = () => {
 
           try {
             const data = JSON.parse(message)
-            if (cur_res_id === null) {
+            updateMessage(data.response, cur_res_id)
-              appendAiMessage(data.response, data.refs)
-              cur_res_id = conv.value.messages[conv.value.messages.length - 1].id
-            } else {
-              updateMessage(data.response, cur_res_id)
-            }
             conv.value.history = data.history
             buffer = ''
           } catch (e) {

web/src/views/ChatView.vue

@@ -185,6 +185,7 @@ onMounted(() => {
   cursor: pointer;
   width: 100%;
   user-select: none;
+  transition: border-left 0.1s;
 
   &__title {
     white-space: nowrap; /* 禁止换行 */
@@ -203,6 +204,7 @@ onMounted(() => {
   }
 
   &.active {
+    border-left: 4px solid var(--main-color);
     background-color: #EDF4F5;
   }
 

web/src/views/DataBaseView.vue

@@ -47,7 +47,7 @@
     </div>
     <h2>图数据库</h2>
     <p>基于 neo4j 构建的图数据库。</p>
-    <div :class="{'graphloading': graphloading}">
+    <div :class="{'graphloading': graphloading}" v-if="graph">
       <div class="dbcard graphbase" @click="navigateToGraph">
         <div class="top">
           <div class="icon"><AppstoreFilled /></div>

web/src/views/GraphView.vue

@@ -1,48 +1,271 @@
 <template>
   <div class="graph-container">
+    <div class="info">
+      <h1>Neo4j 图数据库</h1>
+      <p>基于 Neo4j 构建的图数据库。</p>
+      </div>
+    <div class="actions">
+      <div class="actions-left">
+        <a-button @click="state.showModal = true">上传文件</a-button>
+        <a-modal v-model:open="state.showModal" title="上传文件" @ok="handleUpload">
+          <div class="upload">
+            <a-upload-dragger
+              class="upload-dragger"
+              v-model:fileList="fileList"
+              name="file"
+              :max-count="1"
+              :disabled="state.precessing"
+              action="/api/database/upload"
+              @change="handleFileUpload"
+              @drop="handleDrop"
+            >
+              <p class="ant-upload-text">点击或者把文件拖拽到这里上传</p>
+              <p class="ant-upload-hint">
+                目前仅支持上传文本文件，如 .pdf, .txt, .md。且同名文件无法重复添加。
+              </p>
+            </a-upload-dragger>
+          </div>
+          <a-button
+            type="primary"
+            @click="addDocumentByFile"
+            :loading="state.precessing"
+            :disabled="fileList.length === 0"
+            style="margin: 0px 20px 20px 0;"
+          >
+            添加到图数据库
+          </a-button>
+          <a-button @click="handleRefresh" :loading="state.refrashing">刷新状态</a-button>
+        </a-modal>
+      </div>
+      <div class="action-right">
+        <input
+          v-model="state.searchInput"
+          placeholder="输入要查询的实体"
+          style="width: 200px"
+        />
+        <a-button
+          type="primary"
+          :loading="state.searchLoading"
+          @click="onSearch"
+        >
+          检索实体
+        </a-button>
+      </div>
+    </div>
     <div class="main" id="container"></div>
+
   </div>
 </template>
 
 <script setup>
 import { Graph } from "@antv/g6";
-import { onMounted } from 'vue';
+import { computed, onMounted, reactive, ref } from 'vue';
+import { message } from "ant-design-vue";
+
+let graphInstance
+const fileList = ref([]);
+const subgraph = reactive({
+  nodes: [
+    { id: '1', name: 'node1' },
+    { id: '2', name: 'node2' },
+    { id: '3', name: 'node3' },
+    { id: '4', name: 'node4' },
+    { id: '5', name: 'node5' },
+  ],
+  edges: [
+    { id: 'e1', source_id: '1', target_id: '2', type: 'edge1' },
+    { id: 'e2', source_id: '1', target_id: '3', type: 'edge2' },
+    { id: 'e3', source_id: '2', target_id: '4', type: 'edge3' },
+    { id: 'e4', source_id: '2', target_id: '5', type: 'edge4' },
+  ],
+});
+
+const state = reactive({
+  searchInput: '',
+  searchLoading: false,
+  showModal: false,
+  precessing: false,
+})
 
 const getCurWidth = () => document.getElementById("container").offsetWidth
 const getCurHeight = () => document.getElementById("container").offsetHeight
 
+const graphData = computed(() => {
+  return {
+    nodes: subgraph.nodes.map(node => {
+      return {
+        id: node.id,
+        data: {
+          label: node.name
+        },
+      }
+    }),
+    edges: subgraph.edges.map(edge => {
+      return {
+        source: edge.source_id,
+        target: edge.target_id,
+        data: {
+          label: edge.type
+        }
+      }
+    }),
+  }
+})
+
+const addDocumentByFile = () => {
+  state.precessing = true
+  const files = fileList.value.filter(file => file.status === 'done').map(file => file.response.file_path)
+  fetch('/api/database/graph/add', {
+    method: 'POST',
+    body: JSON.stringify({
+      file_path: files[0]
+    }),
+  })
+  // .then(response => response.json())
+  //   .then((data) => {
+  //     message.success(data.message);
+  //   })
+  //   .catch((error) => {
+  //     message.error(error.message);
+  //   })
+  //   .finally(() => state.precessing = false)
+};
+
+const onSearch = () => {
+  if (!state.searchInput) {
+    message.error('请输入要查询的实体')
+    return
+  }
+
+  state.searchLoading = true
+  fetch(`/api/database/graph/node?entity_name=${state.searchInput}`)
+    .then((res) => {
+      if (res.ok) {
+        return res.json();
+      } else {
+        throw new Error("查询失败");
+      }
+    })
+    .then((data) => {
+      subgraph.nodes = data.result.nodes
+      subgraph.edges = data.result.edges
+      console.log(data)
+      console.log(subgraph)
+      randerGraph()
+    })
+    .catch((error) => {
+      message.error(error.message);
+    })
+    .finally(() => state.searchLoading = false)
+};
+
+const randerGraph = () => {
+  graphInstance.setData(graphData.value);
+  graphInstance.render();
+}
+
+
+
 
 onMounted(() => {
-  const graph = new Graph({
+  graphInstance = new Graph({
     container: document.getElementById("container"),
     width: getCurWidth(),
     height: getCurHeight(),
-    data: {
+    autoFit: true,
-      nodes: [
+    autoResize: true,
-        {
+    layout: {
-          id: "node-1",
+      type: 'force-atlas2',
-          style: { x: 50, y: 100 },
+      preventOverlap: true,
-        },
+      kr: 100,
-        {
-          id: "node-2",
-          style: { x: 150, y: 100 },
-        },
-      ],
-      edges: [{ id: "edge-1", source: "node-1", target: "node-2" }],
     },
-    behaviors: ["drag-canvas", "zoom-canvas", "drag-element"],
+    node: {
+      type: 'circle',
+      style: {
+        labelText: (d) => d.data.label,
+        size: 40,
+      },
+      palette: {
+        field: 'label',
+        color: 'tableau',
+      },
+    },
+    edge: {
+      type: 'line',
+      style: {
+        labelText: (d) => d.data.label,
+        labelBackground: '#fff',
+      },
+    },
+    behaviors: ['drag-element'],
   });
-
+  graphInstance.setData(graphData.value);
-  graph.render();
+  graphInstance.render();
+  window.addEventListener('resize', randerGraph);
 });
-</script>
 
-<style scoped>
-.graph-container {}
 
-#container {
+const handleFileUpload = (event) => {
-  width: 100%;
+  console.log(event)
-  height: 100%;
+  console.log(fileList.value)
 }
 
+const handleDrop = (event) => {
+  console.log(event)
+  console.log(fileList.value)
+}
+
+</script>
+
+<style lang="less" scoped>
+.graph-container {
+  padding: 20px;
+}
+
+.actions {
+  display: flex;
+  justify-content: space-between;
+  margin-bottom: 20px;
+
+  input {
+    margin-right: 10px;
+    border-radius: 8px;
+    padding: 4px 12px;
+    border: 2px solid #d9d9d9;
+    outline: none;
+    height: 42px;
+
+    &:focus {
+      border-color: var(--main-color);
+    }
+  }
+
+  button {
+    height: 40px;
+    box-shadow: none;
+  }
+}
+
+
+.upload {
+  margin-bottom: 20px;
+
+  .upload-dragger {
+    margin: 0px;
+  }
+}
+
+
+#container {
+  background: #F7F7F7;
+  margin: 20px 0;
+  border-radius: 16px;
+  width: 100%;
+  height: 400px;
+}
+
+
+
+
+
 </style>