You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
95 lines
2.4 KiB
95 lines
2.4 KiB
package ai
|
|
|
|
import (
|
|
"context"
|
|
"encoding/csv"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/youruser/base/internal/svc"
|
|
"github.com/youruser/base/internal/types"
|
|
"github.com/youruser/base/model"
|
|
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
)
|
|
|
|
type AiUsageExportLogic struct {
|
|
logx.Logger
|
|
ctx context.Context
|
|
svcCtx *svc.ServiceContext
|
|
}
|
|
|
|
func NewAiUsageExportLogic(ctx context.Context, svcCtx *svc.ServiceContext) *AiUsageExportLogic {
|
|
return &AiUsageExportLogic{
|
|
Logger: logx.WithContext(ctx),
|
|
ctx: ctx,
|
|
svcCtx: svcCtx,
|
|
}
|
|
}
|
|
|
|
func (l *AiUsageExportLogic) AiUsageExport(req *types.AIUsageRecordListRequest, w http.ResponseWriter) error {
|
|
// Get current user from context
|
|
userId := req.UserId
|
|
if userId == 0 {
|
|
if uid, ok := l.ctx.Value("userId").(int64); ok {
|
|
userId = uid
|
|
} else if uidJson, ok2 := l.ctx.Value("userId").(float64); ok2 {
|
|
userId = int64(uidJson)
|
|
}
|
|
}
|
|
|
|
// Fetch all records (large page size for export)
|
|
records, _, err := model.AIUsageRecordFindList(l.ctx, l.svcCtx.DB, userId, req.ModelId, req.Status, 1, 10000)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Build lookup caches
|
|
userMap := make(map[int64]string)
|
|
providerMap := make(map[int64]string)
|
|
for _, r := range records {
|
|
if _, ok := userMap[r.UserId]; !ok {
|
|
user, err := model.FindOne(l.ctx, l.svcCtx.DB, r.UserId)
|
|
if err == nil {
|
|
userMap[r.UserId] = user.Username
|
|
}
|
|
}
|
|
if _, ok := providerMap[r.ProviderId]; !ok {
|
|
p, err := model.AIProviderFindOne(l.ctx, l.svcCtx.DB, r.ProviderId)
|
|
if err == nil {
|
|
providerMap[r.ProviderId] = p.DisplayName
|
|
}
|
|
}
|
|
}
|
|
|
|
// Write CSV response
|
|
filename := fmt.Sprintf("ai-usage-%s.csv", time.Now().Format("20060102"))
|
|
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
|
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filename))
|
|
// BOM for Excel UTF-8 support
|
|
w.Write([]byte{0xEF, 0xBB, 0xBF})
|
|
|
|
writer := csv.NewWriter(w)
|
|
defer writer.Flush()
|
|
|
|
// Header row
|
|
writer.Write([]string{"时间", "用户", "平台", "模型", "输入Tokens", "输出Tokens", "费用(¥)", "延迟(ms)", "状态", "错误信息"})
|
|
|
|
for _, r := range records {
|
|
writer.Write([]string{
|
|
r.CreatedAt.Format("2006-01-02 15:04:05"),
|
|
userMap[r.UserId],
|
|
providerMap[r.ProviderId],
|
|
r.ModelId,
|
|
fmt.Sprintf("%d", r.InputTokens),
|
|
fmt.Sprintf("%d", r.OutputTokens),
|
|
fmt.Sprintf("%.4f", r.Cost),
|
|
fmt.Sprintf("%d", r.LatencyMs),
|
|
r.Status,
|
|
r.ErrorMessage,
|
|
})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|