new-api-custom/model/vendor_management.go

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
}