274 lines
11 KiB
Go
274 lines
11 KiB
Go
package service
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"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 {
|
|
PolicyFingerprint string `json:"policy_fingerprint"`
|
|
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)
|
|
// Snapshot the routing policy before selection. An older in-flight request
|
|
// may finish after a manual switch, but its binding must not override the
|
|
// new policy on the next request.
|
|
policy := RequestRoutingPolicy(c)
|
|
var fingerprintPool []model.TokenGroupEntry
|
|
if policy.Mode != "smart" {
|
|
fingerprintPool = ResolveRequestPool(c, modelName)
|
|
}
|
|
// Automatic rankings change with live metrics and availability. They must
|
|
// not invalidate a healthy session binding; explicit policy edits still do.
|
|
policyJSON, err := common.Marshal(struct {
|
|
Policy *model.TokenRoutingPolicy
|
|
Pool []model.TokenGroupEntry
|
|
}{policy, fingerprintPool})
|
|
if err != nil {
|
|
return nil, ""
|
|
}
|
|
policyFingerprint := fmt.Sprintf("%x", sha256.Sum256(policyJSON))
|
|
c.Set("pool_affinity_policy_fingerprint", policyFingerprint)
|
|
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 {
|
|
Admin struct {
|
|
Details struct {
|
|
PolicyFingerprint string `json:"routing_policy_fingerprint"`
|
|
} `json:"request_details"`
|
|
} `json:"admin_info"`
|
|
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, PolicyFingerprint: other.Admin.Details.PolicyFingerprint}
|
|
}
|
|
// Unversioned legacy bindings cannot prove they belong to this policy.
|
|
// This check also prevents historical logs from restoring an old route.
|
|
if binding.PolicyFingerprint != policyFingerprint {
|
|
ClearPoolAffinity(c)
|
|
return nil, ""
|
|
}
|
|
for _, entry := range ResolveRequestPool(c, modelName) {
|
|
if entry.Group != binding.Group {
|
|
continue
|
|
}
|
|
if requestGroupDegraded(c, modelName, entry.Group) {
|
|
ClearPoolAffinity(c)
|
|
recordDirectSelection(c, nil, binding.Channel, entry.Group, modelName, "pool_group_degraded")
|
|
return nil, ""
|
|
}
|
|
if poolGroupSkipped(c, modelName, entry.Group) || (policy.Mode != "smart" && poolHigherPriorityRecovered(c, modelName, entry.Group)) {
|
|
ClearPoolAffinity(c)
|
|
return nil, ""
|
|
}
|
|
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 {
|
|
if !claimPoolGroup(c, modelName, entry.Group) {
|
|
return nil, ""
|
|
}
|
|
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
|
|
}
|
|
RecordPoolSuccess(c, c.GetString("original_model"), common.GetContextKeyString(c, constant.ContextKeyAutoGroup))
|
|
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), PolicyFingerprint: c.GetString("pool_affinity_policy_fingerprint")}, 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)
|
|
if policy.Mode == "smart" {
|
|
entries := WithImageGenerationGroup(c, modelName, resolveAutomaticPool(c, modelName))
|
|
if len(policy.ExcludedGroups) == 0 {
|
|
return entries
|
|
}
|
|
filtered := make([]model.TokenGroupEntry, 0, len(entries))
|
|
for _, entry := range entries {
|
|
if !policy.ExcludesGroup(entry.Group) {
|
|
filtered = append(filtered, entry)
|
|
}
|
|
}
|
|
return filtered
|
|
}
|
|
inherited := GetUserAutoGroup(userGroup)
|
|
if policy.Mode == "user" {
|
|
inherited = []string{userGroup}
|
|
}
|
|
return WithImageGenerationGroup(c, modelName, 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})
|
|
// Retry once in the current group, including single-channel groups. Preserve
|
|
// every security/protocol filter; remove only the request's attempted list.
|
|
if retryGroup := param.Ctx.GetString("pool_retry_group"); retryGroup != "" && poolFailures(param.Ctx)[retryGroup] == 1 && !poolGroupSkipped(param.Ctx, param.ModelName, retryGroup) && !requestGroupDegraded(param.Ctx, param.ModelName, retryGroup) {
|
|
param.Ctx.Set("pool_retry_group", "")
|
|
for _, entry := range entries {
|
|
if entry.Group != retryGroup {
|
|
continue
|
|
}
|
|
channel, err := model.GetRandomSatisfiedChannel(retryGroup, param.ModelName, 0, constraints.Filters)
|
|
if err != nil {
|
|
return nil, retryGroup, err
|
|
}
|
|
if channel != nil {
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, retryGroup)
|
|
recordDirectSelection(param.Ctx, channel, channel.Id, retryGroup, param.ModelName, "pool_same_group_retry")
|
|
return channel, retryGroup, nil
|
|
}
|
|
}
|
|
}
|
|
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 poolGroupSkipped(param.Ctx, param.ModelName, group) {
|
|
recordDirectSelection(param.Ctx, nil, 0, group, param.ModelName, "pool_group_cooldown")
|
|
continue
|
|
}
|
|
if requestGroupDegraded(param.Ctx, param.ModelName, group) {
|
|
recordDirectSelection(param.Ctx, nil, 0, group, param.ModelName, "pool_group_degraded")
|
|
continue
|
|
}
|
|
trace := &model.SelectionTrace{}
|
|
channel, err := model.GetRandomSatisfiedChannel(group, param.ModelName, 0, filters, trace)
|
|
if err != nil {
|
|
return nil, group, err
|
|
}
|
|
if channel != nil {
|
|
if !claimPoolGroup(param.Ctx, param.ModelName, group) {
|
|
continue
|
|
}
|
|
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
|
|
}
|
|
if param.Ctx.GetBool("pool_degradation_unavailable") {
|
|
return nil, param.TokenGroup, fmt.Errorf("group quality temporarily unavailable")
|
|
}
|
|
return nil, param.TokenGroup, nil
|
|
}
|