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,498 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
var providerCodePattern = regexp.MustCompile(`^[a-z][a-z0-9_-]{1,63}$`)
|
||||
|
||||
type AdminHTTPHandler struct {
|
||||
repository *Repository
|
||||
cipher *CredentialCipher
|
||||
identity *identity.Service
|
||||
allowPrivate bool
|
||||
changeHook func(context.Context) error
|
||||
operations AdminOperations
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) SetChangeHook(hook func(context.Context) error) {
|
||||
h.changeHook = hook
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) SetOperations(operations AdminOperations) {
|
||||
h.operations = operations
|
||||
}
|
||||
|
||||
type providerInput struct {
|
||||
Code string `json:"code"`
|
||||
Adapter string `json:"adapter"`
|
||||
BaseURL string `json:"base_url"`
|
||||
APIKey *string `json:"api_key"`
|
||||
Capabilities []string `json:"capabilities"`
|
||||
Config json.RawMessage `json:"config"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
type modelRouteInput struct {
|
||||
Name string `json:"name"`
|
||||
SourceModel string `json:"source_model"`
|
||||
TargetModel string `json:"target_model"`
|
||||
ProviderID string `json:"provider_id"`
|
||||
Weight int `json:"weight"`
|
||||
Priority int `json:"priority"`
|
||||
Conditions json.RawMessage `json:"conditions"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
type modelRouteConditions struct {
|
||||
Endpoints []string `json:"endpoints"`
|
||||
APIKeyIDs []string `json:"api_key_ids"`
|
||||
TenantIDs []string `json:"tenant_ids"`
|
||||
}
|
||||
|
||||
func NewAdminHTTPHandler(repository *Repository, cipher *CredentialCipher, identityService *identity.Service, allowPrivate bool) *AdminHTTPHandler {
|
||||
handler := &AdminHTTPHandler{
|
||||
repository: repository, cipher: cipher, identity: identityService,
|
||||
allowPrivate: allowPrivate, mux: http.NewServeMux(),
|
||||
}
|
||||
handler.mux.HandleFunc("GET /api/v1/admin/providers", handler.list)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers", handler.create)
|
||||
handler.mux.HandleFunc("PUT /api/v1/admin/providers/{provider_id}", handler.update)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers/{provider_id}/test", handler.testConnection)
|
||||
handler.mux.HandleFunc("GET /api/v1/admin/providers/{provider_id}/models", handler.listModels)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers/{provider_id}/models/sync", handler.syncModels)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/providers/credentials/rotate", handler.rotateCredentials)
|
||||
handler.mux.HandleFunc("GET /api/v1/admin/model-routes", handler.listModelRoutes)
|
||||
handler.mux.HandleFunc("POST /api/v1/admin/model-routes", handler.createModelRoute)
|
||||
handler.mux.HandleFunc("PUT /api/v1/admin/model-routes/{model_route_id}", handler.updateModelRoute)
|
||||
handler.mux.HandleFunc("DELETE /api/v1/admin/model-routes/{model_route_id}", handler.deleteModelRoute)
|
||||
return handler
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
|
||||
h.mux.ServeHTTP(writer, request)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) list(writer http.ResponseWriter, request *http.Request) {
|
||||
if _, ok := h.requirePermission(writer, request, identity.PermissionProviderRead); !ok {
|
||||
return
|
||||
}
|
||||
records, err := h.repository.List(request.Context())
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
items := make([]map[string]any, 0, len(records))
|
||||
for _, record := range records {
|
||||
item, err := h.view(record)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
apiresponse.OK(writer, items)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) create(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
input, record, err := h.decodeRecord(writer, request)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiKey := ""
|
||||
if input.APIKey != nil {
|
||||
apiKey = strings.TrimSpace(*input.APIKey)
|
||||
}
|
||||
if err := h.setCredentials(&record, apiKey); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
created, err := h.repository.Create(request.Context(), record, actor.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if h.changeHook != nil {
|
||||
_ = h.changeHook(request.Context())
|
||||
}
|
||||
view, err := h.view(created)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, view)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) update(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
input, record, err := h.decodeRecord(writer, request)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
record.ID = request.PathValue("provider_id")
|
||||
if record.ID == "" {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "provider_id 不能为空")
|
||||
return
|
||||
}
|
||||
replaceCredentials := input.APIKey != nil
|
||||
if replaceCredentials {
|
||||
if err := h.setCredentials(&record, strings.TrimSpace(*input.APIKey)); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := h.repository.Update(request.Context(), record, actor.ID, replaceCredentials); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
updated, err := h.repository.Get(request.Context(), record.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if h.changeHook != nil {
|
||||
_ = h.changeHook(request.Context())
|
||||
}
|
||||
view, err := h.view(updated)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, view)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) decodeRecord(writer http.ResponseWriter, request *http.Request) (providerInput, Record, error) {
|
||||
var input providerInput
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if err := decoder.Decode(&input); err != nil {
|
||||
return input, Record{}, errors.New("请求格式无效")
|
||||
}
|
||||
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
|
||||
if !providerCodePattern.MatchString(input.Code) {
|
||||
return input, Record{}, errors.New("code 必须以小写字母开头,且只能包含小写字母、数字、下划线和连字符")
|
||||
}
|
||||
input.Adapter = strings.TrimSpace(input.Adapter)
|
||||
if input.Adapter == "" {
|
||||
input.Adapter = "openai-compatible"
|
||||
}
|
||||
if input.Adapter != "openai-compatible" {
|
||||
return input, Record{}, errors.New("当前只支持 openai-compatible adapter")
|
||||
}
|
||||
validationCtx, cancel := context.WithTimeout(request.Context(), 3*time.Second)
|
||||
defer cancel()
|
||||
baseURL, err := ValidateBaseURL(validationCtx, input.BaseURL, h.allowPrivate)
|
||||
if err != nil {
|
||||
return input, Record{}, err
|
||||
}
|
||||
if len(input.Capabilities) == 0 {
|
||||
input.Capabilities = []string{"chat", "responses", "embeddings", "models"}
|
||||
}
|
||||
if len(input.Config) == 0 {
|
||||
input.Config = json.RawMessage(`{}`)
|
||||
}
|
||||
var configObject map[string]any
|
||||
if err := json.Unmarshal(input.Config, &configObject); err != nil || configObject == nil {
|
||||
return input, Record{}, errors.New("config 必须是 JSON 对象")
|
||||
}
|
||||
if value, exists := configObject["default"]; exists {
|
||||
if _, ok := value.(bool); !ok {
|
||||
return input, Record{}, errors.New("config.default 必须是布尔值")
|
||||
}
|
||||
}
|
||||
enabled := true
|
||||
if input.Enabled != nil {
|
||||
enabled = *input.Enabled
|
||||
}
|
||||
return input, Record{
|
||||
Code: input.Code, Adapter: input.Adapter, BaseURL: baseURL,
|
||||
Capabilities: input.Capabilities, Config: input.Config, Enabled: enabled,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) setCredentials(record *Record, apiKey string) error {
|
||||
payload, err := json.Marshal(Credentials{APIKey: apiKey})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
encrypted, version, err := h.cipher.Encrypt(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
record.EncryptedCredentials = encrypted
|
||||
record.CredentialKEKVersion = version
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) view(record Record) (map[string]any, error) {
|
||||
plaintext, err := h.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var credentials Credentials
|
||||
if err := json.Unmarshal(plaintext, &credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"id": record.ID, "code": record.Code, "adapter": record.Adapter,
|
||||
"base_url": record.BaseURL, "capabilities": record.Capabilities,
|
||||
"config": record.Config, "enabled": record.Enabled, "revision": record.Revision,
|
||||
"credential_kek_version": record.CredentialKEKVersion,
|
||||
"key_configured": credentials.APIKey != "", "api_key_masked": maskSecret(credentials.APIKey),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) testConnection(writer http.ResponseWriter, request *http.Request) {
|
||||
if _, ok := h.requirePermission(writer, request, identity.PermissionProviderManage); !ok {
|
||||
return
|
||||
}
|
||||
if !h.requireOperations(writer) {
|
||||
return
|
||||
}
|
||||
result, err := h.operations.TestConnection(request.Context(), request.PathValue("provider_id"))
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, result)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listModels(writer http.ResponseWriter, request *http.Request) {
|
||||
if _, ok := h.requirePermission(writer, request, identity.PermissionProviderRead); !ok {
|
||||
return
|
||||
}
|
||||
if !h.requireOperations(writer) {
|
||||
return
|
||||
}
|
||||
models, err := h.operations.ListModels(request.Context(), request.PathValue("provider_id"))
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if models == nil {
|
||||
models = []Model{}
|
||||
}
|
||||
apiresponse.OK(writer, models)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) syncModels(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if !h.requireOperations(writer) {
|
||||
return
|
||||
}
|
||||
result, err := h.operations.SyncModels(request.Context(), request.PathValue("provider_id"), actor.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, result)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) rotateCredentials(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if !h.requireOperations(writer) {
|
||||
return
|
||||
}
|
||||
result, err := h.operations.RotateCredentials(request.Context(), actor.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
if h.changeHook != nil && result.Rotated > 0 {
|
||||
_ = h.changeHook(request.Context())
|
||||
}
|
||||
apiresponse.OK(writer, result)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listModelRoutes(writer http.ResponseWriter, request *http.Request) {
|
||||
if _, ok := h.requirePermission(writer, request, identity.PermissionProviderRead); !ok {
|
||||
return
|
||||
}
|
||||
routes, err := h.repository.ListModelRoutes(request.Context())
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, routes)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) createModelRoute(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
route, err := h.decodeModelRoute(writer, request)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
created, err := h.repository.CreateModelRoute(request.Context(), route, actor.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
h.propagateChange(request.Context())
|
||||
apiresponse.OK(writer, created)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) updateModelRoute(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
route, err := h.decodeModelRoute(writer, request)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
route.ID = request.PathValue("model_route_id")
|
||||
updated, err := h.repository.UpdateModelRoute(request.Context(), route, actor.ID)
|
||||
if err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
h.propagateChange(request.Context())
|
||||
apiresponse.OK(writer, updated)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) deleteModelRoute(writer http.ResponseWriter, request *http.Request) {
|
||||
actor, ok := h.requirePermission(writer, request, identity.PermissionProviderManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.repository.DeleteModelRoute(request.Context(), request.PathValue("model_route_id"), actor.ID); err != nil {
|
||||
h.writeError(writer, err)
|
||||
return
|
||||
}
|
||||
h.propagateChange(request.Context())
|
||||
apiresponse.OK(writer, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) decodeModelRoute(writer http.ResponseWriter, request *http.Request) (ModelRoute, error) {
|
||||
var input modelRouteInput
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if err := decoder.Decode(&input); err != nil {
|
||||
return ModelRoute{}, errors.New("请求格式无效")
|
||||
}
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.SourceModel = strings.TrimSpace(input.SourceModel)
|
||||
input.TargetModel = strings.TrimSpace(input.TargetModel)
|
||||
input.ProviderID = strings.TrimSpace(input.ProviderID)
|
||||
if input.Name == "" || len(input.Name) > 128 || input.SourceModel == "" || len(input.SourceModel) > 512 || input.TargetModel == "" || len(input.TargetModel) > 512 || input.ProviderID == "" {
|
||||
return ModelRoute{}, errors.New("名称、模型别名、上游模型或供应商无效")
|
||||
}
|
||||
if input.Weight == 0 {
|
||||
input.Weight = 100
|
||||
}
|
||||
if input.Weight < 1 || input.Weight > 10000 || input.Priority < -100000 || input.Priority > 100000 {
|
||||
return ModelRoute{}, errors.New("权重或优先级超出允许范围")
|
||||
}
|
||||
if len(input.Conditions) == 0 || string(input.Conditions) == "null" {
|
||||
input.Conditions = json.RawMessage(`{}`)
|
||||
}
|
||||
var conditions modelRouteConditions
|
||||
if err := json.Unmarshal(input.Conditions, &conditions); err != nil {
|
||||
return ModelRoute{}, errors.New("conditions 必须是 JSON 对象")
|
||||
}
|
||||
allowedEndpoints := map[string]bool{"/v1/chat/completions": true, "/v1/responses": true, "/v1/embeddings": true, "/v1/messages": true}
|
||||
for _, endpoint := range conditions.Endpoints {
|
||||
if !allowedEndpoints[endpoint] {
|
||||
return ModelRoute{}, errors.New("conditions.endpoints 包含不支持的网关端点")
|
||||
}
|
||||
}
|
||||
enabled := true
|
||||
if input.Enabled != nil {
|
||||
enabled = *input.Enabled
|
||||
}
|
||||
return ModelRoute{Name: input.Name, SourceModel: input.SourceModel, TargetModel: input.TargetModel, ProviderID: input.ProviderID, Weight: input.Weight, Priority: input.Priority, Conditions: input.Conditions, Enabled: enabled}, nil
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) propagateChange(ctx context.Context) {
|
||||
if h.changeHook != nil {
|
||||
_ = h.changeHook(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) requireOperations(writer http.ResponseWriter) bool {
|
||||
if h.operations == nil {
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "供应商控制面服务暂不可用")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) (identity.Account, bool) {
|
||||
account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
if errors.Is(err, identity.ErrInvalidSession) || errors.Is(err, identity.ErrNotFound) {
|
||||
apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
|
||||
} else if errors.Is(err, identity.ErrAccountDisabled) {
|
||||
apiresponse.Error(writer, http.StatusForbidden, "管理员账号已被停用")
|
||||
} else {
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "身份服务暂不可用")
|
||||
}
|
||||
return identity.Account{}, false
|
||||
}
|
||||
if !identity.HasPermission(account, permission) {
|
||||
apiresponse.Error(writer, http.StatusForbidden, "缺少模型供应商操作权限")
|
||||
return identity.Account{}, false
|
||||
}
|
||||
return account, true
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) writeError(writer http.ResponseWriter, err error) {
|
||||
switch {
|
||||
case errors.Is(err, ErrProviderNotFound):
|
||||
apiresponse.Error(writer, http.StatusNotFound, "供应商不存在")
|
||||
case errors.Is(err, ErrProviderExists):
|
||||
apiresponse.Error(writer, http.StatusConflict, "供应商 code 已存在")
|
||||
case errors.Is(err, ErrMultipleDefaults):
|
||||
apiresponse.Error(writer, http.StatusConflict, "只能启用一个默认供应商")
|
||||
case errors.Is(err, ErrModelRouteNotFound):
|
||||
apiresponse.Error(writer, http.StatusNotFound, "模型路由不存在")
|
||||
case errors.Is(err, ErrModelRouteExists):
|
||||
apiresponse.Error(writer, http.StatusConflict, "相同模型、供应商和上游模型的路由已存在")
|
||||
case errors.Is(err, ErrProviderStore), errors.Is(err, ErrCredentialKeyUnavailable):
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "供应商配置服务暂不可用")
|
||||
case errors.Is(err, ErrProviderUpstream):
|
||||
apiresponse.Error(writer, http.StatusBadGateway, "无法从上游供应商获取模型信息")
|
||||
default:
|
||||
apiresponse.Error(writer, http.StatusBadRequest, err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func maskSecret(secret string) string {
|
||||
if secret == "" {
|
||||
return ""
|
||||
}
|
||||
if len(secret) <= 10 {
|
||||
return "••••••••"
|
||||
}
|
||||
return secret[:4] + "••••••" + secret[len(secret)-4:]
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
package controlplane
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/provider"
|
||||
provideropenai "aigateway.local/core/internal/provider/openai"
|
||||
)
|
||||
|
||||
const maxModelsResponseBytes int64 = 4 << 20
|
||||
|
||||
type Service struct {
|
||||
repository *provider.Repository
|
||||
cipher *provider.CredentialCipher
|
||||
allowPrivate bool
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func NewService(repository *provider.Repository, cipher *provider.CredentialCipher, allowPrivate bool) *Service {
|
||||
return &Service{
|
||||
repository: repository,
|
||||
cipher: cipher,
|
||||
allowPrivate: allowPrivate,
|
||||
client: &http.Client{
|
||||
Timeout: 10 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
DialContext: safeDialContext(allowPrivate),
|
||||
ForceAttemptHTTP2: true,
|
||||
MaxIdleConns: 32,
|
||||
MaxIdleConnsPerHost: 4,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
ResponseHeaderTimeout: 8 * time.Second,
|
||||
},
|
||||
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
|
||||
return errors.New("upstream redirects are not allowed")
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) TestConnection(ctx context.Context, providerID string) (provider.ConnectionTestResult, error) {
|
||||
record, adapter, err := s.load(ctx, providerID)
|
||||
if err != nil {
|
||||
return provider.ConnectionTestResult{}, err
|
||||
}
|
||||
started := time.Now()
|
||||
response, err := s.doModelsRequest(ctx, adapter)
|
||||
latency := time.Since(started).Milliseconds()
|
||||
if err != nil {
|
||||
return provider.ConnectionTestResult{}, fmt.Errorf("%w: request provider %s: %v", provider.ErrProviderUpstream, record.Code, err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 32<<10))
|
||||
result := provider.ConnectionTestResult{
|
||||
Connected: response.StatusCode >= 200 && response.StatusCode < 300,
|
||||
StatusCode: response.StatusCode,
|
||||
LatencyMS: latency,
|
||||
}
|
||||
if result.Connected {
|
||||
result.Message = "连接成功"
|
||||
} else {
|
||||
result.Message = "上游返回 HTTP " + strconv.Itoa(response.StatusCode)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) SyncModels(ctx context.Context, providerID, actorID string) (provider.ModelSyncResult, error) {
|
||||
record, adapter, err := s.load(ctx, providerID)
|
||||
if err != nil {
|
||||
return provider.ModelSyncResult{}, err
|
||||
}
|
||||
if !hasCapability(record.Capabilities, string(provider.CapabilityModels)) {
|
||||
return provider.ModelSyncResult{}, errors.New("供应商未启用 models 能力")
|
||||
}
|
||||
response, err := s.doModelsRequest(ctx, adapter)
|
||||
if err != nil {
|
||||
return provider.ModelSyncResult{}, fmt.Errorf("%w: request provider %s: %v", provider.ErrProviderUpstream, record.Code, err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 32<<10))
|
||||
return provider.ModelSyncResult{}, fmt.Errorf("%w: provider %s returned HTTP %d", provider.ErrProviderUpstream, record.Code, response.StatusCode)
|
||||
}
|
||||
payload, err := io.ReadAll(io.LimitReader(response.Body, maxModelsResponseBytes+1))
|
||||
if err != nil {
|
||||
return provider.ModelSyncResult{}, fmt.Errorf("%w: read provider %s response: %v", provider.ErrProviderUpstream, record.Code, err)
|
||||
}
|
||||
if int64(len(payload)) > maxModelsResponseBytes {
|
||||
return provider.ModelSyncResult{}, fmt.Errorf("%w: provider %s model response exceeds 4 MiB", provider.ErrProviderUpstream, record.Code)
|
||||
}
|
||||
models, err := decodeModels(payload)
|
||||
if err != nil {
|
||||
return provider.ModelSyncResult{}, fmt.Errorf("%w: provider %s returned invalid model data: %v", provider.ErrProviderUpstream, record.Code, err)
|
||||
}
|
||||
return s.repository.SyncModels(ctx, providerID, actorID, models)
|
||||
}
|
||||
|
||||
func (s *Service) ListModels(ctx context.Context, providerID string) ([]provider.Model, error) {
|
||||
if _, err := s.repository.Get(ctx, providerID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.repository.ListModels(ctx, providerID)
|
||||
}
|
||||
|
||||
func (s *Service) RotateCredentials(ctx context.Context, actorID string) (provider.CredentialRotationResult, error) {
|
||||
result := provider.CredentialRotationResult{
|
||||
ActiveVersion: s.cipher.ActiveVersion(), LoadedVersions: s.cipher.Versions(),
|
||||
}
|
||||
records, err := s.repository.List(ctx)
|
||||
if err != nil {
|
||||
return provider.CredentialRotationResult{}, err
|
||||
}
|
||||
rotations := make([]provider.CredentialRotation, 0, len(records))
|
||||
for _, record := range records {
|
||||
if record.CredentialKEKVersion == result.ActiveVersion {
|
||||
result.Skipped++
|
||||
continue
|
||||
}
|
||||
plaintext, err := s.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
|
||||
if err != nil {
|
||||
return provider.CredentialRotationResult{}, fmt.Errorf("provider %s credentials cannot be decrypted: %w", record.Code, err)
|
||||
}
|
||||
ciphertext, version, err := s.cipher.Encrypt(plaintext)
|
||||
if err != nil {
|
||||
return provider.CredentialRotationResult{}, err
|
||||
}
|
||||
rotations = append(rotations, provider.CredentialRotation{
|
||||
ProviderID: record.ID, FromVersion: record.CredentialKEKVersion,
|
||||
ToVersion: version, Ciphertext: ciphertext,
|
||||
})
|
||||
}
|
||||
if err := s.repository.RotateCredentials(ctx, actorID, rotations); err != nil {
|
||||
return provider.CredentialRotationResult{}, err
|
||||
}
|
||||
result.Rotated = len(rotations)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) load(ctx context.Context, providerID string) (provider.Record, provider.Adapter, error) {
|
||||
record, err := s.repository.Get(ctx, providerID)
|
||||
if err != nil {
|
||||
return provider.Record{}, nil, err
|
||||
}
|
||||
validatedURL, err := provider.ValidateBaseURL(ctx, record.BaseURL, s.allowPrivate)
|
||||
if err != nil {
|
||||
return provider.Record{}, nil, fmt.Errorf("供应商地址校验失败: %w", err)
|
||||
}
|
||||
plaintext, err := s.cipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion)
|
||||
if err != nil {
|
||||
return provider.Record{}, nil, err
|
||||
}
|
||||
var credentials provider.Credentials
|
||||
if err := json.Unmarshal(plaintext, &credentials); err != nil {
|
||||
return provider.Record{}, nil, errors.New("供应商凭据格式无效")
|
||||
}
|
||||
switch record.Adapter {
|
||||
case "openai-compatible":
|
||||
adapter, err := provideropenai.New(validatedURL, credentials.APIKey)
|
||||
return record, adapter, err
|
||||
default:
|
||||
return provider.Record{}, nil, fmt.Errorf("不支持的供应商适配器 %q", record.Adapter)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) doModelsRequest(ctx context.Context, adapter provider.Adapter) (*http.Response, error) {
|
||||
target := adapter.Target()
|
||||
target.Path = strings.TrimRight(target.Path, "/") + "/v1/models"
|
||||
target.RawPath = ""
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Accept", "application/json")
|
||||
adapter.Prepare(request)
|
||||
return s.client.Do(request)
|
||||
}
|
||||
|
||||
func decodeModels(payload []byte) ([]provider.DiscoveredModel, error) {
|
||||
var response struct {
|
||||
Data []json.RawMessage `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(payload, &response); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
seen := make(map[string]struct{}, len(response.Data))
|
||||
models := make([]provider.DiscoveredModel, 0, len(response.Data))
|
||||
for _, raw := range response.Data {
|
||||
var item struct {
|
||||
ID string `json:"id"`
|
||||
OwnedBy string `json:"owned_by"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item.ID = strings.TrimSpace(item.ID)
|
||||
if item.ID == "" || len(item.ID) > 512 {
|
||||
return nil, errors.New("model id must contain 1 to 512 characters")
|
||||
}
|
||||
if _, exists := seen[item.ID]; exists {
|
||||
continue
|
||||
}
|
||||
seen[item.ID] = struct{}{}
|
||||
models = append(models, provider.DiscoveredModel{
|
||||
ProviderModelID: item.ID,
|
||||
OwnedBy: strings.TrimSpace(item.OwnedBy),
|
||||
Metadata: append(json.RawMessage(nil), raw...),
|
||||
})
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func hasCapability(capabilities []string, expected string) bool {
|
||||
for _, capability := range capabilities {
|
||||
if capability == expected {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func safeDialContext(allowPrivate bool) func(context.Context, string, string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
|
||||
if allowPrivate {
|
||||
return dialer.DialContext
|
||||
}
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(addresses) == 0 {
|
||||
return nil, errors.New("upstream host did not resolve")
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if !isPublicAddress(address.IP) {
|
||||
return nil, fmt.Errorf("upstream resolved to blocked address %s", address.IP)
|
||||
}
|
||||
}
|
||||
return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
|
||||
}
|
||||
}
|
||||
|
||||
func isPublicAddress(ip net.IP) bool {
|
||||
return ip != nil && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast() &&
|
||||
!ip.IsLinkLocalMulticast() && !ip.IsMulticast() && !ip.IsUnspecified()
|
||||
}
|
||||
|
||||
var _ provider.AdminOperations = (*Service)(nil)
|
||||
@@ -0,0 +1,36 @@
|
||||
package controlplane
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDecodeModelsDeduplicatesAndRetainsMetadata(t *testing.T) {
|
||||
models, err := decodeModels([]byte(`{
|
||||
"object":"list",
|
||||
"data":[
|
||||
{"id":"gpt-5","owned_by":"openai","context_window":400000},
|
||||
{"id":"gpt-5","owned_by":"duplicate"},
|
||||
{"id":"text-embedding-3-small","owned_by":"openai"}
|
||||
]
|
||||
}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(models) != 2 || models[0].ProviderModelID != "gpt-5" || models[1].ProviderModelID != "text-embedding-3-small" {
|
||||
t.Fatalf("unexpected models: %#v", models)
|
||||
}
|
||||
var metadata map[string]any
|
||||
if err := json.Unmarshal(models[0].Metadata, &metadata); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if metadata["context_window"] != float64(400000) {
|
||||
t.Fatalf("metadata was not retained: %#v", metadata)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeModelsRejectsMissingID(t *testing.T) {
|
||||
if _, err := decodeModels([]byte(`{"data":[{"owned_by":"openai"}]}`)); err == nil {
|
||||
t.Fatal("expected missing model id to be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package provider
|
||||
|
||||
import "aigateway.local/core/internal/platform/cryptox"
|
||||
|
||||
var ErrCredentialKeyUnavailable = cryptox.ErrKeyUnavailable
|
||||
|
||||
type CredentialCipher = cryptox.Keyring
|
||||
|
||||
func NewCredentialCipher(encodedKey string, version int, encodedKeyring ...string) (*CredentialCipher, error) {
|
||||
keyring := ""
|
||||
if len(encodedKeyring) > 0 {
|
||||
keyring = encodedKeyring[0]
|
||||
}
|
||||
return cryptox.NewKeyring(encodedKey, version, keyring, "provider-credentials")
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCredentialCipherRoundTripAndAuthentication(t *testing.T) {
|
||||
key := base64.StdEncoding.EncodeToString(make([]byte, 32))
|
||||
cipher, err := NewCredentialCipher(key, 7)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
encrypted, version, err := cipher.Encrypt([]byte(`{"api_key":"secret"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
plaintext, err := cipher.Decrypt(encrypted, version)
|
||||
if err != nil || string(plaintext) != `{"api_key":"secret"}` {
|
||||
t.Fatalf("unexpected decrypt result %q: %v", plaintext, err)
|
||||
}
|
||||
encrypted[len(encrypted)-1] ^= 1
|
||||
if _, err := cipher.Decrypt(encrypted, version); err == nil {
|
||||
t.Fatal("tampered ciphertext must fail authentication")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrModelRouteNotFound = errors.New("model route not found")
|
||||
ErrModelRouteExists = errors.New("model route already exists")
|
||||
)
|
||||
|
||||
type ModelRoute struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
SourceModel string `json:"source_model"`
|
||||
TargetModel string `json:"target_model"`
|
||||
ProviderID string `json:"provider_id"`
|
||||
ProviderCode string `json:"provider_code"`
|
||||
Weight int `json:"weight"`
|
||||
Priority int `json:"priority"`
|
||||
Conditions json.RawMessage `json:"conditions"`
|
||||
Enabled bool `json:"enabled"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
func (r *Repository) ListModelRoutes(ctx context.Context) ([]ModelRoute, error) {
|
||||
if r.pool == nil {
|
||||
return nil, ErrProviderStore
|
||||
}
|
||||
rows, err := r.pool.Query(ctx, `
|
||||
SELECT mr.id::text, mr.name, mr.source_model, mr.target_model, mr.provider_id::text,
|
||||
p.code, mr.weight, mr.priority, mr.conditions, mr.enabled, mr.created_at, mr.updated_at
|
||||
FROM gateway.model_routes mr JOIN gateway.providers p ON p.id = mr.provider_id
|
||||
ORDER BY mr.source_model, mr.priority DESC, mr.name, p.code`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
routes := make([]ModelRoute, 0)
|
||||
for rows.Next() {
|
||||
var route ModelRoute
|
||||
if err := rows.Scan(&route.ID, &route.Name, &route.SourceModel, &route.TargetModel, &route.ProviderID,
|
||||
&route.ProviderCode, &route.Weight, &route.Priority, &route.Conditions, &route.Enabled, &route.CreatedAt, &route.UpdatedAt); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
routes = append(routes, route)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return routes, nil
|
||||
}
|
||||
|
||||
func (r *Repository) CreateModelRoute(ctx context.Context, route ModelRoute, actorID string) (ModelRoute, error) {
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return ModelRoute{}, err
|
||||
}
|
||||
route.ID = id
|
||||
return r.saveModelRoute(ctx, route, actorID, true)
|
||||
}
|
||||
|
||||
func (r *Repository) UpdateModelRoute(ctx context.Context, route ModelRoute, actorID string) (ModelRoute, error) {
|
||||
return r.saveModelRoute(ctx, route, actorID, false)
|
||||
}
|
||||
|
||||
func (r *Repository) saveModelRoute(ctx context.Context, route ModelRoute, actorID string, create bool) (ModelRoute, error) {
|
||||
if r.pool == nil {
|
||||
return ModelRoute{}, ErrProviderStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return ModelRoute{}, err
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return ModelRoute{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
var providerExists bool
|
||||
if err := tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.providers WHERE id=$1 AND tenant_id IS NULL)`, route.ProviderID).Scan(&providerExists); err != nil {
|
||||
return ModelRoute{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if !providerExists {
|
||||
return ModelRoute{}, ErrProviderNotFound
|
||||
}
|
||||
if create {
|
||||
err = tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway.model_routes
|
||||
(id,name,source_model,target_model,provider_id,weight,priority,conditions,enabled,created_by)
|
||||
SELECT $1,$2,$3,$4,p.id,$6,$7,$8,$9,$10 FROM gateway.providers p WHERE p.id=$5 AND p.tenant_id IS NULL
|
||||
RETURNING created_at,updated_at`, route.ID, route.Name, route.SourceModel, route.TargetModel,
|
||||
route.ProviderID, route.Weight, route.Priority, route.Conditions, route.Enabled, actorID,
|
||||
).Scan(&route.CreatedAt, &route.UpdatedAt)
|
||||
} else {
|
||||
err = tx.QueryRow(ctx, `
|
||||
UPDATE gateway.model_routes
|
||||
SET name=$2,source_model=$3,target_model=$4,provider_id=$5,weight=$6,priority=$7,
|
||||
conditions=$8,enabled=$9,updated_at=clock_timestamp()
|
||||
WHERE id=$1
|
||||
RETURNING created_at,updated_at`, route.ID, route.Name, route.SourceModel, route.TargetModel,
|
||||
route.ProviderID, route.Weight, route.Priority, route.Conditions, route.Enabled,
|
||||
).Scan(&route.CreatedAt, &route.UpdatedAt)
|
||||
}
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
if create {
|
||||
return ModelRoute{}, ErrProviderNotFound
|
||||
}
|
||||
return ModelRoute{}, ErrModelRouteNotFound
|
||||
}
|
||||
if err != nil {
|
||||
var pgErr *pgconn.PgError
|
||||
if errors.As(err, &pgErr) {
|
||||
if pgErr.Code == "23505" {
|
||||
return ModelRoute{}, ErrModelRouteExists
|
||||
}
|
||||
if pgErr.Code == "23503" {
|
||||
return ModelRoute{}, ErrProviderNotFound
|
||||
}
|
||||
}
|
||||
return ModelRoute{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if err := tx.QueryRow(ctx, `SELECT code FROM gateway.providers WHERE id=$1`, route.ProviderID).Scan(&route.ProviderCode); err != nil {
|
||||
return ModelRoute{}, ErrProviderNotFound
|
||||
}
|
||||
eventType := "model_route.updated"
|
||||
if create {
|
||||
eventType = "model_route.created"
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"model_route_id": route.ID, "source_model": route.SourceModel, "target_model": route.TargetModel, "provider_id": route.ProviderID, "actor_id": actorID})
|
||||
if _, err := tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'model_route',$3,$4)`, eventID, eventType, route.ID, payload); err != nil {
|
||||
return ModelRoute{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return ModelRoute{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return route, nil
|
||||
}
|
||||
|
||||
func (r *Repository) DeleteModelRoute(ctx context.Context, id, actorID string) error {
|
||||
if r.pool == nil {
|
||||
return ErrProviderStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
result, err := tx.Exec(ctx, `DELETE FROM gateway.model_routes WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if result.RowsAffected() == 0 {
|
||||
return ErrModelRouteNotFound
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"model_route_id": id, "actor_id": actorID})
|
||||
if _, err := tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,'model_route.deleted',1,'model_route',$2,$3)`, eventID, id, payload); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"aigateway.local/core/internal/provider"
|
||||
)
|
||||
|
||||
type Adapter struct {
|
||||
target *url.URL
|
||||
apiKey string
|
||||
}
|
||||
|
||||
func New(baseURL, apiKey string) (*Adapter, error) {
|
||||
target, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Adapter{target: target, apiKey: apiKey}, nil
|
||||
}
|
||||
|
||||
func (a *Adapter) Name() string { return "openai-compatible" }
|
||||
|
||||
func (a *Adapter) Target() *url.URL {
|
||||
copy := *a.target
|
||||
return ©
|
||||
}
|
||||
|
||||
func (a *Adapter) Capabilities() []provider.Capability {
|
||||
return []provider.Capability{
|
||||
provider.CapabilityChat,
|
||||
provider.CapabilityResponses,
|
||||
provider.CapabilityEmbeddings,
|
||||
provider.CapabilityMessages,
|
||||
provider.CapabilityModels,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adapter) Prepare(request *http.Request) {
|
||||
request.Header.Del("X-Gateway-API-Key")
|
||||
request.Header.Del("X-Gateway-Provider")
|
||||
basePath := strings.TrimRight(a.target.Path, "/")
|
||||
duplicatedVersionPrefix := basePath + "/v1/"
|
||||
if strings.HasSuffix(basePath, "/v1") && strings.HasPrefix(request.URL.Path, duplicatedVersionPrefix) {
|
||||
request.URL.Path = basePath + "/" + strings.TrimPrefix(request.URL.Path, duplicatedVersionPrefix)
|
||||
request.URL.RawPath = ""
|
||||
}
|
||||
if a.apiKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+a.apiKey)
|
||||
}
|
||||
}
|
||||
|
||||
var _ provider.Adapter = (*Adapter)(nil)
|
||||
@@ -0,0 +1,21 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPrepareAvoidsDuplicatedV1BasePath(t *testing.T) {
|
||||
adapter, err := New("https://example.com/proxy/v1", "secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := httptest.NewRequest("POST", "https://example.com/proxy/v1/v1/chat/completions", nil)
|
||||
adapter.Prepare(request)
|
||||
if request.URL.Path != "/proxy/v1/chat/completions" {
|
||||
t.Fatalf("unexpected path %q", request.URL.Path)
|
||||
}
|
||||
if request.Header.Get("Authorization") != "Bearer secret" {
|
||||
t.Fatal("upstream credential was not applied")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Capability string
|
||||
|
||||
const (
|
||||
CapabilityChat Capability = "chat"
|
||||
CapabilityResponses Capability = "responses"
|
||||
CapabilityEmbeddings Capability = "embeddings"
|
||||
CapabilityMessages Capability = "messages"
|
||||
CapabilityModels Capability = "models"
|
||||
)
|
||||
|
||||
type Adapter interface {
|
||||
Name() string
|
||||
Target() *url.URL
|
||||
Capabilities() []Capability
|
||||
Prepare(*http.Request)
|
||||
}
|
||||
|
||||
type Credentials struct {
|
||||
APIKey string `json:"api_key"`
|
||||
}
|
||||
|
||||
type ConnectionTestResult struct {
|
||||
Connected bool `json:"connected"`
|
||||
StatusCode int `json:"status_code,omitempty"`
|
||||
LatencyMS int64 `json:"latency_ms"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type Model struct {
|
||||
ID string `json:"id"`
|
||||
ProviderID string `json:"provider_id"`
|
||||
ProviderModelID string `json:"provider_model_id"`
|
||||
OwnedBy string `json:"owned_by,omitempty"`
|
||||
Metadata json.RawMessage `json:"metadata"`
|
||||
Enabled bool `json:"enabled"`
|
||||
DiscoveredAt time.Time `json:"discovered_at"`
|
||||
LastSeenAt time.Time `json:"last_seen_at"`
|
||||
}
|
||||
|
||||
type DiscoveredModel struct {
|
||||
ProviderModelID string
|
||||
OwnedBy string
|
||||
Metadata json.RawMessage
|
||||
}
|
||||
|
||||
type ModelSyncResult struct {
|
||||
Discovered int `json:"discovered"`
|
||||
Active int `json:"active"`
|
||||
Disabled int `json:"disabled"`
|
||||
SyncedAt time.Time `json:"synced_at"`
|
||||
}
|
||||
|
||||
type CredentialRotation struct {
|
||||
ProviderID string
|
||||
FromVersion int
|
||||
ToVersion int
|
||||
Ciphertext []byte
|
||||
}
|
||||
|
||||
type CredentialRotationResult struct {
|
||||
ActiveVersion int `json:"active_version"`
|
||||
LoadedVersions []int `json:"loaded_versions"`
|
||||
Rotated int `json:"rotated"`
|
||||
Skipped int `json:"skipped"`
|
||||
}
|
||||
|
||||
type AdminOperations interface {
|
||||
TestConnection(context.Context, string) (ConnectionTestResult, error)
|
||||
SyncModels(context.Context, string, string) (ModelSyncResult, error)
|
||||
ListModels(context.Context, string) ([]Model, error)
|
||||
RotateCredentials(context.Context, string) (CredentialRotationResult, error)
|
||||
}
|
||||
@@ -0,0 +1,357 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrProviderNotFound = errors.New("provider not found")
|
||||
ErrProviderExists = errors.New("provider already exists")
|
||||
ErrMultipleDefaults = errors.New("multiple default providers")
|
||||
ErrProviderStore = errors.New("provider store unavailable")
|
||||
ErrProviderUpstream = errors.New("provider upstream unavailable")
|
||||
)
|
||||
|
||||
type Record struct {
|
||||
ID string
|
||||
TenantID *string
|
||||
Code string
|
||||
Adapter string
|
||||
BaseURL string
|
||||
EncryptedCredentials []byte
|
||||
CredentialKEKVersion int
|
||||
Capabilities []string
|
||||
Config json.RawMessage
|
||||
Enabled bool
|
||||
Revision int64
|
||||
}
|
||||
|
||||
func (r *Repository) ListModels(ctx context.Context, providerID string) ([]Model, error) {
|
||||
if r.pool == nil {
|
||||
return nil, ErrProviderStore
|
||||
}
|
||||
rows, err := r.pool.Query(ctx, `
|
||||
SELECT id::text, provider_id::text, provider_model_id, owned_by, metadata,
|
||||
enabled, discovered_at, last_seen_at
|
||||
FROM gateway.provider_models
|
||||
WHERE provider_id = $1
|
||||
ORDER BY enabled DESC, provider_model_id`, providerID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
models := make([]Model, 0)
|
||||
for rows.Next() {
|
||||
var model Model
|
||||
if err := rows.Scan(
|
||||
&model.ID, &model.ProviderID, &model.ProviderModelID, &model.OwnedBy,
|
||||
&model.Metadata, &model.Enabled, &model.DiscoveredAt, &model.LastSeenAt,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
models = append(models, model)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func (r *Repository) SyncModels(ctx context.Context, providerID, actorID string, models []DiscoveredModel) (ModelSyncResult, error) {
|
||||
if r.pool == nil {
|
||||
return ModelSyncResult{}, ErrProviderStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return ModelSyncResult{}, err
|
||||
}
|
||||
transaction, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = transaction.Rollback(ctx) }()
|
||||
var exists bool
|
||||
if err := transaction.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM gateway.providers WHERE id = $1 AND tenant_id IS NULL
|
||||
)`, providerID).Scan(&exists); err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if !exists {
|
||||
return ModelSyncResult{}, ErrProviderNotFound
|
||||
}
|
||||
if _, err := transaction.Exec(ctx, `
|
||||
UPDATE gateway.provider_models
|
||||
SET enabled = false, updated_at = clock_timestamp()
|
||||
WHERE provider_id = $1 AND enabled = true`, providerID); err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
for _, model := range models {
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return ModelSyncResult{}, err
|
||||
}
|
||||
if _, err := transaction.Exec(ctx, `
|
||||
INSERT INTO gateway.provider_models
|
||||
(id, provider_id, provider_model_id, owned_by, metadata, enabled)
|
||||
VALUES ($1, $2, $3, $4, $5, true)
|
||||
ON CONFLICT (provider_id, provider_model_id) DO UPDATE
|
||||
SET owned_by = EXCLUDED.owned_by,
|
||||
metadata = EXCLUDED.metadata,
|
||||
enabled = true,
|
||||
last_seen_at = clock_timestamp(),
|
||||
updated_at = clock_timestamp()`,
|
||||
id, providerID, model.ProviderModelID, model.OwnedBy, model.Metadata,
|
||||
); err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
}
|
||||
result := ModelSyncResult{Discovered: len(models), SyncedAt: time.Now().UTC()}
|
||||
if err := transaction.QueryRow(ctx, `
|
||||
SELECT count(*) FILTER (WHERE enabled), count(*) FILTER (WHERE NOT enabled)
|
||||
FROM gateway.provider_models
|
||||
WHERE provider_id = $1`, providerID).Scan(&result.Active, &result.Disabled); err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"provider_id": providerID, "discovered": result.Discovered,
|
||||
"active": result.Active, "disabled": result.Disabled, "actor_id": actorID,
|
||||
})
|
||||
if _, err := transaction.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events
|
||||
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'provider.models_synced', 1, 'provider', $2, $3)`, eventID, providerID, payload); err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if err := transaction.Commit(ctx); err != nil {
|
||||
return ModelSyncResult{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r *Repository) RotateCredentials(ctx context.Context, actorID string, rotations []CredentialRotation) error {
|
||||
if r.pool == nil {
|
||||
return ErrProviderStore
|
||||
}
|
||||
if len(rotations) == 0 {
|
||||
return nil
|
||||
}
|
||||
transaction, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = transaction.Rollback(ctx) }()
|
||||
for _, rotation := range rotations {
|
||||
result, err := transaction.Exec(ctx, `
|
||||
UPDATE gateway.providers
|
||||
SET encrypted_credentials = $2,
|
||||
credential_kek_version = $3,
|
||||
revision = revision + 1,
|
||||
updated_at = clock_timestamp()
|
||||
WHERE id = $1
|
||||
AND tenant_id IS NULL
|
||||
AND credential_kek_version = $4`,
|
||||
rotation.ProviderID, rotation.Ciphertext, rotation.ToVersion, rotation.FromVersion,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if result.RowsAffected() != 1 {
|
||||
return fmt.Errorf("%w: provider %s changed during credential rotation", ErrProviderStore, rotation.ProviderID)
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"provider_id": rotation.ProviderID, "from_version": rotation.FromVersion,
|
||||
"to_version": rotation.ToVersion, "actor_id": actorID,
|
||||
})
|
||||
if _, err := transaction.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events
|
||||
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'provider.credentials_rotated', 1, 'provider', $2, $3)`,
|
||||
eventID, rotation.ProviderID, payload,
|
||||
); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
}
|
||||
if err := transaction.Commit(ctx); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type Repository struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
func NewRepository(pool *pgxpool.Pool) *Repository {
|
||||
return &Repository{pool: pool}
|
||||
}
|
||||
|
||||
func (r *Repository) List(ctx context.Context) ([]Record, error) {
|
||||
if r.pool == nil {
|
||||
return nil, ErrProviderStore
|
||||
}
|
||||
rows, err := r.pool.Query(ctx, `
|
||||
SELECT id::text, tenant_id::text, code, adapter, base_url,
|
||||
encrypted_credentials, credential_kek_version, capabilities,
|
||||
config, enabled, revision
|
||||
FROM gateway.providers
|
||||
WHERE tenant_id IS NULL
|
||||
ORDER BY code`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
var records []Record
|
||||
for rows.Next() {
|
||||
var record Record
|
||||
if err := rows.Scan(
|
||||
&record.ID, &record.TenantID, &record.Code, &record.Adapter, &record.BaseURL,
|
||||
&record.EncryptedCredentials, &record.CredentialKEKVersion, &record.Capabilities,
|
||||
&record.Config, &record.Enabled, &record.Revision,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
records = append(records, record)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return records, nil
|
||||
}
|
||||
|
||||
func (r *Repository) Get(ctx context.Context, id string) (Record, error) {
|
||||
if r.pool == nil {
|
||||
return Record{}, ErrProviderStore
|
||||
}
|
||||
var record Record
|
||||
err := r.pool.QueryRow(ctx, `
|
||||
SELECT id::text, tenant_id::text, code, adapter, base_url,
|
||||
encrypted_credentials, credential_kek_version, capabilities,
|
||||
config, enabled, revision
|
||||
FROM gateway.providers
|
||||
WHERE id = $1 AND tenant_id IS NULL`, id).Scan(
|
||||
&record.ID, &record.TenantID, &record.Code, &record.Adapter, &record.BaseURL,
|
||||
&record.EncryptedCredentials, &record.CredentialKEKVersion, &record.Capabilities,
|
||||
&record.Config, &record.Enabled, &record.Revision,
|
||||
)
|
||||
return record, mapProviderError(err)
|
||||
}
|
||||
|
||||
func (r *Repository) Create(ctx context.Context, record Record, actorID string) (Record, error) {
|
||||
if r.pool == nil {
|
||||
return Record{}, ErrProviderStore
|
||||
}
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Record{}, err
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Record{}, err
|
||||
}
|
||||
record.ID = id
|
||||
transaction, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Record{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = transaction.Rollback(ctx) }()
|
||||
err = transaction.QueryRow(ctx, `
|
||||
INSERT INTO gateway.providers
|
||||
(id, code, adapter, base_url, encrypted_credentials,
|
||||
credential_kek_version, capabilities, config, enabled, created_by)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
RETURNING revision`, record.ID, record.Code, record.Adapter, record.BaseURL,
|
||||
record.EncryptedCredentials, record.CredentialKEKVersion, record.Capabilities,
|
||||
record.Config, record.Enabled, actorID).Scan(&record.Revision)
|
||||
if err != nil {
|
||||
return Record{}, mapProviderError(err)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"provider_id": record.ID, "code": record.Code, "revision": record.Revision})
|
||||
if _, err := transaction.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events
|
||||
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'provider.created', 1, 'provider', $2, $3)`, eventID, record.ID, payload); err != nil {
|
||||
return Record{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if err := transaction.Commit(ctx); err != nil {
|
||||
return Record{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (r *Repository) Update(ctx context.Context, record Record, actorID string, replaceCredentials bool) (Record, error) {
|
||||
if r.pool == nil {
|
||||
return Record{}, ErrProviderStore
|
||||
}
|
||||
eventID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Record{}, err
|
||||
}
|
||||
transaction, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Record{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
defer func() { _ = transaction.Rollback(ctx) }()
|
||||
err = transaction.QueryRow(ctx, `
|
||||
UPDATE gateway.providers
|
||||
SET code = $2,
|
||||
adapter = $3,
|
||||
base_url = $4,
|
||||
encrypted_credentials = CASE WHEN $10 THEN $5 ELSE encrypted_credentials END,
|
||||
credential_kek_version = CASE WHEN $10 THEN $6 ELSE credential_kek_version END,
|
||||
capabilities = $7,
|
||||
config = $8,
|
||||
enabled = $9,
|
||||
revision = revision + 1,
|
||||
updated_at = clock_timestamp()
|
||||
WHERE id = $1 AND tenant_id IS NULL
|
||||
RETURNING revision`, record.ID, record.Code, record.Adapter, record.BaseURL,
|
||||
record.EncryptedCredentials, record.CredentialKEKVersion, record.Capabilities,
|
||||
record.Config, record.Enabled, replaceCredentials).Scan(&record.Revision)
|
||||
if err != nil {
|
||||
return Record{}, mapProviderError(err)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"provider_id": record.ID, "code": record.Code, "revision": record.Revision, "actor_id": actorID,
|
||||
})
|
||||
if _, err := transaction.Exec(ctx, `
|
||||
INSERT INTO gateway.outbox_events
|
||||
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
|
||||
VALUES ($1, 'provider.updated', 1, 'provider', $2, $3)`, eventID, record.ID, payload); err != nil {
|
||||
return Record{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
if err := transaction.Commit(ctx); err != nil {
|
||||
return Record{}, fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func mapProviderError(err error) error {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return ErrProviderNotFound
|
||||
}
|
||||
var pgError *pgconn.PgError
|
||||
if errors.As(err, &pgError) && pgError.Code == "23505" {
|
||||
if pgError.ConstraintName == "providers_single_global_default_idx" {
|
||||
return ErrMultipleDefaults
|
||||
}
|
||||
return ErrProviderExists
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrProviderStore, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/provider"
|
||||
provideropenai "aigateway.local/core/internal/provider/openai"
|
||||
)
|
||||
|
||||
func TestBuildRuntimeAdapter(t *testing.T) {
|
||||
key := base64.StdEncoding.EncodeToString(make([]byte, 32))
|
||||
cipher, err := provider.NewCredentialCipher(key, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
encrypted, version, err := cipher.Encrypt([]byte(`{"api_key":"upstream-secret"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fallback, _ := provideropenai.New("https://fallback.example", "")
|
||||
resolver := NewResolver(nil, cipher, fallback, time.Second, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
resolved, defaultProvider, err := resolver.build(provider.Record{
|
||||
Code: "primary", Adapter: "openai-compatible", BaseURL: "https://upstream.example/v1",
|
||||
EncryptedCredentials: encrypted, CredentialKEKVersion: version,
|
||||
Capabilities: []string{"chat", "models"}, Config: json.RawMessage(`{"default":true}`), Revision: 3,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resolved.Code != "primary" || resolved.Revision != 3 || !resolved.Capabilities[provider.CapabilityChat] || !defaultProvider {
|
||||
t.Fatalf("unexpected resolved adapter: %#v, default=%v", resolved, defaultProvider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRouteUsesConditionsPriorityAndExplicitProvider(t *testing.T) {
|
||||
first, _ := provideropenai.New("https://first.example", "")
|
||||
second, _ := provideropenai.New("https://second.example", "")
|
||||
resolver := &Resolver{}
|
||||
resolver.snapshot.Store(&snapshot{
|
||||
adapters: map[string]gateway.ResolvedAdapter{
|
||||
"first": {Code: "first", Adapter: first}, "second": {Code: "second", Adapter: second},
|
||||
},
|
||||
routes: map[string][]compiledRoute{"public-chat": {
|
||||
{providerCode: "first", targetModel: "model-a", weight: 100, priority: 10},
|
||||
{providerCode: "second", targetModel: "model-b", weight: 100, priority: 20, conditions: routeConditions{Endpoints: []string{"/v1/chat/completions"}}},
|
||||
}},
|
||||
knownModels: map[string]bool{"public-chat": true},
|
||||
})
|
||||
result, err := resolver.ResolveModelRoute(gateway.ModelRouteQuery{Model: "public-chat", Endpoint: "/v1/chat/completions", Seed: "request-1"})
|
||||
if err != nil || !result.Matched || result.Code != "second" || result.TargetModel != "model-b" {
|
||||
t.Fatalf("unexpected priority route: %#v %v", result, err)
|
||||
}
|
||||
result, err = resolver.ResolveModelRoute(gateway.ModelRouteQuery{ProviderCode: "first", Model: "public-chat", Endpoint: "/v1/chat/completions", Seed: "request-1"})
|
||||
if err != nil || !result.Matched || result.Code != "first" {
|
||||
t.Fatalf("explicit provider was not respected: %#v %v", result, err)
|
||||
}
|
||||
result, err = resolver.ResolveModelRoute(gateway.ModelRouteQuery{ProviderCode: "second", Model: "public-chat", Endpoint: "/v1/embeddings", Seed: "request-1"})
|
||||
if err != nil || result.Matched || !result.Known {
|
||||
t.Fatalf("known alias with unmatched conditions must fail closed: %#v %v", result, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolverStartsWithEnvironmentFallback(t *testing.T) {
|
||||
fallback, _ := provideropenai.New("https://fallback.example", "")
|
||||
resolver := NewResolver(nil, nil, fallback, time.Second, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
resolved, err := resolver.Resolve("")
|
||||
if err != nil || resolved.Code != "environment" {
|
||||
t.Fatalf("fallback resolve: %#v %v", resolved, err)
|
||||
}
|
||||
if _, err := resolver.Resolve("missing"); err == nil {
|
||||
t.Fatal("unknown provider must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func ValidateBaseURL(ctx context.Context, raw string, allowPrivate bool) (string, error) {
|
||||
parsed, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Hostname() == "" {
|
||||
return "", errors.New("base_url must be an absolute http(s) URL")
|
||||
}
|
||||
if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
|
||||
return "", errors.New("base_url must not contain user info, query, or fragment")
|
||||
}
|
||||
if !allowPrivate {
|
||||
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, parsed.Hostname())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve base_url host: %w", err)
|
||||
}
|
||||
if len(addresses) == 0 {
|
||||
return "", errors.New("base_url host did not resolve")
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if !isPublicAddress(address.IP) {
|
||||
return "", fmt.Errorf("base_url resolves to blocked address %s", address.IP)
|
||||
}
|
||||
}
|
||||
}
|
||||
parsed.Path = strings.TrimRight(parsed.Path, "/")
|
||||
return parsed.String(), nil
|
||||
}
|
||||
|
||||
func isPublicAddress(ip net.IP) bool {
|
||||
return ip != nil && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast() &&
|
||||
!ip.IsLinkLocalMulticast() && !ip.IsMulticast() && !ip.IsUnspecified()
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestValidateBaseURLBlocksPrivateLiteral(t *testing.T) {
|
||||
if _, err := ValidateBaseURL(context.Background(), "http://127.0.0.1:8080/v1", false); err == nil {
|
||||
t.Fatal("expected loopback URL to be blocked")
|
||||
}
|
||||
got, err := ValidateBaseURL(context.Background(), "http://127.0.0.1:8080/v1/", true)
|
||||
if err != nil || got != "http://127.0.0.1:8080/v1" {
|
||||
t.Fatalf("unexpected private URL result %q: %v", got, err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user