Compare commits
15 Commits
main
...
defeng-pat
| Author | SHA1 | Date | |
|---|---|---|---|
| 7b3df5b161 | |||
| 49747f42ec | |||
| ced323c740 | |||
| 3ba17b3d97 | |||
| b107830d76 | |||
| 71c0938ecb | |||
| 5cb1690b35 | |||
| 45e9eb3a19 | |||
| 80f0cefddb | |||
| b0c4dc6fb5 | |||
| ca08759825 | |||
| 2ceb6dbaf5 | |||
| 84fa3fde1f | |||
| 80c3344f83 | |||
| b919725117 |
File diff suppressed because one or more lines are too long
Binary file not shown.
Binary file not shown.
Binary file not shown.
@ -898,7 +898,7 @@ def find_ship_info_by_hull(json_file_path, data):
|
||||
except json.JSONDecodeError:
|
||||
print("错误:JSON 文件格式不正确")
|
||||
return None
|
||||
def merge_short_slices(slices, min_length=30,filename="163-06A0014-B01001_发动机-维修手册.pdf"):
|
||||
def merge_short_slices(slices, min_length=30,filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
|
||||
"""
|
||||
合并过短的切片:
|
||||
- 如果某切片 content 长度 <= min_length,则将其合并到下一个切片的开头
|
||||
@ -908,6 +908,8 @@ def merge_short_slices(slices, min_length=30,filename="163-06A0014-B01001_发动
|
||||
Args:
|
||||
slices: [{"content": str, "positions": [...]}]
|
||||
min_length: 短切片的字符长度阈值(含)
|
||||
filename: 用于实体提取的文件名
|
||||
|
||||
Returns:
|
||||
合并后的 slices 列表
|
||||
"""
|
||||
@ -934,6 +936,7 @@ def merge_short_slices(slices, min_length=30,filename="163-06A0014-B01001_发动
|
||||
positions = pending_positions + positions
|
||||
pending_contents = []
|
||||
pending_positions = []
|
||||
|
||||
result.append({
|
||||
"content": content,
|
||||
"positions": positions
|
||||
@ -953,35 +956,43 @@ def merge_short_slices(slices, min_length=30,filename="163-06A0014-B01001_发动
|
||||
"content": merged_suffix,
|
||||
"positions": pending_positions
|
||||
})
|
||||
|
||||
try:
|
||||
final_result = get_entity(filename)
|
||||
print(final_result)
|
||||
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME,final_result)
|
||||
if xinghao.get('model_name'):
|
||||
model_name = xinghao['model_name']
|
||||
else:
|
||||
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
|
||||
if xinghao is None: # 确保 xinghao 为 None 时不会报错
|
||||
xinghao = {"model_name": "", "ship_name": ""}
|
||||
print("未找到匹配的舰船信息")
|
||||
|
||||
# 安全处理 xinghao 为 None 的情况
|
||||
model_name = ""
|
||||
# lower_entity = final_result.get("low_level")
|
||||
otehr_data = format_entity_text(final_result)
|
||||
logger.info(f"文件名实体提取成功: {otehr_data}")
|
||||
if xinghao and 'model_name' in xinghao:
|
||||
model_name = xinghao['model_name']
|
||||
|
||||
other_data = format_entity_text(final_result)
|
||||
logger.info(f"文件名实体提取成功: {other_data}")
|
||||
except Exception as e:
|
||||
logger.warning(f"文件名实体提取失败: {e}")
|
||||
otehr_data = "" # 建议加个兜底,避免后面 NameError
|
||||
other_data = ""
|
||||
model_name = "" # 确保异常时 model_name 有定义
|
||||
|
||||
# 将提取的信息(如舰艇名、型号)注入到每个切片的开头
|
||||
for ins in result:
|
||||
content = ins['content']
|
||||
if not otehr_data:
|
||||
continue
|
||||
# 如果没有提取到有效数据,则跳过注入
|
||||
|
||||
newline_idx = content.find('\n')
|
||||
# 拼接成一个括号:otehr_data + 型号为xxx
|
||||
info = otehr_data
|
||||
info = other_data
|
||||
if model_name:
|
||||
info += f',型号为{model_name}'
|
||||
suffix = f'({info})'
|
||||
|
||||
# 将信息插入到第一行末尾
|
||||
if newline_idx == -1:
|
||||
ins['content'] = content + suffix
|
||||
else:
|
||||
ins['content'] = content[:newline_idx] + suffix + content[newline_idx:]
|
||||
|
||||
return result
|
||||
def find_and_read_content_list(directory, original_filename, encoding='utf-8'):
|
||||
"""
|
||||
@ -1133,8 +1144,3 @@ if __name__ == "__main__":
|
||||
|
||||
for ins in slices_2:
|
||||
print(ins)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
85
config.py
85
config.py
@ -17,6 +17,7 @@ EMBEDDING_CONFIG = {
|
||||
"model": "bge-m3",
|
||||
"base_url": "http://192.168.0.46:59700/v1/embeddings",
|
||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
||||
"base_url_v2" : "http://192.168.0.46:59700/v1",
|
||||
}
|
||||
|
||||
# ==================== Neo4j 图数据库配置 ====================
|
||||
@ -75,3 +76,87 @@ API_OTHER_URLS = [
|
||||
]
|
||||
|
||||
SHIP_MODEL_NAME = "E:\ZKYNLP\Hjunproject\project0506\kgrag\ship_xinghao.json"
|
||||
TREE_JSON_PATH = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\tree_data.json"
|
||||
root@70964b104c5b:/app# pwd
|
||||
/app
|
||||
root@70964b104c5b:/app# vim config.py
|
||||
root@70964b104c5b:/app# cat config.py
|
||||
# app/model_config.py
|
||||
|
||||
import os
|
||||
|
||||
# ==================== 大模型配置 ====================
|
||||
LLM_CONFIG = {
|
||||
"model": "46-qwen3.5-35B",
|
||||
"base_url": "http://192.168.0.46:59800/v1",
|
||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 55096,
|
||||
"timeout": 30
|
||||
}
|
||||
|
||||
# ==================== 嵌入模型配置 ====================
|
||||
EMBEDDING_CONFIG = {
|
||||
"model": "bge-m3",
|
||||
"base_url": "http://192.168.0.46:59700/v1/embeddings",
|
||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
||||
"base_url_v2" : "http://192.168.0.46:59700/v1",
|
||||
}
|
||||
|
||||
# ==================== Neo4j 图数据库配置 ====================
|
||||
NEO4J_CONFIG = {
|
||||
# 优先读取环境变量,如果没有则使用默认值
|
||||
# 注意:本地开发通常用 localhost,部署时用服务名 neo4j
|
||||
"uri": "neo4j://192.168.0.46:57687",
|
||||
"username": "neo4j",
|
||||
"password": "zdht123@",
|
||||
"database": "neo4j", # 新增 database 配置
|
||||
}
|
||||
|
||||
# ==================== 检索与索引配置 ====================
|
||||
SEARCH_CONFIG = {
|
||||
"name_property": "名称",
|
||||
"fulltext_property": "fulltext",
|
||||
"embedding_property": "embedding",
|
||||
"search_label": "Searchable",
|
||||
"vector_index_name": "global_searchable_embedding",
|
||||
"fulltext_index_name": "global_searchable_content_search",
|
||||
"fulltext_field_index_name": "global_searchable_fulltext_search",
|
||||
"global_entity_embedding" : "global_entity_embedding",
|
||||
"global_entity_content_search": "global_entity_content_search",
|
||||
"excluded_business_labels": [
|
||||
"Entity", "Chunk", "Document", "_Bloom_Perspective_", "Searchable"
|
||||
],
|
||||
"vector_dimension": 1024,
|
||||
"jieba_pos_whitelist": ["n", "nr", "ns", "nt", "nz", "vn", "eng"],
|
||||
"embed_batch_size" : 64,
|
||||
"embed_thread_num" : 4,
|
||||
"node_fetch_page" : 10000,
|
||||
}
|
||||
|
||||
INDEXING_ULR = "http://192.168.0.46:59085/neo4j_indexing"
|
||||
|
||||
URL = {
|
||||
"indexing_url" : "http://192.168.0.46:59085/neo4j_indexing",
|
||||
"prefix_url" : "http://192.168.0.46:59085",
|
||||
}
|
||||
|
||||
API_URLS = [
|
||||
"http://192.168.0.64:59988/analyze-pdf",
|
||||
# "http://192.168.0.111:9977/analyze-pdf",
|
||||
# "http://192.168.0.111:9978/analyze-pdf",
|
||||
#"http://192.168.0.111:9979/analyze-pdf",
|
||||
#"http://192.168.0.111:9980/analyze-pdf",
|
||||
# "http://192.168.0.111:9975/analyze-pdf",
|
||||
]
|
||||
API_OTHER_URLS = [
|
||||
"http://192.168.0.46:59988/analyze-otherfile",
|
||||
# "http://192.168.0.111:9977/analyze-otherfile",
|
||||
# "http://192.168.0.111:9978/analyze-otherfile",
|
||||
#"http://192.168.0.111:9979/analyze-otherfile",
|
||||
#"http://192.168.0.111:9980/analyze-otherfile",
|
||||
# "http://192.168.0.111:9975/analyze-otherfile",
|
||||
]
|
||||
|
||||
SHIP_MODEL_NAME = "/app/ship_xinghao.json"
|
||||
TREE_JSON_PATH = "/app/tree_data.json"
|
||||
@ -13,6 +13,7 @@ import json
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
from openai import OpenAI
|
||||
from typing import Dict, List, Any
|
||||
from config import TREE_JSON_PATH
|
||||
|
||||
app = FastAPI(title="PDF Upload Service")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
|
||||
@ -540,8 +541,8 @@ def fetch_graph_sample1(driver):
|
||||
# 文件名处理方法
|
||||
# =============================
|
||||
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
||||
# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
# TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
||||
|
||||
|
||||
def _load_tree_data() -> Dict:
|
||||
|
||||
@ -117,9 +117,10 @@ class NodeRetrieval:
|
||||
tasks = []
|
||||
for label, entity, query_text, query_vector in zip(labels, entities, query_texts, query_vectors):
|
||||
try:
|
||||
vector_index_name = label.lower() + "_vector"
|
||||
fulltext_index_name = label.lower() + "_fulltext"
|
||||
|
||||
# vector_index_name = label.lower() + "_vector"
|
||||
# fulltext_index_name = label.lower() + "_fulltext"
|
||||
vector_index_name = "global_searchable_embedding"
|
||||
fulltext_index_name = "global_searchable_fulltext_search"
|
||||
retriever = HybridRetriever(
|
||||
self.driver,
|
||||
vector_index_name=vector_index_name,
|
||||
|
||||
@ -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"])
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
Loading…
x
Reference in New Issue
Block a user