feat: 暂时禁用透传功能

Modify the default gemini API to v1beta

.env.example

@@ -0,0 +1,71 @@
+# 端口号
+# PORT=3000
+# 前端基础URL
+# FRONTEND_BASE_URL=https://your-frontend-url.com
+
+
+# 调试相关配置
+# 启用pprof
+# ENABLE_PPROF=true
+
+# 数据库相关配置
+# 数据库连接字符串
+# SQL_DSN=mysql://user:password@tcp(127.0.0.1:3306)/dbname?parseTime=true
+# 日志数据库连接字符串
+# LOG_SQL_DSN=mysql://user:password@tcp(127.0.0.1:3306)/logdb?parseTime=true
+# SQLite数据库路径
+# SQLITE_PATH=/path/to/sqlite.db
+# 数据库最大空闲连接数
+# SQL_MAX_IDLE_CONNS=100
+# 数据库最大打开连接数
+# SQL_MAX_OPEN_CONNS=1000
+# 数据库连接最大生命周期（秒）
+# SQL_MAX_LIFETIME=60
+
+
+# 缓存相关配置
+# Redis连接字符串
+# REDIS_CONN_STRING=redis://user:password@localhost:6379/0
+# 同步频率（单位：秒）
+# SYNC_FREQUENCY=60
+# 内存缓存启用
+# MEMORY_CACHE_ENABLED=true
+# 渠道更新频率（单位：秒）
+# CHANNEL_UPDATE_FREQUENCY=30
+# 批量更新启用
+# BATCH_UPDATE_ENABLED=true
+# 批量更新间隔（单位：秒）
+# BATCH_UPDATE_INTERVAL=5
+
+# 任务和功能配置
+# 更新任务启用
+# UPDATE_TASK=true
+
+
+# 会话密钥
+# SESSION_SECRET=random_string
+
+# 其他配置
+# 渠道测试频率（单位：秒）
+# CHANNEL_TEST_FREQUENCY=10
+# 生成默认token
+# GENERATE_DEFAULT_TOKEN=false
+# Gemini 安全设置
+# GEMINI_SAFETY_SETTING=BLOCK_NONE
+# Gemini版本设置
+# GEMINI_MODEL_MAP=gemini-1.0-pro:v1
+# Cohere 安全设置
+# COHERE_SAFETY_SETTING=NONE
+# 是否统计图片token
+# GET_MEDIA_TOKEN=true
+# 是否在非流（stream=false）情况下统计图片token
+# GET_MEDIA_TOKEN_NOT_STREAM=true
+# 设置 Dify 渠道是否输出工作流和节点信息到客户端
+# DIFY_DEBUG=true
+# 设置流式一次回复的超时时间
+# STREAMING_TIMEOUT=90
+
+
+# 节点类型
+# 如果是主节点则为master
+# NODE_TYPE=master

.github/FUNDING.yml

@@ -0,0 +1,12 @@
+# These are supported funding model platforms
+
+github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
+patreon: # Replace with a single Patreon username
+open_collective: # Replace with a single Open Collective username
+ko_fi: # Replace with a single Ko-fi username
+tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
+community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
+liberapay: # Replace with a single Liberapay username
+issuehunt: # Replace with a single IssueHunt username
+otechie: # Replace with a single Otechie username
+custom: ['https://afdian.com/a/new-api'] # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']

.gitignore

@@ -6,4 +6,5 @@ upload
 build
 *.db-journal
 logs
-web/dist
+web/dist
+.env

README.md

@@ -1,6 +1,13 @@
+<div align="center">
+
+![new-api](/web/public/logo.png)
 
 # New API
 
+<a href="https://trendshift.io/repositories/8227" target="_blank"><img src="https://trendshift.io/api/badge/repositories/8227" alt="Calcium-Ion%2Fnew-api | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
+
+</div>
+
 > [!NOTE]
 > 本项目为开源项目，在[One API](https://github.com/songquanpeng/one-api)的基础上进行二次开发
 
@@ -41,6 +48,7 @@
     4. Telegram Bot 名称是bot username 去掉@后的字符串
 13. 添加 [Suno API](https://github.com/Suno-API/Suno-API)接口的支持，[对接文档](Suno.md)
 14. 支持Rerank模型，目前仅兼容Cohere和Jina，可接入Dify，[对接文档](Rerank.md)
+15. **[OpenAI Realtime API](https://platform.openai.com/docs/guides/realtime/integration)** - 支持OpenAI的Realtime API，支持Azure渠道。
 
 ## 模型支持
 此版本额外支持以下模型：
@@ -59,13 +67,14 @@
 您可以在渠道中添加自定义模型gpt-4-gizmo-*，此模型并非OpenAI官方模型，而是第三方模型，使用官方key无法调用。
 
 ## 比原版One API多出的配置
+- `GENERATE_DEFAULT_TOKEN`：是否为新注册用户生成初始令牌，默认为 `false`。
 - `STREAMING_TIMEOUT`：设置流式一次回复的超时时间，默认为 30 秒。
 - `DIFY_DEBUG`：设置 Dify 渠道是否输出工作流和节点信息到客户端，默认为 `true`。
 - `FORCE_STREAM_OPTION`：是否覆盖客户端stream_options参数，请求上游返回流模式usage，默认为 `true`，建议开启，不影响客户端传入stream_options参数返回结果。
 - `GET_MEDIA_TOKEN`：是统计图片token，默认为 `true`，关闭后将不再在本地计算图片token，可能会导致和上游计费不同，此项覆盖 `GET_MEDIA_TOKEN_NOT_STREAM` 选项作用。
 - `GET_MEDIA_TOKEN_NOT_STREAM`：是否在非流（`stream=false`）情况下统计图片token，默认为 `true`。
 - `UPDATE_TASK`：是否更新异步任务（Midjourney、Suno），默认为 `true`，关闭后将不会更新任务进度。
-- `GEMINI_MODEL_MAP`：Gemini模型指定版本(v1/v1beta)，使用“模型:版本”指定，","分隔，例如：-e GEMINI_MODEL_MAP="gemini-1.5-pro-latest:v1beta,gemini-1.5-pro-001:v1beta"，为空则使用默认配置
+- `GEMINI_MODEL_MAP`：Gemini模型指定版本(v1/v1beta)，使用“模型:版本”指定，","分隔，例如：-e GEMINI_MODEL_MAP="gemini-1.5-pro-latest:v1beta,gemini-1.5-pro-001:v1beta"，为空则使用默认配置(v1beta)
 - `COHERE_SAFETY_SETTING`：Cohere模型[安全设置](https://docs.cohere.com/docs/safety-modes#overview)，可选值为 `NONE`, `CONTEXTUAL`，`STRICT`，默认为 `NONE`。
 ## 部署
 ### 部署要求
@@ -115,24 +124,19 @@ docker run --name new-api -d --restart always -p 3000:3000 -e SQL_DSN="root:1234
 ## Suno接口设置文档
 [对接文档](Suno.md)
 
-## 交流群
-<img src="https://github.com/Calcium-Ion/new-api/assets/61247483/de536a8a-0161-47a7-a0a2-66ef6de81266" width="300">
-
 ## 界面截图
+![796df8d287b7b7bd7853b2497e7df511](https://github.com/user-attachments/assets/255b5e97-2d3a-4434-b4fa-e922ad88ff5a)
+
 ![image](https://github.com/Calcium-Ion/new-api/assets/61247483/ad0e7aae-0203-471c-9716-2d83768927d4)
 
-![image](https://github.com/Calcium-Ion/new-api/assets/61247483/d1ac216e-0804-4105-9fdc-66b35022d861)
-
-![image](https://github.com/Calcium-Ion/new-api/assets/61247483/3ca0b282-00ff-4c96-bf9d-e29ef615c605)  
-![image](https://github.com/Calcium-Ion/new-api/assets/61247483/f4f40ed4-8ccb-43d7-a580-90677827646d)  
-![image](https://github.com/Calcium-Ion/new-api/assets/61247483/90d7d763-6a77-4b36-9f76-2bb30f18583d)
-![image](https://github.com/Calcium-Ion/new-api/assets/61247483/e414228a-3c35-429a-b298-6451d76d9032)
+![image](https://github.com/Calcium-Ion/new-api/assets/61247483/3ca0b282-00ff-4c96-bf9d-e29ef615c605)
 夜间模式  
 ![image](https://github.com/Calcium-Ion/new-api/assets/61247483/1c66b593-bb9e-4757-9720-ff2759539242)
-
-![image](https://github.com/Calcium-Ion/new-api/assets/61247483/5b3228e8-2556-44f7-97d6-4f8d8ee6effa)  
 ![image](https://github.com/Calcium-Ion/new-api/assets/61247483/af9a07ee-5101-4b3d-8bd9-ae21a4fd7e9e)
 
+## 交流群
+<img src="https://github.com/Calcium-Ion/new-api/assets/61247483/de536a8a-0161-47a7-a0a2-66ef6de81266" width="200">
+
 ## 相关项目
 - [One API](https://github.com/songquanpeng/one-api)：原版项目
 - [Midjourney-Proxy](https://github.com/novicezk/midjourney-proxy)：Midjourney接口支持

common/constants.go

@@ -41,6 +41,7 @@ var PasswordLoginEnabled = true
 var PasswordRegisterEnabled = true
 var EmailVerificationEnabled = false
 var GitHubOAuthEnabled = false
+var LinuxDOOAuthEnabled = false
 var WeChatAuthEnabled = false
 var TelegramOAuthEnabled = false
 var TurnstileCheckEnabled = false
@@ -75,6 +76,9 @@ var SMTPToken = ""
 var GitHubClientId = ""
 var GitHubClientSecret = ""
 
+var LinuxDOClientId = ""
+var LinuxDOClientSecret = ""
+
 var WeChatServerAddress = ""
 var WeChatServerToken = ""
 var WeChatAccountQRCodeImageURL = ""
@@ -126,6 +130,10 @@ const (
 	RoleRootUser   = 100
 )
 
+func IsValidateRole(role int) bool {
+	return role == RoleGuestUser || role == RoleCommonUser || role == RoleAdminUser || role == RoleRootUser
+}
+
 var (
 	FileUploadPermission    = RoleGuestUser
 	FileDownloadPermission  = RoleGuestUser
@@ -136,11 +144,13 @@ var (
 // All duration's unit is seconds
 // Shouldn't larger then RateLimitKeyExpirationDuration
 var (
-	GlobalApiRateLimitNum            = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 180)
-	GlobalApiRateLimitDuration int64 = 3 * 60
+	GlobalApiRateLimitEnable   = GetEnvOrDefaultBool("GLOBAL_API_RATE_LIMIT_ENABLE", true)
+	GlobalApiRateLimitNum      = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 180)
+	GlobalApiRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_API_RATE_LIMIT_DURATION", 180))
 
-	GlobalWebRateLimitNum            = GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT", 60)
-	GlobalWebRateLimitDuration int64 = 3 * 60
+	GlobalWebRateLimitEnable   = GetEnvOrDefaultBool("GLOBAL_WEB_RATE_LIMIT_ENABLE", true)
+	GlobalWebRateLimitNum      = GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT", 60)
+	GlobalWebRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT_DURATION", 180))
 
 	UploadRateLimitNum            = 10
 	UploadRateLimitDuration int64 = 60
@@ -218,6 +228,7 @@ const (
 	ChannelCloudflare         = 39
 	ChannelTypeSiliconFlow    = 40
 	ChannelTypeVertexAi       = 41
+	ChannelTypeMistral        = 42
 
 	ChannelTypeDummy // this one is only for count, do not add any channel after this
 
@@ -266,4 +277,5 @@ var ChannelBaseURLs = []string{
 	"https://api.cloudflare.com",                //39
 	"https://api.siliconflow.cn",                //40
 	"",                                          //41
+	"https://api.mistral.ai",                    //42
 }

common/model-ratio.go

@@ -32,26 +32,29 @@ var defaultModelRatio = map[string]float64{
 	"gpt-4-0613": 15,
 	"gpt-4-32k":  30,
 	//"gpt-4-32k-0314":               30, //deprecated
-	"gpt-4-32k-0613":            30,
-	"gpt-4-1106-preview":        5,    // $0.01 / 1K tokens
-	"gpt-4-0125-preview":        5,    // $0.01 / 1K tokens
-	"gpt-4-turbo-preview":       5,    // $0.01 / 1K tokens
-	"gpt-4-vision-preview":      5,    // $0.01 / 1K tokens
-	"gpt-4-1106-vision-preview": 5,    // $0.01 / 1K tokens
-	"chatgpt-4o-latest":         2.5,  // $0.01 / 1K tokens
-	"gpt-4o":                    2.5,  // $0.01 / 1K tokens
-	"gpt-4o-2024-05-13":         2.5,  // $0.01 / 1K tokens
-	"gpt-4o-2024-08-06":         1.25, // $0.01 / 1K tokens
-	"o1-preview":                7.5,
-	"o1-preview-2024-09-12":     7.5,
-	"o1-mini":                   1.5,
-	"o1-mini-2024-09-12":        1.5,
-	"gpt-4o-mini":               0.075,
-	"gpt-4o-mini-2024-07-18":    0.075,
-	"gpt-4-turbo":               5,    // $0.01 / 1K tokens
-	"gpt-4-turbo-2024-04-09":    5,    // $0.01 / 1K tokens
-	"gpt-3.5-turbo":             0.25, // $0.0015 / 1K tokens
+	"gpt-4-32k-0613":                  30,
+	"gpt-4-1106-preview":              5,    // $0.01 / 1K tokens
+	"gpt-4-0125-preview":              5,    // $0.01 / 1K tokens
+	"gpt-4-turbo-preview":             5,    // $0.01 / 1K tokens
+	"gpt-4-vision-preview":            5,    // $0.01 / 1K tokens
+	"gpt-4-1106-vision-preview":       5,    // $0.01 / 1K tokens
+	"chatgpt-4o-latest":               2.5,  // $0.01 / 1K tokens
+	"gpt-4o":                          1.25, // $0.01 / 1K tokens
+	"gpt-4o-audio-preview":            1.25, // $0.0015 / 1K tokens
+	"gpt-4o-audio-preview-2024-10-01": 1.25, // $0.0015 / 1K tokens
+	"gpt-4o-2024-08-06":               1.25, // $0.01 / 1K tokens
+	"gpt-4o-2024-05-13":               2.5,
+	"gpt-4o-realtime-preview":         2.5,
+	"o1-preview":                      7.5,
+	"o1-preview-2024-09-12":           7.5,
+	"o1-mini":                         1.5,
+	"o1-mini-2024-09-12":              1.5,
+	"gpt-4o-mini":                     0.075,
+	"gpt-4o-mini-2024-07-18":          0.075,
+	"gpt-4-turbo":                     5, // $0.01 / 1K tokens
+	"gpt-4-turbo-2024-04-09":          5, // $0.01 / 1K tokens
 	//"gpt-3.5-turbo-0301":           0.75, //deprecated
+	"gpt-3.5-turbo":          0.25,
 	"gpt-3.5-turbo-0613":     0.75,
 	"gpt-3.5-turbo-16k":      1.5, // $0.003 / 1K tokens
 	"gpt-3.5-turbo-16k-0613": 1.5,
@@ -86,8 +89,10 @@ var defaultModelRatio = map[string]float64{
 	"claude-2.0":                     4,     // $8 / 1M tokens
 	"claude-2.1":                     4,     // $8 / 1M tokens
 	"claude-3-haiku-20240307":        0.125, // $0.25 / 1M tokens
+	"claude-3-5-haiku-20241022":      0.5,   // $1 / 1M tokens
 	"claude-3-sonnet-20240229":       1.5,   // $3 / 1M tokens
 	"claude-3-5-sonnet-20240620":     1.5,
+	"claude-3-5-sonnet-20241022":     1.5,
 	"claude-3-opus-20240229":         7.5, // $15 / 1M tokens
 	"ERNIE-4.0-8K":                   0.120 * RMB,
 	"ERNIE-3.5-8K":                   0.012 * RMB,
@@ -336,13 +341,13 @@ func GetCompletionRatio(name string) float64 {
 		name = "gpt-4o-gizmo-*"
 	}
 	if strings.HasPrefix(name, "gpt-4") && !strings.HasSuffix(name, "-all") && !strings.HasSuffix(name, "-gizmo-*") {
-		if strings.HasPrefix(name, "gpt-4-turbo") || strings.HasSuffix(name, "preview") {
-			return 3
-		}
 		if strings.HasPrefix(name, "gpt-4o") {
-			if strings.HasPrefix(name, "gpt-4o-mini") || name == "gpt-4o-2024-08-06" {
-				return 4
+			if name == "gpt-4o-2024-05-13" {
+				return 3
 			}
+			return 4
+		}
+		if strings.HasPrefix(name, "gpt-4-turbo") || strings.HasSuffix(name, "preview") {
 			return 3
 		}
 		return 2
@@ -375,7 +380,7 @@ func GetCompletionRatio(name string) float64 {
 		return 3
 	}
 	if strings.HasPrefix(name, "gemini-") {
-		return 3
+		return 4
 	}
 	if strings.HasPrefix(name, "command") {
 		switch name {
@@ -417,6 +422,36 @@ func GetCompletionRatio(name string) float64 {
 	return 1
 }
 
+func GetAudioRatio(name string) float64 {
+	if strings.HasPrefix(name, "gpt-4o-realtime") {
+		return 20
+	} else if strings.HasPrefix(name, "gpt-4o-audio") {
+		return 40
+	}
+	return 20
+}
+
+func GetAudioCompletionRatio(name string) float64 {
+	if strings.HasPrefix(name, "gpt-4o-realtime") {
+		return 2
+	}
+	return 2
+}
+
+//func GetAudioPricePerMinute(name string) float64 {
+//	if strings.HasPrefix(name, "gpt-4o-realtime") {
+//		return 0.06
+//	}
+//	return 0.06
+//}
+//
+//func GetAudioCompletionPricePerMinute(name string) float64 {
+//	if strings.HasPrefix(name, "gpt-4o-realtime") {
+//		return 0.24
+//	}
+//	return 0.24
+//}
+
 func GetCompletionRatioMap() map[string]float64 {
 	if CompletionRatio == nil {
 		CompletionRatio = defaultCompletionRatio

common/user_groups.go

@@ -0,0 +1,46 @@
+package common
+
+import (
+	"encoding/json"
+)
+
+var UserUsableGroups = map[string]string{
+	"default": "默认分组",
+	"vip":     "vip分组",
+}
+
+func UserUsableGroups2JSONString() string {
+	jsonBytes, err := json.Marshal(UserUsableGroups)
+	if err != nil {
+		SysError("error marshalling user groups: " + err.Error())
+	}
+	return string(jsonBytes)
+}
+
+func UpdateUserUsableGroupsByJSONString(jsonStr string) error {
+	UserUsableGroups = make(map[string]string)
+	return json.Unmarshal([]byte(jsonStr), &UserUsableGroups)
+}
+
+func GetUserUsableGroups(userGroup string) map[string]string {
+	if userGroup == "" {
+		// 如果userGroup为空，返回UserUsableGroups
+		return UserUsableGroups
+	}
+	// 如果userGroup不在UserUsableGroups中，返回UserUsableGroups + userGroup
+	if _, ok := UserUsableGroups[userGroup]; !ok {
+		appendUserUsableGroups := make(map[string]string)
+		for k, v := range UserUsableGroups {
+			appendUserUsableGroups[k] = v
+		}
+		appendUserUsableGroups[userGroup] = "用户分组"
+		return appendUserUsableGroups
+	}
+	// 如果userGroup在UserUsableGroups中，返回UserUsableGroups
+	return UserUsableGroups
+}
+
+func GroupInUserUsableGroups(groupName string) bool {
+	_, ok := UserUsableGroups[groupName]
+	return ok
+}

common/utils.go

@@ -1,10 +1,13 @@
 package common
 
 import (
+	crand "crypto/rand"
+	"encoding/base64"
 	"fmt"
 	"github.com/google/uuid"
 	"html/template"
 	"log"
+	"math/big"
 	"math/rand"
 	"net"
 	"os/exec"
@@ -128,6 +131,11 @@ func IntMax(a int, b int) int {
 	}
 }
 
+func IsIP(s string) bool {
+	ip := net.ParseIP(s)
+	return ip != nil
+}
+
 func GetUUID() string {
 	code := uuid.New().String()
 	code = strings.Replace(code, "-", "", -1)
@@ -137,24 +145,35 @@ func GetUUID() string {
 const keyChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
 
 func init() {
-	rand.Seed(time.Now().UnixNano())
+	rand.New(rand.NewSource(time.Now().UnixNano()))
 }
 
-func GenerateKey() string {
-	//rand.Seed(time.Now().UnixNano())
-	key := make([]byte, 48)
-	for i := 0; i < 16; i++ {
-		key[i] = keyChars[rand.Intn(len(keyChars))]
-	}
-	uuid_ := GetUUID()
-	for i := 0; i < 32; i++ {
-		c := uuid_[i]
-		if i%2 == 0 && c >= 'a' && c <= 'z' {
-			c = c - 'a' + 'A'
+func GenerateRandomCharsKey(length int) (string, error) {
+	b := make([]byte, length)
+	maxI := big.NewInt(int64(len(keyChars)))
+
+	for i := range b {
+		n, err := crand.Int(crand.Reader, maxI)
+		if err != nil {
+			return "", err
 		}
-		key[i+16] = c
+		b[i] = keyChars[n.Int64()]
 	}
-	return string(key)
+
+	return string(b), nil
+}
+
+func GenerateRandomKey(length int) (string, error) {
+	bytes := make([]byte, length*3/4) // 对于48位的输出，这里应该是36
+	if _, err := crand.Read(bytes); err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(bytes), nil
+}
+
+func GenerateKey() (string, error) {
+	//rand.Seed(time.Now().UnixNano())
+	return GenerateRandomCharsKey(48)
 }
 
 func GetRandomInt(max int) int {

constant/chat.go

@@ -0,0 +1,35 @@
+package constant
+
+import (
+	"encoding/json"
+	"one-api/common"
+)
+
+var Chats = []map[string]string{
+	{
+		"ChatGPT Next Web 官方示例": "https://app.nextchat.dev/#/?settings={\"key\":\"{key}\",\"url\":\"{address}\"}",
+	},
+	{
+		"Lobe Chat 官方示例": "https://chat-preview.lobehub.com/?settings={\"keyVaults\":{\"openai\":{\"apiKey\":\"{key}\",\"baseURL\":\"{address}/v1\"}}}",
+	},
+	{
+		"AMA 问天": "ama://set-api-key?server={address}&key={key}",
+	},
+	{
+		"OpenCat": "opencat://team/join?domain={address}&token={key}",
+	},
+}
+
+func UpdateChatsByJsonString(jsonString string) error {
+	Chats = make([]map[string]string, 0)
+	return json.Unmarshal([]byte(jsonString), &Chats)
+}
+
+func Chats2JsonString() string {
+	jsonBytes, err := json.Marshal(Chats)
+	if err != nil {
+		common.SysError("error marshalling chats: " + err.Error())
+		return "[]"
+	}
+	return string(jsonBytes)
+}

constant/env.go

@@ -20,16 +20,7 @@ var GetMediaTokenNotStream = common.GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STR
 var UpdateTask = common.GetEnvOrDefaultBool("UPDATE_TASK", true)
 
 var GeminiModelMap = map[string]string{
-	"gemini-1.5-pro-latest":     "v1beta",
-	"gemini-1.5-pro-001":        "v1beta",
-	"gemini-1.5-pro":            "v1beta",
-	"gemini-1.5-pro-exp-0801":   "v1beta",
-	"gemini-1.5-pro-exp-0827":   "v1beta",
-	"gemini-1.5-flash-latest":   "v1beta",
-	"gemini-1.5-flash-exp-0827": "v1beta",
-	"gemini-1.5-flash-001":      "v1beta",
-	"gemini-1.5-flash":          "v1beta",
-	"gemini-ultra":              "v1beta",
+	"gemini-1.0-pro":     "v1",
 }
 
 func InitEnv() {
@@ -46,3 +37,6 @@ func InitEnv() {
 		}
 	}
 }
+
+// 是否生成初始令牌，默认关闭。
+var GenerateDefaultToken = common.GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false)

controller/channel-test.go

@@ -20,6 +20,7 @@ import (
 	"one-api/relay/constant"
 	"one-api/service"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -81,8 +82,7 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr
 		return fmt.Errorf("invalid api type: %d, adaptor is nil", apiType), nil
 	}
 
-	request := buildTestRequest()
-	request.Model = testModel
+	request := buildTestRequest(testModel)
 	meta.UpstreamModelName = testModel
 	common.SysLog(fmt.Sprintf("testing channel %d with model %s", channel.Id, testModel))
 
@@ -102,17 +102,22 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr
 	if err != nil {
 		return err, nil
 	}
-	if resp != nil && resp.StatusCode != http.StatusOK {
-		err := service.RelayErrorHandler(resp)
-		return fmt.Errorf("status code %d: %s", resp.StatusCode, err.Error.Message), err
+	var httpResp *http.Response
+	if resp != nil {
+		httpResp = resp.(*http.Response)
+		if httpResp.StatusCode != http.StatusOK {
+			err := service.RelayErrorHandler(httpResp)
+			return fmt.Errorf("status code %d: %s", httpResp.StatusCode, err.Error.Message), err
+		}
 	}
-	usage, respErr := adaptor.DoResponse(c, resp, meta)
+	usageA, respErr := adaptor.DoResponse(c, httpResp, meta)
 	if respErr != nil {
 		return fmt.Errorf("%s", respErr.Error.Message), respErr
 	}
-	if usage == nil {
+	if usageA == nil {
 		return errors.New("usage is nil"), nil
 	}
+	usage := usageA.(*dto.Usage)
 	result := w.Result()
 	respBody, err := io.ReadAll(result.Body)
 	if err != nil {
@@ -141,17 +146,22 @@ func testChannel(channel *model.Channel, testModel string) (err error, openAIErr
 	return nil, nil
 }
 
-func buildTestRequest() *dto.GeneralOpenAIRequest {
+func buildTestRequest(model string) *dto.GeneralOpenAIRequest {
 	testRequest := &dto.GeneralOpenAIRequest{
-		Model:     "", // this will be set later
-		MaxTokens: 1,
-		Stream:    false,
+		Model:  "", // this will be set later
+		Stream: false,
+	}
+	if strings.HasPrefix(model, "o1-") {
+		testRequest.MaxCompletionTokens = 1
+	} else {
+		testRequest.MaxTokens = 1
 	}
 	content, _ := json.Marshal("hi")
 	testMessage := dto.Message{
 		Role:    "user",
 		Content: content,
 	}
+	testRequest.Model = model
 	testRequest.Messages = append(testRequest.Messages, testMessage)
 	return testRequest
 }
@@ -226,26 +236,22 @@ func testAllChannels(notify bool) error {
 			tok := time.Now()
 			milliseconds := tok.Sub(tik).Milliseconds()
 
-			ban := false
-			if milliseconds > disableThreshold {
-				err = errors.New(fmt.Sprintf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
-				ban = true
-			}
+			shouldBanChannel := false
 
 			// request error disables the channel
 			if openaiWithStatusErr != nil {
 				oaiErr := openaiWithStatusErr.Error
 				err = errors.New(fmt.Sprintf("type %s, httpCode %d, code %v, message %s", oaiErr.Type, openaiWithStatusErr.StatusCode, oaiErr.Code, oaiErr.Message))
-				ban = service.ShouldDisableChannel(channel.Type, openaiWithStatusErr)
+				shouldBanChannel = service.ShouldDisableChannel(channel.Type, openaiWithStatusErr)
 			}
 
-			// parse *int to bool
-			if !channel.GetAutoBan() {
-				ban = false
+			if milliseconds > disableThreshold {
+				err = errors.New(fmt.Sprintf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
+				shouldBanChannel = true
 			}
 
 			// disable channel
-			if ban && isChannelEnabled {
+			if isChannelEnabled && shouldBanChannel && channel.GetAutoBan() {
 				service.DisableChannel(channel.Id, channel.Name, err.Error())
 			}
 

controller/github.go

@@ -112,7 +112,9 @@ func GitHubOAuth(c *gin.Context) {
 	user := model.User{
 		GitHubId: githubUser.Login,
 	}
+	// IsGitHubIdAlreadyTaken is unscoped
 	if model.IsGitHubIdAlreadyTaken(user.GitHubId) {
+		// FillUserByGitHubId is scoped
 		err := user.FillUserByGitHubId()
 		if err != nil {
 			c.JSON(http.StatusOK, gin.H{
@@ -121,6 +123,14 @@ func GitHubOAuth(c *gin.Context) {
 			})
 			return
 		}
+		// if user.Id == 0 , user has been deleted
+		if user.Id == 0 {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "用户已注销",
+			})
+			return
+		}
 	} else {
 		if common.RegisterEnabled {
 			user.Username = "github_" + strconv.Itoa(model.GetMaxUserId()+1)
@@ -132,8 +142,13 @@ func GitHubOAuth(c *gin.Context) {
 			user.Email = githubUser.Email
 			user.Role = common.RoleCommonUser
 			user.Status = common.UserStatusEnabled
+			affCode := session.Get("aff")
+			inviterId := 0
+			if affCode != nil {
+				inviterId, _ = model.GetUserIdByAffCode(affCode.(string))
+			}
 
-			if err := user.Insert(0); err != nil {
+			if err := user.Insert(inviterId); err != nil {
 				c.JSON(http.StatusOK, gin.H{
 					"success": false,
 					"message": err.Error(),
@@ -217,6 +232,10 @@ func GitHubBind(c *gin.Context) {
 func GenerateOAuthCode(c *gin.Context) {
 	session := sessions.Default(c)
 	state := common.GetRandomString(12)
+	affCode := c.Query("aff")
+	if affCode != "" {
+		session.Set("aff", affCode)
+	}
 	session.Set("oauth_state", state)
 	err := session.Save()
 	if err != nil {

controller/group.go

@@ -4,6 +4,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"net/http"
 	"one-api/common"
+	"one-api/model"
 )
 
 func GetGroups(c *gin.Context) {
@@ -17,3 +18,22 @@ func GetGroups(c *gin.Context) {
 		"data":    groupNames,
 	})
 }
+
+func GetUserGroups(c *gin.Context) {
+	usableGroups := make(map[string]string)
+	userGroup := ""
+	userId := c.GetInt("id")
+	userGroup, _ = model.CacheGetUserGroup(userId)
+	for groupName, _ := range common.GroupRatio {
+		// UserUsableGroups contains the groups that the user can use
+		userUsableGroups := common.GetUserUsableGroups(userGroup)
+		if _, ok := userUsableGroups[groupName]; ok {
+			usableGroups[groupName] = userUsableGroups[groupName]
+		}
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    usableGroups,
+	})
+}

controller/linuxdo.go

@@ -0,0 +1,271 @@
+package controller
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"one-api/common"
+	"one-api/model"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/gin-contrib/sessions"
+	"github.com/gin-gonic/gin"
+)
+
+type LinuxdoUser struct {
+	Id         int    `json:"id"`
+	Username   string `json:"username"`
+	Name       string `json:"name"`
+	Active     bool   `json:"active"`
+	TrustLevel int    `json:"trust_level"`
+	Silenced   bool   `json:"silenced"`
+}
+
+func LinuxDoBind(c *gin.Context) {
+	if !common.LinuxDOOAuthEnabled {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "管理员未开启通过 Linux DO 登录以及注册",
+		})
+		return
+	}
+
+	code := c.Query("code")
+	linuxdoUser, err := getLinuxdoUserInfoByCode(code, c)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+
+	user := model.User{
+		LinuxDOId: strconv.Itoa(linuxdoUser.Id),
+	}
+
+	if model.IsLinuxDOIdAlreadyTaken(user.LinuxDOId) {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "该 Linux DO 账户已被绑定",
+		})
+		return
+	}
+
+	session := sessions.Default(c)
+	id := session.Get("id")
+	user.Id = id.(int)
+
+	err = user.FillUserById()
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+
+	user.LinuxDOId = strconv.Itoa(linuxdoUser.Id)
+	err = user.Update(false)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "bind",
+	})
+}
+
+func getLinuxdoUserInfoByCode(code string, c *gin.Context) (*LinuxdoUser, error) {
+	if code == "" {
+		return nil, errors.New("invalid code")
+	}
+
+	// Get access token using Basic auth
+	tokenEndpoint := "https://connect.linux.do/oauth2/token"
+	credentials := common.LinuxDOClientId + ":" + common.LinuxDOClientSecret
+	basicAuth := "Basic " + base64.StdEncoding.EncodeToString([]byte(credentials))
+
+	// Get redirect URI from request
+	scheme := "http"
+	if c.Request.TLS != nil {
+		scheme = "https"
+	}
+	redirectURI := fmt.Sprintf("%s://%s/api/oauth/linuxdo", scheme, c.Request.Host)
+
+	data := url.Values{}
+	data.Set("grant_type", "authorization_code")
+	data.Set("code", code)
+	data.Set("redirect_uri", redirectURI)
+
+	req, err := http.NewRequest("POST", tokenEndpoint, strings.NewReader(data.Encode()))
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Authorization", basicAuth)
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	req.Header.Set("Accept", "application/json")
+
+	client := http.Client{Timeout: 5 * time.Second}
+	res, err := client.Do(req)
+	if err != nil {
+		return nil, errors.New("failed to connect to Linux DO server")
+	}
+	defer res.Body.Close()
+
+	var tokenRes struct {
+		AccessToken string `json:"access_token"`
+		Message     string `json:"message"`
+	}
+	if err := json.NewDecoder(res.Body).Decode(&tokenRes); err != nil {
+		return nil, err
+	}
+
+	if tokenRes.AccessToken == "" {
+		return nil, fmt.Errorf("failed to get access token: %s", tokenRes.Message)
+	}
+
+	// Get user info
+	userEndpoint := "https://connect.linux.do/api/user"
+	req, err = http.NewRequest("GET", userEndpoint, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", "Bearer "+tokenRes.AccessToken)
+	req.Header.Set("Accept", "application/json")
+
+	res2, err := client.Do(req)
+	if err != nil {
+		return nil, errors.New("failed to get user info from Linux DO")
+	}
+	defer res2.Body.Close()
+
+	var linuxdoUser LinuxdoUser
+	if err := json.NewDecoder(res2.Body).Decode(&linuxdoUser); err != nil {
+		return nil, err
+	}
+
+	if linuxdoUser.Id == 0 {
+		return nil, errors.New("invalid user info returned")
+	}
+
+	return &linuxdoUser, nil
+}
+
+func LinuxdoOAuth(c *gin.Context) {
+	session := sessions.Default(c)
+
+	errorCode := c.Query("error")
+	if errorCode != "" {
+		errorDescription := c.Query("error_description")
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": errorDescription,
+		})
+		return
+	}
+
+	state := c.Query("state")
+	if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) {
+		c.JSON(http.StatusForbidden, gin.H{
+			"success": false,
+			"message": "state is empty or not same",
+		})
+		return
+	}
+
+	username := session.Get("username")
+	if username != nil {
+		LinuxDoBind(c)
+		return
+	}
+
+	if !common.LinuxDOOAuthEnabled {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "管理员未开启通过 Linux DO 登录以及注册",
+		})
+		return
+	}
+
+	code := c.Query("code")
+	linuxdoUser, err := getLinuxdoUserInfoByCode(code, c)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+
+	user := model.User{
+		LinuxDOId: strconv.Itoa(linuxdoUser.Id),
+	}
+
+	// Check if user exists
+	if model.IsLinuxDOIdAlreadyTaken(user.LinuxDOId) {
+		err := user.FillUserByLinuxDOId()
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": err.Error(),
+			})
+			return
+		}
+		if user.Id == 0 {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "用户已注销",
+			})
+			return
+		}
+	} else {
+		if common.RegisterEnabled {
+			user.Username = "linuxdo_" + strconv.Itoa(model.GetMaxUserId()+1)
+			user.DisplayName = linuxdoUser.Name
+			user.Role = common.RoleCommonUser
+			user.Status = common.UserStatusEnabled
+
+			affCode := session.Get("aff")
+			inviterId := 0
+			if affCode != nil {
+				inviterId, _ = model.GetUserIdByAffCode(affCode.(string))
+			}
+
+			if err := user.Insert(inviterId); err != nil {
+				c.JSON(http.StatusOK, gin.H{
+					"success": false,
+					"message": err.Error(),
+				})
+				return
+			}
+		} else {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "管理员关闭了新用户注册",
+			})
+			return
+		}
+	}
+
+	if user.Status != common.UserStatusEnabled {
+		c.JSON(http.StatusOK, gin.H{
+			"message": "用户已被封禁",
+			"success": false,
+		})
+		return
+	}
+
+	setupLogin(&user, c)
+}

controller/misc.go

@@ -38,6 +38,8 @@ func GetStatus(c *gin.Context) {
 			"email_verification":       common.EmailVerificationEnabled,
 			"github_oauth":             common.GitHubOAuthEnabled,
 			"github_client_id":         common.GitHubClientId,
+			"linuxdo_oauth":            common.LinuxDOOAuthEnabled,
+			"linuxdo_client_id":        common.LinuxDOClientId,
 			"telegram_oauth":           common.TelegramOAuthEnabled,
 			"telegram_bot_name":        common.TelegramBotName,
 			"system_name":              common.SystemName,
@@ -63,6 +65,7 @@ func GetStatus(c *gin.Context) {
 			"default_collapse_sidebar": common.DefaultCollapseSidebar,
 			"enable_online_topup":      constant.PayAddress != "" && constant.EpayId != "" && constant.EpayKey != "",
 			"mj_notify_enabled":        constant.MjNotifyEnabled,
+			"chats":                    constant.Chats,
 		},
 	})
 	return

controller/model.go

@@ -137,31 +137,63 @@ func init() {
 }
 
 func ListModels(c *gin.Context) {
-	userId := c.GetInt("id")
-	user, err := model.GetUserById(userId, true)
-	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
-		return
-	}
-	models := model.GetGroupModels(user.Group)
 	userOpenAiModels := make([]dto.OpenAIModels, 0)
 	permission := getPermission()
-	for _, s := range models {
-		if _, ok := openAIModelsMap[s]; ok {
-			userOpenAiModels = append(userOpenAiModels, openAIModelsMap[s])
+
+	modelLimitEnable := c.GetBool("token_model_limit_enabled")
+	if modelLimitEnable {
+		s, ok := c.Get("token_model_limit")
+		var tokenModelLimit map[string]bool
+		if ok {
+			tokenModelLimit = s.(map[string]bool)
 		} else {
-			userOpenAiModels = append(userOpenAiModels, dto.OpenAIModels{
-				Id:         s,
-				Object:     "model",
-				Created:    1626777600,
-				OwnedBy:    "custom",
-				Permission: permission,
-				Root:       s,
-				Parent:     nil,
+			tokenModelLimit = map[string]bool{}
+		}
+		for allowModel, _ := range tokenModelLimit {
+			if _, ok := openAIModelsMap[allowModel]; ok {
+				userOpenAiModels = append(userOpenAiModels, openAIModelsMap[allowModel])
+			} else {
+				userOpenAiModels = append(userOpenAiModels, dto.OpenAIModels{
+					Id:         allowModel,
+					Object:     "model",
+					Created:    1626777600,
+					OwnedBy:    "custom",
+					Permission: permission,
+					Root:       allowModel,
+					Parent:     nil,
+				})
+			}
+		}
+	} else {
+		userId := c.GetInt("id")
+		userGroup, err := model.GetUserGroup(userId)
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "get user group failed",
 			})
+			return
+		}
+		group := userGroup
+		tokenGroup := c.GetString("token_group")
+		if tokenGroup != "" {
+			group = tokenGroup
+		}
+		models := model.GetGroupModels(group)
+		for _, s := range models {
+			if _, ok := openAIModelsMap[s]; ok {
+				userOpenAiModels = append(userOpenAiModels, openAIModelsMap[s])
+			} else {
+				userOpenAiModels = append(userOpenAiModels, dto.OpenAIModels{
+					Id:         s,
+					Object:     "model",
+					Created:    1626777600,
+					OwnedBy:    "custom",
+					Permission: permission,
+					Root:       s,
+					Parent:     nil,
+				})
+			}
 		}
 	}
 	c.JSON(200, gin.H{

controller/option.go

@@ -50,6 +50,14 @@ func UpdateOption(c *gin.Context) {
 			})
 			return
 		}
+	case "LinuxDOOAuthEnabled":
+		if option.Value == "true" && common.LinuxDOClientId == "" {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "无法启用 LinuxDO OAuth，请先填入 LinuxDO Client Id 以及 LinuxDO Client Secret！",
+			})
+			return
+		}
 	case "EmailDomainRestrictionEnabled":
 		if option.Value == "true" && len(common.EmailDomainWhitelist) == 0 {
 			c.JSON(http.StatusOK, gin.H{

controller/pricing.go

@@ -7,18 +7,11 @@ import (
 )
 
 func GetPricing(c *gin.Context) {
-	userId := c.GetInt("id")
-	// if no login, get default group ratio
-	groupRatio := common.GetGroupRatio("default")
-	group, err := model.CacheGetUserGroup(userId)
-	if err == nil {
-		groupRatio = common.GetGroupRatio(group)
-	}
-	pricing := model.GetPricing(group)
+	pricing := model.GetPricing()
 	c.JSON(200, gin.H{
 		"success":     true,
 		"data":        pricing,
-		"group_ratio": groupRatio,
+		"group_ratio": common.GroupRatio,
 	})
 }
 

controller/relay.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"io"
 	"log"
 	"net/http"
@@ -38,6 +39,67 @@ func relayHandler(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode
 	return err
 }
 
+func wsHandler(c *gin.Context, ws *websocket.Conn, relayMode int) *dto.OpenAIErrorWithStatusCode {
+	var err *dto.OpenAIErrorWithStatusCode
+	switch relayMode {
+	default:
+		err = relay.TextHelper(c)
+	}
+	return err
+}
+
+func Playground(c *gin.Context) {
+	var openaiErr *dto.OpenAIErrorWithStatusCode
+
+	defer func() {
+		if openaiErr != nil {
+			c.JSON(openaiErr.StatusCode, gin.H{
+				"error": openaiErr.Error,
+			})
+		}
+	}()
+
+	useAccessToken := c.GetBool("use_access_token")
+	if useAccessToken {
+		openaiErr = service.OpenAIErrorWrapperLocal(errors.New("暂不支持使用 access token"), "access_token_not_supported", http.StatusBadRequest)
+		return
+	}
+
+	playgroundRequest := &dto.PlayGroundRequest{}
+	err := common.UnmarshalBodyReusable(c, playgroundRequest)
+	if err != nil {
+		openaiErr = service.OpenAIErrorWrapperLocal(err, "unmarshal_request_failed", http.StatusBadRequest)
+		return
+	}
+
+	if playgroundRequest.Model == "" {
+		openaiErr = service.OpenAIErrorWrapperLocal(errors.New("请选择模型"), "model_required", http.StatusBadRequest)
+		return
+	}
+	c.Set("original_model", playgroundRequest.Model)
+	group := playgroundRequest.Group
+	userGroup := c.GetString("group")
+
+	if group == "" {
+		group = userGroup
+	} else {
+		if !common.GroupInUserUsableGroups(group) && group != userGroup {
+			openaiErr = service.OpenAIErrorWrapperLocal(errors.New("无权访问该分组"), "group_not_allowed", http.StatusForbidden)
+			return
+		}
+		c.Set("group", group)
+	}
+	c.Set("token_name", "playground-"+group)
+	channel, err := model.CacheGetRandomSatisfiedChannel(group, playgroundRequest.Model, 0)
+	if err != nil {
+		message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", group, playgroundRequest.Model)
+		openaiErr = service.OpenAIErrorWrapperLocal(errors.New(message), "get_playground_channel_failed", http.StatusInternalServerError)
+		return
+	}
+	middleware.SetupContextForSelectedChannel(c, channel, playgroundRequest.Model)
+	Relay(c)
+}
+
 func Relay(c *gin.Context) {
 	relayMode := constant.Path2RelayMode(c.Request.URL.Path)
 	requestId := c.GetString(common.RequestIdKey)
@@ -82,6 +144,67 @@ func Relay(c *gin.Context) {
 	}
 }
 
+var upgrader = websocket.Upgrader{
+	Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
+	CheckOrigin: func(r *http.Request) bool {
+		return true // 允许跨域
+	},
+}
+
+func WssRelay(c *gin.Context) {
+	// 将 HTTP 连接升级为 WebSocket 连接
+
+	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	defer ws.Close()
+
+	if err != nil {
+		openaiErr := service.OpenAIErrorWrapper(err, "get_channel_failed", http.StatusInternalServerError)
+		service.WssError(c, ws, openaiErr.Error)
+		return
+	}
+
+	relayMode := constant.Path2RelayMode(c.Request.URL.Path)
+	requestId := c.GetString(common.RequestIdKey)
+	group := c.GetString("group")
+	//wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01
+	originalModel := c.GetString("original_model")
+	var openaiErr *dto.OpenAIErrorWithStatusCode
+
+	for i := 0; i <= common.RetryTimes; i++ {
+		channel, err := getChannel(c, group, originalModel, i)
+		if err != nil {
+			common.LogError(c, err.Error())
+			openaiErr = service.OpenAIErrorWrapperLocal(err, "get_channel_failed", http.StatusInternalServerError)
+			break
+		}
+
+		openaiErr = wssRequest(c, ws, relayMode, channel)
+
+		if openaiErr == nil {
+			return // 成功处理请求，直接返回
+		}
+
+		go processChannelError(c, channel.Id, channel.Type, channel.Name, channel.GetAutoBan(), openaiErr)
+
+		if !shouldRetry(c, openaiErr, common.RetryTimes-i) {
+			break
+		}
+	}
+	useChannel := c.GetStringSlice("use_channel")
+	if len(useChannel) > 1 {
+		retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
+		common.LogInfo(c, retryLogStr)
+	}
+
+	if openaiErr != nil {
+		if openaiErr.StatusCode == http.StatusTooManyRequests {
+			openaiErr.Error.Message = "当前分组上游负载已饱和，请稍后再试"
+		}
+		openaiErr.Error.Message = common.MessageWithRequestId(openaiErr.Error.Message, requestId)
+		service.WssError(c, ws, openaiErr.Error)
+	}
+}
+
 func relayRequest(c *gin.Context, relayMode int, channel *model.Channel) *dto.OpenAIErrorWithStatusCode {
 	addUsedChannel(c, channel.Id)
 	requestBody, _ := common.GetRequestBody(c)
@@ -89,6 +212,13 @@ func relayRequest(c *gin.Context, relayMode int, channel *model.Channel) *dto.Op
 	return relayHandler(c, relayMode)
 }
 
+func wssRequest(c *gin.Context, ws *websocket.Conn, relayMode int, channel *model.Channel) *dto.OpenAIErrorWithStatusCode {
+	addUsedChannel(c, channel.Id)
+	requestBody, _ := common.GetRequestBody(c)
+	c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
+	return relay.WssHelper(c, ws)
+}
+
 func addUsedChannel(c *gin.Context, channelId int) {
 	useChannel := c.GetStringSlice("use_channel")
 	useChannel = append(useChannel, fmt.Sprintf("%d", channelId))

controller/telegram.go

@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"io"
+	"net/http"
 	"one-api/common"
 	"one-api/model"
 	"sort"
@@ -48,6 +49,13 @@ func TelegramBind(c *gin.Context) {
 		})
 		return
 	}
+	if user.Id == 0 {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "用户已注销",
+		})
+		return
+	}
 	user.TelegramId = telegramId
 	if err := user.Update(false); err != nil {
 		c.JSON(200, gin.H{

controller/token.go

@@ -123,10 +123,19 @@ func AddToken(c *gin.Context) {
 		})
 		return
 	}
+	key, err := common.GenerateKey()
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "生成令牌失败",
+		})
+		common.SysError("failed to generate token key: " + err.Error())
+		return
+	}
 	cleanToken := model.Token{
 		UserId:             c.GetInt("id"),
 		Name:               token.Name,
-		Key:                common.GenerateKey(),
+		Key:                key,
 		CreatedTime:        common.GetTimestamp(),
 		AccessedTime:       common.GetTimestamp(),
 		ExpiredTime:        token.ExpiredTime,
@@ -134,6 +143,8 @@ func AddToken(c *gin.Context) {
 		UnlimitedQuota:     token.UnlimitedQuota,
 		ModelLimitsEnabled: token.ModelLimitsEnabled,
 		ModelLimits:        token.ModelLimits,
+		AllowIps:           token.AllowIps,
+		Group:              token.Group,
 	}
 	err = cleanToken.Insert()
 	if err != nil {
@@ -221,6 +232,8 @@ func UpdateToken(c *gin.Context) {
 		cleanToken.UnlimitedQuota = token.UnlimitedQuota
 		cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
 		cleanToken.ModelLimits = token.ModelLimits
+		cleanToken.AllowIps = token.AllowIps
+		cleanToken.Group = token.Group
 	}
 	err = cleanToken.Update()
 	if err != nil {

controller/user.go

@@ -7,10 +7,12 @@ import (
 	"one-api/common"
 	"one-api/model"
 	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-gonic/gin"
+	"one-api/constant"
 )
 
 type LoginRequest struct {
@@ -66,6 +68,7 @@ func setupLogin(user *model.User, c *gin.Context) {
 	session.Set("username", user.Username)
 	session.Set("role", user.Role)
 	session.Set("status", user.Status)
+	session.Set("group", user.Group)
 	err := session.Save()
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
@@ -157,8 +160,9 @@ func Register(c *gin.Context) {
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
-			"message": err.Error(),
+			"message": "数据库错误，请稍后重试",
 		})
+		common.SysError(fmt.Sprintf("CheckUserExistOrDeleted error: %v", err))
 		return
 	}
 	if exist {
@@ -186,6 +190,48 @@ func Register(c *gin.Context) {
 		})
 		return
 	}
+
+	// 获取插入后的用户ID
+	var insertedUser model.User
+	if err := model.DB.Where("username = ?", cleanUser.Username).First(&insertedUser).Error; err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "用户注册失败或用户ID获取失败",
+		})
+		return
+	}
+	// 生成默认令牌
+	if constant.GenerateDefaultToken {
+		key, err := common.GenerateKey()
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "生成默认令牌失败",
+			})
+			common.SysError("failed to generate token key: " + err.Error())
+			return
+		}
+		// 生成默认令牌
+		token := model.Token{
+			UserId:             insertedUser.Id, // 使用插入后的用户ID
+			Name:               cleanUser.Username + "的初始令牌",
+			Key:                key,
+			CreatedTime:        common.GetTimestamp(),
+			AccessedTime:       common.GetTimestamp(),
+			ExpiredTime:        -1,     // 永不过期
+			RemainQuota:        500000, // 示例额度
+			UnlimitedQuota:     true,
+			ModelLimitsEnabled: false,
+		}
+		if err := token.Insert(); err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "创建默认令牌失败",
+			})
+			return
+		}
+	}
+
 	c.JSON(http.StatusOK, gin.H{
 		"success": true,
 		"message": "",
@@ -276,7 +322,18 @@ func GenerateAccessToken(c *gin.Context) {
 		})
 		return
 	}
-	user.AccessToken = common.GetUUID()
+	// get rand int 28-32
+	randI := common.GetRandomInt(4)
+	key, err := common.GenerateRandomKey(29 + randI)
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "生成失败",
+		})
+		common.SysError("failed to generate key: " + err.Error())
+		return
+	}
+	user.SetAccessToken(key)
 
 	if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 {
 		c.JSON(http.StatusOK, gin.H{
@@ -582,6 +639,7 @@ func DeleteSelf(c *gin.Context) {
 func CreateUser(c *gin.Context) {
 	var user model.User
 	err := json.NewDecoder(c.Request.Body).Decode(&user)
+	user.Username = strings.TrimSpace(user.Username)
 	if err != nil || user.Username == "" || user.Password == "" {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
@@ -629,8 +687,8 @@ func CreateUser(c *gin.Context) {
 }
 
 type ManageRequest struct {
-	Username string `json:"username"`
-	Action   string `json:"action"`
+	Id     int    `json:"id"`
+	Action string `json:"action"`
 }
 
 // ManageUser Only admin user can do this
@@ -646,7 +704,7 @@ func ManageUser(c *gin.Context) {
 		return
 	}
 	user := model.User{
-		Username: req.Username,
+		Id: req.Id,
 	}
 	// Fill attributes
 	model.DB.Unscoped().Where(&user).First(&user)

controller/wechat.go

@@ -78,6 +78,13 @@ func WeChatAuth(c *gin.Context) {
 			})
 			return
 		}
+		if user.Id == 0 {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "用户已注销",
+			})
+			return
+		}
 	} else {
 		if common.RegisterEnabled {
 			user.Username = "wechat_" + strconv.Itoa(model.GetMaxUserId()+1)

dto/openai_request.go

@@ -26,6 +26,7 @@ type GeneralOpenAIRequest struct {
 	FrequencyPenalty    float64        `json:"frequency_penalty,omitempty"`
 	PresencePenalty     float64        `json:"presence_penalty,omitempty"`
 	ResponseFormat      any            `json:"response_format,omitempty"`
+	EncodingFormat      any            `json:"encoding_format,omitempty"`
 	Seed                float64        `json:"seed,omitempty"`
 	Tools               []ToolCall     `json:"tools,omitempty"`
 	ToolChoice          any            `json:"tool_choice,omitempty"`
@@ -33,6 +34,8 @@ type GeneralOpenAIRequest struct {
 	LogProbs            bool           `json:"logprobs,omitempty"`
 	TopLogProbs         int            `json:"top_logprobs,omitempty"`
 	Dimensions          int            `json:"dimensions,omitempty"`
+	Modalities          any            `json:"modalities,omitempty"`
+	Audio               any            `json:"audio,omitempty"`
 }
 
 type OpenAITools struct {
@@ -82,9 +85,10 @@ type Message struct {
 }
 
 type MediaMessage struct {
-	Type     string `json:"type"`
-	Text     string `json:"text"`
-	ImageUrl any    `json:"image_url,omitempty"`
+	Type       string `json:"type"`
+	Text       string `json:"text"`
+	ImageUrl   any    `json:"image_url,omitempty"`
+	InputAudio any    `json:"input_audio,omitempty"`
 }
 
 type MessageImageUrl struct {
@@ -92,9 +96,15 @@ type MessageImageUrl struct {
 	Detail string `json:"detail"`
 }
 
+type MessageInputAudio struct {
+	Data   string `json:"data"` //base64
+	Format string `json:"format"`
+}
+
 const (
-	ContentTypeText     = "text"
-	ContentTypeImageURL = "image_url"
+	ContentTypeText       = "text"
+	ContentTypeImageURL   = "image_url"
+	ContentTypeInputAudio = "input_audio"
 )
 
 func (m Message) StringContent() string {
@@ -167,11 +177,19 @@ func (m Message) ParseContent() []MediaMessage {
 						},
 					})
 				}
-
+			case ContentTypeInputAudio:
+				if subObj, ok := contentMap["input_audio"].(map[string]any); ok {
+					contentList = append(contentList, MediaMessage{
+						Type: ContentTypeInputAudio,
+						InputAudio: MessageInputAudio{
+							Data:   subObj["data"].(string),
+							Format: subObj["format"].(string),
+						},
+					})
+				}
 			}
 		}
 		return contentList
 	}
-
 	return nil
 }

dto/openai_response.go

@@ -128,7 +128,9 @@ type CompletionsStreamResponse struct {
 }
 
 type Usage struct {
-	PromptTokens     int `json:"prompt_tokens"`
-	CompletionTokens int `json:"completion_tokens"`
-	TotalTokens      int `json:"total_tokens"`
+	PromptTokens           int                `json:"prompt_tokens"`
+	CompletionTokens       int                `json:"completion_tokens"`
+	TotalTokens            int                `json:"total_tokens"`
+	PromptTokensDetails    InputTokenDetails  `json:"prompt_tokens_details"`
+	CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"`
 }

dto/playground.go

@@ -0,0 +1,6 @@
+package dto
+
+type PlayGroundRequest struct {
+	Model string `json:"model,omitempty"`
+	Group string `json:"group,omitempty"`
+}

dto/realtime.go

@@ -0,0 +1,98 @@
+package dto
+
+const (
+	RealtimeEventTypeError              = "error"
+	RealtimeEventTypeSessionUpdate      = "session.update"
+	RealtimeEventTypeConversationCreate = "conversation.item.create"
+	RealtimeEventTypeResponseCreate     = "response.create"
+	RealtimeEventInputAudioBufferAppend = "input_audio_buffer.append"
+)
+
+const (
+	RealtimeEventTypeResponseDone                   = "response.done"
+	RealtimeEventTypeSessionUpdated                 = "session.updated"
+	RealtimeEventTypeSessionCreated                 = "session.created"
+	RealtimeEventResponseAudioDelta                 = "response.audio.delta"
+	RealtimeEventResponseAudioTranscriptionDelta    = "response.audio_transcript.delta"
+	RealtimeEventResponseFunctionCallArgumentsDelta = "response.function_call_arguments.delta"
+	RealtimeEventResponseFunctionCallArgumentsDone  = "response.function_call_arguments.done"
+	RealtimeEventConversationItemCreated            = "conversation.item.created"
+)
+
+type RealtimeEvent struct {
+	EventId string `json:"event_id"`
+	Type    string `json:"type"`
+	//PreviousItemId string `json:"previous_item_id"`
+	Session  *RealtimeSession  `json:"session,omitempty"`
+	Item     *RealtimeItem     `json:"item,omitempty"`
+	Error    *OpenAIError      `json:"error,omitempty"`
+	Response *RealtimeResponse `json:"response,omitempty"`
+	Delta    string            `json:"delta,omitempty"`
+	Audio    string            `json:"audio,omitempty"`
+}
+
+type RealtimeResponse struct {
+	Usage *RealtimeUsage `json:"usage"`
+}
+
+type RealtimeUsage struct {
+	TotalTokens        int                `json:"total_tokens"`
+	InputTokens        int                `json:"input_tokens"`
+	OutputTokens       int                `json:"output_tokens"`
+	InputTokenDetails  InputTokenDetails  `json:"input_token_details"`
+	OutputTokenDetails OutputTokenDetails `json:"output_token_details"`
+}
+
+type InputTokenDetails struct {
+	CachedTokens int `json:"cached_tokens"`
+	TextTokens   int `json:"text_tokens"`
+	AudioTokens  int `json:"audio_tokens"`
+	ImageTokens  int `json:"image_tokens"`
+}
+
+type OutputTokenDetails struct {
+	TextTokens  int `json:"text_tokens"`
+	AudioTokens int `json:"audio_tokens"`
+}
+
+type RealtimeSession struct {
+	Modalities              []string                `json:"modalities"`
+	Instructions            string                  `json:"instructions"`
+	Voice                   string                  `json:"voice"`
+	InputAudioFormat        string                  `json:"input_audio_format"`
+	OutputAudioFormat       string                  `json:"output_audio_format"`
+	InputAudioTranscription InputAudioTranscription `json:"input_audio_transcription"`
+	TurnDetection           interface{}             `json:"turn_detection"`
+	Tools                   []RealTimeTool          `json:"tools"`
+	ToolChoice              string                  `json:"tool_choice"`
+	Temperature             float64                 `json:"temperature"`
+	//MaxResponseOutputTokens int                     `json:"max_response_output_tokens"`
+}
+
+type InputAudioTranscription struct {
+	Model string `json:"model"`
+}
+
+type RealTimeTool struct {
+	Type        string `json:"type"`
+	Name        string `json:"name"`
+	Description string `json:"description"`
+	Parameters  any    `json:"parameters"`
+}
+
+type RealtimeItem struct {
+	Id        string            `json:"id"`
+	Type      string            `json:"type"`
+	Status    string            `json:"status"`
+	Role      string            `json:"role"`
+	Content   []RealtimeContent `json:"content"`
+	Name      *string           `json:"name,omitempty"`
+	ToolCalls any               `json:"tool_calls,omitempty"`
+	CallId    string            `json:"call_id,omitempty"`
+}
+type RealtimeContent struct {
+	Type       string `json:"type"`
+	Text       string `json:"text,omitempty"`
+	Audio      string `json:"audio,omitempty"` // Base64-encoded audio bytes.
+	Transcript string `json:"transcript,omitempty"`
+}

go.mod

@@ -63,6 +63,7 @@ require (
 	github.com/jackc/puddle/v2 v2.2.1 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
+	github.com/joho/godotenv v1.5.1 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.4 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect

go.sum

@@ -111,6 +111,8 @@ github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkr
 github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
 github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
 github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
+github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
+github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
 github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=

main.go

@@ -3,10 +3,6 @@ package main
 import (
 	"embed"
 	"fmt"
-	"github.com/bytedance/gopkg/util/gopool"
-	"github.com/gin-contrib/sessions"
-	"github.com/gin-contrib/sessions/cookie"
-	"github.com/gin-gonic/gin"
 	"log"
 	"net/http"
 	"one-api/common"
@@ -19,6 +15,12 @@ import (
 	"os"
 	"strconv"
 
+	"github.com/bytedance/gopkg/util/gopool"
+	"github.com/gin-contrib/sessions"
+	"github.com/gin-contrib/sessions/cookie"
+	"github.com/gin-gonic/gin"
+	"github.com/joho/godotenv"
+
 	_ "net/http/pprof"
 )
 
@@ -29,6 +31,11 @@ var buildFS embed.FS
 var indexPage []byte
 
 func main() {
+	err := godotenv.Load(".env")
+	if err != nil {
+		common.SysLog("Can't load .env file")
+	}
+
 	common.SetupLogger()
 	common.SysLog("New API " + common.Version + " started")
 	if os.Getenv("GIN_MODE") != "debug" {
@@ -38,7 +45,7 @@ func main() {
 		common.SysLog("running in debug mode")
 	}
 	// Initialize SQL Database
-	err := model.InitDB()
+	err = model.InitDB()
 	if err != nil {
 		common.FatalLog("failed to initialize database: " + err.Error())
 	}

middleware/auth.go

@@ -10,6 +10,17 @@ import (
 	"strings"
 )
 
+func validUserInfo(username string, role int) bool {
+	// check username is empty
+	if strings.TrimSpace(username) == "" {
+		return false
+	}
+	if !common.IsValidateRole(role) {
+		return false
+	}
+	return true
+}
+
 func authHelper(c *gin.Context, minRole int) {
 	session := sessions.Default(c)
 	username := session.Get("username")
@@ -30,6 +41,14 @@ func authHelper(c *gin.Context, minRole int) {
 		}
 		user := model.ValidateAccessToken(accessToken)
 		if user != nil && user.Username != "" {
+			if !validUserInfo(user.Username, user.Role) {
+				c.JSON(http.StatusOK, gin.H{
+					"success": false,
+					"message": "无权进行此操作，用户信息无效",
+				})
+				c.Abort()
+				return
+			}
 			// Token is valid
 			username = user.Username
 			role = user.Role
@@ -91,9 +110,19 @@ func authHelper(c *gin.Context, minRole int) {
 		c.Abort()
 		return
 	}
+	if !validUserInfo(username.(string), role.(int)) {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": "无权进行此操作，用户信息无效",
+		})
+		c.Abort()
+		return
+	}
 	c.Set("username", username)
 	c.Set("role", role)
 	c.Set("id", id)
+	c.Set("group", session.Get("group"))
+	c.Set("use_access_token", useAccessToken)
 	c.Next()
 }
 
@@ -126,8 +155,27 @@ func RootAuth() func(c *gin.Context) {
 	}
 }
 
+func WssAuth(c *gin.Context) {
+
+}
+
 func TokenAuth() func(c *gin.Context) {
 	return func(c *gin.Context) {
+		// 先检测是否为ws
+		if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" {
+			// Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1
+			// read sk from Sec-WebSocket-Protocol
+			key := c.Request.Header.Get("Sec-WebSocket-Protocol")
+			parts := strings.Split(key, ",")
+			for _, part := range parts {
+				part = strings.TrimSpace(part)
+				if strings.HasPrefix(part, "openai-insecure-api-key") {
+					key = strings.TrimPrefix(part, "openai-insecure-api-key.")
+					break
+				}
+			}
+			c.Request.Header.Set("Authorization", "Bearer "+key)
+		}
 		key := c.Request.Header.Get("Authorization")
 		parts := make([]string, 0)
 		key = strings.TrimPrefix(key, "Bearer ")
@@ -175,6 +223,8 @@ func TokenAuth() func(c *gin.Context) {
 		} else {
 			c.Set("token_model_limit_enabled", false)
 		}
+		c.Set("allow_ips", token.GetIpLimitsMap())
+		c.Set("token_group", token.Group)
 		if len(parts) > 1 {
 			if model.IsAdmin(token.UserId) {
 				c.Set("specific_channel_id", parts[1])

middleware/distributor.go

@@ -22,6 +22,14 @@ type ModelRequest struct {
 
 func Distribute() func(c *gin.Context) {
 	return func(c *gin.Context) {
+		allowIpsMap := c.GetStringMap("allow_ips")
+		if len(allowIpsMap) != 0 {
+			clientIp := c.ClientIP()
+			if _, ok := allowIpsMap[clientIp]; !ok {
+				abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中")
+				return
+			}
+		}
 		userId := c.GetInt("id")
 		var channel *model.Channel
 		channelId, ok := c.Get("specific_channel_id")
@@ -31,6 +39,20 @@ func Distribute() func(c *gin.Context) {
 			return
 		}
 		userGroup, _ := model.CacheGetUserGroup(userId)
+		tokenGroup := c.GetString("token_group")
+		if tokenGroup != "" {
+			// check common.UserUsableGroups[userGroup]
+			if _, ok := common.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
+				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("令牌分组 %s 已被禁用", tokenGroup))
+				return
+			}
+			// check group in common.GroupRatio
+			if _, ok := common.GroupRatio[tokenGroup]; !ok {
+				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
+				return
+			}
+			userGroup = tokenGroup
+		}
 		c.Set("group", userGroup)
 		if ok {
 			id, err := strconv.Atoi(channelId.(string))
@@ -148,6 +170,10 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
 		abortWithOpenAiMessage(c, http.StatusBadRequest, "无效的请求, "+err.Error())
 		return nil, false, errors.New("无效的请求, " + err.Error())
 	}
+	if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") {
+		//wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01
+		modelRequest.Model = c.Query("model")
+	}
 	if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
 		if modelRequest.Model == "" {
 			modelRequest.Model = "text-moderation-stable"

middleware/rate-limit.go

@@ -13,6 +13,10 @@ var timeFormat = "2006-01-02T15:04:05.000Z"
 
 var inMemoryRateLimiter common.InMemoryRateLimiter
 
+var defNext = func(c *gin.Context) {
+	c.Next()
+}
+
 func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) {
 	ctx := context.Background()
 	rdb := common.RDB
@@ -83,11 +87,17 @@ func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gi
 }
 
 func GlobalWebRateLimit() func(c *gin.Context) {
-	return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
+	if common.GlobalWebRateLimitEnable {
+		return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
+	}
+	return defNext
 }
 
 func GlobalAPIRateLimit() func(c *gin.Context) {
-	return rateLimitFactory(common.GlobalApiRateLimitNum, common.GlobalApiRateLimitDuration, "GA")
+	if common.GlobalApiRateLimitEnable {
+		return rateLimitFactory(common.GlobalApiRateLimitNum, common.GlobalApiRateLimitDuration, "GA")
+	}
+	return defNext
 }
 
 func CriticalRateLimit() func(c *gin.Context) {

model/ability.go

@@ -36,6 +36,12 @@ func GetEnabledModels() []string {
 	return models
 }
 
+func GetAllEnableAbilities() []Ability {
+	var abilities []Ability
+	DB.Find(&abilities, "enabled = ?", true)
+	return abilities
+}
+
 func getPriority(group string, model string, retry int) (int, error) {
 	groupCol := "`group`"
 	trueVal := "1"

model/log.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"one-api/common"
+	"os"
 	"strings"
 	"time"
 
@@ -39,7 +40,15 @@ const (
 )
 
 func GetLogByKey(key string) (logs []*Log, err error) {
-	err = LOG_DB.Joins("left join tokens on tokens.id = logs.token_id").Where("tokens.key = ?", strings.TrimPrefix(key, "sk-")).Find(&logs).Error
+	if os.Getenv("LOG_SQL_DSN") != "" {
+		var tk Token
+		if err = DB.Model(&Token{}).Where("`key`=?", strings.TrimPrefix(key, "sk-")).First(&tk).Error; err != nil {
+			return nil, err
+		}
+		err = LOG_DB.Model(&Log{}).Where("token_id=?", tk.Id).Find(&logs).Error
+	} else {
+		err = LOG_DB.Joins("left join tokens on tokens.id = logs.token_id").Where("tokens.key = ?", strings.TrimPrefix(key, "sk-")).Find(&logs).Error
+	}
 	return logs, err
 }
 

model/main.go

@@ -32,7 +32,7 @@ func createRootAccountIfNeed() error {
 			Role:        common.RoleRootUser,
 			Status:      common.UserStatusEnabled,
 			DisplayName: "Root User",
-			AccessToken: common.GetUUID(),
+			AccessToken: nil,
 			Quota:       100000000,
 		}
 		DB.Create(&rootUser)

model/option.go

@@ -31,6 +31,7 @@ func InitOptionMap() {
 	common.OptionMap["PasswordRegisterEnabled"] = strconv.FormatBool(common.PasswordRegisterEnabled)
 	common.OptionMap["EmailVerificationEnabled"] = strconv.FormatBool(common.EmailVerificationEnabled)
 	common.OptionMap["GitHubOAuthEnabled"] = strconv.FormatBool(common.GitHubOAuthEnabled)
+	common.OptionMap["LinuxDOOAuthEnabled"] = strconv.FormatBool(common.LinuxDOOAuthEnabled)
 	common.OptionMap["TelegramOAuthEnabled"] = strconv.FormatBool(common.TelegramOAuthEnabled)
 	common.OptionMap["WeChatAuthEnabled"] = strconv.FormatBool(common.WeChatAuthEnabled)
 	common.OptionMap["TurnstileCheckEnabled"] = strconv.FormatBool(common.TurnstileCheckEnabled)
@@ -69,6 +70,7 @@ func InitOptionMap() {
 	common.OptionMap["Price"] = strconv.FormatFloat(constant.Price, 'f', -1, 64)
 	common.OptionMap["MinTopUp"] = strconv.Itoa(constant.MinTopUp)
 	common.OptionMap["TopupGroupRatio"] = common.TopupGroupRatio2JSONString()
+	common.OptionMap["Chats"] = constant.Chats2JsonString()
 	common.OptionMap["GitHubClientId"] = ""
 	common.OptionMap["GitHubClientSecret"] = ""
 	common.OptionMap["TelegramBotToken"] = ""
@@ -86,6 +88,7 @@ func InitOptionMap() {
 	common.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
 	common.OptionMap["ModelPrice"] = common.ModelPrice2JSONString()
 	common.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
+	common.OptionMap["UserUsableGroups"] = common.UserUsableGroups2JSONString()
 	common.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString()
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMap["ChatLink"] = common.ChatLink
@@ -173,6 +176,8 @@ func updateOptionMap(key string, value string) (err error) {
 			common.EmailVerificationEnabled = boolValue
 		case "GitHubOAuthEnabled":
 			common.GitHubOAuthEnabled = boolValue
+		case "LinuxDOOAuthEnabled":
+			common.LinuxDOOAuthEnabled = boolValue
 		case "WeChatAuthEnabled":
 			common.WeChatAuthEnabled = boolValue
 		case "TelegramOAuthEnabled":
@@ -247,6 +252,8 @@ func updateOptionMap(key string, value string) (err error) {
 		constant.WorkerValidKey = value
 	case "PayAddress":
 		constant.PayAddress = value
+	case "Chats":
+		err = constant.UpdateChatsByJsonString(value)
 	case "CustomCallbackAddress":
 		constant.CustomCallbackAddress = value
 	case "EpayId":
@@ -263,6 +270,10 @@ func updateOptionMap(key string, value string) (err error) {
 		common.GitHubClientId = value
 	case "GitHubClientSecret":
 		common.GitHubClientSecret = value
+	case "LinuxDOClientId":
+		common.LinuxDOClientId = value
+	case "LinuxDOClientSecret":
+		common.LinuxDOClientSecret = value
 	case "Footer":
 		common.Footer = value
 	case "SystemName":
@@ -303,6 +314,8 @@ func updateOptionMap(key string, value string) (err error) {
 		err = common.UpdateModelRatioByJSONString(value)
 	case "GroupRatio":
 		err = common.UpdateGroupRatioByJSONString(value)
+	case "UserUsableGroups":
+		err = common.UpdateUserUsableGroupsByJSONString(value)
 	case "CompletionRatio":
 		err = common.UpdateCompletionRatioByJSONString(value)
 	case "ModelPrice":

model/pricing.go

@@ -7,14 +7,13 @@ import (
 )
 
 type Pricing struct {
-	Available       bool     `json:"available"`
 	ModelName       string   `json:"model_name"`
 	QuotaType       int      `json:"quota_type"`
 	ModelRatio      float64  `json:"model_ratio"`
 	ModelPrice      float64  `json:"model_price"`
 	OwnerBy         string   `json:"owner_by"`
 	CompletionRatio float64  `json:"completion_ratio"`
-	EnableGroup     []string `json:"enable_group,omitempty"`
+	EnableGroup     []string `json:"enable_groups,omitempty"`
 }
 
 var (
@@ -23,40 +22,47 @@ var (
 	updatePricingLock  sync.Mutex
 )
 
-func GetPricing(group string) []Pricing {
+func GetPricing() []Pricing {
 	updatePricingLock.Lock()
 	defer updatePricingLock.Unlock()
 
 	if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
 		updatePricing()
 	}
-	if group != "" {
-		userPricingMap := make([]Pricing, 0)
-		models := GetGroupModels(group)
-		for _, pricing := range pricingMap {
-			if !common.StringsContains(models, pricing.ModelName) {
-				pricing.Available = false
-			}
-			userPricingMap = append(userPricingMap, pricing)
-		}
-		return userPricingMap
-	}
+	//if group != "" {
+	//	userPricingMap := make([]Pricing, 0)
+	//	models := GetGroupModels(group)
+	//	for _, pricing := range pricingMap {
+	//		if !common.StringsContains(models, pricing.ModelName) {
+	//			pricing.Available = false
+	//		}
+	//		userPricingMap = append(userPricingMap, pricing)
+	//	}
+	//	return userPricingMap
+	//}
 	return pricingMap
 }
 
 func updatePricing() {
 	//modelRatios := common.GetModelRatios()
-	enabledModels := GetEnabledModels()
-	allModels := make(map[string]int)
-	for i, model := range enabledModels {
-		allModels[model] = i
+	enableAbilities := GetAllEnableAbilities()
+	modelGroupsMap := make(map[string][]string)
+	for _, ability := range enableAbilities {
+		groups := modelGroupsMap[ability.Model]
+		if groups == nil {
+			groups = make([]string, 0)
+		}
+		if !common.StringsContains(groups, ability.Group) {
+			groups = append(groups, ability.Group)
+		}
+		modelGroupsMap[ability.Model] = groups
 	}
 
 	pricingMap = make([]Pricing, 0)
-	for model, _ := range allModels {
+	for model, groups := range modelGroupsMap {
 		pricing := Pricing{
-			Available: true,
-			ModelName: model,
+			ModelName:   model,
+			EnableGroup: groups,
 		}
 		modelPrice, findPrice := common.GetModelPrice(model, false)
 		if findPrice {

model/token.go

@@ -6,6 +6,7 @@ import (
 	"gorm.io/gorm"
 	"one-api/common"
 	"one-api/constant"
+	relaycommon "one-api/relay/common"
 	"strconv"
 	"strings"
 )
@@ -23,10 +24,34 @@ type Token struct {
 	UnlimitedQuota     bool           `json:"unlimited_quota" gorm:"default:false"`
 	ModelLimitsEnabled bool           `json:"model_limits_enabled" gorm:"default:false"`
 	ModelLimits        string         `json:"model_limits" gorm:"type:varchar(1024);default:''"`
+	AllowIps           *string        `json:"allow_ips" gorm:"default:''"`
 	UsedQuota          int            `json:"used_quota" gorm:"default:0"` // used quota
+	Group              string         `json:"group" gorm:"default:''"`
 	DeletedAt          gorm.DeletedAt `gorm:"index"`
 }
 
+func (token *Token) GetIpLimitsMap() map[string]any {
+	// delete empty spaces
+	//split with \n
+	ipLimitsMap := make(map[string]any)
+	if token.AllowIps == nil {
+		return ipLimitsMap
+	}
+	cleanIps := strings.ReplaceAll(*token.AllowIps, " ", "")
+	if cleanIps == "" {
+		return ipLimitsMap
+	}
+	ips := strings.Split(cleanIps, "\n")
+	for _, ip := range ips {
+		ip = strings.TrimSpace(ip)
+		ip = strings.ReplaceAll(ip, ",", "")
+		if common.IsIP(ip) {
+			ipLimitsMap[ip] = true
+		}
+	}
+	return ipLimitsMap
+}
+
 func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
 	var tokens []*Token
 	var err error
@@ -130,7 +155,8 @@ func (token *Token) Insert() error {
 // Update Make sure your token's fields is completed, because this will update non-zero values
 func (token *Token) Update() error {
 	var err error
-	err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", "model_limits_enabled", "model_limits").Updates(token).Error
+	err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
+		"model_limits_enabled", "model_limits", "allow_ips", "group").Updates(token).Error
 	return err
 }
 
@@ -232,51 +258,56 @@ func decreaseTokenQuota(id int, quota int) (err error) {
 	return err
 }
 
-func PreConsumeTokenQuota(tokenId int, quota int) (userQuota int, err error) {
+func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) (userQuota int, err error) {
 	if quota < 0 {
 		return 0, errors.New("quota 不能为负数！")
 	}
-	token, err := GetTokenById(tokenId)
-	if err != nil {
-		return 0, err
+	if !relayInfo.IsPlayground {
+		token, err := GetTokenById(relayInfo.TokenId)
+		if err != nil {
+			return 0, err
+		}
+		if !token.UnlimitedQuota && token.RemainQuota < quota {
+			return 0, errors.New("令牌额度不足")
+		}
 	}
-	if !token.UnlimitedQuota && token.RemainQuota < quota {
-		return 0, errors.New("令牌额度不足")
-	}
-	userQuota, err = GetUserQuota(token.UserId)
+	userQuota, err = GetUserQuota(relayInfo.UserId)
 	if err != nil {
 		return 0, err
 	}
 	if userQuota < quota {
 		return 0, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
 	}
-	err = DecreaseTokenQuota(tokenId, quota)
-	if err != nil {
-		return 0, err
+	if !relayInfo.IsPlayground {
+		err = DecreaseTokenQuota(relayInfo.TokenId, quota)
+		if err != nil {
+			return 0, err
+		}
 	}
-	err = DecreaseUserQuota(token.UserId, quota)
+	err = DecreaseUserQuota(relayInfo.UserId, quota)
 	return userQuota - quota, err
 }
 
-func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
-	token, err := GetTokenById(tokenId)
+func PostConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
 
 	if quota > 0 {
-		err = DecreaseUserQuota(token.UserId, quota)
+		err = DecreaseUserQuota(relayInfo.UserId, quota)
 	} else {
-		err = IncreaseUserQuota(token.UserId, -quota)
+		err = IncreaseUserQuota(relayInfo.UserId, -quota)
 	}
 	if err != nil {
 		return err
 	}
 
-	if quota > 0 {
-		err = DecreaseTokenQuota(tokenId, quota)
-	} else {
-		err = IncreaseTokenQuota(tokenId, -quota)
-	}
-	if err != nil {
-		return err
+	if !relayInfo.IsPlayground {
+		if quota > 0 {
+			err = DecreaseTokenQuota(relayInfo.TokenId, quota)
+		} else {
+			err = IncreaseTokenQuota(relayInfo.TokenId, -quota)
+		}
+		if err != nil {
+			return err
+		}
 	}
 
 	if sendEmail {
@@ -285,7 +316,7 @@ func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuo
 			noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0
 			if quotaTooLow || noMoreQuota {
 				go func() {
-					email, err := GetUserEmail(token.UserId)
+					email, err := GetUserEmail(relayInfo.UserId)
 					if err != nil {
 						common.SysError("failed to fetch user email: " + err.Error())
 					}

model/user.go

@@ -25,7 +25,7 @@ type User struct {
 	WeChatId         string         `json:"wechat_id" gorm:"column:wechat_id;index"`
 	TelegramId       string         `json:"telegram_id" gorm:"column:telegram_id;index"`
 	VerificationCode string         `json:"verification_code" gorm:"-:all"`                                    // this field is only for Email verification, don't save it to database!
-	AccessToken      string         `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
+	AccessToken      *string        `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
 	Quota            int            `json:"quota" gorm:"type:int;default:0"`
 	UsedQuota        int            `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
 	RequestCount     int            `json:"request_count" gorm:"type:int;default:0;"`               // request number
@@ -36,6 +36,18 @@ type User struct {
 	AffHistoryQuota  int            `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度
 	InviterId        int            `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
 	DeletedAt        gorm.DeletedAt `gorm:"index"`
+	LinuxDOId        string         `json:"linux_do_id" gorm:"column:linux_do_id;index"`
+}
+
+func (user *User) GetAccessToken() string {
+	if user.AccessToken == nil {
+		return ""
+	}
+	return *user.AccessToken
+}
+
+func (user *User) SetAccessToken(token string) {
+	user.AccessToken = &token
 }
 
 // CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
@@ -201,7 +213,7 @@ func (user *User) Insert(inviterId int) error {
 		}
 	}
 	user.Quota = common.QuotaForNewUser
-	user.AccessToken = common.GetUUID()
+	//user.SetAccessToken(common.GetUUID())
 	user.AffCode = common.GetRandomString(4)
 	result := DB.Create(user)
 	if result.Error != nil {
@@ -295,11 +307,12 @@ func (user *User) ValidateAndFill() (err error) {
 	// that means if your field’s value is 0, '', false or other zero values,
 	// it won’t be used to build query conditions
 	password := user.Password
-	if user.Username == "" || password == "" {
+	username := strings.TrimSpace(user.Username)
+	if username == "" || password == "" {
 		return errors.New("用户名或密码为空")
 	}
 	// find buy username or email
-	DB.Where("username = ? OR email = ?", user.Username, user.Username).First(user)
+	DB.Where("username = ? OR email = ?", username, username).First(user)
 	okay := common.ValidatePasswordAndHash(password, user.Password)
 	if !okay || user.Status != common.UserStatusEnabled {
 		return errors.New("用户名或密码错误，或用户已被封禁")
@@ -339,14 +352,6 @@ func (user *User) FillUserByWeChatId() error {
 	return nil
 }
 
-func (user *User) FillUserByUsername() error {
-	if user.Username == "" {
-		return errors.New("username 为空！")
-	}
-	DB.Where(User{Username: user.Username}).First(user)
-	return nil
-}
-
 func (user *User) FillUserByTelegramId() error {
 	if user.TelegramId == "" {
 		return errors.New("Telegram id 为空！")
@@ -359,23 +364,19 @@ func (user *User) FillUserByTelegramId() error {
 }
 
 func IsEmailAlreadyTaken(email string) bool {
-	return DB.Where("email = ?", email).Find(&User{}).RowsAffected == 1
+	return DB.Unscoped().Where("email = ?", email).Find(&User{}).RowsAffected == 1
 }
 
 func IsWeChatIdAlreadyTaken(wechatId string) bool {
-	return DB.Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
+	return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
 }
 
 func IsGitHubIdAlreadyTaken(githubId string) bool {
-	return DB.Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
-}
-
-func IsUsernameAlreadyTaken(username string) bool {
-	return DB.Where("username = ?", username).Find(&User{}).RowsAffected == 1
+	return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
 }
 
 func IsTelegramIdAlreadyTaken(telegramId string) bool {
-	return DB.Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
+	return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
 }
 
 func ResetUserPasswordByEmail(email string, password string) error {
@@ -537,3 +538,17 @@ func GetUsernameById(id int) (username string, err error) {
 	err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
 	return username, err
 }
+
+func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
+	var user User
+	err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error
+	return !errors.Is(err, gorm.ErrRecordNotFound)
+}
+
+func (u *User) FillUserByLinuxDOId() error {
+	if u.LinuxDOId == "" {
+		return errors.New("linux do id is empty")
+	}
+	err := DB.Where("linux_do_id = ?", u.LinuxDOId).First(u).Error
+	return err
+}

relay/channel/adapter.go

@@ -12,13 +12,13 @@ type Adaptor interface {
 	// Init IsStream bool
 	Init(info *relaycommon.RelayInfo)
 	GetRequestURL(info *relaycommon.RelayInfo) (string, error)
-	SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error
+	SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error
 	ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error)
 	ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error)
 	ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error)
 	ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error)
-	DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error)
-	DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode)
+	DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error)
+	DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode)
 	GetModelList() []string
 	GetChannelName() string
 }

relay/channel/ali/adaptor.go

@@ -32,14 +32,14 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fullRequestURL, nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", "Bearer "+info.ApiKey)
+	req.Set("Authorization", "Bearer "+info.ApiKey)
 	if info.IsStream {
-		req.Header.Set("X-DashScope-SSE", "enable")
+		req.Set("X-DashScope-SSE", "enable")
 	}
 	if c.GetString("plugin") != "" {
-		req.Header.Set("X-DashScope-Plugin", c.GetString("plugin"))
+		req.Set("X-DashScope-Plugin", c.GetString("plugin"))
 	}
 	return nil
 }
@@ -72,11 +72,11 @@ func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInf
 	return nil, errors.New("not implemented")
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	switch info.RelayMode {
 	case constant.RelayModeImagesGenerations:
 		err, usage = aliImageHandler(c, resp, info)

relay/channel/api_request.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"io"
 	"net/http"
 	"one-api/relay/common"
@@ -11,14 +12,16 @@ import (
 	"one-api/service"
 )
 
-func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Request) {
+func SetupApiRequestHeader(info *common.RelayInfo, c *gin.Context, req *http.Header) {
 	if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
 		// multipart/form-data
+	} else if info.RelayMode == constant.RelayModeRealtime {
+		// websocket
 	} else {
-		req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
-		req.Header.Set("Accept", c.Request.Header.Get("Accept"))
+		req.Set("Content-Type", c.Request.Header.Get("Content-Type"))
+		req.Set("Accept", c.Request.Header.Get("Accept"))
 		if info.IsStream && c.Request.Header.Get("Accept") == "" {
-			req.Header.Set("Accept", "text/event-stream")
+			req.Set("Accept", "text/event-stream")
 		}
 	}
 }
@@ -32,7 +35,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
 	if err != nil {
 		return nil, fmt.Errorf("new request failed: %w", err)
 	}
-	err = a.SetupRequestHeader(c, req, info)
+	err = a.SetupRequestHeader(c, &req.Header, info)
 	if err != nil {
 		return nil, fmt.Errorf("setup request header failed: %w", err)
 	}
@@ -55,7 +58,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
 	// set form data
 	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
 
-	err = a.SetupRequestHeader(c, req, info)
+	err = a.SetupRequestHeader(c, &req.Header, info)
 	if err != nil {
 		return nil, fmt.Errorf("setup request header failed: %w", err)
 	}
@@ -66,6 +69,27 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
 	return resp, nil
 }
 
+func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*websocket.Conn, error) {
+	fullRequestURL, err := a.GetRequestURL(info)
+	if err != nil {
+		return nil, fmt.Errorf("get request url failed: %w", err)
+	}
+	targetHeader := http.Header{}
+	err = a.SetupRequestHeader(c, &targetHeader, info)
+	if err != nil {
+		return nil, fmt.Errorf("setup request header failed: %w", err)
+	}
+	targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type"))
+	targetConn, _, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader)
+	if err != nil {
+		return nil, fmt.Errorf("dial failed to %s: %w", fullRequestURL, err)
+	}
+	// send request body
+	//all, err := io.ReadAll(requestBody)
+	//err = service.WssString(c, targetConn, string(all))
+	return targetConn, nil
+}
+
 func doRequest(c *gin.Context, req *http.Request) (*http.Response, error) {
 	resp, err := service.GetHttpClient().Do(req)
 	if err != nil {

relay/channel/aws/adaptor.go

@@ -8,7 +8,6 @@ import (
 	"one-api/dto"
 	"one-api/relay/channel/claude"
 	relaycommon "one-api/relay/common"
-	"strings"
 )
 
 const (
@@ -31,18 +30,14 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
 }
 
 func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
-	if strings.HasPrefix(info.UpstreamModelName, "claude-3") {
-		a.RequestMode = RequestModeMessage
-	} else {
-		a.RequestMode = RequestModeCompletion
-	}
+	a.RequestMode = RequestModeMessage
 }
 
 func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return "", nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	return nil
 }
 
@@ -53,11 +48,8 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re
 
 	var claudeReq *claude.ClaudeRequest
 	var err error
-	if a.RequestMode == RequestModeCompletion {
-		claudeReq = claude.RequestOpenAI2ClaudeComplete(*request)
-	} else {
-		claudeReq, err = claude.RequestOpenAI2ClaudeMessage(*request)
-	}
+	claudeReq, err = claude.RequestOpenAI2ClaudeMessage(*request)
+
 	c.Set("request_model", request.Model)
 	c.Set("converted_request", claudeReq)
 	return claudeReq, err
@@ -67,11 +59,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return nil, nil
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = awsStreamHandler(c, resp, info, a.RequestMode)
 	} else {

relay/channel/aws/constants.go

@@ -1,13 +1,14 @@
 package aws
 
 var awsModelIDMap = map[string]string{
-	"claude-instant-1.2":       "anthropic.claude-instant-v1",
-	"claude-2.0":               "anthropic.claude-v2",
-	"claude-2.1":               "anthropic.claude-v2:1",
-	"claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0",
-	"claude-3-opus-20240229":   "anthropic.claude-3-opus-20240229-v1:0",
-	"claude-3-haiku-20240307":  "anthropic.claude-3-haiku-20240307-v1:0",
+	"claude-instant-1.2":         "anthropic.claude-instant-v1",
+	"claude-2.0":                 "anthropic.claude-v2",
+	"claude-2.1":                 "anthropic.claude-v2:1",
+	"claude-3-sonnet-20240229":   "anthropic.claude-3-sonnet-20240229-v1:0",
+	"claude-3-opus-20240229":     "anthropic.claude-3-opus-20240229-v1:0",
+	"claude-3-haiku-20240307":    "anthropic.claude-3-haiku-20240307-v1:0",
 	"claude-3-5-sonnet-20240620": "anthropic.claude-3-5-sonnet-20240620-v1:0",
+	"claude-3-5-sonnet-20241022": "anthropic.claude-3-5-sonnet-20241022-v2:0",
 }
 
 var ChannelName = "aws"

relay/channel/aws/dto.go

@@ -1,15 +1,34 @@
 package aws
 
-import "one-api/relay/channel/claude"
+import (
+	"one-api/relay/channel/claude"
+)
 
 type AwsClaudeRequest struct {
 	// AnthropicVersion should be "bedrock-2023-05-31"
 	AnthropicVersion string                 `json:"anthropic_version"`
-	System           string                 `json:"system"`
+	System           string                 `json:"system,omitempty"`
 	Messages         []claude.ClaudeMessage `json:"messages"`
-	MaxTokens        int                    `json:"max_tokens,omitempty"`
+	MaxTokens        uint                   `json:"max_tokens,omitempty"`
 	Temperature      float64                `json:"temperature,omitempty"`
 	TopP             float64                `json:"top_p,omitempty"`
 	TopK             int                    `json:"top_k,omitempty"`
 	StopSequences    []string               `json:"stop_sequences,omitempty"`
+	Tools            []claude.Tool          `json:"tools,omitempty"`
+	ToolChoice       any                    `json:"tool_choice,omitempty"`
+}
+
+func copyRequest(req *claude.ClaudeRequest) *AwsClaudeRequest {
+	return &AwsClaudeRequest{
+		AnthropicVersion: "bedrock-2023-05-31",
+		System:           req.System,
+		Messages:         req.Messages,
+		MaxTokens:        req.MaxTokens,
+		Temperature:      req.Temperature,
+		TopP:             req.TopP,
+		TopK:             req.TopK,
+		StopSequences:    req.StopSequences,
+		Tools:            req.Tools,
+		ToolChoice:       req.ToolChoice,
+	}
 }

relay/channel/aws/relay-aws.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/gin-gonic/gin"
-	"github.com/jinzhu/copier"
 	"github.com/pkg/errors"
 	"io"
 	"net/http"
@@ -53,7 +52,7 @@ func awsModelID(requestModel string) (string, error) {
 		return awsModelID, nil
 	}
 
-	return "", errors.Errorf("model %s not found", requestModel)
+	return requestModel, nil
 }
 
 func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode int) (*relaymodel.OpenAIErrorWithStatusCode, *relaymodel.Usage) {
@@ -78,13 +77,7 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode int) (*
 		return wrapErr(errors.New("request not found")), nil
 	}
 	claudeReq := claudeReq_.(*claude.ClaudeRequest)
-	awsClaudeReq := &AwsClaudeRequest{
-		AnthropicVersion: "bedrock-2023-05-31",
-	}
-	if err = copier.Copy(awsClaudeReq, claudeReq); err != nil {
-		return wrapErr(errors.Wrap(err, "copy request")), nil
-	}
-
+	awsClaudeReq := copyRequest(claudeReq)
 	awsReq.Body, err = json.Marshal(awsClaudeReq)
 	if err != nil {
 		return wrapErr(errors.Wrap(err, "marshal request")), nil
@@ -136,12 +129,7 @@ func awsStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
 	}
 	claudeReq := claudeReq_.(*claude.ClaudeRequest)
 
-	awsClaudeReq := &AwsClaudeRequest{
-		AnthropicVersion: "bedrock-2023-05-31",
-	}
-	if err = copier.Copy(awsClaudeReq, claudeReq); err != nil {
-		return wrapErr(errors.Wrap(err, "copy request")), nil
-	}
+	awsClaudeReq := copyRequest(claudeReq)
 	awsReq.Body, err = json.Marshal(awsClaudeReq)
 	if err != nil {
 		return wrapErr(errors.Wrap(err, "marshal request")), nil

relay/channel/baidu/adaptor.go

@@ -98,9 +98,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fullRequestURL, nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", "Bearer "+info.ApiKey)
+	req.Set("Authorization", "Bearer "+info.ApiKey)
 	return nil
 }
 
@@ -122,11 +122,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = baiduStreamHandler(c, resp)
 	} else {

relay/channel/claude/adaptor.go

@@ -47,14 +47,14 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	}
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("x-api-key", info.ApiKey)
+	req.Set("x-api-key", info.ApiKey)
 	anthropicVersion := c.Request.Header.Get("anthropic-version")
 	if anthropicVersion == "" {
 		anthropicVersion = "2023-06-01"
 	}
-	req.Header.Set("anthropic-version", anthropicVersion)
+	req.Set("anthropic-version", anthropicVersion)
 	return nil
 }
 
@@ -73,11 +73,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = ClaudeStreamHandler(c, resp, info, a.RequestMode)
 	} else {

relay/channel/claude/constants.go

@@ -8,7 +8,9 @@ var ModelList = []string{
 	"claude-3-sonnet-20240229",
 	"claude-3-opus-20240229",
 	"claude-3-haiku-20240307",
+	"claude-3-5-haiku-20241022",
 	"claude-3-5-sonnet-20240620",
+	"claude-3-5-sonnet-20241022",
 }
 
 var ChannelName = "claude"

relay/channel/claude/relay-claude.go

@@ -509,7 +509,7 @@ func ClaudeHandler(c *gin.Context, resp *http.Response, requestMode int, info *r
 		}, nil
 	}
 	fullTextResponse := ResponseClaude2OpenAI(requestMode, &claudeResponse)
-	completionTokens, err := service.CountTokenText(claudeResponse.Completion, info.OriginModelName)
+	completionTokens, err := service.CountTextToken(claudeResponse.Completion, info.OriginModelName)
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "count_token_text_failed", http.StatusInternalServerError), nil
 	}

relay/channel/cloudflare/adaptor.go

@@ -30,9 +30,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	}
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
+	req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
 	return nil
 }
 
@@ -48,7 +48,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re
 	}
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
@@ -78,7 +78,7 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
 	return nil, errors.New("not implemented")
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	switch info.RelayMode {
 	case constant.RelayModeEmbeddings:
 		fallthrough

relay/channel/cloudflare/relay_cloudflare.go

@@ -149,7 +149,7 @@ func cfSTTHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayIn
 
 	usage := &dto.Usage{}
 	usage.PromptTokens = info.PromptTokens
-	usage.CompletionTokens, _ = service.CountTokenText(cfResp.Result.Text, info.UpstreamModelName)
+	usage.CompletionTokens, _ = service.CountTextToken(cfResp.Result.Text, info.UpstreamModelName)
 	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 
 	return nil, usage

relay/channel/cohere/adaptor.go

@@ -36,9 +36,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	}
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
+	req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
 	return nil
 }
 
@@ -46,7 +46,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re
 	return requestOpenAI2Cohere(*request), nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
@@ -54,7 +54,7 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return requestConvertRerank2Cohere(request), nil
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.RelayMode == constant.RelayModeRerank {
 		err, usage = cohereRerankHandler(c, resp, info)
 	} else {

relay/channel/dify/adaptor.go

@@ -31,9 +31,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/v1/chat-messages", info.BaseUrl), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", "Bearer "+info.ApiKey)
+	req.Set("Authorization", "Bearer "+info.ApiKey)
 	return nil
 }
 
@@ -48,11 +48,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = difyStreamHandler(c, resp, info)
 	} else {

relay/channel/dify/relay-dify.go

@@ -108,7 +108,7 @@ func difyStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Re
 	}
 	if usage.TotalTokens == 0 {
 		usage.PromptTokens = info.PromptTokens
-		usage.CompletionTokens, _ = service.CountTokenText("gpt-3.5-turbo", responseText)
+		usage.CompletionTokens, _ = service.CountTextToken("gpt-3.5-turbo", responseText)
 		usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 	}
 	return nil, usage

relay/channel/gemini/adaptor.go

@@ -30,13 +30,13 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
 }
 
 func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
-	// 从映射中获取模型名称对应的版本，如果找不到就使用 info.ApiVersion 或默认的版本 "v1"
+	// 从映射中获取模型名称对应的版本，如果找不到就使用 info.ApiVersion 或默认的版本 "v1beta"
 	version, beta := constant.GeminiModelMap[info.UpstreamModelName]
 	if !beta {
 		if info.ApiVersion != "" {
 			version = info.ApiVersion
 		} else {
-			version = "v1"
+			version = "v1beta"
 		}
 	}
 
@@ -47,9 +47,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/%s/models/%s:%s", info.BaseUrl, version, info.UpstreamModelName, action), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("x-goog-api-key", info.ApiKey)
+	req.Set("x-goog-api-key", info.ApiKey)
 	return nil
 }
 
@@ -64,11 +64,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = GeminiChatStreamHandler(c, resp, info)
 	} else {

relay/channel/gemini/constant.go

@@ -7,6 +7,7 @@ const (
 var ModelList = []string{
 	"gemini-1.0-pro-latest", "gemini-1.0-pro-001", "gemini-1.5-pro-latest", "gemini-1.5-flash-latest", "gemini-ultra",
 	"gemini-1.0-pro-vision-latest", "gemini-1.0-pro-vision-001", "gemini-1.5-pro-exp-0827", "gemini-1.5-flash-exp-0827",
+	"gemini-exp-1114",
 }
 
 var ChannelName = "google gemini"

relay/channel/jina/adaptor.go

@@ -37,9 +37,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return "", errors.New("invalid relay mode")
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
+	req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
 	return nil
 }
 
@@ -47,7 +47,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re
 	return request, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
@@ -55,7 +55,7 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return request, nil
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.RelayMode == constant.RelayModeRerank {
 		err, usage = jinaRerankHandler(c, resp)
 	} else if info.RelayMode == constant.RelayModeEmbeddings {

relay/channel/mistral/adaptor.go

@@ -0,0 +1,72 @@
+package mistral
+
+import (
+	"errors"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/dto"
+	"one-api/relay/channel"
+	"one-api/relay/channel/openai"
+	relaycommon "one-api/relay/common"
+)
+
+type Adaptor struct {
+}
+
+func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
+	//TODO implement me
+	return nil, errors.New("not implemented")
+}
+
+func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
+	//TODO implement me
+	return nil, errors.New("not implemented")
+}
+
+func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
+}
+
+func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
+	return relaycommon.GetFullRequestURL(info.BaseUrl, info.RequestURLPath, info.ChannelType), nil
+}
+
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
+	channel.SetupApiRequestHeader(info, c, req)
+	req.Set("Authorization", "Bearer "+info.ApiKey)
+	return nil
+}
+
+func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
+	if request == nil {
+		return nil, errors.New("request is nil")
+	}
+	mistralReq := requestOpenAI2Mistral(*request)
+	//common.LogJson(c, "body", mistralReq)
+	return mistralReq, nil
+}
+
+func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
+	return nil, nil
+}
+
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
+	return channel.DoApiRequest(a, c, info, requestBody)
+}
+
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
+	if info.IsStream {
+		err, usage = openai.OaiStreamHandler(c, resp, info)
+	} else {
+		err, usage = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
+	}
+	return
+}
+
+func (a *Adaptor) GetModelList() []string {
+	return ModelList
+}
+
+func (a *Adaptor) GetChannelName() string {
+	return ChannelName
+}

relay/channel/mistral/constants.go

@@ -0,0 +1,12 @@
+package mistral
+
+var ModelList = []string{
+	"open-mistral-7b",
+	"open-mixtral-8x7b",
+	"mistral-small-latest",
+	"mistral-medium-latest",
+	"mistral-large-latest",
+	"mistral-embed",
+}
+
+var ChannelName = "mistral"

relay/channel/mistral/text.go

@@ -0,0 +1,40 @@
+package mistral
+
+import (
+	"encoding/json"
+	"one-api/dto"
+)
+
+func requestOpenAI2Mistral(request dto.GeneralOpenAIRequest) *dto.GeneralOpenAIRequest {
+	messages := make([]dto.Message, 0, len(request.Messages))
+	for _, message := range request.Messages {
+		if !message.IsStringContent() {
+			mediaMessages := message.ParseContent()
+			for j, mediaMessage := range mediaMessages {
+				if mediaMessage.Type == dto.ContentTypeImageURL {
+					imageUrl := mediaMessage.ImageUrl.(dto.MessageImageUrl)
+					mediaMessage.ImageUrl = imageUrl.Url
+					mediaMessages[j] = mediaMessage
+				}
+			}
+			messageRaw, _ := json.Marshal(mediaMessages)
+			message.Content = messageRaw
+		}
+		messages = append(messages, dto.Message{
+			Role:       message.Role,
+			Content:    message.Content,
+			ToolCalls:  message.ToolCalls,
+			ToolCallId: message.ToolCallId,
+		})
+	}
+	return &dto.GeneralOpenAIRequest{
+		Model:       request.Model,
+		Stream:      request.Stream,
+		Messages:    messages,
+		Temperature: request.Temperature,
+		TopP:        request.TopP,
+		MaxTokens:   request.MaxTokens,
+		Tools:       request.Tools,
+		ToolChoice:  request.ToolChoice,
+	}
+}

relay/channel/ollama/adaptor.go

@@ -31,13 +31,13 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
 func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	switch info.RelayMode {
 	case relayconstant.RelayModeEmbeddings:
-		return info.BaseUrl + "/api/embeddings", nil
+		return info.BaseUrl + "/api/embed", nil
 	default:
 		return relaycommon.GetFullRequestURL(info.BaseUrl, info.RequestURLPath, info.ChannelType), nil
 	}
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
 	return nil
 }
@@ -58,11 +58,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = openai.OaiStreamHandler(c, resp, info)
 	} else {

relay/channel/ollama/dto.go

@@ -37,5 +37,5 @@ type OllamaEmbeddingRequest struct {
 type OllamaEmbeddingResponse struct {
 	Error     string    `json:"error,omitempty"`
 	Model     string    `json:"model"`
-	Embedding []float64 `json:"embedding,omitempty"`
+	Embedding [][]float64 `json:"embeddings,omitempty"`
 }

relay/channel/ollama/relay-ollama.go

@@ -73,9 +73,10 @@ func ollamaEmbeddingHandler(c *gin.Context, resp *http.Response, promptTokens in
 	if ollamaEmbeddingResponse.Error != "" {
 		return service.OpenAIErrorWrapper(err, "ollama_error", resp.StatusCode), nil
 	}
+	flattenedEmbeddings := flattenEmbeddings(ollamaEmbeddingResponse.Embedding)
 	data := make([]dto.OpenAIEmbeddingResponseItem, 0, 1)
 	data = append(data, dto.OpenAIEmbeddingResponseItem{
-		Embedding: ollamaEmbeddingResponse.Embedding,
+		Embedding: flattenedEmbeddings,
 		Object:    "embedding",
 	})
 	usage := &dto.Usage{
@@ -120,3 +121,11 @@ func ollamaEmbeddingHandler(c *gin.Context, resp *http.Response, promptTokens in
 	}
 	return nil, usage
 }
+
+func flattenEmbeddings(embeddings [][]float64) []float64 {
+flattened := []float64{}
+for _, row := range embeddings {
+	flattened = append(flattened, row...)
+}
+return flattened
+}

relay/channel/openai/adaptor.go

@@ -31,6 +31,13 @@ func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
 }
 
 func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
+	if info.RelayMode == constant.RelayModeRealtime {
+		// trim https
+		baseUrl := strings.TrimPrefix(info.BaseUrl, "https://")
+		baseUrl = strings.TrimPrefix(baseUrl, "http://")
+		baseUrl = "wss://" + baseUrl
+		info.BaseUrl = baseUrl
+	}
 	switch info.ChannelType {
 	case common.ChannelTypeAzure:
 		// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
@@ -40,8 +47,10 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 		model_ := info.UpstreamModelName
 		model_ = strings.Replace(model_, ".", "", -1)
 		// https://github.com/songquanpeng/one-api/issues/67
-
 		requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task)
+		if info.RelayMode == constant.RelayModeRealtime {
+			requestURL = fmt.Sprintf("/openai/realtime?deployment=%s&api-version=%s", model_, info.ApiVersion)
+		}
 		return relaycommon.GetFullRequestURL(info.BaseUrl, requestURL, info.ChannelType), nil
 	case common.ChannelTypeMiniMax:
 		return minimax.GetRequestURL(info)
@@ -54,16 +63,34 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	}
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
-	channel.SetupApiRequestHeader(info, c, req)
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info *relaycommon.RelayInfo) error {
+	channel.SetupApiRequestHeader(info, c, header)
 	if info.ChannelType == common.ChannelTypeAzure {
-		req.Header.Set("api-key", info.ApiKey)
+		header.Set("api-key", info.ApiKey)
 		return nil
 	}
 	if info.ChannelType == common.ChannelTypeOpenAI && "" != info.Organization {
-		req.Header.Set("OpenAI-Organization", info.Organization)
+		header.Set("OpenAI-Organization", info.Organization)
+	}
+	if info.RelayMode == constant.RelayModeRealtime {
+		swp := c.Request.Header.Get("Sec-WebSocket-Protocol")
+		if swp != "" {
+			items := []string{
+				"realtime",
+				"openai-insecure-api-key." + info.ApiKey,
+				"openai-beta.realtime-v1",
+			}
+			header.Set("Sec-WebSocket-Protocol", strings.Join(items, ","))
+			//req.Header.Set("Sec-WebSocket-Key", c.Request.Header.Get("Sec-WebSocket-Key"))
+			//req.Header.Set("Sec-Websocket-Extensions", c.Request.Header.Get("Sec-Websocket-Extensions"))
+			//req.Header.Set("Sec-Websocket-Version", c.Request.Header.Get("Sec-Websocket-Version"))
+		} else {
+			header.Set("openai-beta", "realtime=v1")
+			header.Set("Authorization", "Bearer "+info.ApiKey)
+		}
+	} else {
+		header.Set("Authorization", "Bearer "+info.ApiKey)
 	}
-	req.Header.Set("Authorization", "Bearer "+info.ApiKey)
 	//if info.ChannelType == common.ChannelTypeOpenRouter {
 	//	req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api")
 	//	req.Header.Set("X-Title", "One API")
@@ -78,6 +105,12 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re
 	if info.ChannelType != common.ChannelTypeOpenAI {
 		request.StreamOptions = nil
 	}
+	if strings.HasPrefix(request.Model, "o1-") {
+		if request.MaxCompletionTokens == 0 && request.MaxTokens != 0 {
+			request.MaxCompletionTokens = request.MaxTokens
+			request.MaxTokens = 0
+		}
+	}
 	return request, nil
 }
 
@@ -125,16 +158,20 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
 	return request, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	if info.RelayMode == constant.RelayModeAudioTranscription || info.RelayMode == constant.RelayModeAudioTranslation {
 		return channel.DoFormRequest(a, c, info, requestBody)
+	} else if info.RelayMode == constant.RelayModeRealtime {
+		return channel.DoWssRequest(a, c, info, requestBody)
 	} else {
 		return channel.DoApiRequest(a, c, info, requestBody)
 	}
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	switch info.RelayMode {
+	case constant.RelayModeRealtime:
+		err, usage = OpenaiRealtimeHandler(c, info)
 	case constant.RelayModeAudioSpeech:
 		err, usage = OpenaiTTSHandler(c, resp, info)
 	case constant.RelayModeAudioTranslation:

relay/channel/openai/constant.go

@@ -13,6 +13,8 @@ var ModelList = []string{
 	"gpt-4o-mini", "gpt-4o-mini-2024-07-18",
 	"o1-preview", "o1-preview-2024-09-12",
 	"o1-mini", "o1-mini-2024-09-12",
+	"gpt-4o-audio-preview", "gpt-4o-audio-preview-2024-10-01",
+	"gpt-4o-realtime-preview", "gpt-4o-realtime-preview-2024-10-01",
 	"text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large",
 	"text-curie-001", "text-babbage-001", "text-ada-001",
 	"text-moderation-latest", "text-moderation-stable",

relay/channel/openai/relay-openai.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"github.com/bytedance/gopkg/util/gopool"
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"io"
 	"net/http"
 	"one-api/common"
@@ -231,7 +232,7 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model
 	if simpleResponse.Usage.TotalTokens == 0 || (simpleResponse.Usage.PromptTokens == 0 && simpleResponse.Usage.CompletionTokens == 0) {
 		completionTokens := 0
 		for _, choice := range simpleResponse.Choices {
-			ctkm, _ := service.CountTokenText(string(choice.Message.Content), model)
+			ctkm, _ := service.CountTextToken(string(choice.Message.Content), model)
 			completionTokens += ctkm
 		}
 		simpleResponse.Usage = dto.Usage{
@@ -324,7 +325,7 @@ func OpenaiSTTHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
 
 	usage := &dto.Usage{}
 	usage.PromptTokens = info.PromptTokens
-	usage.CompletionTokens, _ = service.CountTokenText(text, info.UpstreamModelName)
+	usage.CompletionTokens, _ = service.CountTextToken(text, info.UpstreamModelName)
 	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 	return nil, usage
 }
@@ -373,3 +374,210 @@ func getTextFromJSON(body []byte) (string, error) {
 	}
 	return whisperResponse.Text, nil
 }
+
+func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, *dto.RealtimeUsage) {
+	info.IsStream = true
+	clientConn := info.ClientWs
+	targetConn := info.TargetWs
+
+	clientClosed := make(chan struct{})
+	targetClosed := make(chan struct{})
+	sendChan := make(chan []byte, 100)
+	receiveChan := make(chan []byte, 100)
+	errChan := make(chan error, 2)
+
+	usage := &dto.RealtimeUsage{}
+	localUsage := &dto.RealtimeUsage{}
+	sumUsage := &dto.RealtimeUsage{}
+
+	gopool.Go(func() {
+		for {
+			select {
+			case <-c.Done():
+				return
+			default:
+				_, message, err := clientConn.ReadMessage()
+				if err != nil {
+					if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
+						errChan <- fmt.Errorf("error reading from client: %v", err)
+					}
+					close(clientClosed)
+					return
+				}
+
+				realtimeEvent := &dto.RealtimeEvent{}
+				err = json.Unmarshal(message, realtimeEvent)
+				if err != nil {
+					errChan <- fmt.Errorf("error unmarshalling message: %v", err)
+					return
+				}
+
+				if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdate {
+					if realtimeEvent.Session != nil {
+						if realtimeEvent.Session.Tools != nil {
+							info.RealtimeTools = realtimeEvent.Session.Tools
+						}
+					}
+				}
+
+				textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
+				if err != nil {
+					errChan <- fmt.Errorf("error counting text token: %v", err)
+					return
+				}
+				common.LogInfo(c, fmt.Sprintf("type: %s, textToken: %d, audioToken: %d", realtimeEvent.Type, textToken, audioToken))
+				localUsage.TotalTokens += textToken + audioToken
+				localUsage.InputTokens += textToken + audioToken
+				localUsage.InputTokenDetails.TextTokens += textToken
+				localUsage.InputTokenDetails.AudioTokens += audioToken
+
+				err = service.WssString(c, targetConn, string(message))
+				if err != nil {
+					errChan <- fmt.Errorf("error writing to target: %v", err)
+					return
+				}
+
+				select {
+				case sendChan <- message:
+				default:
+				}
+			}
+		}
+	})
+
+	gopool.Go(func() {
+		for {
+			select {
+			case <-c.Done():
+				return
+			default:
+				_, message, err := targetConn.ReadMessage()
+				if err != nil {
+					if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
+						errChan <- fmt.Errorf("error reading from target: %v", err)
+					}
+					close(targetClosed)
+					return
+				}
+				info.SetFirstResponseTime()
+				realtimeEvent := &dto.RealtimeEvent{}
+				err = json.Unmarshal(message, realtimeEvent)
+				if err != nil {
+					errChan <- fmt.Errorf("error unmarshalling message: %v", err)
+					return
+				}
+
+				if realtimeEvent.Type == dto.RealtimeEventTypeResponseDone {
+					realtimeUsage := realtimeEvent.Response.Usage
+					if realtimeUsage != nil {
+						usage.TotalTokens += realtimeUsage.TotalTokens
+						usage.InputTokens += realtimeUsage.InputTokens
+						usage.OutputTokens += realtimeUsage.OutputTokens
+						usage.InputTokenDetails.AudioTokens += realtimeUsage.InputTokenDetails.AudioTokens
+						usage.InputTokenDetails.CachedTokens += realtimeUsage.InputTokenDetails.CachedTokens
+						usage.InputTokenDetails.TextTokens += realtimeUsage.InputTokenDetails.TextTokens
+						usage.OutputTokenDetails.AudioTokens += realtimeUsage.OutputTokenDetails.AudioTokens
+						usage.OutputTokenDetails.TextTokens += realtimeUsage.OutputTokenDetails.TextTokens
+						err := preConsumeUsage(c, info, usage, sumUsage)
+						if err != nil {
+							errChan <- fmt.Errorf("error consume usage: %v", err)
+							return
+						}
+						// 本次计费完成，清除
+						usage = &dto.RealtimeUsage{}
+
+						localUsage = &dto.RealtimeUsage{}
+					} else {
+						textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
+						if err != nil {
+							errChan <- fmt.Errorf("error counting text token: %v", err)
+							return
+						}
+						common.LogInfo(c, fmt.Sprintf("type: %s, textToken: %d, audioToken: %d", realtimeEvent.Type, textToken, audioToken))
+						localUsage.TotalTokens += textToken + audioToken
+						info.IsFirstRequest = false
+						localUsage.InputTokens += textToken + audioToken
+						localUsage.InputTokenDetails.TextTokens += textToken
+						localUsage.InputTokenDetails.AudioTokens += audioToken
+						err = preConsumeUsage(c, info, localUsage, sumUsage)
+						if err != nil {
+							errChan <- fmt.Errorf("error consume usage: %v", err)
+							return
+						}
+						// 本次计费完成，清除
+						localUsage = &dto.RealtimeUsage{}
+						// print now usage
+					}
+					//common.LogInfo(c, fmt.Sprintf("realtime streaming sumUsage: %v", sumUsage))
+					//common.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage))
+					//common.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage))
+
+				} else if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdated || realtimeEvent.Type == dto.RealtimeEventTypeSessionCreated {
+					realtimeSession := realtimeEvent.Session
+					if realtimeSession != nil {
+						// update audio format
+						info.InputAudioFormat = common.GetStringIfEmpty(realtimeSession.InputAudioFormat, info.InputAudioFormat)
+						info.OutputAudioFormat = common.GetStringIfEmpty(realtimeSession.OutputAudioFormat, info.OutputAudioFormat)
+					}
+				} else {
+					textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
+					if err != nil {
+						errChan <- fmt.Errorf("error counting text token: %v", err)
+						return
+					}
+					common.LogInfo(c, fmt.Sprintf("type: %s, textToken: %d, audioToken: %d", realtimeEvent.Type, textToken, audioToken))
+					localUsage.TotalTokens += textToken + audioToken
+					localUsage.OutputTokens += textToken + audioToken
+					localUsage.OutputTokenDetails.TextTokens += textToken
+					localUsage.OutputTokenDetails.AudioTokens += audioToken
+				}
+
+				err = service.WssString(c, clientConn, string(message))
+				if err != nil {
+					errChan <- fmt.Errorf("error writing to client: %v", err)
+					return
+				}
+
+				select {
+				case receiveChan <- message:
+				default:
+				}
+			}
+		}
+	})
+
+	select {
+	case <-clientClosed:
+	case <-targetClosed:
+	case err := <-errChan:
+		//return service.OpenAIErrorWrapper(err, "realtime_error", http.StatusInternalServerError), nil
+		common.LogError(c, "realtime error: "+err.Error())
+	case <-c.Done():
+	}
+
+	if usage.TotalTokens != 0 {
+		_ = preConsumeUsage(c, info, usage, sumUsage)
+	}
+
+	if localUsage.TotalTokens != 0 {
+		_ = preConsumeUsage(c, info, localUsage, sumUsage)
+	}
+
+	// check usage total tokens, if 0, use local usage
+
+	return nil, sumUsage
+}
+
+func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.RealtimeUsage, totalUsage *dto.RealtimeUsage) error {
+	totalUsage.TotalTokens += usage.TotalTokens
+	totalUsage.InputTokens += usage.InputTokens
+	totalUsage.OutputTokens += usage.OutputTokens
+	totalUsage.InputTokenDetails.CachedTokens += usage.InputTokenDetails.CachedTokens
+	totalUsage.InputTokenDetails.TextTokens += usage.InputTokenDetails.TextTokens
+	totalUsage.InputTokenDetails.AudioTokens += usage.InputTokenDetails.AudioTokens
+	totalUsage.OutputTokenDetails.TextTokens += usage.OutputTokenDetails.TextTokens
+	totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens
+	// clear usage
+	err := service.PreWssConsumeQuota(ctx, info, usage)
+	return err
+}

relay/channel/palm/adaptor.go

@@ -32,9 +32,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/v1beta2/models/chat-bison-001:generateMessage", info.BaseUrl), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("x-goog-api-key", info.ApiKey)
+	req.Set("x-goog-api-key", info.ApiKey)
 	return nil
 }
 
@@ -49,11 +49,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		var responseText string
 		err, responseText = palmStreamHandler(c, resp)

relay/channel/palm/relay-palm.go

@@ -156,7 +156,7 @@ func palmHandler(c *gin.Context, resp *http.Response, promptTokens int, model st
 		}, nil
 	}
 	fullTextResponse := responsePaLM2OpenAI(&palmResponse)
-	completionTokens, _ := service.CountTokenText(palmResponse.Candidates[0].Content, model)
+	completionTokens, _ := service.CountTextToken(palmResponse.Candidates[0].Content, model)
 	usage := dto.Usage{
 		PromptTokens:     promptTokens,
 		CompletionTokens: completionTokens,

relay/channel/perplexity/adaptor.go

@@ -32,9 +32,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/chat/completions", info.BaseUrl), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", "Bearer "+info.ApiKey)
+	req.Set("Authorization", "Bearer "+info.ApiKey)
 	return nil
 }
 
@@ -52,11 +52,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = openai.OaiStreamHandler(c, resp, info)
 	} else {

relay/channel/siliconflow/adaptor.go

@@ -40,9 +40,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return "", errors.New("invalid relay mode")
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
+	req.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
 	return nil
 }
 
@@ -50,7 +50,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, re
 	return request, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
@@ -58,7 +58,7 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return request, nil
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	switch info.RelayMode {
 	case constant.RelayModeRerank:
 		err, usage = siliconflowRerankHandler(c, resp)

relay/channel/tencent/adaptor.go

@@ -43,12 +43,12 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/", info.BaseUrl), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
-	req.Header.Set("Authorization", a.Sign)
-	req.Header.Set("X-TC-Action", a.Action)
-	req.Header.Set("X-TC-Version", a.Version)
-	req.Header.Set("X-TC-Timestamp", strconv.FormatInt(a.Timestamp, 10))
+	req.Set("Authorization", a.Sign)
+	req.Set("X-TC-Action", a.Action)
+	req.Set("X-TC-Version", a.Version)
+	req.Set("X-TC-Timestamp", strconv.FormatInt(a.Timestamp, 10))
 	return nil
 }
 
@@ -73,11 +73,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		var responseText string
 		err, responseText = tencentStreamHandler(c, resp)

relay/channel/vertex/adaptor.go

@@ -107,13 +107,13 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return "", errors.New("unsupported request mode")
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
 	accessToken, err := getAccessToken(a, info)
 	if err != nil {
 		return err
 	}
-	req.Header.Set("Authorization", "Bearer "+accessToken)
+	req.Set("Authorization", "Bearer "+accessToken)
 	return nil
 }
 
@@ -148,11 +148,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		switch a.RequestMode {
 		case RequestModeClaude:

relay/channel/xunfei/adaptor.go

@@ -33,7 +33,7 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return "", nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
 	return nil
 }
@@ -50,14 +50,14 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	// xunfei's request is not http request, so we don't need to do anything here
 	dummyResp := &http.Response{}
 	dummyResp.StatusCode = http.StatusOK
 	return dummyResp, nil
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	splits := strings.Split(info.ApiKey, "|")
 	if len(splits) != 3 {
 		return nil, service.OpenAIErrorWrapper(errors.New("invalid auth"), "invalid_auth", http.StatusBadRequest)

relay/channel/zhipu/adaptor.go

@@ -35,10 +35,10 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/api/paas/v3/model-api/%s/%s", info.BaseUrl, info.UpstreamModelName, method), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
 	token := getZhipuToken(info.ApiKey)
-	req.Header.Set("Authorization", token)
+	req.Set("Authorization", token)
 	return nil
 }
 
@@ -56,11 +56,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = zhipuStreamHandler(c, resp)
 	} else {

relay/channel/zhipu_4v/adaptor.go

@@ -32,10 +32,10 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	return fmt.Sprintf("%s/api/paas/v4/chat/completions", info.BaseUrl), nil
 }
 
-func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
 	channel.SetupApiRequestHeader(info, c, req)
 	token := getZhipuToken(info.ApiKey)
-	req.Header.Set("Authorization", token)
+	req.Set("Authorization", token)
 	return nil
 }
 
@@ -53,11 +53,11 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 	return nil, nil
 }
 
-func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
 }
 
-func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		err, usage = openai.OaiStreamHandler(c, resp, info)
 	} else {

relay/common/relay_info.go

@@ -2,7 +2,9 @@ package common
 
 import (
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"one-api/common"
+	"one-api/dto"
 	"one-api/relay/constant"
 	"strings"
 	"time"
@@ -20,6 +22,8 @@ type RelayInfo struct {
 	setFirstResponse     bool
 	ApiType              int
 	IsStream             bool
+	IsPlayground         bool
+	UsePrice             bool
 	RelayMode            int
 	UpstreamModelName    string
 	OriginModelName      string
@@ -31,6 +35,22 @@ type RelayInfo struct {
 	BaseUrl              string
 	SupportStreamOptions bool
 	ShouldIncludeUsage   bool
+	ClientWs             *websocket.Conn
+	TargetWs             *websocket.Conn
+	InputAudioFormat     string
+	OutputAudioFormat    string
+	RealtimeTools        []dto.RealTimeTool
+	IsFirstRequest       bool
+	AudioUsage           bool
+}
+
+func GenRelayInfoWs(c *gin.Context, ws *websocket.Conn) *RelayInfo {
+	info := GenRelayInfo(c)
+	info.ClientWs = ws
+	info.InputAudioFormat = "pcm16"
+	info.OutputAudioFormat = "pcm16"
+	info.IsFirstRequest = true
+	return info
 }
 
 func GenRelayInfo(c *gin.Context) *RelayInfo {
@@ -65,6 +85,11 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 		ApiKey:            strings.TrimPrefix(c.Request.Header.Get("Authorization"), "Bearer "),
 		Organization:      c.GetString("channel_organization"),
 	}
+	if strings.HasPrefix(c.Request.URL.Path, "/pg") {
+		info.IsPlayground = true
+		info.RequestURLPath = strings.TrimPrefix(info.RequestURLPath, "/pg")
+		info.RequestURLPath = "/v1" + info.RequestURLPath
+	}
 	if info.BaseUrl == "" {
 		info.BaseUrl = common.ChannelBaseURLs[channelType]
 	}
@@ -146,3 +171,20 @@ func GenTaskRelayInfo(c *gin.Context) *TaskRelayInfo {
 	}
 	return info
 }
+
+func (info *TaskRelayInfo) ToRelayInfo() *RelayInfo {
+	return &RelayInfo{
+		ChannelType:       info.ChannelType,
+		ChannelId:         info.ChannelId,
+		TokenId:           info.TokenId,
+		UserId:            info.UserId,
+		Group:             info.Group,
+		StartTime:         info.StartTime,
+		ApiType:           info.ApiType,
+		RelayMode:         info.RelayMode,
+		UpstreamModelName: info.UpstreamModelName,
+		RequestURLPath:    info.RequestURLPath,
+		ApiKey:            info.ApiKey,
+		BaseUrl:           info.BaseUrl,
+	}
+}

relay/constant/api_type.go

@@ -25,6 +25,7 @@ const (
 	APITypeCloudflare
 	APITypeSiliconFlow
 	APITypeVertexAi
+	APITypeMistral
 
 	APITypeDummy // this one is only for count, do not add any channel after this
 )
@@ -72,6 +73,8 @@ func ChannelType2APIType(channelType int) (int, bool) {
 		apiType = APITypeSiliconFlow
 	case common.ChannelTypeVertexAi:
 		apiType = APITypeVertexAi
+	case common.ChannelTypeMistral:
+		apiType = APITypeMistral
 	}
 	if apiType == -1 {
 		return APITypeOpenAI, false

relay/constant/relay_mode.go

@@ -38,11 +38,13 @@ const (
 	RelayModeSunoSubmit
 
 	RelayModeRerank
+
+	RelayModeRealtime
 )
 
 func Path2RelayMode(path string) int {
 	relayMode := RelayModeUnknown
-	if strings.HasPrefix(path, "/v1/chat/completions") {
+	if strings.HasPrefix(path, "/v1/chat/completions") || strings.HasPrefix(path, "/pg/chat/completions") {
 		relayMode = RelayModeChatCompletions
 	} else if strings.HasPrefix(path, "/v1/completions") {
 		relayMode = RelayModeCompletions
@@ -64,6 +66,8 @@ func Path2RelayMode(path string) int {
 		relayMode = RelayModeAudioTranslation
 	} else if strings.HasPrefix(path, "/v1/rerank") {
 		relayMode = RelayModeRerank
+	} else if strings.HasPrefix(path, "/v1/realtime") {
+		relayMode = RelayModeRealtime
 	}
 	return relayMode
 }

relay/relay-audio.go

@@ -46,7 +46,7 @@ func getAndValidAudioRequest(c *gin.Context, info *relaycommon.RelayInfo) (*dto.
 	return audioRequest, nil
 }
 
-func AudioHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
+func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
 	relayInfo := relaycommon.GenRelayInfo(c)
 	audioRequest, err := getAndValidAudioRequest(c, relayInfo)
 
@@ -58,7 +58,7 @@ func AudioHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 	promptTokens := 0
 	preConsumedTokens := common.PreConsumedQuota
 	if relayInfo.RelayMode == relayconstant.RelayModeAudioSpeech {
-		promptTokens, err = service.CountAudioToken(audioRequest.Input, audioRequest.Model)
+		promptTokens, err = service.CountTTSToken(audioRequest.Input, audioRequest.Model)
 		if err != nil {
 			return service.OpenAIErrorWrapper(err, "count_audio_token_failed", http.StatusInternalServerError)
 		}
@@ -87,11 +87,16 @@ func AudioHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 		preConsumedQuota = 0
 	}
 	if preConsumedQuota > 0 {
-		userQuota, err = model.PreConsumeTokenQuota(relayInfo.TokenId, preConsumedQuota)
+		userQuota, err = model.PreConsumeTokenQuota(relayInfo, preConsumedQuota)
 		if err != nil {
 			return service.OpenAIErrorWrapperLocal(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
 	}
+	defer func() {
+		if openaiErr != nil {
+			returnPreConsumedQuota(c, relayInfo, userQuota, preConsumedQuota)
+		}
+	}()
 
 	// map model name
 	modelMapping := c.GetString("model_mapping")
@@ -122,27 +127,27 @@ func AudioHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
 	}
-
 	statusCodeMappingStr := c.GetString("status_code_mapping")
+
+	var httpResp *http.Response
 	if resp != nil {
-		if resp.StatusCode != http.StatusOK {
-			returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
-			openaiErr := service.RelayErrorHandler(resp)
+		httpResp = resp.(*http.Response)
+		if httpResp.StatusCode != http.StatusOK {
+			openaiErr = service.RelayErrorHandler(httpResp)
 			// reset status code 重置状态码
 			service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 			return openaiErr
 		}
 	}
 
-	usage, openaiErr := adaptor.DoResponse(c, resp, relayInfo)
+	usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo)
 	if openaiErr != nil {
-		returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
 		// reset status code 重置状态码
 		service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 		return openaiErr
 	}
 
-	postConsumeQuota(c, relayInfo, audioRequest.Model, usage, ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, 0, false, "")
+	postConsumeQuota(c, relayInfo, audioRequest.Model, usage.(*dto.Usage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, 0, false, "")
 
 	return nil
 }

relay/relay-image.go

@@ -149,22 +149,24 @@ func ImageHelper(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode {
 	requestBody = bytes.NewBuffer(jsonData)
 
 	statusCodeMappingStr := c.GetString("status_code_mapping")
+
 	resp, err := adaptor.DoRequest(c, relayInfo, requestBody)
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
 	}
-
+	var httpResp *http.Response
 	if resp != nil {
-		relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
-		if resp.StatusCode != http.StatusOK {
-			openaiErr := service.RelayErrorHandler(resp)
+		httpResp = resp.(*http.Response)
+		relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
+		if httpResp.StatusCode != http.StatusOK {
+			openaiErr := service.RelayErrorHandler(httpResp)
 			// reset status code 重置状态码
 			service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 			return openaiErr
 		}
 	}
 
-	_, openaiErr := adaptor.DoResponse(c, resp, relayInfo)
+	_, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo)
 	if openaiErr != nil {
 		// reset status code 重置状态码
 		service.ResetStatusCode(openaiErr, statusCodeMappingStr)

relay/relay-mj.go

@@ -12,6 +12,7 @@ import (
 	"one-api/constant"
 	"one-api/dto"
 	"one-api/model"
+	relaycommon "one-api/relay/common"
 	relayconstant "one-api/relay/constant"
 	"one-api/service"
 	"strconv"
@@ -146,6 +147,7 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
 	userId := c.GetInt("id")
 	group := c.GetString("group")
 	channelId := c.GetInt("channel_id")
+	relayInfo := relaycommon.GenRelayInfo(c)
 	var swapFaceRequest dto.SwapFaceRequest
 	err := common.UnmarshalBodyReusable(c, &swapFaceRequest)
 	if err != nil {
@@ -191,7 +193,7 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
 	}
 	defer func(ctx context.Context) {
 		if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 {
-			err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0, true)
+			err := model.PostConsumeTokenQuota(relayInfo, userQuota, quota, 0, true)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
@@ -356,6 +358,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons
 	userId := c.GetInt("id")
 	group := c.GetString("group")
 	channelId := c.GetInt("channel_id")
+	relayInfo := relaycommon.GenRelayInfo(c)
 	consumeQuota := true
 	var midjRequest dto.MidjourneyRequest
 	err := common.UnmarshalBodyReusable(c, &midjRequest)
@@ -495,7 +498,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons
 
 	defer func(ctx context.Context) {
 		if consumeQuota && midjResponseWithStatus.StatusCode == 200 {
-			err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0, true)
+			err := model.PostConsumeTokenQuota(relayInfo, userQuota, quota, 0, true)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}

relay/relay-text.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/bytedance/sonic"
 	"io"
 	"math"
 	"net/http"
@@ -64,7 +65,7 @@ func getAndValidateTextRequest(c *gin.Context, relayInfo *relaycommon.RelayInfo)
 	return textRequest, nil
 }
 
-func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
+func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
 
 	relayInfo := relaycommon.GenRelayInfo(c)
 
@@ -76,6 +77,7 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 	}
 
 	// map model name
+	//isModelMapped := false
 	modelMapping := c.GetString("model_mapping")
 	//isModelMapped := false
 	if modelMapping != "" && modelMapping != "{}" {
@@ -85,6 +87,7 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 			return service.OpenAIErrorWrapperLocal(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError)
 		}
 		if modelMap[textRequest.Model] != "" {
+			//isModelMapped = true
 			textRequest.Model = modelMap[textRequest.Model]
 			// set upstream model name
 			//isModelMapped = true
@@ -129,7 +132,11 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 	if openaiErr != nil {
 		return openaiErr
 	}
-
+	defer func() {
+		if openaiErr != nil {
+			returnPreConsumedQuota(c, relayInfo, userQuota, preConsumedQuota)
+		}
+	}()
 	includeUsage := false
 	// 判断用户是否需要返回使用情况
 	if textRequest.StreamOptions != nil && textRequest.StreamOptions.IncludeUsage {
@@ -159,41 +166,56 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 	adaptor.Init(relayInfo)
 	var requestBody io.Reader
 
+	//if relayInfo.ChannelType == common.ChannelTypeOpenAI && !isModelMapped {
+	//	body, err := common.GetRequestBody(c)
+	//	if err != nil {
+	//		return service.OpenAIErrorWrapperLocal(err, "get_request_body_failed", http.StatusInternalServerError)
+	//	}
+	//	requestBody = bytes.NewBuffer(body)
+	//} else {
+	//
+	//}
+
 	convertedRequest, err := adaptor.ConvertRequest(c, relayInfo, textRequest)
 	if err != nil {
 		return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError)
 	}
-	jsonData, err := json.Marshal(convertedRequest)
+	jsonData, err := sonic.Marshal(convertedRequest)
 	if err != nil {
 		return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError)
 	}
 	requestBody = bytes.NewBuffer(jsonData)
 
 	statusCodeMappingStr := c.GetString("status_code_mapping")
+	var httpResp *http.Response
 	resp, err := adaptor.DoRequest(c, relayInfo, requestBody)
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
 	}
 
 	if resp != nil {
-		relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
-		if resp.StatusCode != http.StatusOK {
-			returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
-			openaiErr := service.RelayErrorHandler(resp)
+		httpResp = resp.(*http.Response)
+		relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
+		if httpResp.StatusCode != http.StatusOK {
+			openaiErr = service.RelayErrorHandler(httpResp)
 			// reset status code 重置状态码
 			service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 			return openaiErr
 		}
 	}
 
-	usage, openaiErr := adaptor.DoResponse(c, resp, relayInfo)
+	usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo)
 	if openaiErr != nil {
-		returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
 		// reset status code 重置状态码
 		service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 		return openaiErr
 	}
-	postConsumeQuota(c, relayInfo, textRequest.Model, usage, ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, getModelPriceSuccess, "")
+
+	if strings.HasPrefix(relayInfo.UpstreamModelName, "gpt-4o-audio") {
+		service.PostAudioConsumeQuota(c, relayInfo, usage.(*dto.Usage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, getModelPriceSuccess, "")
+	} else {
+		postConsumeQuota(c, relayInfo, textRequest.Model, usage.(*dto.Usage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, getModelPriceSuccess, "")
+	}
 	return nil
 }
 
@@ -266,7 +288,7 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo
 		}
 	}
 	if preConsumedQuota > 0 {
-		userQuota, err = model.PreConsumeTokenQuota(relayInfo.TokenId, preConsumedQuota)
+		userQuota, err = model.PreConsumeTokenQuota(relayInfo, preConsumedQuota)
 		if err != nil {
 			return 0, 0, service.OpenAIErrorWrapperLocal(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
@@ -274,11 +296,11 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo
 	return preConsumedQuota, userQuota, nil
 }
 
-func returnPreConsumedQuota(c *gin.Context, tokenId int, userQuota int, preConsumedQuota int) {
+func returnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, userQuota int, preConsumedQuota int) {
 	if preConsumedQuota != 0 {
 		go func(ctx context.Context) {
 			// return pre-consumed quota
-			err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0, false)
+			err := model.PostConsumeTokenQuota(relayInfo, userQuota, -preConsumedQuota, 0, false)
 			if err != nil {
 				common.SysError("error return pre-consumed quota: " + err.Error())
 			}
@@ -336,7 +358,7 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelN
 		//}
 		quotaDelta := quota - preConsumedQuota
 		if quotaDelta != 0 {
-			err := model.PostConsumeTokenQuota(relayInfo.TokenId, userQuota, quotaDelta, preConsumedQuota, true)
+			err := model.PostConsumeTokenQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true)
 			if err != nil {
 				common.LogError(ctx, "error consuming token remain quota: "+err.Error())
 			}

relay/relay_adaptor.go

@@ -12,6 +12,7 @@ import (
 	"one-api/relay/channel/dify"
 	"one-api/relay/channel/gemini"
 	"one-api/relay/channel/jina"
+	"one-api/relay/channel/mistral"
 	"one-api/relay/channel/ollama"
 	"one-api/relay/channel/openai"
 	"one-api/relay/channel/palm"
@@ -68,6 +69,8 @@ func GetAdaptor(apiType int) channel.Adaptor {
 		return &siliconflow.Adaptor{}
 	case constant.APITypeVertexAi:
 		return &vertex.Adaptor{}
+	case constant.APITypeMistral:
+		return &mistral.Adaptor{}
 	}
 	return nil
 }

relay/relay_rerank.go

@@ -23,7 +23,7 @@ func getRerankPromptToken(rerankRequest dto.RerankRequest) int {
 	return token
 }
 
-func RerankHelper(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode {
+func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWithStatusCode) {
 	relayInfo := relaycommon.GenRelayInfo(c)
 
 	var rerankRequest *dto.RerankRequest
@@ -79,6 +79,12 @@ func RerankHelper(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode
 	if openaiErr != nil {
 		return openaiErr
 	}
+	defer func() {
+		if openaiErr != nil {
+			returnPreConsumedQuota(c, relayInfo, userQuota, preConsumedQuota)
+		}
+	}()
+
 	adaptor := GetAdaptor(relayInfo.ApiType)
 	if adaptor == nil {
 		return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest)
@@ -99,23 +105,24 @@ func RerankHelper(c *gin.Context, relayMode int) *dto.OpenAIErrorWithStatusCode
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
 	}
+
+	var httpResp *http.Response
 	if resp != nil {
-		if resp.StatusCode != http.StatusOK {
-			returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
-			openaiErr := service.RelayErrorHandler(resp)
+		httpResp = resp.(*http.Response)
+		if httpResp.StatusCode != http.StatusOK {
+			openaiErr = service.RelayErrorHandler(httpResp)
 			// reset status code 重置状态码
 			service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 			return openaiErr
 		}
 	}
 
-	usage, openaiErr := adaptor.DoResponse(c, resp, relayInfo)
+	usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo)
 	if openaiErr != nil {
-		returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
 		// reset status code 重置状态码
 		service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 		return openaiErr
 	}
-	postConsumeQuota(c, relayInfo, rerankRequest.Model, usage, ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, success, "")
+	postConsumeQuota(c, relayInfo, rerankRequest.Model, usage.(*dto.Usage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, success, "")
 	return nil
 }

relay/relay_task.go

@@ -111,7 +111,8 @@ func RelayTaskSubmit(c *gin.Context, relayMode int) (taskErr *dto.TaskError) {
 	defer func(ctx context.Context) {
 		// release quota
 		if relayInfo.ConsumeQuota && taskErr == nil {
-			err := model.PostConsumeTokenQuota(relayInfo.TokenId, userQuota, quota, 0, true)
+
+			err := model.PostConsumeTokenQuota(relayInfo.ToRelayInfo(), userQuota, quota, 0, true)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}

relay/websocket.go

@@ -0,0 +1,159 @@
+package relay
+
+import (
+	"encoding/json"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"net/http"
+	"one-api/common"
+	"one-api/dto"
+	relaycommon "one-api/relay/common"
+	"one-api/service"
+)
+
+//func getAndValidateWssRequest(c *gin.Context, ws *websocket.Conn) (*dto.RealtimeEvent, error) {
+//	_, p, err := ws.ReadMessage()
+//	if err != nil {
+//		return nil, err
+//	}
+//	realtimeEvent := &dto.RealtimeEvent{}
+//	err = json.Unmarshal(p, realtimeEvent)
+//	if err != nil {
+//		return nil, err
+//	}
+//	// save the original request
+//	if realtimeEvent.Session == nil {
+//		return nil, errors.New("session object is nil")
+//	}
+//	c.Set("first_wss_request", p)
+//	return realtimeEvent, nil
+//}
+
+func WssHelper(c *gin.Context, ws *websocket.Conn) (openaiErr *dto.OpenAIErrorWithStatusCode) {
+	relayInfo := relaycommon.GenRelayInfoWs(c, ws)
+
+	// get & validate textRequest 获取并验证文本请求
+	//realtimeEvent, err := getAndValidateWssRequest(c, ws)
+	//if err != nil {
+	//	common.LogError(c, fmt.Sprintf("getAndValidateWssRequest failed: %s", err.Error()))
+	//	return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest)
+	//}
+
+	// map model name
+	modelMapping := c.GetString("model_mapping")
+	//isModelMapped := false
+	if modelMapping != "" && modelMapping != "{}" {
+		modelMap := make(map[string]string)
+		err := json.Unmarshal([]byte(modelMapping), &modelMap)
+		if err != nil {
+			return service.OpenAIErrorWrapperLocal(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError)
+		}
+		if modelMap[relayInfo.OriginModelName] != "" {
+			relayInfo.UpstreamModelName = modelMap[relayInfo.OriginModelName]
+			// set upstream model name
+			//isModelMapped = true
+		}
+	}
+	//relayInfo.UpstreamModelName = textRequest.Model
+	modelPrice, getModelPriceSuccess := common.GetModelPrice(relayInfo.UpstreamModelName, false)
+	groupRatio := common.GetGroupRatio(relayInfo.Group)
+
+	var preConsumedQuota int
+	var ratio float64
+	var modelRatio float64
+	//err := service.SensitiveWordsCheck(textRequest)
+
+	//if constant.ShouldCheckPromptSensitive() {
+	//	err = checkRequestSensitive(textRequest, relayInfo)
+	//	if err != nil {
+	//		return service.OpenAIErrorWrapperLocal(err, "sensitive_words_detected", http.StatusBadRequest)
+	//	}
+	//}
+
+	//promptTokens, err := getWssPromptTokens(realtimeEvent, relayInfo)
+	//// count messages token error 计算promptTokens错误
+	//if err != nil {
+	//	return service.OpenAIErrorWrapper(err, "count_token_messages_failed", http.StatusInternalServerError)
+	//}
+	//
+	if !getModelPriceSuccess {
+		preConsumedTokens := common.PreConsumedQuota
+		//if realtimeEvent.Session.MaxResponseOutputTokens != 0 {
+		//	preConsumedTokens = promptTokens + int(realtimeEvent.Session.MaxResponseOutputTokens)
+		//}
+		modelRatio = common.GetModelRatio(relayInfo.UpstreamModelName)
+		ratio = modelRatio * groupRatio
+		preConsumedQuota = int(float64(preConsumedTokens) * ratio)
+	} else {
+		preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatio)
+		relayInfo.UsePrice = true
+	}
+
+	// pre-consume quota 预消耗配额
+	preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, preConsumedQuota, relayInfo)
+	if openaiErr != nil {
+		return openaiErr
+	}
+
+	defer func() {
+		if openaiErr != nil {
+			returnPreConsumedQuota(c, relayInfo, userQuota, preConsumedQuota)
+		}
+	}()
+
+	adaptor := GetAdaptor(relayInfo.ApiType)
+	if adaptor == nil {
+		return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest)
+	}
+	adaptor.Init(relayInfo)
+	//var requestBody io.Reader
+	//firstWssRequest, _ := c.Get("first_wss_request")
+	//requestBody = bytes.NewBuffer(firstWssRequest.([]byte))
+
+	statusCodeMappingStr := c.GetString("status_code_mapping")
+	resp, err := adaptor.DoRequest(c, relayInfo, nil)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+
+	if resp != nil {
+		relayInfo.TargetWs = resp.(*websocket.Conn)
+		defer relayInfo.TargetWs.Close()
+	}
+
+	usage, openaiErr := adaptor.DoResponse(c, nil, relayInfo)
+	if openaiErr != nil {
+		// reset status code 重置状态码
+		service.ResetStatusCode(openaiErr, statusCodeMappingStr)
+		return openaiErr
+	}
+	service.PostWssConsumeQuota(c, relayInfo, relayInfo.UpstreamModelName, usage.(*dto.RealtimeUsage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, getModelPriceSuccess, "")
+	return nil
+}
+
+//func getWssPromptTokens(textRequest *dto.RealtimeEvent, info *relaycommon.RelayInfo) (int, error) {
+//	var promptTokens int
+//	var err error
+//	switch info.RelayMode {
+//	default:
+//		promptTokens, err = service.CountTokenRealtime(*textRequest, info.UpstreamModelName)
+//	}
+//	info.PromptTokens = promptTokens
+//	return promptTokens, err
+//}
+
+//func checkWssRequestSensitive(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) error {
+//	var err error
+//	switch info.RelayMode {
+//	case relayconstant.RelayModeChatCompletions:
+//		err = service.CheckSensitiveMessages(textRequest.Messages)
+//	case relayconstant.RelayModeCompletions:
+//		err = service.CheckSensitiveInput(textRequest.Prompt)
+//	case relayconstant.RelayModeModerations:
+//		err = service.CheckSensitiveInput(textRequest.Input)
+//	case relayconstant.RelayModeEmbeddings:
+//		err = service.CheckSensitiveInput(textRequest.Input)
+//	}
+//	return err
+//}

router/api-router.go

@@ -25,6 +25,7 @@ func SetApiRouter(router *gin.Engine) {
 		apiRouter.GET("/reset_password", middleware.CriticalRateLimit(), middleware.TurnstileCheck(), controller.SendPasswordResetEmail)
 		apiRouter.POST("/user/reset", middleware.CriticalRateLimit(), controller.ResetPassword)
 		apiRouter.GET("/oauth/github", middleware.CriticalRateLimit(), controller.GitHubOAuth)
+		apiRouter.GET("/oauth/linuxdo", middleware.CriticalRateLimit(), controller.LinuxdoOAuth)
 		apiRouter.GET("/oauth/state", middleware.CriticalRateLimit(), controller.GenerateOAuthCode)
 		apiRouter.GET("/oauth/wechat", middleware.CriticalRateLimit(), controller.WeChatAuth)
 		apiRouter.GET("/oauth/wechat/bind", middleware.CriticalRateLimit(), middleware.UserAuth(), controller.WeChatBind)
@@ -39,10 +40,12 @@ func SetApiRouter(router *gin.Engine) {
 			//userRoute.POST("/tokenlog", middleware.CriticalRateLimit(), controller.TokenLog)
 			userRoute.GET("/logout", controller.Logout)
 			userRoute.GET("/epay/notify", controller.EpayNotify)
+			userRoute.GET("/groups", controller.GetUserGroups)
 
 			selfRoute := userRoute.Group("/")
 			selfRoute.Use(middleware.UserAuth())
 			{
+				selfRoute.GET("/self/groups", controller.GetUserGroups)
 				selfRoute.GET("/self", controller.GetSelf)
 				selfRoute.GET("/models", controller.GetUserModels)
 				selfRoute.PUT("/self", controller.UpdateSelf)

router/relay-router.go

@@ -16,33 +16,47 @@ func SetRelayRouter(router *gin.Engine) {
 		modelsRouter.GET("", controller.ListModels)
 		modelsRouter.GET("/:model", controller.RetrieveModel)
 	}
-	relayV1Router := router.Group("/v1")
-	relayV1Router.Use(middleware.TokenAuth(), middleware.Distribute())
+	playgroundRouter := router.Group("/pg")
+	playgroundRouter.Use(middleware.UserAuth())
 	{
-		relayV1Router.POST("/completions", controller.Relay)
-		relayV1Router.POST("/chat/completions", controller.Relay)
-		relayV1Router.POST("/edits", controller.Relay)
-		relayV1Router.POST("/images/generations", controller.Relay)
-		relayV1Router.POST("/images/edits", controller.RelayNotImplemented)
-		relayV1Router.POST("/images/variations", controller.RelayNotImplemented)
-		relayV1Router.POST("/embeddings", controller.Relay)
-		relayV1Router.POST("/engines/:model/embeddings", controller.Relay)
-		relayV1Router.POST("/audio/transcriptions", controller.Relay)
-		relayV1Router.POST("/audio/translations", controller.Relay)
-		relayV1Router.POST("/audio/speech", controller.Relay)
-		relayV1Router.GET("/files", controller.RelayNotImplemented)
-		relayV1Router.POST("/files", controller.RelayNotImplemented)
-		relayV1Router.DELETE("/files/:id", controller.RelayNotImplemented)
-		relayV1Router.GET("/files/:id", controller.RelayNotImplemented)
-		relayV1Router.GET("/files/:id/content", controller.RelayNotImplemented)
-		relayV1Router.POST("/fine-tunes", controller.RelayNotImplemented)
-		relayV1Router.GET("/fine-tunes", controller.RelayNotImplemented)
-		relayV1Router.GET("/fine-tunes/:id", controller.RelayNotImplemented)
-		relayV1Router.POST("/fine-tunes/:id/cancel", controller.RelayNotImplemented)
-		relayV1Router.GET("/fine-tunes/:id/events", controller.RelayNotImplemented)
-		relayV1Router.DELETE("/models/:model", controller.RelayNotImplemented)
-		relayV1Router.POST("/moderations", controller.Relay)
-		relayV1Router.POST("/rerank", controller.Relay)
+		playgroundRouter.POST("/chat/completions", controller.Playground)
+	}
+	relayV1Router := router.Group("/v1")
+	relayV1Router.Use(middleware.TokenAuth())
+	{
+		// WebSocket 路由
+		wsRouter := relayV1Router.Group("")
+		wsRouter.Use(middleware.Distribute())
+		wsRouter.GET("/realtime", controller.WssRelay)
+	}
+	{
+		//http router
+		httpRouter := relayV1Router.Group("")
+		httpRouter.Use(middleware.Distribute())
+		httpRouter.POST("/completions", controller.Relay)
+		httpRouter.POST("/chat/completions", controller.Relay)
+		httpRouter.POST("/edits", controller.Relay)
+		httpRouter.POST("/images/generations", controller.Relay)
+		httpRouter.POST("/images/edits", controller.RelayNotImplemented)
+		httpRouter.POST("/images/variations", controller.RelayNotImplemented)
+		httpRouter.POST("/embeddings", controller.Relay)
+		httpRouter.POST("/engines/:model/embeddings", controller.Relay)
+		httpRouter.POST("/audio/transcriptions", controller.Relay)
+		httpRouter.POST("/audio/translations", controller.Relay)
+		httpRouter.POST("/audio/speech", controller.Relay)
+		httpRouter.GET("/files", controller.RelayNotImplemented)
+		httpRouter.POST("/files", controller.RelayNotImplemented)
+		httpRouter.DELETE("/files/:id", controller.RelayNotImplemented)
+		httpRouter.GET("/files/:id", controller.RelayNotImplemented)
+		httpRouter.GET("/files/:id/content", controller.RelayNotImplemented)
+		httpRouter.POST("/fine-tunes", controller.RelayNotImplemented)
+		httpRouter.GET("/fine-tunes", controller.RelayNotImplemented)
+		httpRouter.GET("/fine-tunes/:id", controller.RelayNotImplemented)
+		httpRouter.POST("/fine-tunes/:id/cancel", controller.RelayNotImplemented)
+		httpRouter.GET("/fine-tunes/:id/events", controller.RelayNotImplemented)
+		httpRouter.DELETE("/models/:model", controller.RelayNotImplemented)
+		httpRouter.POST("/moderations", controller.Relay)
+		httpRouter.POST("/rerank", controller.Relay)
 	}
 
 	relayMjRouter := router.Group("/mj")

service/audio.go

@@ -0,0 +1,31 @@
+package service
+
+import (
+	"encoding/base64"
+	"fmt"
+)
+
+func parseAudio(audioBase64 string, format string) (duration float64, err error) {
+	audioData, err := base64.StdEncoding.DecodeString(audioBase64)
+	if err != nil {
+		return 0, fmt.Errorf("base64 decode error: %v", err)
+	}
+
+	var samplesCount int
+	var sampleRate int
+
+	switch format {
+	case "pcm16":
+		samplesCount = len(audioData) / 2 // 16位 = 2字节每样本
+		sampleRate = 24000                // 24kHz
+	case "g711_ulaw", "g711_alaw":
+		samplesCount = len(audioData) // 8位 = 1字节每样本
+		sampleRate = 8000             // 8kHz
+	default:
+		samplesCount = len(audioData) // 8位 = 1字节每样本
+		sampleRate = 8000             // 8kHz
+	}
+
+	duration = float64(samplesCount) / float64(sampleRate)
+	return duration, nil
+}

service/channel.go

@@ -73,6 +73,15 @@ func ShouldDisableChannel(channelType int, err *relaymodel.OpenAIErrorWithStatus
 	} else if strings.HasPrefix(err.Error.Message, "Permission denied") {
 		return true
 	}
+
+	if strings.Contains(err.Error.Message, "The security token included in the request is invalid") { // anthropic
+		return true
+	} else if strings.Contains(err.Error.Message, "Operation not allowed") {
+		return true
+	} else if strings.Contains(err.Error.Message, "Your account is not authorized") {
+		return true
+	}
+
 	return false
 }
 

service/log.go

@@ -2,6 +2,7 @@ package service
 
 import (
 	"github.com/gin-gonic/gin"
+	"one-api/dto"
 	relaycommon "one-api/relay/common"
 )
 
@@ -17,3 +18,27 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m
 	other["admin_info"] = adminInfo
 	return other
 }
+
+func GenerateWssOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice float64) map[string]interface{} {
+	info := GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, modelPrice)
+	info["ws"] = true
+	info["audio_input"] = usage.InputTokenDetails.AudioTokens
+	info["audio_output"] = usage.OutputTokenDetails.AudioTokens
+	info["text_input"] = usage.InputTokenDetails.TextTokens
+	info["text_output"] = usage.OutputTokenDetails.TextTokens
+	info["audio_ratio"] = audioRatio
+	info["audio_completion_ratio"] = audioCompletionRatio
+	return info
+}
+
+func GenerateAudioOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice float64) map[string]interface{} {
+	info := GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, modelPrice)
+	info["audio"] = true
+	info["audio_input"] = usage.PromptTokensDetails.AudioTokens
+	info["audio_output"] = usage.CompletionTokenDetails.AudioTokens
+	info["text_input"] = usage.PromptTokensDetails.TextTokens
+	info["text_output"] = usage.CompletionTokenDetails.TextTokens
+	info["audio_ratio"] = audioRatio
+	info["audio_completion_ratio"] = audioCompletionRatio
+	return info
+}

service/quota.go

@@ -0,0 +1,212 @@
+package service
+
+import (
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"math"
+	"one-api/common"
+	"one-api/dto"
+	"one-api/model"
+	relaycommon "one-api/relay/common"
+	"strings"
+	"time"
+)
+
+func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error {
+	if relayInfo.UsePrice {
+		return nil
+	}
+	userQuota, err := model.GetUserQuota(relayInfo.UserId)
+	if err != nil {
+		return err
+	}
+
+	token, err := model.CacheGetTokenByKey(strings.TrimLeft(relayInfo.ApiKey, "sk-"))
+	if err != nil {
+		return err
+	}
+
+	modelName := relayInfo.UpstreamModelName
+	textInputTokens := usage.InputTokenDetails.TextTokens
+	textOutTokens := usage.OutputTokenDetails.TextTokens
+	audioInputTokens := usage.InputTokenDetails.AudioTokens
+	audioOutTokens := usage.OutputTokenDetails.AudioTokens
+
+	completionRatio := common.GetCompletionRatio(modelName)
+	audioRatio := common.GetAudioRatio(relayInfo.UpstreamModelName)
+	audioCompletionRatio := common.GetAudioCompletionRatio(modelName)
+	groupRatio := common.GetGroupRatio(relayInfo.Group)
+	modelRatio := common.GetModelRatio(modelName)
+
+	ratio := groupRatio * modelRatio
+
+	quota := textInputTokens + int(math.Round(float64(textOutTokens)*completionRatio))
+	quota += int(math.Round(float64(audioInputTokens)*audioRatio)) + int(math.Round(float64(audioOutTokens)*audioRatio*audioCompletionRatio))
+
+	quota = int(math.Round(float64(quota) * ratio))
+	if ratio != 0 && quota <= 0 {
+		quota = 1
+	}
+
+	if userQuota < quota {
+		return errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
+	}
+
+	if token.RemainQuota < quota {
+		return errors.New(fmt.Sprintf("令牌额度不足，剩余额度为 %d", token.RemainQuota))
+	}
+
+	err = model.PostConsumeTokenQuota(relayInfo, 0, quota, 0, false)
+	if err != nil {
+		return err
+	}
+	common.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
+	err = model.CacheUpdateUserQuota(relayInfo.UserId)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
+	usage *dto.RealtimeUsage, ratio float64, preConsumedQuota int, userQuota int, modelRatio float64,
+	groupRatio float64,
+	modelPrice float64, usePrice bool, extraContent string) {
+
+	useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
+	textInputTokens := usage.InputTokenDetails.TextTokens
+	textOutTokens := usage.OutputTokenDetails.TextTokens
+
+	audioInputTokens := usage.InputTokenDetails.AudioTokens
+	audioOutTokens := usage.OutputTokenDetails.AudioTokens
+
+	tokenName := ctx.GetString("token_name")
+	completionRatio := common.GetCompletionRatio(modelName)
+	audioRatio := common.GetAudioRatio(relayInfo.UpstreamModelName)
+	audioCompletionRatio := common.GetAudioCompletionRatio(modelName)
+
+	quota := 0
+	if !usePrice {
+		quota = int(math.Round(float64(textInputTokens) + float64(textOutTokens)*completionRatio))
+		quota += int(math.Round(float64(audioInputTokens)*audioRatio + float64(audioOutTokens)*audioRatio*audioCompletionRatio))
+		quota = int(math.Round(float64(quota) * ratio))
+		if ratio != 0 && quota <= 0 {
+			quota = 1
+		}
+	} else {
+		quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
+	}
+	totalTokens := usage.TotalTokens
+	var logContent string
+	if !usePrice {
+		logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，音频倍率 %.2f，音频补全倍率 %.2f，分组倍率 %.2f", modelRatio, completionRatio, audioRatio, audioCompletionRatio, groupRatio)
+	} else {
+		logContent = fmt.Sprintf("模型价格 %.2f，分组倍率 %.2f", modelPrice, groupRatio)
+	}
+
+	// record all the consume log even if quota is 0
+	if totalTokens == 0 {
+		// in this case, must be some error happened
+		// we cannot just return, because we may have to return the pre-consumed quota
+		quota = 0
+		logContent += fmt.Sprintf("（可能是上游超时）")
+		common.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
+			"tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, preConsumedQuota))
+	} else {
+		//if sensitiveResp != nil {
+		//	logContent += fmt.Sprintf("，敏感词：%s", strings.Join(sensitiveResp.SensitiveWords, ", "))
+		//}
+		//quotaDelta := quota - preConsumedQuota
+		//if quotaDelta != 0 {
+		//	err := model.PostConsumeTokenQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true)
+		//	if err != nil {
+		//		common.LogError(ctx, "error consuming token remain quota: "+err.Error())
+		//	}
+		//}
+
+		//err := model.CacheUpdateUserQuota(relayInfo.UserId)
+		//if err != nil {
+		//	common.LogError(ctx, "error update user quota cache: "+err.Error())
+		//}
+		model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
+		model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
+	}
+
+	logModel := modelName
+	if extraContent != "" {
+		logContent += ", " + extraContent
+	}
+	other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice)
+	model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.InputTokens, usage.OutputTokens, logModel,
+		tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, other)
+}
+
+func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
+	usage *dto.Usage, ratio float64, preConsumedQuota int, userQuota int, modelRatio float64,
+	groupRatio float64,
+	modelPrice float64, usePrice bool, extraContent string) {
+
+	useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
+	textInputTokens := usage.PromptTokensDetails.TextTokens
+	textOutTokens := usage.CompletionTokenDetails.TextTokens
+
+	audioInputTokens := usage.PromptTokensDetails.AudioTokens
+	audioOutTokens := usage.CompletionTokenDetails.AudioTokens
+
+	tokenName := ctx.GetString("token_name")
+	completionRatio := common.GetCompletionRatio(relayInfo.UpstreamModelName)
+	audioRatio := common.GetAudioRatio(relayInfo.UpstreamModelName)
+	audioCompletionRatio := common.GetAudioCompletionRatio(relayInfo.UpstreamModelName)
+
+	quota := 0
+	if !usePrice {
+		quota = int(math.Round(float64(textInputTokens) + float64(textOutTokens)*completionRatio))
+		quota += int(math.Round(float64(audioInputTokens)*audioRatio + float64(audioOutTokens)*audioRatio*audioCompletionRatio))
+		quota = int(math.Round(float64(quota) * ratio))
+		if ratio != 0 && quota <= 0 {
+			quota = 1
+		}
+	} else {
+		quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
+	}
+	totalTokens := usage.TotalTokens
+	var logContent string
+	if !usePrice {
+		logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，音频倍率 %.2f，音频补全倍率 %.2f，分组倍率 %.2f", modelRatio, completionRatio, audioRatio, audioCompletionRatio, groupRatio)
+	} else {
+		logContent = fmt.Sprintf("模型价格 %.2f，分组倍率 %.2f", modelPrice, groupRatio)
+	}
+
+	// record all the consume log even if quota is 0
+	if totalTokens == 0 {
+		// in this case, must be some error happened
+		// we cannot just return, because we may have to return the pre-consumed quota
+		quota = 0
+		logContent += fmt.Sprintf("（可能是上游超时）")
+		common.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
+			"tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.UpstreamModelName, preConsumedQuota))
+	} else {
+		quotaDelta := quota - preConsumedQuota
+		if quotaDelta != 0 {
+			err := model.PostConsumeTokenQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true)
+			if err != nil {
+				common.LogError(ctx, "error consuming token remain quota: "+err.Error())
+			}
+		}
+		err := model.CacheUpdateUserQuota(relayInfo.UserId)
+		if err != nil {
+			common.LogError(ctx, "error update user quota cache: "+err.Error())
+		}
+		model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
+		model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
+	}
+
+	logModel := relayInfo.UpstreamModelName
+	if extraContent != "" {
+		logContent += ", " + extraContent
+	}
+	other := GenerateAudioOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice)
+	model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.PromptTokens, usage.CompletionTokens, logModel,
+		tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, other)
+}

service/relay.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"net/http"
 	"one-api/common"
 	"one-api/dto"
@@ -42,11 +43,47 @@ func Done(c *gin.Context) {
 	_ = StringData(c, "[DONE]")
 }
 
+func WssString(c *gin.Context, ws *websocket.Conn, str string) error {
+	if ws == nil {
+		common.LogError(c, "websocket connection is nil")
+		return errors.New("websocket connection is nil")
+	}
+	//common.LogInfo(c, fmt.Sprintf("sending message: %s", str))
+	return ws.WriteMessage(1, []byte(str))
+}
+
+func WssObject(c *gin.Context, ws *websocket.Conn, object interface{}) error {
+	jsonData, err := json.Marshal(object)
+	if err != nil {
+		return fmt.Errorf("error marshalling object: %w", err)
+	}
+	if ws == nil {
+		common.LogError(c, "websocket connection is nil")
+		return errors.New("websocket connection is nil")
+	}
+	//common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData))
+	return ws.WriteMessage(1, jsonData)
+}
+
+func WssError(c *gin.Context, ws *websocket.Conn, openaiError dto.OpenAIError) {
+	errorObj := &dto.RealtimeEvent{
+		Type:    "error",
+		EventId: GetLocalRealtimeID(c),
+		Error:   &openaiError,
+	}
+	_ = WssObject(c, ws, errorObj)
+}
+
 func GetResponseID(c *gin.Context) string {
-	logID := c.GetString("X-Oneapi-Request-Id")
+	logID := c.GetString(common.RequestIdKey)
 	return fmt.Sprintf("chatcmpl-%s", logID)
 }
 
+func GetLocalRealtimeID(c *gin.Context) string {
+	logID := c.GetString(common.RequestIdKey)
+	return fmt.Sprintf("evt_%s", logID)
+}
+
 func GenerateStopResponse(id string, createAt int64, model string, finishReason string) *dto.ChatCompletionsStreamResponse {
 	return &dto.ChatCompletionsStreamResponse{
 		Id:                id,

service/token_counter.go

@@ -11,6 +11,7 @@ import (
 	"one-api/common"
 	"one-api/constant"
 	"one-api/dto"
+	relaycommon "one-api/relay/common"
 	"strings"
 	"unicode/utf8"
 )
@@ -191,6 +192,72 @@ func CountTokenChatRequest(request dto.GeneralOpenAIRequest, model string) (int,
 	return tkm, nil
 }
 
+func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent, model string) (int, int, error) {
+	audioToken := 0
+	textToken := 0
+	switch request.Type {
+	case dto.RealtimeEventTypeSessionUpdate:
+		if request.Session != nil {
+			msgTokens, err := CountTextToken(request.Session.Instructions, model)
+			if err != nil {
+				return 0, 0, err
+			}
+			textToken += msgTokens
+		}
+	case dto.RealtimeEventResponseAudioDelta:
+		// count audio token
+		atk, err := CountAudioTokenOutput(request.Delta, info.OutputAudioFormat)
+		if err != nil {
+			return 0, 0, fmt.Errorf("error counting audio token: %v", err)
+		}
+		audioToken += atk
+	case dto.RealtimeEventResponseAudioTranscriptionDelta, dto.RealtimeEventResponseFunctionCallArgumentsDelta:
+		// count text token
+		tkm, err := CountTextToken(request.Delta, model)
+		if err != nil {
+			return 0, 0, fmt.Errorf("error counting text token: %v", err)
+		}
+		textToken += tkm
+	case dto.RealtimeEventInputAudioBufferAppend:
+		// count audio token
+		atk, err := CountAudioTokenInput(request.Audio, info.InputAudioFormat)
+		if err != nil {
+			return 0, 0, fmt.Errorf("error counting audio token: %v", err)
+		}
+		audioToken += atk
+	case dto.RealtimeEventConversationItemCreated:
+		if request.Item != nil {
+			switch request.Item.Type {
+			case "message":
+				for _, content := range request.Item.Content {
+					if content.Type == "input_text" {
+						tokens, err := CountTextToken(content.Text, model)
+						if err != nil {
+							return 0, 0, err
+						}
+						textToken += tokens
+					}
+				}
+			}
+		}
+	case dto.RealtimeEventTypeResponseDone:
+		// count tools token
+		if !info.IsFirstRequest {
+			if info.RealtimeTools != nil && len(info.RealtimeTools) > 0 {
+				for _, tool := range info.RealtimeTools {
+					toolTokens, err := CountTokenInput(tool, model)
+					if err != nil {
+						return 0, 0, err
+					}
+					textToken += 8
+					textToken += toolTokens
+				}
+			}
+		}
+	}
+	return textToken, audioToken, nil
+}
+
 func CountTokenMessages(messages []dto.Message, model string, stream bool) (int, error) {
 	//recover when panic
 	tokenEncoder := getTokenEncoder(model)
@@ -223,7 +290,7 @@ func CountTokenMessages(messages []dto.Message, model string, stream bool) (int,
 			} else {
 				arrayContent := message.ParseContent()
 				for _, m := range arrayContent {
-					if m.Type == "image_url" {
+					if m.Type == dto.ContentTypeImageURL {
 						imageUrl := m.ImageUrl.(dto.MessageImageUrl)
 						imageTokenNum, err := getImageToken(&imageUrl, model, stream)
 						if err != nil {
@@ -231,6 +298,9 @@ func CountTokenMessages(messages []dto.Message, model string, stream bool) (int,
 						}
 						tokenNum += imageTokenNum
 						log.Printf("image token num: %d", imageTokenNum)
+					} else if m.Type == dto.ContentTypeInputAudio {
+						// TODO: 音频token数量计算
+						tokenNum += 100
 					} else {
 						tokenNum += getTokenNum(tokenEncoder, m.Text)
 					}
@@ -245,13 +315,13 @@ func CountTokenMessages(messages []dto.Message, model string, stream bool) (int,
 func CountTokenInput(input any, model string) (int, error) {
 	switch v := input.(type) {
 	case string:
-		return CountTokenText(v, model)
+		return CountTextToken(v, model)
 	case []string:
 		text := ""
 		for _, s := range v {
 			text += s
 		}
-		return CountTokenText(text, model)
+		return CountTextToken(text, model)
 	}
 	return CountTokenInput(fmt.Sprintf("%v", input), model)
 }
@@ -273,16 +343,44 @@ func CountTokenStreamChoices(messages []dto.ChatCompletionsStreamResponseChoice,
 	return tokens
 }
 
-func CountAudioToken(text string, model string) (int, error) {
+func CountTTSToken(text string, model string) (int, error) {
 	if strings.HasPrefix(model, "tts") {
 		return utf8.RuneCountInString(text), nil
 	} else {
-		return CountTokenText(text, model)
+		return CountTextToken(text, model)
 	}
 }
 
-// CountTokenText 统计文本的token数量，仅当文本包含敏感词，返回错误，同时返回token数量
-func CountTokenText(text string, model string) (int, error) {
+func CountAudioTokenInput(audioBase64 string, audioFormat string) (int, error) {
+	if audioBase64 == "" {
+		return 0, nil
+	}
+	duration, err := parseAudio(audioBase64, audioFormat)
+	if err != nil {
+		return 0, err
+	}
+	return int(duration / 60 * 100 / 0.06), nil
+}
+
+func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error) {
+	if audioBase64 == "" {
+		return 0, nil
+	}
+	duration, err := parseAudio(audioBase64, audioFormat)
+	if err != nil {
+		return 0, err
+	}
+	return int(duration / 60 * 200 / 0.24), nil
+}
+
+//func CountAudioToken(sec float64, audioType string) {
+//	if audioType == "input" {
+//
+//	}
+//}
+
+// CountTextToken 统计文本的token数量，仅当文本包含敏感词，返回错误，同时返回token数量
+func CountTextToken(text string, model string) (int, error) {
 	var err error
 	tokenEncoder := getTokenEncoder(model)
 	return getTokenNum(tokenEncoder, text), err

service/usage_helpr.go

@@ -19,7 +19,7 @@ import (
 func ResponseText2Usage(responseText string, modeName string, promptTokens int) (*dto.Usage, error) {
 	usage := &dto.Usage{}
 	usage.PromptTokens = promptTokens
-	ctkm, err := CountTokenText(responseText, modeName)
+	ctkm, err := CountTextToken(responseText, modeName)
 	usage.CompletionTokens = ctkm
 	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 	return usage, err

web/package.json

@@ -4,8 +4,8 @@
   "private": true,
   "type": "module",
   "dependencies": {
-    "@douyinfe/semi-icons": "^2.46.1",
-    "@douyinfe/semi-ui": "^2.55.3",
+    "@douyinfe/semi-icons": "^2.63.1",
+    "@douyinfe/semi-ui": "^2.63.1",
     "@visactor/react-vchart": "~1.8.8",
     "@visactor/vchart": "~1.8.8",
     "@visactor/vchart-semi-theme": "~1.8.8",
@@ -22,7 +22,8 @@
     "react-toastify": "^9.0.8",
     "react-turnstile": "^1.0.5",
     "semantic-ui-offline": "^2.5.0",
-    "semantic-ui-react": "^2.1.3"
+    "semantic-ui-react": "^2.1.3",
+    "sse": "github:mpetazzoni/sse.js"
   },
   "scripts": {
     "dev": "vite",