Skip to content

Commit

Permalink
feat: update Gemini adaptor to support custom response format
Browse files Browse the repository at this point in the history
The code changes in the `main.go` file update the Gemini adaptor to support custom response formats. It adds a `mimeTypeMap` variable to map response format types to MIME types. The `ConvertRequest` function now checks if the `ResponseFormat` field is provided in the request and sets the appropriate response MIME type and schema in the `geminiRequest` object.

Co-authored-by: JustSong <[email protected]>
  • Loading branch information
mxdlzg and songquanpeng committed Oct 24, 2024
1 parent 489334d commit fd2846d
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 7 deletions.
17 changes: 16 additions & 1 deletion relay/adaptor/gemini/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,12 @@ import (
"bufio"
"encoding/json"
"fmt"
"github.com/songquanpeng/one-api/common/render"
"io"
"net/http"
"strings"

"github.com/songquanpeng/one-api/common/render"

"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper"
Expand All @@ -28,6 +29,11 @@ const (
VisionMaxImageNum = 16
)

var mimeTypeMap = map[string]string{
"json_object": "application/json",
"text": "text/plain",
}

// Setting safety to the lowest possible values since Gemini is already powerless enough
func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
geminiRequest := ChatRequest{
Expand Down Expand Up @@ -56,6 +62,15 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
MaxOutputTokens: textRequest.MaxTokens,
},
}
if textRequest.ResponseFormat != nil {
if mimeType, ok := mimeTypeMap[textRequest.ResponseFormat.Type]; ok {
geminiRequest.GenerationConfig.ResponseMimeType = mimeType
}
if textRequest.ResponseFormat.JsonSchema != nil {
geminiRequest.GenerationConfig.ResponseSchema = textRequest.ResponseFormat.JsonSchema.Schema
geminiRequest.GenerationConfig.ResponseMimeType = mimeTypeMap["json_object"]
}
}
if textRequest.Tools != nil {
functions := make([]model.Function, 0, len(textRequest.Tools))
for _, tool := range textRequest.Tools {
Expand Down
14 changes: 8 additions & 6 deletions relay/adaptor/gemini/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,12 @@ type ChatTools struct {
}

type ChatGenerationConfig struct {
Temperature float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"`
TopK float64 `json:"topK,omitempty"`
MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
CandidateCount int `json:"candidateCount,omitempty"`
StopSequences []string `json:"stopSequences,omitempty"`
ResponseMimeType string `json:"responseMimeType,omitempty"`
ResponseSchema any `json:"responseSchema,omitempty"`
Temperature float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"`
TopK float64 `json:"topK,omitempty"`
MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
CandidateCount int `json:"candidateCount,omitempty"`
StopSequences []string `json:"stopSequences,omitempty"`
}

0 comments on commit fd2846d

Please sign in to comment.