update graph page
This commit is contained in:
Wenjie Zhang
parent
9bdae3e164
commit
4e1be33ebb
@ -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}")
|
||||
|
||||
|
||||
@ -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()
|
||||
triples_path = oneke.processing_text_to_kg(text_path, output_path)
|
||||
global UIE_MODEL
|
||||
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))
|
||||
|
||||
|
||||
|
||||
@ -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']:
|
||||
return jsonify({'message': 'Unsupported file type'}), 400
|
||||
# elif file.filename.split('.')[-1] not in ['pdf', 'txt', 'md']:
|
||||
# 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
|
||||
|
||||
@ -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) {
|
||||
appendAiMessage(data.response, data.refs)
|
||||
cur_res_id = conv.value.messages[conv.value.messages.length - 1].id
|
||||
} else {
|
||||
updateMessage(data.response, cur_res_id)
|
||||
}
|
||||
updateMessage(data.response, cur_res_id)
|
||||
conv.value.history = data.history
|
||||
buffer = ''
|
||||
} catch (e) {
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
|
||||
@ -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>
|
||||
|
||||
@ -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: {
|
||||
nodes: [
|
||||
{
|
||||
id: "node-1",
|
||||
style: { x: 50, y: 100 },
|
||||
},
|
||||
{
|
||||
id: "node-2",
|
||||
style: { x: 150, y: 100 },
|
||||
},
|
||||
],
|
||||
edges: [{ id: "edge-1", source: "node-1", target: "node-2" }],
|
||||
autoFit: true,
|
||||
autoResize: true,
|
||||
layout: {
|
||||
type: 'force-atlas2',
|
||||
preventOverlap: true,
|
||||
kr: 100,
|
||||
},
|
||||
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'],
|
||||
});
|
||||
|
||||
graph.render();
|
||||
graphInstance.setData(graphData.value);
|
||||
graphInstance.render();
|
||||
window.addEventListener('resize', randerGraph);
|
||||
});
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
.graph-container {}
|
||||
|
||||
#container {
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
const handleFileUpload = (event) => {
|
||||
console.log(event)
|
||||
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>
|
||||
|
||||
Loading…
Reference in New Issue
Block a user