Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Update the Gemini API call method

Browse Source
This commit is contained in:
devocy 2024-02-26 15:47:52 +08:00
parent 7038d2a71b
commit 8bab658cf8
2 changed files with 17 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
relay/channel/gemini/adaptor.go
View File

@ -12,23 +12,27 @@ import (
"github.com/songquanpeng/one-api/relay/util" "github.com/songquanpeng/one-api/relay/util"
"io" "io"
"net/http" "net/http"
"net/url"
) )
type Adaptor struct { type Adaptor struct {
} }
func (a *Adaptor) Init(meta *util.RelayMeta) { func (a *Adaptor) Init(meta *util.RelayMeta) {}
fmt.Println(meta.APIVersion)
}
func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) {
version := helper.AssignOrDefault(meta.APIVersion, "v1") version := helper.AssignOrDefault(meta.APIVersion, "v1")
action := "generateContent" var action string
if relaymode.RelayModeEmbeddings == meta.Mode { switch meta.Mode {
case relaymode.RelayModeEmbeddings:
action = "batchEmbedContents" action = "batchEmbedContents"
} else if meta.IsStream { default:
action = "streamGenerateContent" if meta.IsStream {
action = "streamGenerateContent"
} else {
action = "generateContent"
}
} }
return fmt.Sprintf("%s/%s/models/%s:%s", meta.BaseURL, version, meta.ActualModelName, action), nil return fmt.Sprintf("%s/%s/models/%s:%s", meta.BaseURL, version, meta.ActualModelName, action), nil
@ -36,7 +40,11 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) {
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error {
channelhelper.SetupCommonRequestHeader(c, req, meta) channelhelper.SetupCommonRequestHeader(c, req, meta)
req.Header.Set("x-goog-api-key", meta.APIKey)
req.URL.RawQuery = url.Values{
"key": {meta.APIKey},
}.Encode()
return nil return nil
} }

2
relay/channel/gemini/constants.go
View File

@ -3,5 +3,5 @@ package gemini
var ModelList = []string{ var ModelList = []string{
"gemini-pro", "gemini-pro",
"gemini-pro-vision", "gemini-pro-vision",
"(Gemini)embedding-001", "embedding-001",
} }