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 }