Skip to content

Commit e504665

Browse files
committed
feat: 优化Gemini模型版本获取逻辑
1 parent 54657ec commit e504665

4 files changed

Lines changed: 35 additions & 26 deletions

File tree

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@
6464
- `GET_MEDIA_TOKEN`:是统计图片token,默认为 `true`,关闭后将不再在本地计算图片token,可能会导致和上游计费不同,此项覆盖 `GET_MEDIA_TOKEN_NOT_STREAM` 选项作用。
6565
- `GET_MEDIA_TOKEN_NOT_STREAM`:是否在非流(`stream=false`)情况下统计图片token,默认为 `true`。
6666
- `UPDATE_TASK`:是否更新异步任务(Midjourney、Suno),默认为 `true`,关闭后将不会更新任务进度。
67-
- `GEMINI_MODEL_API`:Gemini模型指定版本(v1/v1beta),如果配置会覆盖默认配置需要完整给出全部v1beta模型,使用模型:版本指定,","分隔,例如:-e GEMINI_MODEL_API="gemini-1.5-pro-latest:v1beta,gemini-1.5-pro-001:v1beta",为空则使用默认配置:gemini-1.5-pro-latest,gemini-1.5-pro-001,gemini-1.5-pro,gemini-1.5-pro-exp-0801,gemini-1.5-flash-latest,gemini-1.5-flash-001,gemini-1.5-flash,gemini-ultra模型为v1beta,其他为v1。
67+
- `GEMINI_MODEL_MAP`:Gemini模型指定版本(v1/v1beta),使用模型:版本指定,","分隔,例如:-e GEMINI_MODEL_API="gemini-1.5-pro-latest:v1beta,gemini-1.5-pro-001:v1beta",为空则使用默认配置
6868
6969
## 部署
7070
### 部署要求

constant/env.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
package constant
22

33
import (
4+
"fmt"
45
"one-api/common"
6+
"os"
7+
"strings"
58
)
69

710
var StreamingTimeout = common.GetEnvOrDefault("STREAMING_TIMEOUT", 30)
@@ -15,3 +18,29 @@ var GetMediaToken = common.GetEnvOrDefaultBool("GET_MEDIA_TOKEN", true)
1518
var GetMediaTokenNotStream = common.GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", true)
1619

1720
var UpdateTask = common.GetEnvOrDefaultBool("UPDATE_TASK", true)
21+
22+
var GeminiModelMap = map[string]string{
23+
"gemini-1.5-pro-latest": "v1beta",
24+
"gemini-1.5-pro-001": "v1beta",
25+
"gemini-1.5-pro": "v1beta",
26+
"gemini-1.5-pro-exp-0801": "v1beta",
27+
"gemini-1.5-flash-latest": "v1beta",
28+
"gemini-1.5-flash-001": "v1beta",
29+
"gemini-1.5-flash": "v1beta",
30+
"gemini-ultra": "v1beta",
31+
}
32+
33+
func InitEnv() {
34+
modelVersionMapStr := strings.TrimSpace(os.Getenv("GEMINI_MODEL_MAP"))
35+
if modelVersionMapStr == "" {
36+
return
37+
}
38+
for _, pair := range strings.Split(modelVersionMapStr, ",") {
39+
parts := strings.Split(pair, ":")
40+
if len(parts) == 2 {
41+
GeminiModelMap[parts[0]] = parts[1]
42+
} else {
43+
common.SysError(fmt.Sprintf("invalid model version map: %s", pair))
44+
}
45+
}
46+
}

main.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ func main() {
5555
common.FatalLog("failed to initialize Redis: " + err.Error())
5656
}
5757

58+
// Initialize constants
59+
constant.InitEnv()
5860
// Initialize options
5961
model.InitOptionMap()
6062
if common.RedisEnabled {

relay/channel/gemini/adaptor.go

Lines changed: 3 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,13 @@ import (
66
"github.com/gin-gonic/gin"
77
"io"
88
"net/http"
9-
"os"
9+
"one-api/constant"
1010
"one-api/dto"
1111
"one-api/relay/channel"
12-
"strings"
1312
relaycommon "one-api/relay/common"
1413
)
1514

1615
type Adaptor struct {
17-
modelVersionMap map[string]string
1816
}
1917

2018
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
@@ -28,32 +26,12 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
2826
}
2927

3028
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
31-
modelVersionMapStr := os.Getenv("GEMINI_MODEL_API")
32-
if modelVersionMapStr == "" {
33-
a.modelVersionMap = map[string]string{
34-
"gemini-1.5-pro-latest": "v1beta",
35-
"gemini-1.5-pro-001": "v1beta",
36-
"gemini-1.5-pro": "v1beta",
37-
"gemini-1.5-pro-exp-0801": "v1beta",
38-
"gemini-1.5-flash-latest": "v1beta",
39-
"gemini-1.5-flash-001": "v1beta",
40-
"gemini-1.5-flash": "v1beta",
41-
"gemini-ultra": "v1beta",
42-
}
43-
return
44-
}
45-
a.modelVersionMap = make(map[string]string)
46-
for _, pair := range strings.Split(modelVersionMapStr, ",") {
47-
parts := strings.Split(pair, ":")
48-
if len(parts) == 2 {
49-
a.modelVersionMap[parts[0]] = parts[1]
50-
}
51-
}
29+
5230
}
5331

5432
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
5533
// 从映射中获取模型名称对应的版本,如果找不到就使用 info.ApiVersion 或默认的版本 "v1"
56-
version, beta := a.modelVersionMap[info.UpstreamModelName]
34+
version, beta := constant.GeminiModelMap[info.UpstreamModelName]
5735
if !beta {
5836
if info.ApiVersion != "" {
5937
version = info.ApiVersion

0 commit comments

Comments
 (0)