mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
yinwm
a73d8e1a16
- Add ModelConfig struct with protocol prefix support (openai/, anthropic/, etc.) - Implement GetModelConfig with round-robin load balancing - Add CreateProviderFromConfig factory for protocol-based routing - Add ModelRegistry for thread-safe endpoint selection - Maintain full backward compatibility with legacy providers config - Update README.md and README.zh.md with model_list documentation - Add migration guide at docs/migration/model-list-migration.md Supported protocols: openai, anthropic, antigravity, claude-cli, codex-cli, github-copilot, openrouter, groq, deepseek, cerebras, qwen, zhipu, gemini Closes #283 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
114 lines
3.0 KiB
Go
114 lines
3.0 KiB
Go
// PicoClaw - Ultra-lightweight personal AI agent
|
|
// License: MIT
|
|
//
|
|
// Copyright (c) 2026 PicoClaw contributors
|
|
|
|
package providers
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/config"
|
|
)
|
|
|
|
// ModelRegistry manages model configurations with thread-safe round-robin load balancing.
|
|
// It allows multiple configurations for the same model_name to distribute load across endpoints.
|
|
type ModelRegistry struct {
|
|
configs map[string][]config.ModelConfig // model_name -> []ModelConfig
|
|
counters map[string]*atomic.Uint64 // model_name -> round-robin counter
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// NewModelRegistry creates a new ModelRegistry from a slice of ModelConfig.
|
|
func NewModelRegistry(modelList []config.ModelConfig) *ModelRegistry {
|
|
r := &ModelRegistry{
|
|
configs: make(map[string][]config.ModelConfig),
|
|
counters: make(map[string]*atomic.Uint64),
|
|
}
|
|
|
|
for _, cfg := range modelList {
|
|
r.configs[cfg.ModelName] = append(r.configs[cfg.ModelName], cfg)
|
|
}
|
|
|
|
// Initialize counters for models with multiple configs
|
|
for name, cfgs := range r.configs {
|
|
if len(cfgs) > 1 {
|
|
r.counters[name] = &atomic.Uint64{}
|
|
}
|
|
}
|
|
|
|
return r
|
|
}
|
|
|
|
// GetModelConfig returns a ModelConfig for the given model name.
|
|
// If multiple configs exist for the same model_name, it uses round-robin selection.
|
|
// Returns an error if the model is not found.
|
|
func (r *ModelRegistry) GetModelConfig(modelName string) (*config.ModelConfig, error) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
configs, ok := r.configs[modelName]
|
|
if !ok || len(configs) == 0 {
|
|
return nil, fmt.Errorf("model %q not found", modelName)
|
|
}
|
|
|
|
// Single config - return directly
|
|
if len(configs) == 1 {
|
|
return &configs[0], nil
|
|
}
|
|
|
|
// Multiple configs - use round-robin for load balancing
|
|
counter, ok := r.counters[modelName]
|
|
if !ok {
|
|
// Should not happen, but handle gracefully
|
|
return &configs[0], nil
|
|
}
|
|
|
|
idx := counter.Add(1) % uint64(len(configs))
|
|
return &configs[idx], nil
|
|
}
|
|
|
|
// AddConfig adds a new ModelConfig to the registry.
|
|
func (r *ModelRegistry) AddConfig(cfg config.ModelConfig) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
r.configs[cfg.ModelName] = append(r.configs[cfg.ModelName], cfg)
|
|
|
|
// Initialize counter if we now have multiple configs
|
|
if len(r.configs[cfg.ModelName]) > 1 && r.counters[cfg.ModelName] == nil {
|
|
r.counters[cfg.ModelName] = &atomic.Uint64{}
|
|
}
|
|
}
|
|
|
|
// RemoveConfig removes all configs with the given model_name.
|
|
func (r *ModelRegistry) RemoveConfig(modelName string) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
delete(r.configs, modelName)
|
|
delete(r.counters, modelName)
|
|
}
|
|
|
|
// ListModels returns all unique model names in the registry.
|
|
func (r *ModelRegistry) ListModels() []string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
names := make([]string, 0, len(r.configs))
|
|
for name := range r.configs {
|
|
names = append(names, name)
|
|
}
|
|
return names
|
|
}
|
|
|
|
// ConfigCount returns the number of configurations for a given model name.
|
|
func (r *ModelRegistry) ConfigCount(modelName string) int {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
return len(r.configs[modelName])
|
|
}
|