mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-19 04:40:52 +08:00
feat(selector): add priority support for auth selection
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -103,13 +104,29 @@ func (e *modelCooldownError) Headers() http.Header {
|
||||
return headers
|
||||
}
|
||||
|
||||
func collectAvailable(auths []*Auth, model string, now time.Time) (available []*Auth, cooldownCount int, earliest time.Time) {
|
||||
available = make([]*Auth, 0, len(auths))
|
||||
func authPriority(auth *Auth) int {
|
||||
if auth == nil || auth.Attributes == nil {
|
||||
return 0
|
||||
}
|
||||
raw := strings.TrimSpace(auth.Attributes["priority"])
|
||||
if raw == "" {
|
||||
return 0
|
||||
}
|
||||
parsed, err := strconv.Atoi(raw)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (available map[int][]*Auth, cooldownCount int, earliest time.Time) {
|
||||
available = make(map[int][]*Auth)
|
||||
for i := 0; i < len(auths); i++ {
|
||||
candidate := auths[i]
|
||||
blocked, reason, next := isAuthBlockedForModel(candidate, model, now)
|
||||
if !blocked {
|
||||
available = append(available, candidate)
|
||||
priority := authPriority(candidate)
|
||||
available[priority] = append(available[priority], candidate)
|
||||
continue
|
||||
}
|
||||
if reason == blockReasonCooldown {
|
||||
@@ -119,9 +136,6 @@ func collectAvailable(auths []*Auth, model string, now time.Time) (available []*
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(available) > 1 {
|
||||
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
||||
}
|
||||
return available, cooldownCount, earliest
|
||||
}
|
||||
|
||||
@@ -130,18 +144,35 @@ func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]
|
||||
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
|
||||
}
|
||||
|
||||
available, cooldownCount, earliest := collectAvailable(auths, model, now)
|
||||
if len(available) == 0 {
|
||||
availableByPriority, cooldownCount, earliest := collectAvailableByPriority(auths, model, now)
|
||||
if len(availableByPriority) == 0 {
|
||||
if cooldownCount == len(auths) && !earliest.IsZero() {
|
||||
providerForError := provider
|
||||
if providerForError == "mixed" {
|
||||
providerForError = ""
|
||||
}
|
||||
resetIn := earliest.Sub(now)
|
||||
if resetIn < 0 {
|
||||
resetIn = 0
|
||||
}
|
||||
return nil, newModelCooldownError(model, provider, resetIn)
|
||||
return nil, newModelCooldownError(model, providerForError, resetIn)
|
||||
}
|
||||
return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
|
||||
}
|
||||
|
||||
bestPriority := 0
|
||||
found := false
|
||||
for priority := range availableByPriority {
|
||||
if !found || priority > bestPriority {
|
||||
bestPriority = priority
|
||||
found = true
|
||||
}
|
||||
}
|
||||
|
||||
available := availableByPriority[bestPriority]
|
||||
if len(available) > 1 {
|
||||
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
||||
}
|
||||
return available, nil
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user