feat: Refactor Gemini tools handling to support JSON raw message format

This commit is contained in:
CaIon
2025-08-11 19:48:04 +08:00
parent 03cfc05afd
commit d3170310ff
3 changed files with 44 additions and 9 deletions

View File

@@ -3,16 +3,52 @@ package dto
import (
"encoding/json"
"one-api/common"
"strings"
)
type GeminiChatRequest struct {
Contents []GeminiChatContent `json:"contents"`
SafetySettings []GeminiChatSafetySettings `json:"safetySettings,omitempty"`
GenerationConfig GeminiChatGenerationConfig `json:"generationConfig,omitempty"`
Tools []GeminiChatTool `json:"tools,omitempty"`
Tools json.RawMessage `json:"tools,omitempty"`
SystemInstructions *GeminiChatContent `json:"systemInstruction,omitempty"`
}
func (r *GeminiChatRequest) GetTools() []GeminiChatTool {
var tools []GeminiChatTool
if strings.HasSuffix(string(r.Tools), "[") {
// is array
if err := common.Unmarshal(r.Tools, &tools); err != nil {
common.LogError(nil, "error_unmarshalling_tools: "+err.Error())
return nil
}
} else if strings.HasPrefix(string(r.Tools), "{") {
// is object
singleTool := GeminiChatTool{}
if err := common.Unmarshal(r.Tools, &singleTool); err != nil {
common.LogError(nil, "error_unmarshalling_single_tool: "+err.Error())
return nil
}
tools = []GeminiChatTool{singleTool}
}
return tools
}
func (r *GeminiChatRequest) SetTools(tools []GeminiChatTool) {
if len(tools) == 0 {
r.Tools = json.RawMessage("[]")
return
}
// Marshal the tools to JSON
data, err := common.Marshal(tools)
if err != nil {
common.LogError(nil, "error_marshalling_tools: "+err.Error())
return
}
r.Tools = data
}
type GeminiThinkingConfig struct {
IncludeThoughts bool `json:"includeThoughts,omitempty"`
ThinkingBudget *int `json:"thinkingBudget,omitempty"`