248 lines
7.2 KiB
Go
248 lines
7.2 KiB
Go
package ai
|
||
|
||
import (
|
||
"context"
|
||
"crypto/hmac"
|
||
"crypto/rand"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"fmt"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
ProviderXAI = "xai"
|
||
ProviderOpenCodeGo = "opencode-go"
|
||
|
||
// ModelsCacheTTL is the fresh model-list window.
|
||
ModelsCacheTTL = 5 * time.Minute
|
||
// ModelsCacheStaleTTL is the maximum stale fallback window.
|
||
ModelsCacheStaleTTL = 15 * time.Minute
|
||
|
||
modelsRefreshLockTTL = 30 * time.Second
|
||
modelsRefreshWait = 25 * time.Millisecond
|
||
modelsRefreshChecks = 20
|
||
)
|
||
|
||
// ModelsCacheKey identifies one provider credential without containing the credential.
|
||
type ModelsCacheKey struct {
|
||
Provider string
|
||
CredentialFingerprint string
|
||
}
|
||
|
||
// ModelsCacheValue is the JSON payload persisted by a ModelsCache.
|
||
type ModelsCacheValue struct {
|
||
Models []string `json:"models"`
|
||
FetchedAt int64 `json:"fetchedAt"`
|
||
}
|
||
|
||
// ModelsCache stores model lists and coordinates refreshes across instances.
|
||
type ModelsCache interface {
|
||
Get(ctx context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error)
|
||
Set(ctx context.Context, key ModelsCacheKey, value ModelsCacheValue, ttl time.Duration) error
|
||
AcquireRefresh(ctx context.Context, key ModelsCacheKey, token string, ttl time.Duration) (bool, error)
|
||
ReleaseRefresh(ctx context.Context, key ModelsCacheKey, token string) error
|
||
}
|
||
|
||
// Registry holds the immutable provider client registry and a shared models cache.
|
||
type Registry struct {
|
||
clients map[string]Client
|
||
cache ModelsCache
|
||
fingerprintSecret []byte
|
||
}
|
||
|
||
func NewRegistry(cache ModelsCache, fingerprintSecret string) *Registry {
|
||
return newRegistry(cache, fingerprintSecret, map[string]Client{
|
||
ProviderXAI: NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
|
||
ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
|
||
})
|
||
}
|
||
|
||
// NewRegistryWithClients injects provider clients for focused tests.
|
||
func NewRegistryWithClients(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
|
||
return newRegistry(cache, fingerprintSecret, clients)
|
||
}
|
||
|
||
func newRegistry(cache ModelsCache, fingerprintSecret string, clients map[string]Client) *Registry {
|
||
immutableClients := make(map[string]Client, len(clients))
|
||
for provider, client := range clients {
|
||
immutableClients[NormalizeProvider(provider)] = client
|
||
}
|
||
return &Registry{clients: immutableClients, cache: cache, fingerprintSecret: []byte(fingerprintSecret)}
|
||
}
|
||
|
||
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 per provider+credential.
|
||
// 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)
|
||
if _, err := r.Client(id); err != nil {
|
||
return nil, false, err
|
||
}
|
||
key := r.cacheKey(id, apiKey)
|
||
now := time.Now().UTC()
|
||
stale, hasStale := r.cached(ctx, key, now)
|
||
if hasStale && cacheAge(now, stale.FetchedAt) <= ModelsCacheTTL {
|
||
return cloneModels(stale.Models), true, nil
|
||
}
|
||
|
||
token, tokenErr := refreshToken()
|
||
locked := false
|
||
if tokenErr != nil {
|
||
err = tokenErr
|
||
} else if r.cache != nil {
|
||
locked, err = r.cache.AcquireRefresh(ctx, key, token, modelsRefreshLockTTL)
|
||
if err != nil {
|
||
locked = false
|
||
}
|
||
}
|
||
if locked {
|
||
defer func() {
|
||
releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second)
|
||
defer cancel()
|
||
_ = r.cache.ReleaseRefresh(releaseCtx, key, token)
|
||
}()
|
||
} else if r.cache != nil && err == nil {
|
||
for range modelsRefreshChecks {
|
||
select {
|
||
case <-ctx.Done():
|
||
return nil, false, ctx.Err()
|
||
case <-time.After(modelsRefreshWait):
|
||
}
|
||
if current, ok := r.cached(ctx, key, time.Now().UTC()); ok && cacheAge(time.Now().UTC(), current.FetchedAt) <= ModelsCacheTTL {
|
||
return cloneModels(current.Models), true, nil
|
||
}
|
||
}
|
||
}
|
||
|
||
c, err := r.Client(id)
|
||
list, err := c.ListModels(ctx, apiKey)
|
||
if err != nil {
|
||
if hasStale {
|
||
return cloneModels(stale.Models), true, nil
|
||
}
|
||
return nil, false, err
|
||
}
|
||
|
||
if r.cache != nil {
|
||
_ = r.cache.Set(ctx, key, ModelsCacheValue{
|
||
Models: cloneModels(list),
|
||
FetchedAt: time.Now().UTC().UnixNano(),
|
||
}, ModelsCacheStaleTTL)
|
||
}
|
||
return list, false, nil
|
||
}
|
||
|
||
func (r *Registry) cached(ctx context.Context, key ModelsCacheKey, now time.Time) (ModelsCacheValue, bool) {
|
||
if r.cache == nil {
|
||
return ModelsCacheValue{}, false
|
||
}
|
||
value, ok, err := r.cache.Get(ctx, key)
|
||
if err != nil || !ok || value.FetchedAt <= 0 {
|
||
return ModelsCacheValue{}, false
|
||
}
|
||
if cacheAge(now, value.FetchedAt) > ModelsCacheStaleTTL {
|
||
return ModelsCacheValue{}, false
|
||
}
|
||
return value, true
|
||
}
|
||
|
||
func (r *Registry) cacheKey(provider, apiKey string) ModelsCacheKey {
|
||
mac := hmac.New(sha256.New, r.fingerprintSecret)
|
||
_, _ = mac.Write([]byte(apiKey))
|
||
return ModelsCacheKey{Provider: provider, CredentialFingerprint: hex.EncodeToString(mac.Sum(nil))}
|
||
}
|
||
|
||
func refreshToken() (string, error) {
|
||
var token [16]byte
|
||
if _, err := rand.Read(token[:]); err != nil {
|
||
return "", err
|
||
}
|
||
return hex.EncodeToString(token[:]), nil
|
||
}
|
||
|
||
func cacheAge(now time.Time, fetchedAt int64) time.Duration {
|
||
age := now.Sub(time.Unix(0, fetchedAt).UTC())
|
||
if age < 0 {
|
||
return 0
|
||
}
|
||
return age
|
||
}
|
||
|
||
func cloneModels(models []string) []string {
|
||
return append([]string(nil), models...)
|
||
}
|
||
|
||
// 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
|
||
}
|