diff --git a/server/routers/chat_router.py b/server/routers/chat_router.py index 0d8058a3..9c5fb82d 100644 --- a/server/routers/chat_router.py +++ b/server/routers/chat_router.py @@ -21,7 +21,7 @@ async def chat_get(): return "Chat Get!" @chat.post("/") -def chat_post( +async def chat_post( query: str = Body(...), meta: dict = Body(None), history: list[dict] | None = Body(None), diff --git a/server/routers/data_router.py b/server/routers/data_router.py index 275237f8..699ed968 100644 --- a/server/routers/data_router.py +++ b/server/routers/data_router.py @@ -171,11 +171,15 @@ async def get_graph_nodes(kgdb_name: str, num: int): @data.post("/graph/add-by-jsonl") async def add_graph_entity(file_path: str = Body(...), kgdb_name: Optional[str] = Body(None)): if not config.enable_knowledge_graph: - raise HTTPException(status_code=400, detail="Knowledge graph is not enabled") + return {"message": "知识图谱未启用", "status": "failed"} if not file_path.endswith('.jsonl'): - raise HTTPException(status_code=400, detail="file_path must be a jsonl file") + return {"message": "文件格式错误，请上传jsonl文件", "status": "failed"} - graph_base.jsonl_file_add_entity(file_path, kgdb_name) - return {"message": "Entity successfully added"} + try: + await graph_base.jsonl_file_add_entity(file_path, kgdb_name) + return {"message": "实体添加成功", "status": "success"} + except Exception as e: + logger.error(f"添加实体失败: {e}, {traceback.format_exc()}") + return {"message": f"添加实体失败: {e}", "status": "failed"} diff --git a/src/core/graphbase.py b/src/core/graphbase.py index f8274568..3862312f 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -106,7 +106,7 @@ class GraphDatabase: with self.driver.session() as session: session.execute_write(create, triples) - def txt_add_vector_entity(self, triples, kgdb_name='neo4j'): + async def txt_add_vector_entity(self, triples, kgdb_name='neo4j'): """添加实体三元组""" self.use_database(kgdb_name) def _index_exists(tx, index_name): @@ -140,6 +140,31 @@ class GraphDatabase: }} }}; """) + def _get_nodes_without_embedding(tx, entity_names): + """获取没有embedding的节点列表""" + # 构建参数字典，将列表转换为"param0"、"param1"等键值对形式 + params = {f"param{i}": name for i, name in enumerate(entity_names)} + + # 构建查询参数列表 + param_placeholders = ", ".join([f"${key}" for key in params.keys()]) + + # 执行查询 + result = tx.run(f""" + MATCH (n:Entity) + WHERE n.name IN [{param_placeholders}] AND n.embedding IS NULL + RETURN n.name AS name + """, params) + + return [record["name"] for record in result] + + def _batch_set_embeddings(tx, entity_embedding_pairs): + """批量设置实体的嵌入向量""" + for entity_name, embedding in entity_embedding_pairs: + tx.run(""" + MATCH (e:Entity {name: $name}) + CALL db.create.setNodeVectorProperty(e, 'embedding', $embedding) + """, name=entity_name, embedding=embedding) + # 判断模型名称是否匹配 cur_embed_info = config.embed_model_names[config.embed_model] self.embed_model_name = self.embed_model_name or cur_embed_info.get('name') @@ -151,18 +176,44 @@ class GraphDatabase: session.execute_write(_create_graph, triples) logger.info(f"Creating vector index for {kgdb_name} with {config.embed_model}") session.execute_write(_create_vector_index, cur_embed_info['dimension']) - # NOTE 这里需要异步处理 - for i, entry in enumerate(triples): - logger.debug(f"Adding entity {i+1}/{len(triples)}") - embedding_h = self.get_embedding(entry['h']) - embedding_t = self.get_embedding(entry['t']) - session.execute_write(self.set_embedding, entry['h'], embedding_h) - session.execute_write(self.set_embedding, entry['t'], embedding_t) + + # 收集所有需要处理的实体名称，去重 + all_entities = [] + for entry in triples: + if entry['h'] not in all_entities: + all_entities.append(entry['h']) + if entry['t'] not in all_entities: + all_entities.append(entry['t']) + + # 筛选出没有embedding的节点 + nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities) + if not nodes_without_embedding: + logger.info(f"所有实体已有embedding，无需重新计算") + return + + logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities)}个实体计算embedding") + + # 批量处理实体 + max_batch_size = 1024 # 限制此部分的主要是内存大小 1024 * 1024 * 4 / 1024 / 1024 = 4GB + total_entities = len(nodes_without_embedding) + + for i in range(0, total_entities, max_batch_size): + batch_entities = nodes_without_embedding[i:i+max_batch_size] + logger.debug(f"Processing entities batch {i//max_batch_size + 1}/{(total_entities-1)//max_batch_size + 1} ({len(batch_entities)} entities)") + + # 批量获取嵌入向量 + batch_embeddings = await self.aget_embedding(batch_entities) + + # 将实体名称和嵌入向量配对 + entity_embedding_pairs = list(zip(batch_entities, batch_embeddings)) + + # 批量写入数据库 + session.execute_write(_batch_set_embeddings, entity_embedding_pairs) # 数据添加完成后保存图信息 self.save_graph_info() - def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'): + async def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'): self.status = "processing" kgdb_name = kgdb_name or 'neo4j' self.use_database(kgdb_name) # 切换到指定数据库 @@ -176,7 +227,7 @@ class GraphDatabase: triples = list(read_triples(file_path)) - self.txt_add_vector_entity(triples, kgdb_name) + await self.txt_add_vector_entity(triples, kgdb_name) self.status = "open" # 更新并保存图数据库信息 @@ -355,9 +406,23 @@ class GraphDatabase: with self.driver.session() as session: return session.execute_read(query, node_name, hops) + async def aget_embedding(self, text): + from src import knowledge_base + + if isinstance(text, list): + outputs = await knowledge_base.embed_model.abatch_encode(text, batch_size=40) + return outputs + else: + outputs = await knowledge_base.embed_model.aencode(text) + return outputs + def get_embedding(self, text): - with torch.no_grad(): - from src import knowledge_base + from src import knowledge_base + + if isinstance(text, list): + outputs = knowledge_base.embed_model.batch_encode(text, batch_size=40) + return outputs + else: outputs = knowledge_base.embed_model.encode([text])[0] return outputs diff --git a/src/models/embedding.py b/src/models/embedding.py index 7ea83dd6..59a1f4bd 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -1,6 +1,7 @@ import os import json import requests +import asyncio from FlagEmbedding import FlagModel from zhipuai import ZhipuAI @@ -26,6 +27,15 @@ class BaseEmbeddingModel: def encode_queries(self, queries): return self.predict(queries) + async def aencode(self, message): + return await asyncio.to_thread(self.encode, message) + + async def aencode_queries(self, queries): + return await asyncio.to_thread(self.encode_queries, queries) + + async def abatch_encode(self, messages, batch_size=20): + return await asyncio.to_thread(self.batch_encode, messages, batch_size) + def batch_encode(self, messages, batch_size=20): logger.info(f"Batch encoding {len(messages)} messages") data = [] diff --git a/web/src/views/GraphView.vue b/web/src/views/GraphView.vue index b8fdecb4..08f84e44 100644 --- a/web/src/views/GraphView.vue +++ b/web/src/views/GraphView.vue @@ -198,10 +198,15 @@ const addDocumentByFile = () => { }) .then(response => response.json()) .then((data) => { - message.success(data.message); - state.showModal = false; + if (data.status === 'success') { + message.success(data.message); + state.showModal = false; + } else { + throw new Error(data.message); + } }) .catch((error) => { + console.error(error) message.error(error.message); }) .finally(() => state.precessing = false)