renjue
e761074fb9
CI / docker (push) Successful in 3m46s
Merge model_id/alias in the image provider UI, and ensure ComfyUI never receives seed -1 after cloneMap type coercion. Co-authored-by: Cursor <cursoragent@cursor.com>
425 lines
13 KiB
Go
425 lines
13 KiB
Go
package admin
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"slices"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/rose_cat707/luminary/internal/auth"
|
|
"github.com/rose_cat707/luminary/internal/model"
|
|
"github.com/rose_cat707/luminary/internal/prediction"
|
|
"github.com/rose_cat707/luminary/internal/provider"
|
|
"github.com/rose_cat707/luminary/internal/store/cache"
|
|
)
|
|
|
|
func (h *Handler) DeleteProvider(c *gin.Context) {
|
|
id := parseUint(c.Param("id"))
|
|
h.db.DB().Where("provider_id = ?", id).Delete(&model.ProviderKey{})
|
|
h.db.DB().Where("provider_id = ?", id).Delete(&model.ModelEntry{})
|
|
h.db.DB().Delete(&model.Provider{}, id)
|
|
h.cache.DeleteProvider(id)
|
|
c.JSON(http.StatusOK, gin.H{"ok": true})
|
|
}
|
|
|
|
func (h *Handler) SyncModels(c *gin.Context) {
|
|
id := parseUint(c.Param("id"))
|
|
p, ok := h.cache.GetProvider(id)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "provider not found"})
|
|
return
|
|
}
|
|
allowed := provider.ParseAllowedEndpoints(p.AllowedEndpointsJSON, p.Category)
|
|
if !slices.Contains(allowed, string(provider.EndpointListModels)) {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "list_models endpoint not enabled"})
|
|
return
|
|
}
|
|
models, err := h.gateway.FetchModels(c.Request.Context(), id)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
existing := h.cache.GetModels(id)
|
|
existingByModelID := make(map[string]model.ModelEntry, len(existing))
|
|
for _, e := range existing {
|
|
existingByModelID[e.ModelID] = e
|
|
}
|
|
upstreamIDs := make(map[string]bool, len(models))
|
|
var entries []model.ModelEntry
|
|
for _, m := range models {
|
|
upstreamIDs[m.ID] = true
|
|
if old, ok := existingByModelID[m.ID]; ok {
|
|
old.DisplayName = m.DisplayName
|
|
if old.DisplayName == "" {
|
|
old.DisplayName = m.ID
|
|
}
|
|
if err := h.db.DB().Save(&old).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
entries = append(entries, old)
|
|
continue
|
|
}
|
|
displayName := m.DisplayName
|
|
if displayName == "" {
|
|
displayName = m.ID
|
|
}
|
|
entry := model.ModelEntry{
|
|
ProviderID: id,
|
|
ModelID: m.ID,
|
|
DisplayName: displayName,
|
|
Enabled: true,
|
|
}
|
|
if err := h.db.DB().Create(&entry).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
entries = append(entries, entry)
|
|
}
|
|
for _, e := range existing {
|
|
if !upstreamIDs[e.ModelID] {
|
|
entries = append(entries, e)
|
|
}
|
|
}
|
|
h.cache.SetModels(id, entries)
|
|
c.JSON(http.StatusOK, gin.H{"synced": len(models), "models": entries})
|
|
}
|
|
|
|
func (h *Handler) ListProviderKeys(c *gin.Context) {
|
|
id := parseUint(c.Param("id"))
|
|
c.JSON(http.StatusOK, gin.H{"data": h.cache.ListProviderKeys(id)})
|
|
}
|
|
|
|
func (h *Handler) CreateProviderKey(c *gin.Context) {
|
|
providerID := parseUint(c.Param("id"))
|
|
var req struct {
|
|
Name string `json:"name"`
|
|
APIKey string `json:"api_key"`
|
|
Weight float64 `json:"weight"`
|
|
Models []string `json:"models"`
|
|
Enabled *bool `json:"enabled"`
|
|
MaxRequestsPerDay int `json:"max_requests_per_day"`
|
|
MaxTokensPerDay int64 `json:"max_tokens_per_day"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil || req.Name == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
|
|
return
|
|
}
|
|
if req.APIKey == "" {
|
|
var p model.Provider
|
|
h.db.DB().First(&p, providerID)
|
|
if p.Type != model.ProviderOllama {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "api_key required"})
|
|
return
|
|
}
|
|
req.APIKey = "ollama"
|
|
}
|
|
enc, err := auth.Encrypt(req.APIKey, h.settings.EncryptionKey())
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "encrypt failed"})
|
|
return
|
|
}
|
|
modelsJSON := cache.ModelsToJSON(req.Models)
|
|
if req.Models == nil {
|
|
modelsJSON = `["*"]`
|
|
}
|
|
k := model.ProviderKey{
|
|
ProviderID: providerID,
|
|
Name: req.Name,
|
|
APIKeyEnc: enc,
|
|
Weight: req.Weight,
|
|
ModelsJSON: modelsJSON,
|
|
Enabled: true,
|
|
MaxRequestsPerDay: req.MaxRequestsPerDay,
|
|
MaxTokensPerDay: req.MaxTokensPerDay,
|
|
}
|
|
if req.Weight <= 0 {
|
|
k.Weight = 1
|
|
}
|
|
if req.Enabled != nil {
|
|
k.Enabled = *req.Enabled
|
|
}
|
|
if err := h.db.DB().Create(&k).Error; err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
h.cache.SetProviderKey(k)
|
|
c.JSON(http.StatusCreated, k)
|
|
}
|
|
|
|
func (h *Handler) UpdateProviderKey(c *gin.Context) {
|
|
keyID := parseUint(c.Param("keyId"))
|
|
k, ok := h.cache.GetProviderKey(keyID)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
|
|
return
|
|
}
|
|
var req struct {
|
|
Name string `json:"name"`
|
|
APIKey string `json:"api_key"`
|
|
Weight float64 `json:"weight"`
|
|
Models []string `json:"models"`
|
|
Enabled *bool `json:"enabled"`
|
|
MaxRequestsPerDay int `json:"max_requests_per_day"`
|
|
MaxTokensPerDay int64 `json:"max_tokens_per_day"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
|
|
return
|
|
}
|
|
if req.Name != "" {
|
|
k.Name = req.Name
|
|
}
|
|
if req.APIKey != "" {
|
|
enc, err := auth.Encrypt(req.APIKey, h.settings.EncryptionKey())
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "encrypt failed"})
|
|
return
|
|
}
|
|
k.APIKeyEnc = enc
|
|
}
|
|
if req.Weight > 0 {
|
|
k.Weight = req.Weight
|
|
}
|
|
if req.Models != nil {
|
|
k.ModelsJSON = cache.ModelsToJSON(req.Models)
|
|
}
|
|
if req.Enabled != nil {
|
|
k.Enabled = *req.Enabled
|
|
}
|
|
if req.MaxRequestsPerDay >= 0 {
|
|
k.MaxRequestsPerDay = req.MaxRequestsPerDay
|
|
}
|
|
if req.MaxTokensPerDay >= 0 {
|
|
k.MaxTokensPerDay = req.MaxTokensPerDay
|
|
}
|
|
h.db.DB().Save(&k)
|
|
h.cache.SetProviderKey(k)
|
|
c.JSON(http.StatusOK, k)
|
|
}
|
|
|
|
func (h *Handler) DeleteProviderKey(c *gin.Context) {
|
|
keyID := parseUint(c.Param("keyId"))
|
|
h.db.DB().Delete(&model.ProviderKey{}, keyID)
|
|
h.cache.DeleteProviderKey(keyID)
|
|
c.JSON(http.StatusOK, gin.H{"ok": true})
|
|
}
|
|
|
|
func (h *Handler) ListModels(c *gin.Context) {
|
|
id := parseUint(c.Param("id"))
|
|
if _, ok := h.cache.GetProvider(id); !ok {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "provider not found"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"data": h.cache.GetModels(id)})
|
|
}
|
|
|
|
func (h *Handler) CreateModel(c *gin.Context) {
|
|
providerID := parseUint(c.Param("id"))
|
|
prov, ok := h.cache.GetProvider(providerID)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "provider not found"})
|
|
return
|
|
}
|
|
var req struct {
|
|
ModelID string `json:"model_id"`
|
|
Alias string `json:"alias"`
|
|
DisplayName string `json:"display_name"`
|
|
WorkflowType model.WorkflowType `json:"workflow_type"`
|
|
WorkflowJSON string `json:"workflow_json"`
|
|
InputBindingJSON string `json:"input_binding_json"`
|
|
InputBinding map[string]*prediction.NodeField `json:"input_binding"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil || req.ModelID == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
|
|
return
|
|
}
|
|
for _, existing := range h.cache.GetModels(providerID) {
|
|
if existing.ModelID == req.ModelID {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "model already exists"})
|
|
return
|
|
}
|
|
}
|
|
if msg := validateModelAlias(h, req.Alias, 0); msg != "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": msg})
|
|
return
|
|
}
|
|
if prov.Category == model.CategoryImage {
|
|
if err := validateImageModelCreate(struct {
|
|
ModelID, DisplayName string
|
|
WorkflowType model.WorkflowType
|
|
WorkflowJSON, InputBindingJSON string
|
|
InputBinding map[string]*prediction.NodeField
|
|
}{req.ModelID, req.DisplayName, req.WorkflowType, req.WorkflowJSON, req.InputBindingJSON, req.InputBinding}); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
}
|
|
bindingJSON := req.InputBindingJSON
|
|
if bindingJSON == "" && req.InputBinding != nil {
|
|
b, _ := json.Marshal(req.InputBinding)
|
|
bindingJSON = string(b)
|
|
}
|
|
wt := req.WorkflowType
|
|
if wt == "" {
|
|
wt = model.WorkflowTxt2Img
|
|
}
|
|
m := model.ModelEntry{
|
|
ProviderID: providerID,
|
|
ModelID: req.ModelID,
|
|
Alias: normalizeAlias(req.Alias),
|
|
DisplayName: req.DisplayName,
|
|
WorkflowType: wt,
|
|
WorkflowJSON: req.WorkflowJSON,
|
|
InputBindingJSON: bindingJSON,
|
|
VersionHash: prediction.HashWorkflowJSON(req.WorkflowJSON),
|
|
Enabled: true,
|
|
}
|
|
if prov.Category == model.CategoryImage {
|
|
m.Alias = ""
|
|
}
|
|
if m.DisplayName == "" {
|
|
m.DisplayName = req.ModelID
|
|
}
|
|
if err := h.db.DB().Create(&m).Error; err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
models := h.cache.GetModels(providerID)
|
|
models = append(models, m)
|
|
h.cache.SetModels(providerID, models)
|
|
c.JSON(http.StatusCreated, m)
|
|
}
|
|
|
|
func (h *Handler) UpdateModel(c *gin.Context) {
|
|
providerID := parseUint(c.Param("id"))
|
|
modelID := parseUint(c.Param("modelId"))
|
|
var req struct {
|
|
Enabled *bool `json:"enabled"`
|
|
DisplayName string `json:"display_name"`
|
|
Alias *string `json:"alias"`
|
|
ModelID string `json:"model_id"`
|
|
WorkflowType model.WorkflowType `json:"workflow_type"`
|
|
WorkflowJSON string `json:"workflow_json"`
|
|
InputBindingJSON string `json:"input_binding_json"`
|
|
InputBinding map[string]*prediction.NodeField `json:"input_binding"`
|
|
}
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
|
|
return
|
|
}
|
|
m, ok := h.findModelEntry(providerID, modelID, req.ModelID)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "model not found"})
|
|
return
|
|
}
|
|
if req.Enabled != nil {
|
|
m.Enabled = *req.Enabled
|
|
}
|
|
if req.DisplayName != "" {
|
|
m.DisplayName = req.DisplayName
|
|
}
|
|
if req.Alias != nil {
|
|
normalized := normalizeAlias(*req.Alias)
|
|
if msg := validateModelAlias(h, normalized, m.ID); msg != "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": msg})
|
|
return
|
|
}
|
|
m.Alias = normalized
|
|
}
|
|
if req.ModelID != "" && req.ModelID != m.ModelID {
|
|
for _, existing := range h.cache.GetModels(m.ProviderID) {
|
|
if existing.ID != m.ID && existing.ModelID == req.ModelID {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "model already exists"})
|
|
return
|
|
}
|
|
}
|
|
m.ModelID = req.ModelID
|
|
}
|
|
if req.WorkflowType != "" {
|
|
m.WorkflowType = req.WorkflowType
|
|
}
|
|
if req.WorkflowJSON != "" {
|
|
m.WorkflowJSON = req.WorkflowJSON
|
|
m.VersionHash = prediction.HashWorkflowJSON(req.WorkflowJSON)
|
|
}
|
|
if req.InputBindingJSON != "" {
|
|
m.InputBindingJSON = req.InputBindingJSON
|
|
} else if req.InputBinding != nil {
|
|
b, _ := json.Marshal(req.InputBinding)
|
|
m.InputBindingJSON = string(b)
|
|
}
|
|
prov, _ := h.cache.GetProvider(m.ProviderID)
|
|
if prov.Category == model.CategoryImage {
|
|
m.Alias = ""
|
|
}
|
|
if prov.Category == model.CategoryImage && m.WorkflowJSON != "" {
|
|
if err := validateImageModelCreate(struct {
|
|
ModelID, DisplayName string
|
|
WorkflowType model.WorkflowType
|
|
WorkflowJSON, InputBindingJSON string
|
|
InputBinding map[string]*prediction.NodeField
|
|
}{m.ModelID, m.DisplayName, m.WorkflowType, m.WorkflowJSON, m.InputBindingJSON, nil}); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
}
|
|
if err := h.db.DB().Save(&m).Error; err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
models := h.cache.GetModels(m.ProviderID)
|
|
for i := range models {
|
|
if models[i].ID == m.ID {
|
|
models[i] = m
|
|
break
|
|
}
|
|
}
|
|
h.cache.SetModels(m.ProviderID, models)
|
|
c.JSON(http.StatusOK, m)
|
|
}
|
|
|
|
func (h *Handler) findModelEntry(providerID, entryID uint, modelID string) (model.ModelEntry, bool) {
|
|
var m model.ModelEntry
|
|
if entryID > 0 {
|
|
if err := h.db.DB().First(&m, entryID).Error; err == nil {
|
|
if providerID > 0 && m.ProviderID != providerID {
|
|
return model.ModelEntry{}, false
|
|
}
|
|
return m, true
|
|
}
|
|
}
|
|
if providerID > 0 && modelID != "" {
|
|
if err := h.db.DB().Where("provider_id = ? AND model_id = ?", providerID, modelID).First(&m).Error; err == nil {
|
|
return m, true
|
|
}
|
|
}
|
|
return model.ModelEntry{}, false
|
|
}
|
|
|
|
func (h *Handler) DeleteModel(c *gin.Context) {
|
|
modelID := parseUint(c.Param("modelId"))
|
|
var m model.ModelEntry
|
|
if err := h.db.DB().First(&m, modelID).Error; err != nil {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
|
|
return
|
|
}
|
|
if providerID := parseUint(c.Param("id")); providerID > 0 && m.ProviderID != providerID {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
|
|
return
|
|
}
|
|
if err := h.db.DB().Delete(&m).Error; err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
models := h.cache.GetModels(m.ProviderID)
|
|
filtered := models[:0]
|
|
for _, item := range models {
|
|
if item.ID != modelID {
|
|
filtered = append(filtered, item)
|
|
}
|
|
}
|
|
h.cache.SetModels(m.ProviderID, filtered)
|
|
c.JSON(http.StatusOK, gin.H{"ok": true})
|
|
}
|