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