AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,330 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/provider"
|
||||
provideropenai "aigateway.local/core/internal/provider/openai"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
const changeChannel = "gateway:providers:changed"
|
||||
|
||||
type snapshot struct {
|
||||
adapters map[string]gateway.ResolvedAdapter
|
||||
defaultCode string
|
||||
routes map[string][]compiledRoute
|
||||
knownModels map[string]bool
|
||||
}
|
||||
|
||||
type routeConditions struct {
|
||||
Endpoints []string `json:"endpoints"`
|
||||
APIKeyIDs []string `json:"api_key_ids"`
|
||||
TenantIDs []string `json:"tenant_ids"`
|
||||
}
|
||||
|
||||
type compiledRoute struct {
|
||||
providerCode string
|
||||
targetModel string
|
||||
weight int
|
||||
priority int
|
||||
conditions routeConditions
|
||||
}
|
||||
|
||||
type Resolver struct {
|
||||
repository *provider.Repository
|
||||
cipher *provider.CredentialCipher
|
||||
fallback gateway.ResolvedAdapter
|
||||
interval time.Duration
|
||||
logger *slog.Logger
|
||||
redis *redis.Client
|
||||
snapshot atomic.Pointer[snapshot]
|
||||
}
|
||||
|
||||
func (r *Resolver) SetNotificationClient(client *redis.Client) {
|
||||
r.redis = client
|
||||
}
|
||||
|
||||
func (r *Resolver) Notify(ctx context.Context) error {
|
||||
if r.redis == nil {
|
||||
return nil
|
||||
}
|
||||
return r.redis.Publish(ctx, changeChannel, time.Now().UTC().Format(time.RFC3339Nano)).Err()
|
||||
}
|
||||
|
||||
func NewResolver(repository *provider.Repository, cipher *provider.CredentialCipher, fallback provider.Adapter, interval time.Duration, logger *slog.Logger) *Resolver {
|
||||
capabilities := make(map[provider.Capability]bool)
|
||||
if fallback != nil {
|
||||
for _, capability := range fallback.Capabilities() {
|
||||
capabilities[capability] = true
|
||||
}
|
||||
}
|
||||
resolver := &Resolver{
|
||||
repository: repository, cipher: cipher, interval: interval, logger: logger,
|
||||
fallback: gateway.ResolvedAdapter{Code: "environment", Adapter: fallback, Capabilities: capabilities},
|
||||
}
|
||||
initial := &snapshot{adapters: make(map[string]gateway.ResolvedAdapter), routes: make(map[string][]compiledRoute), knownModels: make(map[string]bool)}
|
||||
if fallback != nil {
|
||||
initial.adapters["environment"] = resolver.fallback
|
||||
initial.defaultCode = "environment"
|
||||
}
|
||||
resolver.snapshot.Store(initial)
|
||||
return resolver
|
||||
}
|
||||
|
||||
func (r *Resolver) ResolveModelRoute(query gateway.ModelRouteQuery) (gateway.ModelRouteResult, error) {
|
||||
current := r.snapshot.Load()
|
||||
if current == nil {
|
||||
return gateway.ModelRouteResult{}, gateway.ErrProviderUnavailable
|
||||
}
|
||||
candidates := current.routes[query.Model]
|
||||
known := current.knownModels[query.Model]
|
||||
matched := make([]compiledRoute, 0, len(candidates))
|
||||
highestPriority := -100001
|
||||
for _, route := range candidates {
|
||||
if query.ProviderCode != "" && route.providerCode != query.ProviderCode || !route.conditions.matches(query) {
|
||||
continue
|
||||
}
|
||||
if route.priority > highestPriority {
|
||||
highestPriority = route.priority
|
||||
matched = matched[:0]
|
||||
}
|
||||
if route.priority == highestPriority {
|
||||
matched = append(matched, route)
|
||||
}
|
||||
}
|
||||
if len(matched) == 0 {
|
||||
return gateway.ModelRouteResult{Known: known}, nil
|
||||
}
|
||||
total := 0
|
||||
for _, route := range matched {
|
||||
total += route.weight
|
||||
}
|
||||
hasher := fnv.New64a()
|
||||
_, _ = hasher.Write([]byte(query.Seed + "\x00" + query.Model))
|
||||
selected := int(hasher.Sum64() % uint64(total))
|
||||
chosen := matched[len(matched)-1]
|
||||
for _, route := range matched {
|
||||
if selected < route.weight {
|
||||
chosen = route
|
||||
break
|
||||
}
|
||||
selected -= route.weight
|
||||
}
|
||||
resolved, ok := current.adapters[chosen.providerCode]
|
||||
if !ok {
|
||||
return gateway.ModelRouteResult{}, gateway.ErrProviderUnavailable
|
||||
}
|
||||
return gateway.ModelRouteResult{ResolvedAdapter: resolved, TargetModel: chosen.targetModel, Matched: true, Known: true}, nil
|
||||
}
|
||||
|
||||
func (r *Resolver) ModelRoutingEnabled() bool {
|
||||
current := r.snapshot.Load()
|
||||
return current != nil && len(current.knownModels) > 0
|
||||
}
|
||||
|
||||
func (c routeConditions) matches(query gateway.ModelRouteQuery) bool {
|
||||
return containsOrEmpty(c.Endpoints, query.Endpoint) && containsOrEmpty(c.APIKeyIDs, query.APIKeyID) && containsOrEmpty(c.TenantIDs, query.TenantID)
|
||||
}
|
||||
|
||||
func containsOrEmpty(values []string, target string) bool {
|
||||
if len(values) == 0 {
|
||||
return true
|
||||
}
|
||||
for _, value := range values {
|
||||
if value == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (r *Resolver) Resolve(code string) (gateway.ResolvedAdapter, error) {
|
||||
current := r.snapshot.Load()
|
||||
if current == nil {
|
||||
return gateway.ResolvedAdapter{}, gateway.ErrProviderUnavailable
|
||||
}
|
||||
code = strings.ToLower(strings.TrimSpace(code))
|
||||
if code == "" {
|
||||
if current.defaultCode == "" {
|
||||
return gateway.ResolvedAdapter{}, gateway.ErrProviderUnavailable
|
||||
}
|
||||
code = current.defaultCode
|
||||
}
|
||||
resolved, ok := current.adapters[code]
|
||||
if !ok {
|
||||
return gateway.ResolvedAdapter{}, gateway.ErrProviderNotFound
|
||||
}
|
||||
return resolved, nil
|
||||
}
|
||||
|
||||
func (r *Resolver) Reload(ctx context.Context) error {
|
||||
records, err := r.repository.List(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
routeRecords, err := r.repository.ListModelRoutes(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
next := &snapshot{adapters: make(map[string]gateway.ResolvedAdapter), routes: make(map[string][]compiledRoute), knownModels: make(map[string]bool)}
|
||||
previous := r.snapshot.Load()
|
||||
enabledCount := 0
|
||||
explicitDefault := false
|
||||
for _, record := range records {
|
||||
if !record.Enabled {
|
||||
continue
|
||||
}
|
||||
enabledCount++
|
||||
resolved, defaultProvider, err := r.build(record)
|
||||
if err != nil {
|
||||
r.logger.Error("provider snapshot entry rejected", "provider", record.Code, "revision", record.Revision, "error", err)
|
||||
if previous != nil {
|
||||
if lastValid, ok := previous.adapters[record.Code]; ok {
|
||||
next.adapters[record.Code] = lastValid
|
||||
if next.defaultCode == "" || previous.defaultCode == record.Code {
|
||||
next.defaultCode = record.Code
|
||||
}
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
next.adapters[record.Code] = resolved
|
||||
if defaultProvider {
|
||||
if explicitDefault {
|
||||
return errors.New("multiple enabled providers declare config.default=true; retaining previous snapshot")
|
||||
}
|
||||
explicitDefault = true
|
||||
next.defaultCode = record.Code
|
||||
} else if next.defaultCode == "" {
|
||||
next.defaultCode = record.Code
|
||||
}
|
||||
}
|
||||
if enabledCount > 0 && len(next.adapters) == 0 {
|
||||
return errors.New("all enabled provider records are invalid; retaining previous snapshot")
|
||||
}
|
||||
if len(next.adapters) == 0 {
|
||||
if r.fallback.Adapter != nil {
|
||||
next.adapters["environment"] = r.fallback
|
||||
next.defaultCode = "environment"
|
||||
}
|
||||
}
|
||||
for _, route := range routeRecords {
|
||||
if !route.Enabled {
|
||||
continue
|
||||
}
|
||||
if _, ok := next.adapters[route.ProviderCode]; !ok {
|
||||
continue
|
||||
}
|
||||
// Only routes that can actually be served make the model "known".
|
||||
// Marking it before the Enabled check above meant a disabled route
|
||||
// still advertised the model as known, and the proxy hard-blocks any
|
||||
// known-but-unmatched model with a 400 — so disabling the only route
|
||||
// for a model made that model unusable on every other provider.
|
||||
next.knownModels[route.SourceModel] = true
|
||||
var conditions routeConditions
|
||||
if err := json.Unmarshal(route.Conditions, &conditions); err != nil {
|
||||
r.logger.Error("model route snapshot entry rejected", "route", route.ID, "error", err)
|
||||
continue
|
||||
}
|
||||
next.routes[route.SourceModel] = append(next.routes[route.SourceModel], compiledRoute{
|
||||
providerCode: route.ProviderCode, targetModel: route.TargetModel, weight: route.Weight,
|
||||
priority: route.Priority, conditions: conditions,
|
||||
})
|
||||
}
|
||||
r.snapshot.Store(next)
|
||||
r.logger.Info("provider snapshot refreshed", "providers", len(next.adapters), "model_routes", len(routeRecords), "default", next.defaultCode)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Resolver) Run(ctx context.Context) {
|
||||
if err := r.Reload(ctx); err != nil {
|
||||
r.logger.Warn("initial provider snapshot refresh failed; using last valid snapshot", "error", err)
|
||||
}
|
||||
interval := r.interval
|
||||
if interval <= 0 {
|
||||
interval = 5 * time.Second
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
var changes <-chan *redis.Message
|
||||
var subscription *redis.PubSub
|
||||
if r.redis != nil {
|
||||
subscription = r.redis.Subscribe(ctx, changeChannel)
|
||||
defer subscription.Close()
|
||||
changes = subscription.Channel(redis.WithChannelSize(32))
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
refreshCtx, cancel := context.WithTimeout(ctx, interval)
|
||||
err := r.Reload(refreshCtx)
|
||||
cancel()
|
||||
if err != nil && ctx.Err() == nil {
|
||||
r.logger.Warn("provider snapshot refresh failed; retaining last valid snapshot", "error", err)
|
||||
}
|
||||
case _, ok := <-changes:
|
||||
if !ok {
|
||||
changes = nil
|
||||
continue
|
||||
}
|
||||
refreshCtx, cancel := context.WithTimeout(ctx, interval)
|
||||
err := r.Reload(refreshCtx)
|
||||
cancel()
|
||||
if err != nil && ctx.Err() == nil {
|
||||
r.logger.Warn("provider notification refresh failed; retaining last valid snapshot", "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Resolver) build(record provider.Record) (gateway.ResolvedAdapter, bool, error) {
|
||||
plaintext, err := r.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
|
||||
if err != nil {
|
||||
return gateway.ResolvedAdapter{}, false, fmt.Errorf("decrypt credentials: %w", err)
|
||||
}
|
||||
var credentials provider.Credentials
|
||||
if err := json.Unmarshal(plaintext, &credentials); err != nil {
|
||||
return gateway.ResolvedAdapter{}, false, fmt.Errorf("decode credentials: %w", err)
|
||||
}
|
||||
var adapter provider.Adapter
|
||||
switch record.Adapter {
|
||||
case "openai-compatible":
|
||||
adapter, err = provideropenai.New(record.BaseURL, credentials.APIKey)
|
||||
default:
|
||||
err = fmt.Errorf("unsupported adapter %q", record.Adapter)
|
||||
}
|
||||
if err != nil {
|
||||
return gateway.ResolvedAdapter{}, false, err
|
||||
}
|
||||
capabilities := make(map[provider.Capability]bool, len(record.Capabilities))
|
||||
for _, value := range record.Capabilities {
|
||||
capability := provider.Capability(value)
|
||||
switch capability {
|
||||
case provider.CapabilityChat, provider.CapabilityResponses, provider.CapabilityEmbeddings, provider.CapabilityMessages, provider.CapabilityModels:
|
||||
capabilities[capability] = true
|
||||
default:
|
||||
return gateway.ResolvedAdapter{}, false, fmt.Errorf("unsupported capability %q", value)
|
||||
}
|
||||
}
|
||||
var configuration struct {
|
||||
Default bool `json:"default"`
|
||||
}
|
||||
if len(record.Config) > 0 {
|
||||
if err := json.Unmarshal(record.Config, &configuration); err != nil {
|
||||
return gateway.ResolvedAdapter{}, false, fmt.Errorf("decode config: %w", err)
|
||||
}
|
||||
}
|
||||
return gateway.ResolvedAdapter{Code: record.Code, Revision: record.Revision, Adapter: adapter, Capabilities: capabilities}, configuration.Default, nil
|
||||
}
|
||||
Reference in New Issue
Block a user