Merge branch 'upstream-main' into fix/pr-2540

# Conflicts:
#	relay/channel/gemini/relay-gemini.go
This commit is contained in:
Seefs
2026-01-25 14:14:05 +08:00
178 changed files with 22831 additions and 1856 deletions

View File

@@ -6,4 +6,5 @@
Makefile
docs
.eslintcache
.gocache
.gocache
/web/node_modules

View File

@@ -57,6 +57,9 @@
# 流模式无响应超时时间,单位秒,如果出现空补全可以尝试改为更大值
# STREAMING_TIMEOUT=300
# TLS / HTTP 跳过验证设置
# TLS_INSECURE_SKIP_VERIFY=false
# Gemini 识别图片 最大图片数量
# GEMINI_VISION_MAX_IMAGE_NUM=16

5
.gitignore vendored
View File

@@ -19,8 +19,11 @@ tiktoken_cache
.gomodcache/
.cache
web/bun.lock
plans
electron/node_modules
electron/dist
data/
.gomodcache/
.gomodcache/
.gocache-temp
.gopath

View File

@@ -213,9 +213,11 @@ docker run --name new-api -d --restart always \
- 🚦 User-level model rate limiting
**Format Conversion:**
- 🔄 OpenAI ⇄ Claude Messages
- 🔄 OpenAI Gemini Chat
- 🔄 Thinking-to-content functionality
- 🔄 **OpenAI Compatible ⇄ Claude Messages**
- 🔄 **OpenAI Compatible → Google Gemini**
- 🔄 **Google Gemini → OpenAI Compatible** - Text only, function calling not supported yet
- 🚧 **OpenAI Compatible ⇄ OpenAI Responses** - In development
- 🔄 **Thinking-to-content functionality**
**Reasoning Effort Support:**

View File

@@ -212,9 +212,11 @@ docker run --name new-api -d --restart always \
- 🚦 Limitation du débit du modèle pour les utilisateurs
**Conversion de format:**
- 🔄 OpenAI ⇄ Claude Messages
- 🔄 OpenAI Gemini Chat
- 🔄 Fonctionnalité de la pensée au contenu
- 🔄 **OpenAI Compatible ⇄ Claude Messages**
- 🔄 **OpenAI Compatible → Google Gemini**
- 🔄 **Google Gemini → OpenAI Compatible** - Texte uniquement, les appels de fonction ne sont pas encore pris en charge
- 🚧 **OpenAI Compatible ⇄ OpenAI Responses** - En développement
- 🔄 **Fonctionnalité de la pensée au contenu**
**Prise en charge de l'effort de raisonnement:**

View File

@@ -218,9 +218,11 @@ docker run --name new-api -d --restart always \
- 🚦 ユーザーレベルモデルレート制限
**フォーマット変換:**
- 🔄 OpenAI ⇄ Claude Messages
- 🔄 OpenAI Gemini Chat
- 🔄 思考からコンテンツへの機能
- 🔄 **OpenAI Compatible ⇄ Claude Messages**
- 🔄 **OpenAI Compatible → Google Gemini**
- 🔄 **Google Gemini → OpenAI Compatible** - テキストのみ、関数呼び出しはまだサポートされていません
- 🚧 **OpenAI Compatible ⇄ OpenAI Responses** - 開発中
- 🔄 **思考からコンテンツへの機能**
**Reasoning Effort サポート:**

View File

@@ -214,9 +214,11 @@ docker run --name new-api -d --restart always \
- 🚦 用户级别模型限流
**格式转换:**
- 🔄 OpenAI ⇄ Claude Messages
- 🔄 OpenAI Gemini Chat
- 🔄 思考转内容功能
- 🔄 **OpenAI Compatible ⇄ Claude Messages**
- 🔄 **OpenAI Compatible → Google Gemini**
- 🔄 **Google Gemini → OpenAI Compatible** - 仅支持文本,暂不支持函数调用
- 🚧 **OpenAI Compatible ⇄ OpenAI Responses** - 开发中
- 🔄 **思考转内容功能**
**Reasoning Effort 支持:**

View File

@@ -73,6 +73,8 @@ func ChannelType2APIType(channelType int) (int, bool) {
apiType = constant.APITypeMiniMax
case constant.ChannelTypeReplicate:
apiType = constant.APITypeReplicate
case constant.ChannelTypeCodex:
apiType = constant.APITypeCodex
}
if apiType == -1 {
return constant.APITypeOpenAI, false

View File

@@ -1,6 +1,7 @@
package common
import (
"crypto/tls"
//"os"
//"strconv"
"sync"
@@ -73,6 +74,9 @@ var MemoryCacheEnabled bool
var LogConsumeEnabled = true
var TLSInsecureSkipVerify bool
var InsecureTLSConfig = &tls.Config{InsecureSkipVerify: true}
var SMTPServer = ""
var SMTPPort = 587
var SMTPSSLEnabled = false

View File

@@ -40,7 +40,7 @@ func GetRequestBody(c *gin.Context) ([]byte, error) {
}
}
maxMB := constant.MaxRequestBodyMB
if maxMB < 0 {
if maxMB <= 0 {
// no limit
body, err := io.ReadAll(c.Request.Body)
_ = c.Request.Body.Close()

View File

@@ -4,6 +4,7 @@ import (
"flag"
"fmt"
"log"
"net/http"
"os"
"path/filepath"
"strconv"
@@ -81,6 +82,16 @@ func InitEnv() {
DebugEnabled = os.Getenv("DEBUG") == "true"
MemoryCacheEnabled = os.Getenv("MEMORY_CACHE_ENABLED") == "true"
IsMasterNode = os.Getenv("NODE_TYPE") != "slave"
TLSInsecureSkipVerify = GetEnvOrDefaultBool("TLS_INSECURE_SKIP_VERIFY", false)
if TLSInsecureSkipVerify {
if tr, ok := http.DefaultTransport.(*http.Transport); ok && tr != nil {
if tr.TLSClientConfig != nil {
tr.TLSClientConfig.InsecureSkipVerify = true
} else {
tr.TLSClientConfig = InsecureTLSConfig
}
}
}
// Parse requestInterval and set RequestInterval
requestInterval, _ = strconv.Atoi(os.Getenv("POLLING_INTERVAL"))
@@ -115,10 +126,10 @@ func InitEnv() {
func initConstantEnv() {
constant.StreamingTimeout = GetEnvOrDefault("STREAMING_TIMEOUT", 300)
constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", true)
constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", 20)
constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", 64)
constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", 64)
// MaxRequestBodyMB 请求体最大大小(解压后),用于防止超大请求/zip bomb导致内存暴涨
constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", 64)
constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", 128)
// ForceStreamOption 覆盖请求参数强制返回usage信息
constant.ForceStreamOption = GetEnvOrDefaultBool("FORCE_STREAM_OPTION", true)
constant.CountToken = GetEnvOrDefaultBool("CountToken", true)

View File

@@ -16,6 +16,8 @@ var (
maskURLPattern = regexp.MustCompile(`(http|https)://[^\s/$.?#].[^\s]*`)
maskDomainPattern = regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`)
maskIPPattern = regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}\b`)
// maskApiKeyPattern matches patterns like 'api_key:xxx' or "api_key:xxx" to mask the API key value
maskApiKeyPattern = regexp.MustCompile(`(['"]?)api_key:([^\s'"]+)(['"]?)`)
)
func GetStringIfEmpty(str string, defaultValue string) string {
@@ -235,5 +237,8 @@ func MaskSensitiveInfo(str string) string {
// Mask IP addresses
str = maskIPPattern.ReplaceAllString(str, "***.***.***.***")
// Mask API keys (e.g., "api_key:AIzaSyAAAaUooTUni8AdaOkSRMda30n_Q4vrV70" -> "api_key:***")
str = maskApiKeyPattern.ReplaceAllString(str, "${1}api_key:***${3}")
return str
}

View File

@@ -263,7 +263,7 @@ func GetTimestamp() int64 {
}
func GetTimeString() string {
now := time.Now()
now := time.Now().UTC()
return fmt.Sprintf("%s%d", now.Format("20060102150405"), now.UnixNano()%1e9)
}

View File

@@ -35,5 +35,6 @@ const (
APITypeSubmodel
APITypeMiniMax
APITypeReplicate
APITypeCodex
APITypeDummy // this one is only for count, do not add any channel after this
)

View File

@@ -54,6 +54,7 @@ const (
ChannelTypeDoubaoVideo = 54
ChannelTypeSora = 55
ChannelTypeReplicate = 56
ChannelTypeCodex = 57
ChannelTypeDummy // this one is only for count, do not add any channel after this
)
@@ -116,6 +117,7 @@ var ChannelBaseURLs = []string{
"https://ark.cn-beijing.volces.com", //54
"https://api.openai.com", //55
"https://api.replicate.com", //56
"https://chatgpt.com", //57
}
var ChannelTypeNames = map[int]string{
@@ -172,6 +174,7 @@ var ChannelTypeNames = map[int]string{
ChannelTypeDoubaoVideo: "DoubaoVideo",
ChannelTypeSora: "Sora",
ChannelTypeReplicate: "Replicate",
ChannelTypeCodex: "Codex",
}
func GetChannelTypeName(channelType int) string {

View File

@@ -193,6 +193,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string)
}
}
info.IsChannelTest = true
info.InitChannelMeta(c)
err = helper.ModelMappedHelper(c, info, request)
@@ -309,8 +310,29 @@ func testChannel(channel *model.Channel, testModel string, endpointType string)
newAPIError: types.NewError(err, types.ErrorCodeJsonMarshalFailed),
}
}
//jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings)
//if err != nil {
// return testResult{
// context: c,
// localErr: err,
// newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed),
// }
//}
if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil {
return testResult{
context: c,
localErr: err,
newAPIError: types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid),
}
}
}
requestBody := bytes.NewBuffer(jsonData)
c.Request.Body = io.NopCloser(requestBody)
c.Request.Body = io.NopCloser(bytes.NewBuffer(jsonData))
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
return testResult{

View File

@@ -1,26 +1,31 @@
package controller
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel/gemini"
"github.com/QuantumNous/new-api/relay/channel/ollama"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type OpenAIModel struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
OwnedBy string `json:"owned_by"`
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
OwnedBy string `json:"owned_by"`
Metadata map[string]any `json:"metadata,omitempty"`
Permission []struct {
ID string `json:"id"`
Object string `json:"object"`
@@ -207,11 +212,88 @@ func FetchUpstreamModels(c *gin.Context) {
baseURL = channel.GetBaseURL()
}
// 对于 Ollama 渠道,使用特殊处理
if channel.Type == constant.ChannelTypeOllama {
key := strings.Split(channel.Key, "\n")[0]
models, err := ollama.FetchOllamaModels(baseURL, key)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("获取Ollama模型失败: %s", err.Error()),
})
return
}
result := OpenAIModelsResponse{
Data: make([]OpenAIModel, 0, len(models)),
}
for _, modelInfo := range models {
metadata := map[string]any{}
if modelInfo.Size > 0 {
metadata["size"] = modelInfo.Size
}
if modelInfo.Digest != "" {
metadata["digest"] = modelInfo.Digest
}
if modelInfo.ModifiedAt != "" {
metadata["modified_at"] = modelInfo.ModifiedAt
}
details := modelInfo.Details
if details.ParentModel != "" || details.Format != "" || details.Family != "" || len(details.Families) > 0 || details.ParameterSize != "" || details.QuantizationLevel != "" {
metadata["details"] = modelInfo.Details
}
if len(metadata) == 0 {
metadata = nil
}
result.Data = append(result.Data, OpenAIModel{
ID: modelInfo.Name,
Object: "model",
Created: 0,
OwnedBy: "ollama",
Metadata: metadata,
})
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": result.Data,
})
return
}
// 对于 Gemini 渠道,使用特殊处理
if channel.Type == constant.ChannelTypeGemini {
// 获取用于请求的可用密钥(多密钥渠道优先使用启用状态的密钥)
key, _, apiErr := channel.GetNextEnabledKey()
if apiErr != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("获取渠道密钥失败: %s", apiErr.Error()),
})
return
}
key = strings.TrimSpace(key)
models, err := gemini.FetchGeminiModels(baseURL, key, channel.GetSetting().Proxy)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("获取Gemini模型失败: %s", err.Error()),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": models,
})
return
}
var url string
switch channel.Type {
case constant.ChannelTypeGemini:
// curl https://example.com/v1beta/models?key=$GEMINI_API_KEY
url = fmt.Sprintf("%s/v1beta/openai/models", baseURL) // Remove key in url since we need to use AuthHeader
case constant.ChannelTypeAli:
url = fmt.Sprintf("%s/compatible-mode/v1/models", baseURL)
case constant.ChannelTypeZhipu_v4:
@@ -524,9 +606,60 @@ func validateChannel(channel *model.Channel, isAdd bool) error {
}
}
// Codex OAuth key validation (optional, only when JSON object is provided)
if channel.Type == constant.ChannelTypeCodex {
trimmedKey := strings.TrimSpace(channel.Key)
if isAdd || trimmedKey != "" {
if !strings.HasPrefix(trimmedKey, "{") {
return fmt.Errorf("Codex key must be a valid JSON object")
}
var keyMap map[string]any
if err := common.Unmarshal([]byte(trimmedKey), &keyMap); err != nil {
return fmt.Errorf("Codex key must be a valid JSON object")
}
if v, ok := keyMap["access_token"]; !ok || v == nil || strings.TrimSpace(fmt.Sprintf("%v", v)) == "" {
return fmt.Errorf("Codex key JSON must include access_token")
}
if v, ok := keyMap["account_id"]; !ok || v == nil || strings.TrimSpace(fmt.Sprintf("%v", v)) == "" {
return fmt.Errorf("Codex key JSON must include account_id")
}
}
}
return nil
}
func RefreshCodexChannelCredential(c *gin.Context) {
channelId, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
return
}
ctx, cancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
defer cancel()
oauthKey, ch, err := service.RefreshCodexChannelCredential(ctx, channelId, service.CodexCredentialRefreshOptions{ResetCaches: true})
if err != nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "refreshed",
"data": gin.H{
"expires_at": oauthKey.Expired,
"last_refresh": oauthKey.LastRefresh,
"account_id": oauthKey.AccountID,
"email": oauthKey.Email,
"channel_id": ch.Id,
"channel_type": ch.Type,
"channel_name": ch.Name,
},
})
}
type AddChannelRequest struct {
Mode string `json:"mode"`
MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
@@ -917,9 +1050,6 @@ func UpdateChannel(c *gin.Context) {
// 单个JSON密钥
newKeys = []string{channel.Key}
}
// 合并密钥
allKeys := append(existingKeys, newKeys...)
channel.Key = strings.Join(allKeys, "\n")
} else {
// 普通渠道的处理
inputKeys := strings.Split(channel.Key, "\n")
@@ -929,10 +1059,31 @@ func UpdateChannel(c *gin.Context) {
newKeys = append(newKeys, key)
}
}
// 合并密钥
allKeys := append(existingKeys, newKeys...)
channel.Key = strings.Join(allKeys, "\n")
}
seen := make(map[string]struct{}, len(existingKeys)+len(newKeys))
for _, key := range existingKeys {
normalized := strings.TrimSpace(key)
if normalized == "" {
continue
}
seen[normalized] = struct{}{}
}
dedupedNewKeys := make([]string, 0, len(newKeys))
for _, key := range newKeys {
normalized := strings.TrimSpace(key)
if normalized == "" {
continue
}
if _, ok := seen[normalized]; ok {
continue
}
seen[normalized] = struct{}{}
dedupedNewKeys = append(dedupedNewKeys, normalized)
}
allKeys := append(existingKeys, dedupedNewKeys...)
channel.Key = strings.Join(allKeys, "\n")
}
case "replace":
// 覆盖模式:直接使用新密钥(默认行为,不需要特殊处理)
@@ -975,6 +1126,49 @@ func FetchModels(c *gin.Context) {
baseURL = constant.ChannelBaseURLs[req.Type]
}
// remove line breaks and extra spaces.
key := strings.TrimSpace(req.Key)
key = strings.Split(key, "\n")[0]
if req.Type == constant.ChannelTypeOllama {
models, err := ollama.FetchOllamaModels(baseURL, key)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("获取Ollama模型失败: %s", err.Error()),
})
return
}
names := make([]string, 0, len(models))
for _, modelInfo := range models {
names = append(names, modelInfo.Name)
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": names,
})
return
}
if req.Type == constant.ChannelTypeGemini {
models, err := gemini.FetchGeminiModels(baseURL, key, "")
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("获取Gemini模型失败: %s", err.Error()),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": models,
})
return
}
client := &http.Client{}
url := fmt.Sprintf("%s/v1/models", baseURL)
@@ -987,10 +1181,6 @@ func FetchModels(c *gin.Context) {
return
}
// remove line breaks and extra spaces.
key := strings.TrimSpace(req.Key)
// If the key contains a line break, only take the first part.
key = strings.Split(key, "\n")[0]
request.Header.Set("Authorization", "Bearer "+key)
response, err := client.Do(request)
@@ -1640,3 +1830,262 @@ func ManageMultiKeys(c *gin.Context) {
return
}
}
// OllamaPullModel 拉取 Ollama 模型
func OllamaPullModel(c *gin.Context) {
var req struct {
ChannelID int `json:"channel_id"`
ModelName string `json:"model_name"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Invalid request parameters",
})
return
}
if req.ChannelID == 0 || req.ModelName == "" {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Channel ID and model name are required",
})
return
}
// 获取渠道信息
channel, err := model.GetChannelById(req.ChannelID, true)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{
"success": false,
"message": "Channel not found",
})
return
}
// 检查是否是 Ollama 渠道
if channel.Type != constant.ChannelTypeOllama {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "This operation is only supported for Ollama channels",
})
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
key := strings.Split(channel.Key, "\n")[0]
err = ollama.PullOllamaModel(baseURL, key, req.ModelName)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{
"success": false,
"message": fmt.Sprintf("Failed to pull model: %s", err.Error()),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": fmt.Sprintf("Model %s pulled successfully", req.ModelName),
})
}
// OllamaPullModelStream 流式拉取 Ollama 模型
func OllamaPullModelStream(c *gin.Context) {
var req struct {
ChannelID int `json:"channel_id"`
ModelName string `json:"model_name"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Invalid request parameters",
})
return
}
if req.ChannelID == 0 || req.ModelName == "" {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Channel ID and model name are required",
})
return
}
// 获取渠道信息
channel, err := model.GetChannelById(req.ChannelID, true)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{
"success": false,
"message": "Channel not found",
})
return
}
// 检查是否是 Ollama 渠道
if channel.Type != constant.ChannelTypeOllama {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "This operation is only supported for Ollama channels",
})
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
// 设置 SSE 头部
c.Header("Content-Type", "text/event-stream")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
c.Header("Access-Control-Allow-Origin", "*")
key := strings.Split(channel.Key, "\n")[0]
// 创建进度回调函数
progressCallback := func(progress ollama.OllamaPullResponse) {
data, _ := json.Marshal(progress)
fmt.Fprintf(c.Writer, "data: %s\n\n", string(data))
c.Writer.Flush()
}
// 执行拉取
err = ollama.PullOllamaModelStream(baseURL, key, req.ModelName, progressCallback)
if err != nil {
errorData, _ := json.Marshal(gin.H{
"error": err.Error(),
})
fmt.Fprintf(c.Writer, "data: %s\n\n", string(errorData))
} else {
successData, _ := json.Marshal(gin.H{
"message": fmt.Sprintf("Model %s pulled successfully", req.ModelName),
})
fmt.Fprintf(c.Writer, "data: %s\n\n", string(successData))
}
// 发送结束标志
fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
c.Writer.Flush()
}
// OllamaDeleteModel 删除 Ollama 模型
func OllamaDeleteModel(c *gin.Context) {
var req struct {
ChannelID int `json:"channel_id"`
ModelName string `json:"model_name"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Invalid request parameters",
})
return
}
if req.ChannelID == 0 || req.ModelName == "" {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Channel ID and model name are required",
})
return
}
// 获取渠道信息
channel, err := model.GetChannelById(req.ChannelID, true)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{
"success": false,
"message": "Channel not found",
})
return
}
// 检查是否是 Ollama 渠道
if channel.Type != constant.ChannelTypeOllama {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "This operation is only supported for Ollama channels",
})
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
key := strings.Split(channel.Key, "\n")[0]
err = ollama.DeleteOllamaModel(baseURL, key, req.ModelName)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{
"success": false,
"message": fmt.Sprintf("Failed to delete model: %s", err.Error()),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": fmt.Sprintf("Model %s deleted successfully", req.ModelName),
})
}
// OllamaVersion 获取 Ollama 服务版本信息
func OllamaVersion(c *gin.Context) {
id, err := strconv.Atoi(c.Param("id"))
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "Invalid channel id",
})
return
}
channel, err := model.GetChannelById(id, true)
if err != nil {
c.JSON(http.StatusNotFound, gin.H{
"success": false,
"message": "Channel not found",
})
return
}
if channel.Type != constant.ChannelTypeOllama {
c.JSON(http.StatusBadRequest, gin.H{
"success": false,
"message": "This operation is only supported for Ollama channels",
})
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
key := strings.Split(channel.Key, "\n")[0]
version, err := ollama.FetchOllamaVersion(baseURL, key)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("获取Ollama版本失败: %s", err.Error()),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"version": version,
},
})
}

72
controller/checkin.go Normal file
View File

@@ -0,0 +1,72 @@
package controller
import (
"fmt"
"net/http"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
)
// GetCheckinStatus 获取用户签到状态和历史记录
func GetCheckinStatus(c *gin.Context) {
setting := operation_setting.GetCheckinSetting()
if !setting.Enabled {
common.ApiErrorMsg(c, "签到功能未启用")
return
}
userId := c.GetInt("id")
// 获取月份参数,默认为当前月份
month := c.DefaultQuery("month", time.Now().Format("2006-01"))
stats, err := model.GetUserCheckinStats(userId, month)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"enabled": setting.Enabled,
"min_quota": setting.MinQuota,
"max_quota": setting.MaxQuota,
"stats": stats,
},
})
}
// DoCheckin 执行用户签到
func DoCheckin(c *gin.Context) {
setting := operation_setting.GetCheckinSetting()
if !setting.Enabled {
common.ApiErrorMsg(c, "签到功能未启用")
return
}
userId := c.GetInt("id")
checkin, err := model.UserCheckin(userId)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
model.RecordLog(userId, model.LogTypeSystem, fmt.Sprintf("用户签到,获得额度 %s", logger.LogQuota(checkin.QuotaAwarded)))
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "签到成功",
"data": gin.H{
"quota_awarded": checkin.QuotaAwarded,
"checkin_date": checkin.CheckinDate},
})
}

243
controller/codex_oauth.go Normal file
View File

@@ -0,0 +1,243 @@
package controller
import (
"context"
"errors"
"fmt"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel/codex"
"github.com/QuantumNous/new-api/service"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
)
type codexOAuthCompleteRequest struct {
Input string `json:"input"`
}
func codexOAuthSessionKey(channelID int, field string) string {
return fmt.Sprintf("codex_oauth_%s_%d", field, channelID)
}
func parseCodexAuthorizationInput(input string) (code string, state string, err error) {
v := strings.TrimSpace(input)
if v == "" {
return "", "", errors.New("empty input")
}
if strings.Contains(v, "#") {
parts := strings.SplitN(v, "#", 2)
code = strings.TrimSpace(parts[0])
state = strings.TrimSpace(parts[1])
return code, state, nil
}
if strings.Contains(v, "code=") {
u, parseErr := url.Parse(v)
if parseErr == nil {
q := u.Query()
code = strings.TrimSpace(q.Get("code"))
state = strings.TrimSpace(q.Get("state"))
return code, state, nil
}
q, parseErr := url.ParseQuery(v)
if parseErr == nil {
code = strings.TrimSpace(q.Get("code"))
state = strings.TrimSpace(q.Get("state"))
return code, state, nil
}
}
code = v
return code, "", nil
}
func StartCodexOAuth(c *gin.Context) {
startCodexOAuthWithChannelID(c, 0)
}
func StartCodexOAuthForChannel(c *gin.Context) {
channelID, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
return
}
startCodexOAuthWithChannelID(c, channelID)
}
func startCodexOAuthWithChannelID(c *gin.Context, channelID int) {
if channelID > 0 {
ch, err := model.GetChannelById(channelID, false)
if err != nil {
common.ApiError(c, err)
return
}
if ch == nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel not found"})
return
}
if ch.Type != constant.ChannelTypeCodex {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
return
}
}
flow, err := service.CreateCodexOAuthAuthorizationFlow()
if err != nil {
common.ApiError(c, err)
return
}
session := sessions.Default(c)
session.Set(codexOAuthSessionKey(channelID, "state"), flow.State)
session.Set(codexOAuthSessionKey(channelID, "verifier"), flow.Verifier)
session.Set(codexOAuthSessionKey(channelID, "created_at"), time.Now().Unix())
_ = session.Save()
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": gin.H{
"authorize_url": flow.AuthorizeURL,
},
})
}
func CompleteCodexOAuth(c *gin.Context) {
completeCodexOAuthWithChannelID(c, 0)
}
func CompleteCodexOAuthForChannel(c *gin.Context) {
channelID, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
return
}
completeCodexOAuthWithChannelID(c, channelID)
}
func completeCodexOAuthWithChannelID(c *gin.Context, channelID int) {
req := codexOAuthCompleteRequest{}
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiError(c, err)
return
}
code, state, err := parseCodexAuthorizationInput(req.Input)
if err != nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
return
}
if strings.TrimSpace(code) == "" {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "missing authorization code"})
return
}
if strings.TrimSpace(state) == "" {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "missing state in input"})
return
}
if channelID > 0 {
ch, err := model.GetChannelById(channelID, false)
if err != nil {
common.ApiError(c, err)
return
}
if ch == nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel not found"})
return
}
if ch.Type != constant.ChannelTypeCodex {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
return
}
}
session := sessions.Default(c)
expectedState, _ := session.Get(codexOAuthSessionKey(channelID, "state")).(string)
verifier, _ := session.Get(codexOAuthSessionKey(channelID, "verifier")).(string)
if strings.TrimSpace(expectedState) == "" || strings.TrimSpace(verifier) == "" {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "oauth flow not started or session expired"})
return
}
if state != expectedState {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "state mismatch"})
return
}
ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
defer cancel()
tokenRes, err := service.ExchangeCodexAuthorizationCode(ctx, code, verifier)
if err != nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
return
}
accountID, ok := service.ExtractCodexAccountIDFromJWT(tokenRes.AccessToken)
if !ok {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "failed to extract account_id from access_token"})
return
}
email, _ := service.ExtractEmailFromJWT(tokenRes.AccessToken)
key := codex.OAuthKey{
AccessToken: tokenRes.AccessToken,
RefreshToken: tokenRes.RefreshToken,
AccountID: accountID,
LastRefresh: time.Now().Format(time.RFC3339),
Expired: tokenRes.ExpiresAt.Format(time.RFC3339),
Email: email,
Type: "codex",
}
encoded, err := common.Marshal(key)
if err != nil {
common.ApiError(c, err)
return
}
session.Delete(codexOAuthSessionKey(channelID, "state"))
session.Delete(codexOAuthSessionKey(channelID, "verifier"))
session.Delete(codexOAuthSessionKey(channelID, "created_at"))
_ = session.Save()
if channelID > 0 {
if err := model.DB.Model(&model.Channel{}).Where("id = ?", channelID).Update("key", string(encoded)).Error; err != nil {
common.ApiError(c, err)
return
}
model.InitChannelCache()
service.ResetProxyClientCache()
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "saved",
"data": gin.H{
"channel_id": channelID,
"account_id": accountID,
"email": email,
"expires_at": key.Expired,
"last_refresh": key.LastRefresh,
},
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "generated",
"data": gin.H{
"key": string(encoded),
"account_id": accountID,
"email": email,
"expires_at": key.Expired,
"last_refresh": key.LastRefresh,
},
})
}

124
controller/codex_usage.go Normal file
View File

@@ -0,0 +1,124 @@
package controller
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel/codex"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
func GetCodexChannelUsage(c *gin.Context) {
channelId, err := strconv.Atoi(c.Param("id"))
if err != nil {
common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
return
}
ch, err := model.GetChannelById(channelId, true)
if err != nil {
common.ApiError(c, err)
return
}
if ch == nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel not found"})
return
}
if ch.Type != constant.ChannelTypeCodex {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
return
}
if ch.ChannelInfo.IsMultiKey {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "multi-key channel is not supported"})
return
}
oauthKey, err := codex.ParseOAuthKey(strings.TrimSpace(ch.Key))
if err != nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
return
}
accessToken := strings.TrimSpace(oauthKey.AccessToken)
accountID := strings.TrimSpace(oauthKey.AccountID)
if accessToken == "" {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "codex channel: access_token is required"})
return
}
if accountID == "" {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "codex channel: account_id is required"})
return
}
client, err := service.NewProxyHttpClient(ch.GetSetting().Proxy)
if err != nil {
common.ApiError(c, err)
return
}
ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
defer cancel()
statusCode, body, err := service.FetchCodexWhamUsage(ctx, client, ch.GetBaseURL(), accessToken, accountID)
if err != nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
return
}
if (statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden) && strings.TrimSpace(oauthKey.RefreshToken) != "" {
refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
defer refreshCancel()
res, refreshErr := service.RefreshCodexOAuthToken(refreshCtx, oauthKey.RefreshToken)
if refreshErr == nil {
oauthKey.AccessToken = res.AccessToken
oauthKey.RefreshToken = res.RefreshToken
oauthKey.LastRefresh = time.Now().Format(time.RFC3339)
oauthKey.Expired = res.ExpiresAt.Format(time.RFC3339)
if strings.TrimSpace(oauthKey.Type) == "" {
oauthKey.Type = "codex"
}
encoded, encErr := common.Marshal(oauthKey)
if encErr == nil {
_ = model.DB.Model(&model.Channel{}).Where("id = ?", ch.Id).Update("key", string(encoded)).Error
model.InitChannelCache()
service.ResetProxyClientCache()
}
ctx2, cancel2 := context.WithTimeout(c.Request.Context(), 15*time.Second)
defer cancel2()
statusCode, body, err = service.FetchCodexWhamUsage(ctx2, client, ch.GetBaseURL(), oauthKey.AccessToken, accountID)
if err != nil {
c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
return
}
}
}
var payload any
if json.Unmarshal(body, &payload) != nil {
payload = string(body)
}
ok := statusCode >= 200 && statusCode < 300
resp := gin.H{
"success": ok,
"message": "",
"upstream_status": statusCode,
"data": payload,
}
if !ok {
resp["message"] = fmt.Sprintf("upstream status: %d", statusCode)
}
c.JSON(http.StatusOK, resp)
}

810
controller/deployment.go Normal file
View File

@@ -0,0 +1,810 @@
package controller
import (
"bytes"
"encoding/json"
"fmt"
"strconv"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/pkg/ionet"
"github.com/gin-gonic/gin"
)
func getIoAPIKey(c *gin.Context) (string, bool) {
common.OptionMapRWMutex.RLock()
enabled := common.OptionMap["model_deployment.ionet.enabled"] == "true"
apiKey := common.OptionMap["model_deployment.ionet.api_key"]
common.OptionMapRWMutex.RUnlock()
if !enabled || strings.TrimSpace(apiKey) == "" {
common.ApiErrorMsg(c, "io.net model deployment is not enabled or api key missing")
return "", false
}
return apiKey, true
}
func GetModelDeploymentSettings(c *gin.Context) {
common.OptionMapRWMutex.RLock()
enabled := common.OptionMap["model_deployment.ionet.enabled"] == "true"
hasAPIKey := strings.TrimSpace(common.OptionMap["model_deployment.ionet.api_key"]) != ""
common.OptionMapRWMutex.RUnlock()
common.ApiSuccess(c, gin.H{
"provider": "io.net",
"enabled": enabled,
"configured": hasAPIKey,
"can_connect": enabled && hasAPIKey,
})
}
func getIoClient(c *gin.Context) (*ionet.Client, bool) {
apiKey, ok := getIoAPIKey(c)
if !ok {
return nil, false
}
return ionet.NewClient(apiKey), true
}
func getIoEnterpriseClient(c *gin.Context) (*ionet.Client, bool) {
apiKey, ok := getIoAPIKey(c)
if !ok {
return nil, false
}
return ionet.NewEnterpriseClient(apiKey), true
}
func TestIoNetConnection(c *gin.Context) {
var req struct {
APIKey string `json:"api_key"`
}
rawBody, err := c.GetRawData()
if err != nil {
common.ApiError(c, err)
return
}
if len(bytes.TrimSpace(rawBody)) > 0 {
if err := json.Unmarshal(rawBody, &req); err != nil {
common.ApiErrorMsg(c, "invalid request payload")
return
}
}
apiKey := strings.TrimSpace(req.APIKey)
if apiKey == "" {
common.OptionMapRWMutex.RLock()
storedKey := strings.TrimSpace(common.OptionMap["model_deployment.ionet.api_key"])
common.OptionMapRWMutex.RUnlock()
if storedKey == "" {
common.ApiErrorMsg(c, "api_key is required")
return
}
apiKey = storedKey
}
client := ionet.NewEnterpriseClient(apiKey)
result, err := client.GetMaxGPUsPerContainer()
if err != nil {
if apiErr, ok := err.(*ionet.APIError); ok {
message := strings.TrimSpace(apiErr.Message)
if message == "" {
message = "failed to validate api key"
}
common.ApiErrorMsg(c, message)
return
}
common.ApiError(c, err)
return
}
totalHardware := 0
totalAvailable := 0
if result != nil {
totalHardware = len(result.Hardware)
totalAvailable = result.Total
if totalAvailable == 0 {
for _, hw := range result.Hardware {
totalAvailable += hw.Available
}
}
}
common.ApiSuccess(c, gin.H{
"hardware_count": totalHardware,
"total_available": totalAvailable,
})
}
func requireDeploymentID(c *gin.Context) (string, bool) {
deploymentID := strings.TrimSpace(c.Param("id"))
if deploymentID == "" {
common.ApiErrorMsg(c, "deployment ID is required")
return "", false
}
return deploymentID, true
}
func requireContainerID(c *gin.Context) (string, bool) {
containerID := strings.TrimSpace(c.Param("container_id"))
if containerID == "" {
common.ApiErrorMsg(c, "container ID is required")
return "", false
}
return containerID, true
}
func mapIoNetDeployment(d ionet.Deployment) map[string]interface{} {
var created int64
if d.CreatedAt.IsZero() {
created = time.Now().Unix()
} else {
created = d.CreatedAt.Unix()
}
timeRemainingHours := d.ComputeMinutesRemaining / 60
timeRemainingMins := d.ComputeMinutesRemaining % 60
var timeRemaining string
if timeRemainingHours > 0 {
timeRemaining = fmt.Sprintf("%d hour %d minutes", timeRemainingHours, timeRemainingMins)
} else if timeRemainingMins > 0 {
timeRemaining = fmt.Sprintf("%d minutes", timeRemainingMins)
} else {
timeRemaining = "completed"
}
hardwareInfo := fmt.Sprintf("%s %s x%d", d.BrandName, d.HardwareName, d.HardwareQuantity)
return map[string]interface{}{
"id": d.ID,
"deployment_name": d.Name,
"container_name": d.Name,
"status": strings.ToLower(d.Status),
"type": "Container",
"time_remaining": timeRemaining,
"time_remaining_minutes": d.ComputeMinutesRemaining,
"hardware_info": hardwareInfo,
"hardware_name": d.HardwareName,
"brand_name": d.BrandName,
"hardware_quantity": d.HardwareQuantity,
"completed_percent": d.CompletedPercent,
"compute_minutes_served": d.ComputeMinutesServed,
"compute_minutes_remaining": d.ComputeMinutesRemaining,
"created_at": created,
"updated_at": created,
"model_name": "",
"model_version": "",
"instance_count": d.HardwareQuantity,
"resource_config": map[string]interface{}{
"cpu": "",
"memory": "",
"gpu": strconv.Itoa(d.HardwareQuantity),
},
"description": "",
"provider": "io.net",
}
}
func computeStatusCounts(total int, deployments []ionet.Deployment) map[string]int64 {
counts := map[string]int64{
"all": int64(total),
}
for _, status := range []string{"running", "completed", "failed", "deployment requested", "termination requested", "destroyed"} {
counts[status] = 0
}
for _, d := range deployments {
status := strings.ToLower(strings.TrimSpace(d.Status))
counts[status] = counts[status] + 1
}
return counts
}
func GetAllDeployments(c *gin.Context) {
pageInfo := common.GetPageQuery(c)
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
status := c.Query("status")
opts := &ionet.ListDeploymentsOptions{
Status: strings.ToLower(strings.TrimSpace(status)),
Page: pageInfo.GetPage(),
PageSize: pageInfo.GetPageSize(),
SortBy: "created_at",
SortOrder: "desc",
}
dl, err := client.ListDeployments(opts)
if err != nil {
common.ApiError(c, err)
return
}
items := make([]map[string]interface{}, 0, len(dl.Deployments))
for _, d := range dl.Deployments {
items = append(items, mapIoNetDeployment(d))
}
data := gin.H{
"page": pageInfo.GetPage(),
"page_size": pageInfo.GetPageSize(),
"total": dl.Total,
"items": items,
"status_counts": computeStatusCounts(dl.Total, dl.Deployments),
}
common.ApiSuccess(c, data)
}
func SearchDeployments(c *gin.Context) {
pageInfo := common.GetPageQuery(c)
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
status := strings.ToLower(strings.TrimSpace(c.Query("status")))
keyword := strings.TrimSpace(c.Query("keyword"))
dl, err := client.ListDeployments(&ionet.ListDeploymentsOptions{
Status: status,
Page: pageInfo.GetPage(),
PageSize: pageInfo.GetPageSize(),
SortBy: "created_at",
SortOrder: "desc",
})
if err != nil {
common.ApiError(c, err)
return
}
filtered := make([]ionet.Deployment, 0, len(dl.Deployments))
if keyword == "" {
filtered = dl.Deployments
} else {
kw := strings.ToLower(keyword)
for _, d := range dl.Deployments {
if strings.Contains(strings.ToLower(d.Name), kw) {
filtered = append(filtered, d)
}
}
}
items := make([]map[string]interface{}, 0, len(filtered))
for _, d := range filtered {
items = append(items, mapIoNetDeployment(d))
}
total := dl.Total
if keyword != "" {
total = len(filtered)
}
data := gin.H{
"page": pageInfo.GetPage(),
"page_size": pageInfo.GetPageSize(),
"total": total,
"items": items,
}
common.ApiSuccess(c, data)
}
func GetDeployment(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
details, err := client.GetDeployment(deploymentID)
if err != nil {
common.ApiError(c, err)
return
}
data := map[string]interface{}{
"id": details.ID,
"deployment_name": details.ID,
"model_name": "",
"model_version": "",
"status": strings.ToLower(details.Status),
"instance_count": details.TotalContainers,
"hardware_id": details.HardwareID,
"resource_config": map[string]interface{}{
"cpu": "",
"memory": "",
"gpu": strconv.Itoa(details.TotalGPUs),
},
"created_at": details.CreatedAt.Unix(),
"updated_at": details.CreatedAt.Unix(),
"description": "",
"amount_paid": details.AmountPaid,
"completed_percent": details.CompletedPercent,
"gpus_per_container": details.GPUsPerContainer,
"total_gpus": details.TotalGPUs,
"total_containers": details.TotalContainers,
"hardware_name": details.HardwareName,
"brand_name": details.BrandName,
"compute_minutes_served": details.ComputeMinutesServed,
"compute_minutes_remaining": details.ComputeMinutesRemaining,
"locations": details.Locations,
"container_config": details.ContainerConfig,
}
common.ApiSuccess(c, data)
}
func UpdateDeploymentName(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
var req struct {
Name string `json:"name" binding:"required"`
}
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiError(c, err)
return
}
updateReq := &ionet.UpdateClusterNameRequest{
Name: strings.TrimSpace(req.Name),
}
if updateReq.Name == "" {
common.ApiErrorMsg(c, "deployment name cannot be empty")
return
}
available, err := client.CheckClusterNameAvailability(updateReq.Name)
if err != nil {
common.ApiError(c, fmt.Errorf("failed to check name availability: %w", err))
return
}
if !available {
common.ApiErrorMsg(c, "deployment name is not available, please choose a different name")
return
}
resp, err := client.UpdateClusterName(deploymentID, updateReq)
if err != nil {
common.ApiError(c, err)
return
}
data := gin.H{
"status": resp.Status,
"message": resp.Message,
"id": deploymentID,
"name": updateReq.Name,
}
common.ApiSuccess(c, data)
}
func UpdateDeployment(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
var req ionet.UpdateDeploymentRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiError(c, err)
return
}
resp, err := client.UpdateDeployment(deploymentID, &req)
if err != nil {
common.ApiError(c, err)
return
}
data := gin.H{
"status": resp.Status,
"deployment_id": resp.DeploymentID,
}
common.ApiSuccess(c, data)
}
func ExtendDeployment(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
var req ionet.ExtendDurationRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiError(c, err)
return
}
details, err := client.ExtendDeployment(deploymentID, &req)
if err != nil {
common.ApiError(c, err)
return
}
data := mapIoNetDeployment(ionet.Deployment{
ID: details.ID,
Status: details.Status,
Name: deploymentID,
CompletedPercent: float64(details.CompletedPercent),
HardwareQuantity: details.TotalGPUs,
BrandName: details.BrandName,
HardwareName: details.HardwareName,
ComputeMinutesServed: details.ComputeMinutesServed,
ComputeMinutesRemaining: details.ComputeMinutesRemaining,
CreatedAt: details.CreatedAt,
})
common.ApiSuccess(c, data)
}
func DeleteDeployment(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
resp, err := client.DeleteDeployment(deploymentID)
if err != nil {
common.ApiError(c, err)
return
}
data := gin.H{
"status": resp.Status,
"deployment_id": resp.DeploymentID,
"message": "Deployment termination requested successfully",
}
common.ApiSuccess(c, data)
}
func CreateDeployment(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
var req ionet.DeploymentRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiError(c, err)
return
}
resp, err := client.DeployContainer(&req)
if err != nil {
common.ApiError(c, err)
return
}
data := gin.H{
"deployment_id": resp.DeploymentID,
"status": resp.Status,
"message": "Deployment created successfully",
}
common.ApiSuccess(c, data)
}
func GetHardwareTypes(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
hardwareTypes, totalAvailable, err := client.ListHardwareTypes()
if err != nil {
common.ApiError(c, err)
return
}
data := gin.H{
"hardware_types": hardwareTypes,
"total": len(hardwareTypes),
"total_available": totalAvailable,
}
common.ApiSuccess(c, data)
}
func GetLocations(c *gin.Context) {
client, ok := getIoClient(c)
if !ok {
return
}
locationsResp, err := client.ListLocations()
if err != nil {
common.ApiError(c, err)
return
}
total := locationsResp.Total
if total == 0 {
total = len(locationsResp.Locations)
}
data := gin.H{
"locations": locationsResp.Locations,
"total": total,
}
common.ApiSuccess(c, data)
}
func GetAvailableReplicas(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
hardwareIDStr := c.Query("hardware_id")
gpuCountStr := c.Query("gpu_count")
if hardwareIDStr == "" {
common.ApiErrorMsg(c, "hardware_id parameter is required")
return
}
hardwareID, err := strconv.Atoi(hardwareIDStr)
if err != nil || hardwareID <= 0 {
common.ApiErrorMsg(c, "invalid hardware_id parameter")
return
}
gpuCount := 1
if gpuCountStr != "" {
if parsed, err := strconv.Atoi(gpuCountStr); err == nil && parsed > 0 {
gpuCount = parsed
}
}
replicas, err := client.GetAvailableReplicas(hardwareID, gpuCount)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, replicas)
}
func GetPriceEstimation(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
var req ionet.PriceEstimationRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ApiError(c, err)
return
}
priceResp, err := client.GetPriceEstimation(&req)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, priceResp)
}
func CheckClusterNameAvailability(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
clusterName := strings.TrimSpace(c.Query("name"))
if clusterName == "" {
common.ApiErrorMsg(c, "name parameter is required")
return
}
available, err := client.CheckClusterNameAvailability(clusterName)
if err != nil {
common.ApiError(c, err)
return
}
data := gin.H{
"available": available,
"name": clusterName,
}
common.ApiSuccess(c, data)
}
func GetDeploymentLogs(c *gin.Context) {
client, ok := getIoClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
containerID := c.Query("container_id")
if containerID == "" {
common.ApiErrorMsg(c, "container_id parameter is required")
return
}
level := c.Query("level")
stream := c.Query("stream")
cursor := c.Query("cursor")
limitStr := c.Query("limit")
follow := c.Query("follow") == "true"
var limit int = 100
if limitStr != "" {
if parsedLimit, err := strconv.Atoi(limitStr); err == nil && parsedLimit > 0 {
limit = parsedLimit
if limit > 1000 {
limit = 1000
}
}
}
opts := &ionet.GetLogsOptions{
Level: level,
Stream: stream,
Limit: limit,
Cursor: cursor,
Follow: follow,
}
if startTime := c.Query("start_time"); startTime != "" {
if t, err := time.Parse(time.RFC3339, startTime); err == nil {
opts.StartTime = &t
}
}
if endTime := c.Query("end_time"); endTime != "" {
if t, err := time.Parse(time.RFC3339, endTime); err == nil {
opts.EndTime = &t
}
}
rawLogs, err := client.GetContainerLogsRaw(deploymentID, containerID, opts)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, rawLogs)
}
func ListDeploymentContainers(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
containers, err := client.ListContainers(deploymentID)
if err != nil {
common.ApiError(c, err)
return
}
items := make([]map[string]interface{}, 0)
if containers != nil {
items = make([]map[string]interface{}, 0, len(containers.Workers))
for _, ctr := range containers.Workers {
events := make([]map[string]interface{}, 0, len(ctr.ContainerEvents))
for _, event := range ctr.ContainerEvents {
events = append(events, map[string]interface{}{
"time": event.Time.Unix(),
"message": event.Message,
})
}
items = append(items, map[string]interface{}{
"container_id": ctr.ContainerID,
"device_id": ctr.DeviceID,
"status": strings.ToLower(strings.TrimSpace(ctr.Status)),
"hardware": ctr.Hardware,
"brand_name": ctr.BrandName,
"created_at": ctr.CreatedAt.Unix(),
"uptime_percent": ctr.UptimePercent,
"gpus_per_container": ctr.GPUsPerContainer,
"public_url": ctr.PublicURL,
"events": events,
})
}
}
response := gin.H{
"total": 0,
"containers": items,
}
if containers != nil {
response["total"] = containers.Total
}
common.ApiSuccess(c, response)
}
func GetContainerDetails(c *gin.Context) {
client, ok := getIoEnterpriseClient(c)
if !ok {
return
}
deploymentID, ok := requireDeploymentID(c)
if !ok {
return
}
containerID, ok := requireContainerID(c)
if !ok {
return
}
details, err := client.GetContainerDetails(deploymentID, containerID)
if err != nil {
common.ApiError(c, err)
return
}
if details == nil {
common.ApiErrorMsg(c, "container details not found")
return
}
events := make([]map[string]interface{}, 0, len(details.ContainerEvents))
for _, event := range details.ContainerEvents {
events = append(events, map[string]interface{}{
"time": event.Time.Unix(),
"message": event.Message,
})
}
data := gin.H{
"deployment_id": deploymentID,
"container_id": details.ContainerID,
"device_id": details.DeviceID,
"status": strings.ToLower(strings.TrimSpace(details.Status)),
"hardware": details.Hardware,
"brand_name": details.BrandName,
"created_at": details.CreatedAt.Unix(),
"uptime_percent": details.UptimePercent,
"gpus_per_container": details.GPUsPerContainer,
"public_url": details.PublicURL,
"events": events,
}
common.ApiSuccess(c, data)
}

View File

@@ -114,6 +114,7 @@ func GetStatus(c *gin.Context) {
"setup": constant.Setup,
"user_agreement_enabled": legalSetting.UserAgreement != "",
"privacy_policy_enabled": legalSetting.PrivacyPolicy != "",
"checkin_enabled": operation_setting.GetCheckinSetting().Enabled,
}
// 根据启用状态注入可选内容

View File

@@ -99,6 +99,9 @@ func newHTTPClient() *http.Client {
ExpectContinueTimeout: 1 * time.Second,
ResponseHeaderTimeout: time.Duration(timeoutSec) * time.Second,
}
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
}
transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
host, _, err := net.SplitHostPort(addr)
if err != nil {
@@ -115,7 +118,17 @@ func newHTTPClient() *http.Client {
return &http.Client{Transport: transport}
}
var httpClient = newHTTPClient()
var (
httpClientOnce sync.Once
httpClient *http.Client
)
func getHTTPClient() *http.Client {
httpClientOnce.Do(func() {
httpClient = newHTTPClient()
})
return httpClient
}
func fetchJSON[T any](ctx context.Context, url string, out *upstreamEnvelope[T]) error {
var lastErr error
@@ -138,7 +151,7 @@ func fetchJSON[T any](ctx context.Context, url string, out *upstreamEnvelope[T])
}
cacheMutex.RUnlock()
resp, err := httpClient.Do(req)
resp, err := getHTTPClient().Do(req)
if err != nil {
lastErr = err
// backoff with jitter

View File

@@ -10,6 +10,7 @@ import (
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/console_setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/setting/system_setting"
@@ -20,7 +21,11 @@ func GetOptions(c *gin.Context) {
var options []*model.Option
common.OptionMapRWMutex.Lock()
for k, v := range common.OptionMap {
if strings.HasSuffix(k, "Token") || strings.HasSuffix(k, "Secret") || strings.HasSuffix(k, "Key") {
if strings.HasSuffix(k, "Token") ||
strings.HasSuffix(k, "Secret") ||
strings.HasSuffix(k, "Key") ||
strings.HasSuffix(k, "secret") ||
strings.HasSuffix(k, "api_key") {
continue
}
options = append(options, &model.Option{
@@ -173,6 +178,24 @@ func UpdateOption(c *gin.Context) {
})
return
}
case "AutomaticDisableStatusCodes":
_, err = operation_setting.ParseHTTPStatusCodeRanges(option.Value.(string))
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
case "AutomaticRetryStatusCodes":
_, err = operation_setting.ParseHTTPStatusCodeRanges(option.Value.(string))
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
case "console_setting.api_info":
err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo")
if err != nil {

View File

@@ -11,6 +11,7 @@ import (
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/dto"
@@ -110,6 +111,9 @@ func FetchUpstreamRatios(c *gin.Context) {
dialer := &net.Dialer{Timeout: 10 * time.Second}
transport := &http.Transport{MaxIdleConns: 100, IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ExpectContinueTimeout: 1 * time.Second, ResponseHeaderTimeout: 10 * time.Second}
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
}
transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
host, _, err := net.SplitHostPort(addr)
if err != nil {

View File

@@ -21,6 +21,7 @@ import (
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/types"
"github.com/bytedance/gopkg/util/gopool"
@@ -316,30 +317,14 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b
if _, ok := c.Get("specific_channel_id"); ok {
return false
}
if openaiErr.StatusCode == http.StatusTooManyRequests {
return true
}
if openaiErr.StatusCode == 307 {
return true
}
if openaiErr.StatusCode/100 == 5 {
// 超时不重试
if openaiErr.StatusCode == 504 || openaiErr.StatusCode == 524 {
return false
}
return true
}
if openaiErr.StatusCode == http.StatusBadRequest {
code := openaiErr.StatusCode
if code >= 200 && code < 300 {
return false
}
if openaiErr.StatusCode == 408 {
// azure处理超时不重试
return false
if code < 100 || code > 599 {
return true
}
if openaiErr.StatusCode/100 == 2 {
return false
}
return true
return operation_setting.ShouldRetryByStatusCode(code)
}
func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) {
@@ -348,7 +333,7 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
if service.ShouldDisableChannel(channelError.ChannelType, err) && channelError.AutoBan {
gopool.Go(func() {
service.DisableChannel(channelError, err.Error())
service.DisableChannel(channelError, err.ErrorWithStatusCode())
})
}
@@ -378,7 +363,7 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t
adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex)
}
other["admin_info"] = adminInfo
model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveError(), tokenId, 0, false, userGroup, other)
model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, 0, false, userGroup, other)
}
}

View File

@@ -74,7 +74,13 @@ func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, cha
logger.LogError(ctx, fmt.Sprintf("Task %s not found in taskM", taskId))
return fmt.Errorf("task %s not found", taskId)
}
resp, err := adaptor.FetchTask(baseURL, channel.Key, map[string]any{
key := channel.Key
privateData := task.PrivateData
if privateData.Key != "" {
key = privateData.Key
}
resp, err := adaptor.FetchTask(baseURL, key, map[string]any{
"task_id": taskId,
"action": task.Action,
}, proxy)

View File

@@ -1,6 +1,7 @@
package controller
import (
"fmt"
"net/http"
"strconv"
"strings"
@@ -149,6 +150,24 @@ func AddToken(c *gin.Context) {
})
return
}
// 非无限额度时,检查额度值是否超出有效范围
if !token.UnlimitedQuota {
if token.RemainQuota < 0 {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "额度值不能为负数",
})
return
}
maxQuotaValue := int((1000000000 * common.QuotaPerUnit))
if token.RemainQuota > maxQuotaValue {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("额度值超出有效范围,最大值为 %d", maxQuotaValue),
})
return
}
}
key, err := common.GenerateKey()
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -216,6 +235,23 @@ func UpdateToken(c *gin.Context) {
})
return
}
if !token.UnlimitedQuota {
if token.RemainQuota < 0 {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "额度值不能为负数",
})
return
}
maxQuotaValue := int((1000000000 * common.QuotaPerUnit))
if token.RemainQuota > maxQuotaValue {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("额度值超出有效范围,最大值为 %d", maxQuotaValue),
})
return
}
}
cleanToken, err := model.GetTokenByIds(token.Id, userId)
if err != nil {
common.ApiError(c, err)
@@ -261,7 +297,6 @@ func UpdateToken(c *gin.Context) {
"message": "",
"data": cleanToken,
})
return
}
type TokenBatch struct {

7
docs/ionet-client.md Normal file
View File

@@ -0,0 +1,7 @@
Request URL
https://api.io.solutions/v1/io-cloud/clusters/654fc0a9-0d4a-4db4-9b95-3f56189348a2/update-name
Request Method
PUT
{"status":"succeeded","message":"Cluster name updated successfully"}

View File

@@ -26,7 +26,8 @@ type GeneralErrorResponse struct {
Msg string `json:"msg"`
Err string `json:"err"`
ErrorMsg string `json:"error_msg"`
Metadata json.RawMessage `json:"metadata,omitempty"`
Metadata json.RawMessage `json:"metadata,omitempty"`
Detail string `json:"detail,omitempty"`
Header struct {
Message string `json:"message"`
} `json:"header"`
@@ -79,6 +80,9 @@ func (e GeneralErrorResponse) ToMessage() string {
if e.ErrorMsg != "" {
return e.ErrorMsg
}
if e.Detail != "" {
return e.Detail
}
if e.Header.Message != "" {
return e.Header.Message
}

View File

@@ -126,7 +126,7 @@ func (r *GeminiChatRequest) SetModelName(modelName string) {
func (r *GeminiChatRequest) GetTools() []GeminiChatTool {
var tools []GeminiChatTool
if strings.HasSuffix(string(r.Tools), "[") {
if strings.HasPrefix(string(r.Tools), "[") {
// is array
if err := common.Unmarshal(r.Tools, &tools); err != nil {
logger.LogError(nil, "error_unmarshalling_tools: "+err.Error())
@@ -341,6 +341,88 @@ type GeminiChatGenerationConfig struct {
ImageConfig json.RawMessage `json:"imageConfig,omitempty"` // RawMessage to allow flexible image config
}
// UnmarshalJSON allows GeminiChatGenerationConfig to accept both snake_case and camelCase fields.
func (c *GeminiChatGenerationConfig) UnmarshalJSON(data []byte) error {
type Alias GeminiChatGenerationConfig
var aux struct {
Alias
TopPSnake float64 `json:"top_p,omitempty"`
TopKSnake float64 `json:"top_k,omitempty"`
MaxOutputTokensSnake uint `json:"max_output_tokens,omitempty"`
CandidateCountSnake int `json:"candidate_count,omitempty"`
StopSequencesSnake []string `json:"stop_sequences,omitempty"`
ResponseMimeTypeSnake string `json:"response_mime_type,omitempty"`
ResponseSchemaSnake any `json:"response_schema,omitempty"`
ResponseJsonSchemaSnake json.RawMessage `json:"response_json_schema,omitempty"`
PresencePenaltySnake *float32 `json:"presence_penalty,omitempty"`
FrequencyPenaltySnake *float32 `json:"frequency_penalty,omitempty"`
ResponseLogprobsSnake bool `json:"response_logprobs,omitempty"`
MediaResolutionSnake MediaResolution `json:"media_resolution,omitempty"`
ResponseModalitiesSnake []string `json:"response_modalities,omitempty"`
ThinkingConfigSnake *GeminiThinkingConfig `json:"thinking_config,omitempty"`
SpeechConfigSnake json.RawMessage `json:"speech_config,omitempty"`
ImageConfigSnake json.RawMessage `json:"image_config,omitempty"`
}
if err := common.Unmarshal(data, &aux); err != nil {
return err
}
*c = GeminiChatGenerationConfig(aux.Alias)
// Prioritize snake_case if present
if aux.TopPSnake != 0 {
c.TopP = aux.TopPSnake
}
if aux.TopKSnake != 0 {
c.TopK = aux.TopKSnake
}
if aux.MaxOutputTokensSnake != 0 {
c.MaxOutputTokens = aux.MaxOutputTokensSnake
}
if aux.CandidateCountSnake != 0 {
c.CandidateCount = aux.CandidateCountSnake
}
if len(aux.StopSequencesSnake) > 0 {
c.StopSequences = aux.StopSequencesSnake
}
if aux.ResponseMimeTypeSnake != "" {
c.ResponseMimeType = aux.ResponseMimeTypeSnake
}
if aux.ResponseSchemaSnake != nil {
c.ResponseSchema = aux.ResponseSchemaSnake
}
if len(aux.ResponseJsonSchemaSnake) > 0 {
c.ResponseJsonSchema = aux.ResponseJsonSchemaSnake
}
if aux.PresencePenaltySnake != nil {
c.PresencePenalty = aux.PresencePenaltySnake
}
if aux.FrequencyPenaltySnake != nil {
c.FrequencyPenalty = aux.FrequencyPenaltySnake
}
if aux.ResponseLogprobsSnake {
c.ResponseLogprobs = aux.ResponseLogprobsSnake
}
if aux.MediaResolutionSnake != "" {
c.MediaResolution = aux.MediaResolutionSnake
}
if len(aux.ResponseModalitiesSnake) > 0 {
c.ResponseModalities = aux.ResponseModalitiesSnake
}
if aux.ThinkingConfigSnake != nil {
c.ThinkingConfig = aux.ThinkingConfigSnake
}
if len(aux.SpeechConfigSnake) > 0 {
c.SpeechConfig = aux.SpeechConfigSnake
}
if len(aux.ImageConfigSnake) > 0 {
c.ImageConfig = aux.ImageConfigSnake
}
return nil
}
type MediaResolution string
type GeminiChatCandidate struct {

View File

@@ -167,9 +167,9 @@ func (i *ImageRequest) SetModelName(modelName string) {
}
type ImageResponse struct {
Data []ImageData `json:"data"`
Created int64 `json:"created"`
Extra any `json:"extra,omitempty"`
Data []ImageData `json:"data"`
Created int64 `json:"created"`
Metadata json.RawMessage `json:"metadata,omitempty"`
}
type ImageData struct {
Url string `json:"url"`

View File

@@ -23,6 +23,8 @@ type FormatJsonSchema struct {
Strict json.RawMessage `json:"strict,omitempty"`
}
// GeneralOpenAIRequest represents a general request structure for OpenAI-compatible APIs.
// 参数增加规范无引用的参数必须使用json.RawMessage类型并添加omitempty标签
type GeneralOpenAIRequest struct {
Model string `json:"model,omitempty"`
Messages []Message `json:"messages,omitempty"`
@@ -82,8 +84,9 @@ type GeneralOpenAIRequest struct {
Reasoning json.RawMessage `json:"reasoning,omitempty"`
// Ali Qwen Params
VlHighResolutionImages json.RawMessage `json:"vl_high_resolution_images,omitempty"`
EnableThinking any `json:"enable_thinking,omitempty"`
EnableThinking json.RawMessage `json:"enable_thinking,omitempty"`
ChatTemplateKwargs json.RawMessage `json:"chat_template_kwargs,omitempty"`
EnableSearch json.RawMessage `json:"enable_search,omitempty"`
// ollama Params
Think json.RawMessage `json:"think,omitempty"`
// baidu v2
@@ -805,11 +808,11 @@ type OpenAIResponsesRequest struct {
PromptCacheKey json.RawMessage `json:"prompt_cache_key,omitempty"`
PromptCacheRetention json.RawMessage `json:"prompt_cache_retention,omitempty"`
Stream bool `json:"stream,omitempty"`
Temperature float64 `json:"temperature,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
Text json.RawMessage `json:"text,omitempty"`
ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
Tools json.RawMessage `json:"tools,omitempty"` // 需要处理的参数很少MCP 参数太多不确定,所以用 map
TopP float64 `json:"top_p,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Truncation string `json:"truncation,omitempty"`
User string `json:"user,omitempty"`
MaxToolCalls uint `json:"max_tool_calls,omitempty"`

View File

@@ -334,13 +334,16 @@ type IncompleteDetails struct {
}
type ResponsesOutput struct {
Type string `json:"type"`
ID string `json:"id"`
Status string `json:"status"`
Role string `json:"role"`
Content []ResponsesOutputContent `json:"content"`
Quality string `json:"quality"`
Size string `json:"size"`
Type string `json:"type"`
ID string `json:"id"`
Status string `json:"status"`
Role string `json:"role"`
Content []ResponsesOutputContent `json:"content"`
Quality string `json:"quality"`
Size string `json:"size"`
CallId string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
}
type ResponsesOutputContent struct {
@@ -369,6 +372,10 @@ type ResponsesStreamResponse struct {
Response *OpenAIResponsesResponse `json:"response,omitempty"`
Delta string `json:"delta,omitempty"`
Item *ResponsesOutput `json:"item,omitempty"`
// - response.function_call_arguments.delta
// - response.function_call_arguments.done
OutputIndex *int `json:"output_index,omitempty"`
ItemID string `json:"item_id,omitempty"`
}
// GetOpenAIError 从动态错误类型中提取OpenAIError结构

55
dto/values.go Normal file
View File

@@ -0,0 +1,55 @@
package dto
import (
"encoding/json"
"strconv"
)
type IntValue int
func (i *IntValue) UnmarshalJSON(b []byte) error {
var n int
if err := json.Unmarshal(b, &n); err == nil {
*i = IntValue(n)
return nil
}
var s string
if err := json.Unmarshal(b, &s); err != nil {
return err
}
v, err := strconv.Atoi(s)
if err != nil {
return err
}
*i = IntValue(v)
return nil
}
func (i IntValue) MarshalJSON() ([]byte, error) {
return json.Marshal(int(i))
}
type BoolValue bool
func (b *BoolValue) UnmarshalJSON(data []byte) error {
var boolean bool
if err := json.Unmarshal(data, &boolean); err == nil {
*b = BoolValue(boolean)
return nil
}
var str string
if err := json.Unmarshal(data, &str); err != nil {
return err
}
if str == "true" {
*b = BoolValue(true)
} else if str == "false" {
*b = BoolValue(false)
} else {
return json.Unmarshal(data, &boolean)
}
return nil
}
func (b BoolValue) MarshalJSON() ([]byte, error) {
return json.Marshal(bool(b))
}

4
go.mod
View File

@@ -37,6 +37,7 @@ require (
github.com/samber/lo v1.52.0
github.com/shirou/gopsutil v3.21.11+incompatible
github.com/shopspring/decimal v1.4.0
github.com/stretchr/testify v1.11.1
github.com/stripe/stripe-go/v81 v81.4.0
github.com/tcolgate/mp3 v0.0.0-20170426193717-e79c5a46d300
github.com/thanhpk/randstr v1.0.6
@@ -63,6 +64,7 @@ require (
github.com/bytedance/sonic/loader v0.3.0 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cloudwego/base64x v0.1.6 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
github.com/dlclark/regexp2 v1.11.5 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
@@ -103,7 +105,9 @@ require (
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/ncruces/go-strftime v0.1.9 // indirect
github.com/pelletier/go-toml/v2 v2.2.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/stretchr/objx v0.5.2 // indirect
github.com/tidwall/match v1.1.1 // indirect
github.com/tidwall/pretty v1.2.0 // indirect
github.com/tklauser/go-sysconf v0.3.12 // indirect

View File

@@ -102,6 +102,9 @@ func main() {
go controller.AutomaticallyTestChannels()
// Codex credential auto-refresh check every 10 minutes, refresh when expires within 1 day
service.StartCodexCredentialAutoRefreshTask()
if common.IsMasterNode && constant.UpdateTask {
gopool.Go(func() {
controller.UpdateMidjourneyTaskBulk()
@@ -188,6 +191,7 @@ func InjectUmamiAnalytics() {
analyticsInjectBuilder.WriteString(umamiSiteID)
analyticsInjectBuilder.WriteString("\"></script>")
}
analyticsInjectBuilder.WriteString("<!--Umami QuantumNous-->\n")
analyticsInject := analyticsInjectBuilder.String()
indexPage = bytes.ReplaceAll(indexPage, []byte("<!--umami-->\n"), []byte(analyticsInject))
}
@@ -209,6 +213,7 @@ func InjectGoogleAnalytics() {
analyticsInjectBuilder.WriteString("');")
analyticsInjectBuilder.WriteString("</script>")
}
analyticsInjectBuilder.WriteString("<!--Google Analytics QuantumNous-->\n")
analyticsInject := analyticsInjectBuilder.String()
indexPage = bytes.ReplaceAll(indexPage, []byte("<!--Google Analytics-->\n"), []byte(analyticsInject))
}

View File

@@ -13,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
@@ -195,8 +196,8 @@ func TokenAuth() func(c *gin.Context) {
}
c.Request.Header.Set("Authorization", "Bearer "+key)
}
// 检查path包含/v1/messages
if strings.Contains(c.Request.URL.Path, "/v1/messages") {
// 检查path包含/v1/messages 或 /v1/models
if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") {
anthropicKey := c.Request.Header.Get("x-api-key")
if anthropicKey != "" {
c.Request.Header.Set("Authorization", "Bearer "+anthropicKey)
@@ -256,7 +257,7 @@ func TokenAuth() func(c *gin.Context) {
return
}
if common.IsIpInCIDRList(ip, allowIps) == false {
abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中")
abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中", types.ErrorCodeAccessDenied)
return
}
logger.LogDebug(c, "Client IP %s passed the token IP restrictions check", clientIp)

View File

@@ -114,11 +114,11 @@ func Distribute() func(c *gin.Context) {
// common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id))
// message = "数据库一致性已被破坏,请联系管理员"
//}
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, string(types.ErrorCodeModelNotFound))
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound)
return
}
if channel == nil {
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("分组 %s 下模型 %s 无可用渠道distributor", usingGroup, modelRequest.Model), string(types.ErrorCodeModelNotFound))
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("分组 %s 下模型 %s 无可用渠道distributor", usingGroup, modelRequest.Model), types.ErrorCodeModelNotFound)
return
}
}

View File

@@ -5,13 +5,14 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string, code ...string) {
func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string, code ...types.ErrorCode) {
codeStr := ""
if len(code) > 0 {
codeStr = code[0]
codeStr = string(code[0])
}
userId := c.GetInt("id")
c.JSON(statusCode, gin.H{

179
model/checkin.go Normal file
View File

@@ -0,0 +1,179 @@
package model
import (
"errors"
"math/rand"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/setting/operation_setting"
"gorm.io/gorm"
)
// Checkin 签到记录
type Checkin struct {
Id int `json:"id" gorm:"primaryKey;autoIncrement"`
UserId int `json:"user_id" gorm:"not null;uniqueIndex:idx_user_checkin_date"`
CheckinDate string `json:"checkin_date" gorm:"type:varchar(10);not null;uniqueIndex:idx_user_checkin_date"` // 格式: YYYY-MM-DD
QuotaAwarded int `json:"quota_awarded" gorm:"not null"`
CreatedAt int64 `json:"created_at" gorm:"bigint"`
}
// CheckinRecord 用于API返回的签到记录不包含敏感字段
type CheckinRecord struct {
CheckinDate string `json:"checkin_date"`
QuotaAwarded int `json:"quota_awarded"`
}
func (Checkin) TableName() string {
return "checkins"
}
// GetUserCheckinRecords 获取用户在指定日期范围内的签到记录
func GetUserCheckinRecords(userId int, startDate, endDate string) ([]Checkin, error) {
var records []Checkin
err := DB.Where("user_id = ? AND checkin_date >= ? AND checkin_date <= ?",
userId, startDate, endDate).
Order("checkin_date DESC").
Find(&records).Error
return records, err
}
// HasCheckedInToday 检查用户今天是否已签到
func HasCheckedInToday(userId int) (bool, error) {
today := time.Now().Format("2006-01-02")
var count int64
err := DB.Model(&Checkin{}).
Where("user_id = ? AND checkin_date = ?", userId, today).
Count(&count).Error
return count > 0, err
}
// UserCheckin 执行用户签到
// MySQL 和 PostgreSQL 使用事务保证原子性
// SQLite 不支持嵌套事务,使用顺序操作 + 手动回滚
func UserCheckin(userId int) (*Checkin, error) {
setting := operation_setting.GetCheckinSetting()
if !setting.Enabled {
return nil, errors.New("签到功能未启用")
}
// 检查今天是否已签到
hasChecked, err := HasCheckedInToday(userId)
if err != nil {
return nil, err
}
if hasChecked {
return nil, errors.New("今日已签到")
}
// 计算随机额度奖励
quotaAwarded := setting.MinQuota
if setting.MaxQuota > setting.MinQuota {
quotaAwarded = setting.MinQuota + rand.Intn(setting.MaxQuota-setting.MinQuota+1)
}
today := time.Now().Format("2006-01-02")
checkin := &Checkin{
UserId: userId,
CheckinDate: today,
QuotaAwarded: quotaAwarded,
CreatedAt: time.Now().Unix(),
}
// 根据数据库类型选择不同的策略
if common.UsingSQLite {
// SQLite 不支持嵌套事务,使用顺序操作 + 手动回滚
return userCheckinWithoutTransaction(checkin, userId, quotaAwarded)
}
// MySQL 和 PostgreSQL 支持事务,使用事务保证原子性
return userCheckinWithTransaction(checkin, userId, quotaAwarded)
}
// userCheckinWithTransaction 使用事务执行签到(适用于 MySQL 和 PostgreSQL
func userCheckinWithTransaction(checkin *Checkin, userId int, quotaAwarded int) (*Checkin, error) {
err := DB.Transaction(func(tx *gorm.DB) error {
// 步骤1: 创建签到记录
// 数据库有唯一约束 (user_id, checkin_date),可以防止并发重复签到
if err := tx.Create(checkin).Error; err != nil {
return errors.New("签到失败,请稍后重试")
}
// 步骤2: 在事务中增加用户额度
if err := tx.Model(&User{}).Where("id = ?", userId).
Update("quota", gorm.Expr("quota + ?", quotaAwarded)).Error; err != nil {
return errors.New("签到失败:更新额度出错")
}
return nil
})
if err != nil {
return nil, err
}
// 事务成功后,异步更新缓存
go func() {
_ = cacheIncrUserQuota(userId, int64(quotaAwarded))
}()
return checkin, nil
}
// userCheckinWithoutTransaction 不使用事务执行签到(适用于 SQLite
func userCheckinWithoutTransaction(checkin *Checkin, userId int, quotaAwarded int) (*Checkin, error) {
// 步骤1: 创建签到记录
// 数据库有唯一约束 (user_id, checkin_date),可以防止并发重复签到
if err := DB.Create(checkin).Error; err != nil {
return nil, errors.New("签到失败,请稍后重试")
}
// 步骤2: 增加用户额度
// 使用 db=true 强制直接写入数据库,不使用批量更新
if err := IncreaseUserQuota(userId, quotaAwarded, true); err != nil {
// 如果增加额度失败,需要回滚签到记录
DB.Delete(checkin)
return nil, errors.New("签到失败:更新额度出错")
}
return checkin, nil
}
// GetUserCheckinStats 获取用户签到统计信息
func GetUserCheckinStats(userId int, month string) (map[string]interface{}, error) {
// 获取指定月份的所有签到记录
startDate := month + "-01"
endDate := month + "-31"
records, err := GetUserCheckinRecords(userId, startDate, endDate)
if err != nil {
return nil, err
}
// 转换为不包含敏感字段的记录
checkinRecords := make([]CheckinRecord, len(records))
for i, r := range records {
checkinRecords[i] = CheckinRecord{
CheckinDate: r.CheckinDate,
QuotaAwarded: r.QuotaAwarded,
}
}
// 检查今天是否已签到
hasCheckedToday, _ := HasCheckedInToday(userId)
// 获取用户所有时间的签到统计
var totalCheckins int64
var totalQuota int64
DB.Model(&Checkin{}).Where("user_id = ?", userId).Count(&totalCheckins)
DB.Model(&Checkin{}).Where("user_id = ?", userId).Select("COALESCE(SUM(quota_awarded), 0)").Scan(&totalQuota)
return map[string]interface{}{
"total_quota": totalQuota, // 所有时间累计获得的额度
"total_checkins": totalCheckins, // 所有时间累计签到次数
"checkin_count": len(records), // 本月签到次数
"checked_in_today": hasCheckedToday, // 今天是否已签到
"records": checkinRecords, // 本月签到记录详情不含id和user_id
}, nil
}

View File

@@ -56,8 +56,9 @@ func formatUserLogs(logs []*Log) {
var otherMap map[string]interface{}
otherMap, _ = common.StrToMap(logs[i].Other)
if otherMap != nil {
// delete admin
// Remove admin-only debug fields.
delete(otherMap, "admin_info")
delete(otherMap, "request_conversion")
}
logs[i].Other = common.MapToJsonStr(otherMap)
logs[i].Id = logs[i].Id % 1024

View File

@@ -267,6 +267,7 @@ func migrateDB() error {
&Setup{},
&TwoFA{},
&TwoFABackupCode{},
&Checkin{},
)
if err != nil {
return err
@@ -300,6 +301,7 @@ func migrateDBFast() error {
{&Setup{}, "Setup"},
{&TwoFA{}, "TwoFA"},
{&TwoFABackupCode{}, "TwoFABackupCode"},
{&Checkin{}, "Checkin"},
}
// 动态计算migration数量确保errChan缓冲区足够大
errChan := make(chan error, len(migrations))

View File

@@ -143,6 +143,8 @@ func InitOptionMap() {
common.OptionMap["SensitiveWords"] = setting.SensitiveWordsToString()
common.OptionMap["StreamCacheQueueLength"] = strconv.Itoa(setting.StreamCacheQueueLength)
common.OptionMap["AutomaticDisableKeywords"] = operation_setting.AutomaticDisableKeywordsToString()
common.OptionMap["AutomaticDisableStatusCodes"] = operation_setting.AutomaticDisableStatusCodesToString()
common.OptionMap["AutomaticRetryStatusCodes"] = operation_setting.AutomaticRetryStatusCodesToString()
common.OptionMap["ExposeRatioEnabled"] = strconv.FormatBool(ratio_setting.IsExposeRatioEnabled())
// 自动添加所有注册的模型配置
@@ -444,6 +446,10 @@ func updateOptionMap(key string, value string) (err error) {
setting.SensitiveWordsFromString(value)
case "AutomaticDisableKeywords":
operation_setting.AutomaticDisableKeywordsFromString(value)
case "AutomaticDisableStatusCodes":
err = operation_setting.AutomaticDisableStatusCodesFromString(value)
case "AutomaticRetryStatusCodes":
err = operation_setting.AutomaticRetryStatusCodesFromString(value)
case "StreamCacheQueueLength":
setting.StreamCacheQueueLength, _ = strconv.Atoi(value)
case "PayMethods":

View File

@@ -26,7 +26,7 @@ type Token struct {
AllowIps *string `json:"allow_ips" gorm:"default:''"`
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
Group string `json:"group" gorm:"default:''"`
CrossGroupRetry bool `json:"cross_group_retry" gorm:"default:false"` // 跨分组重试仅auto分组有效
CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试仅auto分组有效
DeletedAt gorm.DeletedAt `gorm:"index"`
}

219
pkg/ionet/client.go Normal file
View File

@@ -0,0 +1,219 @@
package ionet
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"net/url"
"strconv"
"time"
)
const (
DefaultEnterpriseBaseURL = "https://api.io.solutions/enterprise/v1/io-cloud/caas"
DefaultBaseURL = "https://api.io.solutions/v1/io-cloud/caas"
DefaultTimeout = 30 * time.Second
)
// DefaultHTTPClient is the default HTTP client implementation
type DefaultHTTPClient struct {
client *http.Client
}
// NewDefaultHTTPClient creates a new default HTTP client
func NewDefaultHTTPClient(timeout time.Duration) *DefaultHTTPClient {
return &DefaultHTTPClient{
client: &http.Client{
Timeout: timeout,
},
}
}
// Do executes an HTTP request
func (c *DefaultHTTPClient) Do(req *HTTPRequest) (*HTTPResponse, error) {
httpReq, err := http.NewRequest(req.Method, req.URL, bytes.NewReader(req.Body))
if err != nil {
return nil, fmt.Errorf("failed to create HTTP request: %w", err)
}
// Set headers
for key, value := range req.Headers {
httpReq.Header.Set(key, value)
}
resp, err := c.client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("HTTP request failed: %w", err)
}
defer resp.Body.Close()
// Read response body
var body bytes.Buffer
_, err = body.ReadFrom(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}
// Convert headers
headers := make(map[string]string)
for key, values := range resp.Header {
if len(values) > 0 {
headers[key] = values[0]
}
}
return &HTTPResponse{
StatusCode: resp.StatusCode,
Headers: headers,
Body: body.Bytes(),
}, nil
}
// NewEnterpriseClient creates a new IO.NET API client targeting the enterprise API base URL.
func NewEnterpriseClient(apiKey string) *Client {
return NewClientWithConfig(apiKey, DefaultEnterpriseBaseURL, nil)
}
// NewClient creates a new IO.NET API client targeting the public API base URL.
func NewClient(apiKey string) *Client {
return NewClientWithConfig(apiKey, DefaultBaseURL, nil)
}
// NewClientWithConfig creates a new IO.NET API client with custom configuration
func NewClientWithConfig(apiKey, baseURL string, httpClient HTTPClient) *Client {
if baseURL == "" {
baseURL = DefaultBaseURL
}
if httpClient == nil {
httpClient = NewDefaultHTTPClient(DefaultTimeout)
}
return &Client{
BaseURL: baseURL,
APIKey: apiKey,
HTTPClient: httpClient,
}
}
// makeRequest performs an HTTP request and handles common response processing
func (c *Client) makeRequest(method, endpoint string, body interface{}) (*HTTPResponse, error) {
var reqBody []byte
var err error
if body != nil {
reqBody, err = json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
}
headers := map[string]string{
"X-API-KEY": c.APIKey,
"Content-Type": "application/json",
}
req := &HTTPRequest{
Method: method,
URL: c.BaseURL + endpoint,
Headers: headers,
Body: reqBody,
}
resp, err := c.HTTPClient.Do(req)
if err != nil {
return nil, fmt.Errorf("request failed: %w", err)
}
// Handle API errors
if resp.StatusCode >= 400 {
var apiErr APIError
if len(resp.Body) > 0 {
// Try to parse the actual error format: {"detail": "message"}
var errorResp struct {
Detail string `json:"detail"`
}
if err := json.Unmarshal(resp.Body, &errorResp); err == nil && errorResp.Detail != "" {
apiErr = APIError{
Code: resp.StatusCode,
Message: errorResp.Detail,
}
} else {
// Fallback: use raw body as details
apiErr = APIError{
Code: resp.StatusCode,
Message: fmt.Sprintf("API request failed with status %d", resp.StatusCode),
Details: string(resp.Body),
}
}
} else {
apiErr = APIError{
Code: resp.StatusCode,
Message: fmt.Sprintf("API request failed with status %d", resp.StatusCode),
}
}
return nil, &apiErr
}
return resp, nil
}
// buildQueryParams builds query parameters for GET requests
func buildQueryParams(params map[string]interface{}) string {
if len(params) == 0 {
return ""
}
values := url.Values{}
for key, value := range params {
if value == nil {
continue
}
switch v := value.(type) {
case string:
if v != "" {
values.Add(key, v)
}
case int:
if v != 0 {
values.Add(key, strconv.Itoa(v))
}
case int64:
if v != 0 {
values.Add(key, strconv.FormatInt(v, 10))
}
case float64:
if v != 0 {
values.Add(key, strconv.FormatFloat(v, 'f', -1, 64))
}
case bool:
values.Add(key, strconv.FormatBool(v))
case time.Time:
if !v.IsZero() {
values.Add(key, v.Format(time.RFC3339))
}
case *time.Time:
if v != nil && !v.IsZero() {
values.Add(key, v.Format(time.RFC3339))
}
case []int:
if len(v) > 0 {
if encoded, err := json.Marshal(v); err == nil {
values.Add(key, string(encoded))
}
}
case []string:
if len(v) > 0 {
if encoded, err := json.Marshal(v); err == nil {
values.Add(key, string(encoded))
}
}
default:
values.Add(key, fmt.Sprint(v))
}
}
if len(values) > 0 {
return "?" + values.Encode()
}
return ""
}

302
pkg/ionet/container.go Normal file
View File

@@ -0,0 +1,302 @@
package ionet
import (
"encoding/json"
"fmt"
"strings"
"time"
"github.com/samber/lo"
)
// ListContainers retrieves all containers for a specific deployment
func (c *Client) ListContainers(deploymentID string) (*ContainerList, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s/containers", deploymentID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to list containers: %w", err)
}
var containerList ContainerList
if err := decodeDataWithFlexibleTimes(resp.Body, &containerList); err != nil {
return nil, fmt.Errorf("failed to parse containers list: %w", err)
}
return &containerList, nil
}
// GetContainerDetails retrieves detailed information about a specific container
func (c *Client) GetContainerDetails(deploymentID, containerID string) (*Container, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return nil, fmt.Errorf("container ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s/container/%s", deploymentID, containerID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get container details: %w", err)
}
// API response format not documented, assuming direct format
var container Container
if err := decodeWithFlexibleTimes(resp.Body, &container); err != nil {
return nil, fmt.Errorf("failed to parse container details: %w", err)
}
return &container, nil
}
// GetContainerJobs retrieves containers jobs for a specific container (similar to containers endpoint)
func (c *Client) GetContainerJobs(deploymentID, containerID string) (*ContainerList, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return nil, fmt.Errorf("container ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s/containers-jobs/%s", deploymentID, containerID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get container jobs: %w", err)
}
var containerList ContainerList
if err := decodeDataWithFlexibleTimes(resp.Body, &containerList); err != nil {
return nil, fmt.Errorf("failed to parse container jobs: %w", err)
}
return &containerList, nil
}
// buildLogEndpoint constructs the request path for fetching logs
func buildLogEndpoint(deploymentID, containerID string, opts *GetLogsOptions) (string, error) {
if deploymentID == "" {
return "", fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return "", fmt.Errorf("container ID cannot be empty")
}
params := make(map[string]interface{})
if opts != nil {
if opts.Level != "" {
params["level"] = opts.Level
}
if opts.Stream != "" {
params["stream"] = opts.Stream
}
if opts.Limit > 0 {
params["limit"] = opts.Limit
}
if opts.Cursor != "" {
params["cursor"] = opts.Cursor
}
if opts.Follow {
params["follow"] = true
}
if opts.StartTime != nil {
params["start_time"] = opts.StartTime
}
if opts.EndTime != nil {
params["end_time"] = opts.EndTime
}
}
endpoint := fmt.Sprintf("/deployment/%s/log/%s", deploymentID, containerID)
endpoint += buildQueryParams(params)
return endpoint, nil
}
// GetContainerLogs retrieves logs for containers in a deployment and normalizes them
func (c *Client) GetContainerLogs(deploymentID, containerID string, opts *GetLogsOptions) (*ContainerLogs, error) {
raw, err := c.GetContainerLogsRaw(deploymentID, containerID, opts)
if err != nil {
return nil, err
}
logs := &ContainerLogs{
ContainerID: containerID,
}
if raw == "" {
return logs, nil
}
normalized := strings.ReplaceAll(raw, "\r\n", "\n")
lines := strings.Split(normalized, "\n")
logs.Logs = lo.FilterMap(lines, func(line string, _ int) (LogEntry, bool) {
if strings.TrimSpace(line) == "" {
return LogEntry{}, false
}
return LogEntry{Message: line}, true
})
return logs, nil
}
// GetContainerLogsRaw retrieves the raw text logs for a specific container
func (c *Client) GetContainerLogsRaw(deploymentID, containerID string, opts *GetLogsOptions) (string, error) {
endpoint, err := buildLogEndpoint(deploymentID, containerID, opts)
if err != nil {
return "", err
}
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return "", fmt.Errorf("failed to get container logs: %w", err)
}
return string(resp.Body), nil
}
// StreamContainerLogs streams real-time logs for a specific container
// This method uses a callback function to handle incoming log entries
func (c *Client) StreamContainerLogs(deploymentID, containerID string, opts *GetLogsOptions, callback func(*LogEntry) error) error {
if deploymentID == "" {
return fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return fmt.Errorf("container ID cannot be empty")
}
if callback == nil {
return fmt.Errorf("callback function cannot be nil")
}
// Set follow to true for streaming
if opts == nil {
opts = &GetLogsOptions{}
}
opts.Follow = true
endpoint, err := buildLogEndpoint(deploymentID, containerID, opts)
if err != nil {
return err
}
// Note: This is a simplified implementation. In a real scenario, you might want to use
// Server-Sent Events (SSE) or WebSocket for streaming logs
for {
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return fmt.Errorf("failed to stream container logs: %w", err)
}
var logs ContainerLogs
if err := decodeWithFlexibleTimes(resp.Body, &logs); err != nil {
return fmt.Errorf("failed to parse container logs: %w", err)
}
// Call the callback for each log entry
for _, logEntry := range logs.Logs {
if err := callback(&logEntry); err != nil {
return fmt.Errorf("callback error: %w", err)
}
}
// If there are no more logs or we have a cursor, continue polling
if !logs.HasMore && logs.NextCursor == "" {
break
}
// Update cursor for next request
if logs.NextCursor != "" {
opts.Cursor = logs.NextCursor
endpoint, err = buildLogEndpoint(deploymentID, containerID, opts)
if err != nil {
return err
}
}
// Wait a bit before next poll to avoid overwhelming the API
time.Sleep(2 * time.Second)
}
return nil
}
// RestartContainer restarts a specific container (if supported by the API)
func (c *Client) RestartContainer(deploymentID, containerID string) error {
if deploymentID == "" {
return fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return fmt.Errorf("container ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s/container/%s/restart", deploymentID, containerID)
_, err := c.makeRequest("POST", endpoint, nil)
if err != nil {
return fmt.Errorf("failed to restart container: %w", err)
}
return nil
}
// StopContainer stops a specific container (if supported by the API)
func (c *Client) StopContainer(deploymentID, containerID string) error {
if deploymentID == "" {
return fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return fmt.Errorf("container ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s/container/%s/stop", deploymentID, containerID)
_, err := c.makeRequest("POST", endpoint, nil)
if err != nil {
return fmt.Errorf("failed to stop container: %w", err)
}
return nil
}
// ExecuteInContainer executes a command in a specific container (if supported by the API)
func (c *Client) ExecuteInContainer(deploymentID, containerID string, command []string) (string, error) {
if deploymentID == "" {
return "", fmt.Errorf("deployment ID cannot be empty")
}
if containerID == "" {
return "", fmt.Errorf("container ID cannot be empty")
}
if len(command) == 0 {
return "", fmt.Errorf("command cannot be empty")
}
reqBody := map[string]interface{}{
"command": command,
}
endpoint := fmt.Sprintf("/deployment/%s/container/%s/exec", deploymentID, containerID)
resp, err := c.makeRequest("POST", endpoint, reqBody)
if err != nil {
return "", fmt.Errorf("failed to execute command in container: %w", err)
}
var result map[string]interface{}
if err := json.Unmarshal(resp.Body, &result); err != nil {
return "", fmt.Errorf("failed to parse execution result: %w", err)
}
if output, ok := result["output"].(string); ok {
return output, nil
}
return string(resp.Body), nil
}

377
pkg/ionet/deployment.go Normal file
View File

@@ -0,0 +1,377 @@
package ionet
import (
"encoding/json"
"fmt"
"strings"
"github.com/samber/lo"
)
// DeployContainer deploys a new container with the specified configuration
func (c *Client) DeployContainer(req *DeploymentRequest) (*DeploymentResponse, error) {
if req == nil {
return nil, fmt.Errorf("deployment request cannot be nil")
}
// Validate required fields
if req.ResourcePrivateName == "" {
return nil, fmt.Errorf("resource_private_name is required")
}
if len(req.LocationIDs) == 0 {
return nil, fmt.Errorf("location_ids is required")
}
if req.HardwareID <= 0 {
return nil, fmt.Errorf("hardware_id is required")
}
if req.RegistryConfig.ImageURL == "" {
return nil, fmt.Errorf("registry_config.image_url is required")
}
if req.GPUsPerContainer < 1 {
return nil, fmt.Errorf("gpus_per_container must be at least 1")
}
if req.DurationHours < 1 {
return nil, fmt.Errorf("duration_hours must be at least 1")
}
if req.ContainerConfig.ReplicaCount < 1 {
return nil, fmt.Errorf("container_config.replica_count must be at least 1")
}
resp, err := c.makeRequest("POST", "/deploy", req)
if err != nil {
return nil, fmt.Errorf("failed to deploy container: %w", err)
}
// API returns direct format:
// {"status": "string", "deployment_id": "..."}
var deployResp DeploymentResponse
if err := json.Unmarshal(resp.Body, &deployResp); err != nil {
return nil, fmt.Errorf("failed to parse deployment response: %w", err)
}
return &deployResp, nil
}
// ListDeployments retrieves a list of deployments with optional filtering
func (c *Client) ListDeployments(opts *ListDeploymentsOptions) (*DeploymentList, error) {
params := make(map[string]interface{})
if opts != nil {
params["status"] = opts.Status
params["location_id"] = opts.LocationID
params["page"] = opts.Page
params["page_size"] = opts.PageSize
params["sort_by"] = opts.SortBy
params["sort_order"] = opts.SortOrder
}
endpoint := "/deployments" + buildQueryParams(params)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to list deployments: %w", err)
}
var deploymentList DeploymentList
if err := decodeData(resp.Body, &deploymentList); err != nil {
return nil, fmt.Errorf("failed to parse deployments list: %w", err)
}
deploymentList.Deployments = lo.Map(deploymentList.Deployments, func(deployment Deployment, _ int) Deployment {
deployment.GPUCount = deployment.HardwareQuantity
deployment.Replicas = deployment.HardwareQuantity // Assuming 1:1 mapping for now
return deployment
})
return &deploymentList, nil
}
// GetDeployment retrieves detailed information about a specific deployment
func (c *Client) GetDeployment(deploymentID string) (*DeploymentDetail, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s", deploymentID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get deployment details: %w", err)
}
var deploymentDetail DeploymentDetail
if err := decodeDataWithFlexibleTimes(resp.Body, &deploymentDetail); err != nil {
return nil, fmt.Errorf("failed to parse deployment details: %w", err)
}
return &deploymentDetail, nil
}
// UpdateDeployment updates the configuration of an existing deployment
func (c *Client) UpdateDeployment(deploymentID string, req *UpdateDeploymentRequest) (*UpdateDeploymentResponse, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
if req == nil {
return nil, fmt.Errorf("update request cannot be nil")
}
endpoint := fmt.Sprintf("/deployment/%s", deploymentID)
resp, err := c.makeRequest("PATCH", endpoint, req)
if err != nil {
return nil, fmt.Errorf("failed to update deployment: %w", err)
}
// API returns direct format:
// {"status": "string", "deployment_id": "..."}
var updateResp UpdateDeploymentResponse
if err := json.Unmarshal(resp.Body, &updateResp); err != nil {
return nil, fmt.Errorf("failed to parse update deployment response: %w", err)
}
return &updateResp, nil
}
// ExtendDeployment extends the duration of an existing deployment
func (c *Client) ExtendDeployment(deploymentID string, req *ExtendDurationRequest) (*DeploymentDetail, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
if req == nil {
return nil, fmt.Errorf("extend request cannot be nil")
}
if req.DurationHours < 1 {
return nil, fmt.Errorf("duration_hours must be at least 1")
}
endpoint := fmt.Sprintf("/deployment/%s/extend", deploymentID)
resp, err := c.makeRequest("POST", endpoint, req)
if err != nil {
return nil, fmt.Errorf("failed to extend deployment: %w", err)
}
var deploymentDetail DeploymentDetail
if err := decodeDataWithFlexibleTimes(resp.Body, &deploymentDetail); err != nil {
return nil, fmt.Errorf("failed to parse extended deployment details: %w", err)
}
return &deploymentDetail, nil
}
// DeleteDeployment deletes an active deployment
func (c *Client) DeleteDeployment(deploymentID string) (*UpdateDeploymentResponse, error) {
if deploymentID == "" {
return nil, fmt.Errorf("deployment ID cannot be empty")
}
endpoint := fmt.Sprintf("/deployment/%s", deploymentID)
resp, err := c.makeRequest("DELETE", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to delete deployment: %w", err)
}
// API returns direct format:
// {"status": "string", "deployment_id": "..."}
var deleteResp UpdateDeploymentResponse
if err := json.Unmarshal(resp.Body, &deleteResp); err != nil {
return nil, fmt.Errorf("failed to parse delete deployment response: %w", err)
}
return &deleteResp, nil
}
// GetPriceEstimation calculates the estimated cost for a deployment
func (c *Client) GetPriceEstimation(req *PriceEstimationRequest) (*PriceEstimationResponse, error) {
if req == nil {
return nil, fmt.Errorf("price estimation request cannot be nil")
}
// Validate required fields
if len(req.LocationIDs) == 0 {
return nil, fmt.Errorf("location_ids is required")
}
if req.HardwareID == 0 {
return nil, fmt.Errorf("hardware_id is required")
}
if req.ReplicaCount < 1 {
return nil, fmt.Errorf("replica_count must be at least 1")
}
currency := strings.TrimSpace(req.Currency)
if currency == "" {
currency = "usdc"
}
durationType := strings.TrimSpace(req.DurationType)
if durationType == "" {
durationType = "hour"
}
durationType = strings.ToLower(durationType)
apiDurationType := ""
durationQty := req.DurationQty
if durationQty < 1 {
durationQty = req.DurationHours
}
if durationQty < 1 {
return nil, fmt.Errorf("duration_qty must be at least 1")
}
hardwareQty := req.HardwareQty
if hardwareQty < 1 {
hardwareQty = req.GPUsPerContainer
}
if hardwareQty < 1 {
return nil, fmt.Errorf("hardware_qty must be at least 1")
}
durationHoursForRate := req.DurationHours
if durationHoursForRate < 1 {
durationHoursForRate = durationQty
}
switch durationType {
case "hour", "hours", "hourly":
durationHoursForRate = durationQty
apiDurationType = "hourly"
case "day", "days", "daily":
durationHoursForRate = durationQty * 24
apiDurationType = "daily"
case "week", "weeks", "weekly":
durationHoursForRate = durationQty * 24 * 7
apiDurationType = "weekly"
case "month", "months", "monthly":
durationHoursForRate = durationQty * 24 * 30
apiDurationType = "monthly"
}
if durationHoursForRate < 1 {
durationHoursForRate = 1
}
if apiDurationType == "" {
apiDurationType = "hourly"
}
params := map[string]interface{}{
"location_ids": req.LocationIDs,
"hardware_id": req.HardwareID,
"hardware_qty": hardwareQty,
"gpus_per_container": req.GPUsPerContainer,
"duration_type": apiDurationType,
"duration_qty": durationQty,
"duration_hours": req.DurationHours,
"replica_count": req.ReplicaCount,
"currency": currency,
}
endpoint := "/price" + buildQueryParams(params)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get price estimation: %w", err)
}
// Parse according to the actual API response format from docs:
// {
// "data": {
// "replica_count": 0,
// "gpus_per_container": 0,
// "available_replica_count": [0],
// "discount": 0,
// "ionet_fee": 0,
// "ionet_fee_percent": 0,
// "currency_conversion_fee": 0,
// "currency_conversion_fee_percent": 0,
// "total_cost_usdc": 0
// }
// }
var pricingData struct {
ReplicaCount int `json:"replica_count"`
GPUsPerContainer int `json:"gpus_per_container"`
AvailableReplicaCount []int `json:"available_replica_count"`
Discount float64 `json:"discount"`
IonetFee float64 `json:"ionet_fee"`
IonetFeePercent float64 `json:"ionet_fee_percent"`
CurrencyConversionFee float64 `json:"currency_conversion_fee"`
CurrencyConversionFeePercent float64 `json:"currency_conversion_fee_percent"`
TotalCostUSDC float64 `json:"total_cost_usdc"`
}
if err := decodeData(resp.Body, &pricingData); err != nil {
return nil, fmt.Errorf("failed to parse price estimation response: %w", err)
}
// Convert to our internal format
durationHoursFloat := float64(durationHoursForRate)
if durationHoursFloat <= 0 {
durationHoursFloat = 1
}
priceResp := &PriceEstimationResponse{
EstimatedCost: pricingData.TotalCostUSDC,
Currency: strings.ToUpper(currency),
EstimationValid: true,
PriceBreakdown: PriceBreakdown{
ComputeCost: pricingData.TotalCostUSDC - pricingData.IonetFee - pricingData.CurrencyConversionFee,
TotalCost: pricingData.TotalCostUSDC,
HourlyRate: pricingData.TotalCostUSDC / durationHoursFloat,
},
}
return priceResp, nil
}
// CheckClusterNameAvailability checks if a cluster name is available
func (c *Client) CheckClusterNameAvailability(clusterName string) (bool, error) {
if clusterName == "" {
return false, fmt.Errorf("cluster name cannot be empty")
}
params := map[string]interface{}{
"cluster_name": clusterName,
}
endpoint := "/clusters/check_cluster_name_availability" + buildQueryParams(params)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return false, fmt.Errorf("failed to check cluster name availability: %w", err)
}
var availabilityResp bool
if err := json.Unmarshal(resp.Body, &availabilityResp); err != nil {
return false, fmt.Errorf("failed to parse cluster name availability response: %w", err)
}
return availabilityResp, nil
}
// UpdateClusterName updates the name of an existing cluster/deployment
func (c *Client) UpdateClusterName(clusterID string, req *UpdateClusterNameRequest) (*UpdateClusterNameResponse, error) {
if clusterID == "" {
return nil, fmt.Errorf("cluster ID cannot be empty")
}
if req == nil {
return nil, fmt.Errorf("update cluster name request cannot be nil")
}
if req.Name == "" {
return nil, fmt.Errorf("cluster name cannot be empty")
}
endpoint := fmt.Sprintf("/clusters/%s/update-name", clusterID)
resp, err := c.makeRequest("PUT", endpoint, req)
if err != nil {
return nil, fmt.Errorf("failed to update cluster name: %w", err)
}
// Parse the response directly without data wrapper based on API docs
var updateResp UpdateClusterNameResponse
if err := json.Unmarshal(resp.Body, &updateResp); err != nil {
return nil, fmt.Errorf("failed to parse update cluster name response: %w", err)
}
return &updateResp, nil
}

202
pkg/ionet/hardware.go Normal file
View File

@@ -0,0 +1,202 @@
package ionet
import (
"encoding/json"
"fmt"
"strings"
"github.com/samber/lo"
)
// GetAvailableReplicas retrieves available replicas per location for specified hardware
func (c *Client) GetAvailableReplicas(hardwareID int, gpuCount int) (*AvailableReplicasResponse, error) {
if hardwareID <= 0 {
return nil, fmt.Errorf("hardware_id must be greater than 0")
}
if gpuCount < 1 {
return nil, fmt.Errorf("gpu_count must be at least 1")
}
params := map[string]interface{}{
"hardware_id": hardwareID,
"hardware_qty": gpuCount,
}
endpoint := "/available-replicas" + buildQueryParams(params)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get available replicas: %w", err)
}
type availableReplicaPayload struct {
ID int `json:"id"`
ISO2 string `json:"iso2"`
Name string `json:"name"`
AvailableReplicas int `json:"available_replicas"`
}
var payload []availableReplicaPayload
if err := decodeData(resp.Body, &payload); err != nil {
return nil, fmt.Errorf("failed to parse available replicas response: %w", err)
}
replicas := lo.Map(payload, func(item availableReplicaPayload, _ int) AvailableReplica {
return AvailableReplica{
LocationID: item.ID,
LocationName: item.Name,
HardwareID: hardwareID,
HardwareName: "",
AvailableCount: item.AvailableReplicas,
MaxGPUs: gpuCount,
}
})
return &AvailableReplicasResponse{Replicas: replicas}, nil
}
// GetMaxGPUsPerContainer retrieves the maximum number of GPUs available per hardware type
func (c *Client) GetMaxGPUsPerContainer() (*MaxGPUResponse, error) {
resp, err := c.makeRequest("GET", "/hardware/max-gpus-per-container", nil)
if err != nil {
return nil, fmt.Errorf("failed to get max GPUs per container: %w", err)
}
var maxGPUResp MaxGPUResponse
if err := decodeData(resp.Body, &maxGPUResp); err != nil {
return nil, fmt.Errorf("failed to parse max GPU response: %w", err)
}
return &maxGPUResp, nil
}
// ListHardwareTypes retrieves available hardware types using the max GPUs endpoint
func (c *Client) ListHardwareTypes() ([]HardwareType, int, error) {
maxGPUResp, err := c.GetMaxGPUsPerContainer()
if err != nil {
return nil, 0, fmt.Errorf("failed to list hardware types: %w", err)
}
mapped := lo.Map(maxGPUResp.Hardware, func(hw MaxGPUInfo, _ int) HardwareType {
name := strings.TrimSpace(hw.HardwareName)
if name == "" {
name = fmt.Sprintf("Hardware %d", hw.HardwareID)
}
return HardwareType{
ID: hw.HardwareID,
Name: name,
GPUType: "",
GPUMemory: 0,
MaxGPUs: hw.MaxGPUsPerContainer,
CPU: "",
Memory: 0,
Storage: 0,
HourlyRate: 0,
Available: hw.Available > 0,
BrandName: strings.TrimSpace(hw.BrandName),
AvailableCount: hw.Available,
}
})
totalAvailable := maxGPUResp.Total
if totalAvailable == 0 {
totalAvailable = lo.SumBy(maxGPUResp.Hardware, func(hw MaxGPUInfo) int {
return hw.Available
})
}
return mapped, totalAvailable, nil
}
// ListLocations retrieves available deployment locations (if supported by the API)
func (c *Client) ListLocations() (*LocationsResponse, error) {
resp, err := c.makeRequest("GET", "/locations", nil)
if err != nil {
return nil, fmt.Errorf("failed to list locations: %w", err)
}
var locations LocationsResponse
if err := decodeData(resp.Body, &locations); err != nil {
return nil, fmt.Errorf("failed to parse locations response: %w", err)
}
locations.Locations = lo.Map(locations.Locations, func(location Location, _ int) Location {
location.ISO2 = strings.ToUpper(strings.TrimSpace(location.ISO2))
return location
})
if locations.Total == 0 {
locations.Total = lo.SumBy(locations.Locations, func(location Location) int {
return location.Available
})
}
return &locations, nil
}
// GetHardwareType retrieves details about a specific hardware type
func (c *Client) GetHardwareType(hardwareID int) (*HardwareType, error) {
if hardwareID <= 0 {
return nil, fmt.Errorf("hardware ID must be greater than 0")
}
endpoint := fmt.Sprintf("/hardware/types/%d", hardwareID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get hardware type: %w", err)
}
// API response format not documented, assuming direct format
var hardwareType HardwareType
if err := json.Unmarshal(resp.Body, &hardwareType); err != nil {
return nil, fmt.Errorf("failed to parse hardware type: %w", err)
}
return &hardwareType, nil
}
// GetLocation retrieves details about a specific location
func (c *Client) GetLocation(locationID int) (*Location, error) {
if locationID <= 0 {
return nil, fmt.Errorf("location ID must be greater than 0")
}
endpoint := fmt.Sprintf("/locations/%d", locationID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get location: %w", err)
}
// API response format not documented, assuming direct format
var location Location
if err := json.Unmarshal(resp.Body, &location); err != nil {
return nil, fmt.Errorf("failed to parse location: %w", err)
}
return &location, nil
}
// GetLocationAvailability retrieves real-time availability for a specific location
func (c *Client) GetLocationAvailability(locationID int) (*LocationAvailability, error) {
if locationID <= 0 {
return nil, fmt.Errorf("location ID must be greater than 0")
}
endpoint := fmt.Sprintf("/locations/%d/availability", locationID)
resp, err := c.makeRequest("GET", endpoint, nil)
if err != nil {
return nil, fmt.Errorf("failed to get location availability: %w", err)
}
// API response format not documented, assuming direct format
var availability LocationAvailability
if err := json.Unmarshal(resp.Body, &availability); err != nil {
return nil, fmt.Errorf("failed to parse location availability: %w", err)
}
return &availability, nil
}

96
pkg/ionet/jsonutil.go Normal file
View File

@@ -0,0 +1,96 @@
package ionet
import (
"encoding/json"
"strings"
"time"
"github.com/samber/lo"
)
// decodeWithFlexibleTimes unmarshals API responses while tolerating timestamp strings
// that omit timezone information by normalizing them to RFC3339Nano.
func decodeWithFlexibleTimes(data []byte, target interface{}) error {
var intermediate interface{}
if err := json.Unmarshal(data, &intermediate); err != nil {
return err
}
normalized := normalizeTimeValues(intermediate)
reencoded, err := json.Marshal(normalized)
if err != nil {
return err
}
return json.Unmarshal(reencoded, target)
}
func decodeData[T any](data []byte, target *T) error {
var wrapper struct {
Data T `json:"data"`
}
if err := json.Unmarshal(data, &wrapper); err != nil {
return err
}
*target = wrapper.Data
return nil
}
func decodeDataWithFlexibleTimes[T any](data []byte, target *T) error {
var wrapper struct {
Data T `json:"data"`
}
if err := decodeWithFlexibleTimes(data, &wrapper); err != nil {
return err
}
*target = wrapper.Data
return nil
}
func normalizeTimeValues(value interface{}) interface{} {
switch v := value.(type) {
case map[string]interface{}:
return lo.MapValues(v, func(val interface{}, _ string) interface{} {
return normalizeTimeValues(val)
})
case []interface{}:
return lo.Map(v, func(item interface{}, _ int) interface{} {
return normalizeTimeValues(item)
})
case string:
if normalized, changed := normalizeTimeString(v); changed {
return normalized
}
return v
default:
return value
}
}
func normalizeTimeString(input string) (string, bool) {
trimmed := strings.TrimSpace(input)
if trimmed == "" {
return input, false
}
if _, err := time.Parse(time.RFC3339Nano, trimmed); err == nil {
return trimmed, trimmed != input
}
if _, err := time.Parse(time.RFC3339, trimmed); err == nil {
return trimmed, trimmed != input
}
layouts := []string{
"2006-01-02T15:04:05.999999999",
"2006-01-02T15:04:05.999999",
"2006-01-02T15:04:05",
}
for _, layout := range layouts {
if parsed, err := time.Parse(layout, trimmed); err == nil {
return parsed.UTC().Format(time.RFC3339Nano), true
}
}
return input, false
}

353
pkg/ionet/types.go Normal file
View File

@@ -0,0 +1,353 @@
package ionet
import (
"time"
)
// Client represents the IO.NET API client
type Client struct {
BaseURL string
APIKey string
HTTPClient HTTPClient
}
// HTTPClient interface for making HTTP requests
type HTTPClient interface {
Do(req *HTTPRequest) (*HTTPResponse, error)
}
// HTTPRequest represents an HTTP request
type HTTPRequest struct {
Method string
URL string
Headers map[string]string
Body []byte
}
// HTTPResponse represents an HTTP response
type HTTPResponse struct {
StatusCode int
Headers map[string]string
Body []byte
}
// DeploymentRequest represents a container deployment request
type DeploymentRequest struct {
ResourcePrivateName string `json:"resource_private_name"`
DurationHours int `json:"duration_hours"`
GPUsPerContainer int `json:"gpus_per_container"`
HardwareID int `json:"hardware_id"`
LocationIDs []int `json:"location_ids"`
ContainerConfig ContainerConfig `json:"container_config"`
RegistryConfig RegistryConfig `json:"registry_config"`
}
// ContainerConfig represents container configuration
type ContainerConfig struct {
ReplicaCount int `json:"replica_count"`
EnvVariables map[string]string `json:"env_variables,omitempty"`
SecretEnvVariables map[string]string `json:"secret_env_variables,omitempty"`
Entrypoint []string `json:"entrypoint,omitempty"`
TrafficPort int `json:"traffic_port,omitempty"`
Args []string `json:"args,omitempty"`
}
// RegistryConfig represents registry configuration
type RegistryConfig struct {
ImageURL string `json:"image_url"`
RegistryUsername string `json:"registry_username,omitempty"`
RegistrySecret string `json:"registry_secret,omitempty"`
}
// DeploymentResponse represents the response from deployment creation
type DeploymentResponse struct {
DeploymentID string `json:"deployment_id"`
Status string `json:"status"`
}
// DeploymentDetail represents detailed deployment information
type DeploymentDetail struct {
ID string `json:"id"`
Status string `json:"status"`
CreatedAt time.Time `json:"created_at"`
StartedAt *time.Time `json:"started_at,omitempty"`
FinishedAt *time.Time `json:"finished_at,omitempty"`
AmountPaid float64 `json:"amount_paid"`
CompletedPercent float64 `json:"completed_percent"`
TotalGPUs int `json:"total_gpus"`
GPUsPerContainer int `json:"gpus_per_container"`
TotalContainers int `json:"total_containers"`
HardwareName string `json:"hardware_name"`
HardwareID int `json:"hardware_id"`
Locations []DeploymentLocation `json:"locations"`
BrandName string `json:"brand_name"`
ComputeMinutesServed int `json:"compute_minutes_served"`
ComputeMinutesRemaining int `json:"compute_minutes_remaining"`
ContainerConfig DeploymentContainerConfig `json:"container_config"`
}
// DeploymentLocation represents a location in deployment details
type DeploymentLocation struct {
ID int `json:"id"`
ISO2 string `json:"iso2"`
Name string `json:"name"`
}
// DeploymentContainerConfig represents container config in deployment details
type DeploymentContainerConfig struct {
Entrypoint []string `json:"entrypoint"`
EnvVariables map[string]interface{} `json:"env_variables"`
TrafficPort int `json:"traffic_port"`
ImageURL string `json:"image_url"`
}
// Container represents a container within a deployment
type Container struct {
DeviceID string `json:"device_id"`
ContainerID string `json:"container_id"`
Hardware string `json:"hardware"`
BrandName string `json:"brand_name"`
CreatedAt time.Time `json:"created_at"`
UptimePercent int `json:"uptime_percent"`
GPUsPerContainer int `json:"gpus_per_container"`
Status string `json:"status"`
ContainerEvents []ContainerEvent `json:"container_events"`
PublicURL string `json:"public_url"`
}
// ContainerEvent represents a container event
type ContainerEvent struct {
Time time.Time `json:"time"`
Message string `json:"message"`
}
// ContainerList represents a list of containers
type ContainerList struct {
Total int `json:"total"`
Workers []Container `json:"workers"`
}
// Deployment represents a deployment in the list
type Deployment struct {
ID string `json:"id"`
Status string `json:"status"`
Name string `json:"name"`
CompletedPercent float64 `json:"completed_percent"`
HardwareQuantity int `json:"hardware_quantity"`
BrandName string `json:"brand_name"`
HardwareName string `json:"hardware_name"`
Served string `json:"served"`
Remaining string `json:"remaining"`
ComputeMinutesServed int `json:"compute_minutes_served"`
ComputeMinutesRemaining int `json:"compute_minutes_remaining"`
CreatedAt time.Time `json:"created_at"`
GPUCount int `json:"-"` // Derived from HardwareQuantity
Replicas int `json:"-"` // Derived from HardwareQuantity
}
// DeploymentList represents a list of deployments with pagination
type DeploymentList struct {
Deployments []Deployment `json:"deployments"`
Total int `json:"total"`
Statuses []string `json:"statuses"`
}
// AvailableReplica represents replica availability for a location
type AvailableReplica struct {
LocationID int `json:"location_id"`
LocationName string `json:"location_name"`
HardwareID int `json:"hardware_id"`
HardwareName string `json:"hardware_name"`
AvailableCount int `json:"available_count"`
MaxGPUs int `json:"max_gpus"`
}
// AvailableReplicasResponse represents the response for available replicas
type AvailableReplicasResponse struct {
Replicas []AvailableReplica `json:"replicas"`
}
// MaxGPUResponse represents the response for maximum GPUs per container
type MaxGPUResponse struct {
Hardware []MaxGPUInfo `json:"hardware"`
Total int `json:"total"`
}
// MaxGPUInfo represents max GPU information for a hardware type
type MaxGPUInfo struct {
MaxGPUsPerContainer int `json:"max_gpus_per_container"`
Available int `json:"available"`
HardwareID int `json:"hardware_id"`
HardwareName string `json:"hardware_name"`
BrandName string `json:"brand_name"`
}
// PriceEstimationRequest represents a price estimation request
type PriceEstimationRequest struct {
LocationIDs []int `json:"location_ids"`
HardwareID int `json:"hardware_id"`
GPUsPerContainer int `json:"gpus_per_container"`
DurationHours int `json:"duration_hours"`
ReplicaCount int `json:"replica_count"`
Currency string `json:"currency"`
DurationType string `json:"duration_type"`
DurationQty int `json:"duration_qty"`
HardwareQty int `json:"hardware_qty"`
}
// PriceEstimationResponse represents the price estimation response
type PriceEstimationResponse struct {
EstimatedCost float64 `json:"estimated_cost"`
Currency string `json:"currency"`
PriceBreakdown PriceBreakdown `json:"price_breakdown"`
EstimationValid bool `json:"estimation_valid"`
}
// PriceBreakdown represents detailed cost breakdown
type PriceBreakdown struct {
ComputeCost float64 `json:"compute_cost"`
NetworkCost float64 `json:"network_cost,omitempty"`
StorageCost float64 `json:"storage_cost,omitempty"`
TotalCost float64 `json:"total_cost"`
HourlyRate float64 `json:"hourly_rate"`
}
// ContainerLogs represents container log entries
type ContainerLogs struct {
ContainerID string `json:"container_id"`
Logs []LogEntry `json:"logs"`
HasMore bool `json:"has_more"`
NextCursor string `json:"next_cursor,omitempty"`
}
// LogEntry represents a single log entry
type LogEntry struct {
Timestamp time.Time `json:"timestamp"`
Level string `json:"level,omitempty"`
Message string `json:"message"`
Source string `json:"source,omitempty"`
}
// UpdateDeploymentRequest represents request to update deployment configuration
type UpdateDeploymentRequest struct {
EnvVariables map[string]string `json:"env_variables,omitempty"`
SecretEnvVariables map[string]string `json:"secret_env_variables,omitempty"`
Entrypoint []string `json:"entrypoint,omitempty"`
TrafficPort *int `json:"traffic_port,omitempty"`
ImageURL string `json:"image_url,omitempty"`
RegistryUsername string `json:"registry_username,omitempty"`
RegistrySecret string `json:"registry_secret,omitempty"`
Args []string `json:"args,omitempty"`
Command string `json:"command,omitempty"`
}
// ExtendDurationRequest represents request to extend deployment duration
type ExtendDurationRequest struct {
DurationHours int `json:"duration_hours"`
}
// UpdateDeploymentResponse represents response from deployment update
type UpdateDeploymentResponse struct {
Status string `json:"status"`
DeploymentID string `json:"deployment_id"`
}
// UpdateClusterNameRequest represents request to update cluster name
type UpdateClusterNameRequest struct {
Name string `json:"cluster_name"`
}
// UpdateClusterNameResponse represents response from cluster name update
type UpdateClusterNameResponse struct {
Status string `json:"status"`
Message string `json:"message"`
}
// APIError represents an API error response
type APIError struct {
Code int `json:"code"`
Message string `json:"message"`
Details string `json:"details,omitempty"`
}
// Error implements the error interface
func (e *APIError) Error() string {
if e.Details != "" {
return e.Message + ": " + e.Details
}
return e.Message
}
// ListDeploymentsOptions represents options for listing deployments
type ListDeploymentsOptions struct {
Status string `json:"status,omitempty"` // filter by status
LocationID int `json:"location_id,omitempty"` // filter by location
Page int `json:"page,omitempty"` // pagination
PageSize int `json:"page_size,omitempty"` // pagination
SortBy string `json:"sort_by,omitempty"` // sort field
SortOrder string `json:"sort_order,omitempty"` // asc/desc
}
// GetLogsOptions represents options for retrieving container logs
type GetLogsOptions struct {
StartTime *time.Time `json:"start_time,omitempty"`
EndTime *time.Time `json:"end_time,omitempty"`
Level string `json:"level,omitempty"` // filter by log level
Stream string `json:"stream,omitempty"` // filter by stdout/stderr streams
Limit int `json:"limit,omitempty"` // max number of log entries
Cursor string `json:"cursor,omitempty"` // pagination cursor
Follow bool `json:"follow,omitempty"` // stream logs
}
// HardwareType represents a hardware type available for deployment
type HardwareType struct {
ID int `json:"id"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
GPUType string `json:"gpu_type"`
GPUMemory int `json:"gpu_memory"` // in GB
MaxGPUs int `json:"max_gpus"`
CPU string `json:"cpu,omitempty"`
Memory int `json:"memory,omitempty"` // in GB
Storage int `json:"storage,omitempty"` // in GB
HourlyRate float64 `json:"hourly_rate"`
Available bool `json:"available"`
BrandName string `json:"brand_name,omitempty"`
AvailableCount int `json:"available_count,omitempty"`
}
// Location represents a deployment location
type Location struct {
ID int `json:"id"`
Name string `json:"name"`
ISO2 string `json:"iso2,omitempty"`
Region string `json:"region,omitempty"`
Country string `json:"country,omitempty"`
Latitude float64 `json:"latitude,omitempty"`
Longitude float64 `json:"longitude,omitempty"`
Available int `json:"available,omitempty"`
Description string `json:"description,omitempty"`
}
// LocationsResponse represents the list of locations and aggregated metadata.
type LocationsResponse struct {
Locations []Location `json:"locations"`
Total int `json:"total"`
}
// LocationAvailability represents real-time availability for a location
type LocationAvailability struct {
LocationID int `json:"location_id"`
LocationName string `json:"location_name"`
Available bool `json:"available"`
HardwareAvailability []HardwareAvailability `json:"hardware_availability"`
UpdatedAt time.Time `json:"updated_at"`
}
// HardwareAvailability represents availability for specific hardware at a location
type HardwareAvailability struct {
HardwareID int `json:"hardware_id"`
HardwareName string `json:"hardware_name"`
AvailableCount int `json:"available_count"`
MaxGPUs int `json:"max_gpus"`
}

View File

@@ -70,7 +70,7 @@ func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
if usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0 {
service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "")
} else {
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
}
return nil

View File

@@ -13,12 +13,37 @@ import (
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
IsSyncImageModel bool
}
/*
var syncModels = []string{
"z-image",
"qwen-image",
"wan2.6",
}
*/
func supportsAliAnthropicMessages(modelName string) bool {
// Only models with the "qwen" designation can use the Claude-compatible interface; others require conversion.
return strings.Contains(strings.ToLower(modelName), "qwen")
}
var syncModels = []string{
"z-image",
"qwen-image",
"wan2.6",
}
func isSyncImageModel(modelName string) bool {
return model_setting.IsSyncImageModel(modelName)
}
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
@@ -27,7 +52,18 @@ func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dt
}
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.ClaudeRequest) (any, error) {
return req, nil
if supportsAliAnthropicMessages(info.UpstreamModelName) {
return req, nil
}
oaiReq, err := service.ClaudeToOpenAIRequest(*req, info)
if err != nil {
return nil, err
}
if info.SupportStreamOptions && info.IsStream {
oaiReq.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
}
return a.ConvertOpenAIRequest(c, info, oaiReq)
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
@@ -37,7 +73,11 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
var fullRequestURL string
switch info.RelayFormat {
case types.RelayFormatClaude:
fullRequestURL = fmt.Sprintf("%s/api/v2/apps/claude-code-proxy/v1/messages", info.ChannelBaseUrl)
if supportsAliAnthropicMessages(info.UpstreamModelName) {
fullRequestURL = fmt.Sprintf("%s/apps/anthropic/v1/messages", info.ChannelBaseUrl)
} else {
fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/chat/completions", info.ChannelBaseUrl)
}
default:
switch info.RelayMode {
case constant.RelayModeEmbeddings:
@@ -45,10 +85,16 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
case constant.RelayModeRerank:
fullRequestURL = fmt.Sprintf("%s/api/v1/services/rerank/text-rerank/text-rerank", info.ChannelBaseUrl)
case constant.RelayModeImagesGenerations:
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/text2image/image-synthesis", info.ChannelBaseUrl)
if isSyncImageModel(info.OriginModelName) {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/multimodal-generation/generation", info.ChannelBaseUrl)
} else {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/text2image/image-synthesis", info.ChannelBaseUrl)
}
case constant.RelayModeImagesEdits:
if isWanModel(info.OriginModelName) {
if isOldWanModel(info.OriginModelName) {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/image2image/image-synthesis", info.ChannelBaseUrl)
} else if isWanModel(info.OriginModelName) {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/image-generation/generation", info.ChannelBaseUrl)
} else {
fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/multimodal-generation/generation", info.ChannelBaseUrl)
}
@@ -72,7 +118,11 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *rel
req.Set("X-DashScope-Plugin", c.GetString("plugin"))
}
if info.RelayMode == constant.RelayModeImagesGenerations {
req.Set("X-DashScope-Async", "enable")
if isSyncImageModel(info.OriginModelName) {
} else {
req.Set("X-DashScope-Async", "enable")
}
}
if info.RelayMode == constant.RelayModeImagesEdits {
if isWanModel(info.OriginModelName) {
@@ -108,15 +158,25 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
if info.RelayMode == constant.RelayModeImagesGenerations {
aliRequest, err := oaiImage2Ali(request)
if isSyncImageModel(info.OriginModelName) {
a.IsSyncImageModel = true
}
aliRequest, err := oaiImage2AliImageRequest(info, request, a.IsSyncImageModel)
if err != nil {
return nil, fmt.Errorf("convert image request failed: %w", err)
return nil, fmt.Errorf("convert image request to async ali image request failed: %w", err)
}
return aliRequest, nil
} else if info.RelayMode == constant.RelayModeImagesEdits {
if isWanModel(info.OriginModelName) {
if isOldWanModel(info.OriginModelName) {
return oaiFormEdit2WanxImageEdit(c, info, request)
}
if isSyncImageModel(info.OriginModelName) {
if isWanModel(info.OriginModelName) {
a.IsSyncImageModel = false
} else {
a.IsSyncImageModel = true
}
}
// ali image edit https://bailian.console.aliyun.com/?tab=api#/api/?type=model&url=2976416
// 如果用户使用表单,则需要解析表单数据
if strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
@@ -126,9 +186,9 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
}
return aliRequest, nil
} else {
aliRequest, err := oaiImage2Ali(request)
aliRequest, err := oaiImage2AliImageRequest(info, request, a.IsSyncImageModel)
if err != nil {
return nil, fmt.Errorf("convert image request failed: %w", err)
return nil, fmt.Errorf("convert image request to async ali image request failed: %w", err)
}
return aliRequest, nil
}
@@ -150,7 +210,7 @@ func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInf
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
// TODO implement me
//TODO implement me
return nil, errors.New("not implemented")
}
@@ -161,21 +221,22 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayFormat {
case types.RelayFormatClaude:
if info.IsStream {
return claude.ClaudeStreamHandler(c, resp, info, claude.RequestModeMessage)
} else {
if supportsAliAnthropicMessages(info.UpstreamModelName) {
if info.IsStream {
return claude.ClaudeStreamHandler(c, resp, info, claude.RequestModeMessage)
}
return claude.ClaudeHandler(c, resp, info, claude.RequestModeMessage)
}
adaptor := openai.Adaptor{}
return adaptor.DoResponse(c, resp, info)
default:
switch info.RelayMode {
case constant.RelayModeImagesGenerations:
err, usage = aliImageHandler(c, resp, info)
err, usage = aliImageHandler(a, c, resp, info)
case constant.RelayModeImagesEdits:
if isWanModel(info.OriginModelName) {
err, usage = aliImageHandler(c, resp, info)
} else {
err, usage = aliImageEditHandler(c, resp, info)
}
err, usage = aliImageHandler(a, c, resp, info)
case constant.RelayModeRerank:
err, usage = RerankHandler(c, resp, info)
default:

View File

@@ -1,6 +1,13 @@
package ali
import "github.com/QuantumNous/new-api/dto"
import (
"strings"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
type AliMessage struct {
Content any `json:"content"`
@@ -65,6 +72,7 @@ type AliUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
ImageCount int `json:"image_count,omitempty"`
}
type TaskResult struct {
@@ -75,14 +83,78 @@ type TaskResult struct {
}
type AliOutput struct {
TaskId string `json:"task_id,omitempty"`
TaskStatus string `json:"task_status,omitempty"`
Text string `json:"text"`
FinishReason string `json:"finish_reason"`
Message string `json:"message,omitempty"`
Code string `json:"code,omitempty"`
Results []TaskResult `json:"results,omitempty"`
Choices []map[string]any `json:"choices,omitempty"`
TaskId string `json:"task_id,omitempty"`
TaskStatus string `json:"task_status,omitempty"`
Text string `json:"text"`
FinishReason string `json:"finish_reason"`
Message string `json:"message,omitempty"`
Code string `json:"code,omitempty"`
Results []TaskResult `json:"results,omitempty"`
Choices []struct {
FinishReason string `json:"finish_reason,omitempty"`
Message struct {
Role string `json:"role,omitempty"`
Content []AliMediaContent `json:"content,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
} `json:"message,omitempty"`
} `json:"choices,omitempty"`
}
func (o *AliOutput) ChoicesToOpenAIImageDate(c *gin.Context, responseFormat string) []dto.ImageData {
var imageData []dto.ImageData
if len(o.Choices) > 0 {
for _, choice := range o.Choices {
var data dto.ImageData
for _, content := range choice.Message.Content {
if content.Image != "" {
if strings.HasPrefix(content.Image, "http") {
var b64Json string
if responseFormat == "b64_json" {
_, b64, err := service.GetImageFromUrl(content.Image)
if err != nil {
logger.LogError(c, "get_image_data_failed: "+err.Error())
continue
}
b64Json = b64
}
data.Url = content.Image
data.B64Json = b64Json
} else {
data.B64Json = content.Image
}
} else if content.Text != "" {
data.RevisedPrompt = content.Text
}
}
imageData = append(imageData, data)
}
}
return imageData
}
func (o *AliOutput) ResultToOpenAIImageDate(c *gin.Context, responseFormat string) []dto.ImageData {
var imageData []dto.ImageData
for _, data := range o.Results {
var b64Json string
if responseFormat == "b64_json" {
_, b64, err := service.GetImageFromUrl(data.Url)
if err != nil {
logger.LogError(c, "get_image_data_failed: "+err.Error())
continue
}
b64Json = b64
} else {
b64Json = data.B64Image
}
imageData = append(imageData, dto.ImageData{
Url: data.Url,
B64Json: b64Json,
RevisedPrompt: "",
})
}
return imageData
}
type AliResponse struct {
@@ -92,18 +164,26 @@ type AliResponse struct {
}
type AliImageRequest struct {
Model string `json:"model"`
Input any `json:"input"`
Parameters any `json:"parameters,omitempty"`
ResponseFormat string `json:"response_format,omitempty"`
Model string `json:"model"`
Input any `json:"input"`
Parameters AliImageParameters `json:"parameters,omitempty"`
ResponseFormat string `json:"response_format,omitempty"`
}
type AliImageParameters struct {
Size string `json:"size,omitempty"`
N int `json:"n,omitempty"`
Steps string `json:"steps,omitempty"`
Scale string `json:"scale,omitempty"`
Watermark *bool `json:"watermark,omitempty"`
Size string `json:"size,omitempty"`
N int `json:"n,omitempty"`
Steps string `json:"steps,omitempty"`
Scale string `json:"scale,omitempty"`
Watermark *bool `json:"watermark,omitempty"`
PromptExtend *bool `json:"prompt_extend,omitempty"`
}
func (p *AliImageParameters) PromptExtendValue() bool {
if p != nil && p.PromptExtend != nil {
return *p.PromptExtend
}
return false
}
type AliImageInput struct {

View File

@@ -1,7 +1,6 @@
package ali
import (
"context"
"encoding/base64"
"errors"
"fmt"
@@ -21,17 +20,23 @@ import (
"github.com/gin-gonic/gin"
)
func oaiImage2Ali(request dto.ImageRequest) (*AliImageRequest, error) {
func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequest, isSync bool) (*AliImageRequest, error) {
var imageRequest AliImageRequest
imageRequest.Model = request.Model
imageRequest.ResponseFormat = request.ResponseFormat
logger.LogJson(context.Background(), "oaiImage2Ali request extra", request.Extra)
if request.Extra != nil {
if val, ok := request.Extra["parameters"]; ok {
err := common.Unmarshal(val, &imageRequest.Parameters)
if err != nil {
return nil, fmt.Errorf("invalid parameters field: %w", err)
}
} else {
// 兼容没有parameters字段的情况从openai标准字段中提取参数
imageRequest.Parameters = AliImageParameters{
Size: strings.Replace(request.Size, "x", "*", -1),
N: int(request.N),
Watermark: request.Watermark,
}
}
if val, ok := request.Extra["input"]; ok {
err := common.Unmarshal(val, &imageRequest.Input)
@@ -41,23 +46,44 @@ func oaiImage2Ali(request dto.ImageRequest) (*AliImageRequest, error) {
}
}
if imageRequest.Parameters == nil {
imageRequest.Parameters = AliImageParameters{
Size: strings.Replace(request.Size, "x", "*", -1),
N: int(request.N),
Watermark: request.Watermark,
if strings.Contains(request.Model, "z-image") {
// z-image 开启prompt_extend后按2倍计费
if imageRequest.Parameters.PromptExtendValue() {
info.PriceData.AddOtherRatio("prompt_extend", 2)
}
}
if imageRequest.Input == nil {
imageRequest.Input = AliImageInput{
Prompt: request.Prompt,
// 检查n参数
if imageRequest.Parameters.N != 0 {
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
}
// 同步图片模型和异步图片模型请求格式不一样
if isSync {
if imageRequest.Input == nil {
imageRequest.Input = AliImageInput{
Messages: []AliMessage{
{
Role: "user",
Content: []AliMediaContent{
{
Text: request.Prompt,
},
},
},
},
}
}
} else {
if imageRequest.Input == nil {
imageRequest.Input = AliImageInput{
Prompt: request.Prompt,
}
}
}
return &imageRequest, nil
}
func getImageBase64sFromForm(c *gin.Context, fieldName string) ([]string, error) {
mf := c.Request.MultipartForm
if mf == nil {
@@ -199,6 +225,8 @@ func asyncTaskWait(c *gin.Context, info *relaycommon.RelayInfo, taskID string) (
var taskResponse AliResponse
var responseBody []byte
time.Sleep(time.Duration(5) * time.Second)
for {
logger.LogDebug(c, fmt.Sprintf("asyncTaskWait step %d/%d, wait %d seconds", step, maxStep, waitSeconds))
step++
@@ -238,32 +266,17 @@ func responseAli2OpenAIImage(c *gin.Context, response *AliResponse, originBody [
Created: info.StartTime.Unix(),
}
for _, data := range response.Output.Results {
var b64Json string
if responseFormat == "b64_json" {
_, b64, err := service.GetImageFromUrl(data.Url)
if err != nil {
logger.LogError(c, "get_image_data_failed: "+err.Error())
continue
}
b64Json = b64
} else {
b64Json = data.B64Image
}
imageResponse.Data = append(imageResponse.Data, dto.ImageData{
Url: data.Url,
B64Json: b64Json,
RevisedPrompt: "",
})
if len(response.Output.Results) > 0 {
imageResponse.Data = response.Output.ResultToOpenAIImageDate(c, responseFormat)
} else if len(response.Output.Choices) > 0 {
imageResponse.Data = response.Output.ChoicesToOpenAIImageDate(c, responseFormat)
}
var mapResponse map[string]any
_ = common.Unmarshal(originBody, &mapResponse)
imageResponse.Extra = mapResponse
imageResponse.Metadata = originBody
return &imageResponse
}
func aliImageHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) {
func aliImageHandler(a *Adaptor, c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) {
responseFormat := c.GetString("response_format")
var aliTaskResponse AliResponse
@@ -282,66 +295,49 @@ func aliImageHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rela
return types.NewError(errors.New(aliTaskResponse.Message), types.ErrorCodeBadResponse), nil
}
aliResponse, originRespBody, err := asyncTaskWait(c, info, aliTaskResponse.Output.TaskId)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponse), nil
}
var (
aliResponse *AliResponse
originRespBody []byte
)
if aliResponse.Output.TaskStatus != "SUCCEEDED" {
return types.WithOpenAIError(types.OpenAIError{
Message: aliResponse.Output.Message,
Type: "ali_error",
Param: "",
Code: aliResponse.Output.Code,
}, resp.StatusCode), nil
}
fullTextResponse := responseAli2OpenAIImage(c, aliResponse, originRespBody, info, responseFormat)
jsonResponse, err := common.Marshal(fullTextResponse)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.IOCopyBytesGracefully(c, resp, jsonResponse)
return nil, &dto.Usage{}
}
func aliImageEditHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) {
var aliResponse AliResponse
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError), nil
}
service.CloseResponseBodyGracefully(resp)
err = common.Unmarshal(responseBody, &aliResponse)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError), nil
}
if aliResponse.Message != "" {
logger.LogError(c, "ali_task_failed: "+aliResponse.Message)
return types.NewError(errors.New(aliResponse.Message), types.ErrorCodeBadResponse), nil
}
var fullTextResponse dto.ImageResponse
if len(aliResponse.Output.Choices) > 0 {
fullTextResponse = dto.ImageResponse{
Created: info.StartTime.Unix(),
Data: []dto.ImageData{
{
Url: aliResponse.Output.Choices[0]["message"].(map[string]any)["content"].([]any)[0].(map[string]any)["image"].(string),
B64Json: "",
},
},
if a.IsSyncImageModel {
aliResponse = &aliTaskResponse
originRespBody = responseBody
} else {
// 异步图片模型需要轮询任务结果
aliResponse, originRespBody, err = asyncTaskWait(c, info, aliTaskResponse.Output.TaskId)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponse), nil
}
if aliResponse.Output.TaskStatus != "SUCCEEDED" {
return types.WithOpenAIError(types.OpenAIError{
Message: aliResponse.Output.Message,
Type: "ali_error",
Param: "",
Code: aliResponse.Output.Code,
}, resp.StatusCode), nil
}
}
var mapResponse map[string]any
_ = common.Unmarshal(responseBody, &mapResponse)
fullTextResponse.Extra = mapResponse
jsonResponse, err := common.Marshal(fullTextResponse)
//logger.LogDebug(c, "ali_async_task_result: "+string(originRespBody))
if a.IsSyncImageModel {
logger.LogDebug(c, "ali_sync_image_result: "+string(originRespBody))
} else {
logger.LogDebug(c, "ali_async_image_result: "+string(originRespBody))
}
imageResponses := responseAli2OpenAIImage(c, aliResponse, originRespBody, info, responseFormat)
// 可能生成多张图片修正计费数量n
if aliResponse.Usage.ImageCount != 0 {
info.PriceData.AddOtherRatio("n", float64(aliResponse.Usage.ImageCount))
} else if len(imageResponses.Data) != 0 {
info.PriceData.AddOtherRatio("n", float64(len(imageResponses.Data)))
}
jsonResponse, err := common.Marshal(imageResponses)
if err != nil {
return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
service.IOCopyBytesGracefully(c, resp, jsonResponse)
return nil, &dto.Usage{}
}

View File

@@ -26,14 +26,22 @@ func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, requ
if wanInput.Images, err = getImageBase64sFromForm(c, "image"); err != nil {
return nil, fmt.Errorf("get image base64s from form failed: %w", err)
}
wanParams := WanImageParameters{
//wanParams := WanImageParameters{
// N: int(request.N),
//}
imageRequest.Input = wanInput
imageRequest.Parameters = AliImageParameters{
N: int(request.N),
}
imageRequest.Input = wanInput
imageRequest.Parameters = wanParams
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
return &imageRequest, nil
}
func isOldWanModel(modelName string) bool {
return strings.Contains(modelName, "wan") && !strings.Contains(modelName, "wan2.6")
}
func isWanModel(modelName string) bool {
return strings.Contains(modelName, "wan")
}

View File

@@ -1,11 +1,13 @@
package aws
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
@@ -37,6 +39,13 @@ func getAwsErrorStatusCode(err error) int {
return http.StatusInternalServerError
}
func newAwsInvokeContext() (context.Context, context.CancelFunc) {
if common.RelayTimeout <= 0 {
return context.Background(), func() {}
}
return context.WithTimeout(context.Background(), time.Duration(common.RelayTimeout)*time.Second)
}
func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) {
var (
httpClient *http.Client
@@ -117,6 +126,7 @@ func doAwsClientRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor,
return nil, types.NewError(errors.Wrap(err, "marshal nova request"), types.ErrorCodeBadResponseBody)
}
awsReq.Body = reqBody
a.AwsReq = awsReq
return nil, nil
} else {
awsClaudeReq, err := formatRequest(requestBody, requestHeader)
@@ -201,7 +211,10 @@ func getAwsModelID(requestModel string) string {
func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
awsResp, err := a.AwsClient.InvokeModel(c.Request.Context(), a.AwsReq.(*bedrockruntime.InvokeModelInput))
ctx, cancel := newAwsInvokeContext()
defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil {
statusCode := getAwsErrorStatusCode(err)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil
@@ -228,7 +241,10 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types
}
func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(c.Request.Context(), a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
ctx, cancel := newAwsInvokeContext()
defer cancel()
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
if err != nil {
statusCode := getAwsErrorStatusCode(err)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModelWithResponseStream"), types.ErrorCodeAwsInvokeError, statusCode), nil
@@ -268,7 +284,10 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
// Nova模型处理函数
func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
awsResp, err := a.AwsClient.InvokeModel(c.Request.Context(), a.AwsReq.(*bedrockruntime.InvokeModelInput))
ctx, cancel := newAwsInvokeContext()
defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil {
statusCode := getAwsErrorStatusCode(err)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil

View File

@@ -0,0 +1,164 @@
package codex
import (
"encoding/json"
"errors"
"io"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
type Adaptor struct {
}
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
}
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
return nil, errors.New("codex channel: endpoint not supported")
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
if info != nil && info.ChannelSetting.SystemPrompt != "" {
systemPrompt := info.ChannelSetting.SystemPrompt
if len(request.Instructions) == 0 {
if b, err := common.Marshal(systemPrompt); err == nil {
request.Instructions = b
} else {
return nil, err
}
} else if info.ChannelSetting.SystemPromptOverride {
var existing string
if err := common.Unmarshal(request.Instructions, &existing); err == nil {
existing = strings.TrimSpace(existing)
if existing == "" {
if b, err := common.Marshal(systemPrompt); err == nil {
request.Instructions = b
} else {
return nil, err
}
} else {
if b, err := common.Marshal(systemPrompt + "\n" + existing); err == nil {
request.Instructions = b
} else {
return nil, err
}
}
} else {
if b, err := common.Marshal(systemPrompt); err == nil {
request.Instructions = b
} else {
return nil, err
}
}
}
}
// codex: store must be false
request.Store = json.RawMessage("false")
// rm max_output_tokens
request.MaxOutputTokens = 0
request.Temperature = nil
return request, 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 *types.NewAPIError) {
if info.RelayMode != relayconstant.RelayModeResponses {
return nil, types.NewError(errors.New("codex channel: endpoint not supported"), types.ErrorCodeInvalidRequest)
}
if info.IsStream {
return openai.OaiResponsesStreamHandler(c, info, resp)
}
return openai.OaiResponsesHandler(c, info, resp)
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.RelayMode != relayconstant.RelayModeResponses {
return "", errors.New("codex channel: only /v1/responses is supported")
}
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, "/backend-api/codex/responses", info.ChannelType), nil
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
key := strings.TrimSpace(info.ApiKey)
if !strings.HasPrefix(key, "{") {
return errors.New("codex channel: key must be a JSON object")
}
oauthKey, err := ParseOAuthKey(key)
if err != nil {
return err
}
accessToken := strings.TrimSpace(oauthKey.AccessToken)
accountID := strings.TrimSpace(oauthKey.AccountID)
if accessToken == "" {
return errors.New("codex channel: access_token is required")
}
if accountID == "" {
return errors.New("codex channel: account_id is required")
}
req.Set("Authorization", "Bearer "+accessToken)
req.Set("chatgpt-account-id", accountID)
if req.Get("OpenAI-Beta") == "" {
req.Set("OpenAI-Beta", "responses=experimental")
}
if req.Get("originator") == "" {
req.Set("originator", "codex_cli_rs")
}
return nil
}

View File

@@ -0,0 +1,9 @@
package codex
var ModelList = []string{
"gpt-5", "gpt-5-codex", "gpt-5-codex-mini",
"gpt-5.1", "gpt-5.1-codex", "gpt-5.1-codex-max", "gpt-5.1-codex-mini",
"gpt-5.2", "gpt-5.2-codex",
}
const ChannelName = "codex"

View File

@@ -0,0 +1,30 @@
package codex
import (
"errors"
"github.com/QuantumNous/new-api/common"
)
type OAuthKey struct {
IDToken string `json:"id_token,omitempty"`
AccessToken string `json:"access_token,omitempty"`
RefreshToken string `json:"refresh_token,omitempty"`
AccountID string `json:"account_id,omitempty"`
LastRefresh string `json:"last_refresh,omitempty"`
Email string `json:"email,omitempty"`
Type string `json:"type,omitempty"`
Expired string `json:"expired,omitempty"`
}
func ParseOAuthKey(raw string) (*OAuthKey, error) {
if raw == "" {
return nil, errors.New("codex channel: empty oauth key")
}
var key OAuthKey
if err := common.Unmarshal([]byte(raw), &key); err != nil {
return nil, errors.New("codex channel: invalid oauth key json")
}
return &key, nil
}

View File

@@ -1,6 +1,7 @@
package gemini
import (
"context"
"encoding/json"
"errors"
"fmt"
@@ -8,6 +9,7 @@ import (
"net/http"
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
@@ -32,6 +34,7 @@ var geminiSupportedMimeTypes = map[string]bool{
"audio/wav": true,
"image/png": true,
"image/jpeg": true,
"image/jpg": true, // support old image/jpeg
"image/webp": true,
"text/plain": true,
"video/mov": true,
@@ -381,7 +384,7 @@ func CovertOpenAI2Gemini(c *gin.Context, textRequest dto.GeneralOpenAIRequest, i
var system_content []string
//shouldAddDummyModelMessage := false
for _, message := range textRequest.Messages {
if message.Role == "system" {
if message.Role == "system" || message.Role == "developer" {
system_content = append(system_content, message.StringContent())
continue
} else if message.Role == "tool" || message.Role == "function" {
@@ -659,101 +662,84 @@ func getSupportedMimeTypesList() []string {
return keys
}
var geminiOpenAPISchemaAllowedFields = map[string]struct{}{
"anyOf": {},
"default": {},
"description": {},
"enum": {},
"example": {},
"format": {},
"items": {},
"maxItems": {},
"maxLength": {},
"maxProperties": {},
"maximum": {},
"minItems": {},
"minLength": {},
"minProperties": {},
"minimum": {},
"nullable": {},
"pattern": {},
"properties": {},
"propertyOrdering": {},
"required": {},
"title": {},
"type": {},
}
const geminiFunctionSchemaMaxDepth = 64
// cleanFunctionParameters recursively removes unsupported fields from Gemini function parameters.
func cleanFunctionParameters(params interface{}) interface{} {
return cleanFunctionParametersWithDepth(params, 0)
}
func cleanFunctionParametersWithDepth(params interface{}, depth int) interface{} {
if params == nil {
return nil
}
if depth >= geminiFunctionSchemaMaxDepth {
return cleanFunctionParametersShallow(params)
}
switch v := params.(type) {
case map[string]interface{}:
// Create a copy to avoid modifying the original
cleanedMap := make(map[string]interface{})
// Keep only Gemini-supported OpenAPI schema subset fields (per official SDK Schema).
cleanedMap := make(map[string]interface{}, len(v))
for k, val := range v {
cleanedMap[k] = val
}
// Remove unsupported root-level fields
delete(cleanedMap, "default")
delete(cleanedMap, "exclusiveMaximum")
delete(cleanedMap, "exclusiveMinimum")
delete(cleanedMap, "$schema")
delete(cleanedMap, "additionalProperties")
// Check and clean 'format' for string types
if propType, typeExists := cleanedMap["type"].(string); typeExists && propType == "string" {
if formatValue, formatExists := cleanedMap["format"].(string); formatExists {
if formatValue != "enum" && formatValue != "date-time" {
delete(cleanedMap, "format")
}
if _, ok := geminiOpenAPISchemaAllowedFields[k]; ok {
cleanedMap[k] = val
}
}
normalizeGeminiSchemaTypeAndNullable(cleanedMap)
// Clean properties
if props, ok := cleanedMap["properties"].(map[string]interface{}); ok && props != nil {
cleanedProps := make(map[string]interface{})
for propName, propValue := range props {
cleanedProps[propName] = cleanFunctionParameters(propValue)
cleanedProps[propName] = cleanFunctionParametersWithDepth(propValue, depth+1)
}
cleanedMap["properties"] = cleanedProps
}
// Recursively clean items in arrays
if items, ok := cleanedMap["items"].(map[string]interface{}); ok && items != nil {
cleanedMap["items"] = cleanFunctionParameters(items)
cleanedMap["items"] = cleanFunctionParametersWithDepth(items, depth+1)
}
// Also handle items if it's an array of schemas
if itemsArray, ok := cleanedMap["items"].([]interface{}); ok {
cleanedItemsArray := make([]interface{}, len(itemsArray))
for i, item := range itemsArray {
cleanedItemsArray[i] = cleanFunctionParameters(item)
}
cleanedMap["items"] = cleanedItemsArray
// OpenAPI tuple-style items is not supported by Gemini SDK Schema; keep first to avoid API rejection.
if itemsArray, ok := cleanedMap["items"].([]interface{}); ok && len(itemsArray) > 0 {
cleanedMap["items"] = cleanFunctionParametersWithDepth(itemsArray[0], depth+1)
}
// Recursively clean other schema composition keywords
for _, field := range []string{"allOf", "anyOf", "oneOf"} {
if nested, ok := cleanedMap[field].([]interface{}); ok {
cleanedNested := make([]interface{}, len(nested))
for i, item := range nested {
cleanedNested[i] = cleanFunctionParameters(item)
}
cleanedMap[field] = cleanedNested
}
}
// Recursively clean patternProperties
if patternProps, ok := cleanedMap["patternProperties"].(map[string]interface{}); ok {
cleanedPatternProps := make(map[string]interface{})
for pattern, schema := range patternProps {
cleanedPatternProps[pattern] = cleanFunctionParameters(schema)
}
cleanedMap["patternProperties"] = cleanedPatternProps
}
// Recursively clean definitions
if definitions, ok := cleanedMap["definitions"].(map[string]interface{}); ok {
cleanedDefinitions := make(map[string]interface{})
for defName, defSchema := range definitions {
cleanedDefinitions[defName] = cleanFunctionParameters(defSchema)
}
cleanedMap["definitions"] = cleanedDefinitions
}
// Recursively clean $defs (newer JSON Schema draft)
if defs, ok := cleanedMap["$defs"].(map[string]interface{}); ok {
cleanedDefs := make(map[string]interface{})
for defName, defSchema := range defs {
cleanedDefs[defName] = cleanFunctionParameters(defSchema)
}
cleanedMap["$defs"] = cleanedDefs
}
// Clean conditional keywords
for _, field := range []string{"if", "then", "else", "not"} {
if nested, ok := cleanedMap[field]; ok {
cleanedMap[field] = cleanFunctionParameters(nested)
// Recursively clean anyOf
if nested, ok := cleanedMap["anyOf"].([]interface{}); ok && nested != nil {
cleanedNested := make([]interface{}, len(nested))
for i, item := range nested {
cleanedNested[i] = cleanFunctionParametersWithDepth(item, depth+1)
}
cleanedMap["anyOf"] = cleanedNested
}
return cleanedMap
@@ -762,7 +748,7 @@ func cleanFunctionParameters(params interface{}) interface{} {
// Handle arrays of schemas
cleanedArray := make([]interface{}, len(v))
for i, item := range v {
cleanedArray[i] = cleanFunctionParameters(item)
cleanedArray[i] = cleanFunctionParametersWithDepth(item, depth+1)
}
return cleanedArray
@@ -772,6 +758,91 @@ func cleanFunctionParameters(params interface{}) interface{} {
}
}
func cleanFunctionParametersShallow(params interface{}) interface{} {
switch v := params.(type) {
case map[string]interface{}:
cleanedMap := make(map[string]interface{}, len(v))
for k, val := range v {
if _, ok := geminiOpenAPISchemaAllowedFields[k]; ok {
cleanedMap[k] = val
}
}
normalizeGeminiSchemaTypeAndNullable(cleanedMap)
// Stop recursion and avoid retaining huge nested structures.
delete(cleanedMap, "properties")
delete(cleanedMap, "items")
delete(cleanedMap, "anyOf")
return cleanedMap
case []interface{}:
// Prefer an empty list over deep recursion on attacker-controlled inputs.
return []interface{}{}
default:
return params
}
}
func normalizeGeminiSchemaTypeAndNullable(schema map[string]interface{}) {
rawType, ok := schema["type"]
if !ok || rawType == nil {
return
}
normalize := func(t string) (string, bool) {
switch strings.ToLower(strings.TrimSpace(t)) {
case "object":
return "OBJECT", false
case "array":
return "ARRAY", false
case "string":
return "STRING", false
case "integer":
return "INTEGER", false
case "number":
return "NUMBER", false
case "boolean":
return "BOOLEAN", false
case "null":
return "", true
default:
return t, false
}
}
switch t := rawType.(type) {
case string:
normalized, isNull := normalize(t)
if isNull {
schema["nullable"] = true
delete(schema, "type")
return
}
schema["type"] = normalized
case []interface{}:
nullable := false
var chosen string
for _, item := range t {
if s, ok := item.(string); ok {
normalized, isNull := normalize(s)
if isNull {
nullable = true
continue
}
if chosen == "" {
chosen = normalized
}
}
}
if nullable {
schema["nullable"] = true
}
if chosen != "" {
schema["type"] = chosen
} else {
delete(schema, "type")
}
}
}
func removeAdditionalPropertiesWithDepth(schema interface{}, depth int) interface{} {
if depth >= 5 {
return schema
@@ -1183,6 +1254,8 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
id := helper.GetResponseID(c)
createAt := common.GetTimestamp()
finishReason := constant.FinishReasonStop
toolCallIndexByChoice := make(map[int]map[string]int)
nextToolCallIndexByChoice := make(map[int]int)
usage, err := geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool {
response, isStop := streamResponseGeminiChat2OpenAI(geminiResponse)
@@ -1190,6 +1263,28 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
response.Id = id
response.Created = createAt
response.Model = info.UpstreamModelName
for choiceIdx := range response.Choices {
choiceKey := response.Choices[choiceIdx].Index
for toolIdx := range response.Choices[choiceIdx].Delta.ToolCalls {
tool := &response.Choices[choiceIdx].Delta.ToolCalls[toolIdx]
if tool.ID == "" {
continue
}
m := toolCallIndexByChoice[choiceKey]
if m == nil {
m = make(map[string]int)
toolCallIndexByChoice[choiceKey] = m
}
if idx, ok := m[tool.ID]; ok {
tool.SetIndex(idx)
continue
}
idx := nextToolCallIndexByChoice[choiceKey]
nextToolCallIndexByChoice[choiceKey] = idx + 1
m[tool.ID] = idx
tool.SetIndex(idx)
}
}
logger.LogDebug(c, fmt.Sprintf("info.SendResponseCount = %d", info.SendResponseCount))
if info.SendResponseCount == 0 {
@@ -1417,6 +1512,79 @@ func GeminiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
return usage, nil
}
type GeminiModelsResponse struct {
Models []dto.GeminiModel `json:"models"`
NextPageToken string `json:"nextPageToken"`
}
func FetchGeminiModels(baseURL, apiKey, proxyURL string) ([]string, error) {
client, err := service.GetHttpClientWithProxy(proxyURL)
if err != nil {
return nil, fmt.Errorf("创建HTTP客户端失败: %v", err)
}
allModels := make([]string, 0)
nextPageToken := ""
maxPages := 100 // Safety limit to prevent infinite loops
for page := 0; page < maxPages; page++ {
url := fmt.Sprintf("%s/v1beta/models", baseURL)
if nextPageToken != "" {
url = fmt.Sprintf("%s?pageToken=%s", url, nextPageToken)
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
request, err := http.NewRequestWithContext(ctx, "GET", url, nil)
if err != nil {
cancel()
return nil, fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("x-goog-api-key", apiKey)
response, err := client.Do(request)
if err != nil {
cancel()
return nil, fmt.Errorf("请求失败: %v", err)
}
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
response.Body.Close()
cancel()
return nil, fmt.Errorf("服务器返回错误 %d: %s", response.StatusCode, string(body))
}
body, err := io.ReadAll(response.Body)
response.Body.Close()
cancel()
if err != nil {
return nil, fmt.Errorf("读取响应失败: %v", err)
}
var modelsResponse GeminiModelsResponse
if err = common.Unmarshal(body, &modelsResponse); err != nil {
return nil, fmt.Errorf("解析响应失败: %v", err)
}
for _, model := range modelsResponse.Models {
modelNameValue, ok := model.Name.(string)
if !ok {
continue
}
modelName := strings.TrimPrefix(modelNameValue, "models/")
allModels = append(allModels, modelName)
}
nextPageToken = modelsResponse.NextPageToken
if nextPageToken == "" {
break
}
}
return allModels, nil
}
// convertToolChoiceToGeminiConfig converts OpenAI tool_choice to Gemini toolConfig
// OpenAI tool_choice values:
// - "auto": Let the model decide (default)

View File

@@ -14,6 +14,9 @@ var ModelList = []string{
"speech-02-turbo",
"speech-01-hd",
"speech-01-turbo",
"MiniMax-M2.1",
"MiniMax-M2.1-lightning",
"MiniMax-M2",
}
var ChannelName = "minimax"

View File

@@ -67,3 +67,40 @@ type OllamaEmbeddingResponse struct {
Embeddings [][]float64 `json:"embeddings"`
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
}
type OllamaTagsResponse struct {
Models []OllamaModel `json:"models"`
}
type OllamaModel struct {
Name string `json:"name"`
Size int64 `json:"size"`
Digest string `json:"digest,omitempty"`
ModifiedAt string `json:"modified_at"`
Details OllamaModelDetail `json:"details,omitempty"`
}
type OllamaModelDetail struct {
ParentModel string `json:"parent_model,omitempty"`
Format string `json:"format,omitempty"`
Family string `json:"family,omitempty"`
Families []string `json:"families,omitempty"`
ParameterSize string `json:"parameter_size,omitempty"`
QuantizationLevel string `json:"quantization_level,omitempty"`
}
type OllamaPullRequest struct {
Name string `json:"name"`
Stream bool `json:"stream,omitempty"`
}
type OllamaPullResponse struct {
Status string `json:"status"`
Digest string `json:"digest,omitempty"`
Total int64 `json:"total,omitempty"`
Completed int64 `json:"completed,omitempty"`
}
type OllamaDeleteRequest struct {
Name string `json:"name"`
}

View File

@@ -1,11 +1,13 @@
package ollama
import (
"bufio"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
@@ -283,3 +285,246 @@ func ollamaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h
service.IOCopyBytesGracefully(c, resp, out)
return usage, nil
}
func FetchOllamaModels(baseURL, apiKey string) ([]OllamaModel, error) {
url := fmt.Sprintf("%s/api/tags", baseURL)
client := &http.Client{}
request, err := http.NewRequest("GET", url, nil)
if err != nil {
return nil, fmt.Errorf("创建请求失败: %v", err)
}
// Ollama 通常不需要 Bearer token但为了兼容性保留
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return nil, fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return nil, fmt.Errorf("服务器返回错误 %d: %s", response.StatusCode, string(body))
}
var tagsResponse OllamaTagsResponse
body, err := io.ReadAll(response.Body)
if err != nil {
return nil, fmt.Errorf("读取响应失败: %v", err)
}
err = common.Unmarshal(body, &tagsResponse)
if err != nil {
return nil, fmt.Errorf("解析响应失败: %v", err)
}
return tagsResponse.Models, nil
}
// 拉取 Ollama 模型 (非流式)
func PullOllamaModel(baseURL, apiKey, modelName string) error {
url := fmt.Sprintf("%s/api/pull", baseURL)
pullRequest := OllamaPullRequest{
Name: modelName,
Stream: false, // 非流式,简化处理
}
requestBody, err := common.Marshal(pullRequest)
if err != nil {
return fmt.Errorf("序列化请求失败: %v", err)
}
client := &http.Client{
Timeout: 30 * 60 * 1000 * time.Millisecond, // 30分钟超时支持大模型
}
request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
if err != nil {
return fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("Content-Type", "application/json")
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
}
return nil
}
// 流式拉取 Ollama 模型 (支持进度回调)
func PullOllamaModelStream(baseURL, apiKey, modelName string, progressCallback func(OllamaPullResponse)) error {
url := fmt.Sprintf("%s/api/pull", baseURL)
pullRequest := OllamaPullRequest{
Name: modelName,
Stream: true, // 启用流式
}
requestBody, err := common.Marshal(pullRequest)
if err != nil {
return fmt.Errorf("序列化请求失败: %v", err)
}
client := &http.Client{
Timeout: 60 * 60 * 1000 * time.Millisecond, // 1小时超时支持超大模型
}
request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
if err != nil {
return fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("Content-Type", "application/json")
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
}
// 读取流式响应
scanner := bufio.NewScanner(response.Body)
successful := false
for scanner.Scan() {
line := scanner.Text()
if strings.TrimSpace(line) == "" {
continue
}
var pullResponse OllamaPullResponse
if err := common.Unmarshal([]byte(line), &pullResponse); err != nil {
continue // 忽略解析失败的行
}
if progressCallback != nil {
progressCallback(pullResponse)
}
// 检查是否出现错误或完成
if strings.EqualFold(pullResponse.Status, "error") {
return fmt.Errorf("拉取模型失败: %s", strings.TrimSpace(line))
}
if strings.EqualFold(pullResponse.Status, "success") {
successful = true
break
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("读取流式响应失败: %v", err)
}
if !successful {
return fmt.Errorf("拉取模型未完成: 未收到成功状态")
}
return nil
}
// 删除 Ollama 模型
func DeleteOllamaModel(baseURL, apiKey, modelName string) error {
url := fmt.Sprintf("%s/api/delete", baseURL)
deleteRequest := OllamaDeleteRequest{
Name: modelName,
}
requestBody, err := common.Marshal(deleteRequest)
if err != nil {
return fmt.Errorf("序列化请求失败: %v", err)
}
client := &http.Client{}
request, err := http.NewRequest("DELETE", url, strings.NewReader(string(requestBody)))
if err != nil {
return fmt.Errorf("创建请求失败: %v", err)
}
request.Header.Set("Content-Type", "application/json")
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
body, _ := io.ReadAll(response.Body)
return fmt.Errorf("删除模型失败 %d: %s", response.StatusCode, string(body))
}
return nil
}
func FetchOllamaVersion(baseURL, apiKey string) (string, error) {
trimmedBase := strings.TrimRight(baseURL, "/")
if trimmedBase == "" {
return "", fmt.Errorf("baseURL 为空")
}
url := fmt.Sprintf("%s/api/version", trimmedBase)
client := &http.Client{Timeout: 10 * time.Second}
request, err := http.NewRequest("GET", url, nil)
if err != nil {
return "", fmt.Errorf("创建请求失败: %v", err)
}
if apiKey != "" {
request.Header.Set("Authorization", "Bearer "+apiKey)
}
response, err := client.Do(request)
if err != nil {
return "", fmt.Errorf("请求失败: %v", err)
}
defer response.Body.Close()
body, err := io.ReadAll(response.Body)
if err != nil {
return "", fmt.Errorf("读取响应失败: %v", err)
}
if response.StatusCode != http.StatusOK {
return "", fmt.Errorf("查询版本失败 %d: %s", response.StatusCode, string(body))
}
var versionResp struct {
Version string `json:"version"`
}
if err := json.Unmarshal(body, &versionResp); err != nil {
return "", fmt.Errorf("解析响应失败: %v", err)
}
if versionResp.Version == "" {
return "", fmt.Errorf("未返回版本信息")
}
return versionResp.Version, nil
}

View File

@@ -0,0 +1,369 @@
package openai
import (
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
if resp == nil || resp.Body == nil {
return nil, types.NewOpenAIError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
defer service.CloseResponseBodyGracefully(resp)
var responsesResp dto.OpenAIResponsesResponse
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError)
}
if err := common.Unmarshal(body, &responsesResp); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if oaiError := responsesResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" {
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
}
chatId := helper.GetResponseID(c)
chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
if usage == nil || usage.TotalTokens == 0 {
text := service.ExtractOutputTextFromResponses(&responsesResp)
usage = service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens())
chatResp.Usage = *usage
}
chatBody, err := common.Marshal(chatResp)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError)
}
service.IOCopyBytesGracefully(c, resp, chatBody)
return usage, nil
}
func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
if resp == nil || resp.Body == nil {
return nil, types.NewOpenAIError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
defer service.CloseResponseBodyGracefully(resp)
responseId := helper.GetResponseID(c)
createAt := time.Now().Unix()
model := info.UpstreamModelName
var (
usage = &dto.Usage{}
outputText strings.Builder
usageText strings.Builder
sentStart bool
sentStop bool
sawToolCall bool
streamErr *types.NewAPIError
)
toolCallIndexByID := make(map[string]int)
toolCallNameByID := make(map[string]string)
toolCallArgsByID := make(map[string]string)
toolCallNameSent := make(map[string]bool)
toolCallCanonicalIDByItemID := make(map[string]string)
sendStartIfNeeded := func() bool {
if sentStart {
return true
}
if err := helper.ObjectData(c, helper.GenerateStartEmptyResponse(responseId, createAt, model, nil)); err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
sentStart = true
return true
}
sendToolCallDelta := func(callID string, name string, argsDelta string) bool {
if callID == "" {
return true
}
if outputText.Len() > 0 {
// Prefer streaming assistant text over tool calls to match non-stream behavior.
return true
}
if !sendStartIfNeeded() {
return false
}
idx, ok := toolCallIndexByID[callID]
if !ok {
idx = len(toolCallIndexByID)
toolCallIndexByID[callID] = idx
}
if name != "" {
toolCallNameByID[callID] = name
}
if toolCallNameByID[callID] != "" {
name = toolCallNameByID[callID]
}
tool := dto.ToolCallResponse{
ID: callID,
Type: "function",
Function: dto.FunctionResponse{
Arguments: argsDelta,
},
}
tool.SetIndex(idx)
if name != "" && !toolCallNameSent[callID] {
tool.Function.Name = name
toolCallNameSent[callID] = true
}
chunk := &dto.ChatCompletionsStreamResponse{
Id: responseId,
Object: "chat.completion.chunk",
Created: createAt,
Model: model,
Choices: []dto.ChatCompletionsStreamResponseChoice{
{
Index: 0,
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{
ToolCalls: []dto.ToolCallResponse{tool},
},
},
},
}
if err := helper.ObjectData(c, chunk); err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
sawToolCall = true
// Include tool call data in the local builder for fallback token estimation.
if tool.Function.Name != "" {
usageText.WriteString(tool.Function.Name)
}
if argsDelta != "" {
usageText.WriteString(argsDelta)
}
return true
}
helper.StreamScannerHandler(c, resp, info, func(data string) bool {
if streamErr != nil {
return false
}
var streamResp dto.ResponsesStreamResponse
if err := common.UnmarshalJsonStr(data, &streamResp); err != nil {
logger.LogError(c, "failed to unmarshal responses stream event: "+err.Error())
return true
}
switch streamResp.Type {
case "response.created":
if streamResp.Response != nil {
if streamResp.Response.Model != "" {
model = streamResp.Response.Model
}
if streamResp.Response.CreatedAt != 0 {
createAt = int64(streamResp.Response.CreatedAt)
}
}
case "response.output_text.delta":
if !sendStartIfNeeded() {
return false
}
if streamResp.Delta != "" {
outputText.WriteString(streamResp.Delta)
usageText.WriteString(streamResp.Delta)
delta := streamResp.Delta
chunk := &dto.ChatCompletionsStreamResponse{
Id: responseId,
Object: "chat.completion.chunk",
Created: createAt,
Model: model,
Choices: []dto.ChatCompletionsStreamResponseChoice{
{
Index: 0,
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{
Content: &delta,
},
},
},
}
if err := helper.ObjectData(c, chunk); err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
}
case "response.output_item.added", "response.output_item.done":
if streamResp.Item == nil {
break
}
if streamResp.Item.Type != "function_call" {
break
}
itemID := strings.TrimSpace(streamResp.Item.ID)
callID := strings.TrimSpace(streamResp.Item.CallId)
if callID == "" {
callID = itemID
}
if itemID != "" && callID != "" {
toolCallCanonicalIDByItemID[itemID] = callID
}
name := strings.TrimSpace(streamResp.Item.Name)
if name != "" {
toolCallNameByID[callID] = name
}
newArgs := streamResp.Item.Arguments
prevArgs := toolCallArgsByID[callID]
argsDelta := ""
if newArgs != "" {
if strings.HasPrefix(newArgs, prevArgs) {
argsDelta = newArgs[len(prevArgs):]
} else {
argsDelta = newArgs
}
toolCallArgsByID[callID] = newArgs
}
if !sendToolCallDelta(callID, name, argsDelta) {
return false
}
case "response.function_call_arguments.delta":
itemID := strings.TrimSpace(streamResp.ItemID)
callID := toolCallCanonicalIDByItemID[itemID]
if callID == "" {
callID = itemID
}
if callID == "" {
break
}
toolCallArgsByID[callID] += streamResp.Delta
if !sendToolCallDelta(callID, "", streamResp.Delta) {
return false
}
case "response.function_call_arguments.done":
case "response.completed":
if streamResp.Response != nil {
if streamResp.Response.Model != "" {
model = streamResp.Response.Model
}
if streamResp.Response.CreatedAt != 0 {
createAt = int64(streamResp.Response.CreatedAt)
}
if streamResp.Response.Usage != nil {
if streamResp.Response.Usage.InputTokens != 0 {
usage.PromptTokens = streamResp.Response.Usage.InputTokens
usage.InputTokens = streamResp.Response.Usage.InputTokens
}
if streamResp.Response.Usage.OutputTokens != 0 {
usage.CompletionTokens = streamResp.Response.Usage.OutputTokens
usage.OutputTokens = streamResp.Response.Usage.OutputTokens
}
if streamResp.Response.Usage.TotalTokens != 0 {
usage.TotalTokens = streamResp.Response.Usage.TotalTokens
} else {
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
}
if streamResp.Response.Usage.InputTokensDetails != nil {
usage.PromptTokensDetails.CachedTokens = streamResp.Response.Usage.InputTokensDetails.CachedTokens
usage.PromptTokensDetails.ImageTokens = streamResp.Response.Usage.InputTokensDetails.ImageTokens
usage.PromptTokensDetails.AudioTokens = streamResp.Response.Usage.InputTokensDetails.AudioTokens
}
if streamResp.Response.Usage.CompletionTokenDetails.ReasoningTokens != 0 {
usage.CompletionTokenDetails.ReasoningTokens = streamResp.Response.Usage.CompletionTokenDetails.ReasoningTokens
}
}
}
if !sendStartIfNeeded() {
return false
}
if !sentStop {
finishReason := "stop"
if sawToolCall && outputText.Len() == 0 {
finishReason = "tool_calls"
}
stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason)
if err := helper.ObjectData(c, stop); err != nil {
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
}
sentStop = true
}
case "response.error", "response.failed":
if streamResp.Response != nil {
if oaiErr := streamResp.Response.GetOpenAIError(); oaiErr != nil && oaiErr.Type != "" {
streamErr = types.WithOpenAIError(*oaiErr, http.StatusInternalServerError)
return false
}
}
streamErr = types.NewOpenAIError(fmt.Errorf("responses stream error: %s", streamResp.Type), types.ErrorCodeBadResponse, http.StatusInternalServerError)
return false
default:
}
return true
})
if streamErr != nil {
return nil, streamErr
}
if usage.TotalTokens == 0 {
usage = service.ResponseText2Usage(c, usageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
}
if !sentStart {
if err := helper.ObjectData(c, helper.GenerateStartEmptyResponse(responseId, createAt, model, nil)); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
}
if !sentStop {
finishReason := "stop"
if sawToolCall && outputText.Len() == 0 {
finishReason = "tool_calls"
}
stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason)
if err := helper.ObjectData(c, stop); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
}
if info.ShouldIncludeUsage && usage != nil {
if err := helper.ObjectData(c, helper.GenerateFinalUsageResponse(responseId, createAt, model, *usage)); err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
}
helper.Done(c)
return usage, nil
}

View File

@@ -208,7 +208,6 @@ func HandleFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, lastStream
helper.Done(c)
case types.RelayFormatClaude:
info.ClaudeConvertInfo.Done = true
var streamResponse dto.ChatCompletionsStreamResponse
if err := common.Unmarshal(common.StringToByteSlice(lastStreamData), &streamResponse); err != nil {
common.SysLog("error unmarshalling stream response: " + err.Error())
@@ -221,6 +220,7 @@ func HandleFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, lastStream
for _, resp := range claudeResponses {
_ = helper.ClaudeData(c, *resp)
}
info.ClaudeConvertInfo.Done = true
case types.RelayFormatGemini:
var streamResponse dto.ChatCompletionsStreamResponse

View File

@@ -186,7 +186,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
usage.CompletionTokens += toolCount * 7
}
applyUsagePostProcessing(info, usage, nil)
applyUsagePostProcessing(info, usage, common.StringToByteSlice(lastStreamData))
HandleFinalResponse(c, info, lastStreamData, responseId, createAt, model, systemFingerprint, usage, containStreamUsage)
@@ -596,7 +596,8 @@ func applyUsagePostProcessing(info *relaycommon.RelayInfo, usage *dto.Usage, res
if usage.PromptTokensDetails.CachedTokens == 0 && usage.PromptCacheHitTokens != 0 {
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
}
case constant.ChannelTypeZhipu_v4, constant.ChannelTypeMoonshot:
case constant.ChannelTypeZhipu_v4:
// 智普的cached_tokens在标准位置: usage.prompt_tokens_details.cached_tokens
if usage.PromptTokensDetails.CachedTokens == 0 {
if usage.InputTokensDetails != nil && usage.InputTokensDetails.CachedTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens
@@ -606,6 +607,19 @@ func applyUsagePostProcessing(info *relaycommon.RelayInfo, usage *dto.Usage, res
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
}
}
case constant.ChannelTypeMoonshot:
// Moonshot的cached_tokens在非标准位置: choices[].usage.cached_tokens
if usage.PromptTokensDetails.CachedTokens == 0 {
if usage.InputTokensDetails != nil && usage.InputTokensDetails.CachedTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens
} else if cachedTokens, ok := extractMoonshotCachedTokensFromBody(responseBody); ok {
usage.PromptTokensDetails.CachedTokens = cachedTokens
} else if cachedTokens, ok := extractCachedTokensFromBody(responseBody); ok {
usage.PromptTokensDetails.CachedTokens = cachedTokens
} else if usage.PromptCacheHitTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
}
}
}
}
@@ -639,3 +653,32 @@ func extractCachedTokensFromBody(body []byte) (int, bool) {
}
return 0, false
}
// extractMoonshotCachedTokensFromBody 从Moonshot的非标准位置提取cached_tokens
// Moonshot的流式响应格式: {"choices":[{"usage":{"cached_tokens":111}}]}
func extractMoonshotCachedTokensFromBody(body []byte) (int, bool) {
if len(body) == 0 {
return 0, false
}
var payload struct {
Choices []struct {
Usage struct {
CachedTokens *int `json:"cached_tokens"`
} `json:"usage"`
} `json:"choices"`
}
if err := common.Unmarshal(body, &payload); err != nil {
return 0, false
}
// 遍历choices查找cached_tokens
for _, choice := range payload.Choices {
if choice.Usage.CachedTokens != nil && *choice.Usage.CachedTokens > 0 {
return *choice.Usage.CachedTokens, true
}
}
return 0, false
}

View File

@@ -192,6 +192,10 @@ func sizeToResolution(size string) (string, error) {
func ProcessAliOtherRatios(aliReq *AliVideoRequest) (map[string]float64, error) {
otherRatios := make(map[string]float64)
aliRatios := map[string]map[string]float64{
"wan2.6-i2v": {
"720P": 1,
"1080P": 1 / 0.6,
},
"wan2.5-t2v-preview": {
"480P": 1,
"720P": 2,
@@ -287,7 +291,9 @@ func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relay
aliReq.Parameters.Size = "1280*720"
}
} else {
if strings.HasPrefix(req.Model, "wan2.5") {
if strings.HasPrefix(req.Model, "wan2.6") {
aliReq.Parameters.Resolution = "1080P"
} else if strings.HasPrefix(req.Model, "wan2.5") {
aliReq.Parameters.Resolution = "1080P"
} else if strings.HasPrefix(req.Model, "wan2.2-i2v-flash") {
aliReq.Parameters.Resolution = "720P"

View File

@@ -6,6 +6,9 @@ import (
"fmt"
"io"
"net/http"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
@@ -23,18 +26,36 @@ import (
// ============================
type ContentItem struct {
Type string `json:"type"` // "text" or "image_url"
Text string `json:"text,omitempty"` // for text type
ImageURL *ImageURL `json:"image_url,omitempty"` // for image_url type
Type string `json:"type"` // "text", "image_url" or "video"
Text string `json:"text,omitempty"` // for text type
ImageURL *ImageURL `json:"image_url,omitempty"` // for image_url type
Video *VideoReference `json:"video,omitempty"` // for video (sample) type
}
type ImageURL struct {
URL string `json:"url"`
}
type VideoReference struct {
URL string `json:"url"` // Draft video URL
}
type requestPayload struct {
Model string `json:"model"`
Content []ContentItem `json:"content"`
Model string `json:"model"`
Content []ContentItem `json:"content"`
CallbackURL string `json:"callback_url,omitempty"`
ReturnLastFrame *dto.BoolValue `json:"return_last_frame,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
ExecutionExpiresAfter dto.IntValue `json:"execution_expires_after,omitempty"`
GenerateAudio *dto.BoolValue `json:"generate_audio,omitempty"`
Draft *dto.BoolValue `json:"draft,omitempty"`
Resolution string `json:"resolution,omitempty"`
Ratio string `json:"ratio,omitempty"`
Duration dto.IntValue `json:"duration,omitempty"`
Frames dto.IntValue `json:"frames,omitempty"`
Seed dto.IntValue `json:"seed,omitempty"`
CameraFixed *dto.BoolValue `json:"camera_fixed,omitempty"`
Watermark *dto.BoolValue `json:"watermark,omitempty"`
}
type responsePayload struct {
@@ -53,6 +74,7 @@ type responseTask struct {
Duration int `json:"duration"`
Ratio string `json:"ratio"`
FramesPerSecond int `json:"framespersecond"`
ServiceTier string `json:"service_tier"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
@@ -98,16 +120,16 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
// BuildRequestBody converts request into Doubao specific format.
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
v, exists := c.Get("task_request")
if !exists {
return nil, fmt.Errorf("request not found in context")
req, err := relaycommon.GetTaskRequest(c)
if err != nil {
return nil, err
}
req := v.(relaycommon.TaskSubmitReq)
body, err := a.convertToRequestPayload(&req)
if err != nil {
return nil, errors.Wrap(err, "convert request payload failed")
}
info.UpstreamModelName = body.Model
data, err := json.Marshal(body)
if err != nil {
return nil, err
@@ -141,7 +163,13 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
c.JSON(http.StatusOK, gin.H{"task_id": dResp.ID})
ov := dto.NewOpenAIVideo()
ov.ID = dResp.ID
ov.TaskID = dResp.ID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
return dResp.ID, responseBody, nil
}
@@ -204,12 +232,15 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
}
}
// TODO: Add support for additional parameters from metadata
// such as ratio, duration, seed, etc.
// metadata := req.Metadata
// if metadata != nil {
// // Parse and apply metadata parameters
// }
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
return &r, nil
}
@@ -229,7 +260,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
case "pending", "queued":
taskResult.Status = model.TaskStatusQueued
taskResult.Progress = "10%"
case "processing":
case "processing", "running":
taskResult.Status = model.TaskStatusInProgress
taskResult.Progress = "50%"
case "succeeded":
@@ -251,3 +282,30 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
return &taskResult, nil
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var dResp responseTask
if err := json.Unmarshal(originTask.Data, &dResp); err != nil {
return nil, errors.Wrap(err, "unmarshal doubao task data failed")
}
openAIVideo := dto.NewOpenAIVideo()
openAIVideo.ID = originTask.TaskID
openAIVideo.TaskID = originTask.TaskID
openAIVideo.Status = originTask.Status.ToVideoStatus()
openAIVideo.SetProgressStr(originTask.Progress)
openAIVideo.SetMetadata("url", dResp.Content.VideoURL)
openAIVideo.CreatedAt = originTask.CreatedAt
openAIVideo.CompletedAt = originTask.UpdatedAt
openAIVideo.Model = originTask.Properties.OriginModelName
if dResp.Status == "failed" {
openAIVideo.Error = &dto.OpenAIVideoError{
Message: "task failed",
Code: "failed",
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
}

View File

@@ -4,6 +4,7 @@ var ModelList = []string{
"doubao-seedance-1-0-pro-250528",
"doubao-seedance-1-0-lite-t2v",
"doubao-seedance-1-0-lite-i2v",
"doubao-seedance-1-5-pro-251215",
}
var ChannelName = "doubao-video"

View File

@@ -17,6 +17,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
"github.com/pkg/errors"
@@ -409,14 +410,15 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
// 即梦视频3.0 ReqKey转换
// https://www.volcengine.com/docs/85621/1792707
imageLen := lo.Max([]int{len(req.Images), len(r.BinaryDataBase64), len(r.ImageUrls)})
if strings.Contains(r.ReqKey, "jimeng_v30") {
if r.ReqKey == "jimeng_v30_pro" {
// 3.0 pro只有固定的jimeng_ti2v_v30_pro
r.ReqKey = "jimeng_ti2v_v30_pro"
} else if len(req.Images) > 1 {
} else if imageLen > 1 {
// 多张图片:首尾帧生成
r.ReqKey = strings.TrimSuffix(strings.Replace(r.ReqKey, "jimeng_v30", "jimeng_i2v_first_tail_v30", 1), "p")
} else if len(req.Images) == 1 {
} else if imageLen == 1 {
// 单张图片:图生视频
r.ReqKey = strings.TrimSuffix(strings.Replace(r.ReqKey, "jimeng_v30", "jimeng_i2v_first_v30", 1), "p")
} else {

View File

@@ -346,7 +346,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
}
taskInfo.Code = resPayload.Code
taskInfo.TaskID = resPayload.Data.TaskId
taskInfo.Reason = resPayload.Message
taskInfo.Reason = resPayload.Data.TaskStatusMsg
//任务状态枚举值submitted已提交、processing处理中、succeed成功、failed失败
status := resPayload.Data.TaskStatus
switch status {

View File

@@ -40,6 +40,7 @@ var claudeModelMap = map[string]string{
"claude-opus-4-20250514": "claude-opus-4@20250514",
"claude-opus-4-1-20250805": "claude-opus-4-1@20250805",
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5@20250929",
"claude-haiku-4-5-20251001": "claude-haiku-4-5@20251001",
"claude-opus-4-5-20251101": "claude-opus-4-5@20251101",
}

View File

@@ -270,6 +270,8 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
// return fmt.Sprintf("%s/api/v3/images/edits", baseUrl), nil
case constant.RelayModeRerank:
return fmt.Sprintf("%s/api/v3/rerank", baseUrl), nil
case constant.RelayModeResponses:
return fmt.Sprintf("%s/api/v3/responses", baseUrl), nil
case constant.RelayModeAudioSpeech:
if baseUrl == channelconstant.ChannelBaseURLs[channelconstant.ChannelTypeVolcEngine] {
return "wss://openspeech.bytedance.com/api/v1/tts/ws_binary", nil
@@ -323,7 +325,7 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
return nil, errors.New("not implemented")
return request, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {

View File

@@ -0,0 +1,162 @@
package relay
import (
"bytes"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
openaichannel "github.com/QuantumNous/new-api/relay/channel/openai"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) {
if info == nil || request == nil {
return
}
if info.ChannelSetting.SystemPrompt == "" {
return
}
systemRole := request.GetSystemRoleName()
containSystemPrompt := false
for _, message := range request.Messages {
if message.Role == systemRole {
containSystemPrompt = true
break
}
}
if !containSystemPrompt {
systemMessage := dto.Message{
Role: systemRole,
Content: info.ChannelSetting.SystemPrompt,
}
request.Messages = append([]dto.Message{systemMessage}, request.Messages...)
return
}
if !info.ChannelSetting.SystemPromptOverride {
return
}
common.SetContextKey(c, constant.ContextKeySystemPromptOverride, true)
for i, message := range request.Messages {
if message.Role != systemRole {
continue
}
if message.IsStringContent() {
request.Messages[i].SetStringContent(info.ChannelSetting.SystemPrompt + "\n" + message.StringContent())
return
}
contents := message.ParseContent()
contents = append([]dto.MediaContent{
{
Type: dto.ContentTypeText,
Text: info.ChannelSetting.SystemPrompt,
},
}, contents...)
request.Messages[i].Content = contents
return
}
}
func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.NewAPIError) {
overrideCtx := relaycommon.BuildParamOverrideContext(info)
chatJSON, err := common.Marshal(request)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
chatJSON, err = relaycommon.RemoveDisabledFields(chatJSON, info.ChannelOtherSettings)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
if len(info.ParamOverride) > 0 {
chatJSON, err = relaycommon.ApplyParamOverride(chatJSON, info.ParamOverride, overrideCtx)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry())
}
}
var overriddenChatReq dto.GeneralOpenAIRequest
if err := common.Unmarshal(chatJSON, &overriddenChatReq); err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry())
}
responsesReq, err := service.ChatCompletionsRequestToResponsesRequest(&overriddenChatReq)
if err != nil {
return nil, types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
info.AppendRequestConversion(types.RelayFormatOpenAIResponses)
savedRelayMode := info.RelayMode
savedRequestURLPath := info.RequestURLPath
defer func() {
info.RelayMode = savedRelayMode
info.RequestURLPath = savedRequestURLPath
}()
info.RelayMode = relayconstant.RelayModeResponses
info.RequestURLPath = "/v1/responses"
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *responsesReq)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
var httpResp *http.Response
resp, err := adaptor.DoRequest(c, info, bytes.NewBuffer(jsonData))
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
}
if resp == nil {
return nil, types.NewOpenAIError(nil, types.ErrorCodeBadResponse, http.StatusInternalServerError)
}
statusCodeMappingStr := c.GetString("status_code_mapping")
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newApiErr := service.RelayErrorHandler(c.Request.Context(), httpResp, false)
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
return nil, newApiErr
}
if info.IsStream {
usage, newApiErr := openaichannel.OaiResponsesToChatStreamHandler(c, info, httpResp)
if newApiErr != nil {
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
return nil, newApiErr
}
return usage, nil
}
usage, newApiErr := openaichannel.OaiResponsesToChatHandler(c, info, httpResp)
if newApiErr != nil {
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
return nil, newApiErr
}
return usage, nil
}

View File

@@ -110,6 +110,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())

View File

@@ -23,7 +23,7 @@ type ConditionOperation struct {
type ParamOperation struct {
Path string `json:"path"`
Mode string `json:"mode"` // delete, set, move, prepend, append
Mode string `json:"mode"` // delete, set, move, copy, prepend, append, trim_prefix, trim_suffix, ensure_prefix, ensure_suffix, trim_space, to_lower, to_upper, replace, regex_replace
Value interface{} `json:"value"`
KeepOrigin bool `json:"keep_origin"`
From string `json:"from,omitempty"`
@@ -330,8 +330,6 @@ func applyOperations(jsonStr string, operations []ParamOperation, conditionConte
}
// 处理路径中的负数索引
opPath := processNegativeIndex(result, op.Path)
opFrom := processNegativeIndex(result, op.From)
opTo := processNegativeIndex(result, op.To)
switch op.Mode {
case "delete":
@@ -342,11 +340,38 @@ func applyOperations(jsonStr string, operations []ParamOperation, conditionConte
}
result, err = sjson.Set(result, opPath, op.Value)
case "move":
opFrom := processNegativeIndex(result, op.From)
opTo := processNegativeIndex(result, op.To)
result, err = moveValue(result, opFrom, opTo)
case "copy":
if op.From == "" || op.To == "" {
return "", fmt.Errorf("copy from/to is required")
}
opFrom := processNegativeIndex(result, op.From)
opTo := processNegativeIndex(result, op.To)
result, err = copyValue(result, opFrom, opTo)
case "prepend":
result, err = modifyValue(result, opPath, op.Value, op.KeepOrigin, true)
case "append":
result, err = modifyValue(result, opPath, op.Value, op.KeepOrigin, false)
case "trim_prefix":
result, err = trimStringValue(result, opPath, op.Value, true)
case "trim_suffix":
result, err = trimStringValue(result, opPath, op.Value, false)
case "ensure_prefix":
result, err = ensureStringAffix(result, opPath, op.Value, true)
case "ensure_suffix":
result, err = ensureStringAffix(result, opPath, op.Value, false)
case "trim_space":
result, err = transformStringValue(result, opPath, strings.TrimSpace)
case "to_lower":
result, err = transformStringValue(result, opPath, strings.ToLower)
case "to_upper":
result, err = transformStringValue(result, opPath, strings.ToUpper)
case "replace":
result, err = replaceStringValue(result, opPath, op.From, op.To)
case "regex_replace":
result, err = regexReplaceStringValue(result, opPath, op.From, op.To)
default:
return "", fmt.Errorf("unknown operation: %s", op.Mode)
}
@@ -369,6 +394,14 @@ func moveValue(jsonStr, fromPath, toPath string) (string, error) {
return sjson.Delete(result, fromPath)
}
func copyValue(jsonStr, fromPath, toPath string) (string, error) {
sourceValue := gjson.Get(jsonStr, fromPath)
if !sourceValue.Exists() {
return jsonStr, fmt.Errorf("source path does not exist: %s", fromPath)
}
return sjson.Set(jsonStr, toPath, sourceValue.Value())
}
func modifyValue(jsonStr, path string, value interface{}, keepOrigin, isPrepend bool) (string, error) {
current := gjson.Get(jsonStr, path)
switch {
@@ -422,6 +455,88 @@ func modifyString(jsonStr, path string, value interface{}, isPrepend bool) (stri
return sjson.Set(jsonStr, path, newStr)
}
func trimStringValue(jsonStr, path string, value interface{}, isPrefix bool) (string, error) {
current := gjson.Get(jsonStr, path)
if current.Type != gjson.String {
return jsonStr, fmt.Errorf("operation not supported for type: %v", current.Type)
}
if value == nil {
return jsonStr, fmt.Errorf("trim value is required")
}
valueStr := fmt.Sprintf("%v", value)
var newStr string
if isPrefix {
newStr = strings.TrimPrefix(current.String(), valueStr)
} else {
newStr = strings.TrimSuffix(current.String(), valueStr)
}
return sjson.Set(jsonStr, path, newStr)
}
func ensureStringAffix(jsonStr, path string, value interface{}, isPrefix bool) (string, error) {
current := gjson.Get(jsonStr, path)
if current.Type != gjson.String {
return jsonStr, fmt.Errorf("operation not supported for type: %v", current.Type)
}
if value == nil {
return jsonStr, fmt.Errorf("ensure value is required")
}
valueStr := fmt.Sprintf("%v", value)
if valueStr == "" {
return jsonStr, fmt.Errorf("ensure value is required")
}
currentStr := current.String()
if isPrefix {
if strings.HasPrefix(currentStr, valueStr) {
return jsonStr, nil
}
return sjson.Set(jsonStr, path, valueStr+currentStr)
}
if strings.HasSuffix(currentStr, valueStr) {
return jsonStr, nil
}
return sjson.Set(jsonStr, path, currentStr+valueStr)
}
func transformStringValue(jsonStr, path string, transform func(string) string) (string, error) {
current := gjson.Get(jsonStr, path)
if current.Type != gjson.String {
return jsonStr, fmt.Errorf("operation not supported for type: %v", current.Type)
}
return sjson.Set(jsonStr, path, transform(current.String()))
}
func replaceStringValue(jsonStr, path, from, to string) (string, error) {
current := gjson.Get(jsonStr, path)
if current.Type != gjson.String {
return jsonStr, fmt.Errorf("operation not supported for type: %v", current.Type)
}
if from == "" {
return jsonStr, fmt.Errorf("replace from is required")
}
return sjson.Set(jsonStr, path, strings.ReplaceAll(current.String(), from, to))
}
func regexReplaceStringValue(jsonStr, path, pattern, replacement string) (string, error) {
current := gjson.Get(jsonStr, path)
if current.Type != gjson.String {
return jsonStr, fmt.Errorf("operation not supported for type: %v", current.Type)
}
if pattern == "" {
return jsonStr, fmt.Errorf("regex pattern is required")
}
re, err := regexp.Compile(pattern)
if err != nil {
return jsonStr, err
}
return sjson.Set(jsonStr, path, re.ReplaceAllString(current.String(), replacement))
}
func mergeObjects(jsonStr, path string, value interface{}, keepOrigin bool) (string, error) {
current := gjson.Get(jsonStr, path)
var currentMap, newMap map[string]interface{}
@@ -455,18 +570,19 @@ func mergeObjects(jsonStr, path string, value interface{}, keepOrigin bool) (str
// BuildParamOverrideContext 提供 ApplyParamOverride 可用的上下文信息。
// 目前内置以下字段:
// - model优先使用上游模型名UpstreamModelName若不存在则回落到原始模型名OriginModelName
// - upstream_model始终为通道映射后的上游模型名。
// - upstream_model/model始终为通道映射后的上游模型名
// - original_model请求最初指定的模型名。
// - request_path请求路径
// - is_channel_test是否为渠道测试请求同 is_test
func BuildParamOverrideContext(info *RelayInfo) map[string]interface{} {
if info == nil || info.ChannelMeta == nil {
if info == nil {
return nil
}
ctx := make(map[string]interface{})
if info.UpstreamModelName != "" {
ctx["model"] = info.UpstreamModelName
ctx["upstream_model"] = info.UpstreamModelName
if info.ChannelMeta != nil && info.ChannelMeta.UpstreamModelName != "" {
ctx["model"] = info.ChannelMeta.UpstreamModelName
ctx["upstream_model"] = info.ChannelMeta.UpstreamModelName
}
if info.OriginModelName != "" {
ctx["original_model"] = info.OriginModelName
@@ -475,8 +591,13 @@ func BuildParamOverrideContext(info *RelayInfo) map[string]interface{} {
}
}
if len(ctx) == 0 {
return nil
if info.RequestURLPath != "" {
requestPath := info.RequestURLPath
if requestPath != "" {
ctx["request_path"] = requestPath
}
}
ctx["is_channel_test"] = info.IsChannelTest
return ctx
}

View File

@@ -0,0 +1,791 @@
package common
import (
"encoding/json"
"reflect"
"testing"
)
func TestApplyParamOverrideTrimPrefix(t *testing.T) {
// trim_prefix example:
// {"operations":[{"path":"model","mode":"trim_prefix","value":"openai/"}]}
input := []byte(`{"model":"openai/gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "trim_prefix",
"value": "openai/",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideTrimSuffix(t *testing.T) {
// trim_suffix example:
// {"operations":[{"path":"model","mode":"trim_suffix","value":"-latest"}]}
input := []byte(`{"model":"gpt-4-latest","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "trim_suffix",
"value": "-latest",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideTrimNoop(t *testing.T) {
// trim_prefix no-op example:
// {"operations":[{"path":"model","mode":"trim_prefix","value":"openai/"}]}
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "trim_prefix",
"value": "openai/",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideTrimRequiresValue(t *testing.T) {
// trim_prefix requires value example:
// {"operations":[{"path":"model","mode":"trim_prefix"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "trim_prefix",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideReplace(t *testing.T) {
// replace example:
// {"operations":[{"path":"model","mode":"replace","from":"openai/","to":""}]}
input := []byte(`{"model":"openai/gpt-4o-mini","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "replace",
"from": "openai/",
"to": "",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4o-mini","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideRegexReplace(t *testing.T) {
// regex_replace example:
// {"operations":[{"path":"model","mode":"regex_replace","from":"^gpt-","to":"openai/gpt-"}]}
input := []byte(`{"model":"gpt-4o-mini","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "regex_replace",
"from": "^gpt-",
"to": "openai/gpt-",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"openai/gpt-4o-mini","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideReplaceRequiresFrom(t *testing.T) {
// replace requires from example:
// {"operations":[{"path":"model","mode":"replace"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "replace",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideRegexReplaceRequiresPattern(t *testing.T) {
// regex_replace requires from(pattern) example:
// {"operations":[{"path":"model","mode":"regex_replace"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "regex_replace",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideDelete(t *testing.T) {
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "delete",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
var got map[string]interface{}
if err := json.Unmarshal(out, &got); err != nil {
t.Fatalf("failed to unmarshal output JSON: %v", err)
}
if _, exists := got["temperature"]; exists {
t.Fatalf("expected temperature to be deleted")
}
}
func TestApplyParamOverrideSet(t *testing.T) {
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.1}`, string(out))
}
func TestApplyParamOverrideSetKeepOrigin(t *testing.T) {
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"keep_origin": true,
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideMove(t *testing.T) {
input := []byte(`{"model":"gpt-4","meta":{"x":1}}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "move",
"from": "model",
"to": "meta.model",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"meta":{"x":1,"model":"gpt-4"}}`, string(out))
}
func TestApplyParamOverrideMoveMissingSource(t *testing.T) {
input := []byte(`{"meta":{"x":1}}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "move",
"from": "model",
"to": "meta.model",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverridePrependAppendString(t *testing.T) {
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "prepend",
"value": "openai/",
},
map[string]interface{}{
"path": "model",
"mode": "append",
"value": "-latest",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"openai/gpt-4-latest"}`, string(out))
}
func TestApplyParamOverridePrependAppendArray(t *testing.T) {
input := []byte(`{"arr":[1,2]}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "arr",
"mode": "prepend",
"value": 0,
},
map[string]interface{}{
"path": "arr",
"mode": "append",
"value": []interface{}{3, 4},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"arr":[0,1,2,3,4]}`, string(out))
}
func TestApplyParamOverrideAppendObjectMergeKeepOrigin(t *testing.T) {
input := []byte(`{"obj":{"a":1}}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "obj",
"mode": "append",
"keep_origin": true,
"value": map[string]interface{}{
"a": 2,
"b": 3,
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"obj":{"a":1,"b":3}}`, string(out))
}
func TestApplyParamOverrideAppendObjectMergeOverride(t *testing.T) {
input := []byte(`{"obj":{"a":1}}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "obj",
"mode": "append",
"value": map[string]interface{}{
"a": 2,
"b": 3,
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"obj":{"a":2,"b":3}}`, string(out))
}
func TestApplyParamOverrideConditionORDefault(t *testing.T) {
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "prefix",
"value": "gpt",
},
map[string]interface{}{
"path": "model",
"mode": "prefix",
"value": "claude",
},
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.1}`, string(out))
}
func TestApplyParamOverrideConditionAND(t *testing.T) {
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"logic": "AND",
"conditions": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "prefix",
"value": "gpt",
},
map[string]interface{}{
"path": "temperature",
"mode": "gt",
"value": 0.5,
},
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.1}`, string(out))
}
func TestApplyParamOverrideConditionInvert(t *testing.T) {
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "prefix",
"value": "gpt",
"invert": true,
},
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideConditionPassMissingKey(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "prefix",
"value": "gpt",
"pass_missing_key": true,
},
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func TestApplyParamOverrideConditionFromContext(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "prefix",
"value": "gpt",
},
},
},
},
}
ctx := map[string]interface{}{
"model": "gpt-4",
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func TestApplyParamOverrideNegativeIndexPath(t *testing.T) {
input := []byte(`{"arr":[{"model":"a"},{"model":"b"}]}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "arr.-1.model",
"mode": "set",
"value": "c",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"arr":[{"model":"a"},{"model":"c"}]}`, string(out))
}
func TestApplyParamOverrideRegexReplaceInvalidPattern(t *testing.T) {
// regex_replace invalid pattern example:
// {"operations":[{"path":"model","mode":"regex_replace","from":"(","to":"x"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "regex_replace",
"from": "(",
"to": "x",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideCopy(t *testing.T) {
// copy example:
// {"operations":[{"mode":"copy","from":"model","to":"original_model"}]}
input := []byte(`{"model":"gpt-4","temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "copy",
"from": "model",
"to": "original_model",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4","original_model":"gpt-4","temperature":0.7}`, string(out))
}
func TestApplyParamOverrideCopyMissingSource(t *testing.T) {
// copy missing source example:
// {"operations":[{"mode":"copy","from":"model","to":"original_model"}]}
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "copy",
"from": "model",
"to": "original_model",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideCopyRequiresFromTo(t *testing.T) {
// copy requires from/to example:
// {"operations":[{"mode":"copy"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "copy",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideEnsurePrefix(t *testing.T) {
// ensure_prefix example:
// {"operations":[{"path":"model","mode":"ensure_prefix","value":"openai/"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "ensure_prefix",
"value": "openai/",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"openai/gpt-4"}`, string(out))
}
func TestApplyParamOverrideEnsurePrefixNoop(t *testing.T) {
// ensure_prefix no-op example:
// {"operations":[{"path":"model","mode":"ensure_prefix","value":"openai/"}]}
input := []byte(`{"model":"openai/gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "ensure_prefix",
"value": "openai/",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"openai/gpt-4"}`, string(out))
}
func TestApplyParamOverrideEnsureSuffix(t *testing.T) {
// ensure_suffix example:
// {"operations":[{"path":"model","mode":"ensure_suffix","value":"-latest"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "ensure_suffix",
"value": "-latest",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4-latest"}`, string(out))
}
func TestApplyParamOverrideEnsureSuffixNoop(t *testing.T) {
// ensure_suffix no-op example:
// {"operations":[{"path":"model","mode":"ensure_suffix","value":"-latest"}]}
input := []byte(`{"model":"gpt-4-latest"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "ensure_suffix",
"value": "-latest",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4-latest"}`, string(out))
}
func TestApplyParamOverrideEnsureRequiresValue(t *testing.T) {
// ensure_prefix requires value example:
// {"operations":[{"path":"model","mode":"ensure_prefix"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "ensure_prefix",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideTrimSpace(t *testing.T) {
// trim_space example:
// {"operations":[{"path":"model","mode":"trim_space"}]}
input := []byte("{\"model\":\" gpt-4 \\n\"}")
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "trim_space",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4"}`, string(out))
}
func TestApplyParamOverrideToLower(t *testing.T) {
// to_lower example:
// {"operations":[{"path":"model","mode":"to_lower"}]}
input := []byte(`{"model":"GPT-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "to_lower",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"gpt-4"}`, string(out))
}
func TestApplyParamOverrideToUpper(t *testing.T) {
// to_upper example:
// {"operations":[{"path":"model","mode":"to_upper"}]}
input := []byte(`{"model":"gpt-4"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "model",
"mode": "to_upper",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"model":"GPT-4"}`, string(out))
}
func assertJSONEqual(t *testing.T, want, got string) {
t.Helper()
var wantObj interface{}
var gotObj interface{}
if err := json.Unmarshal([]byte(want), &wantObj); err != nil {
t.Fatalf("failed to unmarshal want JSON: %v", err)
}
if err := json.Unmarshal([]byte(got), &gotObj); err != nil {
t.Fatalf("failed to unmarshal got JSON: %v", err)
}
if !reflect.DeepEqual(wantObj, gotObj) {
t.Fatalf("json not equal\nwant: %s\ngot: %s", want, got)
}
}

View File

@@ -115,11 +115,16 @@ type RelayInfo struct {
SendResponseCount int
FinalPreConsumedQuota int // 最终预消耗的配额
IsClaudeBetaQuery bool // /v1/messages?beta=true
IsChannelTest bool // channel test request
PriceData types.PriceData
Request dto.Request
// RequestConversionChain records request format conversions in order, e.g.
// ["openai", "openai_responses"] or ["openai", "claude"].
RequestConversionChain []types.RelayFormat
ThinkingContentInfo
TokenCountMeta
*ClaudeConvertInfo
@@ -273,6 +278,7 @@ var streamSupportedChannels = map[int]bool{
constant.ChannelTypeZhipu_v4: true,
constant.ChannelTypeAli: true,
constant.ChannelTypeSubmodel: true,
constant.ChannelTypeCodex: true,
}
func GenRelayInfoWs(c *gin.Context, ws *websocket.Conn) *RelayInfo {
@@ -446,38 +452,83 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
}
func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Request, ws *websocket.Conn) (*RelayInfo, error) {
var info *RelayInfo
var err error
switch relayFormat {
case types.RelayFormatOpenAI:
return GenRelayInfoOpenAI(c, request), nil
info = GenRelayInfoOpenAI(c, request)
case types.RelayFormatOpenAIAudio:
return GenRelayInfoOpenAIAudio(c, request), nil
info = GenRelayInfoOpenAIAudio(c, request)
case types.RelayFormatOpenAIImage:
return GenRelayInfoImage(c, request), nil
info = GenRelayInfoImage(c, request)
case types.RelayFormatOpenAIRealtime:
return GenRelayInfoWs(c, ws), nil
info = GenRelayInfoWs(c, ws)
case types.RelayFormatClaude:
return GenRelayInfoClaude(c, request), nil
info = GenRelayInfoClaude(c, request)
case types.RelayFormatRerank:
if request, ok := request.(*dto.RerankRequest); ok {
return GenRelayInfoRerank(c, request), nil
info = GenRelayInfoRerank(c, request)
break
}
return nil, errors.New("request is not a RerankRequest")
err = errors.New("request is not a RerankRequest")
case types.RelayFormatGemini:
return GenRelayInfoGemini(c, request), nil
info = GenRelayInfoGemini(c, request)
case types.RelayFormatEmbedding:
return GenRelayInfoEmbedding(c, request), nil
info = GenRelayInfoEmbedding(c, request)
case types.RelayFormatOpenAIResponses:
if request, ok := request.(*dto.OpenAIResponsesRequest); ok {
return GenRelayInfoResponses(c, request), nil
info = GenRelayInfoResponses(c, request)
break
}
return nil, errors.New("request is not a OpenAIResponsesRequest")
err = errors.New("request is not a OpenAIResponsesRequest")
case types.RelayFormatTask:
return genBaseRelayInfo(c, nil), nil
info = genBaseRelayInfo(c, nil)
case types.RelayFormatMjProxy:
return genBaseRelayInfo(c, nil), nil
info = genBaseRelayInfo(c, nil)
default:
return nil, errors.New("invalid relay format")
err = errors.New("invalid relay format")
}
if err != nil {
return nil, err
}
if info == nil {
return nil, errors.New("failed to build relay info")
}
info.InitRequestConversionChain()
return info, nil
}
func (info *RelayInfo) InitRequestConversionChain() {
if info == nil {
return
}
if len(info.RequestConversionChain) > 0 {
return
}
if info.RelayFormat == "" {
return
}
info.RequestConversionChain = []types.RelayFormat{info.RelayFormat}
}
func (info *RelayInfo) AppendRequestConversion(format types.RelayFormat) {
if info == nil {
return
}
if format == "" {
return
}
if len(info.RequestConversionChain) == 0 {
info.RequestConversionChain = []types.RelayFormat{format}
return
}
last := info.RequestConversionChain[len(info.RequestConversionChain)-1]
if last == format {
return
}
info.RequestConversionChain = append(info.RequestConversionChain, format)
}
//func (info *RelayInfo) SetPromptTokens(promptTokens int) {

View File

@@ -0,0 +1,40 @@
package common
import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/types"
)
func GuessRelayFormatFromRequest(req any) (types.RelayFormat, bool) {
switch req.(type) {
case *dto.GeneralOpenAIRequest, dto.GeneralOpenAIRequest:
return types.RelayFormatOpenAI, true
case *dto.OpenAIResponsesRequest, dto.OpenAIResponsesRequest:
return types.RelayFormatOpenAIResponses, true
case *dto.ClaudeRequest, dto.ClaudeRequest:
return types.RelayFormatClaude, true
case *dto.GeminiChatRequest, dto.GeminiChatRequest:
return types.RelayFormatGemini, true
case *dto.EmbeddingRequest, dto.EmbeddingRequest:
return types.RelayFormatEmbedding, true
case *dto.RerankRequest, dto.RerankRequest:
return types.RelayFormatRerank, true
case *dto.ImageRequest, dto.ImageRequest:
return types.RelayFormatOpenAIImage, true
case *dto.AudioRequest, dto.AudioRequest:
return types.RelayFormatOpenAIAudio, true
default:
return "", false
}
}
func AppendRequestConversionFromRequest(info *RelayInfo, req any) {
if info == nil {
return
}
format, ok := GuessRelayFormatFromRequest(req)
if !ok {
return
}
info.AppendRequestConversion(format)
}

View File

@@ -14,10 +14,12 @@ import (
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
"github.com/shopspring/decimal"
@@ -72,9 +74,32 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
}
adaptor.Init(info)
passThroughGlobal := model_setting.GetGlobalSettings().PassThroughRequestEnabled
if info.RelayMode == relayconstant.RelayModeChatCompletions &&
!passThroughGlobal &&
!info.ChannelSetting.PassThroughBodyEnabled &&
shouldChatCompletionsViaResponses(info) {
applySystemPromptIfNeeded(c, info, request)
usage, newApiErr := chatCompletionsViaResponses(c, info, adaptor, request)
if newApiErr != nil {
return newApiErr
}
var containAudioTokens = usage.CompletionTokenDetails.AudioTokens > 0 || usage.PromptTokensDetails.AudioTokens > 0
var containsAudioRatios = ratio_setting.ContainsAudioRatio(info.OriginModelName) || ratio_setting.ContainsAudioCompletionRatio(info.OriginModelName)
if containAudioTokens && containsAudioRatios {
service.PostAudioConsumeQuota(c, info, usage, "")
} else {
postConsumeQuota(c, info, usage)
}
return nil
}
var requestBody io.Reader
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
if passThroughGlobal || info.ChannelSetting.PassThroughBodyEnabled {
body, err := common.GetRequestBody(c)
if err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
@@ -88,6 +113,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
if info.ChannelSetting.SystemPrompt != "" {
// 如果有系统提示,则将其添加到请求中
@@ -181,22 +207,35 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
return newApiErr
}
if usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0 {
var containAudioTokens = usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0
var containsAudioRatios = ratio_setting.ContainsAudioRatio(info.OriginModelName) || ratio_setting.ContainsAudioCompletionRatio(info.OriginModelName)
if containAudioTokens && containsAudioRatios {
service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "")
} else {
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
}
return nil
}
func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) {
func shouldChatCompletionsViaResponses(info *relaycommon.RelayInfo) bool {
if info == nil {
return false
}
if info.RelayMode != relayconstant.RelayModeChatCompletions {
return false
}
return service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.OriginModelName)
}
func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent ...string) {
if usage == nil {
usage = &dto.Usage{
PromptTokens: relayInfo.GetEstimatePromptTokens(),
CompletionTokens: 0,
TotalTokens: relayInfo.GetEstimatePromptTokens(),
}
extraContent += "(可能是请求出错)"
extraContent = append(extraContent, "上游无计费信息")
}
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
promptTokens := usage.PromptTokens
@@ -246,8 +285,8 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
dWebSearchQuota = decimal.NewFromFloat(webSearchPrice).
Mul(decimal.NewFromInt(int64(webSearchTool.CallCount))).
Div(decimal.NewFromInt(1000)).Mul(dGroupRatio).Mul(dQuotaPerUnit)
extraContent += fmt.Sprintf("Web Search 调用 %d 次,上下文大小 %s调用花费 %s",
webSearchTool.CallCount, webSearchTool.SearchContextSize, dWebSearchQuota.String())
extraContent = append(extraContent, fmt.Sprintf("Web Search 调用 %d 次,上下文大小 %s调用花费 %s",
webSearchTool.CallCount, webSearchTool.SearchContextSize, dWebSearchQuota.String()))
}
} else if strings.HasSuffix(modelName, "search-preview") {
// search-preview 模型不支持 response api
@@ -258,8 +297,8 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
webSearchPrice = operation_setting.GetWebSearchPricePerThousand(modelName, searchContextSize)
dWebSearchQuota = decimal.NewFromFloat(webSearchPrice).
Div(decimal.NewFromInt(1000)).Mul(dGroupRatio).Mul(dQuotaPerUnit)
extraContent += fmt.Sprintf("Web Search 调用 1 次,上下文大小 %s调用花费 %s",
searchContextSize, dWebSearchQuota.String())
extraContent = append(extraContent, fmt.Sprintf("Web Search 调用 1 次,上下文大小 %s调用花费 %s",
searchContextSize, dWebSearchQuota.String()))
}
// claude web search tool 计费
var dClaudeWebSearchQuota decimal.Decimal
@@ -269,8 +308,8 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
claudeWebSearchPrice = operation_setting.GetClaudeWebSearchPricePerThousand()
dClaudeWebSearchQuota = decimal.NewFromFloat(claudeWebSearchPrice).
Div(decimal.NewFromInt(1000)).Mul(dGroupRatio).Mul(dQuotaPerUnit).Mul(decimal.NewFromInt(int64(claudeWebSearchCallCount)))
extraContent += fmt.Sprintf("Claude Web Search 调用 %d 次,调用花费 %s",
claudeWebSearchCallCount, dClaudeWebSearchQuota.String())
extraContent = append(extraContent, fmt.Sprintf("Claude Web Search 调用 %d 次,调用花费 %s",
claudeWebSearchCallCount, dClaudeWebSearchQuota.String()))
}
// file search tool 计费
var dFileSearchQuota decimal.Decimal
@@ -281,8 +320,8 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
dFileSearchQuota = decimal.NewFromFloat(fileSearchPrice).
Mul(decimal.NewFromInt(int64(fileSearchTool.CallCount))).
Div(decimal.NewFromInt(1000)).Mul(dGroupRatio).Mul(dQuotaPerUnit)
extraContent += fmt.Sprintf("File Search 调用 %d 次,调用花费 %s",
fileSearchTool.CallCount, dFileSearchQuota.String())
extraContent = append(extraContent, fmt.Sprintf("File Search 调用 %d 次,调用花费 %s",
fileSearchTool.CallCount, dFileSearchQuota.String()))
}
}
var dImageGenerationCallQuota decimal.Decimal
@@ -290,13 +329,14 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
if ctx.GetBool("image_generation_call") {
imageGenerationCallPrice = operation_setting.GetGPTImage1PriceOnceCall(ctx.GetString("image_generation_call_quality"), ctx.GetString("image_generation_call_size"))
dImageGenerationCallQuota = decimal.NewFromFloat(imageGenerationCallPrice).Mul(dGroupRatio).Mul(dQuotaPerUnit)
extraContent += fmt.Sprintf("Image Generation Call 花费 %s", dImageGenerationCallQuota.String())
extraContent = append(extraContent, fmt.Sprintf("Image Generation Call 花费 %s", dImageGenerationCallQuota.String()))
}
var quotaCalculateDecimal decimal.Decimal
var audioInputQuota decimal.Decimal
var audioInputPrice float64
isClaudeUsageSemantic := relayInfo.ChannelType == constant.ChannelTypeAnthropic
if !relayInfo.PriceData.UsePrice {
baseTokens := dPromptTokens
// 减去 cached tokens
@@ -304,14 +344,14 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
// OpenAI/OpenRouter 等 API 的 prompt_tokens 包含缓存 tokens需要减去
var cachedTokensWithRatio decimal.Decimal
if !dCacheTokens.IsZero() {
if relayInfo.ChannelType != constant.ChannelTypeAnthropic {
if !isClaudeUsageSemantic {
baseTokens = baseTokens.Sub(dCacheTokens)
}
cachedTokensWithRatio = dCacheTokens.Mul(dCacheRatio)
}
var dCachedCreationTokensWithRatio decimal.Decimal
if !dCachedCreationTokens.IsZero() {
if relayInfo.ChannelType != constant.ChannelTypeAnthropic {
if !isClaudeUsageSemantic {
baseTokens = baseTokens.Sub(dCachedCreationTokens)
}
dCachedCreationTokensWithRatio = dCachedCreationTokens.Mul(dCachedCreationRatio)
@@ -331,7 +371,7 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
// 重新计算 base tokens
baseTokens = baseTokens.Sub(dAudioTokens)
audioInputQuota = decimal.NewFromFloat(audioInputPrice).Div(decimal.NewFromInt(1000000)).Mul(dAudioTokens).Mul(dGroupRatio).Mul(dQuotaPerUnit)
extraContent += fmt.Sprintf("Audio Input 花费 %s", audioInputQuota.String())
extraContent = append(extraContent, fmt.Sprintf("Audio Input 花费 %s", audioInputQuota.String()))
}
}
promptQuota := baseTokens.Add(cachedTokensWithRatio).
@@ -356,17 +396,25 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
// 添加 image generation call 计费
quotaCalculateDecimal = quotaCalculateDecimal.Add(dImageGenerationCallQuota)
if len(relayInfo.PriceData.OtherRatios) > 0 {
for key, otherRatio := range relayInfo.PriceData.OtherRatios {
dOtherRatio := decimal.NewFromFloat(otherRatio)
quotaCalculateDecimal = quotaCalculateDecimal.Mul(dOtherRatio)
extraContent = append(extraContent, fmt.Sprintf("其他倍率 %s: %f", key, otherRatio))
}
}
quota := int(quotaCalculateDecimal.Round(0).IntPart())
totalTokens := promptTokens + completionTokens
var logContent string
//var logContent string
// 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("(可能是上游超时)")
extraContent = append(extraContent, "上游没有返回计费信息,无法扣费(可能是上游超时)")
logger.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, relayInfo.FinalPreConsumedQuota))
} else {
@@ -405,16 +453,19 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
logModel := modelName
if strings.HasPrefix(logModel, "gpt-4-gizmo") {
logModel = "gpt-4-gizmo-*"
logContent += fmt.Sprintf("模型 %s", modelName)
extraContent = append(extraContent, fmt.Sprintf("模型 %s", modelName))
}
if strings.HasPrefix(logModel, "gpt-4o-gizmo") {
logModel = "gpt-4o-gizmo-*"
logContent += fmt.Sprintf("模型 %s", modelName)
}
if extraContent != "" {
logContent += ", " + extraContent
extraContent = append(extraContent, fmt.Sprintf("模型 %s", modelName))
}
logContent := strings.Join(extraContent, ", ")
other := service.GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, cacheTokens, cacheRatio, modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
// For chat-based calls to the Claude model, tagging is required. Using Claude's rendering logs, the two approaches handle input rendering differently.
if isClaudeUsageSemantic {
other["claude"] = true
other["usage_semantic"] = "anthropic"
}
if imageTokens != 0 {
other["image"] = true
other["image_ratio"] = imageRatio

View File

@@ -45,6 +45,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := json.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
@@ -82,6 +83,6 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
}
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
return nil
}

View File

@@ -149,6 +149,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
@@ -193,7 +194,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
return openaiErr
}
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
return nil
}
@@ -292,6 +293,6 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
return openaiErr
}
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
return nil
}

View File

@@ -57,6 +57,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed)
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
switch convertedRequest.(type) {
case *bytes.Buffer:
@@ -124,12 +125,18 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
quality = "hd"
}
var logContent string
var logContent []string
if len(request.Size) > 0 {
logContent = fmt.Sprintf("大小 %s, 品质 %s, 张数 %d", request.Size, quality, request.N)
logContent = append(logContent, fmt.Sprintf("大小 %s", request.Size))
}
if len(quality) > 0 {
logContent = append(logContent, fmt.Sprintf("品质 %s", quality))
}
if request.N > 0 {
logContent = append(logContent, fmt.Sprintf("生成数量 %d", request.N))
}
postConsumeQuota(c, info, usage.(*dto.Usage), logContent)
postConsumeQuota(c, info, usage.(*dto.Usage), logContent...)
return nil
}

View File

@@ -11,6 +11,7 @@ import (
"github.com/QuantumNous/new-api/relay/channel/baidu_v2"
"github.com/QuantumNous/new-api/relay/channel/claude"
"github.com/QuantumNous/new-api/relay/channel/cloudflare"
"github.com/QuantumNous/new-api/relay/channel/codex"
"github.com/QuantumNous/new-api/relay/channel/cohere"
"github.com/QuantumNous/new-api/relay/channel/coze"
"github.com/QuantumNous/new-api/relay/channel/deepseek"
@@ -117,6 +118,8 @@ func GetAdaptor(apiType int) channel.Adaptor {
return &minimax.Adaptor{}
case constant.APITypeReplicate:
return &replicate.Adaptor{}
case constant.APITypeCodex:
return &codex.Adaptor{}
}
return nil
}
@@ -148,7 +151,7 @@ func GetTaskAdaptor(platform constant.TaskPlatform) channel.TaskAdaptor {
return &taskvertex.TaskAdaptor{}
case constant.ChannelTypeVidu:
return &taskVidu.TaskAdaptor{}
case constant.ChannelTypeDoubaoVideo:
case constant.ChannelTypeDoubaoVideo, constant.ChannelTypeVolcEngine:
return &taskdoubao.TaskAdaptor{}
case constant.ChannelTypeSora, constant.ChannelTypeOpenAI:
return &tasksora.TaskAdaptor{}

View File

@@ -150,6 +150,14 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.
}
}
// 处理 auto 分组:从 context 获取实际选中的分组
// 当使用 auto 分组时Distribute 中间件会将实际选中的分组存储在 ContextKeyAutoGroup 中
if autoGroup, exists := common.GetContextKey(c, constant.ContextKeyAutoGroup); exists {
if groupStr, ok := autoGroup.(string); ok && groupStr != "" {
info.UsingGroup = groupStr
}
}
// 预扣
groupRatio := ratio_setting.GetGroupRatio(info.UsingGroup)
var ratio float64

View File

@@ -53,6 +53,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
@@ -95,6 +96,6 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
}
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
return nil
}

View File

@@ -53,6 +53,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
}
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := common.Marshal(convertedRequest)
if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
@@ -107,7 +108,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if strings.HasPrefix(info.OriginModelName, "gpt-4o-audio") {
service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "")
} else {
postConsumeQuota(c, info, usage.(*dto.Usage), "")
postConsumeQuota(c, info, usage.(*dto.Usage))
}
return nil
}

View File

@@ -93,6 +93,10 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.POST("/2fa/enable", controller.Enable2FA)
selfRoute.POST("/2fa/disable", controller.Disable2FA)
selfRoute.POST("/2fa/backup_codes", controller.RegenerateBackupCodes)
// Check-in routes
selfRoute.GET("/checkin", controller.GetCheckinStatus)
selfRoute.POST("/checkin", middleware.TurnstileCheck(), controller.DoCheckin)
}
adminRoute := userRoute.Group("/")
@@ -152,6 +156,16 @@ func SetApiRouter(router *gin.Engine) {
channelRoute.POST("/fix", controller.FixChannelsAbilities)
channelRoute.GET("/fetch_models/:id", controller.FetchUpstreamModels)
channelRoute.POST("/fetch_models", controller.FetchModels)
channelRoute.POST("/codex/oauth/start", controller.StartCodexOAuth)
channelRoute.POST("/codex/oauth/complete", controller.CompleteCodexOAuth)
channelRoute.POST("/:id/codex/oauth/start", controller.StartCodexOAuthForChannel)
channelRoute.POST("/:id/codex/oauth/complete", controller.CompleteCodexOAuthForChannel)
channelRoute.POST("/:id/codex/refresh", controller.RefreshCodexChannelCredential)
channelRoute.GET("/:id/codex/usage", controller.GetCodexChannelUsage)
channelRoute.POST("/ollama/pull", controller.OllamaPullModel)
channelRoute.POST("/ollama/pull/stream", controller.OllamaPullModelStream)
channelRoute.DELETE("/ollama/delete", controller.OllamaDeleteModel)
channelRoute.GET("/ollama/version/:id", controller.OllamaVersion)
channelRoute.POST("/batch/tag", controller.BatchSetChannelTag)
channelRoute.GET("/tag/models", controller.GetTagModels)
channelRoute.POST("/copy/:id", controller.CopyChannel)
@@ -256,5 +270,31 @@ func SetApiRouter(router *gin.Engine) {
modelsRoute.PUT("/", controller.UpdateModelMeta)
modelsRoute.DELETE("/:id", controller.DeleteModelMeta)
}
// Deployments (model deployment management)
deploymentsRoute := apiRouter.Group("/deployments")
deploymentsRoute.Use(middleware.AdminAuth())
{
deploymentsRoute.GET("/settings", controller.GetModelDeploymentSettings)
deploymentsRoute.POST("/settings/test-connection", controller.TestIoNetConnection)
deploymentsRoute.GET("/", controller.GetAllDeployments)
deploymentsRoute.GET("/search", controller.SearchDeployments)
deploymentsRoute.POST("/test-connection", controller.TestIoNetConnection)
deploymentsRoute.GET("/hardware-types", controller.GetHardwareTypes)
deploymentsRoute.GET("/locations", controller.GetLocations)
deploymentsRoute.GET("/available-replicas", controller.GetAvailableReplicas)
deploymentsRoute.POST("/price-estimation", controller.GetPriceEstimation)
deploymentsRoute.GET("/check-name", controller.CheckClusterNameAvailability)
deploymentsRoute.POST("/", controller.CreateDeployment)
deploymentsRoute.GET("/:id", controller.GetDeployment)
deploymentsRoute.GET("/:id/logs", controller.GetDeploymentLogs)
deploymentsRoute.GET("/:id/containers", controller.ListDeploymentContainers)
deploymentsRoute.GET("/:id/containers/:container_id", controller.GetContainerDetails)
deploymentsRoute.PUT("/:id", controller.UpdateDeployment)
deploymentsRoute.PUT("/:id/name", controller.UpdateDeploymentName)
deploymentsRoute.POST("/:id/extend", controller.ExtendDeployment)
deploymentsRoute.DELETE("/:id", controller.DeleteDeployment)
}
}
}

View File

@@ -57,9 +57,12 @@ func ShouldDisableChannel(channelType int, err *types.NewAPIError) bool {
if types.IsSkipRetryError(err) {
return false
}
if err.StatusCode == http.StatusUnauthorized {
if operation_setting.ShouldDisableByStatusCode(err.StatusCode) {
return true
}
//if err.StatusCode == http.StatusUnauthorized {
// return true
//}
if err.StatusCode == http.StatusForbidden {
switch channelType {
case constant.ChannelTypeGemini:

View File

@@ -0,0 +1,104 @@
package service
import (
"context"
"errors"
"fmt"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
)
type CodexCredentialRefreshOptions struct {
ResetCaches bool
}
type CodexOAuthKey struct {
IDToken string `json:"id_token,omitempty"`
AccessToken string `json:"access_token,omitempty"`
RefreshToken string `json:"refresh_token,omitempty"`
AccountID string `json:"account_id,omitempty"`
LastRefresh string `json:"last_refresh,omitempty"`
Email string `json:"email,omitempty"`
Type string `json:"type,omitempty"`
Expired string `json:"expired,omitempty"`
}
func parseCodexOAuthKey(raw string) (*CodexOAuthKey, error) {
if strings.TrimSpace(raw) == "" {
return nil, errors.New("codex channel: empty oauth key")
}
var key CodexOAuthKey
if err := common.Unmarshal([]byte(raw), &key); err != nil {
return nil, errors.New("codex channel: invalid oauth key json")
}
return &key, nil
}
func RefreshCodexChannelCredential(ctx context.Context, channelID int, opts CodexCredentialRefreshOptions) (*CodexOAuthKey, *model.Channel, error) {
ch, err := model.GetChannelById(channelID, true)
if err != nil {
return nil, nil, err
}
if ch == nil {
return nil, nil, fmt.Errorf("channel not found")
}
if ch.Type != constant.ChannelTypeCodex {
return nil, nil, fmt.Errorf("channel type is not Codex")
}
oauthKey, err := parseCodexOAuthKey(strings.TrimSpace(ch.Key))
if err != nil {
return nil, nil, err
}
if strings.TrimSpace(oauthKey.RefreshToken) == "" {
return nil, nil, fmt.Errorf("codex channel: refresh_token is required to refresh credential")
}
refreshCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
res, err := RefreshCodexOAuthToken(refreshCtx, oauthKey.RefreshToken)
if err != nil {
return nil, nil, err
}
oauthKey.AccessToken = res.AccessToken
oauthKey.RefreshToken = res.RefreshToken
oauthKey.LastRefresh = time.Now().Format(time.RFC3339)
oauthKey.Expired = res.ExpiresAt.Format(time.RFC3339)
if strings.TrimSpace(oauthKey.Type) == "" {
oauthKey.Type = "codex"
}
if strings.TrimSpace(oauthKey.AccountID) == "" {
if accountID, ok := ExtractCodexAccountIDFromJWT(oauthKey.AccessToken); ok {
oauthKey.AccountID = accountID
}
}
if strings.TrimSpace(oauthKey.Email) == "" {
if email, ok := ExtractEmailFromJWT(oauthKey.AccessToken); ok {
oauthKey.Email = email
}
}
encoded, err := common.Marshal(oauthKey)
if err != nil {
return nil, nil, err
}
if err := model.DB.Model(&model.Channel{}).Where("id = ?", ch.Id).Update("key", string(encoded)).Error; err != nil {
return nil, nil, err
}
if opts.ResetCaches {
model.InitChannelCache()
ResetProxyClientCache()
}
return oauthKey, ch, nil
}

View File

@@ -0,0 +1,140 @@
package service
import (
"context"
"fmt"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/bytedance/gopkg/util/gopool"
)
const (
codexCredentialRefreshTickInterval = 10 * time.Minute
codexCredentialRefreshThreshold = 24 * time.Hour
codexCredentialRefreshBatchSize = 200
codexCredentialRefreshTimeout = 15 * time.Second
)
var (
codexCredentialRefreshOnce sync.Once
codexCredentialRefreshRunning atomic.Bool
)
func StartCodexCredentialAutoRefreshTask() {
codexCredentialRefreshOnce.Do(func() {
if !common.IsMasterNode {
return
}
gopool.Go(func() {
logger.LogInfo(context.Background(), fmt.Sprintf("codex credential auto-refresh task started: tick=%s threshold=%s", codexCredentialRefreshTickInterval, codexCredentialRefreshThreshold))
ticker := time.NewTicker(codexCredentialRefreshTickInterval)
defer ticker.Stop()
runCodexCredentialAutoRefreshOnce()
for range ticker.C {
runCodexCredentialAutoRefreshOnce()
}
})
})
}
func runCodexCredentialAutoRefreshOnce() {
if !codexCredentialRefreshRunning.CompareAndSwap(false, true) {
return
}
defer codexCredentialRefreshRunning.Store(false)
ctx := context.Background()
now := time.Now()
var refreshed int
var scanned int
offset := 0
for {
var channels []*model.Channel
err := model.DB.
Select("id", "name", "key", "status", "channel_info").
Where("type = ? AND status = 1", constant.ChannelTypeCodex).
Order("id asc").
Limit(codexCredentialRefreshBatchSize).
Offset(offset).
Find(&channels).Error
if err != nil {
logger.LogError(ctx, fmt.Sprintf("codex credential auto-refresh: query channels failed: %v", err))
return
}
if len(channels) == 0 {
break
}
offset += codexCredentialRefreshBatchSize
for _, ch := range channels {
if ch == nil {
continue
}
scanned++
if ch.ChannelInfo.IsMultiKey {
continue
}
rawKey := strings.TrimSpace(ch.Key)
if rawKey == "" {
continue
}
oauthKey, err := parseCodexOAuthKey(rawKey)
if err != nil {
continue
}
refreshToken := strings.TrimSpace(oauthKey.RefreshToken)
if refreshToken == "" {
continue
}
expiredAtRaw := strings.TrimSpace(oauthKey.Expired)
expiredAt, err := time.Parse(time.RFC3339, expiredAtRaw)
if err == nil && !expiredAt.IsZero() && expiredAt.Sub(now) > codexCredentialRefreshThreshold {
continue
}
refreshCtx, cancel := context.WithTimeout(ctx, codexCredentialRefreshTimeout)
newKey, _, err := RefreshCodexChannelCredential(refreshCtx, ch.Id, CodexCredentialRefreshOptions{ResetCaches: false})
cancel()
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("codex credential auto-refresh: channel_id=%d name=%s refresh failed: %v", ch.Id, ch.Name, err))
continue
}
refreshed++
logger.LogInfo(ctx, fmt.Sprintf("codex credential auto-refresh: channel_id=%d name=%s refreshed, expires_at=%s", ch.Id, ch.Name, newKey.Expired))
}
}
if refreshed > 0 {
func() {
defer func() {
if r := recover(); r != nil {
logger.LogWarn(ctx, fmt.Sprintf("codex credential auto-refresh: InitChannelCache panic: %v", r))
}
}()
model.InitChannelCache()
}()
ResetProxyClientCache()
}
if common.DebugEnabled {
logger.LogDebug(ctx, "codex credential auto-refresh: scanned=%d refreshed=%d", scanned, refreshed)
}
}

288
service/codex_oauth.go Normal file
View File

@@ -0,0 +1,288 @@
package service
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/url"
"strings"
"time"
)
const (
codexOAuthClientID = "app_EMoamEEZ73f0CkXaXp7hrann"
codexOAuthAuthorizeURL = "https://auth.openai.com/oauth/authorize"
codexOAuthTokenURL = "https://auth.openai.com/oauth/token"
codexOAuthRedirectURI = "http://localhost:1455/auth/callback"
codexOAuthScope = "openid profile email offline_access"
codexJWTClaimPath = "https://api.openai.com/auth"
defaultHTTPTimeout = 20 * time.Second
)
type CodexOAuthTokenResult struct {
AccessToken string
RefreshToken string
ExpiresAt time.Time
}
type CodexOAuthAuthorizationFlow struct {
State string
Verifier string
Challenge string
AuthorizeURL string
}
func RefreshCodexOAuthToken(ctx context.Context, refreshToken string) (*CodexOAuthTokenResult, error) {
client := &http.Client{Timeout: defaultHTTPTimeout}
return refreshCodexOAuthToken(ctx, client, codexOAuthTokenURL, codexOAuthClientID, refreshToken)
}
func ExchangeCodexAuthorizationCode(ctx context.Context, code string, verifier string) (*CodexOAuthTokenResult, error) {
client := &http.Client{Timeout: defaultHTTPTimeout}
return exchangeCodexAuthorizationCode(ctx, client, codexOAuthTokenURL, codexOAuthClientID, code, verifier, codexOAuthRedirectURI)
}
func CreateCodexOAuthAuthorizationFlow() (*CodexOAuthAuthorizationFlow, error) {
state, err := createStateHex(16)
if err != nil {
return nil, err
}
verifier, challenge, err := generatePKCEPair()
if err != nil {
return nil, err
}
u, err := buildCodexAuthorizeURL(state, challenge)
if err != nil {
return nil, err
}
return &CodexOAuthAuthorizationFlow{
State: state,
Verifier: verifier,
Challenge: challenge,
AuthorizeURL: u,
}, nil
}
func refreshCodexOAuthToken(
ctx context.Context,
client *http.Client,
tokenURL string,
clientID string,
refreshToken string,
) (*CodexOAuthTokenResult, error) {
rt := strings.TrimSpace(refreshToken)
if rt == "" {
return nil, errors.New("empty refresh_token")
}
form := url.Values{}
form.Set("grant_type", "refresh_token")
form.Set("refresh_token", rt)
form.Set("client_id", clientID)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(form.Encode()))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
resp, err := client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
var payload struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ExpiresIn int `json:"expires_in"`
}
if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("codex oauth refresh failed: status=%d", resp.StatusCode)
}
if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 {
return nil, errors.New("codex oauth refresh response missing fields")
}
return &CodexOAuthTokenResult{
AccessToken: strings.TrimSpace(payload.AccessToken),
RefreshToken: strings.TrimSpace(payload.RefreshToken),
ExpiresAt: time.Now().Add(time.Duration(payload.ExpiresIn) * time.Second),
}, nil
}
func exchangeCodexAuthorizationCode(
ctx context.Context,
client *http.Client,
tokenURL string,
clientID string,
code string,
verifier string,
redirectURI string,
) (*CodexOAuthTokenResult, error) {
c := strings.TrimSpace(code)
v := strings.TrimSpace(verifier)
if c == "" {
return nil, errors.New("empty authorization code")
}
if v == "" {
return nil, errors.New("empty code_verifier")
}
form := url.Values{}
form.Set("grant_type", "authorization_code")
form.Set("client_id", clientID)
form.Set("code", c)
form.Set("code_verifier", v)
form.Set("redirect_uri", redirectURI)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(form.Encode()))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
resp, err := client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
var payload struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ExpiresIn int `json:"expires_in"`
}
if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("codex oauth code exchange failed: status=%d", resp.StatusCode)
}
if strings.TrimSpace(payload.AccessToken) == "" || strings.TrimSpace(payload.RefreshToken) == "" || payload.ExpiresIn <= 0 {
return nil, errors.New("codex oauth token response missing fields")
}
return &CodexOAuthTokenResult{
AccessToken: strings.TrimSpace(payload.AccessToken),
RefreshToken: strings.TrimSpace(payload.RefreshToken),
ExpiresAt: time.Now().Add(time.Duration(payload.ExpiresIn) * time.Second),
}, nil
}
func buildCodexAuthorizeURL(state string, challenge string) (string, error) {
u, err := url.Parse(codexOAuthAuthorizeURL)
if err != nil {
return "", err
}
q := u.Query()
q.Set("response_type", "code")
q.Set("client_id", codexOAuthClientID)
q.Set("redirect_uri", codexOAuthRedirectURI)
q.Set("scope", codexOAuthScope)
q.Set("code_challenge", challenge)
q.Set("code_challenge_method", "S256")
q.Set("state", state)
q.Set("id_token_add_organizations", "true")
q.Set("codex_cli_simplified_flow", "true")
q.Set("originator", "codex_cli_rs")
u.RawQuery = q.Encode()
return u.String(), nil
}
func createStateHex(nBytes int) (string, error) {
if nBytes <= 0 {
return "", errors.New("invalid state bytes length")
}
b := make([]byte, nBytes)
if _, err := rand.Read(b); err != nil {
return "", err
}
return fmt.Sprintf("%x", b), nil
}
func generatePKCEPair() (verifier string, challenge string, err error) {
b := make([]byte, 32)
if _, err := rand.Read(b); err != nil {
return "", "", err
}
verifier = base64.RawURLEncoding.EncodeToString(b)
sum := sha256.Sum256([]byte(verifier))
challenge = base64.RawURLEncoding.EncodeToString(sum[:])
return verifier, challenge, nil
}
func ExtractCodexAccountIDFromJWT(token string) (string, bool) {
claims, ok := decodeJWTClaims(token)
if !ok {
return "", false
}
raw, ok := claims[codexJWTClaimPath]
if !ok {
return "", false
}
obj, ok := raw.(map[string]any)
if !ok {
return "", false
}
v, ok := obj["chatgpt_account_id"]
if !ok {
return "", false
}
s, ok := v.(string)
if !ok {
return "", false
}
s = strings.TrimSpace(s)
if s == "" {
return "", false
}
return s, true
}
func ExtractEmailFromJWT(token string) (string, bool) {
claims, ok := decodeJWTClaims(token)
if !ok {
return "", false
}
v, ok := claims["email"]
if !ok {
return "", false
}
s, ok := v.(string)
if !ok {
return "", false
}
s = strings.TrimSpace(s)
if s == "" {
return "", false
}
return s, true
}
func decodeJWTClaims(token string) (map[string]any, bool) {
parts := strings.Split(token, ".")
if len(parts) != 3 {
return nil, false
}
payloadRaw, err := base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return nil, false
}
var claims map[string]any
if err := json.Unmarshal(payloadRaw, &claims); err != nil {
return nil, false
}
return claims, true
}

View File

@@ -0,0 +1,56 @@
package service
import (
"context"
"fmt"
"io"
"net/http"
"strings"
)
func FetchCodexWhamUsage(
ctx context.Context,
client *http.Client,
baseURL string,
accessToken string,
accountID string,
) (statusCode int, body []byte, err error) {
if client == nil {
return 0, nil, fmt.Errorf("nil http client")
}
bu := strings.TrimRight(strings.TrimSpace(baseURL), "/")
if bu == "" {
return 0, nil, fmt.Errorf("empty baseURL")
}
at := strings.TrimSpace(accessToken)
aid := strings.TrimSpace(accountID)
if at == "" {
return 0, nil, fmt.Errorf("empty accessToken")
}
if aid == "" {
return 0, nil, fmt.Errorf("empty accountID")
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, bu+"/backend-api/wham/usage", nil)
if err != nil {
return 0, nil, err
}
req.Header.Set("Authorization", "Bearer "+at)
req.Header.Set("chatgpt-account-id", aid)
req.Header.Set("Accept", "application/json")
if req.Header.Get("originator") == "" {
req.Header.Set("originator", "codex_cli_rs")
}
resp, err := client.Do(req)
if err != nil {
return 0, nil, err
}
defer resp.Body.Close()
body, err = io.ReadAll(resp.Body)
if err != nil {
return resp.StatusCode, nil, err
}
return resp.StatusCode, body, nil
}

View File

@@ -674,20 +674,21 @@ func GeminiToOpenAIRequest(geminiRequest *dto.GeminiChatRequest, info *relaycomm
var tools []dto.ToolCallRequest
for _, tool := range geminiRequest.GetTools() {
if tool.FunctionDeclarations != nil {
// 将 Gemini 的 FunctionDeclarations 转换为 OpenAI 的 ToolCallRequest
functionDeclarations, ok := tool.FunctionDeclarations.([]dto.FunctionRequest)
if ok {
for _, function := range functionDeclarations {
openAITool := dto.ToolCallRequest{
Type: "function",
Function: dto.FunctionRequest{
Name: function.Name,
Description: function.Description,
Parameters: function.Parameters,
},
}
tools = append(tools, openAITool)
functionDeclarations, err := common.Any2Type[[]dto.FunctionRequest](tool.FunctionDeclarations)
if err != nil {
common.SysError(fmt.Sprintf("failed to parse gemini function declarations: %v (type=%T)", err, tool.FunctionDeclarations))
continue
}
for _, function := range functionDeclarations {
openAITool := dto.ToolCallRequest{
Type: "function",
Function: dto.FunctionRequest{
Name: function.Name,
Description: function.Description,
Parameters: function.Parameters,
},
}
tools = append(tools, openAITool)
}
}
}

View File

@@ -57,4 +57,5 @@ func IOCopyBytesGracefully(c *gin.Context, src *http.Response, data []byte) {
if err != nil {
logger.LogError(c, fmt.Sprintf("failed to copy response body: %s", err.Error()))
}
c.Writer.Flush()
}

View File

@@ -38,6 +38,10 @@ func InitHttpClient() {
MaxIdleConns: common.RelayMaxIdleConns,
MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost,
ForceAttemptHTTP2: true,
Proxy: http.ProxyFromEnvironment, // Support HTTP_PROXY, HTTPS_PROXY, NO_PROXY env vars
}
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
}
if common.RelayTimeout == 0 {
@@ -81,6 +85,9 @@ func ResetProxyClientCache() {
// NewProxyHttpClient 创建支持代理的 HTTP 客户端
func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
if proxyURL == "" {
if client := GetHttpClient(); client != nil {
return client, nil
}
return http.DefaultClient, nil
}
@@ -98,13 +105,17 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
switch parsedURL.Scheme {
case "http", "https":
transport := &http.Transport{
MaxIdleConns: common.RelayMaxIdleConns,
MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost,
ForceAttemptHTTP2: true,
Proxy: http.ProxyURL(parsedURL),
}
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
}
client := &http.Client{
Transport: &http.Transport{
MaxIdleConns: common.RelayMaxIdleConns,
MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost,
ForceAttemptHTTP2: true,
Proxy: http.ProxyURL(parsedURL),
},
Transport: transport,
CheckRedirect: checkRedirect,
}
client.Timeout = time.Duration(common.RelayTimeout) * time.Second
@@ -133,17 +144,19 @@ func NewProxyHttpClient(proxyURL string) (*http.Client, error) {
return nil, err
}
client := &http.Client{
Transport: &http.Transport{
MaxIdleConns: common.RelayMaxIdleConns,
MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost,
ForceAttemptHTTP2: true,
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return dialer.Dial(network, addr)
},
transport := &http.Transport{
MaxIdleConns: common.RelayMaxIdleConns,
MaxIdleConnsPerHost: common.RelayMaxIdleConnsPerHost,
ForceAttemptHTTP2: true,
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return dialer.Dial(network, addr)
},
CheckRedirect: checkRedirect,
}
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
}
client := &http.Client{Transport: transport, CheckRedirect: checkRedirect}
client.Timeout = time.Duration(common.RelayTimeout) * time.Second
proxyClientLock.Lock()
proxyClients[proxyURL] = client

View File

@@ -70,9 +70,38 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m
other["admin_info"] = adminInfo
appendRequestPath(ctx, relayInfo, other)
appendRequestConversionChain(relayInfo, other)
return other
}
func appendRequestConversionChain(relayInfo *relaycommon.RelayInfo, other map[string]interface{}) {
if relayInfo == nil || other == nil {
return
}
if len(relayInfo.RequestConversionChain) == 0 {
return
}
chain := make([]string, 0, len(relayInfo.RequestConversionChain))
for _, f := range relayInfo.RequestConversionChain {
switch f {
case types.RelayFormatOpenAI:
chain = append(chain, "OpenAI Compatible")
case types.RelayFormatClaude:
chain = append(chain, "Claude Messages")
case types.RelayFormatGemini:
chain = append(chain, "Google Gemini")
case types.RelayFormatOpenAIResponses:
chain = append(chain, "OpenAI Responses")
default:
chain = append(chain, string(f))
}
}
if len(chain) == 0 {
return
}
other["request_conversion"] = chain
}
func GenerateWssOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage, modelRatio, groupRatio, completionRatio, audioRatio, audioCompletionRatio, modelPrice, userGroupRatio float64) map[string]interface{} {
info := GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, 0, 0.0, modelPrice, userGroupRatio)
info["ws"] = true

View File

@@ -0,0 +1,18 @@
package service
import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/service/openaicompat"
)
func ChatCompletionsRequestToResponsesRequest(req *dto.GeneralOpenAIRequest) (*dto.OpenAIResponsesRequest, error) {
return openaicompat.ChatCompletionsRequestToResponsesRequest(req)
}
func ResponsesResponseToChatCompletionsResponse(resp *dto.OpenAIResponsesResponse, id string) (*dto.OpenAITextResponse, *dto.Usage, error) {
return openaicompat.ResponsesResponseToChatCompletionsResponse(resp, id)
}
func ExtractOutputTextFromResponses(resp *dto.OpenAIResponsesResponse) string {
return openaicompat.ExtractOutputTextFromResponses(resp)
}

View File

@@ -0,0 +1,14 @@
package service
import (
"github.com/QuantumNous/new-api/service/openaicompat"
"github.com/QuantumNous/new-api/setting/model_setting"
)
func ShouldChatCompletionsUseResponsesPolicy(policy model_setting.ChatCompletionsToResponsesPolicy, channelID int, model string) bool {
return openaicompat.ShouldChatCompletionsUseResponsesPolicy(policy, channelID, model)
}
func ShouldChatCompletionsUseResponsesGlobal(channelID int, model string) bool {
return openaicompat.ShouldChatCompletionsUseResponsesGlobal(channelID, model)
}

Some files were not shown because too many files have changed in this diff Show More