xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-11 02:43:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: add tool choice configuration and update steam handling in Gemini

Browse Source 
update ToolConfig to use pointer type in ChatRequest
This commit is contained in:
mxdlzg
2024-12-23 12:58:57 +08:00
parent a858292b54
commit ba50a137ea
2 changed files with 50 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
relay/adaptor/gemini/main.go
View File

@@ -34,6 +34,12 @@ var mimeTypeMap = map[string]string{
"text": "text/plain",
}
var toolChoiceTypeMap = map[string]string{
"none": "NONE",
"auto": "AUTO",
"required": "ANY",
}
// Setting safety to the lowest possible values since Gemini is already powerless enough
func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
geminiRequest := ChatRequest{
@@ -92,6 +98,24 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
},
}
}
if textRequest.ToolChoice != nil {
geminiRequest.ToolConfig = &ToolConfig{
FunctionCallingConfig: FunctionCallingConfig{
Mode: "auto",
},
}
switch mode := textRequest.ToolChoice.(type) {
case string:
geminiRequest.ToolConfig.FunctionCallingConfig.Mode = toolChoiceTypeMap[mode]
case map[string]interface{}:
geminiRequest.ToolConfig.FunctionCallingConfig.Mode = "ANY"
if fn, ok := mode["function"].(map[string]interface{}); ok {
if name, ok := fn["name"].(string); ok {
geminiRequest.ToolConfig.FunctionCallingConfig.AllowedFunctionNames = []string{name}
}
}
}
}
shouldAddDummyModelMessage := false
for _, message := range textRequest.Messages {
content := ChatContent{
@@ -186,10 +210,16 @@ func (g *ChatResponse) GetResponseText() string {
if g == nil {
return ""
}
if len(g.Candidates) > 0 && len(g.Candidates[0].Content.Parts) > 0 {
return g.Candidates[0].Content.Parts[0].Text
var builder strings.Builder
for _, candidate := range g.Candidates {
for idx, part := range candidate.Content.Parts {
if idx > 0 {
builder.WriteString("\n")
}
builder.WriteString(part.Text)
}
}
return ""
return builder.String()
}
type ChatCandidate struct {
@@ -252,8 +282,8 @@ func responseGeminiChat2OpenAI(response *ChatResponse) *openai.TextResponse {
choice.Message.ToolCalls = getToolCalls(&candidate)
} else {
var builder strings.Builder
for _, part := range candidate.Content.Parts {
if i > 0 {
for idx, part := range candidate.Content.Parts {
if idx > 0 {
builder.WriteString("\n")
}
builder.WriteString(part.Text)