new-api-custom/service/automatic_group.go

package service

import (
	"cmp"
	"context"
	"encoding/json"
	"math"
	"slices"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
	"github.com/gin-gonic/gin"
	"github.com/samber/hot"
)

var automaticMetricCache = hot.NewHotCache[string, perfmetrics.QueryResult](hot.LRU, 1000).WithTTL(time.Minute).Build()
var automaticMetricMu sync.Mutex

func automaticModelMetrics(modelName string) perfmetrics.QueryResult {
	automaticMetricMu.Lock()
	defer automaticMetricMu.Unlock()
	if value, found, err := automaticMetricCache.Get(modelName); err == nil && found {
		return value
	}
	value, err := perfmetrics.Query(perfmetrics.QueryParams{Model: modelName, Hours: 24})
	if err == nil {
		automaticMetricCache.Set(modelName, value)
	}
	return value
}

type automaticCandidate struct {
	group                                  string
	cost, ttft, throughput, success, ratio float64
}

func resolveAutomaticPool(c *gin.Context, modelName string) []model.TokenGroupEntry {
	// Freeze ranking within a request: retries must not reorder their own pool.
	key := "automatic_pool:" + modelName
	if cached, ok := c.Get(key); ok {
		return cached.([]model.TokenGroupEntry)
	}
	result := []model.TokenGroupEntry{}
	userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
	var abilities []model.Ability
	query := model.DB.WithContext(c.Request.Context()).Where("enabled = ?", true)
	if modelName != "" {
		query = query.Where("model = ?", modelName)
	}
	if err := query.Find(&abilities).Error; err != nil {
		return result
	}
	allowed := GetUserUsableGroups(userGroup)
	names := []string{}
	seen := map[string]bool{}
	for _, ability := range abilities {
		group := ability.Group
		if _, ok := allowed[group]; !ok || seen[group] || !IsUserSelectableGroup(userGroup, group) || RequestRoutingPolicy(c).ExcludesGroup(group) {
			continue
		}
		channel, err := model.CacheGetChannel(ability.ChannelId)
		if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled {
			continue
		}
		seen[group] = true
		names = append(names, group)
	}
	if len(names) == 0 {
		c.Set(key, result)
		return result
	}
	var performance perfmetrics.QueryResult
	if modelName != "" {
		performance = automaticModelMetrics(modelName)
	}
	testSpeeds := map[int]float64{}
	if RequestRoutingPolicy(c).AutoPreference == "speed" {
		testSpeeds = latestModelTestSpeeds(c.Request.Context(), modelName)
	}
	metrics := map[string]perfmetrics.GroupResult{}
	for _, item := range performance.Groups {
		metrics[item.Group] = item
	}
	candidates := []automaticCandidate{}
	for _, group := range names {
		if requestGroupDegraded(c, modelName, group) {
			continue
		}
		ratio := model.EffectiveGroupRatio(c, userGroup, group).GroupRatio
		item := automaticCandidate{group: group, ttft: math.Inf(1), success: -1, ratio: ratio}
		if RequestRoutingPolicy(c).AutoPreference == "throughput" {
			metric, err := GroupModelPerformance(c.Request.Context(), group, modelName)
			if err != nil {
				common.SysError("automatic group performance query failed: " + err.Error())
			} else if metric.Throughput != nil {
				item.throughput = metric.Throughput.Value
			}
		}
		if metric, ok := metrics[group]; ok && len(metric.Series) > 0 {
			item.success = metric.SuccessRate

		}
		total, count := 0.0, 0
		for _, ability := range abilities {
			if ability.Group != group {
				continue
			}
			channel, err := model.CacheGetChannel(ability.ChannelId)
			if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled {
				continue
			}
			if speed, ok := testSpeeds[ability.ChannelId]; ok {
				total += speed
				count++
			}
		}
		if count > 0 {
			item.ttft = total / float64(count)
		}
		candidates = append(candidates, item)
	}
	sortAutomaticCandidates(candidates, RequestRoutingPolicy(c).AutoPreference)
	for priority, item := range candidates {
		result = append(result, model.TokenGroupEntry{Group: item.group, Priority: priority})
	}
	c.Set(key, result)
	return result
}

func sortAutomaticCandidates(items []automaticCandidate, preference string) {
	slices.SortFunc(items, func(a, b automaticCandidate) int {
		if preference == "throughput" {
			if n := cmp.Compare(b.throughput, a.throughput); n != 0 {
				return n
			}
		}
		if preference == "speed" {
			if n := cmp.Compare(a.ttft, b.ttft); n != 0 {
				return n
			}
		}
		if preference == "stable" {
			if n := cmp.Compare(b.success, a.success); n != 0 {
				return n
			}
		}
		// Effective customer multipliers are independent of traffic sample size.
		if n := cmp.Compare(a.ratio, b.ratio); n != 0 {
			return n
		}
		if n := cmp.Compare(b.success, a.success); n != 0 {
			return n
		}
		if n := cmp.Compare(a.ttft, b.ttft); n != 0 {
			return n
		}
		if n := cmp.Compare(a.ratio, b.ratio); n != 0 {
			return n
		}
		return cmp.Compare(a.group, b.group)
	})
}

// AutomaticPoolModels is the discoverable model union, independent of ranking,
// temporary cooldowns and quality checks used when routing individual requests.
func AutomaticPoolModels(c *gin.Context) []string {
	userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
	var abilities []model.Ability
	if err := model.DB.WithContext(c.Request.Context()).Where("enabled = ?", true).Find(&abilities).Error; err != nil {
		common.SysError("automatic pool model discovery failed: " + err.Error())
		return []string{}
	}
	names := map[string]bool{}
	for _, ability := range abilities {
		if !IsUserSelectableGroup(userGroup, ability.Group) || RequestRoutingPolicy(c).ExcludesGroup(ability.Group) {
			continue
		}
		channel, err := model.CacheGetChannel(ability.ChannelId)
		if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled {
			continue
		}
		names[ability.Model] = true
	}
	result := make([]string, 0, len(names))
	for name := range names {
		result = append(result, name)
	}
	slices.Sort(result)
	return result
}

// Model tests use token_id=0 and a server-owned content marker. User traffic and
// another model's timing must not masquerade as a speed measurement.
func latestModelTestSpeeds(ctx context.Context, modelName string) map[int]float64 {
	result := map[int]float64{}
	if model.LOG_DB == nil {
		return result
	}
	var rows []model.Log
	err := model.LOG_DB.WithContext(ctx).Select("id", "channel_id", "other").Where("type = ? AND token_id = ? AND content = ? AND model_name = ? AND is_stream = ? AND created_at >= ?", model.LogTypeConsume, 0, "模型测试", modelName, true, time.Now().Add(-time.Hour).Unix()).Order("created_at desc, id desc").Find(&rows).Error
	if err != nil {
		common.SysError("automatic routing test measurements unavailable: " + err.Error())
		return result
	}
	for _, row := range rows {
		if _, ok := result[row.ChannelId]; ok {
			continue
		}
		var data struct {
			FRT float64 `json:"frt"`
		}
		if json.Unmarshal([]byte(row.Other), &data) != nil || data.FRT <= 0 || math.IsInf(data.FRT, 0) || math.IsNaN(data.FRT) {
			continue
		}
		result[row.ChannelId] = data.FRT
	}
	return result
}