Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 18 additions & 18 deletions common/init.go
Original file line number Diff line number Diff line change
Expand Up @@ -146,30 +146,30 @@ func InitEnv() {
}

func initConstantEnv() {
constant.StreamingTimeout = GetEnvOrDefault("STREAMING_TIMEOUT", 300)
constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", true)
constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", 64)
constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", 128)
constant.StreamingTimeout = GetEnvOrDefault("STREAMING_TIMEOUT", constant.DefaultStreamingTimeout)
constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", constant.DefaultDifyDebug)
constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", constant.DefaultMaxFileDownloadMB)
constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", constant.DefaultStreamScannerMaxBufferMB)
// MaxRequestBodyMB 请求体最大大小(解压后),用于防止超大请求/zip bomb导致内存暴涨
constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", 128)
constant.AnonymousRequestBodyLimitKB = GetEnvOrDefault("ANONYMOUS_REQUEST_BODY_LIMIT_KB", 512)
constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", constant.DefaultMaxRequestBodyMB)
constant.AnonymousRequestBodyLimitKB = GetEnvOrDefault("ANONYMOUS_REQUEST_BODY_LIMIT_KB", constant.DefaultAnonymousRequestBodyLimitKB)
// ForceStreamOption 覆盖请求参数,强制返回usage信息
constant.ForceStreamOption = GetEnvOrDefaultBool("FORCE_STREAM_OPTION", true)
constant.CountToken = GetEnvOrDefaultBool("CountToken", true)
constant.GetMediaToken = GetEnvOrDefaultBool("GET_MEDIA_TOKEN", true)
constant.GetMediaTokenNotStream = GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", false)
constant.UpdateTask = GetEnvOrDefaultBool("UPDATE_TASK", true)
constant.AzureDefaultAPIVersion = GetEnvOrDefaultString("AZURE_DEFAULT_API_VERSION", "2025-04-01-preview")
constant.NotifyLimitCount = GetEnvOrDefault("NOTIFY_LIMIT_COUNT", 2)
constant.NotificationLimitDurationMinute = GetEnvOrDefault("NOTIFICATION_LIMIT_DURATION_MINUTE", 10)
constant.ForceStreamOption = GetEnvOrDefaultBool("FORCE_STREAM_OPTION", constant.DefaultForceStreamOption)
constant.CountToken = GetEnvOrDefaultBool("CountToken", constant.DefaultCountToken)
constant.GetMediaToken = GetEnvOrDefaultBool("GET_MEDIA_TOKEN", constant.DefaultGetMediaToken)
constant.GetMediaTokenNotStream = GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", constant.DefaultGetMediaTokenNotStream)
constant.UpdateTask = GetEnvOrDefaultBool("UPDATE_TASK", constant.DefaultUpdateTask)
constant.AzureDefaultAPIVersion = GetEnvOrDefaultString("AZURE_DEFAULT_API_VERSION", constant.DefaultAzureDefaultAPIVersion)
constant.NotifyLimitCount = GetEnvOrDefault("NOTIFY_LIMIT_COUNT", constant.DefaultNotifyLimitCount)
constant.NotificationLimitDurationMinute = GetEnvOrDefault("NOTIFICATION_LIMIT_DURATION_MINUTE", constant.DefaultNotificationLimitDurationMinute)
// GenerateDefaultToken 是否生成初始令牌,默认关闭。
constant.GenerateDefaultToken = GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false)
constant.GenerateDefaultToken = GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", constant.DefaultGenerateDefaultToken)
// 是否启用错误日志
constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", false)
constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", constant.DefaultErrorLogEnabled)
// 任务轮询时查询的最大数量
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000)
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", constant.DefaultTaskQueryLimit)
// 异步任务超时时间(分钟),超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440)
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", constant.DefaultTaskTimeoutMinutes)

soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "")
if soraPatchStr != "" {
Expand Down
57 changes: 39 additions & 18 deletions constant/env.go
Original file line number Diff line number Diff line change
@@ -1,23 +1,44 @@
package constant

var StreamingTimeout int
var DifyDebug bool
var MaxFileDownloadMB int
var StreamScannerMaxBufferMB int
var ForceStreamOption bool
var CountToken bool
var GetMediaToken bool
var GetMediaTokenNotStream bool
var UpdateTask bool
var MaxRequestBodyMB int
var AnonymousRequestBodyLimitKB int
var AzureDefaultAPIVersion string
var NotifyLimitCount int
var NotificationLimitDurationMinute int
var GenerateDefaultToken bool
var ErrorLogEnabled bool
var TaskQueryLimit int
var TaskTimeoutMinutes int
const (
DefaultStreamingTimeout = 300
DefaultMaxFileDownloadMB = 64
DefaultStreamScannerMaxBufferMB = 128
DefaultMaxRequestBodyMB = 128
DefaultAnonymousRequestBodyLimitKB = 512
DefaultAzureDefaultAPIVersion = "2025-04-01-preview"
DefaultNotifyLimitCount = 2
DefaultNotificationLimitDurationMinute = 10
DefaultTaskQueryLimit = 1000
DefaultTaskTimeoutMinutes = 1440
DefaultDifyDebug = true
DefaultForceStreamOption = true
DefaultCountToken = true
DefaultGetMediaToken = true
DefaultGetMediaTokenNotStream = false
DefaultUpdateTask = true
DefaultGenerateDefaultToken = false
DefaultErrorLogEnabled = false
)

var StreamingTimeout = DefaultStreamingTimeout
var DifyDebug = DefaultDifyDebug
var MaxFileDownloadMB = DefaultMaxFileDownloadMB
var StreamScannerMaxBufferMB = DefaultStreamScannerMaxBufferMB
var ForceStreamOption = DefaultForceStreamOption
var CountToken = DefaultCountToken
var GetMediaToken = DefaultGetMediaToken
var GetMediaTokenNotStream = DefaultGetMediaTokenNotStream
var UpdateTask = DefaultUpdateTask
var MaxRequestBodyMB = DefaultMaxRequestBodyMB
var AnonymousRequestBodyLimitKB = DefaultAnonymousRequestBodyLimitKB
var AzureDefaultAPIVersion = DefaultAzureDefaultAPIVersion
var NotifyLimitCount = DefaultNotifyLimitCount
var NotificationLimitDurationMinute = DefaultNotificationLimitDurationMinute
var GenerateDefaultToken = DefaultGenerateDefaultToken
var ErrorLogEnabled = DefaultErrorLogEnabled
var TaskQueryLimit = DefaultTaskQueryLimit
var TaskTimeoutMinutes = DefaultTaskTimeoutMinutes
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// temporary variable for sora patch, will be removed in future
var TaskPricePatches []string
Expand Down
31 changes: 31 additions & 0 deletions constant/env_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
package constant

import "testing"

func TestEnvironmentBackedDefaultsAreInitialized(t *testing.T) {
checks := map[string]bool{
"StreamingTimeout": StreamingTimeout == DefaultStreamingTimeout,
"DifyDebug": DifyDebug == DefaultDifyDebug,
"MaxFileDownloadMB": MaxFileDownloadMB == DefaultMaxFileDownloadMB,
"StreamScannerMaxBufferMB": StreamScannerMaxBufferMB == DefaultStreamScannerMaxBufferMB,
"ForceStreamOption": ForceStreamOption == DefaultForceStreamOption,
"CountToken": CountToken == DefaultCountToken,
"GetMediaToken": GetMediaToken == DefaultGetMediaToken,
"GetMediaTokenNotStream": GetMediaTokenNotStream == DefaultGetMediaTokenNotStream,
"UpdateTask": UpdateTask == DefaultUpdateTask,
"MaxRequestBodyMB": MaxRequestBodyMB == DefaultMaxRequestBodyMB,
"AnonymousRequestBodyLimitKB": AnonymousRequestBodyLimitKB == DefaultAnonymousRequestBodyLimitKB,
"AzureDefaultAPIVersion": AzureDefaultAPIVersion == DefaultAzureDefaultAPIVersion,
"NotifyLimitCount": NotifyLimitCount == DefaultNotifyLimitCount,
"NotificationLimitDurationMinute": NotificationLimitDurationMinute == DefaultNotificationLimitDurationMinute,
"GenerateDefaultToken": GenerateDefaultToken == DefaultGenerateDefaultToken,
"ErrorLogEnabled": ErrorLogEnabled == DefaultErrorLogEnabled,
"TaskQueryLimit": TaskQueryLimit == DefaultTaskQueryLimit,
"TaskTimeoutMinutes": TaskTimeoutMinutes == DefaultTaskTimeoutMinutes,
}
for name, ok := range checks {
if !ok {
t.Fatalf("%s package default does not match expected init default", name)
}
}
}
58 changes: 48 additions & 10 deletions model/log.go
Original file line number Diff line number Diff line change
Expand Up @@ -95,16 +95,35 @@ func applyLogTypeFilter(tx *gorm.DB, logType int) *gorm.DB {
return tx.Where("logs.type = ?", logType)
}

func applyRetryLogFilter(tx *gorm.DB) *gorm.DB {
return tx.Where("logs.is_retry = ?", true)
var ensureLogRetryMarkerBackfillCompletedForRead = ensureLogRetryMarkerBackfillCompletedForReadDefault

func applyRetryLogFilter(tx *gorm.DB) (*gorm.DB, error) {
if err := ensureLogRetryMarkerBackfillCompletedForRead(); err != nil {
return nil, err
}
return tx.Where("logs.is_retry = ?", true), nil
}

func applyLogFilter(tx *gorm.DB, filter string) *gorm.DB {
func ensureLogRetryMarkerBackfillCompletedForReadDefault() error {
completed, err := isLogRetryMarkerBackfillCompleted()
if err != nil {
return fmt.Errorf("failed to check log retry marker backfill status before retry log read: %w", err)
}
if completed {
return nil
}
if err := backfillLogRetryMarker(); err != nil {
return fmt.Errorf("failed to backfill log retry markers before retry log read: %w", err)
}
return nil
}

func applyLogFilter(tx *gorm.DB, filter string) (*gorm.DB, error) {
switch filter {
case LogFilterRetry:
return applyRetryLogFilter(tx)
default:
return tx
return tx, nil
}
}

Expand Down Expand Up @@ -522,7 +541,10 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
}

func GetAllLogs(logType int, logFilter string, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
tx := applyLogFilter(applyLogTypeFilter(LOG_DB, logType), logFilter)
tx, err := applyLogFilter(applyLogTypeFilter(LOG_DB, logType), logFilter)
if err != nil {
return nil, 0, err
}

if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil {
return nil, 0, err
Expand Down Expand Up @@ -607,7 +629,10 @@ func GetAllLogs(logType int, logFilter string, startTimestamp int64, endTimestam
const logSearchCountLimit = 10000

func GetUserLogs(userId int, logType int, logFilter string, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
tx := applyLogFilter(applyLogTypeFilter(LOG_DB.Where("logs.user_id = ?", userId), logType), logFilter)
tx, err := applyLogFilter(applyLogTypeFilter(LOG_DB.Where("logs.user_id = ?", userId), logType), logFilter)
if err != nil {
return nil, 0, err
}

if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil {
return nil, 0, err
Expand Down Expand Up @@ -688,8 +713,14 @@ func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimest
// 为rpm和tpm创建单独的查询
rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm")

tx = applyLogFilter(tx, logFilter)
rpmTpmQuery = applyLogFilter(rpmTpmQuery, logFilter)
tx, err = applyLogFilter(tx, logFilter)
if err != nil {
return stat, err
}
rpmTpmQuery, err = applyLogFilter(rpmTpmQuery, logFilter)
if err != nil {
return stat, err
}

if tx, err = applyExplicitLogTextFilter(tx, "username", username); err != nil {
return stat, err
Expand All @@ -703,9 +734,11 @@ func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimest
}
if startTimestamp != 0 {
tx = tx.Where("created_at >= ?", startTimestamp)
rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", startTimestamp)
}
if endTimestamp != 0 {
tx = tx.Where("created_at <= ?", endTimestamp)
rpmTpmQuery = rpmTpmQuery.Where("created_at <= ?", endTimestamp)
}
if tx, err = applyExplicitLogTextFilter(tx, "model_name", modelName); err != nil {
return stat, err
Expand All @@ -722,8 +755,13 @@ func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimest
rpmTpmQuery = rpmTpmQuery.Where(logGroupCol+" = ?", group)
}

tx = tx.Where("type = ?", LogTypeConsume)
rpmTpmQuery = rpmTpmQuery.Where("type = ?", LogTypeConsume)
if logType != LogTypeUnknown {
tx = tx.Where("logs.type = ?", logType)
rpmTpmQuery = rpmTpmQuery.Where("logs.type = ?", logType)
} else if logFilter != LogFilterRetry {
tx = tx.Where("logs.type = ?", LogTypeConsume)
rpmTpmQuery = rpmTpmQuery.Where("logs.type = ?", LogTypeConsume)
}

// 只统计最近60秒的rpm和tpm
rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", time.Now().Add(-60*time.Second).Unix())
Expand Down
Loading
Loading