new-api-custom/service/token_routing.go

package service

import (
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

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

type poolBinding struct {
	Channel int    `json:"channel"`
	Group   string `json:"group"`
}

var poolBindingOnce sync.Once
var poolBindingCache *cachex.HybridCache[poolBinding]

func poolBindings() *cachex.HybridCache[poolBinding] {
	poolBindingOnce.Do(func() {
		poolBindingCache = cachex.NewHybridCache[poolBinding](cachex.HybridCacheConfig[poolBinding]{Namespace: cachex.Namespace("new-api:pool-affinity:v1"), Redis: common.RDB, RedisEnabled: func() bool { return common.RedisEnabled && common.RDB != nil }, RedisCodec: cachex.JSONCodec[poolBinding]{}, Memory: func() *hot.HotCache[string, poolBinding] {
			return hot.NewHotCache[string, poolBinding](hot.LRU, 100000).WithTTL(time.Hour).Build()
		}})
	})
	return poolBindingCache
}

func PreferredPoolChannel(c *gin.Context, modelName string) (*model.Channel, string) {
	// Resolve the existing session-key rules, but isolate the binding by owner,
	// token and original model rather than the currently selected group.
	GetPreferredChannelByAffinity(c, modelName, "pool")
	meta, ok := getChannelAffinityMeta(c)
	if !ok || RequestPolicy(c).SessionMode == "off" {
		return nil, ""
	}
	key := common.GenerateHMAC(fmt.Sprintf("%d:%d:%d:%s:%s", c.GetInt("id"), c.GetInt("token_id"), len(modelName), modelName, meta.CacheKey))
	c.Set("pool_affinity_key", key)
	c.Set("pool_affinity_ttl", meta.TTLSeconds)
	binding, found, err := poolBindings().Get(key)
	if err != nil {
		return nil, ""
	}
	if !found {
		// Adopt only a successful old binding attributable to this exact owner,
		// token, model and session. Never trust a legacy unscoped cache alone.
		appendRequestDetails(c, nil, model.NewLogOther())
		correlation := c.GetString("log_correlation_key")
		if correlation == "" || model.LOG_DB == nil {
			return nil, ""
		}
		var previous model.Log
		if model.LOG_DB.Where("user_id = ? AND token_id = ? AND model_name = ? AND correlation_key = ? AND type = ? AND created_at > ?", c.GetInt("id"), c.GetInt("token_id"), modelName, correlation, model.LogTypeConsume, time.Now().Unix()-int64(meta.TTLSeconds)).Order("id desc").First(&previous).Error != nil {
			return nil, ""
		}
		var other struct {
			Result string `json:"model_guard_result"`
			Stream struct {
				Status    string `json:"status"`
				EndReason string `json:"end_reason"`
				Response  string `json:"response_status"`
			} `json:"stream_status"`
		}
		_ = common.UnmarshalJsonStr(previous.Other, &other)
		if other.Result != "" {
			return nil, ""
		}
		if previous.IsStream && (other.Stream.Status != "ok" || (other.Stream.EndReason != "done" && other.Stream.EndReason != "eof") || (other.Stream.Response != "" && other.Stream.Response != "completed")) {
			return nil, ""
		}
		binding = poolBinding{Channel: previous.ChannelId, Group: previous.Group}
	}
	for _, entry := range ResolveRequestPool(c, modelName) {
		if entry.Group != binding.Group {
			continue
		}
		channel, err := model.CacheGetChannel(binding.Channel)
		if err == nil && channel.Status == common.ChannelStatusEnabled && model.IsChannelEnabledForGroupModel(entry.Group, modelName, channel.Id) {
			if ok, _ := model.ChannelSatisfiesFilters(channel, modelName, GetChannelConstraints(c).Filters); ok {
				common.SetContextKey(c, constant.ContextKeyAutoGroup, entry.Group)
				MarkChannelAffinityUsed(c, entry.Group, channel.Id)
				recordDirectSelection(c, channel, channel.Id, entry.Group, modelName, "pool_affinity_hit")
				return channel, entry.Group
			}
		}
	}
	ClearPoolAffinity(c)
	return nil, ""
}

func ClearPoolAffinity(c *gin.Context) {
	if key := c.GetString("pool_affinity_key"); key != "" {
		ttl := c.GetInt("pool_affinity_ttl")
		if ttl <= 0 {
			ttl = 3600
		}
		// A tombstone prevents re-adopting the same historical binding after a
		// rejection. A later successful generation replaces it normally.
		_ = poolBindings().SetWithTTL(key, poolBinding{}, time.Duration(ttl)*time.Second)
	}
	ClearCurrentChannelAffinityCache(c)
}

func RecordPoolAffinity(c *gin.Context) {
	if !RequestPolicy(c).Successful {
		return
	}
	key := c.GetString("pool_affinity_key")
	if key == "" {
		return
	}
	ttl := c.GetInt("pool_affinity_ttl")
	if ttl <= 0 {
		ttl = 3600
	}
	_ = poolBindings().SetWithTTL(key, poolBinding{Channel: c.GetInt("channel_id"), Group: common.GetContextKeyString(c, constant.ContextKeyAutoGroup)}, time.Duration(ttl)*time.Second)
}

func RequestRoutingPolicy(c *gin.Context) *model.TokenRoutingPolicy {
	value, _ := c.Get(model.TokenRoutingContextKey)
	policy, _ := value.(*model.TokenRoutingPolicy)
	return policy
}

func ResolveRequestPool(c *gin.Context, modelName string) []model.TokenGroupEntry {
	policy := RequestRoutingPolicy(c)
	if policy == nil {
		return nil
	}
	userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
	inherited := GetUserAutoGroup(userGroup)
	if policy.Mode == "user" {
		inherited = []string{userGroup}
	}
	return policy.Resolve(modelName, inherited, func(group string) bool { return IsUserSelectableGroup(userGroup, group) })
}

func selectTokenRoutingPool(param *RetryParam) (*model.Channel, string, error) {
	entries := ResolveRequestPool(param.Ctx, param.ModelName)
	constraints := GetChannelConstraints(param.Ctx)
	filters := append([]dto.ChannelFilter{}, constraints.Filters...)
	var excluded []int
	for _, value := range param.Ctx.GetStringSlice("use_channel") {
		id, _ := strconv.Atoi(value)
		excluded = append(excluded, id)
	}
	filters = append(filters, dto.ChannelFilter{Kind: dto.FilterAttemptedChannel, ExcludedChannelIDs: excluded})
	for start := 0; start < len(entries); {
		end := start + 1
		for end < len(entries) && entries[end].Priority == entries[start].Priority {
			end++
		}
		groups := entries[start:end]
		var names []string
		for _, entry := range groups {
			names = append(names, entry.Group)
		}
		offset := nextTierIndex(fmt.Sprintf("pool:%d:%s:%s", param.Ctx.GetInt("token_id"), param.ModelName, strings.Join(names, ",")), len(groups))
		for i := range len(groups) {
			group := groups[(offset+i)%len(groups)].Group
			if !IsGroupAvailable(group) {
				continue
			}
			trace := &model.SelectionTrace{}
			channel, err := model.GetRandomSatisfiedChannel(group, param.ModelName, 0, filters, trace)
			if err != nil {
				return nil, group, err
			}
			if channel != nil {
				common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
				recordDirectSelection(param.Ctx, channel, channel.Id, group, param.ModelName, "model_pool_priority")
				return channel, group, nil
			}
		}
		start = end
	}
	return nil, param.TokenGroup, nil
}