feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -0,0 +1,185 @@
|
||||
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
|
||||
}
|
||||
Reference in new issue
Block a user