feat: whitelist of ai_summary and add trace_id

Author: Crnaneo

apis/hole/apis.go

@@ -14,6 +14,7 @@ import (
 	"treehole_next/utils/sensitive"
 
 	"github.com/gofiber/fiber/v2"
+	"github.com/google/uuid"
 	"github.com/opentreehole/go-common"
 	"github.com/rs/zerolog/log"
 	"gorm.io/gorm"
@@ -951,7 +952,13 @@ func DeleteHole(c *fiber.Ctx) error {
 	return c.Status(204).JSON(nil)
 }
 func GenerateSummary(c *fiber.Ctx) error {
-
+	uid, _ := common.GetUserID(c)
+	if config.Config.WhiteListUserIds != nil && !slices.Contains(config.Config.WhiteListUserIds, uid) {
+		return c.Status(404).JSON(fiber.Map{
+			"code":    404,
+			"message": "Cannot GET " + c.Path(),
+		})
+	}
 	id, _ := c.ParamsInt("id")
 	forceRefresh := c.QueryBool("force_refresh")
 	var cachedData Summary
@@ -1041,6 +1048,7 @@ func GenerateSummary(c *fiber.Ctx) error {
 					if cachedData.Data.HoleID != id && cachedData.Data.HoleID != 0 {
 						log.Error().
 							Int("hole_id", id).
+							Str("trace_id", cachedData.TraceID).
 							Int("ai_server_return_id", cachedData.Data.HoleID).
 							Msg("AISummary: hole id error")
 						cachedData.Data.HoleID = id
@@ -1099,9 +1107,10 @@ func GenerateSummary(c *fiber.Ctx) error {
 	}
 
 	requestBody := map[string]any{
-		"floors":  summaryFloors,
-		"content": content,
-		"hole_id": hole.ID,
+		"floors":   summaryFloors,
+		"content":  content,
+		"hole_id":  hole.ID,
+		"trace_id": uuid.NewString(),
 	}
 
 	requestJSON, err := json.Marshal(requestBody)
@@ -1128,6 +1137,7 @@ func GenerateSummary(c *fiber.Ctx) error {
 	if resp.StatusCode != http.StatusOK {
 		log.Error().
 			Str("url", config.Config.AISummaryURL+"/generate_summary").
+			Str("trace_id", requestBody["trace_id"].(string)).
 			Int("status", resp.StatusCode).
 			Str("req_body_base64", func() string {
 				if len(requestJSON) > config.Config.SummaryLogLimit {
@@ -1151,6 +1161,7 @@ func GenerateSummary(c *fiber.Ctx) error {
 	if err != nil {
 		log.Error().
 			Str("url", config.Config.AISummaryURL+"/generate_summary").
+			Str("trace_id", requestBody["trace_id"].(string)).
 			Int("status", resp.StatusCode).
 			Str("req_body_base64", func() string {
 				if len(requestJSON) > config.Config.SummaryLogLimit {
@@ -1184,6 +1195,7 @@ func GenerateSummary(c *fiber.Ctx) error {
 		response.Message = errCode2Message[response.Code]
 	default:
 		log.Error().Str("url", config.Config.AISummaryURL+"/generate_summary").
+			Str("trace_id", requestBody["trace_id"].(string)).
 			Int("status", resp.StatusCode).
 			Str("req_body_base64", func() string {
 				if len(requestJSON) > config.Config.SummaryLogLimit {

apis/hole/schemas.go

@@ -133,6 +133,7 @@ func (body ModifyModel) DoNothing() bool {
 
 type Summary struct {
 	Code    int    `json:"code"`
+	TraceID string `json:"trace_id"`
 	Message string `json:"message"`
 	Data    struct {
 		HoleID   int    `json:"hole_id"`

config/config.go

@@ -60,6 +60,7 @@ var Config struct {
 	SummaryContentLimit          int64    `env:"SUMMARY_CONTENT_LIMIT" envDefault:"500"`
 	SummarySteps                 int      `env:"SUMMARY_STEPS" envDefault:"5"`
 	SummaryLogLimit              int      `env:"SUMMARY_LOG_LIMIT" envDefault:"1000"`
+	WhiteListUserIds             []int    `env:"WHITE_LIST_USER_IDS"`
 }
 
 var DynamicConfig struct {

models/hole.go

@@ -2,7 +2,6 @@ package models
 
 import (
 	"fmt"
-	"strconv"
 	"time"
 
 	"golang.org/x/exp/maps"
@@ -276,7 +275,8 @@ func (holes Holes) Preprocess(c *fiber.Ctx) error {
 		hole.SetHoleFloor()
 		floors = append(floors, hole.Floors...)
 		// set ai_summary_available
-		hole.AISummaryAvailable = true
+		uid, _ := common.GetUserID(c)
+		hole.AISummaryAvailable = config.Config.WhiteListUserIds == nil || slices.Contains(config.Config.WhiteListUserIds, uid)
 		for _, tag := range hole.Tags {
 			if len(tag.Name) > 0 && tag.Name[0] == '*' {
 				hole.AISummaryAvailable = false
@@ -294,14 +294,11 @@ func (holes Holes) Preprocess(c *fiber.Ctx) error {
 					return err
 				}
 
-				var discard any
-				hole.AISummaryAvailable = hole.Reply > config.Config.SummaryFloorLimit || contentSum >= config.Config.SummaryContentLimit || utils.GetCache("AISummary"+strconv.Itoa(hole.ID), &discard)
+				// var discard any
+				hole.AISummaryAvailable = hole.Reply > config.Config.SummaryFloorLimit || contentSum >= config.Config.SummaryContentLimit // || utils.GetCache("AISummary"+strconv.Itoa(hole.ID), &discard)
 
 				if hole.AISummaryAvailable {
-					err = query.Where("is_sensitive = ?", true).Count(&sensitiveCount).Error
-					if err != nil {
-						return err
-					}
+					query.Where("is_sensitive = ? AND is_actual_sensitive IS NULL", true).Count(&sensitiveCount)
 					if sensitiveCount > 0 {
 						hole.AISummaryAvailable = false
 					}