152 lines
3.7 KiB
Go
152 lines
3.7 KiB
Go
package ai
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"fmt"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
ProviderXAI = "xai"
|
||
ProviderOpenCodeGo = "opencode-go"
|
||
|
||
// ModelsCacheTTL — list models cache window
|
||
ModelsCacheTTL = 5 * time.Minute
|
||
)
|
||
|
||
// Registry holds provider clients + models cache.
|
||
type Registry struct {
|
||
clients map[string]Client
|
||
mu sync.Mutex
|
||
cache map[string]modelsCacheEntry
|
||
}
|
||
|
||
type modelsCacheEntry struct {
|
||
Models []string
|
||
ExpiresAt time.Time
|
||
}
|
||
|
||
func NewRegistry() *Registry {
|
||
return &Registry{
|
||
clients: map[string]Client{
|
||
ProviderXAI: NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
|
||
ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
|
||
},
|
||
cache: map[string]modelsCacheEntry{},
|
||
}
|
||
}
|
||
|
||
func NormalizeProvider(p string) string {
|
||
switch strings.ToLower(strings.TrimSpace(p)) {
|
||
case "opencode-go", "opencode", "opencode_go":
|
||
return ProviderOpenCodeGo
|
||
case "xai", "grok", "":
|
||
return ProviderXAI
|
||
default:
|
||
return strings.ToLower(strings.TrimSpace(p))
|
||
}
|
||
}
|
||
|
||
func (r *Registry) Client(provider string) (Client, error) {
|
||
id := NormalizeProvider(provider)
|
||
c, ok := r.clients[id]
|
||
if !ok {
|
||
return nil, fmt.Errorf("unsupported provider: %s (use xai or opencode-go)", provider)
|
||
}
|
||
return c, nil
|
||
}
|
||
|
||
// ListModelsCached fetches models via provider API; caches 5 minutes per provider+key.
|
||
// fromCache indicates whether result was served from cache.
|
||
func (r *Registry) ListModelsCached(ctx context.Context, provider, apiKey string) (models []string, fromCache bool, err error) {
|
||
id := NormalizeProvider(provider)
|
||
key := cacheKey(id, apiKey)
|
||
|
||
r.mu.Lock()
|
||
if ent, ok := r.cache[key]; ok && time.Now().Before(ent.ExpiresAt) {
|
||
out := append([]string(nil), ent.Models...)
|
||
r.mu.Unlock()
|
||
return out, true, nil
|
||
}
|
||
r.mu.Unlock()
|
||
|
||
c, err := r.Client(id)
|
||
if err != nil {
|
||
return nil, false, err
|
||
}
|
||
list, err := c.ListModels(ctx, apiKey)
|
||
if err != nil {
|
||
return nil, false, err
|
||
}
|
||
|
||
r.mu.Lock()
|
||
r.cache[key] = modelsCacheEntry{
|
||
Models: append([]string(nil), list...),
|
||
ExpiresAt: time.Now().Add(ModelsCacheTTL),
|
||
}
|
||
r.mu.Unlock()
|
||
return list, false, nil
|
||
}
|
||
|
||
func cacheKey(provider, apiKey string) string {
|
||
sum := sha256.Sum256([]byte(provider + "\n" + apiKey))
|
||
return hex.EncodeToString(sum[:16])
|
||
}
|
||
|
||
// FallbackModels when no key / provider error (never empty for UI).
|
||
// OpenCode Go 模型 id 見 https://opencode.ai/docs/go/（不是 OpenAI gpt-*）
|
||
func FallbackModels(provider string) []string {
|
||
switch NormalizeProvider(provider) {
|
||
case ProviderOpenCodeGo:
|
||
return []string{
|
||
"kimi-k2.6",
|
||
"glm-5.1",
|
||
"glm-5.2",
|
||
"deepseek-v4-flash",
|
||
"deepseek-v4-pro",
|
||
"qwen3.7-plus",
|
||
"minimax-m2.7",
|
||
"mimo-v2.5",
|
||
}
|
||
default:
|
||
return []string{"grok-3", "grok-3-mini", "grok-2"}
|
||
}
|
||
}
|
||
|
||
// DefaultModel for provider when member settings model empty / invalid.
|
||
func DefaultModel(provider string) string {
|
||
list := FallbackModels(provider)
|
||
if len(list) == 0 {
|
||
return "grok-3"
|
||
}
|
||
return list[0]
|
||
}
|
||
|
||
// SanitizeModel remaps invalid model ids for a provider (e.g. gpt-4o-mini on opencode-go).
|
||
func SanitizeModel(provider, model string) string {
|
||
provider = NormalizeProvider(provider)
|
||
model = strings.TrimSpace(model)
|
||
if model == "" {
|
||
return DefaultModel(provider)
|
||
}
|
||
if provider != ProviderOpenCodeGo {
|
||
// xAI：允許 grok-*；若誤存 opencode 模型名，仍交由 API 驗證
|
||
return model
|
||
}
|
||
// OpenCode Go：常見錯誤 — 從 xAI / OpenAI 帶過來的模型名
|
||
lower := strings.ToLower(model)
|
||
if strings.HasPrefix(lower, "gpt-") ||
|
||
strings.HasPrefix(lower, "o1") ||
|
||
strings.HasPrefix(lower, "o3") ||
|
||
strings.HasPrefix(lower, "grok") ||
|
||
strings.HasPrefix(lower, "claude") ||
|
||
strings.HasPrefix(lower, "gemini") {
|
||
return DefaultModel(provider)
|
||
}
|
||
return model
|
||
}
|