renjue
76ba500417
CI / docker (push) Successful in 2m4s
OpenAI-compatible AI gateway with Vue admin UI, multi-provider egress, ingress key governance, monitoring, and security controls. Co-authored-by: Cursor <cursoragent@cursor.com>
476 lines
16 KiB
Go
476 lines
16 KiB
Go
package gateway
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/rose_cat707/luminary/internal/auth"
|
|
"github.com/rose_cat707/luminary/internal/limiter"
|
|
"github.com/rose_cat707/luminary/internal/model"
|
|
"github.com/rose_cat707/luminary/internal/pricing"
|
|
"github.com/rose_cat707/luminary/internal/provider"
|
|
"github.com/rose_cat707/luminary/internal/provider/anthropic"
|
|
"github.com/rose_cat707/luminary/internal/provider/comfyui"
|
|
"github.com/rose_cat707/luminary/internal/provider/openai"
|
|
"github.com/rose_cat707/luminary/internal/settings"
|
|
"github.com/rose_cat707/luminary/internal/store/cache"
|
|
)
|
|
|
|
type Service struct {
|
|
cache *cache.Store
|
|
settings *settings.Runtime
|
|
openai *openai.Client
|
|
anthropic *anthropic.Client
|
|
streamClient *provider.HTTPClientWrapper
|
|
}
|
|
|
|
func New(c *cache.Store, rt *settings.Runtime) *Service {
|
|
return &Service{
|
|
cache: c,
|
|
settings: rt,
|
|
openai: openai.New(),
|
|
anthropic: anthropic.New(),
|
|
streamClient: &provider.HTTPClientWrapper{Client: provider.NewStreamHTTPClient()},
|
|
}
|
|
}
|
|
|
|
func (s *Service) ResolveIngressKey(plainKey string) (*model.IngressKey, error) {
|
|
k, ok := s.cache.LookupIngressKey(plainKey)
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid api key")
|
|
}
|
|
if k.ExpiresAt != nil && k.ExpiresAt.Before(time.Now()) {
|
|
return nil, fmt.Errorf("api key expired")
|
|
}
|
|
return k, nil
|
|
}
|
|
|
|
func (s *Service) providerConfig(p *model.Provider) provider.ProviderConfig {
|
|
return provider.ProviderConfig{
|
|
Type: p.Type,
|
|
Category: p.Category,
|
|
BaseURL: p.BaseURL,
|
|
EndpointPaths: provider.ParseEndpointPaths(p.EndpointPathsJSON),
|
|
}
|
|
}
|
|
|
|
func (s *Service) DecryptAPIKey(enc string) (string, error) {
|
|
return auth.Decrypt(enc, s.settings.EncryptionKey())
|
|
}
|
|
|
|
func (s *Service) providerAllowsEndpoint(p *model.Provider, endpoint provider.EndpointType) bool {
|
|
allowed := provider.ParseAllowedEndpoints(p.AllowedEndpointsJSON, p.Category)
|
|
return provider.IsEndpointAllowed(allowed, endpoint)
|
|
}
|
|
|
|
func (s *Service) canServeEndpoint(p *model.Provider, endpoint provider.EndpointType) bool {
|
|
if s.providerAllowsEndpoint(p, endpoint) {
|
|
return true
|
|
}
|
|
return endpoint == provider.EndpointResponses && s.providerAllowsEndpoint(p, provider.EndpointChatCompletion)
|
|
}
|
|
|
|
func (s *Service) shouldAdaptResponses(p *model.Provider) bool {
|
|
return !s.providerAllowsEndpoint(p, provider.EndpointResponses) && s.providerAllowsEndpoint(p, provider.EndpointChatCompletion)
|
|
}
|
|
|
|
func (s *Service) forward(ctx context.Context, p *model.Provider, apiKey string, endpoint provider.EndpointType, body []byte) (*provider.ChatResponse, error) {
|
|
var err error
|
|
body, err = s.rewriteBodyForProvider(p.ID, body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
actualEndpoint := endpoint
|
|
actualBody := body
|
|
adaptResponses := false
|
|
|
|
if endpoint == provider.EndpointResponses {
|
|
if s.shouldAdaptResponses(p) {
|
|
adapted, err := provider.ResponsesToChatCompletion(body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
actualBody = adapted
|
|
actualEndpoint = provider.EndpointChatCompletion
|
|
adaptResponses = true
|
|
} else if !s.providerAllowsEndpoint(p, endpoint) {
|
|
return nil, fmt.Errorf("endpoint %s not enabled on provider %s", endpoint, p.Name)
|
|
}
|
|
} else if !s.providerAllowsEndpoint(p, endpoint) {
|
|
return nil, fmt.Errorf("endpoint %s not enabled on provider %s", endpoint, p.Name)
|
|
}
|
|
|
|
resp, err := s.forwardRaw(ctx, p, apiKey, actualEndpoint, actualBody)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if endpoint == provider.EndpointResponses && !adaptResponses && resp != nil && resp.StatusCode == http.StatusNotFound && s.providerAllowsEndpoint(p, provider.EndpointChatCompletion) {
|
|
adaptedBody, adaptErr := provider.ResponsesToChatCompletion(body)
|
|
if adaptErr == nil {
|
|
if resp2, err2 := s.forwardRaw(ctx, p, apiKey, provider.EndpointChatCompletion, adaptedBody); err2 == nil && resp2 != nil {
|
|
resp = resp2
|
|
adaptResponses = true
|
|
}
|
|
}
|
|
}
|
|
|
|
if adaptResponses && resp != nil && resp.StatusCode < 400 {
|
|
converted, convErr := provider.ChatCompletionToResponses(resp.Body)
|
|
if convErr == nil {
|
|
resp.Body = converted
|
|
}
|
|
}
|
|
return resp, nil
|
|
}
|
|
|
|
func (s *Service) forwardRaw(ctx context.Context, p *model.Provider, apiKey string, endpoint provider.EndpointType, body []byte) (*provider.ChatResponse, error) {
|
|
cfg := s.providerConfig(p)
|
|
if p.Type == model.ProviderAnthropic && endpoint == provider.EndpointChatCompletion {
|
|
return s.anthropic.ForwardWithConfig(ctx, apiKey, cfg, endpoint, body)
|
|
}
|
|
return s.openai.ForwardWithConfig(ctx, apiKey, cfg, endpoint, body)
|
|
}
|
|
|
|
func (s *Service) forwardStream(ctx context.Context, p *model.Provider, apiKey string, endpoint provider.EndpointType, body []byte, w io.Writer, flusher http.Flusher) (*provider.StreamResult, error) {
|
|
var err error
|
|
body, err = s.rewriteBodyForProvider(p.ID, body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
actualEndpoint := endpoint
|
|
actualBody := body
|
|
adaptResponses := false
|
|
|
|
if endpoint == provider.EndpointResponses {
|
|
if s.shouldAdaptResponses(p) {
|
|
adapted, err := provider.ResponsesToChatCompletion(body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
actualBody = adapted
|
|
actualEndpoint = provider.EndpointChatCompletion
|
|
adaptResponses = true
|
|
} else if !s.providerAllowsEndpoint(p, endpoint) {
|
|
return nil, fmt.Errorf("endpoint %s not enabled", endpoint)
|
|
}
|
|
} else if !s.providerAllowsEndpoint(p, endpoint) {
|
|
return nil, fmt.Errorf("endpoint %s not enabled", endpoint)
|
|
}
|
|
|
|
cfg := s.providerConfig(p)
|
|
url := cfg.ResolveEndpointURL(actualEndpoint)
|
|
|
|
if adaptResponses {
|
|
pr, pw := io.Pipe()
|
|
errCh := make(chan error, 1)
|
|
go func() {
|
|
_, err := provider.ForwardStreamHTTP(ctx, s.streamClient.Client, "POST", url, apiKey, p.Type, actualBody, pw, nil)
|
|
_ = pw.Close()
|
|
errCh <- err
|
|
}()
|
|
result, err := provider.AdaptResponsesStream(body, pr, w, flusher)
|
|
if streamErr := <-errCh; err == nil && streamErr != nil {
|
|
err = streamErr
|
|
}
|
|
return result, err
|
|
}
|
|
|
|
return provider.ForwardStreamHTTP(ctx, s.streamClient.Client, "POST", url, apiKey, p.Type, actualBody, w, flusher)
|
|
}
|
|
|
|
type forwardParams struct {
|
|
ingress *model.IngressKey
|
|
endpoint provider.EndpointType
|
|
category model.ProviderCategory
|
|
body []byte
|
|
stream bool
|
|
}
|
|
|
|
func (s *Service) prepareForward(params forwardParams) (modelName string, providers []model.Provider, route *routeTarget, err error) {
|
|
if len(params.body) > 0 {
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(params.body, &payload); err != nil {
|
|
return "", nil, nil, fmt.Errorf("%w: %v", ErrInvalidJSON, err)
|
|
}
|
|
modelName, _ = payload["model"].(string)
|
|
}
|
|
if params.endpoint != provider.EndpointListModels && modelName == "" {
|
|
return "", nil, nil, ErrModelRequired
|
|
}
|
|
allowedProviders := cache.ParseJSONList(params.ingress.AllowedProvidersJSON)
|
|
if modelName != "" {
|
|
allowedModels := cache.ParseJSONList(params.ingress.AllowedModelsJSON)
|
|
if !cache.MatchesList(allowedModels, modelName) && !cache.MatchesList(allowedModels, "*") {
|
|
return "", nil, nil, ErrModelNotAllowed
|
|
}
|
|
}
|
|
if !limiter.CheckIngressBudget(params.ingress, 0) {
|
|
return "", nil, nil, ErrBudgetExceeded
|
|
}
|
|
if !limiter.CheckRateLimit(s.cache, params.ingress.ID, params.ingress.RateLimitRPM, params.ingress.RateLimitTPM, 1000) {
|
|
return "", nil, nil, ErrRateLimitExceeded
|
|
}
|
|
route = s.resolveRoute(params.ingress.ID, modelName)
|
|
providers, err = s.findProvidersForModel(modelName, allowedProviders, params.category, route)
|
|
return modelName, providers, route, err
|
|
}
|
|
|
|
func (s *Service) recordUsage(ingress *model.IngressKey, att *attempt, clientIP, modelName, endpoint string, stream bool, tokensIn, tokensOut int, latency int64, status, errMsg, requestBody, responseBody string) {
|
|
var providerKeyID, providerID uint
|
|
if att != nil {
|
|
providerKeyID = att.key.ID
|
|
providerID = att.provider.ID
|
|
}
|
|
cost := pricing.Estimate(modelName, tokensIn, tokensOut, endpoint)
|
|
s.cache.RecordUsage(cache.UsageDelta{
|
|
IngressKeyID: ingress.ID,
|
|
ProviderKeyID: providerKeyID,
|
|
ProviderID: providerID,
|
|
ClientIP: clientIP,
|
|
Endpoint: endpoint,
|
|
Model: modelName,
|
|
TokensIn: tokensIn,
|
|
TokensOut: tokensOut,
|
|
Cost: cost,
|
|
LatencyMs: latency,
|
|
Status: status,
|
|
Stream: stream,
|
|
ErrorMsg: errMsg,
|
|
RequestBody: requestBody,
|
|
ResponseBody: responseBody,
|
|
BudgetUsed: cost,
|
|
Requests: 1,
|
|
Tokens: int64(tokensIn + tokensOut),
|
|
})
|
|
}
|
|
|
|
func (s *Service) ForwardEndpoint(ctx context.Context, ingress *model.IngressKey, endpoint provider.EndpointType, rawBody []byte, clientIP string) (*provider.ChatResponse, error) {
|
|
category := categoryForEndpoint(endpoint)
|
|
modelName, providers, route, err := s.prepareForward(forwardParams{
|
|
ingress: ingress, endpoint: endpoint, category: category, body: rawBody,
|
|
})
|
|
if err != nil {
|
|
s.recordUsage(ingress, nil, clientIP, modelNameOrBody(rawBody, modelName), string(endpoint), false, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
|
|
return nil, err
|
|
}
|
|
attempts, err := s.buildAttempts(providers, modelName, route)
|
|
if err != nil {
|
|
s.recordUsage(ingress, nil, clientIP, modelName, string(endpoint), false, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
|
|
return nil, err
|
|
}
|
|
start := time.Now()
|
|
resp, att, err := s.forwardWithFailover(ctx, attempts, endpoint, rawBody)
|
|
latency := time.Since(start).Milliseconds()
|
|
status := "success"
|
|
if err != nil || (resp != nil && resp.StatusCode >= 400) {
|
|
status = "error"
|
|
}
|
|
errMsg := describeForwardError(err, resp)
|
|
tokensIn, tokensOut := 0, 0
|
|
var respBody []byte
|
|
if resp != nil {
|
|
tokensIn, tokensOut = resp.TokensIn, resp.TokensOut
|
|
respBody = resp.Body
|
|
}
|
|
s.recordUsage(ingress, att, clientIP, modelName, string(endpoint), false, tokensIn, tokensOut, latency, status, errMsg, truncateLogBody(rawBody), truncateLogBody(respBody))
|
|
if err != nil && resp != nil && resp.StatusCode >= 400 {
|
|
return resp, nil
|
|
}
|
|
return resp, err
|
|
}
|
|
|
|
func modelNameOrBody(body []byte, modelName string) string {
|
|
if modelName != "" {
|
|
return modelName
|
|
}
|
|
return extractModelName(body)
|
|
}
|
|
|
|
func (s *Service) ForwardStream(ctx context.Context, ingress *model.IngressKey, endpoint provider.EndpointType, rawBody []byte, clientIP string, w io.Writer, flusher http.Flusher) error {
|
|
category := categoryForEndpoint(endpoint)
|
|
modelName, providers, route, err := s.prepareForward(forwardParams{
|
|
ingress: ingress, endpoint: endpoint, category: category, body: rawBody, stream: true,
|
|
})
|
|
if err != nil {
|
|
s.recordUsage(ingress, nil, clientIP, modelNameOrBody(rawBody, modelName), string(endpoint), true, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
|
|
return err
|
|
}
|
|
attempts, err := s.buildAttempts(providers, modelName, route)
|
|
if err != nil {
|
|
s.recordUsage(ingress, nil, clientIP, modelName, string(endpoint), true, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
|
|
return err
|
|
}
|
|
var lastErr error
|
|
var lastAtt *attempt
|
|
maxRetries := s.settings.GatewayMaxRetries()
|
|
if maxRetries < 1 {
|
|
maxRetries = 1
|
|
}
|
|
backoffMs := s.settings.GatewayRetryBackoffMs()
|
|
for _, att := range attempts {
|
|
if !s.canServeEndpoint(att.provider, endpoint) {
|
|
continue
|
|
}
|
|
for try := 0; try < maxRetries; try++ {
|
|
if try > 0 {
|
|
time.Sleep(time.Duration(backoffMs) * time.Millisecond)
|
|
}
|
|
start := time.Now()
|
|
result, err := s.forwardStream(ctx, att.provider, att.apiKey, endpoint, rawBody, w, flusher)
|
|
latency := time.Since(start).Milliseconds()
|
|
lastAtt = &att
|
|
if err != nil {
|
|
lastErr = err
|
|
continue
|
|
}
|
|
if result == nil {
|
|
lastErr = fmt.Errorf("empty upstream stream response")
|
|
continue
|
|
}
|
|
if result.StatusCode < 400 {
|
|
s.recordUsage(ingress, &att, clientIP, modelName, string(endpoint), true, result.TokensIn, result.TokensOut, latency, "success", "", truncateLogBody(rawBody), "[SSE stream]")
|
|
return nil
|
|
}
|
|
errMsg := describeForwardError(nil, &provider.ChatResponse{StatusCode: result.StatusCode, Body: result.Body})
|
|
lastErr = fmt.Errorf("%s", errMsg)
|
|
respBody := "[SSE stream]"
|
|
if len(result.Body) > 0 {
|
|
respBody = truncateLogBody(result.Body)
|
|
}
|
|
if passThroughUpstreamStatus(result.StatusCode) {
|
|
s.recordUsage(ingress, &att, clientIP, modelName, string(endpoint), true, result.TokensIn, result.TokensOut, latency, "error", errMsg, truncateLogBody(rawBody), respBody)
|
|
return nil
|
|
}
|
|
if tryNextKeyOnUpstreamStatus(result.StatusCode) {
|
|
break
|
|
}
|
|
if retryableUpstreamStatus(result.StatusCode) {
|
|
continue
|
|
}
|
|
s.recordUsage(ingress, &att, clientIP, modelName, string(endpoint), true, result.TokensIn, result.TokensOut, latency, "error", errMsg, truncateLogBody(rawBody), respBody)
|
|
return nil
|
|
}
|
|
}
|
|
if lastErr == nil {
|
|
lastErr = fmt.Errorf("all providers failed")
|
|
}
|
|
s.recordUsage(ingress, lastAtt, clientIP, modelName, string(endpoint), true, 0, 0, 0, "error", lastErr.Error(), truncateLogBody(rawBody), "[SSE stream]")
|
|
return lastErr
|
|
}
|
|
|
|
func categoryForEndpoint(ep provider.EndpointType) model.ProviderCategory {
|
|
switch ep {
|
|
case provider.EndpointEmbeddings:
|
|
return model.CategoryEmbedding
|
|
case provider.EndpointRerank:
|
|
return model.CategoryRerank
|
|
default:
|
|
return model.CategoryLLM
|
|
}
|
|
}
|
|
|
|
func listModelsCategories() []model.ProviderCategory {
|
|
return []model.ProviderCategory{
|
|
model.CategoryLLM,
|
|
model.CategoryEmbedding,
|
|
model.CategoryRerank,
|
|
}
|
|
}
|
|
|
|
func providerInListModelsCategory(cat model.ProviderCategory) bool {
|
|
for _, c := range listModelsCategories() {
|
|
if c == cat {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (s *Service) ListModels(ctx context.Context, ingress *model.IngressKey) ([]map[string]any, error) {
|
|
allowedProviders := cache.ParseJSONList(ingress.AllowedProvidersJSON)
|
|
var result []map[string]any
|
|
var hasLLM bool
|
|
for _, p := range s.cache.ListProviders() {
|
|
if !providerInListModelsCategory(p.Category) {
|
|
continue
|
|
}
|
|
if !cache.MatchesList(allowedProviders, p.Name) && !cache.MatchesList(allowedProviders, string(p.Type)) {
|
|
continue
|
|
}
|
|
if !s.providerAllowsEndpoint(&p, provider.EndpointListModels) {
|
|
continue
|
|
}
|
|
if p.Category == model.CategoryLLM {
|
|
hasLLM = true
|
|
}
|
|
for _, m := range s.cache.GetModels(p.ID) {
|
|
if m.Enabled {
|
|
result = append(result, map[string]any{
|
|
"id": IngressModelID(m), "object": "model", "owned_by": p.Name,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
if len(result) == 0 && hasLLM {
|
|
result = append(result, map[string]any{"id": "gpt-4o-mini", "object": "model", "owned_by": "luminary"})
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (s *Service) FetchModels(ctx context.Context, providerID uint) ([]provider.ModelInfo, error) {
|
|
p, ok := s.cache.GetProvider(providerID)
|
|
if !ok {
|
|
return nil, fmt.Errorf("provider not found")
|
|
}
|
|
if !s.providerAllowsEndpoint(&p, provider.EndpointListModels) {
|
|
return nil, fmt.Errorf("list_models endpoint not enabled")
|
|
}
|
|
keys := s.cache.ListProviderKeys(providerID)
|
|
var firstKey string
|
|
for _, k := range keys {
|
|
if k.Enabled {
|
|
plain, err := s.DecryptAPIKey(k.APIKeyEnc)
|
|
if err == nil {
|
|
firstKey = plain
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if firstKey == "" && p.Type != model.ProviderOllama {
|
|
return nil, fmt.Errorf("no enabled keys")
|
|
}
|
|
cfg := s.providerConfig(&p)
|
|
if p.Type == model.ProviderAnthropic {
|
|
return s.anthropic.ListModelsWithConfig(ctx, firstKey, cfg)
|
|
}
|
|
return s.openai.ListModelsWithConfig(ctx, firstKey, cfg)
|
|
}
|
|
|
|
func (s *Service) HealthCheckProvider(ctx context.Context, providerID uint) (model.ProviderStatus, int64, string) {
|
|
p, ok := s.cache.GetProvider(providerID)
|
|
if !ok {
|
|
return model.ProviderUnhealthy, 0, "not found"
|
|
}
|
|
if !model.IsCategorySupported(p.Category) {
|
|
return model.ProviderUnhealthy, 0, "category not supported"
|
|
}
|
|
if p.Type == model.ProviderComfyUI {
|
|
start := time.Now()
|
|
if err := comfyui.New(p.BaseURL).HealthCheck(ctx); err != nil {
|
|
return model.ProviderUnhealthy, time.Since(start).Milliseconds(), err.Error()
|
|
}
|
|
return model.ProviderHealthy, time.Since(start).Milliseconds(), ""
|
|
}
|
|
keys := s.cache.ListProviderKeys(providerID)
|
|
if len(keys) == 0 && p.Type != model.ProviderOllama {
|
|
return model.ProviderUnhealthy, 0, "no provider keys configured"
|
|
}
|
|
cfg := s.providerConfig(&p)
|
|
return s.runHealthChecks(ctx, &p, keys, cfg)
|
|
}
|