new-api/service/chinamobile_user_asset_group.go

package service

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
)

type chinaMobileAssetGroupCreator func(context.Context, int, *model.Channel) (string, *AssetError)

func IsChinaMobileAssetChannel(channel *model.Channel) bool {
	return channel != nil && channel.Type == constant.ChannelTypeChinaMobileSeedance
}

func IsAssetGroupOperation(operation AssetOperation) bool {
	return strings.HasPrefix(string(operation), "asset_group.")
}

func GetOrCreateChinaMobileUserAssetGroup(ctx context.Context, userId int, channel *model.Channel, createGroup chinaMobileAssetGroupCreator) (string, *AssetError) {
	if channel == nil {
		return "", newAssetError(AssetErrorServer, "China Mobile asset channel is required", http.StatusInternalServerError)
	}
	binding, err := model.GetUserAssetGroup(userId, channel.Id)
	if err != nil {
		return "", newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	if binding != nil {
		return binding.GroupId, nil
	}

	groupId, assetErr := createGroup(ctx, userId, channel)
	if assetErr != nil {
		return "", assetErr
	}
	if strings.TrimSpace(groupId) == "" {
		return "", newAssetError(AssetErrorUpstream, "China Mobile asset group creation returned an empty group ID", http.StatusBadGateway)
	}
	if err := model.CreateUserAssetGroup(userId, channel.Id, groupId); err == nil {
		return groupId, nil
	}

	// Another first request may have persisted the binding while this one created its upstream group.
	binding, readErr := model.GetUserAssetGroup(userId, channel.Id)
	if readErr == nil && binding != nil {
		return binding.GroupId, nil
	}
	if readErr != nil {
		return "", newAssetError(AssetErrorServer, readErr.Error(), http.StatusInternalServerError)
	}
	return "", newAssetError(AssetErrorServer, "failed to persist China Mobile user asset group", http.StatusInternalServerError)
}

func CreateChinaMobileUserAssetGroup(ctx context.Context, userId int, adapter AssetAdapter, channel *model.Channel) (string, *AssetError) {
	spec, ok := ParseAssetAction("CreateAssetGroup")
	if !ok {
		return "", newAssetError(AssetErrorServer, "CreateAssetGroup action is not registered", http.StatusInternalServerError)
	}
	resp, assetErr := adapter.DoAssetRequest(ctx, channel, AssetRequest{
		Action:  spec,
		Version: "2024-01-01",
		Body: map[string]any{
			"GroupType": "AIGC",
			"Name":      fmt.Sprintf("new-api-user-%d-channel-%d", userId, channel.Id),
		},
	})
	if assetErr != nil {
		return "", assetErr
	}
	var payload struct {
		Result struct {
			GroupId string `json:"GroupId"`
		} `json:"Result"`
	}
	if err := common.Unmarshal(resp.Body, &payload); err != nil {
		return "", newAssetError(AssetErrorUpstream, fmt.Sprintf("invalid China Mobile asset group response: %v", err), http.StatusBadGateway)
	}
	return strings.TrimSpace(payload.Result.GroupId), nil
}

func ScopeChinaMobileAssetRequest(req *AssetRequest, groupId string) {
	switch req.Action.Operation {
	case AssetOperationAssetCreate:
		req.Body["GroupId"] = groupId
	case AssetOperationAssetList:
		filter := mapValue(req.Body, "Filter")
		if filter == nil {
			return
		}
		filter["GroupIds"] = []string{groupId}
		req.Body["Filter"] = filter
	}
}

func RequireChinaMobileAssetOwnership(ctx context.Context, adapter AssetAdapter, channel *model.Channel, request AssetRequest, groupId string) *AssetError {
	spec, ok := ParseAssetAction("GetAsset")
	if !ok {
		return newAssetError(AssetErrorServer, "GetAsset action is not registered", http.StatusInternalServerError)
	}
	assetId := strings.TrimSpace(stringValue(request.Body, "Id"))
	resp, assetErr := adapter.DoAssetRequest(ctx, channel, AssetRequest{
		Action:  spec,
		Version: request.Version,
		Body:    map[string]any{"Id": assetId},
	})
	if assetErr != nil {
		if assetErr.HTTPStatus == http.StatusNotFound {
			return newAssetError(AssetErrorNotFound, "asset not found", http.StatusNotFound)
		}
		return assetErr
	}
	var payload struct {
		Result struct {
			GroupId string `json:"GroupId"`
		} `json:"Result"`
	}
	if err := common.Unmarshal(resp.Body, &payload); err != nil {
		return newAssetError(AssetErrorUpstream, fmt.Sprintf("invalid China Mobile asset response: %v", err), http.StatusBadGateway)
	}
	if strings.TrimSpace(payload.Result.GroupId) != groupId {
		return newAssetError(AssetErrorNotFound, "asset not found", http.StatusNotFound)
	}
	return nil
}