159 lines
6.6 KiB
Python
159 lines
6.6 KiB
Python
"""
|
||
入口节点处理模块
|
||
|
||
功能：过滤、排序和格式化入口节点
|
||
"""
|
||
|
||
import logging
|
||
from typing import Dict, List, Tuple, Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def filter_and_sort_entry_nodes(entry_nodes: dict,max_nodes_per_label: int = 5,min_score_threshold: float = 0.2,adaptive_limit: bool = True,min_total_nodes: int = 3,) -> dict:
|
||
"""
|
||
过滤和重排序入口节点（按标签内排序和限制，带简单质量控制）
|
||
|
||
策略：
|
||
1. 对每个标签内的节点按 score 降序排序；
|
||
2. 过滤掉分数过低的节点（min_score_threshold 以下）；
|
||
3. 自适应限制：高分节点多时适当多留，低分节点多时按上限截断；
|
||
4. 如果整体节点数量过少（例如 < min_total_nodes），在内存中做一次
|
||
跨标签的轻量补充，不新增任何数据库查询，几乎不增加运行时间。
|
||
|
||
Args:
|
||
entry_nodes: 入口节点字典，格式如 {"设备": [{"name": "xxx", "score": 0.9}]}
|
||
max_nodes_per_label: 每个标签最大节点数（默认 5）
|
||
min_score_threshold: 入口节点最低得分阈值（默认 0.2）
|
||
adaptive_limit: 是否根据得分分布自适应调整每个标签保留数量
|
||
min_total_nodes: 全部标签合计至少保留的节点数量下限
|
||
|
||
Returns:
|
||
dict: 过滤和排序后的入口节点字典
|
||
"""
|
||
if not entry_nodes:
|
||
return {}
|
||
|
||
filtered: Dict[str, List[Dict[str, Any]]] = {}
|
||
|
||
# ========== 1. 标签内排序 + 质量过滤 ==========
|
||
for label, nodes in entry_nodes.items():
|
||
if not isinstance(nodes, list) or not nodes:
|
||
continue
|
||
|
||
# 1.1 按 score 降序排序（标签内排序）
|
||
sorted_nodes = sorted(nodes, key=lambda x: x.get('score', 0), reverse=True)
|
||
|
||
# 1.2 过滤低质量节点
|
||
high_quality_nodes = [
|
||
n for n in sorted_nodes
|
||
if n.get('score', 0) >= min_score_threshold
|
||
]
|
||
if not high_quality_nodes:
|
||
# 如果全都低于阈值，则保底保留分数最高的 1 个，避免完全丢标签
|
||
high_quality_nodes = sorted_nodes[:1]
|
||
|
||
# 1.3 自适应限制
|
||
if adaptive_limit:
|
||
# 认为 score >= 0.7 为"高置信度"
|
||
strong_nodes = [n for n in high_quality_nodes if n.get('score', 0) >= 0.7]
|
||
if len(strong_nodes) >= 1:
|
||
# 高置信度节点存在时，优先全部保留这些节点；
|
||
# 如果数量仍然超过上限，则再按上限截断。
|
||
if len(strong_nodes) > max_nodes_per_label:
|
||
filtered_nodes = strong_nodes[:max_nodes_per_label]
|
||
else:
|
||
# 高置信度节点数量不多，可以再补少量次优节点
|
||
remain = max_nodes_per_label - len(strong_nodes)
|
||
extra = [n for n in high_quality_nodes if n not in strong_nodes][:remain]
|
||
filtered_nodes = strong_nodes + extra
|
||
else:
|
||
# 没有高置信度节点，直接按上限截断
|
||
filtered_nodes = high_quality_nodes[:max_nodes_per_label]
|
||
else:
|
||
filtered_nodes = high_quality_nodes[:max_nodes_per_label]
|
||
|
||
if filtered_nodes:
|
||
filtered[label] = filtered_nodes
|
||
|
||
total_before = sum(len(nodes) for nodes in entry_nodes.values() if isinstance(nodes, list))
|
||
total_after = sum(len(nodes) for nodes in filtered.values())
|
||
|
||
# ========== 2. 结果过少时的轻量级跨标签补充 ==========
|
||
# 不触发任何新的检索，只在已有 entry_nodes 里做二次筛选，
|
||
# 对整体性能影响可以忽略。
|
||
if total_after < min_total_nodes:
|
||
# 已经存在的 (label, id(node))，用于去重
|
||
existing_keys = set()
|
||
for label, nodes in filtered.items():
|
||
for n in nodes:
|
||
key = (label, n.get('name') or n.get('_name') or (n.get('node', {}).get('名称')) or id(n))
|
||
existing_keys.add(key)
|
||
|
||
# 收集所有可候选的节点
|
||
candidates: List[Tuple[str, Dict[str, Any]]] = []
|
||
for label, nodes in entry_nodes.items():
|
||
if not isinstance(nodes, list):
|
||
continue
|
||
for n in nodes:
|
||
score = n.get('score', 0)
|
||
if score < min_score_threshold:
|
||
continue
|
||
key = (label, n.get('name') or n.get('_name') or (n.get('node', {}).get('名称')) or id(n))
|
||
if key in existing_keys:
|
||
continue
|
||
candidates.append((label, n))
|
||
|
||
# 按 score 降序排序，补充若干个节点，直到达到 min_total_nodes 或候选耗尽
|
||
candidates.sort(key=lambda x: x[1].get('score', 0), reverse=True)
|
||
for label, node in candidates:
|
||
if sum(len(nodes) for nodes in filtered.values()) >= min_total_nodes:
|
||
break
|
||
# 控制单个标签不会被过度补充，这里允许最多是 max_nodes_per_label 的 2 倍
|
||
if len(filtered.get(label, [])) >= max_nodes_per_label * 2:
|
||
continue
|
||
filtered.setdefault(label, []).append(node)
|
||
|
||
total_after = sum(len(nodes) for nodes in filtered.values())
|
||
|
||
logger.debug(
|
||
f"入口节点过滤: 原始 {total_before} 个节点 "
|
||
f"-> 过滤后 {total_after} 个节点 "
|
||
f"({len(filtered)} 个标签，每个标签最多 {max_nodes_per_label} 个, "
|
||
f"min_score_threshold={min_score_threshold}, adaptive_limit={adaptive_limit}, "
|
||
f"min_total_nodes={min_total_nodes})"
|
||
)
|
||
|
||
return filtered
|
||
|
||
|
||
def format_entry_nodes(entry_nodes: dict) -> str:
|
||
"""
|
||
格式化入口节点信息
|
||
|
||
Args:
|
||
entry_nodes: 入口节点字典，格式如 {"设备": [{"name": "xxx", "score": 0.9}]}
|
||
|
||
Returns:
|
||
str: 格式化后的节点信息字符串
|
||
"""
|
||
if not entry_nodes:
|
||
return "无入口节点"
|
||
|
||
formatted = []
|
||
for label, nodes in entry_nodes.items():
|
||
if isinstance(nodes, list) and nodes:
|
||
node_list = []
|
||
# 显示所有节点（因为已经过滤过了）
|
||
for node in nodes:
|
||
name = node.get('name') or node.get('_name') or str(node)
|
||
score = node.get('score', 0)
|
||
node_list.append(f" - {name} (score: {score:.2f})")
|
||
formatted.append(f"{label}:\n" + "\n".join(node_list))
|
||
elif nodes:
|
||
formatted.append(f"{label}: {nodes}")
|
||
|
||
logger.info(f"入口节点: {formatted}")
|
||
return "\n\n".join(formatted) if formatted else "无入口节点"
|
||
|