dify-1.8.1/api/services/segment_mention_service.py

"""
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