dify-1.8.1/api/controllers/console/datasets/segment_mentions.py

"""
Segment Mentions API
片段@提及和通知API接口
"""

from flask import request
from flask_login import current_user
from flask_restx import Resource, fields, marshal_with
from werkzeug.exceptions import BadRequest, NotFound

from controllers.console import api
from controllers.console.wraps import (
    account_initialization_required,
    setup_required,
)
from libs.login import login_required
from services.dataset_service import DatasetService, DocumentService
from services.segment_mention_service import SegmentMentionService


# Response fields
comment_fields = {
    'id': fields.String,
    'content': fields.String,
    'account_id': fields.String,
    'account_name': fields.String,
    'account_avatar': fields.String,
    'mentioned_account_ids': fields.List(fields.String),
    'created_at': fields.String,
}

comment_list_fields = {
    'comments': fields.List(fields.Nested(comment_fields)),
    'total': fields.Integer,
}

notification_fields = {
    'id': fields.String,
    'type': fields.String,
    'title': fields.String,
    'content': fields.String,
    'sender_id': fields.String,
    'sender_name': fields.String,
    'dataset_id': fields.String,
    'dataset_name': fields.String,
    'segment_id': fields.String,
    'segment_preview': fields.String,
    'link_url': fields.String,
    'is_read': fields.Boolean,
    'created_at': fields.String,
}

notification_list_fields = {
    'notifications': fields.List(fields.Nested(notification_fields)),
    'unread_count': fields.Integer,
    'total': fields.Integer,
}


class SegmentCommentsApi(Resource):
    """片段评论API"""

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(comment_list_fields)
    def get(self, dataset_id: str, segment_id: str):
        """
        获取片段的评论列表

        Args:
            dataset_id: 知识库ID
            segment_id: 片段ID
        """
        dataset_id_str = str(dataset_id)
        segment_id_str = str(segment_id)

        # 权限检查
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        DatasetService.check_dataset_permission(dataset, current_user)

        # 获取评论
        limit = request.args.get('limit', 50, type=int)
        comments = SegmentMentionService.get_segment_comments(
            dataset_id=dataset_id_str,
            segment_id=segment_id_str,
            limit=min(limit, 100)  # 最多100条
        )

        return {
            'comments': comments,
            'total': len(comments)
        }

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(comment_fields)
    def post(self, dataset_id: str, segment_id: str):
        """
        创建评论

        Args:
            dataset_id: 知识库ID
            segment_id: 片段ID

        Request Body:
            - content: 评论内容
            - mentioned_account_ids: @提及的用户ID列表
        """
        dataset_id_str = str(dataset_id)
        segment_id_str = str(segment_id)

        # 权限检查
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        DatasetService.check_dataset_permission(dataset, current_user)

        # 获取请求数据
        data = request.get_json()
        content = data.get('content', '').strip()
        mentioned_account_ids = data.get('mentioned_account_ids', [])

        # 调试日志
        import logging
        logger = logging.getLogger(__name__)
        logger.info(f"Creating comment with mentioned_account_ids: {mentioned_account_ids}")

        if not content:
            raise BadRequest("Comment content cannot be empty.")

        if len(content) > 2000:
            raise BadRequest("Comment content is too long (max 2000 characters).")

        # 创建评论
        comment = SegmentMentionService.create_comment(
            tenant_id=str(dataset.tenant_id),
            dataset_id=dataset_id_str,
            segment_id=segment_id_str,
            content=content,
            account_id=str(current_user.id),
            mentioned_account_ids=mentioned_account_ids
        )

        return comment, 201


class SegmentCommentApi(Resource):
    """单个评论API"""

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, dataset_id: str, segment_id: str, comment_id: str):
        """
        删除评论

        Args:
            dataset_id: 知识库ID
            segment_id: 片段ID
            comment_id: 评论ID
        """
        dataset_id_str = str(dataset_id)
        comment_id_str = str(comment_id)

        # 权限检查
        dataset = DatasetService.get_dataset(dataset_id_str)
        if not dataset:
            raise NotFound("Dataset not found.")

        DatasetService.check_dataset_permission(dataset, current_user)

        # 删除评论
        try:
            SegmentMentionService.delete_comment(
                comment_id=comment_id_str,
                account_id=str(current_user.id)
            )
        except PermissionError as e:
            raise BadRequest(str(e))

        return {'result': 'success'}, 200


class UserNotificationsApi(Resource):
    """用户通知API"""

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(notification_list_fields)
    def get(self):
        """
        获取当前用户的通知列表

        Query Parameters:
            - is_read: 是否已读（true/false，不传表示全部）
            - limit: 返回数量限制（默认50，最大100）
        """
        import logging
        logger = logging.getLogger(__name__)

        is_read_param = request.args.get('is_read')
        is_read = None
        if is_read_param is not None:
            is_read = is_read_param.lower() == 'true'

        limit = request.args.get('limit', 50, type=int)

        logger.info(f"[NOTIFY-API] Request params - is_read_param: {is_read_param}, parsed is_read: {is_read}, limit: {limit}, user_id: {current_user.id}")

        result = SegmentMentionService.get_user_notifications(
            account_id=str(current_user.id),
            is_read=is_read,
            limit=min(limit, 100)
        )

        logger.info(f"[NOTIFY-API] Response - notifications count: {len(result.get('notifications', []))}, unread_count: {result.get('unread_count', 0)}")

        return result

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        """
        标记通知为已读

        Request Body:
            - notification_ids: 通知ID列表（可选）
            - mark_all_as_read: 是否标记全部为已读（布尔值）
        """
        data = request.get_json()

        if data.get('mark_all_as_read'):
            SegmentMentionService.mark_all_as_read(
                account_id=str(current_user.id)
            )
        else:
            notification_ids = data.get('notification_ids', [])
            if notification_ids:
                SegmentMentionService.mark_as_read(
                    account_id=str(current_user.id),
                    notification_ids=notification_ids
                )

        return {'result': 'success'}, 200


class UnreadNotificationsCountApi(Resource):
    """未读通知数量API"""

    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        """获取当前用户的未读通知数量"""
        unread_count = SegmentMentionService.get_unread_count(
            account_id=str(current_user.id)
        )

        return {'unread_count': unread_count}, 200


# 注册路由
api.add_resource(
    SegmentCommentsApi,
    '/datasets/<uuid:dataset_id>/segments/<uuid:segment_id>/comments'
)
api.add_resource(
    SegmentCommentApi,
    '/datasets/<uuid:dataset_id>/segments/<uuid:segment_id>/comments/<uuid:comment_id>'
)
api.add_resource(
    UserNotificationsApi,
    '/datasets/notifications'
)
api.add_resource(
    UnreadNotificationsCountApi,
    '/datasets/notifications/unread-count'
)