Merge pull request #401 from slovx2/heihuzi_main
feat(gemini): 为 Gemini 原生平台添加图片计费支持
This commit is contained in:
@@ -931,6 +931,13 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 图片生成计费
|
||||||
|
imageCount := 0
|
||||||
|
imageSize := s.extractImageSize(body)
|
||||||
|
if isImageGenerationModel(originalModel) {
|
||||||
|
imageCount = 1
|
||||||
|
}
|
||||||
|
|
||||||
return &ForwardResult{
|
return &ForwardResult{
|
||||||
RequestID: requestID,
|
RequestID: requestID,
|
||||||
Usage: *usage,
|
Usage: *usage,
|
||||||
@@ -938,6 +945,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
|||||||
Stream: req.Stream,
|
Stream: req.Stream,
|
||||||
Duration: time.Since(startTime),
|
Duration: time.Since(startTime),
|
||||||
FirstTokenMs: firstTokenMs,
|
FirstTokenMs: firstTokenMs,
|
||||||
|
ImageCount: imageCount,
|
||||||
|
ImageSize: imageSize,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1371,6 +1380,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
|||||||
usage = &ClaudeUsage{}
|
usage = &ClaudeUsage{}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 图片生成计费
|
||||||
|
imageCount := 0
|
||||||
|
imageSize := s.extractImageSize(body)
|
||||||
|
if isImageGenerationModel(originalModel) {
|
||||||
|
imageCount = 1
|
||||||
|
}
|
||||||
|
|
||||||
return &ForwardResult{
|
return &ForwardResult{
|
||||||
RequestID: requestID,
|
RequestID: requestID,
|
||||||
Usage: *usage,
|
Usage: *usage,
|
||||||
@@ -1378,6 +1394,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
|||||||
Stream: stream,
|
Stream: stream,
|
||||||
Duration: time.Since(startTime),
|
Duration: time.Since(startTime),
|
||||||
FirstTokenMs: firstTokenMs,
|
FirstTokenMs: firstTokenMs,
|
||||||
|
ImageCount: imageCount,
|
||||||
|
ImageSize: imageSize,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -3031,3 +3049,26 @@ func convertClaudeGenerationConfig(req map[string]any) map[string]any {
|
|||||||
}
|
}
|
||||||
return out
|
return out
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// extractImageSize 从 Gemini 请求中提取 image_size 参数
|
||||||
|
func (s *GeminiMessagesCompatService) extractImageSize(body []byte) string {
|
||||||
|
var req struct {
|
||||||
|
GenerationConfig *struct {
|
||||||
|
ImageConfig *struct {
|
||||||
|
ImageSize string `json:"imageSize"`
|
||||||
|
} `json:"imageConfig"`
|
||||||
|
} `json:"generationConfig"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(body, &req); err != nil {
|
||||||
|
return "2K"
|
||||||
|
}
|
||||||
|
|
||||||
|
if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil {
|
||||||
|
size := strings.ToUpper(strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize))
|
||||||
|
if size == "1K" || size == "2K" || size == "4K" {
|
||||||
|
return size
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return "2K"
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user