BruceSun
/
MaiBot
mirror of https://github.com/Mai-with-u/MaiBot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

🤖 自动格式化代码 [skip ci]

pull/914/head
github-actions[bot] 2025-05-01 23:00:56 +00:00
parent 1b9dc432c2
commit 681814f7d1
2 changed files with 35 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
src/plugins/group_nickname/nickname_processor.py
View File

@ -14,6 +14,7 @@ logger = get_logger("nickname_processor")
_stop_event = threading.Event()
def _upsert_person(collection: Collection, person_id: str, user_id_int: int, platform: str):
"""
确保数据库中存在指定 person_id 的文档 (Upsert)。
@ -41,7 +42,7 @@ def _upsert_person(collection: Collection, person_id: str, user_id_int: int, pla
{
"$setOnInsert": {
"person_id": person_id,
"user_id": user_id_int, # 确保这里使用传入的 user_id_int
"user_id": user_id_int, # 确保这里使用传入的 user_id_int
"platform": platform,
"group_nicknames": [], # 初始化 group_nicknames 数组
}
@ -59,10 +60,10 @@ def _upsert_person(collection: Collection, person_id: str, user_id_int: int, pla
logger.error(
f"数据库操作失败 (DuplicateKeyError): person_id {person_id}. 错误: {dk_err}. 这不应该发生，请检查 person_id 生成逻辑和数据库状态。"
)
raise # 将异常向上抛出，让调用者处理
raise # 将异常向上抛出，让调用者处理
except Exception as e:
logger.exception(f"对 person_id {person_id} 执行 Upsert 时失败: {e}")
raise # 将异常向上抛出
raise # 将异常向上抛出
def _update_group_nickname(collection: Collection, person_id: str, group_id_str: str, nickname: str):
@ -80,9 +81,7 @@ def _update_group_nickname(collection: Collection, person_id: str, group_id_str:
result_inc = collection.update_one(
{
"person_id": person_id,
"group_nicknames": {
"$elemMatch": {"group_id": group_id_str, "nicknames.name": nickname}
},
"group_nicknames": {"$elemMatch": {"group_id": group_id_str, "nicknames.name": nickname}},
},
{"$inc": {"group_nicknames.$[group].nicknames.$[nick].count": 1}},
array_filters=[
@ -92,20 +91,20 @@ def _update_group_nickname(collection: Collection, person_id: str, group_id_str:
)
if result_inc.modified_count > 0:
# logger.debug(f"成功增加 person_id {person_id} 在群组 {group_id_str} 中绰号 '{nickname}' 的计数。")
return # 成功增加计数，操作完成
return # 成功增加计数，操作完成
# 3b. 如果上一步未修改 (绰号不存在于该群组)，尝试将新绰号添加到现有群组
result_push_nick = collection.update_one(
{
"person_id": person_id,
"group_nicknames.group_id": group_id_str, # 检查群组是否存在
"group_nicknames.group_id": group_id_str, # 检查群组是否存在
},
{"$push": {"group_nicknames.$[group].nicknames": {"name": nickname, "count": 1}}},
array_filters=[{"group.group_id": group_id_str}],
)
if result_push_nick.modified_count > 0:
logger.debug(f"成功为 person_id {person_id} 在现有群组 {group_id_str} 中添加新绰号 '{nickname}'。")
return # 成功添加绰号，操作完成
return # 成功添加绰号，操作完成
# 3c. 如果上一步也未修改 (群组条目本身不存在)，则添加新的群组条目和绰号
# 确保 group_nicknames 数组存在 (作为保险措施)
@ -117,7 +116,7 @@ def _update_group_nickname(collection: Collection, person_id: str, group_id_str:
result_push_group = collection.update_one(
{
"person_id": person_id,
"group_nicknames.group_id": {"$ne": group_id_str}, # 确保该群组 ID 尚未存在
"group_nicknames.group_id": {"$ne": group_id_str}, # 确保该群组 ID 尚未存在
},
{
"$push": {
@ -131,9 +130,10 @@ def _update_group_nickname(collection: Collection, person_id: str, group_id_str:
if result_push_group.modified_count > 0:
logger.debug(f"为 person_id {person_id} 添加了新的群组 {group_id_str} 和绰号 '{nickname}'。")
# else:
# 如果连添加群组也失败 (例如 group_id 已存在但之前的步骤都未匹配，理论上不太可能)，
# 可能需要进一步的日志或错误处理，但这通常意味着数据状态异常。
# logger.warning(f"尝试为 person_id {person_id} 添加新群组 {group_id_str} 失败，可能群组已存在但结构不符合预期。")
# 如果连添加群组也失败 (例如 group_id 已存在但之前的步骤都未匹配，理论上不太可能)，
# 可能需要进一步的日志或错误处理，但这通常意味着数据状态异常。
# logger.warning(f"尝试为 person_id {person_id} 添加新群组 {group_id_str} 失败，可能群组已存在但结构不符合预期。")
async def update_nickname_counts(platform: str, group_id: str, nickname_map: Dict[str, str]):
"""
@ -185,7 +185,7 @@ async def update_nickname_counts(platform: str, group_id: str, nickname_map: Dic
_update_group_nickname(person_info_collection, person_id, group_id_str, nickname)
# --- 统一处理数据库操作可能抛出的异常 ---
except (OperationFailure, DuplicateKeyError) as db_err: # 捕获特定的数据库错误
except (OperationFailure, DuplicateKeyError) as db_err: # 捕获特定的数据库错误
logger.exception(
f"数据库操作失败 ({type(db_err).__name__}): 用户 {user_id_str}, 群组 {group_id_str}, 绰号 {nickname}. 错误: {db_err}"
)
@ -193,6 +193,7 @@ async def update_nickname_counts(platform: str, group_id: str, nickname_map: Dic
# 捕获其他所有可能的错误 (例如 person_id 生成、辅助函数内部未捕获的错误等)
logger.exception(f"处理用户 {user_id_str} 的绰号 '{nickname}' 时发生意外错误：{e}")
# --- 使用 queue.Queue ---
queue_max_size = getattr(global_config, "NICKNAME_QUEUE_MAX_SIZE", 100)
nickname_queue: queue.Queue = queue.Queue(maxsize=queue_max_size)

34
src/plugins/group_nickname/nickname_utils.py
View File

@ -55,11 +55,15 @@ def select_nicknames_for_prompt(all_nicknames_info: Dict[str, List[Dict[str, int
# 如果抽样结果数量不足（例如权重问题导致提前退出），可以考虑是否需要补充
if len(selected_candidates_with_weight) < num_to_select:
logger.debug(f"加权随机选择后数量不足 ({len(selected_candidates_with_weight)}/{num_to_select})，补充选择次数最多的。")
logger.debug(
f"加权随机选择后数量不足 ({len(selected_candidates_with_weight)}/{num_to_select})，补充选择次数最多的。"
)
# 筛选出未被选中的候选
selected_ids = set((c[0], c[1]) for c in selected_candidates_with_weight) # 使用 (用户名, 绰号) 作为唯一标识
selected_ids = set(
(c[0], c[1]) for c in selected_candidates_with_weight
) # 使用 (用户名, 绰号) 作为唯一标识
remaining_candidates = [c for c in candidates if (c[0], c[1]) not in selected_ids]
remaining_candidates.sort(key=lambda x: x[2], reverse=True) # 按原始次数排序
remaining_candidates.sort(key=lambda x: x[2], reverse=True) # 按原始次数排序
needed = num_to_select - len(selected_candidates_with_weight)
selected_candidates_with_weight.extend(remaining_candidates[:needed])
@ -67,15 +71,14 @@ def select_nicknames_for_prompt(all_nicknames_info: Dict[str, List[Dict[str, int
# 日志：记录加权随机选择时发生的错误，并回退到简单选择
logger.error(f"绰号加权随机选择时出错: {e}。将回退到选择次数最多的 Top N。", exc_info=True)
# 出错时回退到选择次数最多的 N 个
candidates.sort(key=lambda x: x[2], reverse=True) # 按原始次数排序
candidates.sort(key=lambda x: x[2], reverse=True) # 按原始次数排序
# 注意：这里需要选择包含权重的元组，或者调整后续处理
selected_candidates_with_weight = candidates[:num_to_select]
# 格式化输出结果为 (用户名, 绰号, 次数)，移除权重
result = [(user, nick, count) for user, nick, count, _weight in selected_candidates_with_weight]
result.sort(key=lambda x: x[2], reverse=True) # 按次数降序
result.sort(key=lambda x: x[2], reverse=True) # 按次数降序
logger.debug(f"为 Prompt 选择的绰号: {result}")
return result
@ -285,7 +288,10 @@ async def trigger_nickname_analysis_if_needed(
# 日志：记录触发分析过程中发生的任何其他错误
logger.error(f"{log_prefix} 触发绰号分析时出错: {e}", exc_info=True)
def weighted_sample_without_replacement(candidates: List[Tuple[str, str, int, float]], k: int) -> List[Tuple[str, str, int, float]]:
def weighted_sample_without_replacement(
candidates: List[Tuple[str, str, int, float]], k: int
) -> List[Tuple[str, str, int, float]]:
"""
执行不重复的加权随机抽样。
@ -300,18 +306,18 @@ def weighted_sample_without_replacement(candidates: List[Tuple[str, str, int, fl
return []
if k >= len(candidates):
# 如果需要选择的数量大于或等于候选数量，直接返回所有候选
return candidates[:] # 返回副本以避免修改原始列表
return candidates[:] # 返回副本以避免修改原始列表
pool = candidates[:] # 创建候选列表的副本进行操作
pool = candidates[:] # 创建候选列表的副本进行操作
selected = []
# 注意：原评论代码中计算 total_weight 但未使用，这里也省略。
# random.choices 内部会处理权重的归一化。
for _ in range(min(k, len(pool))): # 确保迭代次数不超过池中剩余元素
if not pool: # 如果池已空，提前结束
for _ in range(min(k, len(pool))): # 确保迭代次数不超过池中剩余元素
if not pool: # 如果池已空，提前结束
break
weights = [c[3] for c in pool] # 获取当前池中所有元素的权重
weights = [c[3] for c in pool] # 获取当前池中所有元素的权重
# 检查权重是否有效
if sum(weights) <= 0:
# 如果所有剩余权重无效，随机选择一个（或根据需要采取其他策略）
@ -322,8 +328,8 @@ def weighted_sample_without_replacement(candidates: List[Tuple[str, str, int, fl
# 使用 random.choices 进行加权抽样，选择 1 个
# random.choices 返回一个列表，所以取第一个元素 [0]
chosen = random.choices(pool, weights=weights, k=1)[0]
pool.remove(chosen) # 从池中移除选中的元素，实现不重复抽样
pool.remove(chosen) # 从池中移除选中的元素，实现不重复抽样
selected.append(chosen)
return selected
return selected