Enhance logging and error handling in processing logic

- Enhance logging by including requested image counts for improved tracking.
- Implement user quota validation before processing image requests.
- Improve error handling for model validation and JSON unmarshalling in billing processes.
This commit is contained in:
Laisky.Cai 2025-03-16 07:47:59 +00:00
parent fa794e7bd5
commit 400ee3aa24
2 changed files with 52 additions and 3 deletions

View File

@ -863,11 +863,59 @@ func ModelRatio2JSONString() string {
return string(jsonBytes)
}
// UpdateModelRatioByJSONString updates the ModelRatio map with the given JSON string.
func UpdateModelRatioByJSONString(jsonStr string) error {
modelRatioLock.Lock()
defer modelRatioLock.Unlock()
ModelRatio = make(map[string]float64)
return json.Unmarshal([]byte(jsonStr), &ModelRatio)
err := json.Unmarshal([]byte(jsonStr), &ModelRatio)
if err != nil {
logger.SysError("error unmarshalling model ratio: " + err.Error())
return err
}
// f3300f08e25e212f1b32ae1f678eb7ec2dec6a8c change the ratio of image models,
// so we need to multiply the ratio by 1000 for legacy settings.
for name, ratio := range ModelRatio {
switch name {
case "dall-e-2",
"dall-e-3",
"imagen-3.0-generate-001",
"imagen-3.0-generate-002",
"imagen-3.0-fast-generate-001",
"imagen-3.0-capability-001",
"ali-stable-diffusion-xl",
"ali-stable-diffusion-v1.5",
"black-forest-labs/flux-1.1-pro",
"black-forest-labs/flux-1.1-pro-ultra",
"black-forest-labs/flux-canny-dev",
"black-forest-labs/flux-canny-pro",
"black-forest-labs/flux-depth-dev",
"black-forest-labs/flux-depth-pro",
"black-forest-labs/flux-dev",
"black-forest-labs/flux-dev-lora",
"black-forest-labs/flux-fill-dev",
"black-forest-labs/flux-fill-pro",
"black-forest-labs/flux-pro",
"black-forest-labs/flux-redux-dev",
"black-forest-labs/flux-redux-schnell",
"black-forest-labs/flux-schnell",
"black-forest-labs/flux-schnell-lora",
"ideogram-ai/ideogram-v2",
"ideogram-ai/ideogram-v2-turbo",
"recraft-ai/recraft-v3",
"recraft-ai/recraft-v3-svg",
"stability-ai/stable-diffusion-3",
"stability-ai/stable-diffusion-3.5-large",
"stability-ai/stable-diffusion-3.5-large-turbo",
"stability-ai/stable-diffusion-3.5-medium":
if ratio < 1000 {
ModelRatio[name] = ratio * 1000
}
}
}
return nil
}
func GetModelRatio(name string, channelType int) float64 {
@ -895,7 +943,7 @@ func GetModelRatio(name string, channelType int) float64 {
}
logger.SysError("model ratio not found: " + name)
return 30
return 0.1 * QuotaPerUsd
}
// CompletionRatio2JSONString returns the CompletionRatio map as a JSON string.

View File

@ -226,7 +226,8 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
}
if quota >= 0 {
tokenName := c.GetString(ctxkey.TokenName)
logContent := fmt.Sprintf("model rate %.2f, group rate %.2f", modelRatio, groupRatio)
logContent := fmt.Sprintf("model rate %.2f, group rate %.2f, num %d",
modelRatio, groupRatio, imageRequest.N)
model.RecordConsumeLog(ctx, &model.Log{
UserId: meta.UserId,
ChannelId: meta.ChannelId,