import json import os import tempfile import traceback import warnings from urllib.parse import urlparse from neo4j import GraphDatabase as GD from src import config from src.models import select_embedding_model from src.storage.minio.client import get_minio_client from src.utils import logger from src.utils.datetime_utils import utc_isoformat warnings.filterwarnings("ignore", category=UserWarning) UIE_MODEL = None class GraphDatabase: def __init__(self): self.driver = None self.files = [] self.status = "closed" self.kgdb_name = "neo4j" self.embed_model_name = os.getenv("GRAPH_EMBED_MODEL_NAME") or "siliconflow/BAAI/bge-m3" self.embed_model = select_embedding_model(self.embed_model_name) self.work_dir = os.path.join(config.save_dir, "knowledge_graph", self.kgdb_name) os.makedirs(self.work_dir, exist_ok=True) # 尝试加载已保存的图数据库信息 if not self.load_graph_info(): logger.debug("创建新的图数据库配置") self.start() def start(self): uri = os.environ.get("NEO4J_URI", "bolt://localhost:7687") username = os.environ.get("NEO4J_USERNAME", "neo4j") password = os.environ.get("NEO4J_PASSWORD", "0123456789") logger.info(f"Connecting to Neo4j: {uri}/{self.kgdb_name}") try: self.driver = GD.driver(f"{uri}/{self.kgdb_name}", auth=(username, password)) self.status = "open" logger.info(f"Connected to Neo4j: {self.get_graph_info(self.kgdb_name)}") # 连接成功后保存图数据库信息 self.save_graph_info(self.kgdb_name) except Exception as e: logger.error(f"Failed to connect to Neo4j: {e}, {uri}, {self.kgdb_name}, {username}, {password}") def close(self): """关闭数据库连接""" assert self.driver is not None, "Database is not connected" self.driver.close() def is_running(self): """检查图数据库是否正在运行""" return self.status == "open" or self.status == "processing" def get_sample_nodes(self, kgdb_name="neo4j", num=50): """获取指定数据库的 num 个节点信息，优先返回连通的节点子图""" assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) def _process_record_props(record): """处理记录中的属性：扁平化 properties 并移除 embedding""" if record is None: return None # 复制一份以避免修改原字典 data = dict(record) props = data.pop("properties", {}) or {} # 移除 embedding if "embedding" in props: del props["embedding"] # 合并属性（优先保留原字典中的 id, name, type 等核心字段） return {**props, **data} def query(tx, num): """Note: 使用连通性查询获取集中的节点子图""" # 首先尝试获取一个连通的子图 query_str = """ // 获取高度数节点作为种子节点 MATCH (seed:Entity) WITH seed, COUNT{(seed)-[]->()} + COUNT{(seed)<-[]-()} as degree WHERE degree > 0 ORDER BY degree DESC LIMIT 5 // 为每个种子节点收集更多邻居节点 UNWIND seed as s MATCH (s)-[*1..1]-(neighbor:Entity) WITH s, neighbor, COUNT{(s)-[]->()} + COUNT{(s)<-[]-()} as s_degree WITH s, s_degree, collect(DISTINCT neighbor) as neighbors // 调整限制比例，允许更多的邻居节点 WITH s, s_degree, neighbors[0..toInteger($num * 0.15)] as limited_neighbors // 从邻居节点扩展到二跳节点，形成开枝散叶结构 UNWIND limited_neighbors as neighbor OPTIONAL MATCH (neighbor)-[*1..1]-(second_hop:Entity) WHERE second_hop <> s // 增加二跳节点的数量 WITH s, limited_neighbors, neighbor, collect(DISTINCT second_hop)[0..5] as second_hops // 收集所有连通节点 WITH collect(DISTINCT s) as seeds, collect(DISTINCT neighbor) as first_hop_nodes, reduce(acc = [], x IN collect(second_hops) | acc + x) as second_hop_nodes WITH seeds + first_hop_nodes + second_hop_nodes as connected_nodes // 确保不会超过请求的节点数量 WITH connected_nodes[0..$num] as final_nodes // 获取这些节点之间的关系，避免双向边 UNWIND final_nodes as n OPTIONAL MATCH (n)-[rel]-(m) WHERE m IN final_nodes AND elementId(n) < elementId(m) RETURN {id: elementId(n), name: n.name, properties: properties(n)} AS h, CASE WHEN rel IS NOT NULL THEN { id: elementId(rel), type: rel.type, source_id: elementId(startNode(rel)), target_id: elementId(endNode(rel)), properties: properties(rel) } ELSE null END AS r, CASE WHEN m IS NOT NULL THEN {id: elementId(m), name: m.name, properties: properties(m)} ELSE null END AS t """ try: results = tx.run(query_str, num=int(num)) formatted_results = {"nodes": [], "edges": []} node_ids = set() for item in results: h_node = _process_record_props(item["h"]) # 始终添加头节点 if h_node["id"] not in node_ids: formatted_results["nodes"].append(h_node) node_ids.add(h_node["id"]) # 只有当边和尾节点都存在时才处理 if item["r"] is not None and item["t"] is not None: t_node = _process_record_props(item["t"]) r_edge = _process_record_props(item["r"]) # 避免重复添加尾节点 if t_node["id"] not in node_ids: formatted_results["nodes"].append(t_node) node_ids.add(t_node["id"]) formatted_results["edges"].append(r_edge) # 如果连通查询返回的节点数不足，补充更多节点 if len(formatted_results["nodes"]) < num: remaining_count = num - len(formatted_results["nodes"]) # 获取额外的节点来补充 supplement_query = """ MATCH (n:Entity) WHERE NOT elementId(n) IN $existing_ids RETURN {id: elementId(n), name: n.name, properties: properties(n)} AS node LIMIT $count """ supplement_results = tx.run(supplement_query, existing_ids=list(node_ids), count=remaining_count) for item in supplement_results: node = _process_record_props(item["node"]) formatted_results["nodes"].append(node) node_ids.add(node["id"]) return formatted_results except Exception as e: # 如果连通查询失败，使用原始查询作为备选 logger.warning(f"Connected subgraph query failed, falling back to simple query: {e}") fallback_query = """ MATCH (n:Entity)-[r]-(m:Entity) WHERE elementId(n) < elementId(m) RETURN {id: elementId(n), name: n.name, properties: properties(n)} AS h, { id: elementId(r), type: r.type, source_id: elementId(startNode(r)), target_id: elementId(endNode(r)), properties: properties(r) } AS r, {id: elementId(m), name: m.name, properties: properties(m)} AS t LIMIT $num """ results = tx.run(fallback_query, num=int(num)) formatted_results = {"nodes": [], "edges": []} node_ids = set() for item in results: h_node = _process_record_props(item["h"]) t_node = _process_record_props(item["t"]) r_edge = _process_record_props(item["r"]) # 避免重复添加节点 if h_node["id"] not in node_ids: formatted_results["nodes"].append(h_node) node_ids.add(h_node["id"]) if t_node["id"] not in node_ids: formatted_results["nodes"].append(t_node) node_ids.add(t_node["id"]) formatted_results["edges"].append(r_edge) return formatted_results with self.driver.session() as session: results = session.execute_read(query, num) return results def create_graph_database(self, kgdb_name): """创建新的数据库，如果已存在则返回已有数据库的名称""" assert self.driver is not None, "Database is not connected" with self.driver.session() as session: existing_databases = session.run("SHOW DATABASES") existing_db_names = [db["name"] for db in existing_databases] if existing_db_names: print(f"已存在数据库: {existing_db_names[0]}") return existing_db_names[0] # 返回所有已有数据库名称 session.run(f"CREATE DATABASE {kgdb_name}") # type: ignore print(f"数据库 '{kgdb_name}' 创建成功.") return kgdb_name # 返回创建的数据库名称 def use_database(self, kgdb_name="neo4j"): """切换到指定数据库""" assert kgdb_name == self.kgdb_name, ( f"传入的数据库名称 '{kgdb_name}' 与当前实例的数据库名称 '{self.kgdb_name}' 不一致" ) if self.status == "closed": self.start() async def txt_add_vector_entity(self, triples, kgdb_name="neo4j"): """添加实体三元组""" assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) def _index_exists(tx, index_name): """检查索引是否存在""" result = tx.run("SHOW INDEXES") for record in result: if record["name"] == index_name: return True return False def _parse_node(node_data): """解析节点数据，返回 (name, props)""" if isinstance(node_data, dict): props = node_data.copy() name = props.pop("name", "") return name, props return str(node_data), {} def _parse_relation(rel_data): """解析关系数据，返回 (type, props)""" if isinstance(rel_data, dict): props = rel_data.copy() rel_type = props.pop("type", "") return rel_type, props return str(rel_data), {} def _create_graph(tx, data): """添加一个三元组""" for entry in data: h_name, h_props = _parse_node(entry.get("h")) t_name, t_props = _parse_node(entry.get("t")) r_type, r_props = _parse_relation(entry.get("r")) if not h_name or not t_name or not r_type: continue tx.run( """ MERGE (h:Entity:Upload {name: $h_name}) SET h += $h_props MERGE (t:Entity:Upload {name: $t_name}) SET t += $t_props MERGE (h)-[r:RELATION {type: $r_type}]->(t) SET r += $r_props """, h_name=h_name, h_props=h_props, t_name=t_name, t_props=t_props, r_type=r_type, r_props=r_props, ) def _create_vector_index(tx, dim): """创建向量索引""" # NOTE 这里是否是会重复构建索引？ index_name = "entityEmbeddings" if not _index_exists(tx, index_name): tx.run(f""" CREATE VECTOR INDEX {index_name} FOR (n: Entity) ON (n.embedding) OPTIONS {{indexConfig: {{ `vector.dimensions`: {dim}, `vector.similarity_function`: 'cosine' }} }}; """) def _get_nodes_without_embedding(tx, entity_names): """获取没有embedding的节点列表""" # 构建参数字典，将列表转换为"param0"、"param1"等键值对形式 params = {f"param{i}": name for i, name in enumerate(entity_names)} if not params: return [] # 构建查询参数列表 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, ) # 判断模型名称是否匹配 self.embed_model_name = self.embed_model_name or config.embed_model cur_embed_info = config.embed_model_names.get(self.embed_model_name) logger.warning(f"embed_model_name={self.embed_model_name}, {cur_embed_info=}") assert self.embed_model_name == config.embed_model or self.embed_model_name is None, ( f"embed_model_name={self.embed_model_name}, {config.embed_model=}" ) with self.driver.session() as session: logger.info(f"Adding entity to {kgdb_name}") 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, getattr(cur_embed_info, "dimension", 1024)) # 收集所有需要处理的实体名称，去重 all_entities = set() for entry in triples: h_name, _ = _parse_node(entry.get("h")) t_name, _ = _parse_node(entry.get("t")) if h_name: all_entities.add(h_name) if t_name: all_entities.add(t_name) all_entities_list = list(all_entities) # 筛选出没有embedding的节点 nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities_list) if not nodes_without_embedding: logger.info("所有实体已有embedding，无需重新计算") return logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities_list)}个实体计算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}/" f"{(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() async def jsonl_file_add_entity(self, file_path, kgdb_name="neo4j"): assert self.driver is not None, "Database is not connected" self.status = "processing" kgdb_name = kgdb_name or "neo4j" self.use_database(kgdb_name) # 切换到指定数据库 logger.info(f"Start adding entity to {kgdb_name} with {file_path}") # 检测 file_path 是否是 URL parsed_url = urlparse(file_path) temp_file_path = None try: if parsed_url.scheme in ("http", "https"): # 如果是 URL logger.info(f"检测到 URL，正在从 MinIO 下载文件: {file_path}") # 从 URL 解析 bucket_name 和 object_name # URL 格式: http://host:port/bucket_name/object_name path_parts = parsed_url.path.lstrip("/").split("/", 1) if len(path_parts) < 2: raise ValueError(f"无法解析 MinIO URL: {file_path}") bucket_name = path_parts[0] object_name = path_parts[1] # 从 MinIO 下载文件 minio_client = get_minio_client() file_data = await minio_client.adownload_file(bucket_name, object_name) # 创建临时文件保存下载的内容 with tempfile.NamedTemporaryFile( mode="w", suffix=".jsonl", delete=False, encoding="utf-8" ) as temp_file: temp_file.write(file_data.decode("utf-8")) temp_file_path = temp_file.name logger.info(f"文件已下载到临时路径: {temp_file_path}") actual_file_path = temp_file_path else: # 本地文件路径 actual_file_path = file_path def read_triples(file_path): with open(file_path, encoding="utf-8") as file: for line in file: if line.strip(): yield json.loads(line.strip()) triples = list(read_triples(actual_file_path)) await self.txt_add_vector_entity(triples, kgdb_name) except Exception as e: logger.error(f"处理文件失败: {e}") raise finally: # 清理临时文件 if temp_file_path and os.path.exists(temp_file_path): try: os.unlink(temp_file_path) logger.info(f"已删除临时文件: {temp_file_path}") except Exception as e: logger.warning(f"删除临时文件失败: {e}") self.status = "open" # 更新并保存图数据库信息 self.save_graph_info() return kgdb_name def delete_entity(self, entity_name=None, kgdb_name="neo4j"): """删除数据库中的指定实体三元组, 参数entity_name为空则删除全部实体""" assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) with self.driver.session() as session: if entity_name: session.execute_write(self._delete_specific_entity, entity_name) else: session.execute_write(self._delete_all_entities) def _delete_specific_entity(self, tx, entity_name): query = """ MATCH (n {name: $entity_name}) DETACH DELETE n """ tx.run(query, entity_name=entity_name) def _delete_all_entities(self, tx): query = """ MATCH (n) DETACH DELETE n """ tx.run(query) def query_node( self, keyword, threshold=0.9, kgdb_name="neo4j", hops=2, max_entities=8, return_format="graph", **kwargs ): """知识图谱查询节点的入口:""" assert self.driver is not None, "Database is not connected" assert self.is_running(), "图数据库未启动" self.use_database(kgdb_name) # 简单空格分词，OR 聚合 tokens = [t for t in str(keyword).split(" ") if t] if not tokens: tokens = [str(keyword)] # name -> score 聚合；向量分数累加，模糊命中给予轻权重 entity_to_score = {} for token in tokens: # 使用向量索引进行查询 results_sim = self._query_with_vector_sim(token, kgdb_name, threshold) for r in results_sim: name = r[0] # 与下方保持统一的 [0] 取 name 的方式 score = 0.0 try: score = float(r["score"]) # neo4j.Record 支持键访问 except Exception: # 兜底：若无法取到score，给个基础分 score = 0.5 entity_to_score[name] = max(entity_to_score.get(name, 0.0), score) # 模糊查询（不区分大小写），命中加一个较小分 results_fuzzy = self._query_with_fuzzy_match(token, kgdb_name) for fr in results_fuzzy: # _query_with_fuzzy_match 返回 values()，形如 [name] name = fr[0] # 给予轻权重，避免覆盖向量高分 entity_to_score[name] = max(entity_to_score.get(name, 0.0), 0.3) # 排序并截断 qualified_entities = [name for name, _ in sorted(entity_to_score.items(), key=lambda x: x[1], reverse=True)][ :max_entities ] logger.debug(f"Graph Query Entities: {keyword}, {qualified_entities=}") # 对每个合格的实体进行查询 all_query_results = {"nodes": [], "edges": [], "triples": []} for entity in qualified_entities: query_result = self._query_specific_entity(entity_name=entity, kgdb_name=kgdb_name, hops=hops) if return_format == "graph": all_query_results["nodes"].extend(query_result["nodes"]) all_query_results["edges"].extend(query_result["edges"]) elif return_format == "triples": all_query_results["triples"].extend(query_result["triples"]) else: raise ValueError(f"Invalid return_format: {return_format}") # 基础去重 if return_format == "graph": seen_node_ids = set() dedup_nodes = [] for n in all_query_results["nodes"]: nid = n.get("id") if isinstance(n, dict) else n if nid not in seen_node_ids: seen_node_ids.add(nid) dedup_nodes.append(n) all_query_results["nodes"] = dedup_nodes seen_edges = set() dedup_edges = [] for e in all_query_results["edges"]: key = (e.get("source_id"), e.get("target_id"), e.get("type")) if key not in seen_edges: seen_edges.add(key) dedup_edges.append(e) all_query_results["edges"] = dedup_edges elif return_format == "triples": seen_triples = set() dedup_triples = [] for t in all_query_results["triples"]: if t not in seen_triples: seen_triples.add(t) dedup_triples.append(t) all_query_results["triples"] = dedup_triples return all_query_results def _query_with_fuzzy_match(self, keyword, kgdb_name="neo4j"): """模糊查询""" assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) def query_fuzzy_match(tx, keyword): result = tx.run( """ MATCH (n:Entity) WHERE toLower(n.name) CONTAINS toLower($keyword) RETURN DISTINCT n.name AS name """, keyword=keyword, ) values = result.values() logger.debug(f"Fuzzy Query Results: {values}") return values with self.driver.session() as session: return session.execute_read(query_fuzzy_match, keyword) def _query_with_vector_sim(self, keyword, kgdb_name="neo4j", threshold=0.9): """向量查询""" assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) def _index_exists(tx, index_name): """检查索引是否存在""" result = tx.run("SHOW INDEXES") for record in result: if record["name"] == index_name: return True return False def query_by_vector(tx, text, threshold): # 首先检查索引是否存在 if not _index_exists(tx, "entityEmbeddings"): raise Exception( "向量索引不存在，请先创建索引，或当前图谱中未上传任何三元组（知识库中自动构建的，不会在此处展示和检索）。" ) embedding = self.get_embedding(text) result = tx.run( """ CALL db.index.vector.queryNodes('entityEmbeddings', 10, $embedding) YIELD node AS similarEntity, score RETURN similarEntity.name AS name, score """, embedding=embedding, ) return [r for r in result if r["score"] > threshold] with self.driver.session() as session: results = session.execute_read(query_by_vector, keyword, threshold=threshold) return results def _query_specific_entity(self, entity_name, kgdb_name="neo4j", hops=2, limit=100): """查询指定实体三元组信息（无向关系）""" assert self.driver is not None, "Database is not connected" if not entity_name: logger.warning("实体名称为空") return [] self.use_database(kgdb_name) def _process_record_props(record): """处理记录中的属性：扁平化 properties 并移除 embedding""" if record is None: return None # 复制一份以避免修改原字典 data = dict(record) props = data.pop("properties", {}) or {} # 移除 embedding if "embedding" in props: del props["embedding"] # 合并属性（优先保留原字典中的 id, name, type 等核心字段） return {**props, **data} def query(tx, entity_name, hops, limit): try: query_str = """ WITH [ // 1跳出边 [(n {name: $entity_name})-[r1]->(m1) | {h: {id: elementId(n), name: n.name, properties: properties(n)}, r: {id: elementId(r1), type: r1.type, source_id: elementId(n), target_id: elementId(m1), properties: properties(r1)}, t: {id: elementId(m1), name: m1.name, properties: properties(m1)}}], // 2跳出边 [(n {name: $entity_name})-[r1]->(m1)-[r2]->(m2) | {h: {id: elementId(m1), name: m1.name, properties: properties(m1)}, r: {id: elementId(r2), type: r2.type, source_id: elementId(m1), target_id: elementId(m2), properties: properties(r2)}, t: {id: elementId(m2), name: m2.name, properties: properties(m2)}}], // 1跳入边 [(m1)-[r1]->(n {name: $entity_name}) | {h: {id: elementId(m1), name: m1.name, properties: properties(m1)}, r: {id: elementId(r1), type: r1.type, source_id: elementId(m1), target_id: elementId(n), properties: properties(r1)}, t: {id: elementId(n), name: n.name, properties: properties(n)}}], // 2跳入边 [(m2)-[r2]->(m1)-[r1]->(n {name: $entity_name}) | {h: {id: elementId(m2), name: m2.name, properties: properties(m2)}, r: {id: elementId(r2), type: r2.type, source_id: elementId(m2), target_id: elementId(m1), properties: properties(r2)}, t: {id: elementId(m1), name: m1.name, properties: properties(m1)}}] ] AS all_results UNWIND all_results AS result_list UNWIND result_list AS item RETURN item.h AS h, item.r AS r, item.t AS t LIMIT $limit """ results = tx.run(query_str, entity_name=entity_name, limit=limit) if not results: logger.info(f"未找到实体 {entity_name} 的相关信息") return {} formatted_results = {"nodes": [], "edges": [], "triples": []} for item in results: h = _process_record_props(item["h"]) r = _process_record_props(item["r"]) t = _process_record_props(item["t"]) formatted_results["nodes"].extend([h, t]) formatted_results["edges"].append(r) formatted_results["triples"].append((h["name"], r["type"], t["name"])) logger.debug(f"Query Results: {results}") return formatted_results except Exception as e: logger.error(f"查询实体 {entity_name} 失败: {str(e)}") return [] try: with self.driver.session() as session: return session.execute_read(query, entity_name, hops, limit) except Exception as e: logger.error(f"数据库会话异常: {str(e)}") return [] async def aget_embedding(self, text): if isinstance(text, list): outputs = await self.embed_model.abatch_encode(text, batch_size=40) return outputs else: outputs = await self.embed_model.aencode(text) return outputs def get_embedding(self, text): if isinstance(text, list): outputs = self.embed_model.batch_encode(text, batch_size=40) return outputs else: outputs = self.embed_model.encode([text])[0] return outputs 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 get_graph_info(self, graph_name="neo4j"): assert self.driver is not None, "Database is not connected" self.use_database(graph_name) def query(tx): # 只统计包含Entity标签的节点 entity_count = tx.run("MATCH (n:Entity) RETURN count(n) AS count").single()["count"] # 只统计包含RELATION标签的关系 relationship_count = tx.run("MATCH ()-[r:RELATION]->() RETURN count(r) AS count").single()["count"] triples_count = tx.run("MATCH (n:Entity)-[r:RELATION]->(m:Entity) RETURN count(n) AS count").single()[ "count" ] # 获取所有标签 labels = tx.run("CALL db.labels() YIELD label RETURN collect(label) AS labels").single()["labels"] return { "graph_name": graph_name, "entity_count": entity_count, "relationship_count": relationship_count, "triples_count": triples_count, "labels": labels, "status": self.status, "embed_model_name": self.embed_model_name, "unindexed_node_count": len(self.query_nodes_without_embedding(graph_name)), } try: if self.is_running(): # 获取数据库信息 with self.driver.session() as session: graph_info = session.execute_read(query) # 添加时间戳 graph_info["last_updated"] = utc_isoformat() return graph_info else: logger.warning(f"图数据库未连接或未运行:{self.status=}") return None except Exception as e: logger.error(f"获取图数据库信息失败：{e}, {traceback.format_exc()}") return None def save_graph_info(self, graph_name="neo4j"): """ 将图数据库的基本信息保存到工作目录中的JSON文件 保存的信息包括：数据库名称、状态、嵌入模型名称等 """ try: graph_info = self.get_graph_info(graph_name) if graph_info is None: logger.error("图数据库信息为空，无法保存") return False info_file_path = os.path.join(self.work_dir, "graph_info.json") with open(info_file_path, "w", encoding="utf-8") as f: json.dump(graph_info, f, ensure_ascii=False, indent=2) # logger.info(f"图数据库信息已保存到：{info_file_path}") return True except Exception as e: logger.error(f"保存图数据库信息失败：{e}") return False def query_nodes_without_embedding(self, kgdb_name="neo4j"): """查询没有嵌入向量的节点 Returns: list: 没有嵌入向量的节点列表 """ assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) def query(tx): result = tx.run(""" MATCH (n:Entity) WHERE n.embedding IS NULL RETURN n.name AS name """) return [record["name"] for record in result] with self.driver.session() as session: return session.execute_read(query) def load_graph_info(self): """ 从工作目录中的JSON文件加载图数据库的基本信息 返回True表示加载成功，False表示加载失败 """ try: info_file_path = os.path.join(self.work_dir, "graph_info.json") if not os.path.exists(info_file_path): logger.debug(f"图数据库信息文件不存在：{info_file_path}") return False with open(info_file_path, encoding="utf-8") as f: graph_info = json.load(f) # 更新对象属性 if graph_info.get("embed_model_name"): self.embed_model_name = graph_info["embed_model_name"] # 如果需要，可以加载更多信息 # 注意：这里不更新self.kgdb_name，因为它是在初始化时设置的 logger.info(f"已加载图数据库信息，最后更新时间：{graph_info.get('last_updated')}") return True except Exception as e: logger.error(f"加载图数据库信息失败：{e}") return False def add_embedding_to_nodes(self, node_names=None, kgdb_name="neo4j"): """为节点添加嵌入向量 Args: node_names (list, optional): 要添加嵌入向量的节点名称列表，None表示所有没有嵌入向量的节点 kgdb_name (str, optional): 图数据库名称，默认为'neo4j' Returns: int: 成功添加嵌入向量的节点数量 """ assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) # 如果node_names为None，则获取所有没有嵌入向量的节点 if node_names is None: node_names = self.query_nodes_without_embedding(kgdb_name) count = 0 with self.driver.session() as session: for node_name in node_names: try: embedding = self.get_embedding(node_name) session.execute_write(self.set_embedding, node_name, embedding) count += 1 except Exception as e: logger.error(f"为节点 '{node_name}' 添加嵌入向量失败: {e}, {traceback.format_exc()}") return count def format_general_results(self, results): nodes = [] edges = [] for item in results: nodes.extend([item["h"], item["t"]]) edges.append(item["r"]) formatted_results = {"nodes": nodes, "edges": edges} return formatted_results def _extract_relationship_info(self, relationship, source_name=None, target_name=None, node_dict=None): """ 提取关系信息并返回格式化的节点和边信息 """ rel_id = relationship.element_id nodes = relationship.nodes if len(nodes) != 2: return None, None source, target = nodes source_id = source.element_id target_id = target.element_id # 如果没有提供 source_name 或 target_name，则需要 node_dict if source_name is None or target_name is None: assert node_dict is not None, "node_dict is required when source_name or target_name is None" source_name = node_dict[source_id]["name"] if source_name is None else source_name target_name = node_dict[target_id]["name"] if target_name is None else target_name relationship_type = relationship._properties.get("type", "unknown") if relationship_type == "unknown": relationship_type = relationship.type edge_info = { "id": rel_id, "type": relationship_type, "source_id": source_id, "target_id": target_id, "source_name": source_name, "target_name": target_name, } node_info = [ {"id": source_id, "name": source_name}, {"id": target_id, "name": target_name}, ] return node_info, edge_info def clean_triples_embedding(triples): for item in triples: if hasattr(item[0], "_properties"): item[0]._properties["embedding"] = None if hasattr(item[2], "_properties"): item[2]._properties["embedding"] = None return triples if __name__ == "__main__": pass