new-api-custom/controller/chain_detection.go

package controller

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"net/http"
	"os"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

var chainEnqueue sync.Mutex
var chainKinds = []string{"fingerprint", "ciallo", "vision"}

// All three child tasks are committed together, using the existing durable ledger.
func enqueueChainTest(ctx context.Context, group string, channelID int) ([]model.DetectionRun, error) {
	chainEnqueue.Lock()
	defer chainEnqueue.Unlock()
	var runs []model.DetectionRun
	err := model.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var active []model.DetectionRun
		if err := tx.Where("group_name = ? AND request_key LIKE ? AND settled = ? AND deadline > ?", group, "chain:%", false, time.Now().Unix()).Order("created_at asc").Limit(1).Find(&active).Error; err != nil {
			return err
		}
		if len(active) > 0 {
			id, _, _ := strings.Cut(active[0].ID, "_")
			return tx.Where("id IN ?", []string{id + "_fingerprint", id + "_ciallo", id + "_vision"}).Find(&runs).Error
		}
		id, err := common.GenerateRandomCharsKey(32)
		if err != nil {
			return err
		}
		now := time.Now().Unix()
		for _, kind := range chainKinds {
			runs = append(runs, model.DetectionRun{ID: id + "_" + kind, RequestKey: "chain:" + id + ":" + kind, ChannelID: channelID, GroupName: group, Kind: kind, Status: "pending", CreatedAt: now, Deadline: now + 600})
		}
		return tx.Create(&runs).Error
	})
	return runs, err
}

func BotChainTest(c *gin.Context) {
	c.Header("Cache-Control", "no-store")
	key := os.Getenv("BOT_BALANCE_KEY")
	derived := sha256.Sum256([]byte("newapi-chaintest-v1:" + key))
	raw, bearer := strings.CutPrefix(c.GetHeader("Authorization"), "Bearer ")
	expected, supplied := sha256.Sum256([]byte(hex.EncodeToString(derived[:]))), sha256.Sum256([]byte(raw))
	if len(key) < 32 || !bearer || subtle.ConstantTimeCompare(expected[:], supplied[:]) != 1 {
		c.JSON(http.StatusUnauthorized, gin.H{"success": false, "code": "unauthorized"})
		return
	}
	var runs []model.DetectionRun
	if c.Request.Method == http.MethodPost {
		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 1024)
		var input struct {
			Group string `json:"group"`
		}
		if c.ShouldBindJSON(&input) != nil || strings.TrimSpace(input.Group) == "" || len(input.Group) > 128 {
			c.JSON(400, gin.H{"success": false, "code": "invalid_group"})
			return
		}
		input.Group = strings.TrimSpace(input.Group)
		channels, err := service.AstraSamplingChannels(c.Request.Context())
		if err != nil {
			c.JSON(503, gin.H{"success": false, "code": "unavailable"})
			return
		}
		slices.SortFunc(channels, func(a, b model.Channel) int { return a.Id - b.Id })
		selected := 0
		for _, channel := range channels {
			if slices.Contains(channel.GetGroups(), input.Group) {
				selected = channel.Id
				break
			}
		}
		if selected == 0 {
			c.JSON(404, gin.H{"success": false, "code": "group_unavailable"})
			return
		}
		runs, err = enqueueChainTest(c.Request.Context(), input.Group, selected)
		if err != nil {
			c.JSON(503, gin.H{"success": false, "code": "unavailable"})
			return
		}
		for _, run := range runs {
			go processDetection(context.Background(), run)
		}
	} else {
		id := c.Param("id")
		if len(id) != 32 || strings.ContainsAny(id, "_%") {
			c.JSON(404, gin.H{"success": false, "code": "not_found"})
			return
		}
		err := model.DB.WithContext(c.Request.Context()).Where("id IN ? AND request_key LIKE ?", []string{id + "_fingerprint", id + "_ciallo", id + "_vision"}, "chain:"+id+":%").Find(&runs).Error
		if err != nil {
			c.JSON(503, gin.H{"success": false, "code": "unavailable"})
			return
		}
	}
	data, err := chainTestResult(runs, time.Now().Unix())
	if err != nil {
		c.JSON(404, gin.H{"success": false, "code": "not_found"})
		return
	}
	c.JSON(200, gin.H{"success": true, "data": data})
}

func chainTestResult(runs []model.DetectionRun, now int64) (gin.H, error) {
	if len(runs) != 3 {
		return nil, errors.New("incomplete chain")
	}
	id, _, _ := strings.Cut(runs[0].ID, "_")
	results := make(map[string]any, 3)
	done := true
	for _, run := range runs {
		run = service.NormalizeVisionRun(run)
		state, text, match := "running", "", ""
		if run.Settled {
			state = "failed"
			var result fingerprint.Result
			if run.Status == "succeeded" && common.UnmarshalJsonStr(run.Result, &result) == nil {
				text = result.Closest
				if run.Kind == "fingerprint" {
					match = service.AstraSamplingResultStatus(result)
					state = match
					if state != "match" && state != "mismatch" && state != "inconclusive" {
						state = "failed"
					}
				} else if strings.TrimSpace(text) != "" {
					state = service.VisionQuality(text)
					if run.Kind == "ciallo" {
						state = service.CialloQuality(text)
					}
				}
			}
		} else if now >= run.Deadline {
			state = "failed"
		} else {
			done = false
		}
		results[run.Kind] = gin.H{"status": state, "text": text}
	}
	return gin.H{"id": id, "group": runs[0].GroupName, "deadline": runs[0].Deadline, "done": done, "results": results}, nil
}