model/vendor_management.go
Rama main · réplica de lectura del repositorio original.
package model
import (
"crypto/sha256"
"errors"
"fmt"
"slices"
"strings"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
"gorm.io/gorm"
)
var ErrVendorConflict = errors.New("vendor data changed; preview again before applying")
// VendorReferenceError reports every referenced vendor in a rejected delete.
type VendorReferenceError struct {
Counts map[int]int64 `json:"reference_counts"`
}
func (e *VendorReferenceError) Error() string {
return "vendors are still referenced by models; transfer or clear their assignments first"
}
func FindMetadataVendor(vendors map[string]*Vendor, name string) *Vendor {
for _, vendor := range vendors {
if strings.EqualFold(strings.TrimSpace(vendor.Name), strings.TrimSpace(name)) {
return vendor
}
}
return nil
}
func validateModelVendor(tx *gorm.DB, id int) error {
if id == 0 {
return nil
}
if id < 0 {
return errors.New("select a saved vendor")
}
var count int64
if err := tx.Model(&Vendor{}).Where("id = ?", id).Count(&count).Error; err != nil {
return err
}
if count == 0 {
return errors.New("vendor does not exist")
}
return nil
}
func validateVendorMetadata(tx *gorm.DB, vendor *Vendor) error {
vendor.Name = strings.TrimSpace(vendor.Name)
vendor.Icon = strings.TrimSpace(vendor.Icon)
if vendor.Name == "" {
return errors.New("vendor name is required")
}
if utf8.RuneCountInString(vendor.Name) > 128 || utf8.RuneCountInString(vendor.Icon) > 128 {
return errors.New("vendor name and icon must not exceed 128 characters")
}
var others []Vendor
if err := tx.Select("id", "name").Where("id <> ?", vendor.Id).Find(&others).Error; err != nil {
return err
}
for _, other := range others {
if strings.EqualFold(strings.TrimSpace(other.Name), vendor.Name) {
return errors.New("vendor name already exists")
}
}
return nil
}
func VendorRecordVersion(v *Vendor) string {
encoded, _ := common.Marshal([]any{v.Id, v.Name, v.Description, v.Icon, v.Status, v.CreatedTime, v.UpdatedTime})
return fmt.Sprintf("%x", sha256.Sum256(encoded))
}
// VendorOperation uses explicit selections for both preview and application.
type VendorOperation struct {
Action string `json:"action"`
VendorIDs []int `json:"vendor_ids"`
ModelIDs []int `json:"model_ids"`
TargetVendorID int `json:"target_vendor_id"`
ExpectedVersion string `json:"expected_version,omitempty"`
}
type VendorAssignmentModel struct {
ID int `json:"id"`
ModelName string `json:"model_name"`
NameRule int `json:"name_rule"`
VendorID int `json:"vendor_id"`
VendorName string `json:"vendor_name"`
UpdatedTime int64 `json:"updated_time"`
}
type VendorOperationPreview struct {
Action string `json:"action"`
Sources []Vendor `json:"sources"`
Target *Vendor `json:"target"`
Models []VendorAssignmentModel `json:"models"`
Version string `json:"version"`
}
type VendorOperationResult struct {
UpdatedModels []int `json:"updated_models"`
DeletedVendors []int `json:"deleted_vendors"`
}
func buildVendorOperationPreview(db *gorm.DB, operation VendorOperation) (*VendorOperationPreview, error) {
if operation.Action != "assign" && operation.Action != "merge" && operation.Action != "delete" {
return nil, errors.New("unsupported vendor operation")
}
ids := append([]int{}, operation.VendorIDs...)
if operation.Action == "assign" {
ids = append([]int{}, operation.ModelIDs...)
}
if len(ids) == 0 || len(ids) > 1000 {
return nil, errors.New("select between 1 and 1000 records")
}
slices.Sort(ids)
for i, id := range ids {
if id <= 0 || i > 0 && ids[i-1] == id {
return nil, errors.New("invalid or duplicate selection")
}
}
preview := &VendorOperationPreview{Action: operation.Action, Sources: []Vendor{}, Models: []VendorAssignmentModel{}}
var vendors []Vendor
if err := db.Session(&gorm.Session{}).Order("id").Find(&vendors).Error; err != nil {
return nil, err
}
byID := make(map[int]*Vendor, len(vendors))
for i := range vendors {
byID[vendors[i].Id] = &vendors[i]
}
if operation.Action != "delete" {
if operation.TargetVendorID < 0 || operation.Action == "merge" && operation.TargetVendorID == 0 {
return nil, errors.New("select a saved target vendor")
}
if operation.TargetVendorID != 0 {
preview.Target = byID[operation.TargetVendorID]
if preview.Target == nil {
return nil, fmt.Errorf("%w: target vendor does not exist", ErrVendorConflict)
}
}
}
var models []Model
query := db.Session(&gorm.Session{}).Model(&Model{}).Order("id")
if operation.Action == "assign" {
query = query.Where("id IN ?", ids)
} else {
for _, id := range ids {
vendor := byID[id]
if vendor == nil {
return nil, fmt.Errorf("%w: source vendor does not exist", ErrVendorConflict)
}
if operation.Action == "merge" && id == operation.TargetVendorID {
return nil, errors.New("target vendor cannot also be a source")
}
preview.Sources = append(preview.Sources, *vendor)
}
query = query.Where("vendor_id IN ?", ids)
}
if err := query.Find(&models).Error; err != nil {
return nil, err
}
if operation.Action == "assign" && len(models) != len(ids) {
return nil, fmt.Errorf("%w: a selected model no longer exists", ErrVendorConflict)
}
seen := make(map[int]bool)
modelVersions := make([]string, 0, len(models))
counts := make(map[int]int64)
for _, item := range models {
modelVersions = append(modelVersions, MetadataRecordVersion(&item, nil, nil))
row := VendorAssignmentModel{ID: item.Id, ModelName: item.ModelName, NameRule: item.NameRule, VendorID: item.VendorID, UpdatedTime: item.UpdatedTime}
if vendor := byID[item.VendorID]; vendor != nil {
row.VendorName = vendor.Name
if operation.Action == "assign" && !seen[vendor.Id] {
preview.Sources = append(preview.Sources, *vendor)
seen[vendor.Id] = true
}
}
preview.Models = append(preview.Models, row)
counts[item.VendorID]++
}
if operation.Action == "delete" && len(models) > 0 {
return nil, &VendorReferenceError{Counts: counts}
}
slices.SortFunc(preview.Sources, func(a, b Vendor) int { return a.Id - b.Id })
encoded, err := common.Marshal([]any{preview, modelVersions})
if err != nil {
return nil, err
}
preview.Version = fmt.Sprintf("%x", sha256.Sum256(encoded))
return preview, nil
}
func PreviewVendorOperation(operation VendorOperation) (*VendorOperationPreview, error) {
return buildVendorOperationPreview(DB, operation)
}
func ApplyVendorOperation(operation VendorOperation) (*VendorOperationResult, error) {
if operation.ExpectedVersion == "" {
return nil, ErrVendorConflict
}
return applyVendorOperation(operation)
}
func DeleteVendors(ids []int) error {
_, err := applyVendorOperation(VendorOperation{Action: "delete", VendorIDs: ids})
return err
}
func applyVendorOperation(operation VendorOperation) (*VendorOperationResult, error) {
result := &VendorOperationResult{UpdatedModels: []int{}, DeletedVendors: []int{}}
err := metadataTransaction(func(tx *gorm.DB) error {
preview, err := buildVendorOperationPreview(lockForUpdate(tx), operation)
if err != nil {
return err
}
if operation.ExpectedVersion != "" && operation.ExpectedVersion != preview.Version {
return ErrVendorConflict
}
for _, item := range preview.Models {
if item.VendorID != operation.TargetVendorID {
result.UpdatedModels = append(result.UpdatedModels, item.ID)
}
}
if len(result.UpdatedModels) > 0 {
if err := tx.Model(&Model{}).Where("id IN ?", result.UpdatedModels).Updates(map[string]any{"vendor_id": operation.TargetVendorID, "updated_time": common.GetTimestamp()}).Error; err != nil {
return err
}
}
if operation.Action == "merge" || operation.Action == "delete" {
for _, vendor := range preview.Sources {
result.DeletedVendors = append(result.DeletedVendors, vendor.Id)
}
if err := tx.Where("id IN ?", result.DeletedVendors).Delete(&Vendor{}).Error; err != nil {
return err
}
}
return nil
})
if err != nil {
return nil, err
}
RefreshPricing()
return result, nil
}