359 lines
14 KiB
Python
359 lines
14 KiB
Python
import time
|
||
import asyncio
|
||
from typing import List, Any, Optional
|
||
from collections import OrderedDict
|
||
from dataclasses import dataclass
|
||
from src.common.logger import get_logger
|
||
from src.config.config import global_config
|
||
from src.chat.message_receive.chat_stream import get_chat_manager
|
||
from src.chat.utils.chat_message_builder import get_raw_msg_by_timestamp_with_chat_inclusive
|
||
from src.bw_learner.expression_learner import expression_learner_manager
|
||
from src.bw_learner.jargon_miner import miner_manager
|
||
from src.person_info.person_info import Person
|
||
|
||
logger = get_logger("bw_learner")
|
||
|
||
|
||
@dataclass
|
||
class PersonInfo:
|
||
"""参与聊天的人物信息"""
|
||
user_id: str
|
||
user_platform: str
|
||
user_nickname: str
|
||
user_cardname: Optional[str]
|
||
person_name: str
|
||
last_seen_time: float # 最后发言时间
|
||
|
||
def get_unique_key(self) -> str:
|
||
"""获取唯一标识（用于去重）"""
|
||
return f"{self.user_platform}:{self.user_id}"
|
||
|
||
|
||
class MessageRecorder:
|
||
"""
|
||
统一的消息记录器，负责管理时间窗口和消息提取，并将消息分发给 expression_learner 和 jargon_miner
|
||
"""
|
||
|
||
def __init__(self, chat_id: str) -> None:
|
||
self.chat_id = chat_id
|
||
self.chat_stream = get_chat_manager().get_stream(chat_id)
|
||
self.chat_name = get_chat_manager().get_stream_name(chat_id) or chat_id
|
||
|
||
# 维护每个chat的上次提取时间
|
||
self.last_extraction_time: float = time.time()
|
||
|
||
# 提取锁，防止并发执行
|
||
self._extraction_lock = asyncio.Lock()
|
||
|
||
# 维护参与该chat_id的人物列表（最多30个，使用OrderedDict保持插入顺序）
|
||
# key: f"{platform}:{user_id}", value: PersonInfo
|
||
self._person_list: OrderedDict[str, PersonInfo] = OrderedDict()
|
||
self._max_person_count = 30
|
||
|
||
# 获取 expression 和 jargon 的配置参数
|
||
self._init_parameters()
|
||
|
||
# 获取 expression_learner 和 jargon_miner 实例
|
||
self.expression_learner = expression_learner_manager.get_expression_learner(chat_id)
|
||
self.jargon_miner = miner_manager.get_miner(chat_id)
|
||
|
||
def _init_parameters(self) -> None:
|
||
"""初始化提取参数"""
|
||
# 获取 expression 配置
|
||
_, self.enable_expression_learning, self.enable_jargon_learning = (
|
||
global_config.expression.get_expression_config_for_chat(self.chat_id)
|
||
)
|
||
self.min_messages_for_extraction = 30
|
||
self.min_extraction_interval = 60
|
||
|
||
logger.debug(
|
||
f"MessageRecorder 初始化: chat_id={self.chat_id}, "
|
||
f"min_messages={self.min_messages_for_extraction}, "
|
||
f"min_interval={self.min_extraction_interval}"
|
||
)
|
||
|
||
def should_trigger_extraction(self) -> bool:
|
||
"""
|
||
检查是否应该触发消息提取
|
||
|
||
Returns:
|
||
bool: 是否应该触发提取
|
||
"""
|
||
# 检查时间间隔
|
||
time_diff = time.time() - self.last_extraction_time
|
||
if time_diff < self.min_extraction_interval:
|
||
return False
|
||
|
||
# 检查消息数量
|
||
recent_messages = get_raw_msg_by_timestamp_with_chat_inclusive(
|
||
chat_id=self.chat_id,
|
||
timestamp_start=self.last_extraction_time,
|
||
timestamp_end=time.time(),
|
||
)
|
||
|
||
if not recent_messages or len(recent_messages) < self.min_messages_for_extraction:
|
||
return False
|
||
|
||
return True
|
||
|
||
async def extract_and_distribute(self) -> None:
|
||
"""
|
||
提取消息并分发给 expression_learner 和 jargon_miner
|
||
"""
|
||
# 使用异步锁防止并发执行
|
||
async with self._extraction_lock:
|
||
# 在锁内检查，避免并发触发
|
||
if not self.should_trigger_extraction():
|
||
return
|
||
|
||
# 检查 chat_stream 是否存在
|
||
if not self.chat_stream:
|
||
return
|
||
|
||
# 记录本次提取的时间窗口，避免重复提取
|
||
extraction_start_time = self.last_extraction_time
|
||
extraction_end_time = time.time()
|
||
|
||
# 立即更新提取时间，防止并发触发
|
||
self.last_extraction_time = extraction_end_time
|
||
|
||
try:
|
||
logger.info(f"在聊天流 {self.chat_name} 开始统一消息提取和分发")
|
||
|
||
# 拉取提取窗口内的消息
|
||
messages = get_raw_msg_by_timestamp_with_chat_inclusive(
|
||
chat_id=self.chat_id,
|
||
timestamp_start=extraction_start_time,
|
||
timestamp_end=extraction_end_time,
|
||
)
|
||
|
||
if not messages:
|
||
logger.debug(f"聊天流 {self.chat_name} 没有新消息，跳过提取")
|
||
return
|
||
|
||
# 按时间排序，确保顺序一致
|
||
messages = sorted(messages, key=lambda msg: msg.time or 0)
|
||
|
||
# 更新参与聊天的人物列表
|
||
self._update_person_list(messages)
|
||
|
||
logger.info(f"聊天流 {self.chat_name} 的人物列表: {self._person_list}")
|
||
|
||
logger.info(
|
||
f"聊天流 {self.chat_name} 提取到 {len(messages)} 条消息，"
|
||
f"时间窗口: {extraction_start_time:.2f} - {extraction_end_time:.2f}"
|
||
)
|
||
|
||
|
||
# 分别触发 expression_learner 和 jargon_miner 的处理
|
||
# 传递提取的消息，避免它们重复获取
|
||
# 触发 expression 学习（如果启用）
|
||
if self.enable_expression_learning:
|
||
asyncio.create_task(
|
||
self._trigger_expression_learning(extraction_start_time, extraction_end_time, messages)
|
||
)
|
||
|
||
# 触发 jargon 提取（如果启用），传递消息
|
||
# if self.enable_jargon_learning:
|
||
# asyncio.create_task(
|
||
# self._trigger_jargon_extraction(extraction_start_time, extraction_end_time, messages)
|
||
# )
|
||
|
||
except Exception as e:
|
||
logger.error(f"为聊天流 {self.chat_name} 提取和分发消息失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
# 即使失败也保持时间戳更新，避免频繁重试
|
||
|
||
async def _trigger_expression_learning(
|
||
self,
|
||
timestamp_start: float,
|
||
timestamp_end: float,
|
||
messages: List[Any]
|
||
) -> None:
|
||
"""
|
||
触发 expression 学习，使用指定的消息列表
|
||
|
||
Args:
|
||
timestamp_start: 开始时间戳
|
||
timestamp_end: 结束时间戳
|
||
messages: 消息列表
|
||
"""
|
||
try:
|
||
# 传递消息和过滤函数给 ExpressionLearner
|
||
learnt_style = await self.expression_learner.learn_and_store(
|
||
messages=messages,
|
||
person_name_filter=self.contains_person_name
|
||
)
|
||
|
||
if learnt_style:
|
||
logger.info(f"聊天流 {self.chat_name} 表达学习完成")
|
||
else:
|
||
logger.debug(f"聊天流 {self.chat_name} 表达学习未获得有效结果")
|
||
except Exception as e:
|
||
logger.error(f"为聊天流 {self.chat_name} 触发表达学习失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
async def _trigger_jargon_extraction(
|
||
self,
|
||
timestamp_start: float,
|
||
timestamp_end: float,
|
||
messages: List[Any]
|
||
) -> None:
|
||
"""
|
||
触发 jargon 提取，使用指定的消息列表
|
||
|
||
Args:
|
||
timestamp_start: 开始时间戳
|
||
timestamp_end: 结束时间戳
|
||
messages: 消息列表
|
||
"""
|
||
try:
|
||
# 传递消息和过滤函数给 JargonMiner
|
||
await self.jargon_miner.run_once(
|
||
messages=messages,
|
||
person_name_filter=self.contains_person_name
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"为聊天流 {self.chat_name} 触发黑话提取失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
def _update_person_list(self, messages: List[Any]) -> None:
|
||
"""
|
||
从消息中提取人物信息并更新人物列表
|
||
|
||
Args:
|
||
messages: 消息列表
|
||
"""
|
||
for msg in messages:
|
||
# 获取消息发送者信息
|
||
# 消息对象可能是 DatabaseMessages，它有 user_info 属性
|
||
if hasattr(msg, 'user_info'):
|
||
# DatabaseMessages 类型
|
||
user_info = msg.user_info
|
||
user_id = getattr(user_info, 'user_id', None) or ''
|
||
user_platform = getattr(user_info, 'platform', None) or ''
|
||
user_nickname = getattr(user_info, 'user_nickname', None) or ''
|
||
user_cardname = getattr(user_info, 'user_cardname', None)
|
||
else:
|
||
# 直接属性访问
|
||
user_id = getattr(msg, 'user_id', None) or ''
|
||
user_platform = getattr(msg, 'user_platform', None) or ''
|
||
user_nickname = getattr(msg, 'user_nickname', None) or ''
|
||
user_cardname = getattr(msg, 'user_cardname', None)
|
||
|
||
msg_time = getattr(msg, 'time', time.time())
|
||
|
||
# 检查必要信息
|
||
if not user_id or not user_platform:
|
||
continue
|
||
|
||
# 获取 person_name
|
||
try:
|
||
person = Person(platform=user_platform, user_id=str(user_id))
|
||
person_name = person.person_name or user_nickname or (user_cardname if user_cardname else "未知用户")
|
||
except Exception as e:
|
||
logger.info(f"获取person_name失败: {e}, 使用nickname")
|
||
person_name = user_nickname or (user_cardname if user_cardname else "未知用户")
|
||
|
||
# 生成唯一key
|
||
unique_key = f"{user_platform}:{user_id}"
|
||
|
||
# 如果已存在，更新最后发言时间
|
||
if unique_key in self._person_list:
|
||
self._person_list[unique_key].last_seen_time = msg_time
|
||
# 移动到末尾（表示最近活跃）
|
||
self._person_list.move_to_end(unique_key)
|
||
else:
|
||
# 如果超过最大数量，移除最早的（最前面的）
|
||
if len(self._person_list) >= self._max_person_count:
|
||
oldest_key = next(iter(self._person_list))
|
||
del self._person_list[oldest_key]
|
||
logger.info(f"人物列表已满，移除最早的人物: {oldest_key}")
|
||
|
||
# 添加新人物
|
||
person_info = PersonInfo(
|
||
user_id=str(user_id),
|
||
user_platform=user_platform,
|
||
user_nickname=user_nickname or "",
|
||
user_cardname=user_cardname,
|
||
person_name=person_name,
|
||
last_seen_time=msg_time
|
||
)
|
||
self._person_list[unique_key] = person_info
|
||
logger.info(f"添加新人物到列表: {unique_key}, person_name={person_name}")
|
||
|
||
def contains_person_name(self, content: str) -> bool:
|
||
"""
|
||
检查内容是否包含任何参与聊天的人物的名称或昵称
|
||
|
||
Args:
|
||
content: 要检查的内容
|
||
|
||
Returns:
|
||
bool: 如果包含任何人物名称或昵称，返回True
|
||
"""
|
||
if not content or not self._person_list:
|
||
return False
|
||
|
||
content_lower = content.strip().lower()
|
||
if not content_lower:
|
||
return False
|
||
|
||
# 检查所有人物
|
||
for person_info in self._person_list.values():
|
||
# 检查 person_name
|
||
if person_info.person_name:
|
||
person_name_lower = person_info.person_name.strip().lower()
|
||
if person_name_lower and person_name_lower in content_lower:
|
||
logger.debug(f"内容包含person_name: {person_info.person_name} in {content}")
|
||
return True
|
||
|
||
# 检查 user_nickname
|
||
if person_info.user_nickname:
|
||
nickname_lower = person_info.user_nickname.strip().lower()
|
||
if nickname_lower and nickname_lower in content_lower:
|
||
logger.debug(f"内容包含nickname: {person_info.user_nickname} in {content}")
|
||
return True
|
||
|
||
# 检查 user_cardname（群昵称）
|
||
if person_info.user_cardname:
|
||
cardname_lower = person_info.user_cardname.strip().lower()
|
||
if cardname_lower and cardname_lower in content_lower:
|
||
logger.debug(f"内容包含cardname: {person_info.user_cardname} in {content}")
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
class MessageRecorderManager:
|
||
"""MessageRecorder 管理器"""
|
||
|
||
def __init__(self) -> None:
|
||
self._recorders: dict[str, MessageRecorder] = {}
|
||
|
||
def get_recorder(self, chat_id: str) -> MessageRecorder:
|
||
"""获取或创建指定 chat_id 的 MessageRecorder"""
|
||
if chat_id not in self._recorders:
|
||
self._recorders[chat_id] = MessageRecorder(chat_id)
|
||
return self._recorders[chat_id]
|
||
|
||
|
||
# 全局管理器实例
|
||
recorder_manager = MessageRecorderManager()
|
||
|
||
|
||
async def extract_and_distribute_messages(chat_id: str) -> None:
|
||
"""
|
||
统一的消息提取和分发入口函数
|
||
|
||
Args:
|
||
chat_id: 聊天流ID
|
||
"""
|
||
recorder = recorder_manager.get_recorder(chat_id)
|
||
await recorder.extract_and_distribute()
|
||
|