geekai/api/handler/admin/moderation_handler.go

package admin

// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
// * Use of this source code is governed by a Apache-2.0 license
// * that can be found in the LICENSE file.
// * @Author yangjian102621@163.com
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import (
	"geekai/core"
	"geekai/core/middleware"
	"geekai/core/types"
	"geekai/handler"
	"geekai/store/model"
	"geekai/utils"
	"geekai/utils/resp"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

type ModerationHandler struct {
	handler.BaseHandler
}

func NewModerationHandler(app *core.AppServer, db *gorm.DB) *ModerationHandler {
	return &ModerationHandler{BaseHandler: handler.BaseHandler{DB: db, App: app}}
}

// RegisterRoutes 注册路由
func (h *ModerationHandler) RegisterRoutes() {
	group := h.App.Engine.Group("/api/admin/moderation/")

	// 需要管理员授权的接口
	group.Use(middleware.AdminAuthMiddleware(h.App.Config.AdminSession.SecretKey, h.App.Redis))
	{
		group.POST("list", h.List)
		group.GET("remove", h.Remove)
		group.POST("batch-remove", h.BatchRemove)
		group.GET("source-list", h.GetSourceList)
	}
}

// List 获取文本审核记录列表
func (h *ModerationHandler) List(c *gin.Context) {
	var data struct {
		Username  string `json:"username"`
		Source    string `json:"source"`
		StartDate string `json:"start_date"`
		EndDate   string `json:"end_date"`
		Page      int    `json:"page"`
		PageSize  int    `json:"page_size"`
	}

	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	session := h.DB.Session(&gorm.Session{})

	// 构建查询条件
	if data.Username != "" {
		// 通过用户名查找用户ID
		var user model.User
		if err := h.DB.Where("username LIKE ?", "%"+data.Username+"%").First(&user).Error; err == nil {
			session = session.Where("user_id", user.Id)
		}
	}

	if data.Source != "" {
		session = session.Where("source", data.Source)
	}

	if data.StartDate != "" && data.EndDate != "" {
		startTime := data.StartDate + " 00:00:00"
		endTime := data.EndDate + " 23:59:59"
		session = session.Where("created_at >= ? AND created_at <= ?", startTime, endTime)
	}

	// 统计总数
	var total int64
	session.Model(&model.Moderation{}).Count(&total)

	// 分页
	page := data.Page
	pageSize := data.PageSize
	if page <= 0 {
		page = 1
	}
	if pageSize <= 0 {
		pageSize = 20
	}

	offset := (page - 1) * pageSize
	session = session.Offset(offset).Limit(pageSize)

	// 查询数据
	var items []model.Moderation
	err := session.Order("id DESC").Find(&items).Error
	if err != nil {
		resp.ERROR(c, err.Error())
		return
	}

	// 获取用户信息
	userIds := make([]uint, 0)
	for _, item := range items {
		userIds = append(userIds, item.UserId)
	}

	var users []model.User
	if len(userIds) > 0 {
		h.DB.Where("id IN ?", userIds).Find(&users)
	}

	userMap := make(map[uint]string)
	for _, user := range users {
		userMap[user.Id] = user.Username
	}

	// 转换为响应数据
	list := make([]map[string]any, 0)
	for _, item := range items {
		var moderation types.ModerationResult
		err := utils.JsonDecode(item.Result, &moderation)
		if err != nil {
			continue
		}
		var result []string
		for value, label := range types.ModerationCategories {
			if moderation.Categories[value] {
				result = append(result, label)
			}
		}
		list = append(list, map[string]any{
			"id":         item.Id,
			"user_id":    item.UserId,
			"username":   userMap[item.UserId],
			"source":     item.Source,
			"input":      item.Input,
			"output":     item.Output,
			"result":     result,
			"created_at": item.CreatedAt.Unix(),
		})
	}

	resp.SUCCESS(c, map[string]any{
		"items":     list,
		"total":     total,
		"page":      page,
		"page_size": pageSize,
	})
}

func (h *ModerationHandler) Remove(c *gin.Context) {
	id := h.GetInt(c, "id", 0)
	if id <= 0 {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	err := h.DB.Where("id", id).Delete(&model.Moderation{}).Error
	if err != nil {
		resp.ERROR(c, err.Error())
		return
	}
	resp.SUCCESS(c)
}

// BatchRemove 批量删除文本审核记录
func (h *ModerationHandler) BatchRemove(c *gin.Context) {
	var data struct {
		Ids []uint `json:"ids"`
	}

	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	if len(data.Ids) == 0 {
		resp.ERROR(c, "请选择要删除的记录")
		return
	}

	err := h.DB.Where("id IN ?", data.Ids).Delete(&model.Moderation{}).Error
	if err != nil {
		resp.ERROR(c, err.Error())
		return
	}

	resp.SUCCESS(c)
}

// 获取 source 列表
func (h *ModerationHandler) GetSourceList(c *gin.Context) {
	sources := []gin.H{
		{
			"id":   types.ModerationSourceChat,
			"name": "AI对话",
		},
		{
			"id":   types.ModerationSourceMJ,
			"name": "Midjourney 绘图",
		},
		{
			"id":   types.ModerationSourceDalle,
			"name": "Dalle 绘图",
		},
		{
			"id":   types.ModerationSourceSD,
			"name": "StableDiffusion 绘图",
		},
		{
			"id":   types.ModerationSourceSuno,
			"name": "Suno 音乐",
		},
		{
			"id":   types.ModerationSourceVideo,
			"name": "视频生成",
		},
		{
			"id":   types.ModerationSourceJiMeng,
			"name": "即梦AI",
		},
	}

	resp.SUCCESS(c, sources)
}