Files
team-api/internal/logic/admin/model_import_export.go
T

382 lines
13 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package admin
import (
"bytes"
"context"
"encoding/json"
"io"
"net/url"
"github.com/gogf/gf/v2/database/gdb"
"github.com/gogf/gf/v2/errors/gcode"
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/gtime"
"github.com/shopspring/decimal"
"github.com/qianfree/team-api/api/admin/v1"
"github.com/qianfree/team-api/internal/consts"
"github.com/qianfree/team-api/internal/dao"
"github.com/qianfree/team-api/internal/logic/billing"
"github.com/qianfree/team-api/internal/logic/relay"
do "github.com/qianfree/team-api/internal/model/do"
)
// ExportModelsJson 导出模型配置为 JSON 文件(含定价数据,用于跨环境迁移)
func (s *sAdmin) ExportModelsJson(ctx context.Context, req *v1.ModelExportJsonReq) (*v1.ModelExportJsonRes, error) {
var models []struct {
ID int64 `orm:"id" json:"id"`
ModelId string `orm:"model_id" json:"model_id"`
ModelName string `orm:"model_name" json:"model_name"`
Category string `orm:"category" json:"category"`
Status string `orm:"status" json:"status"`
MaxContextTokens int `orm:"max_context_tokens" json:"max_context_tokens"`
MaxOutputTokens int `orm:"max_output_tokens" json:"max_output_tokens"`
Description string `orm:"description" json:"description"`
Tags []string `orm:"tags" json:"tags"`
Capabilities string `orm:"capabilities" json:"capabilities"`
SunsetDate *gtime.Time `orm:"sunset_date" json:"sunset_date"`
ReplacementModel string `orm:"replacement_model" json:"replacement_model"`
}
err := dao.MdlModels.Ctx(ctx).
WhereIn("model_id", req.ModelIds).
Scan(&models)
if err != nil {
return nil, err
}
if len(models) == 0 {
return nil, gerror.New("未找到匹配的模型")
}
type exportPricing struct {
BillingMode string `json:"billing_mode"`
MinTokens int64 `json:"min_tokens"`
MaxTokens *int64 `json:"max_tokens"`
InputPrice float64 `json:"input_price"`
OutputPrice float64 `json:"output_price"`
PerRequestPrice *float64 `json:"per_request_price"`
CacheReadPrice float64 `json:"cache_read_price"`
CacheCreationPrice float64 `json:"cache_creation_price"`
}
type exportModel struct {
ModelId string `json:"model_id"`
ModelName string `json:"model_name"`
Category string `json:"category"`
Status string `json:"status"`
MaxContextTokens int `json:"max_context_tokens"`
MaxOutputTokens int `json:"max_output_tokens"`
Description string `json:"description"`
Tags []string `json:"tags"`
Capabilities map[string]bool `json:"capabilities"`
SunsetDate string `json:"sunset_date,omitempty"`
ReplacementModel string `json:"replacement_model,omitempty"`
Pricing []exportPricing `json:"pricing"`
TimeSegments []v1.TimeSegmentItem `json:"time_segments,omitempty"`
}
// 批量查询所有模型的定价行(避免循环内逐模型查询导致 N+1),与 ListModels 的批量模式对齐。
modelIDs := make([]int64, 0, len(models))
for _, m := range models {
modelIDs = append(modelIDs, m.ID)
}
type pricingRow struct {
ModelId int64 `orm:"model_id" json:"model_id"`
BillingMode string `orm:"billing_mode" json:"billing_mode"`
MinTokens int64 `orm:"min_tokens" json:"min_tokens"`
MaxTokens *int64 `orm:"max_tokens" json:"max_tokens"`
InputPrice float64 `orm:"input_price" json:"input_price"`
OutputPrice float64 `orm:"output_price" json:"output_price"`
PerRequestPrice *float64 `orm:"per_request_price" json:"per_request_price"`
CacheReadPrice float64 `orm:"cache_read_price" json:"cache_read_price"`
CacheCreationPrice float64 `orm:"cache_creation_price" json:"cache_creation_price"`
TimeSegments string `orm:"time_segments" json:"time_segments"`
}
var allPricingRows []pricingRow
if len(modelIDs) > 0 {
if err := dao.MdlPricing.Ctx(ctx).WhereIn("model_id", modelIDs).Scan(&allPricingRows); err != nil {
return nil, err
}
}
// 按 model_id 分组,供后续组装;时段定价只认锚点行(min_tokens=0)
pricingByModel := make(map[int64][]exportPricing, len(allPricingRows))
segmentsByModel := make(map[int64][]v1.TimeSegmentItem, len(allPricingRows))
for _, p := range allPricingRows {
pricingByModel[p.ModelId] = append(pricingByModel[p.ModelId], exportPricing{
BillingMode: p.BillingMode,
MinTokens: p.MinTokens,
MaxTokens: p.MaxTokens,
InputPrice: p.InputPrice,
OutputPrice: p.OutputPrice,
PerRequestPrice: p.PerRequestPrice,
CacheReadPrice: p.CacheReadPrice,
CacheCreationPrice: p.CacheCreationPrice,
})
if p.MinTokens == 0 && p.TimeSegments != "" && p.TimeSegments != "null" {
var segs []billing.TimeSegment
if err := json.Unmarshal([]byte(p.TimeSegments), &segs); err == nil {
segmentsByModel[p.ModelId] = timeSegmentsToAPI(segs)
}
}
}
result := make([]exportModel, 0, len(models))
for _, m := range models {
em := exportModel{
ModelId: m.ModelId,
ModelName: m.ModelName,
Category: m.Category,
Status: m.Status,
MaxContextTokens: m.MaxContextTokens,
MaxOutputTokens: m.MaxOutputTokens,
Description: m.Description,
Capabilities: parseCapabilities(m.Capabilities),
ReplacementModel: m.ReplacementModel,
Pricing: pricingByModel[m.ID],
TimeSegments: segmentsByModel[m.ID],
}
if em.Pricing == nil {
em.Pricing = []exportPricing{}
}
if m.SunsetDate != nil {
em.SunsetDate = m.SunsetDate.Format("Y-m-d")
}
em.Tags = m.Tags
result = append(result, em)
}
exportData := g.Map{
"version": "1.0",
"exported_at": gtime.Now().Format("Y-m-d\\TH:i:sP"),
"models": result,
}
jsonBytes, err := json.MarshalIndent(exportData, "", " ")
if err != nil {
return nil, err
}
filename := "models_" + gtime.Now().Format("Ymd_His") + ".json"
r := g.RequestFromCtx(ctx)
r.Response.Header().Set("Content-Type", "application/json; charset=utf-8")
r.Response.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''"+url.QueryEscape(filename))
r.Response.Write(jsonBytes)
return nil, nil
}
// ImportModelsPreview 导入模型预览(解析上传文件,检测冲突)
func (s *sAdmin) ImportModelsPreview(ctx context.Context, req *v1.ModelImportPreviewReq) (*v1.ModelImportPreviewRes, error) {
var data []byte
r := g.RequestFromCtx(ctx)
// 优先通过标准文件上传读取
file := r.GetUploadFile("file")
if file != nil {
f, err := file.Open()
if err != nil {
return nil, gerror.NewCode(gcode.New(consts.CodeModelImportInvalidFile, consts.MsgModelImportInvalidFile, nil), consts.MsgModelImportInvalidFile)
}
defer f.Close()
data, err = io.ReadAll(f)
if err != nil {
return nil, gerror.NewCode(gcode.New(consts.CodeModelImportInvalidFile, consts.MsgModelImportInvalidFile, nil), consts.MsgModelImportInvalidFile)
}
}
// 兜底:从已解析的 MultipartForm.File 中读取(文件在 File 而非 Value 中)
if len(data) == 0 {
mf := r.GetMultipartForm()
if mf != nil {
if files := mf.File["file"]; len(files) > 0 {
f, err := files[0].Open()
if err == nil {
defer f.Close()
data, _ = io.ReadAll(f)
}
}
}
}
if len(data) == 0 {
return nil, gerror.NewCode(gcode.New(consts.CodeModelImportInvalidFile, consts.MsgModelImportInvalidFile, nil), consts.MsgModelImportInvalidFile)
}
var exportData struct {
Version string `json:"version"`
Models []v1.ModelImportPreviewItem `json:"models"`
}
// 使用 json.Decoder 而非 json.Unmarshal,只解析第一个完整 JSON 值,
// 兼容旧版本导出文件末尾可能被中间件追加的标准响应体。
if err := json.NewDecoder(bytes.NewReader(data)).Decode(&exportData); err != nil {
return nil, gerror.NewCode(gcode.New(consts.CodeModelImportInvalidFile, consts.MsgModelImportInvalidFile, nil), consts.MsgModelImportInvalidFile)
}
if exportData.Version != "1.0" {
return nil, gerror.NewCode(gcode.New(consts.CodeModelImportBadVersion, consts.MsgModelImportBadVersion, nil), consts.MsgModelImportBadVersion)
}
for i := range exportData.Models {
count, _ := dao.MdlModels.Ctx(ctx).Where("model_id", exportData.Models[i].ModelId).Count()
if count > 0 {
exportData.Models[i].Conflict = "exists"
}
}
return &v1.ModelImportPreviewRes{Models: exportData.Models}, nil
}
// ImportModels 确认导入模型(事务内执行)
func (s *sAdmin) ImportModels(ctx context.Context, req *v1.ModelImportReq) (*v1.ModelImportRes, error) {
if len(req.Models) > 200 {
return nil, gerror.New("单次导入不能超过 200 个模型")
}
res := &v1.ModelImportRes{}
err := g.DB().Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
for _, item := range req.Models {
var existing *struct {
ID int64 `orm:"id"`
}
err := dao.MdlModels.Ctx(ctx).Where("model_id", item.ModelId).Scan(&existing)
if err != nil {
return err
}
if existing != nil {
if item.ConflictAction == "skip" {
res.Skipped++
continue
}
// overwrite: 更新模型 + 替换定价
updateData := do.MdlModels{
ModelName: item.ModelName,
Category: item.Category,
MaxContextTokens: item.MaxContextTokens,
MaxOutputTokens: item.MaxOutputTokens,
Description: item.Description,
}
if item.Tags != nil {
updateData.Tags = item.Tags
}
if item.Capabilities != nil {
capJson, _ := json.Marshal(item.Capabilities)
updateData.Capabilities = string(capJson)
}
if item.SunsetDate != "" {
updateData.SunsetDate = gtime.NewFromStr(item.SunsetDate)
}
if item.ReplacementModel != "" {
updateData.ReplacementModel = item.ReplacementModel
}
if item.Status != "" {
updateData.Status = item.Status
}
_, err = dao.MdlModels.Ctx(ctx).Where("id", existing.ID).Data(updateData).Update()
if err != nil {
return err
}
_, err = dao.MdlPricing.Ctx(ctx).Where("model_id", existing.ID).Delete()
if err != nil {
return err
}
for _, p := range item.Pricing {
// API 边界 float64 → DO decimal 转换
var perRequestPriceDecimal *decimal.Decimal
if p.PerRequestPrice != nil {
d := billing.NewFromFloat(*p.PerRequestPrice)
perRequestPriceDecimal = &d
}
_, err = dao.MdlPricing.Ctx(ctx).Insert(do.MdlPricing{
ModelId: existing.ID,
BillingMode: p.BillingMode,
MinTokens: p.MinTokens,
MaxTokens: p.MaxTokens,
InputPrice: billing.NewFromFloat(p.InputPrice),
OutputPrice: billing.NewFromFloat(p.OutputPrice),
PerRequestPrice: perRequestPriceDecimal,
CacheReadPrice: billing.NewFromFloat(p.CacheReadPrice),
CacheCreationPrice: billing.NewFromFloat(p.CacheCreationPrice),
})
if err != nil {
return err
}
}
// 时段定价透传(挂锚点行,全量替换语义:导入文件无该字段=清除)
if err := writeTimeSegmentsForModel(ctx, existing.ID, item.TimeSegments); err != nil {
return gerror.Wrapf(err, "模型 %s 时段定价", item.ModelId)
}
relay.NewDataProvider().InvalidateModelCache(item.ModelId)
billing.ClearModelPriceCache(ctx, item.ModelId)
res.Imported++
} else {
// 新建模型
insertData := do.MdlModels{
ModelId: item.ModelId,
ModelName: item.ModelName,
Category: item.Category,
Status: "active",
MaxContextTokens: item.MaxContextTokens,
MaxOutputTokens: item.MaxOutputTokens,
Description: item.Description,
}
if item.Tags != nil {
insertData.Tags = item.Tags
} else {
insertData.Tags = []string{}
}
if item.Capabilities != nil {
capJson, _ := json.Marshal(item.Capabilities)
insertData.Capabilities = string(capJson)
}
id, err := dao.MdlModels.Ctx(ctx).InsertAndGetId(insertData)
if err != nil {
return err
}
for _, p := range item.Pricing {
// API 边界 float64 → DO decimal 转换
var perRequestPriceDecimal *decimal.Decimal
if p.PerRequestPrice != nil {
d := billing.NewFromFloat(*p.PerRequestPrice)
perRequestPriceDecimal = &d
}
_, err = dao.MdlPricing.Ctx(ctx).Insert(do.MdlPricing{
ModelId: id,
BillingMode: p.BillingMode,
MinTokens: p.MinTokens,
MaxTokens: p.MaxTokens,
InputPrice: billing.NewFromFloat(p.InputPrice),
OutputPrice: billing.NewFromFloat(p.OutputPrice),
PerRequestPrice: perRequestPriceDecimal,
CacheReadPrice: billing.NewFromFloat(p.CacheReadPrice),
CacheCreationPrice: billing.NewFromFloat(p.CacheCreationPrice),
})
if err != nil {
return err
}
}
// 时段定价透传(挂锚点行)
if err := writeTimeSegmentsForModel(ctx, id, item.TimeSegments); err != nil {
return gerror.Wrapf(err, "模型 %s 时段定价", item.ModelId)
}
res.Imported++
}
}
return nil
})
return res, err
}