fix: integrate Gemini v2 modalities support and refactor response handling

This commit is contained in:
Laisky.Cai 2025-03-19 00:40:28 +00:00
parent fc6d85932b
commit 2d978686bb
2 changed files with 21 additions and 6 deletions

View File

@ -16,6 +16,7 @@ import (
"github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/random"
"github.com/songquanpeng/one-api/common/render"
"github.com/songquanpeng/one-api/relay/adaptor/geminiv2"
"github.com/songquanpeng/one-api/relay/adaptor/openai"
"github.com/songquanpeng/one-api/relay/constant"
"github.com/songquanpeng/one-api/relay/model"
@ -59,12 +60,10 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
},
},
GenerationConfig: ChatGenerationConfig{
Temperature: textRequest.Temperature,
TopP: textRequest.TopP,
MaxOutputTokens: textRequest.MaxTokens,
ResponseModalities: []string{
"TEXT", "IMAGE",
},
Temperature: textRequest.Temperature,
TopP: textRequest.TopP,
MaxOutputTokens: textRequest.MaxTokens,
ResponseModalities: geminiv2.GetModelModalities(textRequest.Model),
},
}
if textRequest.ResponseFormat != nil {

View File

@ -1,5 +1,7 @@
package geminiv2
import "strings"
// https://ai.google.dev/models/gemini
var ModelList = []string{
@ -14,3 +16,17 @@ var ModelList = []string{
"gemini-2.0-pro-exp-02-05",
"gemini-2.0-flash-exp-image-generation",
}
const (
ModalityText = "TEXT"
ModalityImage = "IMAGE"
)
// GetModelModalities returns the modalities of the model.
func GetModelModalities(model string) []string {
if strings.Contains(model, "-image-generation") {
return []string{ModalityText, ModalityImage}
}
return []string{ModalityText}
}