""" Segment Mention Service 轻量级的知识库片段@提及通知服务 数据存储在Redis中 """ import json import uuid from datetime import datetime, timedelta from typing import Optional from extensions.ext_redis import redis_client from models.account import Account from models.dataset import Dataset, DocumentSegment from extensions.ext_database import db class SegmentMentionService: """片段@提及服务 - 基于Redis""" # Redis Key 前缀 COMMENT_KEY_PREFIX = "segment:comment:" # segment:comment:{comment_id} SEGMENT_COMMENTS_KEY = "segment:comments:{dataset_id}:{segment_id}" # 某个片段的所有评论ID列表 NOTIFICATION_KEY_PREFIX = "user:notification:" # user:notification:{account_id} UNREAD_COUNT_KEY = "user:notification:unread:{account_id}" # 未读数量 # 数据过期时间(30天) EXPIRATION_DAYS = 30 EXPIRATION_SECONDS = EXPIRATION_DAYS * 24 * 60 * 60 @staticmethod def create_comment( tenant_id: str, dataset_id: str, segment_id: str, content: str, account_id: str, mentioned_account_ids: list[str] = None ) -> dict: """ 创建评论 Args: tenant_id: 租户ID dataset_id: 知识库ID segment_id: 片段ID content: 评论内容 account_id: 评论者ID mentioned_account_ids: @提及的用户ID列表 Returns: 评论数据字典 """ comment_id = str(uuid.uuid4()) account = db.session.get(Account, account_id) comment_data = { 'id': comment_id, 'tenant_id': tenant_id, 'dataset_id': dataset_id, 'segment_id': segment_id, 'content': content, 'account_id': account_id, 'account_name': account.name if account else 'Unknown', 'account_avatar': account.avatar if account else '', 'mentioned_account_ids': mentioned_account_ids or [], 'created_at': datetime.utcnow().isoformat(), 'is_deleted': False } # 存储评论 comment_key = f"{SegmentMentionService.COMMENT_KEY_PREFIX}{comment_id}" redis_client.setex( comment_key, SegmentMentionService.EXPIRATION_SECONDS, json.dumps(comment_data) ) # 添加到片段的评论列表 segment_comments_key = SegmentMentionService.SEGMENT_COMMENTS_KEY.format( dataset_id=dataset_id, segment_id=segment_id ) redis_client.lpush(segment_comments_key, comment_id) redis_client.expire(segment_comments_key, SegmentMentionService.EXPIRATION_SECONDS) # 发送@提及通知 if mentioned_account_ids: SegmentMentionService._create_mention_notifications( comment_id=comment_id, comment_content=content, mentioned_account_ids=mentioned_account_ids, sender_id=account_id, sender_name=account.name if account else 'Unknown', dataset_id=dataset_id, segment_id=segment_id ) return comment_data @staticmethod def get_segment_comments( dataset_id: str, segment_id: str, limit: int = 50 ) -> list[dict]: """ 获取片段的所有评论 Args: dataset_id: 知识库ID segment_id: 片段ID limit: 最大返回数量 Returns: 评论列表 """ segment_comments_key = SegmentMentionService.SEGMENT_COMMENTS_KEY.format( dataset_id=dataset_id, segment_id=segment_id ) # 获取评论ID列表(最新的在前) comment_ids = redis_client.lrange(segment_comments_key, 0, limit - 1) comments = [] for comment_id in comment_ids: comment_key = f"{SegmentMentionService.COMMENT_KEY_PREFIX}{comment_id.decode('utf-8')}" comment_data = redis_client.get(comment_key) if comment_data: comment = json.loads(comment_data) if not comment.get('is_deleted'): comments.append(comment) return comments @staticmethod def delete_comment(comment_id: str, account_id: str) -> bool: """ 删除评论(软删除) Args: comment_id: 评论ID account_id: 操作者ID Returns: 是否删除成功 """ comment_key = f"{SegmentMentionService.COMMENT_KEY_PREFIX}{comment_id}" comment_data = redis_client.get(comment_key) if not comment_data: return False comment = json.loads(comment_data) # 只能删除自己的评论 if comment.get('account_id') != account_id: raise PermissionError("You can only delete your own comments") # 软删除 comment['is_deleted'] = True comment['deleted_at'] = datetime.utcnow().isoformat() redis_client.setex( comment_key, SegmentMentionService.EXPIRATION_SECONDS, json.dumps(comment) ) return True @staticmethod def _create_mention_notifications( comment_id: str, comment_content: str, mentioned_account_ids: list[str], sender_id: str, sender_name: str, dataset_id: str, segment_id: str ): """ 创建@提及通知 Args: comment_id: 评论ID comment_content: 评论内容 mentioned_account_ids: 被@的用户ID列表 sender_id: 发送者ID sender_name: 发送者名称 dataset_id: 知识库ID segment_id: 片段ID """ import logging logger = logging.getLogger(__name__) logger.info(f"[MENTION] Creating notifications for comment {comment_id}") logger.info(f"[MENTION] Mentioned account IDs: {mentioned_account_ids}") logger.info(f"[MENTION] Sender: {sender_name} ({sender_id})") # 获取知识库名称 dataset = db.session.get(Dataset, dataset_id) dataset_name = dataset.name if dataset else 'Unknown Dataset' # 获取片段信息 segment = db.session.get(DocumentSegment, segment_id) segment_preview = segment.content[:50] + '...' if segment and len(segment.content) > 50 else (segment.content if segment else '') for recipient_id in mentioned_account_ids: # 不给自己发通知 if recipient_id == sender_id: logger.info(f"[MENTION] Skipping notification to self: {recipient_id}") continue logger.info(f"[MENTION] Creating notification for recipient: {recipient_id}") notification_id = str(uuid.uuid4()) notification_data = { 'id': notification_id, 'type': 'mention', 'title': f'{sender_name} mentioned you in a comment', 'content': comment_content[:200], # 截取前200字符 'sender_id': sender_id, 'sender_name': sender_name, 'dataset_id': dataset_id, 'dataset_name': dataset_name, 'segment_id': segment_id, 'segment_preview': segment_preview, 'comment_id': comment_id, 'link_url': f'/datasets/{dataset_id}/documents', 'is_read': False, 'created_at': datetime.utcnow().isoformat() } # 存储通知 notification_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{recipient_id}:{notification_id}" redis_client.setex( notification_key, SegmentMentionService.EXPIRATION_SECONDS, json.dumps(notification_data) ) logger.info(f"[MENTION] Stored notification in Redis: {notification_key}") # 添加到用户的通知列表 user_notifications_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{recipient_id}:list" redis_client.lpush(user_notifications_key, notification_id) redis_client.expire(user_notifications_key, SegmentMentionService.EXPIRATION_SECONDS) logger.info(f"[MENTION] Added to user notification list: {user_notifications_key}") # 增加未读计数 unread_count_key = SegmentMentionService.UNREAD_COUNT_KEY.format(account_id=recipient_id) redis_client.incr(unread_count_key) redis_client.expire(unread_count_key, SegmentMentionService.EXPIRATION_SECONDS) logger.info(f"[MENTION] Incremented unread count: {unread_count_key}") # 发布实时通知(用于WebSocket/SSE) redis_client.publish( f'segment_notification:{recipient_id}', json.dumps(notification_data) ) logger.info(f"[MENTION] Published notification to channel: segment_notification:{recipient_id}") @staticmethod def get_user_notifications( account_id: str, is_read: Optional[bool] = None, limit: int = 50 ) -> dict: """ 获取用户的通知列表 Args: account_id: 用户ID is_read: 是否已读(None表示全部) limit: 最大返回数量 Returns: 包含通知列表和未读数量的字典 """ import logging logger = logging.getLogger(__name__) logger.info(f"[NOTIFY-GET] Getting notifications for account: {account_id}, is_read filter: {is_read}, limit: {limit}") user_notifications_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{account_id}:list" notification_ids = redis_client.lrange(user_notifications_key, 0, limit - 1) logger.info(f"[NOTIFY-GET] Found {len(notification_ids)} notification IDs in list: {user_notifications_key}") notifications = [] for notification_id in notification_ids: notification_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{account_id}:{notification_id.decode('utf-8')}" notification_data = redis_client.get(notification_key) logger.info(f"[NOTIFY-GET] Checking key: {notification_key}, exists: {notification_data is not None}") if notification_data: notification = json.loads(notification_data) notification_is_read = notification.get('is_read', False) logger.info(f"[NOTIFY-GET] Notification {notification_id.decode('utf-8')} is_read: {notification_is_read}, filter is_read: {is_read}") # 过滤已读/未读 if is_read is not None and notification_is_read != is_read: logger.info(f"[NOTIFY-GET] Filtering out notification {notification_id.decode('utf-8')} (is_read={notification_is_read}, filter={is_read})") continue notifications.append(notification) logger.info(f"[NOTIFY-GET] Added notification: {notification_id.decode('utf-8')}, is_read: {notification_is_read}") # 获取未读数量 unread_count_key = SegmentMentionService.UNREAD_COUNT_KEY.format(account_id=account_id) unread_count = redis_client.get(unread_count_key) unread_count = int(unread_count) if unread_count else 0 logger.info(f"[NOTIFY-GET] Unread count: {unread_count}, Total notifications: {len(notifications)}") return { 'notifications': notifications, 'unread_count': unread_count, 'total': len(notifications) } @staticmethod def mark_as_read(account_id: str, notification_ids: list[str]) -> bool: """ 标记通知为已读 Args: account_id: 用户ID notification_ids: 通知ID列表 Returns: 是否成功 """ for notification_id in notification_ids: notification_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{account_id}:{notification_id}" notification_data = redis_client.get(notification_key) if notification_data: notification = json.loads(notification_data) # 如果之前未读,现在标记为已读 if not notification.get('is_read'): notification['is_read'] = True notification['read_at'] = datetime.utcnow().isoformat() redis_client.setex( notification_key, SegmentMentionService.EXPIRATION_SECONDS, json.dumps(notification) ) # 减少未读计数 unread_count_key = SegmentMentionService.UNREAD_COUNT_KEY.format(account_id=account_id) current_count = redis_client.get(unread_count_key) if current_count and int(current_count) > 0: redis_client.decr(unread_count_key) return True @staticmethod def mark_all_as_read(account_id: str) -> bool: """ 标记所有通知为已读 Args: account_id: 用户ID Returns: 是否成功 """ user_notifications_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{account_id}:list" notification_ids = redis_client.lrange(user_notifications_key, 0, -1) for notification_id in notification_ids: notification_key = f"{SegmentMentionService.NOTIFICATION_KEY_PREFIX}{account_id}:{notification_id.decode('utf-8')}" notification_data = redis_client.get(notification_key) if notification_data: notification = json.loads(notification_data) if not notification.get('is_read'): notification['is_read'] = True notification['read_at'] = datetime.utcnow().isoformat() redis_client.setex( notification_key, SegmentMentionService.EXPIRATION_SECONDS, json.dumps(notification) ) # 重置未读计数 unread_count_key = SegmentMentionService.UNREAD_COUNT_KEY.format(account_id=account_id) redis_client.set(unread_count_key, 0) redis_client.expire(unread_count_key, SegmentMentionService.EXPIRATION_SECONDS) return True @staticmethod def get_unread_count(account_id: str) -> int: """ 获取用户的未读通知数量 Args: account_id: 用户ID Returns: 未读数量 """ import logging logger = logging.getLogger(__name__) unread_count_key = SegmentMentionService.UNREAD_COUNT_KEY.format(account_id=account_id) logger.info(f"[UNREAD] Getting unread count for account: {account_id}") logger.info(f"[UNREAD] Redis key: {unread_count_key}") unread_count = redis_client.get(unread_count_key) logger.info(f"[UNREAD] Raw value from Redis: {unread_count}") result = int(unread_count) if unread_count else 0 logger.info(f"[UNREAD] Returning count: {result}") return result