fix(phase-2): add thread-safety to ModelRegistry and Router reload operations
This commit is contained in:
@@ -1,6 +1,9 @@
|
||||
package models
|
||||
|
||||
import "strings"
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// CanonicalProviders lists the original model creators in priority order.
|
||||
// When a model name exists in multiple providers (e.g. deepseek-v4-pro in
|
||||
@@ -22,6 +25,7 @@ var CanonicalProviders = []string{
|
||||
}
|
||||
|
||||
type ModelRegistry struct {
|
||||
mu sync.RWMutex
|
||||
Providers map[string]ProviderInfo `json:"-"`
|
||||
}
|
||||
|
||||
@@ -173,6 +177,12 @@ func (r *ModelRegistry) findAllForwardFuzzy(modelID string) (*ModelMetadata, boo
|
||||
// etc.) from overriding the original provider's authoritative pricing and
|
||||
// limits.
|
||||
func (r *ModelRegistry) FindModel(modelID string) *ModelMetadata {
|
||||
if r == nil {
|
||||
return nil
|
||||
}
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
// 1. Exact key match — canonical first, then all
|
||||
if m, ok := r.findInCanonical(modelID); ok {
|
||||
return m
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"gophergate/internal/db"
|
||||
)
|
||||
@@ -32,6 +33,7 @@ type ClassifierFunc func(ctx context.Context, selectorModel, systemPrompt, userM
|
||||
|
||||
// Router resolves model groups to concrete models.
|
||||
type Router struct {
|
||||
mu sync.RWMutex
|
||||
groups map[string]db.ModelGroup
|
||||
classify ClassifierFunc
|
||||
}
|
||||
@@ -50,6 +52,8 @@ func New(groups []db.ModelGroup, classify ClassifierFunc) *Router {
|
||||
|
||||
// Groups returns all registered model group IDs.
|
||||
func (r *Router) Groups() []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
ids := make([]string, 0, len(r.groups))
|
||||
for id := range r.groups {
|
||||
ids = append(ids, id)
|
||||
@@ -59,13 +63,17 @@ func (r *Router) Groups() []string {
|
||||
|
||||
// IsGroup returns true if the model name is a group ID.
|
||||
func (r *Router) IsGroup(modelID string) bool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
_, ok := r.groups[modelID]
|
||||
return ok
|
||||
}
|
||||
|
||||
// Route resolves a group to a concrete model.
|
||||
func (r *Router) Route(ctx context.Context, groupID string, routeCtx *RouteContext) (*Decision, error) {
|
||||
r.mu.RLock()
|
||||
group, ok := r.groups[groupID]
|
||||
r.mu.RUnlock()
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown model group: %s", groupID)
|
||||
}
|
||||
@@ -133,8 +141,12 @@ func (r *Router) RouteToConcrete(ctx context.Context, modelID string, routeCtx *
|
||||
|
||||
// Reload replaces the group definitions without recreating the router.
|
||||
func (r *Router) Reload(groups []db.ModelGroup) {
|
||||
r.groups = make(map[string]db.ModelGroup)
|
||||
newGroups := make(map[string]db.ModelGroup)
|
||||
for _, g := range groups {
|
||||
r.groups[g.ID] = g
|
||||
newGroups[g.ID] = g
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.groups = newGroups
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
@@ -628,7 +628,10 @@ func (s *Server) handleChatCompletions(c *gin.Context) {
|
||||
|
||||
// Inject or cap max_tokens from model registry.
|
||||
s.registryMu.RLock()
|
||||
meta := s.registry.FindModel(modelID)
|
||||
var meta *models.ModelMetadata
|
||||
if s.registry != nil {
|
||||
meta = s.registry.FindModel(modelID)
|
||||
}
|
||||
s.registryMu.RUnlock()
|
||||
|
||||
if meta != nil && meta.Limit != nil && meta.Limit.Output > 0 {
|
||||
|
||||
Reference in New Issue
Block a user