new-api-custom/model/token_routing.go

package model

import (
	"fmt"
	"slices"
	"strings"

	"github.com/QuantumNous/new-api/common"
)

const TokenRoutingContextKey = "token_routing_policy"

type ModelRoutingRule struct {
	Model  string            `json:"model"`
	Groups []TokenGroupEntry `json:"groups"`
}

type TokenRoutingPolicy struct {
	Version    int                `json:"version"`
	Mode       string             `json:"mode"`
	Groups     []TokenGroupEntry  `json:"groups"`
	ModelRules []ModelRoutingRule `json:"model_rules"`
}

func (p *TokenRoutingPolicy) Validate(maxGroups int, allowed func(string) bool) error {
	if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user") {
		return fmt.Errorf("invalid routing policy version or mode")
	}
	if p.Mode == "custom" && len(p.Groups) == 0 {
		return fmt.Errorf("a group pool must contain at least one group")
	}
	if p.Mode != "custom" && len(p.Groups) != 0 {
		return fmt.Errorf("inherited pools cannot contain custom groups")
	}
	if len(p.Groups) > maxGroups || len(p.ModelRules) > 128 {
		return fmt.Errorf("routing policy exceeds configured limits")
	}
	pool := map[string]bool{}
	for _, entry := range p.Groups {
		if entry.Group == "" || entry.Group == "auto" || entry.Group != strings.TrimSpace(entry.Group) || pool[entry.Group] || entry.Priority < 0 || entry.Priority > 10000 || !allowed(entry.Group) {
			return fmt.Errorf("invalid or inaccessible pool group: %s", entry.Group)
		}
		pool[entry.Group] = true
	}
	models := map[string]bool{}
	for _, rule := range p.ModelRules {
		if rule.Model == "" || len(rule.Model) > 256 || strings.TrimSpace(rule.Model) != rule.Model || models[rule.Model] || len(rule.Groups) == 0 || len(rule.Groups) > maxGroups {
			return fmt.Errorf("invalid or duplicate model routing rule: %s", rule.Model)
		}
		models[rule.Model] = true
		seen := map[string]bool{}
		for _, entry := range rule.Groups {
			if !pool[entry.Group] || seen[entry.Group] || entry.Priority < 0 || entry.Priority > 10000 {
				return fmt.Errorf("model rule references an invalid pool group: %s", entry.Group)
			}
			seen[entry.Group] = true
		}
	}
	return nil
}

// Resolve is shared by routing and the management preview. Model preferences
// precede remaining defaults; matching priorities stay in the same tier.
func (p *TokenRoutingPolicy) Resolve(modelName string, inherited []string, allowed func(string) bool) []TokenGroupEntry {
	defaults := slices.Clone(p.Groups)
	if p.Mode != "custom" {
		defaults = nil
		for i, group := range inherited {
			defaults = append(defaults, TokenGroupEntry{Group: group, Priority: i})
		}
	}
	slices.SortStableFunc(defaults, func(a, b TokenGroupEntry) int { return a.Priority - b.Priority })
	pool := map[string]bool{}
	for _, entry := range defaults {
		if allowed(entry.Group) {
			pool[entry.Group] = true
		}
	}
	var preferred []TokenGroupEntry
	for _, rule := range p.ModelRules {
		if rule.Model == modelName {
			preferred = slices.Clone(rule.Groups)
			break
		}
	}
	slices.SortStableFunc(preferred, func(a, b TokenGroupEntry) int { return a.Priority - b.Priority })
	result := []TokenGroupEntry{}
	seen := map[string]bool{}
	next := 0
	for _, list := range [][]TokenGroupEntry{preferred, defaults} {
		last, tier := 0, next
		haveLast := false
		for _, entry := range list {
			if !pool[entry.Group] || seen[entry.Group] {
				continue
			}
			if haveLast && last != entry.Priority {
				tier++
			}
			result = append(result, TokenGroupEntry{Group: entry.Group, Priority: tier})
			seen[entry.Group] = true
			last = entry.Priority
			haveLast = true
		}
		if haveLast {
			next = tier + 1
		}
	}
	return result
}

func (token *Token) GetRoutingPolicy() (*TokenRoutingPolicy, error) {
	if token.RoutingPolicy == "" {
		return nil, nil
	}
	var policy TokenRoutingPolicy
	if err := common.UnmarshalJsonStr(token.RoutingPolicy, &policy); err != nil {
		return nil, err
	}
	if policy.Version != 1 {
		return nil, fmt.Errorf("unsupported routing policy version")
	}
	return &policy, nil
}

// LegacyRoutingPolicy preserves effective order, including an explicit old
// pool which already contained the primary group. No model rule is invented.
func (token *Token) LegacyRoutingPolicy() (*TokenRoutingPolicy, error) {
	entries, err := token.GetGroupPool()
	if err != nil {
		return nil, err
	}
	policy := &TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{}, ModelRules: []ModelRoutingRule{}}
	if token.Group == "auto" {
		if len(entries) == 0 {
			policy.Mode = "inherit"
			return policy, nil
		}
		policy.Groups = entries
		return policy, nil
	}
	if token.Group == "" {
		policy.Mode = "user"
		return policy, nil
	}
	if !token.CrossGroupRetry {
		entries = nil
	}
	minPriority := 0
	hasPrimary := false
	for _, entry := range entries {
		minPriority = min(minPriority, entry.Priority)
		hasPrimary = hasPrimary || entry.Group == token.Group
	}
	if !hasPrimary {
		policy.Groups = append(policy.Groups, TokenGroupEntry{Group: token.Group, Priority: minPriority - 1})
	}
	policy.Groups = append(policy.Groups, entries...)
	policy.Groups = policy.Resolve("", nil, func(string) bool { return true })
	return policy, nil
}

func MigrateTokenRoutingPolicies() (int, error) {
	var tokens []Token
	if err := DB.Where("routing_policy IS NULL OR routing_policy = ?", "").Find(&tokens).Error; err != nil {
		return 0, err
	}
	count := 0
	for _, token := range tokens {
		policy, err := token.LegacyRoutingPolicy()
		if err != nil {
			return count, fmt.Errorf("token %d: %w", token.Id, err)
		}
		body, err := common.Marshal(policy)
		if err != nil {
			return count, err
		}
		r := DB.Model(&Token{}).Where("id = ? AND (routing_policy IS NULL OR routing_policy = ?)", token.Id, "").Update("routing_policy", string(body))
		if r.Error != nil {
			return count, r.Error
		}
		count += int(r.RowsAffected)
	}
	return count, nil
}