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>
This commit is contained in:
renjue
@@ -0,0 +1,475 @@
|
||||
package prediction
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"path"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/rose_cat707/luminary/internal/model"
|
||||
"github.com/rose_cat707/luminary/internal/provider/comfyui"
|
||||
"github.com/rose_cat707/luminary/internal/storage/imagestore"
|
||||
"github.com/rose_cat707/luminary/internal/store/cache"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
db *gorm.DB
|
||||
cache *cache.Store
|
||||
images *imagestore.Store
|
||||
hub *Hub
|
||||
publicURL string
|
||||
waitSec int64
|
||||
mu sync.Mutex
|
||||
running map[string]context.CancelFunc
|
||||
}
|
||||
|
||||
func NewService(db *gorm.DB, c *cache.Store, images *imagestore.Store, publicURL string, waitSec int64) *Service {
|
||||
if waitSec <= 0 {
|
||||
waitSec = 60
|
||||
}
|
||||
return &Service{
|
||||
db: db, cache: c, images: images, hub: NewHub(),
|
||||
publicURL: strings.TrimRight(publicURL, "/"), waitSec: waitSec,
|
||||
running: make(map[string]context.CancelFunc),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Hub() *Hub { return s.hub }
|
||||
|
||||
func NewPredictionID() (string, error) {
|
||||
b := make([]byte, 12)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
type CreateRequest struct {
|
||||
Version string `json:"version"`
|
||||
Input map[string]any `json:"input"`
|
||||
}
|
||||
|
||||
func (s *Service) Create(ctx context.Context, ingress *model.IngressKey, req CreateRequest, clientIP string, wait bool) (*model.Prediction, error) {
|
||||
if req.Version == "" {
|
||||
return nil, fmt.Errorf("version is required")
|
||||
}
|
||||
entry, provider, err := s.resolveModel(req.Version, ingress)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !provider.Enabled {
|
||||
return nil, fmt.Errorf("provider disabled")
|
||||
}
|
||||
input, err := ValidateInput(entry.WorkflowType, req.Input)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
binding, err := ParseBinding(entry.InputBindingJSON)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var workflow map[string]any
|
||||
_ = json.Unmarshal([]byte(entry.WorkflowJSON), &workflow)
|
||||
if err := binding.Validate(workflow, RequiredBindings(entry.WorkflowType)); err != nil {
|
||||
return nil, fmt.Errorf("model binding invalid: %w", err)
|
||||
}
|
||||
|
||||
id, err := NewPredictionID()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inputJSON, _ := json.Marshal(input)
|
||||
p := model.Prediction{
|
||||
ID: id,
|
||||
IngressKeyID: ingress.ID,
|
||||
ProviderID: provider.ID,
|
||||
ModelEntryID: entry.ID,
|
||||
Version: req.Version,
|
||||
Model: entry.ModelID,
|
||||
InputJSON: string(inputJSON),
|
||||
Status: model.PredictionStarting,
|
||||
ClientID: "luminary-" + id,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := s.db.Create(&p).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
runCtx, cancel := context.WithCancel(context.Background())
|
||||
s.mu.Lock()
|
||||
s.running[id] = cancel
|
||||
s.mu.Unlock()
|
||||
|
||||
go s.execute(runCtx, &p, entry, &provider, binding, input, clientIP)
|
||||
|
||||
if wait {
|
||||
deadline := time.After(time.Duration(s.waitSec) * time.Second)
|
||||
for {
|
||||
select {
|
||||
case <-deadline:
|
||||
return s.Get(ctx, id, ingress)
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
cur, err := s.Get(ctx, id, ingress)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cur.Status == model.PredictionSucceeded || cur.Status == model.PredictionFailed || cur.Status == model.PredictionCanceled {
|
||||
return cur, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return s.toAPI(&p), nil
|
||||
}
|
||||
|
||||
func (s *Service) resolveModel(version string, ingress *model.IngressKey) (model.ModelEntry, model.Provider, error) {
|
||||
allowedProviders := cache.ParseJSONList(ingress.AllowedProvidersJSON)
|
||||
allowedModels := cache.ParseJSONList(ingress.AllowedModelsJSON)
|
||||
var candidates []model.ModelEntry
|
||||
for _, p := range s.cache.ListProviders() {
|
||||
if p.Category != model.CategoryImage || !p.Enabled {
|
||||
continue
|
||||
}
|
||||
if !cache.MatchesList(allowedProviders, p.Name) && !cache.MatchesList(allowedProviders, string(p.Type)) && !cache.MatchesList(allowedProviders, "*") {
|
||||
continue
|
||||
}
|
||||
for _, m := range s.cache.GetModels(p.ID) {
|
||||
if !m.Enabled {
|
||||
continue
|
||||
}
|
||||
ingressID := m.IngressID()
|
||||
if !m.MatchesIngressName(version) {
|
||||
continue
|
||||
}
|
||||
if !cache.MatchesList(allowedModels, ingressID) && !cache.MatchesList(allowedModels, m.ModelID) && !cache.MatchesList(allowedModels, "*") {
|
||||
continue
|
||||
}
|
||||
candidates = append(candidates, m)
|
||||
}
|
||||
}
|
||||
if len(candidates) == 0 {
|
||||
return model.ModelEntry{}, model.Provider{}, fmt.Errorf("no model found for version %s", version)
|
||||
}
|
||||
m := candidates[0]
|
||||
p, ok := s.cache.GetProvider(m.ProviderID)
|
||||
if !ok {
|
||||
return model.ModelEntry{}, model.Provider{}, fmt.Errorf("provider not found")
|
||||
}
|
||||
return m, p, nil
|
||||
}
|
||||
|
||||
func (s *Service) execute(ctx context.Context, p *model.Prediction, entry model.ModelEntry, provider *model.Provider, binding InputBinding, input map[string]any, clientIP string) {
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
delete(s.running, p.ID)
|
||||
s.mu.Unlock()
|
||||
}()
|
||||
|
||||
start := time.Now()
|
||||
s.updateStatus(p, model.PredictionProcessing, "", nil)
|
||||
s.hub.Publish(p.ID, Event{Status: string(model.PredictionProcessing)})
|
||||
|
||||
client := comfyui.New(provider.BaseURL)
|
||||
runInput := cloneMap(input)
|
||||
if entry.WorkflowType == model.WorkflowImg2Img {
|
||||
imageURL, _ := runInput["image"].(string)
|
||||
data, err := downloadURL(ctx, imageURL)
|
||||
if err != nil {
|
||||
s.fail(p, clientIP, start, fmt.Errorf("download image: %w", err))
|
||||
return
|
||||
}
|
||||
name := path.Base(imageURL)
|
||||
if name == "" || name == "." || name == "/" {
|
||||
name = "input.png"
|
||||
}
|
||||
uploaded, err := client.UploadImage(ctx, data, name)
|
||||
if err != nil {
|
||||
s.fail(p, clientIP, start, fmt.Errorf("upload image: %w", err))
|
||||
return
|
||||
}
|
||||
runInput["image"] = uploaded
|
||||
}
|
||||
|
||||
prompt, err := BuildPrompt(entry.WorkflowJSON, binding, runInput)
|
||||
if err != nil {
|
||||
s.fail(p, clientIP, start, err)
|
||||
return
|
||||
}
|
||||
|
||||
wsCtx, wsCancel := context.WithCancel(ctx)
|
||||
defer wsCancel()
|
||||
go client.Watch(wsCtx, p.ClientID, func(ev comfyui.ProgressEvent) {
|
||||
switch ev.Type {
|
||||
case "progress":
|
||||
if ev.Max > 0 {
|
||||
logs := fmt.Sprintf("progress %.0f/%.0f", ev.Value, ev.Max)
|
||||
s.appendLogs(p, logs)
|
||||
s.hub.Publish(p.ID, Event{Status: string(model.PredictionProcessing), Logs: logs})
|
||||
}
|
||||
case "execution_error":
|
||||
s.appendLogs(p, ev.ErrorText)
|
||||
}
|
||||
})
|
||||
|
||||
submit, err := client.SubmitPrompt(ctx, prompt, p.ClientID)
|
||||
if err != nil {
|
||||
s.fail(p, clientIP, start, err)
|
||||
return
|
||||
}
|
||||
p.ComfyPromptID = submit.PromptID
|
||||
now := time.Now()
|
||||
p.StartedAt = &now
|
||||
s.db.Model(p).Updates(map[string]any{"comfy_prompt_id": p.ComfyPromptID, "started_at": p.StartedAt})
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_ = client.Interrupt(context.Background())
|
||||
s.updateStatus(p, model.PredictionCanceled, "canceled", nil)
|
||||
s.hub.Publish(p.ID, Event{Status: string(model.PredictionCanceled), Done: true})
|
||||
s.recordUsage(ingressByID(s.cache, p.IngressKeyID), provider, clientIP, p, start, "canceled", "canceled")
|
||||
return
|
||||
case <-time.After(2 * time.Second):
|
||||
history, err := client.GetHistory(ctx, p.ComfyPromptID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if _, ok := history[p.ComfyPromptID]; !ok {
|
||||
continue
|
||||
}
|
||||
images := comfyui.CollectOutputImages(history, p.ComfyPromptID)
|
||||
if len(images) == 0 {
|
||||
continue
|
||||
}
|
||||
var urls []string
|
||||
for _, img := range images {
|
||||
data, mime, err := client.DownloadView(ctx, img.Filename, img.Subfolder, img.Type)
|
||||
if err != nil {
|
||||
s.fail(p, clientIP, start, err)
|
||||
return
|
||||
}
|
||||
stored, err := s.images.Save(p.ID, img.Filename, data, mime)
|
||||
if err != nil {
|
||||
s.fail(p, clientIP, start, err)
|
||||
return
|
||||
}
|
||||
urls = append(urls, s.images.SignURL(stored.ID))
|
||||
}
|
||||
outJSON, _ := json.Marshal(urls)
|
||||
p.OutputJSON = string(outJSON)
|
||||
completed := time.Now()
|
||||
p.CompletedAt = &completed
|
||||
metrics, _ := json.Marshal(map[string]any{
|
||||
"predict_time": time.Since(start).Seconds(),
|
||||
"total_time": completed.Sub(p.CreatedAt).Seconds(),
|
||||
})
|
||||
p.MetricsJSON = string(metrics)
|
||||
s.updateStatus(p, model.PredictionSucceeded, "", &completed)
|
||||
s.hub.Publish(p.ID, Event{
|
||||
Status: string(model.PredictionSucceeded), Output: urls,
|
||||
Metrics: map[string]any{"predict_time": time.Since(start).Seconds()},
|
||||
Done: true,
|
||||
})
|
||||
s.recordUsage(ingressByID(s.cache, p.IngressKeyID), provider, clientIP, p, start, "success", "")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) fail(p *model.Prediction, clientIP string, start time.Time, err error) {
|
||||
completed := time.Now()
|
||||
p.CompletedAt = &completed
|
||||
p.Error = err.Error()
|
||||
s.updateStatus(p, model.PredictionFailed, err.Error(), &completed)
|
||||
s.hub.Publish(p.ID, Event{Status: string(model.PredictionFailed), Error: err.Error(), Done: true})
|
||||
prov, _ := s.cache.GetProvider(p.ProviderID)
|
||||
s.recordUsage(ingressByID(s.cache, p.IngressKeyID), &prov, clientIP, p, start, "error", err.Error())
|
||||
}
|
||||
|
||||
func (s *Service) updateStatus(p *model.Prediction, status model.PredictionStatus, errMsg string, completed *time.Time) {
|
||||
p.Status = status
|
||||
if errMsg != "" {
|
||||
p.Error = errMsg
|
||||
}
|
||||
updates := map[string]any{"status": status, "error": p.Error, "logs": p.Logs, "output_json": p.OutputJSON, "metrics_json": p.MetricsJSON}
|
||||
if completed != nil {
|
||||
updates["completed_at"] = completed
|
||||
}
|
||||
s.db.Model(p).Where("id = ?", p.ID).Updates(updates)
|
||||
}
|
||||
|
||||
func (s *Service) appendLogs(p *model.Prediction, line string) {
|
||||
if p.Logs != "" {
|
||||
p.Logs += "\n"
|
||||
}
|
||||
p.Logs += line
|
||||
s.db.Model(p).Where("id = ?", p.ID).Update("logs", p.Logs)
|
||||
}
|
||||
|
||||
func (s *Service) Get(ctx context.Context, id string, ingress *model.IngressKey) (*model.Prediction, error) {
|
||||
var p model.Prediction
|
||||
if err := s.db.First(&p, "id = ?", id).Error; err != nil {
|
||||
return nil, fmt.Errorf("prediction not found")
|
||||
}
|
||||
if ingress != nil && p.IngressKeyID != ingress.ID {
|
||||
return nil, fmt.Errorf("prediction not found")
|
||||
}
|
||||
return s.toAPI(&p), nil
|
||||
}
|
||||
|
||||
func (s *Service) List(ctx context.Context, ingress *model.IngressKey, cursor string, limit int) ([]*model.Prediction, error) {
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 100
|
||||
}
|
||||
q := s.db.Where("ingress_key_id = ?", ingress.ID).Order("created_at desc").Limit(limit)
|
||||
if cursor != "" {
|
||||
q = q.Where("id < ?", cursor)
|
||||
}
|
||||
var rows []model.Prediction
|
||||
if err := q.Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]*model.Prediction, 0, len(rows))
|
||||
for i := range rows {
|
||||
out = append(out, s.toAPI(&rows[i]))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Service) Cancel(ctx context.Context, id string, ingress *model.IngressKey) (*model.Prediction, error) {
|
||||
p, err := s.Get(ctx, id, ingress)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if p.Status == model.PredictionSucceeded || p.Status == model.PredictionFailed || p.Status == model.PredictionCanceled {
|
||||
return p, nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
if cancel, ok := s.running[id]; ok {
|
||||
cancel()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
return s.Get(ctx, id, ingress)
|
||||
}
|
||||
|
||||
func (s *Service) toAPI(p *model.Prediction) *model.Prediction {
|
||||
return p
|
||||
}
|
||||
|
||||
func (s *Service) APIView(p *model.Prediction, baseURL string) map[string]any {
|
||||
var input any
|
||||
_ = json.Unmarshal([]byte(p.InputJSON), &input)
|
||||
var output any
|
||||
if p.OutputJSON != "" {
|
||||
_ = json.Unmarshal([]byte(p.OutputJSON), &output)
|
||||
}
|
||||
var metrics any
|
||||
if p.MetricsJSON != "" {
|
||||
_ = json.Unmarshal([]byte(p.MetricsJSON), &metrics)
|
||||
}
|
||||
base := strings.TrimRight(baseURL, "/")
|
||||
return map[string]any{
|
||||
"id": p.ID,
|
||||
"version": p.Version,
|
||||
"input": input,
|
||||
"output": output,
|
||||
"error": nilIfEmpty(p.Error),
|
||||
"status": p.Status,
|
||||
"created_at": p.CreatedAt.UTC().Format(time.RFC3339Nano),
|
||||
"started_at": formatTime(p.StartedAt),
|
||||
"completed_at": formatTime(p.CompletedAt),
|
||||
"logs": p.Logs,
|
||||
"metrics": metrics,
|
||||
"urls": map[string]string{
|
||||
"get": fmt.Sprintf("%s/v1/predictions/%s", base, p.ID),
|
||||
"cancel": fmt.Sprintf("%s/v1/predictions/%s/cancel", base, p.ID),
|
||||
"stream": fmt.Sprintf("%s/v1/predictions/%s/stream", base, p.ID),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func nilIfEmpty(s string) any {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func formatTime(t *time.Time) any {
|
||||
if t == nil {
|
||||
return nil
|
||||
}
|
||||
return t.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
|
||||
func downloadURL(ctx context.Context, raw string) ([]byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, raw, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("http %d", resp.StatusCode)
|
||||
}
|
||||
return io.ReadAll(resp.Body)
|
||||
}
|
||||
|
||||
func cloneMap(in map[string]any) map[string]any {
|
||||
b, _ := json.Marshal(in)
|
||||
var out map[string]any
|
||||
_ = json.Unmarshal(b, &out)
|
||||
return out
|
||||
}
|
||||
|
||||
func ingressByID(c *cache.Store, id uint) *model.IngressKey {
|
||||
for _, k := range c.ListIngressKeys() {
|
||||
if k.ID == id {
|
||||
kCopy := k
|
||||
return &kCopy
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) recordUsage(ingress *model.IngressKey, provider *model.Provider, clientIP string, p *model.Prediction, start time.Time, status, errMsg string) {
|
||||
if ingress == nil {
|
||||
return
|
||||
}
|
||||
var providerID uint
|
||||
if provider != nil {
|
||||
providerID = provider.ID
|
||||
}
|
||||
s.cache.RecordUsage(cache.UsageDelta{
|
||||
IngressKeyID: ingress.ID,
|
||||
ProviderID: providerID,
|
||||
ClientIP: clientIP,
|
||||
Endpoint: model.EndpointPrediction,
|
||||
Model: p.ID,
|
||||
LatencyMs: time.Since(start).Milliseconds(),
|
||||
Status: status,
|
||||
ErrorMsg: errMsg,
|
||||
RequestBody: p.InputJSON,
|
||||
ResponseBody: p.OutputJSON,
|
||||
Requests: 1,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) GetDBPrediction(id string) (*model.Prediction, error) {
|
||||
var p model.Prediction
|
||||
if err := s.db.First(&p, "id = ?", id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &p, nil
|
||||
}
|
||||
Reference in New Issue
Block a user