From b9197251178840078346fa6c01af342db2d3438b Mon Sep 17 00:00:00 2001 From: Defeng Date: Tue, 7 Jul 2026 16:18:57 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20neo4j=5Findexing.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 检索索引构建冗余代码 --- neo4j_indexing.py | 399 +--------------------------------------------- 1 file changed, 5 insertions(+), 394 deletions(-) diff --git a/neo4j_indexing.py b/neo4j_indexing.py index c490788..46671eb 100644 --- a/neo4j_indexing.py +++ b/neo4j_indexing.py @@ -35,8 +35,6 @@ EMBED_MAX_WORKERS = SEARCH_CONFIG["embed_thread_num"] # embedding 批次并行 NODE_FETCH_PAGE = SEARCH_CONFIG["node_fetch_page"] # 单次拉取待生成 embedding 节点的分页大小 # 全局索引名称（覆盖所有节点标签） -GLOBAL_VECTOR_INDEX_NAME = SEARCH_CONFIG["global_entity_embedding"] -GLOBAL_FULLTEXT_INDEX_NAME = SEARCH_CONFIG["global_entity_content_search"] DEFAULT_URI = NEO4J_CONFIG["uri"] DEFAULT_AUTH = ( NEO4J_CONFIG["username"], @@ -57,8 +55,6 @@ OLD_INDEX_NAMES: List[str] = [ VECTOR_INDEX_NAME, FULLTEXT_INDEX_NAME, FULLTEXT_FIELD_INDEX_NAME, - GLOBAL_VECTOR_INDEX_NAME, - GLOBAL_FULLTEXT_INDEX_NAME, ] @@ -192,390 +188,6 @@ def drop_index_without_constraint(driver): driver.execute_query(f"drop index {record['name']} if exists") -# --------- 创建全局混合索引（覆盖所有节点标签） --------- -def create_global_indexes(driver, force_refresh=False): - """ - 创建覆盖所有节点标签的全局向量索引和全文索引。 - 与按标签创建的索引互补，支持跨标签统一检索。 - - Args: - driver: Neo4j driver - force_refresh: True=先删除旧索引再重建，False=IF NOT EXISTS 跳过已有索引 - """ - if force_refresh: - for index_name in (GLOBAL_VECTOR_INDEX_NAME, GLOBAL_FULLTEXT_INDEX_NAME): - try: - driver.execute_query(f"DROP INDEX {index_name} IF EXISTS") - logger.info(f"强制刷新：已删除全局索引 {index_name}") - except Exception as e: - logger.warning(f"删除全局索引 {index_name} 时出错（可能不存在）: {e}") - - # 全局向量索引（作用于所有节点的 embedding 属性） - try: - driver.execute_query(f""" - CREATE VECTOR INDEX {GLOBAL_VECTOR_INDEX_NAME} IF NOT EXISTS - FOR (n) ON n.embedding - OPTIONS {{indexConfig: {{ - `vector.dimensions`: {vector_dim}, - `vector.similarity_function`: 'cosine' - }}}} - """) - logger.info(f"全局向量索引 '{GLOBAL_VECTOR_INDEX_NAME}' 已就绪") - except Exception as e: - logger.error(f"创建全局向量索引失败: {e}") - raise - - # 全局全文索引（作用于所有节点的 名称 和 内容 属性） - try: - driver.execute_query(f""" - CREATE FULLTEXT INDEX {GLOBAL_FULLTEXT_INDEX_NAME} IF NOT EXISTS - FOR (n) ON EACH [n.名称, n.内容] - """) - logger.info(f"全局全文索引 '{GLOBAL_FULLTEXT_INDEX_NAME}' 已就绪") - except Exception as e: - logger.error(f"创建全局全文索引失败: {e}") - raise - - -# --------- 创建向量索引 --------- -def vector_indexing(driver, label, force_refresh=False): - """ - 创建向量索引（如果节点没有 embedding，使用"名称"属性生成 embedding） - - 说明： - - 待生成 embedding 的节点通过分页循环全部处理，不再受单次 LIMIT 限制， - 避免节点数超过一万时被静默丢弃。 - - embedding 计算采用并行批次以缩短耗时。 - - Args: - driver: Neo4j driver - label: 节点标签 - force_refresh: 是否强制刷新（如果索引已存在，force_refresh=False 时会跳过，True 时会先删除再创建） - """ - - # === 分页处理所有缺失 embedding 的节点 === - # 每轮取一页未生成 embedding 的节点并写回，直到没有剩余节点为止。 - # 这样即使待处理节点远超一万也能全部覆盖。 - total_generated = 0 - while True: - query_no_embedding = f""" - MATCH (n:{label}) - WHERE n.embedding IS NULL - AND n.名称 IS NOT NULL - RETURN elementId(n) AS id, n.名称 AS 名称 - LIMIT {NODE_FETCH_PAGE} - """ - result_no_embedding = driver.execute_query(query_no_embedding) - nodes_to_process = [ - (record["id"], record["名称"]) for record in result_no_embedding.records - ] - - if not nodes_to_process: - break - - logger.info(f"标签 {label} 本轮发现 {len(nodes_to_process)} 个节点需要生成 embedding") - - ids, names = zip(*nodes_to_process) - ids = list(ids) - names = list(names) - - # 并行计算 embedding - logger.info(f"开始为 {len(names)} 个节点计算 embedding（并行批次）...") - try: - all_embeddings = _compute_embeddings_parallel(names) - except Exception as e: - logger.error(f"批量嵌入向量计算失败: {e}") - raise - - # 转换为 numpy 数组并做 L2 归一化 - embeddings = np.array(all_embeddings) - norms = np.linalg.norm(embeddings, axis=1, keepdims=True) - norms = np.where(norms == 0, 1, norms) # 避免除零 - embeddings = embeddings / norms - - # 保存 embedding 到节点 - logger.info(f"写入 {len(ids)} 个节点的 embedding...") - upsert_vectors( - driver, - ids=ids, - embedding_property="embedding", - embeddings=embeddings, - ) - total_generated += len(ids) - logger.info(f"本轮成功为 {len(ids)} 个节点生成并保存 embedding") - - # 不足一页说明已处理完，提前结束 - if len(nodes_to_process) < NODE_FETCH_PAGE: - break - - if total_generated: - logger.info(f"标签 {label} 共生成 {total_generated} 个节点的 embedding") - - # 检查是否有节点存在 embedding（包括刚生成的） - check_query = f"MATCH (n:{label}) WHERE n.embedding IS NOT NULL RETURN count(n) AS count" - result = driver.execute_query(check_query) - node_count = result.records[0]["count"] if result.records else 0 - - if node_count == 0: - logger.warning(f"标签 {label} 没有节点包含 embedding 属性，跳过向量索引创建") - return - - # 创建向量索引 - index_name = f"{label.lower()}_vector" - - # 如果强制刷新且索引已存在，先删除它 - if force_refresh: - try: - driver.execute_query(f"DROP INDEX {index_name} IF EXISTS") - logger.info(f"强制刷新模式：已删除旧索引 {index_name}") - except Exception as e: - logger.warning(f"删除索引 {index_name} 时出错（可能不存在）: {e}") - - try: - create_vector_index( - driver, - name=index_name, # 索引的唯一名称 - label=label, # 要索引的节点标签 - embedding_property="embedding", # 包含嵌入向量值的节点属性键 - dimensions=vector_dim, # 向量嵌入维度，1024与使用的 bge-m3 嵌入模型一致 - similarity_fn="cosine", # 向量相似度函数，可选 "euclidean" 或 "cosine" - ) - logger.info(f"成功为标签 {label} 创建向量索引（共 {node_count} 个节点）") - except Exception as e: - # 如果是增量更新模式且索引已存在，这是正常的，只记录警告 - if not force_refresh and "already exists" in str(e).lower(): - logger.info(f"标签 {label} 的向量索引已存在，跳过创建") - else: - logger.error(f"为标签 {label} 创建向量索引失败: {e}") - raise - - -# --------- 创建全文索引 --------- -def fulltext_indexing(driver, label, force_refresh=False): - """ - 创建全文索引，并添加节点属性（自动获取所有属性） - - 全文索引：词 ---》 id - 后续检索，分词后，根据词反向找到节点 - - Args: - driver: Neo4j driver - label: 节点标签 - force_refresh: 是否强制刷新所有节点的全文索引（True=清空后重新生成，False=只处理缺失的节点） - """ - # 排除内部属性 - excluded_props = {"embedding", "fulltext", "elementId", "id"} - - # 根据 force_refresh 决定查询条件 - if force_refresh: - # 强制刷新：处理所有节点（包括已有 fulltext 的节点） - query = f"""MATCH (n:{label}) - WITH n, elementId(n) AS id, properties(n) AS props - RETURN id, props""" - else: - # 增量更新：只处理没有 fulltext 的节点 - query = f"""MATCH (n:{label}) - WHERE n.fulltext IS NULL - WITH n, elementId(n) AS id, properties(n) AS props - RETURN id, props""" - - # 执行查询 - records = driver.execute_query(query).records - - if not records: - if force_refresh: - logger.info(f"{label} 没有节点需要处理") - else: - logger.info(f"{label} 所有节点皆存在全文索引属性") - return - - # 构建文本描述（使用所有属性） - record_list = [] - for r in records: - node_id = r["id"] - props = r["props"] - # 过滤掉内部属性 - filtered_props = {k: v for k, v in props.items() if k not in excluded_props} - # 使用所有属性构建文本 - text = build_text_from_all_properties(filtered_props) - if text: - record_list.append({"id": node_id, "text": text}) - - if not record_list: - logger.info(f"{label} 没有可索引的文本内容") - return - - # 创建全文索引（在写入数据之前创建索引） - try: - create_fulltext_index( - driver, - name=f"{label.lower()}_fulltext", # 索引的唯一名称 - label=label, # 要创建索引的节点标签 - node_properties=["fulltext"], # 要创建全文索引的节点属性列表 - ) - except Exception as e: - logger.warning(f"创建全文索引可能已存在（继续处理）: {e}") - - # 文本分词，作为全文索引属性 - logger.info(f"计算 {label} ({len(record_list)}) 的全文索引") - pattern = re.compile(r"[a-zA-Z0-9\u4e00-\u9fa5]+") # 匹配英文字母、数字和中文字符 - fulltext_tuple_list = [] - for id_, text in [(r["id"], r["text"]) for r in record_list]: - # 使用 jieba 分词 - words = jieba.lcut(text) - # 过滤并连接 - filtered_words = [ - word.strip() - for word in words - if word.strip() and pattern.fullmatch(word.strip()) - ] - fulltext_value = " ".join(filtered_words) - if fulltext_value: - fulltext_tuple_list.append((id_, fulltext_value)) - - if not fulltext_tuple_list: - logger.warning(f"{label} 分词后没有有效内容") - return - - ids, fulltexts = zip(*fulltext_tuple_list) - ids = list(ids) - fulltexts = list(fulltexts) - - # 按 elementId 添加全文索引属性 - logger.info(f"写入 {label} ({len(fulltexts)}) 的全文索引") - insert_batch_size = 1000 - for i in range(0, len(fulltext_tuple_list), insert_batch_size): - batch_rows = [ - {"id": id_, "fulltext": ft} - for id_, ft in zip( - ids[i: i + insert_batch_size], - fulltexts[i: i + insert_batch_size], - ) - ] - - # UNWIND：将列表数据展开为多行记录 - driver.execute_query( - "UNWIND $rows AS row " # 将传入的 rows 列表展开，每项作为一行数据，命名为 row - "MATCH (n) " - "WHERE elementId(n) = row.id " - "SET n.fulltext = row.fulltext ", - {"rows": batch_rows}, - ) - - -def create_all_indexes(neo4j_uri=None, neo4j_username=None, neo4j_password=None, force_refresh=False, driver=None): - """ - 创建所有索引的主函数（可被外部调用） - - 注意：Neo4j 连接配置优先从 .env 文件中读取（NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD） - 如果 .env 文件中没有配置，则使用参数中的值（如果提供）或默认值 - - Args: - neo4j_uri: Neo4j URI（可选，优先级低于 .env，提供 driver 时忽略） - neo4j_username: Neo4j 用户名（可选，优先级低于 .env，提供 driver 时忽略） - neo4j_password: Neo4j 密码（可选，优先级低于 .env，提供 driver 时忽略） - force_refresh: 是否强制刷新所有索引（True=清空重建，False=增量更新） - driver: 可选的 Neo4j driver 实例，提供则复用外部 driver - - Returns: - dict: 包含执行结果的字典 - """ - use_external_driver = driver is not None - - if not use_external_driver: - # 优先从 .env 文件读取配置，如果 .env 中没有则使用参数或默认值 - neo4j_uri = os.getenv("NEO4J_URI") or neo4j_uri or "bolt://192.168.0.46:57687" - neo4j_username = os.getenv("NEO4J_USERNAME") or neo4j_username or "neo4j" - neo4j_password = os.getenv("NEO4J_PASSWORD") or neo4j_password or "zdht123@" - logger.info(f"连接到 Neo4j: {neo4j_uri}") - else: - logger.info("使用外部提供的 Neo4j driver") - - # 定义内部函数来执行索引创建逻辑（避免代码重复） - def _execute_indexing(driver): - # 1、根据 force_refresh 决定是否清空索引和约束 - if force_refresh: - logger.info("强制刷新模式：清空所有约束...") - drop_constraint(driver) - logger.info("强制刷新模式：清空所有索引...") - drop_index_without_constraint(driver) - logger.info("强制刷新模式：清空所有节点的 fulltext 属性...") - driver.execute_query("MATCH (n) WHERE n.fulltext IS NOT NULL REMOVE n.fulltext") - else: - logger.info("增量更新模式：保留现有索引和约束") - - # 2、自动获取所有节点标签并创建向量索引 - logger.info("=" * 50) - logger.info("开始创建向量索引...") - logger.info("=" * 50) - - all_labels = get_all_labels(driver) - logger.info(f"发现 {len(all_labels)} 个节点标签: {all_labels}") - - vector_index_errors = [] - for label in all_labels: - logger.info(f"处理标签: {label}") - try: - vector_indexing(driver, label, force_refresh=force_refresh) - except Exception as e: - error_msg = f"为标签 {label} 创建向量索引时出错: {e}" - logger.error(error_msg) - vector_index_errors.append(error_msg) - continue - - # 3、创建全文索引（自动获取所有属性） - logger.info("=" * 50) - logger.info("开始创建全文索引...") - logger.info("=" * 50) - - fulltext_index_errors = [] - for label in all_labels: - logger.info(f"处理标签: {label}") - try: - fulltext_indexing(driver, label, force_refresh=force_refresh) - except Exception as e: - error_msg = f"为标签 {label} 创建全文索引时出错: {e}" - logger.error(error_msg) - fulltext_index_errors.append(error_msg) - continue - - logger.info("=" * 50) - logger.info("所有索引创建完成！") - logger.info("=" * 50) - - has_errors = len(vector_index_errors) > 0 or len(fulltext_index_errors) > 0 - mode_text = "强制刷新" if force_refresh else "增量更新" - return { - "success": not has_errors, - "message": f"{mode_text}模式：索引创建完成" if not has_errors - else f"{mode_text}模式：索引创建完成，但部分标签出现错误", - "mode": "force_refresh" if force_refresh else "incremental", - "labels": all_labels, - "vector_index_errors": vector_index_errors, - "fulltext_index_errors": fulltext_index_errors, - } - - try: - if use_external_driver: - # 使用外部 driver，直接执行，不关闭 driver - return _execute_indexing(driver) - else: - # 使用内部 driver，使用上下文管理器自动关闭 - with GraphDatabase.driver(neo4j_uri, auth=(neo4j_username, neo4j_password)) as driver: - return _execute_indexing(driver) - except Exception as e: - error_msg = f"创建索引过程中发生错误: {e}" - logger.error(error_msg, exc_info=True) - return { - "success": False, - "message": error_msg, - "mode": "error", - "labels": [], - "vector_index_errors": [], - "fulltext_index_errors": [], - } - - def print_neo4j_version(driver) -> None: """打印当前 Neo4j 实例的版本和版本类型。""" with driver.session() as session: @@ -795,13 +407,12 @@ build_hyrid_indexes = build_hybrid_indexes if __name__ == "__main__": - # 可以通过命令行参数或环境变量控制是否强制刷新 - import sys - force_refresh = "--force" in sys.argv or os.getenv("FORCE_REFRESH", "false").lower() == "true" + result = build_hybrid_indexes( + drop_old=True + ) - result = create_all_indexes(force_refresh=force_refresh) if result["success"]: - logger.info(f"索引创建成功完成（模式: {result['mode']}）") + logger.info("混合索引创建成功") else: - logger.error(f"索引创建完成但有错误: {result['message']}") + logger.error(result["message"]) \ No newline at end of file