refactor(auth): dedupe route-aware model support checks
This commit is contained in:
@@ -450,6 +450,10 @@ func (m *Manager) selectionModelForAuth(auth *Auth, routeModel string) string {
|
|||||||
return resolvedModel
|
return resolvedModel
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *Manager) selectionModelKeyForAuth(auth *Auth, routeModel string) string {
|
||||||
|
return canonicalModelKey(m.selectionModelForAuth(auth, routeModel))
|
||||||
|
}
|
||||||
|
|
||||||
func (m *Manager) stateModelForExecution(auth *Auth, routeModel, upstreamModel string, pooled bool) string {
|
func (m *Manager) stateModelForExecution(auth *Auth, routeModel, upstreamModel string, pooled bool) string {
|
||||||
stateModel := executionResultModel(routeModel, upstreamModel, pooled)
|
stateModel := executionResultModel(routeModel, upstreamModel, pooled)
|
||||||
selectionModel := m.selectionModelForAuth(auth, routeModel)
|
selectionModel := m.selectionModelForAuth(auth, routeModel)
|
||||||
@@ -507,8 +511,7 @@ func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeMode
|
|||||||
availableByPriority := make(map[int][]*Auth)
|
availableByPriority := make(map[int][]*Auth)
|
||||||
cooldownCount := 0
|
cooldownCount := 0
|
||||||
var earliest time.Time
|
var earliest time.Time
|
||||||
for i := 0; i < len(auths); i++ {
|
for _, candidate := range auths {
|
||||||
candidate := auths[i]
|
|
||||||
checkModel := m.selectionModelForAuth(candidate, routeModel)
|
checkModel := m.selectionModelForAuth(candidate, routeModel)
|
||||||
blocked, reason, next := isAuthBlockedForModel(candidate, checkModel, now)
|
blocked, reason, next := isAuthBlockedForModel(candidate, checkModel, now)
|
||||||
if !blocked {
|
if !blocked {
|
||||||
@@ -555,6 +558,28 @@ func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeMode
|
|||||||
return available, nil
|
return available, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func selectionArgForSelector(selector Selector, routeModel string) string {
|
||||||
|
if isBuiltInSelector(selector) {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return routeModel
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, auth *Auth, routeModel string) bool {
|
||||||
|
if registryRef == nil || auth == nil {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
routeKey := canonicalModelKey(routeModel)
|
||||||
|
if routeKey == "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if registryRef.ClientSupportsModel(auth.ID, routeKey) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
selectionKey := m.selectionModelKeyForAuth(auth, routeModel)
|
||||||
|
return selectionKey != "" && selectionKey != routeKey && registryRef.ClientSupportsModel(auth.ID, selectionKey)
|
||||||
|
}
|
||||||
|
|
||||||
func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) {
|
func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) {
|
||||||
if ch == nil {
|
if ch == nil {
|
||||||
return
|
return
|
||||||
@@ -2353,7 +2378,7 @@ func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) boo
|
|||||||
if auth == nil || strings.TrimSpace(routeModel) == "" {
|
if auth == nil || strings.TrimSpace(routeModel) == "" {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
return canonicalModelKey(m.selectionModelForAuth(auth, routeModel)) != canonicalModelKey(routeModel)
|
return m.selectionModelKeyForAuth(auth, routeModel) != canonicalModelKey(routeModel)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) {
|
func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) {
|
||||||
@@ -2385,18 +2410,9 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op
|
|||||||
if _, used := tried[candidate.ID]; used {
|
if _, used := tried[candidate.ID]; used {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if modelKey != "" && registryRef != nil {
|
if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) {
|
||||||
supportsModel := registryRef.ClientSupportsModel(candidate.ID, modelKey)
|
|
||||||
if !supportsModel {
|
|
||||||
selectionKey := canonicalModelKey(m.selectionModelForAuth(candidate, model))
|
|
||||||
if selectionKey != "" && selectionKey != modelKey {
|
|
||||||
supportsModel = registryRef.ClientSupportsModel(candidate.ID, selectionKey)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if !supportsModel {
|
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
}
|
|
||||||
candidates = append(candidates, candidate)
|
candidates = append(candidates, candidate)
|
||||||
}
|
}
|
||||||
if len(candidates) == 0 {
|
if len(candidates) == 0 {
|
||||||
@@ -2408,7 +2424,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op
|
|||||||
m.mu.RUnlock()
|
m.mu.RUnlock()
|
||||||
return nil, nil, errAvailable
|
return nil, nil, errAvailable
|
||||||
}
|
}
|
||||||
selected, errPick := m.selector.Pick(ctx, provider, "", opts, available)
|
selected, errPick := m.selector.Pick(ctx, provider, selectionArgForSelector(m.selector, model), opts, available)
|
||||||
if errPick != nil {
|
if errPick != nil {
|
||||||
m.mu.RUnlock()
|
m.mu.RUnlock()
|
||||||
return nil, nil, errPick
|
return nil, nil, errPick
|
||||||
@@ -2523,18 +2539,9 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m
|
|||||||
if _, ok := m.executors[providerKey]; !ok {
|
if _, ok := m.executors[providerKey]; !ok {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if modelKey != "" && registryRef != nil {
|
if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) {
|
||||||
supportsModel := registryRef.ClientSupportsModel(candidate.ID, modelKey)
|
|
||||||
if !supportsModel {
|
|
||||||
selectionKey := canonicalModelKey(m.selectionModelForAuth(candidate, model))
|
|
||||||
if selectionKey != "" && selectionKey != modelKey {
|
|
||||||
supportsModel = registryRef.ClientSupportsModel(candidate.ID, selectionKey)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if !supportsModel {
|
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
}
|
|
||||||
candidates = append(candidates, candidate)
|
candidates = append(candidates, candidate)
|
||||||
}
|
}
|
||||||
if len(candidates) == 0 {
|
if len(candidates) == 0 {
|
||||||
@@ -2546,7 +2553,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m
|
|||||||
m.mu.RUnlock()
|
m.mu.RUnlock()
|
||||||
return nil, nil, "", errAvailable
|
return nil, nil, "", errAvailable
|
||||||
}
|
}
|
||||||
selected, errPick := m.selector.Pick(ctx, "mixed", "", opts, available)
|
selected, errPick := m.selector.Pick(ctx, "mixed", selectionArgForSelector(m.selector, model), opts, available)
|
||||||
if errPick != nil {
|
if errPick != nil {
|
||||||
m.mu.RUnlock()
|
m.mu.RUnlock()
|
||||||
return nil, nil, "", errPick
|
return nil, nil, "", errPick
|
||||||
|
|||||||
Reference in New Issue
Block a user