kgrag/filename_proceess_and_kgquery.py
2026-07-29 18:10:19 +08:00

862 lines
28 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
PDF 用户上传接口FastAPI
- 处理文件名和其对应关系
"""
from typing import Dict, Tuple
import os
import re
from typing import Dict
from fastapi import FastAPI, HTTPException
import logging
import json
from modelsAPI.model_api import OpenaiAPI
from openai import OpenAI
from typing import Dict, List, Any
app = FastAPI(title="PDF Upload Service")
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
logger = logging.getLogger("filename_neo4j_process")
# =========================
# 图方法
# =========================
def node_to_json(node):
props = dict(node)
# ks_raw = props.get("knowledge_source")
# knowledge_source = json.loads(ks_raw) if ks_raw else []
ks_raw = props.get("knowledge_source")
knowledge_source = []
if ks_raw and isinstance(ks_raw, str):
try:
ks_list = json.loads(ks_raw)
if isinstance(ks_list, list):
for ks in ks_list:
if isinstance(ks, dict):
ks.pop("info", None)
knowledge_source.append(ks)
except (json.JSONDecodeError, TypeError):
# 解析失败或类型错误时,保持 knowledge_source 为空列表
pass
name = props.get("名称")
props.pop("embedding", None)
props.pop("knowledge_source", None)
props.pop("name", None) # 现在安全地移除
props.pop("fulltext", None)
if "last_updated" in props:
props["最后更新时间"] = props.pop("last_updated")
if "created_at" in props:
props["创建时间"] = props.pop("created_at")
return {
"id": node.element_id,
"name": name, # 使用提前保存的值
"labels": list(node.labels),
"knowledge_source": knowledge_source,
"properties": props,
}
def rel_to_json(rel):
props = dict(rel)
props.pop("embedding", None)
props.pop("fact_timeline", None)
return {"id": rel.element_id, "label": rel.type, "source": rel.start_node.element_id, "target": rel.end_node.element_id, "properties": props}
def is_simple_path(path):
"""
判断是否为无环路径simple path
"""
seen = set()
for n in path.nodes:
nid = n.element_id
if nid in seen:
return False
seen.add(nid)
return True
def build_tree(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
"""
按照固定层级构建树:舰船(根) -> 舰艇 -> 系统 -> 子系统
不再依赖 links 的连线,直接根据 labels 归类,彻底避免环状数据导致的死循环。
默认只展开第一艘舰艇的系统/子系统,其它舰艇 children 置空(前端可点击按需加载)。
"""
# 1. 辅助函数:判断标签是否包含关键字
def has_label(node, keyword):
labels = node.get("labels", [])
if isinstance(labels, list):
return any(keyword in lbl for lbl in labels)
return False
# 2. 辅助函数:构建标准的树节点格式
def make_node(node_data):
return {
"id": node_data["id"],
"name": node_data.get("name", ""),
"labels": node_data.get("labels", []),
"properties": node_data.get("properties", {}),
"color": node_data.get("color"),
"size": node_data.get("size"),
"children": []
}
# 3. 将节点按层级分类存储到字典中,方便快速查找
fleet_roots = {} # 第零层:舰船/舰号(父根节点)
ships = {} # 第一层:舰艇
systems = {} # 第二层:系统(不含子系统)
sub_systems = {} # 第三层:子系统
for node in nodes:
# 注意:判断逻辑要互斥,防止"子系统"被误判为"系统"
if has_label(node, "舰船") or has_label(node, "舰号"):
fleet_roots[node["id"]] = make_node(node)
elif has_label(node, "舰艇"):
ships[node["id"]] = make_node(node)
elif has_label(node, "子系统"):
sub_systems[node["id"]] = make_node(node)
elif has_label(node, "系统"):
systems[node["id"]] = make_node(node)
else:
pass
# 4. 组装系统 -> 子系统、舰艇 -> 系统 关系(先把完整子树挂好)
for link in links:
if link.get("label") != "包含":
continue
src_id = link["source"]
tgt_id = link["target"]
if src_id in ships and tgt_id in systems:
ships[src_id]["children"].append(systems[tgt_id])
elif src_id in systems and tgt_id in sub_systems:
systems[src_id]["children"].append(sub_systems[tgt_id])
elif src_id in ships and tgt_id in sub_systems:
ships[src_id]["children"].append(sub_systems[tgt_id])
# 5. 所有舰艇 children 全部保留,仅用 collapsed 标记默认折叠状态。
# 前端读取 collapsed=False 的舰艇默认展开,其它折叠;
# 用户点击折叠的舰艇时,前端切换 collapsed 即可展示子树。
# 只保留名称包含 "122" 的舰艇(隐藏 JZ贱 等其它分支)
#sorted_ships = sorted(
# (s for s in ships.values() if "122" in s.get("name", "")),
# key=lambda s: s.get("name", ""),
#)
sorted_ships = sorted(ships.values(), key=lambda s: s.get("name", ""))
for idx, ship in enumerate(sorted_ships):
ship["collapsed"] = (idx != 0)
# 6. 挂载到舰船根节点;如果数据里没有舰船根,就创建一个虚拟根
if fleet_roots:
# 取第一个舰船根节点作为父(一般场景下只有一个)
root = next(iter(fleet_roots.values()))
root["children"] = sorted_ships
return [root]
else:
virtual_root = {
"id": "__fleet_root__",
"name": "舰号",
"labels": ["舰船"],
"properties": {},
"color": None,
"size": None,
"children": sorted_ships,
}
return [virtual_root]
def build_tree1(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
"""
从 nodes + links 构建树结构。
- 根节点label 包含 '舰艇' 的节点
- 父子关系links 中 label == '包含' 的边 (source -> target)
- 支持多棵树(多个舰艇根节点)
"""
# 1. 建立 id -> node 的映射
node_map: Dict[str, Dict] = {n["id"]: n for n in nodes}
# 2. 建立 父id -> [子id列表] 的邻接表(只处理"包含"关系)
children_map: Dict[str, List[str]] = {n["id"]: [] for n in nodes}
parent_map: Dict[str, str] = {} # 子id -> 父id用于识别根节点
for link in links:
if link.get("label") == "包含":
src = link["source"]
tgt = link["target"]
if src in children_map:
children_map[src].append(tgt)
if tgt not in parent_map:
parent_map[tgt] = src
# 3. 递归构建子树
def build_subtree(node_id: str, ancestors: set = None) -> Dict[str, Any]:
if ancestors is None:
ancestors = set()
if node_id in ancestors:
return None
ancestors.add(node_id)
node = node_map[node_id]
subtree = {"id": node_id, "name": node.get("name", ""), "labels": node.get("labels", []), "properties": node.get("properties", {}), "color": node.get("color"), "size": node.get("size"), "children": []}
current_ancestors = ancestors | {node_id}
for child_id in children_map.get(node_id, []):
child = build_subtree(child_id, current_ancestors)
if child is not None:
subtree["children"].append(child)
return subtree
# def build_subtree(node_id: str) -> Dict[str, Any]:
# node = node_map[node_id]
# subtree = {
# "id": node_id,
# "name": node.get("name", ""),
# "labels": node.get("labels", []),
# "properties": node.get("properties", {}),
# "color": node.get("color"),
# "size": node.get("size"),
# "children": []
# }
# for child_id in children_map.get(node_id, []):
# subtree["children"].append(build_subtree(child_id))
# return subtree
# 4. 找出所有根节点label 包含 '舰艇',且没有父节点
roots = [n["id"] for n in nodes if "舰艇" in n.get("labels", []) and n["id"] not in parent_map]
# 5. 为每棵树构建结构,并打上树编号
forest = []
for i, root_id in enumerate(roots):
tree = build_subtree(root_id)
tree["tree_index"] = i # 区分不同树
forest.append(tree)
return forest
def _filter_searchable(graph: dict) -> dict:
"""剥掉节点的 Searchable 标签;若剥完后 labels 为空则整体移除该节点及其关联边"""
for n in graph["nodes"]:
n["labels"] = [l for l in n.get("labels", []) if l != "Searchable"]
valid_ids = {n["id"] for n in graph["nodes"] if n["labels"]}
return {
"nodes": [n for n in graph["nodes"] if n["id"] in valid_ids],
"links": [
l for l in graph["links"]
if l["source"] in valid_ids and l["target"] in valid_ids
],
}
def build_graph_from_paths(paths):
"""
核心公共方法
从 Neo4j 路径结果构建 graphData
"""
node_map: Dict[str, dict] = {}
link_map: Dict[str, dict] = {}
def add_edge(rel):
s = rel.start_node
t = rel.end_node
if s.element_id not in node_map:
node_map[s.element_id] = node_to_json(s)
if t.element_id not in node_map:
node_map[t.element_id] = node_to_json(t)
if rel.element_id not in link_map:
link_map[rel.element_id] = rel_to_json(rel)
for record in paths:
p = record["p"]
if not is_simple_path(p):
continue
for r in p.relationships:
add_edge(r)
return _filter_searchable({"nodes": list(node_map.values()), "links": list(link_map.values())})
# def fetch_graph_sample(driver,limit: int = 200):
# """提取采样逻辑为可复用函数"""
# with driver.session() as session:
# result = session.run("""
# MATCH (n)
# WITH n
# ORDER BY elementId(n) ASC
# LIMIT $limit
# MATCH (n)-[rel]->(m)
# RETURN n, rel, m
# """, limit=limit)
# node_map, link_map = {}, {}
# for rec in result:
# n, r, m = rec["n"], rec["rel"], rec["m"]
# node_map[n.element_id] = node_to_json(n)
# node_map[m.element_id] = node_to_json(m)
# link_map[r.element_id] = rel_to_json(r)
# return {
# "nodes": list(node_map.values()),
# "links": list(link_map.values())
# }
# def fetch_graph_sample(driver, limit: int = 200):
# """
# 获取采样数据:
# 1. 选取前 $limit 个种子节点Level 0 或 1取决于业务定义这里定义种子为 Level 0
# 2. 对每个节点向下探索最多 3 跳。
# 3. 每个节点对象中增加 level 字段,表示其距离种子节点的深度。
# """
# with driver.session() as session:
# # Cypher 逻辑:
# # - 选出种子节点 n
# # - 匹配 1-3 跳路径
# # - 返回种子节点 n路径 path以及该路径的长度即层级
# result = session.run("""
# MATCH (n)
# WITH n ORDER BY elementId(n) ASC LIMIT $limit
# OPTIONAL MATCH path = (n)-[*1..3]-(m)
# UNWIND nodes(path) AS x
# WITH n, path, m, COUNT(DISTINCT x) AS nodeCount
# WHERE path IS NULL OR nodeCount = LENGTH(path) + 1
# RETURN n, path, LENGTH(path) AS depth
# """, limit=limit)
# node_map, link_map = {}, {}
# for rec in result:
# seed_node = rec["n"]
# path = rec["path"]
# depth = rec["depth"] or 0 # 如果没有路径,深度为 0
# # 1. 处理种子节点 (Level 0)
# if seed_node.element_id not in node_map:
# node_data = node_to_json(seed_node)
# node_data["level"] = 0 # 种子节点设为 0 层
# node_map[seed_node.element_id] = node_data
# # 2. 如果存在路径,解析路径中的所有节点和关系
# if path:
# # 路径中的节点处理
# # path.nodes 包含了从起始到终点的所有节点
# for i, node in enumerate(path.nodes):
# if node.element_id not in node_map:
# node_data = node_to_json(node)
# # 层级即为该节点在当前路径中的索引
# node_data["level"] = i
# node_map[node.element_id] = node_data
# else:
# # 如果节点已存在,保留最小的 level (即最靠近种子的距离)
# node_map[node.element_id]["level"] = min(node_map[node.element_id].get("level", 3), i)
# # 路径中的关系处理
# for rel in path.relationships:
# if rel.element_id not in link_map:
# link_map[rel.element_id] = rel_to_json(rel)
# return {
# "nodes": list(node_map.values()),
# "links": list(link_map.values())
# }
# def fetch_graph_sample(driver, limit=200, max_depth=3):
# with driver.session() as session:
# result = session.run("""
# MATCH (c)
# WITH c ORDER BY elementId(c) ASC LIMIT $limit
# MATCH path = (c)-[*0..3]-(n)
# WITH n, MIN(length(path)) AS dist
# OPTIONAL MATCH (n)-[r]-(m)
# RETURN n, dist, r, m
# """, limit=limit)
# node_map, link_map = {}, {}
# for rec in result:
# n = rec["n"]
# dist = rec["dist"]
# r = rec["r"]
# m = rec["m"]
# # ---- 节点 ----
# nid = n.element_id
# if nid not in node_map:
# node_data = node_to_json(n)
# node_data["level"] = min(dist, 3) if dist is not None else 0
# node_map[nid] = node_data
# # ---- 边 ----
# if r and m:
# mid = m.element_id
# if mid not in node_map:
# node_map[mid] = node_to_json(m)
# link_map.setdefault(r.element_id, rel_to_json(r))
# return {
# "nodes": list(node_map.values()),
# "links": list(link_map.values())
# }
def fetch_graph_sample(driver):
"""
查询所有舰艇及其关联的系统和子系统
路径结构:(舰艇) -> (系统) -> (子系统)
"""
with driver.session() as session:
# =========================
# 统一查询:舰艇 -> 系统 -> 子系统
# =========================
# 逻辑说明:
# 1. MATCH (ship:舰艇):选中所有舰艇
# 2. -[*1..3]->(target):查找深度为 1 到 3 的路径
# - 深度 1舰艇 -> 系统
# - 深度 2舰艇 -> 系统 -> 子系统
# 3. WHERE target:系统 OR target:子系统:确保终点是我们关心的节点类型
# (防止查询出其他无关的深层节点)
query = """
MATCH path = (ship:舰艇)-[*1..3]->(target)
WHERE target:系统 OR target:子系统
RETURN path
"""
result = session.run(query)
# =========================
# 解析 paths
# =========================
all_nodes = {}
all_links = {}
for record in result:
path = record["path"]
# 解析节点
for node in path.nodes:
# 兼容 ID 获取
node_key = node.element_id if hasattr(node, 'element_id') else node.id
if node_key not in all_nodes:
all_nodes[node_key] = node_to_json(node)
# 解析关系
for rel in path.relationships:
rel_key = rel.element_id if hasattr(rel, 'element_id') else rel.id
if rel_key not in all_links:
all_links[rel_key] = rel_to_json(rel)
return _filter_searchable({
"nodes": list(all_nodes.values()),
"links": list(all_links.values())
})
def fetch_graph_sample1(driver):
"""
查找路径:
1. 优先:舰艇 -> 子系统
2. fallback舰艇 -> 系统
注意:根据数据库实际结构,使用 Label (:舰艇, :子系统) 进行匹配,
而不是 WHERE category = '...'
"""
# 定义采样数量
SAMPLE_LIMIT = 400
with driver.session() as session:
# =========================
# 第一阶段:舰艇 -> 子系统
# =========================
# 修改点:直接使用 (:舰艇) 和 (:子系统) 标签匹配,去掉 WHERE 子句
result = session.run(
"""
MATCH (ship:舰艇)
WITH ship LIMIT $limit
MATCH path = (ship)-[*1..3]->(sub:子系统)
RETURN path
""",
limit=SAMPLE_LIMIT,
)
paths = [record["path"] for record in result]
# =========================
# fallback舰艇 -> 系统
# =========================
if not paths:
# 修改点:直接使用 (:舰艇) 和 (:系统) 标签匹配
result = session.run(
"""
MATCH (ship:舰艇)
WITH ship LIMIT $limit
MATCH path = (ship)-[*1..2]->(sys:系统)
RETURN path
""",
limit=SAMPLE_LIMIT,
)
paths = [record["path"] for record in result]
# =========================
# 解析 paths
# =========================
all_nodes = {}
all_links = {}
for path in paths:
# 节点
for node in path.nodes:
# 兼容 Neo4j 5.x (element_id) 和 4.x (id)
node_key = node.element_id if hasattr(node, "element_id") else node.id
if node_key not in all_nodes:
all_nodes[node_key] = node_to_json(node)
# 关系
for rel in path.relationships:
rel_key = rel.element_id if hasattr(rel, "element_id") else rel.id
if rel_key not in all_links:
all_links[rel_key] = rel_to_json(rel)
return {"nodes": list(all_nodes.values()), "links": list(all_links.values())}
# =============================
# 文件名处理方法
# =============================
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:
# 注意:实际运行时请确保 result.json 存在,此处仅为代码结构展示
if not os.path.exists(TREE_JSON_PATH):
return {}
with open(TREE_JSON_PATH, "r", encoding="utf-8") as f:
return json.load(f)
def _resolve_ship_root(data: Dict, ship_code: str) -> Dict:
raw = data.get(ship_code)
if not isinstance(raw, dict):
return {}
# 结构 1直接是节点 {"code","name","children"}
if {"code", "name", "children"}.issubset(raw.keys()):
return raw
# 结构 2包装层 {"101": {...root node...}}
if ship_code in raw and isinstance(raw[ship_code], dict):
node = raw[ship_code]
if {"code", "name", "children"}.issubset(node.keys()):
return node
# 兜底:取第一个 node-like 值
for value in raw.values():
if isinstance(value, dict) and {"code", "name", "children"}.issubset(value.keys()):
return value
return {}
def _longest_common_substring_len(a: str, b: str) -> int:
if not a or not b:
return 0
# 优化:确保 a 是较短的字符串以节省空间(可选),原逻辑保持不动也没问题
dp = [0] * (len(b) + 1)
best = 0
for i in range(1, len(a) + 1):
prev = 0
for j in range(1, len(b) + 1):
temp = dp[j]
if a[i - 1] == b[j - 1]:
dp[j] = prev + 1
if dp[j] > best:
best = dp[j]
else:
dp[j] = 0
prev = temp
return best
def _match_level4_tool(filename_text: str, device_node: Dict) -> Tuple[str, str]:
"""
匹配 Level 4 零件。
规则:如果文件名与 Level 3 (设备名) 的最大共同字符串长度 >= 与任何 Level 4 (零件名) 的长度,
则只匹配到 Level 3返回空。
只有当 Level 4 的匹配度严格高于 Level 3 时,才返回 Level 4 信息。
"""
children = device_node.get("children") or {}
# 1. 计算文件名与当前设备 (Level 3) 名称的匹配度
device_name = str(device_node.get("name") or "")
level3_score = _longest_common_substring_len(filename_text, device_name)
best_name = ""
best_code = ""
best_score = 0 # 记录 Level 4 中的最佳得分
# 2. 遍历 Level 4 子节点
for child in children.values():
tool_name = str(child.get("name") or "")
tool_code = str(child.get("code") or "")
score = _longest_common_substring_len(filename_text, tool_name)
# 关键修改:只有当 Level 4 的得分严格大于 Level 3 的得分时,才视为有效匹配
# 并且要比当前找到的其他 Level 4 更好
if score > level3_score and score > best_score:
best_score = score
best_name = tool_name
best_code = tool_code
# 如果没有找到比 Level 3 匹配度更高的 Level 4则返回空
if best_score <= level3_score:
return "", ""
return best_name, best_code
def process_filename(filename: str) -> Dict[str, str]:
name, _ = os.path.splitext(os.path.basename(filename))
parts = re.split(r"[_-]", name)
xian_number = ""
level_1_system_name = ""
level_2_system_name = ""
device_name = ""
tool_name = ""
tool_code = ""
if not parts:
return {}
xian_number_code = parts[0].strip()
level1_code = ""
level2_code = ""
device_code = ""
if len(parts) > 1:
part = parts[1].strip()
# 防止索引越界
if len(part) < 2:
return {}
level1_code = part[:2]
level2_code = part[2:4] if len(part) > 2 else ""
device_code = part[4:] if len(part) > 4 else ""
# 若 parts[2] 存在,则用 parts[2] 作为 device_code支持 122-06A0014-B01001_发动机 格式)
if len(parts) > 2:
device_code_alt = parts[2].strip()
if device_code_alt:
device_code = device_code_alt
data = _load_tree_data()
root = _resolve_ship_root(data, xian_number_code)
if not root:
return {}
xian_number = root.get("name", "")
# 一级
s1 = root.get("children", {}).get(level1_code)
if s1:
level_1_system_name = s1.get("name", "")
# 二级
s2 = s1.get("children", {}).get(level2_code)
if s2:
level_2_system_name = s2.get("name", "")
# 三级设备
device = s2.get("children", {}).get(device_code)
if device:
device_name = device.get("name", "")
# 若存在 level4匹配与文件名最大相同字符串的零件
# 内部已包含逻辑:如果设备名匹配度更高,则不返回零件
tool_name, tool_code = _match_level4_tool(name, device)
return {
"Xian_Number": xian_number,
"Xian_Number_code": xian_number_code,
"Level_1_System_Name": level_1_system_name,
"level1_code": level1_code,
"Level_2_System_Name": level_2_system_name,
"level2_code": level2_code,
"Device_Name": device_name,
"device_code": device_code,
"tool_name": tool_name,
"tool_code": tool_code,
}
def get_entity(filename):
result = process_filename(filename)
entities = []
relationships = []
low_level = ""
if result:
xian_number = result.get("Xian_Number")
xian_number_code = result.get("Xian_Number_code")
level_1_system_name = result.get("Level_1_System_Name")
level1_code = result.get("level1_code")
level_2_system_name = result.get("Level_2_System_Name")
level2_code = result.get("level2_code")
device_name = result.get("Device_Name")
device_code = result.get("device_code")
tool_name = result.get("tool_name")
tool_code = result.get("tool_code")
if xian_number:
entities.append(
{
"type": "舰艇",
"properties": {
"名称": xian_number,
"舷号": xian_number_code,
},
}
)
low_level = xian_number
if level_1_system_name:
entities.append(
{
"type": "系统",
"properties": {
"名称": level_1_system_name,
"系统编码": level1_code,
},
}
)
low_level = level_1_system_name
if xian_number and level_1_system_name:
relationships.append(
{
"type": "包含",
"from_entity": xian_number,
"to_entity": level_1_system_name,
}
)
if level_2_system_name:
entities.append(
{
"type": "子系统",
"properties": {
"名称": level_2_system_name,
"子系统编码": level2_code,
},
}
)
low_level = level_2_system_name
if level_1_system_name and level_2_system_name:
relationships.append(
{
"type": "包含",
"from_entity": level_1_system_name,
"to_entity": level_2_system_name,
}
)
if device_name:
entities.append(
{
"type": "设备",
"properties": {
"名称": device_name,
"设备编码": device_code,
},
}
)
low_level = device_name
if level_2_system_name and device_name:
relationships.append(
{
"type": "包含",
"from_entity": level_2_system_name,
"to_entity": device_name,
}
)
if tool_name:
entities.append(
{
"type": "零件",
"properties": {
"名称": tool_name,
"零件编码": tool_code,
},
}
)
low_level = tool_name
if device_name and tool_name:
relationships.append(
{
"type": "包含",
"from_entity": device_name,
"to_entity": tool_name,
}
)
final_result = {
"entities": entities,
"relationships": relationships,
"low_level": low_level,
}
return final_result
# 示例测试逻辑(非必须,仅供验证)
if __name__ == "__main__":
# 阀控蓄电池脉冲式快速充电装置
test_list = [
"777-11A1014-B01001_发动机-维修手册.pdf",
]
for i in test_list:
result = get_entity(i)
lower_entity = result.get("low_level")
print(result)
print(lower_entity)
#from neo4j import GraphDatabase
#NEO4J_URI = "bolt://192.168.0.111:7687"
#NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
#NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD","zdht123@") # 不设默认值,强制要求提供
#NEO4J_DATABASE = os.getenv("NEO4J_DATABASE", "neo4j")
#driver = GraphDatabase.driver(
# NEO4J_URI,
# auth=(NEO4J_USER, NEO4J_PASSWORD),
# database=NEO4J_DATABASE
#)
#res = fetch_graph_sample(driver)
#print(res)