From 4885d224a0bf71cc674b2ab7848f624e1ffefb09 Mon Sep 17 00:00:00 2001 From: zijiren <84728412+zijiren233@users.noreply.github.com> Date: Mon, 10 Mar 2025 15:26:23 +0800 Subject: [PATCH] feat: aiproxy feishu notify (#5445) * feat: get balance notify * feat: balance notify * feat: auto test bannel channel notify * feat: notify lock limit * feat: notify throttle * fix: claude think token * feat: azure support and mode string * feat: feishu notify * fix: banned notify * fix: lint * feat: notify note * feat: notify note * fix: do not expire banned * fix: default disable auto update channel balance * feat: beyond threshold notify * feat: notify error detial * fix: add request script * feat: sync service error notify * feat: baidu error code handler * feat: baidu error code handler * fix: unsupported mode info * refactor: relay error handler * feat: auto reload channel and model config * feat: ollama support * fix: lint * chore: error log message * chore: sync notify throttle * fix: redis script key expire * feat: no permission channel notify * feat: get channel with error rate priority * feat: retry log filed * fix: lint * fix: group and token update sql double where id condition * feat: async batch process usage * fix: rate limit overlimit count * fix: rate limit count * fix: no permission error message * fix: ban expiry * feat: mem model monitor * feat: mem monitor cleanup * feat: guess model type cache * fix: get priority * fix: do not unban when error rate reduce * chore: docker image add curl cli * fix: model not exist status code * fix: error code * fix: record consume notify and clean log batch size * chore: batch record * fix: clean log error notify throttle * chore: rename notify env name --- service/aiproxy/Dockerfile | 2 +- service/aiproxy/common/balance/balance.go | 4 + service/aiproxy/common/config/config.go | 13 +- service/aiproxy/common/consume/consume.go | 12 +- service/aiproxy/common/image/image.go | 1 - service/aiproxy/common/init.go | 32 -- service/aiproxy/common/notify/feishu.go | 185 ++++++ service/aiproxy/common/notify/feishu_test.go | 21 + service/aiproxy/common/notify/notify.go | 64 +++ service/aiproxy/common/notify/std.go | 48 ++ service/aiproxy/common/rpmlimit/rate-limit.go | 86 ++- service/aiproxy/common/trylock/lock.go | 74 +++ service/aiproxy/common/trylock/lock_test.go | 24 + service/aiproxy/controller/channel-billing.go | 47 +- service/aiproxy/controller/channel-test.go | 52 +- service/aiproxy/controller/dashboard.go | 3 +- service/aiproxy/controller/import.go | 8 +- service/aiproxy/controller/monitor.go | 9 + service/aiproxy/controller/relay.go | 219 +++++--- service/aiproxy/main.go | 47 +- service/aiproxy/middleware/auth.go | 12 +- service/aiproxy/middleware/distributor.go | 120 ++-- service/aiproxy/middleware/utils.go | 22 +- service/aiproxy/model/cache.go | 9 +- service/aiproxy/model/channel.go | 113 ++-- service/aiproxy/model/configkey.go | 2 - service/aiproxy/model/group.go | 7 +- service/aiproxy/model/log.go | 63 ++- service/aiproxy/model/modelconfig.go | 18 +- service/aiproxy/model/option.go | 6 +- service/aiproxy/model/owner.go | 1 - service/aiproxy/model/token.go | 7 +- service/aiproxy/model/utils.go | 175 +++++- service/aiproxy/monitor/memmodel.go | 337 +++++++++++ service/aiproxy/monitor/model.go | 528 ++++++++++-------- service/aiproxy/relay/adaptor/ali/adaptor.go | 8 +- .../aiproxy/relay/adaptor/ali/embeddings.go | 2 +- service/aiproxy/relay/adaptor/ali/image.go | 11 +- service/aiproxy/relay/adaptor/ali/rerank.go | 4 + .../relay/adaptor/anthropic/adaptor.go | 4 +- .../aiproxy/relay/adaptor/anthropic/main.go | 36 +- .../aiproxy/relay/adaptor/aws/claude/main.go | 3 +- .../aiproxy/relay/adaptor/aws/llama3/main.go | 3 +- .../aiproxy/relay/adaptor/azure/constants.go | 6 +- service/aiproxy/relay/adaptor/azure/key.go | 4 +- .../aiproxy/relay/adaptor/baidu/adaptor.go | 13 +- .../aiproxy/relay/adaptor/baidu/embeddings.go | 5 +- service/aiproxy/relay/adaptor/baidu/error.go | 48 ++ service/aiproxy/relay/adaptor/baidu/image.go | 3 +- service/aiproxy/relay/adaptor/baidu/main.go | 8 +- service/aiproxy/relay/adaptor/baidu/rerank.go | 3 +- .../aiproxy/relay/adaptor/baiduv2/adaptor.go | 6 +- .../aiproxy/relay/adaptor/cohere/adaptor.go | 4 +- service/aiproxy/relay/adaptor/cohere/main.go | 30 +- service/aiproxy/relay/adaptor/coze/adaptor.go | 13 +- service/aiproxy/relay/adaptor/coze/main.go | 28 +- .../aiproxy/relay/adaptor/doc2x/adaptor.go | 6 +- service/aiproxy/relay/adaptor/doubao/main.go | 5 + .../aiproxy/relay/adaptor/doubaoaudio/main.go | 10 +- .../aiproxy/relay/adaptor/gemini/adaptor.go | 5 +- .../relay/adaptor/gemini/embeddings.go | 4 + service/aiproxy/relay/adaptor/gemini/main.go | 8 + service/aiproxy/relay/adaptor/interface.go | 3 + .../relay/adaptor/lingyiwanwu/adaptor.go | 5 + .../aiproxy/relay/adaptor/minimax/adaptor.go | 5 + service/aiproxy/relay/adaptor/minimax/tts.go | 4 + .../aiproxy/relay/adaptor/moonshot/balance.go | 2 +- .../aiproxy/relay/adaptor/ollama/adaptor.go | 15 +- service/aiproxy/relay/adaptor/ollama/error.go | 24 + service/aiproxy/relay/adaptor/ollama/main.go | 104 ++-- .../aiproxy/relay/adaptor/openai/adaptor.go | 7 +- .../relay/{utils => adaptor/openai}/error.go | 29 +- service/aiproxy/relay/adaptor/openai/image.go | 4 + service/aiproxy/relay/adaptor/openai/main.go | 63 ++- .../relay/adaptor/openai/moderations.go | 4 + .../aiproxy/relay/adaptor/openai/rerank.go | 4 + service/aiproxy/relay/adaptor/openai/stt.go | 4 + service/aiproxy/relay/adaptor/openai/tts.go | 4 + .../relay/adaptor/siliconflow/balance.go | 2 +- .../aiproxy/relay/adaptor/stepfun/adaptor.go | 5 + .../aiproxy/relay/adaptor/tencent/adaptor.go | 5 + .../relay/adaptor/vertexai/claude/adapter.go | 9 +- .../relay/adaptor/vertexai/gemini/adapter.go | 5 +- .../aiproxy/relay/adaptor/vertexai/token.go | 4 +- .../aiproxy/relay/adaptor/xunfei/adaptor.go | 5 + .../aiproxy/relay/adaptor/zhipu/adaptor.go | 5 + service/aiproxy/relay/controller/consume.go | 11 +- service/aiproxy/relay/controller/dohelper.go | 35 +- service/aiproxy/relay/controller/handle.go | 8 +- service/aiproxy/relay/meta/meta.go | 13 +- service/aiproxy/relay/model/misc.go | 8 +- service/aiproxy/relay/relaymode/define.go | 37 +- service/aiproxy/relay/utils/testreq.go | 5 +- service/aiproxy/relay/utils/utils.go | 7 +- service/aiproxy/router/api.go | 4 +- service/aiproxy/router/relay.go | 3 +- 96 files changed, 2285 insertions(+), 818 deletions(-) delete mode 100644 service/aiproxy/common/init.go create mode 100644 service/aiproxy/common/notify/feishu.go create mode 100644 service/aiproxy/common/notify/feishu_test.go create mode 100644 service/aiproxy/common/notify/notify.go create mode 100644 service/aiproxy/common/notify/std.go create mode 100644 service/aiproxy/common/trylock/lock.go create mode 100644 service/aiproxy/common/trylock/lock_test.go create mode 100644 service/aiproxy/monitor/memmodel.go create mode 100644 service/aiproxy/relay/adaptor/baidu/error.go create mode 100644 service/aiproxy/relay/adaptor/ollama/error.go rename service/aiproxy/relay/{utils => adaptor/openai}/error.go (79%) diff --git a/service/aiproxy/Dockerfile b/service/aiproxy/Dockerfile index 8602bc081..69b9899ab 100644 --- a/service/aiproxy/Dockerfile +++ b/service/aiproxy/Dockerfile @@ -9,7 +9,7 @@ ENV FFPROBE_ENABLED=true EXPOSE 3000 -RUN apk add --no-cache ca-certificates tzdata ffmpeg && \ +RUN apk add --no-cache ca-certificates tzdata ffmpeg curl && \ rm -rf /var/cache/apk/* ENTRYPOINT ["/aiproxy"] diff --git a/service/aiproxy/common/balance/balance.go b/service/aiproxy/common/balance/balance.go index 934448d70..553f33b29 100644 --- a/service/aiproxy/common/balance/balance.go +++ b/service/aiproxy/common/balance/balance.go @@ -22,3 +22,7 @@ var ( func MockGetGroupRemainBalance(ctx context.Context, group model.GroupCache) (float64, PostGroupConsumer, error) { return mock.GetGroupRemainBalance(ctx, group) } + +func GetGroupRemainBalance(ctx context.Context, group model.GroupCache) (float64, PostGroupConsumer, error) { + return Default.GetGroupRemainBalance(ctx, group) +} diff --git a/service/aiproxy/common/config/config.go b/service/aiproxy/common/config/config.go index 29fd13163..09b53d2f1 100644 --- a/service/aiproxy/common/config/config.go +++ b/service/aiproxy/common/config/config.go @@ -30,12 +30,13 @@ var ( logDetailResponseBodyMaxSize int64 = 128 * 1024 // 128KB logDetailStorageHours int64 = 3 * 24 // 3 days internalToken atomic.Value + notifyNote atomic.Value ) var ( retryTimes atomic.Int64 enableModelErrorAutoBan atomic.Bool - modelErrorAutoBanRate = math.Float64bits(0.5) + modelErrorAutoBanRate = math.Float64bits(0.3) timeoutWithModelType atomic.Value disableModelConfig = env.Bool("DISABLE_MODEL_CONFIG", false) ) @@ -59,6 +60,7 @@ func init() { geminiSafetySetting.Store("BLOCK_NONE") billingEnabled.Store(true) internalToken.Store(os.Getenv("INTERNAL_TOKEN")) + notifyNote.Store(os.Getenv("NOTIFY_NOTE")) } func GetDisableModelConfig() bool { @@ -231,3 +233,12 @@ func SetInternalToken(token string) { token = env.String("INTERNAL_TOKEN", token) internalToken.Store(token) } + +func GetNotifyNote() string { + return notifyNote.Load().(string) +} + +func SetNotifyNote(note string) { + note = env.String("NOTIFY_NOTE", note) + notifyNote.Store(note) +} diff --git a/service/aiproxy/common/consume/consume.go b/service/aiproxy/common/consume/consume.go index 02dd6822c..ba778a646 100644 --- a/service/aiproxy/common/consume/consume.go +++ b/service/aiproxy/common/consume/consume.go @@ -3,8 +3,10 @@ package consume import ( "context" "sync" + "time" "github.com/labring/sealos/service/aiproxy/common/balance" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" @@ -27,6 +29,7 @@ func AsyncConsume( outputPrice float64, content string, ip string, + retryTimes int, requestDetail *model.RequestDetail, ) { if meta.IsChannelTest { @@ -51,6 +54,7 @@ func AsyncConsume( outputPrice, content, ip, + retryTimes, requestDetail, ) } @@ -65,6 +69,7 @@ func Consume( outputPrice float64, content string, ip string, + retryTimes int, requestDetail *model.RequestDetail, ) { if meta.IsChannelTest { @@ -75,9 +80,10 @@ func Consume( amount = consumeAmount(ctx, amount, postGroupConsumer, meta) - err := recordConsume(meta, code, usage, inputPrice, outputPrice, content, ip, requestDetail, amount) + err := recordConsume(meta, code, usage, inputPrice, outputPrice, content, ip, requestDetail, amount, retryTimes) if err != nil { log.Error("error batch record consume: " + err.Error()) + notify.ErrorThrottle("recordConsume", time.Minute, "record consume failed", err.Error()) } } @@ -155,6 +161,7 @@ func recordConsume( ip string, requestDetail *model.RequestDetail, amount float64, + retryTimes int, ) error { promptTokens := 0 completionTokens := 0 @@ -184,8 +191,9 @@ func recordConsume( outputPrice, meta.Endpoint, content, - meta.Mode, + int(meta.Mode), ip, + retryTimes, requestDetail, ) } diff --git a/service/aiproxy/common/image/image.go b/service/aiproxy/common/image/image.go index 9ec808c0e..8eac5f595 100644 --- a/service/aiproxy/common/image/image.go +++ b/service/aiproxy/common/image/image.go @@ -20,7 +20,6 @@ import ( "strings" "github.com/labring/sealos/service/aiproxy/common" - // import webp decoder _ "golang.org/x/image/webp" ) diff --git a/service/aiproxy/common/init.go b/service/aiproxy/common/init.go deleted file mode 100644 index b17fc1c8a..000000000 --- a/service/aiproxy/common/init.go +++ /dev/null @@ -1,32 +0,0 @@ -package common - -import ( - "flag" - "os" - "path/filepath" - - log "github.com/sirupsen/logrus" -) - -var ( - Port = flag.Int("port", 3000, "the listening port") - LogDir = flag.String("log-dir", "", "specify the log directory") -) - -func Init() { - flag.Parse() - - if *LogDir != "" { - var err error - *LogDir, err = filepath.Abs(*LogDir) - if err != nil { - log.Fatal(err) - } - if _, err := os.Stat(*LogDir); os.IsNotExist(err) { - err = os.Mkdir(*LogDir, 0o777) - if err != nil { - log.Fatal(err) - } - } - } -} diff --git a/service/aiproxy/common/notify/feishu.go b/service/aiproxy/common/notify/feishu.go new file mode 100644 index 000000000..12aaffe3a --- /dev/null +++ b/service/aiproxy/common/notify/feishu.go @@ -0,0 +1,185 @@ +package notify + +import ( + "bytes" + "context" + "errors" + "net/http" + "time" + + "github.com/bytedance/sonic" + "github.com/labring/sealos/service/aiproxy/common/config" + "github.com/labring/sealos/service/aiproxy/common/trylock" +) + +type FeishuNotifier struct { + wh string +} + +func level2Color(level Level) string { + switch level { + case LevelInfo: + return FeishuColorGreen + case LevelError: + return FeishuColorRed + case LevelWarn: + return FeishuColorOrange + default: + return FeishuColorGreen + } +} + +func (f *FeishuNotifier) Notify(level Level, title, message string) { + stdNotifier.Notify(level, title, message) + go func() { + _ = PostToFeiShuv2(context.Background(), level2Color(level), title, message, f.wh) + }() +} + +func (f *FeishuNotifier) NotifyThrottle(level Level, key string, expiration time.Duration, title, message string) { + if trylock.Lock(key, expiration) { + stdNotifier.Notify(level, title, message) + go func() { + _ = PostToFeiShuv2(context.Background(), level2Color(level), title, message, f.wh) + }() + } +} + +func NewFeishuNotify(wh string) Notifier { + return &FeishuNotifier{ + wh: wh, + } +} + +type FSMessagev2 struct { + MsgType string `json:"msg_type"` + Email string `json:"email"` + Card Cards `json:"card"` +} + +type Cards struct { + Config Conf `json:"config"` + Elements []Element `json:"elements"` + Header Headers `json:"header"` +} + +type Conf struct { + WideScreenMode bool `json:"wide_screen_mode"` + EnableForward bool `json:"enable_forward"` +} + +type Te struct { + Content string `json:"content"` + Tag string `json:"tag"` +} + +type Element struct { + Tag string `json:"tag"` + Text Te `json:"text"` + Content string `json:"content"` + Elements []Element `json:"elements"` +} + +type Titles struct { + Content string `json:"content"` + Tag string `json:"tag"` +} + +type Headers struct { + Title Titles `json:"title"` + Template string `json:"template"` +} + +type TenantAccessMeg struct { + AppID string `json:"app_id"` + AppSecret string `json:"app_secret"` +} + +type TenantAccessResp struct { + Code int `json:"code"` + Msg string `json:"msg"` + TenantAccessToken string `json:"tenant_access_token"` +} + +type FeiShuv2Resp struct { + StatusCode int `json:"StatusCode"` + StatusMessage string `json:"StatusMessage"` + Code int `json:"code"` + Data any `json:"data"` + Msg string `json:"msg"` +} + +const ( + FeishuColorOrange = "orange" + FeishuColorGreen = "green" + FeishuColorRed = "red" +) + +func PostToFeiShuv2(ctx context.Context, color, title, text, wh string) error { + if wh == "" { + return errors.New("feishu webhook url is empty") + } + note := config.GetNotifyNote() + if note == "" { + note = "AI Proxy" + } + u := FSMessagev2{ + MsgType: "interactive", + Card: Cards{ + Config: Conf{ + WideScreenMode: true, + EnableForward: true, + }, + Header: Headers{ + Title: Titles{ + Content: title, + Tag: "plain_text", + }, + Template: color, + }, + Elements: []Element{ + { + Tag: "div", + Text: Te{ + Content: text, + Tag: "lark_md", + }, + }, + { + Tag: "hr", + }, + { + Tag: "note", + Elements: []Element{ + { + Content: note, + Tag: "lark_md", + }, + }, + }, + }, + }, + } + + data, err := sonic.ConfigDefault.Marshal(u) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, wh, bytes.NewReader(data)) + if err != nil { + return err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + feishuResp := FeiShuv2Resp{} + if err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&feishuResp); err != nil { + return err + } + if feishuResp.Code != 0 { + return errors.New(feishuResp.Msg) + } + return nil +} diff --git a/service/aiproxy/common/notify/feishu_test.go b/service/aiproxy/common/notify/feishu_test.go new file mode 100644 index 000000000..b63076d8a --- /dev/null +++ b/service/aiproxy/common/notify/feishu_test.go @@ -0,0 +1,21 @@ +package notify_test + +import ( + "context" + "os" + "testing" + + "github.com/labring/sealos/service/aiproxy/common/notify" +) + +func TestPostToFeiShuv2(t *testing.T) { + err := notify.PostToFeiShuv2( + context.Background(), + notify.FeishuColorRed, + "Error", + "Error Message", + os.Getenv("FEISHU_WEBHOOK")) + if err != nil { + t.Error(err) + } +} diff --git a/service/aiproxy/common/notify/notify.go b/service/aiproxy/common/notify/notify.go new file mode 100644 index 000000000..23df50010 --- /dev/null +++ b/service/aiproxy/common/notify/notify.go @@ -0,0 +1,64 @@ +package notify + +import ( + "fmt" + "time" +) + +type Level string + +const ( + LevelInfo Level = "info" + LevelWarn Level = "warn" + LevelError Level = "error" +) + +type Notifier interface { + Notify(level Level, title, message string) + NotifyThrottle(level Level, key string, expiration time.Duration, title, message string) +} + +var ( + stdNotifier Notifier = &StdNotifier{} + defaultNotifier Notifier = stdNotifier +) + +func SetDefaultNotifier(notifier Notifier) { + defaultNotifier = notifier +} + +func Notify(level Level, title, message string) { + defaultNotifier.Notify(level, title, message) +} + +func Info(title, message string) { + defaultNotifier.Notify(LevelInfo, title, message) +} + +func Warn(title, message string) { + defaultNotifier.Notify(LevelWarn, title, message) +} + +func Error(title, message string) { + defaultNotifier.Notify(LevelError, title, message) +} + +func limitKey(level Level, key string) string { + return fmt.Sprintf("notifylimit:%s:%s", level, key) +} + +func NotifyThrottle(level Level, key string, expiration time.Duration, title, message string) { + defaultNotifier.NotifyThrottle(level, limitKey(level, key), expiration, title, message) +} + +func InfoThrottle(key string, expiration time.Duration, title, message string) { + defaultNotifier.NotifyThrottle(LevelInfo, limitKey(LevelInfo, key), expiration, title, message) +} + +func WarnThrottle(key string, expiration time.Duration, title, message string) { + defaultNotifier.NotifyThrottle(LevelWarn, limitKey(LevelWarn, key), expiration, title, message) +} + +func ErrorThrottle(key string, expiration time.Duration, title, message string) { + defaultNotifier.NotifyThrottle(LevelError, limitKey(LevelError, key), expiration, title, message) +} diff --git a/service/aiproxy/common/notify/std.go b/service/aiproxy/common/notify/std.go new file mode 100644 index 000000000..65040f838 --- /dev/null +++ b/service/aiproxy/common/notify/std.go @@ -0,0 +1,48 @@ +package notify + +import ( + "time" + + "github.com/labring/sealos/service/aiproxy/common/config" + "github.com/labring/sealos/service/aiproxy/common/trylock" + log "github.com/sirupsen/logrus" +) + +type StdNotifier struct{} + +var ( + infoLogrus = log.WithField("notify", "std") + warnLogrus = log.WithField("notify", "std") + errorLogrus = log.WithField("notify", "std") +) + +func (n *StdNotifier) Notify(level Level, title, message string) { + note := config.GetNotifyNote() + switch level { + case LevelInfo: + logrus := infoLogrus.WithField("title", title) + if note != "" { + logrus = logrus.WithField("note", note) + } + logrus.Info(message) + case LevelWarn: + logrus := warnLogrus.WithField("title", title) + if note != "" { + logrus = logrus.WithField("note", note) + } + logrus.Warn(message) + case LevelError: + logrus := errorLogrus.WithField("title", title) + if note != "" { + logrus = logrus.WithField("note", note) + } + logrus.Error(message) + } +} + +func (n *StdNotifier) NotifyThrottle(level Level, key string, expiration time.Duration, title, message string) { + if !trylock.MemLock(key, expiration) { + return + } + n.Notify(level, title, message) +} diff --git a/service/aiproxy/common/rpmlimit/rate-limit.go b/service/aiproxy/common/rpmlimit/rate-limit.go index 32b88baa0..32668438e 100644 --- a/service/aiproxy/common/rpmlimit/rate-limit.go +++ b/service/aiproxy/common/rpmlimit/rate-limit.go @@ -2,7 +2,10 @@ package rpmlimit import ( "context" + "errors" "fmt" + "strconv" + "strings" "time" "github.com/labring/sealos/service/aiproxy/common" @@ -13,33 +16,54 @@ var inMemoryRateLimiter InMemoryRateLimiter const ( groupModelRPMKey = "group_model_rpm:%s:%s" + overLimitRPMKey = "over_limit_rpm:%s:%s" ) var pushRequestScript = ` local key = KEYS[1] +local over_limit_key = KEYS[2] +local window = tonumber(ARGV[1]) +local current_time = tonumber(ARGV[2]) +local max_requests = tonumber(ARGV[3]) +local cutoff = current_time - window + +redis.call('ZREMRANGEBYSCORE', key, '-inf', cutoff) +redis.call('ZREMRANGEBYSCORE', over_limit_key, '-inf', cutoff) +local count = redis.call('ZCOUNT', key, cutoff, current_time) +local over_limit_count = redis.call('ZCOUNT', over_limit_key, cutoff, current_time) + +if count <= max_requests then + redis.call('ZADD', key, current_time, current_time) + redis.call('PEXPIRE', key, window / 1000) + count = count + 1 +else + redis.call('ZADD', over_limit_key, current_time, current_time) + redis.call('PEXPIRE', over_limit_key, window / 1000) + over_limit_count = over_limit_count + 1 +end + +return string.format("%d:%d", count, over_limit_count) +` + +var getRequestCountScript = ` +local pattern = KEYS[1] +local over_limit_pattern = KEYS[2] local window = tonumber(ARGV[1]) local current_time = tonumber(ARGV[2]) local cutoff = current_time - window -redis.call('ZREMRANGEBYSCORE', key, '-inf', cutoff) -redis.call('ZADD', key, current_time, current_time) -redis.call('PEXPIRE', key, window) -return redis.call('ZCOUNT', key, cutoff, current_time) -` - -var getRequestCountScript = ` -local pattern = ARGV[1] -local window = tonumber(ARGV[2]) -local current_time = tonumber(ARGV[3]) -local cutoff = current_time - window - -local keys = redis.call('KEYS', pattern) local total = 0 +local keys = redis.call('KEYS', pattern) for _, key in ipairs(keys) do redis.call('ZREMRANGEBYSCORE', key, '-inf', cutoff) - local count = redis.call('ZCOUNT', key, cutoff, current_time) - total = total + count + total = total + redis.call('ZCOUNT', key, cutoff, current_time) +end + +local over_limit_keys = redis.call('KEYS', over_limit_pattern) +for _, key in ipairs(over_limit_keys) do + redis.call('ZREMRANGEBYSCORE', key, '-inf', cutoff) + total = total + redis.call('ZCOUNT', key, cutoff, current_time) end return total @@ -51,22 +75,26 @@ func GetRPM(ctx context.Context, group, model string) (int64, error) { } var pattern string + var overLimitPattern string if group == "" && model == "" { pattern = "group_model_rpm:*:*" + overLimitPattern = "over_limit_rpm:*:*" } else if group == "" { pattern = "group_model_rpm:*:" + model + overLimitPattern = "over_limit_rpm:*:" + model } else if model == "" { pattern = fmt.Sprintf("group_model_rpm:%s:*", group) + overLimitPattern = fmt.Sprintf("over_limit_rpm:%s:*", group) } else { pattern = fmt.Sprintf("group_model_rpm:%s:%s", group, model) + overLimitPattern = fmt.Sprintf("over_limit_rpm:%s:%s", group, model) } rdb := common.RDB result, err := rdb.Eval( ctx, getRequestCountScript, - []string{}, - pattern, + []string{pattern, overLimitPattern}, time.Minute.Microseconds(), time.Now().UnixMicro(), ).Int64() @@ -77,27 +105,41 @@ func GetRPM(ctx context.Context, group, model string) (int64, error) { } func redisRateLimitRequest(ctx context.Context, group, model string, maxRequestNum int64, duration time.Duration) (bool, error) { - result, err := PushRequest(ctx, group, model, duration) + result, _, err := PushRequest(ctx, group, model, maxRequestNum, duration) if err != nil { return false, err } return result <= maxRequestNum, nil } -func PushRequest(ctx context.Context, group, model string, duration time.Duration) (int64, error) { +func PushRequest(ctx context.Context, group, model string, maxRequestNum int64, duration time.Duration) (int64, int64, error) { result, err := common.RDB.Eval( ctx, pushRequestScript, []string{ fmt.Sprintf(groupModelRPMKey, group, model), + fmt.Sprintf(overLimitRPMKey, group, model), }, duration.Microseconds(), time.Now().UnixMicro(), - ).Int64() + maxRequestNum, + ).Text() if err != nil { - return 0, err + return 0, 0, err } - return result, nil + count, overLimitCount, ok := strings.Cut(result, ":") + if !ok { + return 0, 0, errors.New("invalid result") + } + countInt, err := strconv.ParseInt(count, 10, 64) + if err != nil { + return 0, 0, err + } + overLimitCountInt, err := strconv.ParseInt(overLimitCount, 10, 64) + if err != nil { + return 0, 0, err + } + return countInt, overLimitCountInt, nil } func RateLimit(ctx context.Context, group, model string, maxRequestNum int64, duration time.Duration) (bool, error) { diff --git a/service/aiproxy/common/trylock/lock.go b/service/aiproxy/common/trylock/lock.go new file mode 100644 index 000000000..5bad5d072 --- /dev/null +++ b/service/aiproxy/common/trylock/lock.go @@ -0,0 +1,74 @@ +package trylock + +import ( + "context" + "sync" + "time" + + "github.com/labring/sealos/service/aiproxy/common" + log "github.com/sirupsen/logrus" +) + +var memRecord = sync.Map{} + +func init() { + go cleanMemLock() +} + +func cleanMemLock() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + for now := range ticker.C { + memRecord.Range(func(key, value any) bool { + if exp, ok := value.(time.Time); ok { + if now.After(exp) { + memRecord.Delete(key) + } + } else { + memRecord.Delete(key) + } + return true + }) + } +} + +func MemLock(key string, expiration time.Duration) bool { + now := time.Now() + newExpiration := now.Add(expiration) + + for { + actual, loaded := memRecord.LoadOrStore(key, newExpiration) + if !loaded { + return true + } + oldExpiration, ok := actual.(time.Time) + if !ok { + memRecord.Delete(key) + continue + } + if now.After(oldExpiration) { + if memRecord.CompareAndSwap(key, actual, newExpiration) { + return true + } + continue + } + return false + } +} + +func Lock(key string, expiration time.Duration) bool { + if !common.RedisEnabled { + return MemLock(key, expiration) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result, err := common.RDB.SetNX(ctx, key, true, expiration).Result() + if err != nil { + if MemLock("lockerror", time.Second*3) { + log.Errorf("try notify error: %v", err) + } + return MemLock(key, expiration) + } + return result +} diff --git a/service/aiproxy/common/trylock/lock_test.go b/service/aiproxy/common/trylock/lock_test.go new file mode 100644 index 000000000..808d565f7 --- /dev/null +++ b/service/aiproxy/common/trylock/lock_test.go @@ -0,0 +1,24 @@ +package trylock_test + +import ( + "testing" + "time" + + "github.com/labring/sealos/service/aiproxy/common/trylock" +) + +func TestMemLock(t *testing.T) { + if !trylock.MemLock("", time.Second) { + t.Error("Expected true, Got false") + } + if trylock.MemLock("", time.Second) { + t.Error("Expected false, Got true") + } + if trylock.MemLock("", time.Second) { + t.Error("Expected false, Got true") + } + time.Sleep(time.Second) + if !trylock.MemLock("", time.Second) { + t.Error("Expected true, Got false") + } +} diff --git a/service/aiproxy/controller/channel-billing.go b/service/aiproxy/controller/channel-billing.go index cb6cf9d05..eac975a3a 100644 --- a/service/aiproxy/controller/channel-billing.go +++ b/service/aiproxy/controller/channel-billing.go @@ -5,10 +5,12 @@ import ( "fmt" "net/http" "strconv" + "sync" "time" "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common/balance" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/adaptor" @@ -22,20 +24,23 @@ import ( func updateChannelBalance(channel *model.Channel) (float64, error) { adaptorI, ok := channeltype.GetAdaptor(channel.Type) if !ok { - return 0, fmt.Errorf("invalid channel type: %d", channel.Type) + return 0, fmt.Errorf("invalid channel type: %d, channel: %s(%d)", channel.Type, channel.Name, channel.ID) } if getBalance, ok := adaptorI.(adaptor.Balancer); ok { balance, err := getBalance.GetBalance(channel) - if err != nil { - return 0, err + if err != nil && !errors.Is(err, adaptor.ErrGetBalanceNotImplemented) { + return 0, fmt.Errorf("failed to get channel[%d] %s(%d) balance: %s", channel.Type, channel.Name, channel.ID, err.Error()) } - err = channel.UpdateBalance(balance) - if err != nil { - log.Errorf("failed to update channel %s(%d) balance: %s", channel.Name, channel.ID, err.Error()) + if err := channel.UpdateBalance(balance); err != nil { + return 0, fmt.Errorf("failed to update channel [%d] %s(%d) balance: %s", channel.Type, channel.Name, channel.ID, err.Error()) + } + if !errors.Is(err, adaptor.ErrGetBalanceNotImplemented) && + balance < channel.GetBalanceThreshold() { + return 0, fmt.Errorf("channel[%d] %s(%d) balance is less than threshold: %f", channel.Type, channel.Name, channel.ID, balance) } return balance, nil } - return 0, fmt.Errorf("channel type %d does not support get balance", channel.Type) + return 0, nil } func UpdateChannelBalance(c *gin.Context) { @@ -57,6 +62,7 @@ func UpdateChannelBalance(c *gin.Context) { } balance, err := updateChannelBalance(channel) if err != nil { + notify.Error(fmt.Sprintf("check channel[%d] %s(%d) balance error", channel.Type, channel.Name, channel.ID), err.Error()) c.JSON(http.StatusOK, middleware.APIResponse{ Success: false, Message: err.Error(), @@ -75,12 +81,27 @@ func updateAllChannelsBalance() error { if err != nil { return err } + + var wg sync.WaitGroup + semaphore := make(chan struct{}, 10) + for _, channel := range channels { - _, err := updateChannelBalance(channel) - if err != nil { + if !channel.EnabledAutoBalanceCheck { continue } + wg.Add(1) + semaphore <- struct{}{} + go func(ch *model.Channel) { + defer wg.Done() + defer func() { <-semaphore }() + _, err := updateChannelBalance(ch) + if err != nil { + notify.Error(fmt.Sprintf("check channel[%d] %s(%d) balance error", ch.Type, ch.Name, ch.ID), err.Error()) + } + }(channel) } + + wg.Wait() return nil } @@ -93,19 +114,17 @@ func UpdateAllChannelsBalance(c *gin.Context) { middleware.SuccessResponse(c, nil) } -func AutomaticallyUpdateChannels(frequency int) { +func UpdateChannelsBalance(frequency time.Duration) { for { - time.Sleep(time.Duration(frequency) * time.Minute) - log.Info("updating all channels") + time.Sleep(frequency) _ = updateAllChannelsBalance() - log.Info("channels update done") } } // subscription func GetSubscription(c *gin.Context) { group := middleware.GetGroup(c) - b, _, err := balance.Default.GetGroupRemainBalance(c, *group) + b, _, err := balance.GetGroupRemainBalance(c, *group) if err != nil { if errors.Is(err, balance.ErrNoRealNameUsedAmountLimit) { middleware.ErrorResponse(c, http.StatusForbidden, err.Error()) diff --git a/service/aiproxy/controller/channel-test.go b/service/aiproxy/controller/channel-test.go index c7656edf6..e1561068b 100644 --- a/service/aiproxy/controller/channel-test.go +++ b/service/aiproxy/controller/channel-test.go @@ -17,10 +17,13 @@ import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/labring/sealos/service/aiproxy/common/render" + "github.com/labring/sealos/service/aiproxy/common/trylock" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/monitor" + "github.com/labring/sealos/service/aiproxy/relay/channeltype" "github.com/labring/sealos/service/aiproxy/relay/meta" "github.com/labring/sealos/service/aiproxy/relay/relaymode" "github.com/labring/sealos/service/aiproxy/relay/utils" @@ -29,12 +32,40 @@ import ( const channelTestRequestID = "channel-test" +var ( + modelTypeCache map[string]relaymode.Mode = make(map[string]relaymode.Mode) + modelTypeCacheOnce sync.Once +) + +func guessModelType(model string) relaymode.Mode { + modelTypeCacheOnce.Do(func() { + for _, c := range channeltype.ChannelAdaptor { + for _, m := range c.GetModelList() { + if _, ok := modelTypeCache[m.Model]; !ok { + modelTypeCache[m.Model] = m.Type + } + } + } + }) + + if cachedType, ok := modelTypeCache[model]; ok { + return cachedType + } + return relaymode.Unknown +} + // testSingleModel tests a single model in the channel func testSingleModel(mc *model.ModelCaches, channel *model.Channel, modelName string) (*model.ChannelTest, error) { modelConfig, ok := mc.ModelConfig.GetModelConfig(modelName) if !ok { return nil, errors.New(modelName + " model config not found") } + if modelConfig.Type == relaymode.Unknown { + newModelConfig := *modelConfig + newModelConfig.Type = guessModelType(modelName) + modelConfig = &newModelConfig + } + body, mode, err := utils.BuildRequest(modelConfig) if err != nil { return nil, err @@ -187,8 +218,6 @@ func processTestResult(mc *model.ModelCaches, channel *model.Channel, modelName return result } -//nolint:goconst -//nolint:gosec func TestChannelModels(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) if err != nil { @@ -275,8 +304,6 @@ func TestChannelModels(c *gin.Context) { } } -//nolint:goconst -//nolint:gosec func TestAllChannels(c *gin.Context) { testDisabled := c.Query("test_disabled") == "true" var channels []*model.Channel @@ -372,11 +399,15 @@ func TestAllChannels(c *gin.Context) { } } +func tryTestChannel(channelID int, modelName string) bool { + return trylock.Lock(fmt.Sprintf("channel_test_lock:%d:%s", channelID, modelName), 30*time.Second) +} + func AutoTestBannedModels() { log := log.WithFields(log.Fields{ "auto_test_banned_models": "true", }) - channels, err := monitor.GetAllBannedChannels(context.Background()) + channels, err := monitor.GetAllBannedModelChannels(context.Background()) if err != nil { log.Errorf("failed to get banned channels: %s", err.Error()) return @@ -389,6 +420,9 @@ func AutoTestBannedModels() { for modelName, ids := range channels { for _, id := range ids { + if !tryTestChannel(int(id), modelName) { + continue + } channel, err := model.LoadChannelByID(int(id)) if err != nil { log.Errorf("failed to get channel by model %s: %s", modelName, err.Error()) @@ -396,16 +430,18 @@ func AutoTestBannedModels() { } result, err := testSingleModel(mc, channel, modelName) if err != nil { - log.Errorf("failed to test channel %s(%d) model %s: %s", channel.Name, channel.ID, modelName, err.Error()) + notify.Error(fmt.Sprintf("channel[%d] %s(%d) model %s test failed", channel.Type, channel.Name, channel.ID, modelName), err.Error()) + continue } if result.Success { - log.Infof("model %s(%d) test success, unban it", modelName, channel.ID) + notify.Info(fmt.Sprintf("channel[%d] %s(%d) model %s test success", channel.Type, channel.Name, channel.ID, modelName), "unban it") err = monitor.ClearChannelModelErrors(context.Background(), modelName, channel.ID) if err != nil { log.Errorf("clear channel errors failed: %+v", err) } } else { - log.Infof("model %s(%d) test failed", modelName, channel.ID) + notify.Error(fmt.Sprintf("channel[%d] %s(%d) model %s test failed", channel.Type, channel.Name, channel.ID, modelName), + fmt.Sprintf("code: %d, response: %s", result.Code, result.Response)) } } } diff --git a/service/aiproxy/controller/dashboard.go b/service/aiproxy/controller/dashboard.go index 128d2b570..4630d5411 100644 --- a/service/aiproxy/controller/dashboard.go +++ b/service/aiproxy/controller/dashboard.go @@ -11,6 +11,7 @@ import ( "github.com/labring/sealos/service/aiproxy/common/rpmlimit" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/model" + "gorm.io/gorm" ) func getDashboardTime(t string) (time.Time, time.Time, model.TimeSpanType) { @@ -190,7 +191,7 @@ func GetGroupDashboardModels(c *gin.Context) { } groupCache, err := model.CacheGetGroup(group) if err != nil { - if errors.Is(err, model.NotFoundError(model.ErrGroupNotFound)) { + if errors.Is(err, gorm.ErrRecordNotFound) { middleware.SuccessResponse(c, model.LoadModelCaches().EnabledModelConfigs) } else { middleware.ErrorResponse(c, http.StatusOK, fmt.Sprintf("failed to get group: %v", err)) diff --git a/service/aiproxy/controller/import.go b/service/aiproxy/controller/import.go index 94c81df9d..ef1c19f08 100644 --- a/service/aiproxy/controller/import.go +++ b/service/aiproxy/controller/import.go @@ -11,10 +11,10 @@ import ( ) type OneAPIChannel struct { - Type int `json:"type" gorm:"default:0"` - Key string `json:"key" gorm:"type:text"` - Status int `json:"status" gorm:"default:1"` - Name string `json:"name" gorm:"index"` + Type int `gorm:"default:0" json:"type"` + Key string `gorm:"type:text" json:"key"` + Status int `gorm:"default:1" json:"status"` + Name string `gorm:"index" json:"name"` BaseURL string `gorm:"column:base_url;default:''"` Models string `json:"models"` ModelMapping map[string]string `gorm:"type:varchar(1024);serializer:fastjson"` diff --git a/service/aiproxy/controller/monitor.go b/service/aiproxy/controller/monitor.go index 5e5554668..4b7d370aa 100644 --- a/service/aiproxy/controller/monitor.go +++ b/service/aiproxy/controller/monitor.go @@ -81,3 +81,12 @@ func GetModelsErrorRate(c *gin.Context) { } c.JSON(http.StatusOK, rates) } + +func GetAllBannedModelChannels(c *gin.Context) { + channels, err := monitor.GetAllBannedModelChannels(c.Request.Context()) + if err != nil { + middleware.ErrorResponse(c, http.StatusOK, err.Error()) + return + } + c.JSON(http.StatusOK, channels) +} diff --git a/service/aiproxy/controller/relay.go b/service/aiproxy/controller/relay.go index 4171d8501..8d4f255d7 100644 --- a/service/aiproxy/controller/relay.go +++ b/service/aiproxy/controller/relay.go @@ -4,15 +4,18 @@ import ( "bytes" "context" "errors" + "fmt" "io" "math/rand/v2" "net/http" "slices" + "strconv" "time" "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common" "github.com/labring/sealos/service/aiproxy/common/config" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/labring/sealos/service/aiproxy/middleware" dbmodel "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/monitor" @@ -27,7 +30,7 @@ import ( type RelayController func(*meta.Meta, *gin.Context) *model.ErrorWithStatusCode -func relayController(mode int) (RelayController, bool) { +func relayController(mode relaymode.Mode) (RelayController, bool) { var relayController RelayController switch mode { case relaymode.ImagesGenerations, @@ -58,38 +61,67 @@ func relayController(mode int) (RelayController, bool) { } func RelayHelper(meta *meta.Meta, c *gin.Context, relayController RelayController) (*model.ErrorWithStatusCode, bool) { - err := relayController(meta, c) - if err == nil { - if err := monitor.AddRequest( + relayErr := relayController(meta, c) + if relayErr == nil { + if _, _, err := monitor.AddRequest( context.Background(), meta.OriginModel, int64(meta.Channel.ID), false, + false, ); err != nil { log.Errorf("add request failed: %+v", err) } return nil, false } - if shouldErrorMonitor(err.StatusCode) { - if err := monitor.AddRequest( + if shouldErrorMonitor(relayErr.StatusCode) { + hasPermission := channelHasPermission(relayErr.StatusCode) + beyondThreshold, autoBanned, err := monitor.AddRequest( context.Background(), meta.OriginModel, int64(meta.Channel.ID), true, - ); err != nil { + !hasPermission, + ) + if err != nil { log.Errorf("add request failed: %+v", err) } + if autoBanned { + notify.ErrorThrottle( + fmt.Sprintf("autoBanned:%d:%s", meta.Channel.ID, meta.OriginModel), + time.Minute, + fmt.Sprintf("channel[%d] %s(%d) model %s is auto banned", + meta.Channel.Type, meta.Channel.Name, meta.Channel.ID, meta.OriginModel), + relayErr.JSONOrEmpty(), + ) + } else if beyondThreshold { + notify.WarnThrottle( + fmt.Sprintf("beyondThreshold:%d:%s", meta.Channel.ID, meta.OriginModel), + time.Minute, + fmt.Sprintf("channel[%d] %s(%d) model %s error rate is beyond threshold", + meta.Channel.Type, meta.Channel.Name, meta.Channel.ID, meta.OriginModel), + relayErr.JSONOrEmpty(), + ) + } else if !hasPermission { + notify.ErrorThrottle( + fmt.Sprintf("channelHasPermission:%d:%s", meta.Channel.ID, meta.OriginModel), + time.Minute, + fmt.Sprintf("channel[%d] %s(%d) model %s has no permission", + meta.Channel.Type, meta.Channel.Name, meta.Channel.ID, meta.OriginModel), + relayErr.JSONOrEmpty(), + ) + } } - return err, shouldRetry(c, err.StatusCode) + return relayErr, shouldRetry(c, relayErr.StatusCode) } -func filterChannels(channels []*dbmodel.Channel, ignoreChannel ...int) []*dbmodel.Channel { +func filterChannels(channels []*dbmodel.Channel, ignoreChannel ...int64) []*dbmodel.Channel { filtered := make([]*dbmodel.Channel, 0) for _, channel := range channels { if channel.Status != dbmodel.ChannelStatusEnabled { continue } - if slices.Contains(ignoreChannel, channel.ID) { + if slices.Contains(ignoreChannel, int64(channel.ID)) { continue } filtered = append(filtered, channel) @@ -102,12 +134,22 @@ var ( ErrChannelsExhausted = errors.New("channels exhausted") ) -func GetRandomChannel(c *dbmodel.ModelCaches, model string, ignoreChannel ...int) (*dbmodel.Channel, error) { - return getRandomChannel(c.EnabledModel2channels[model], ignoreChannel...) +func GetRandomChannel(c *dbmodel.ModelCaches, model string, errorRates map[int64]float64, ignoreChannel ...int64) (*dbmodel.Channel, error) { + return getRandomChannel(c.EnabledModel2channels[model], errorRates, ignoreChannel...) +} + +func getPriority(channel *dbmodel.Channel, errorRate float64) int32 { + priority := channel.GetPriority() + if errorRate > 1 { + errorRate = 1 + } else if errorRate < 0.1 { + errorRate = 0.1 + } + return int32(float64(priority) / errorRate) } //nolint:gosec -func getRandomChannel(channels []*dbmodel.Channel, ignoreChannel ...int) (*dbmodel.Channel, error) { +func getRandomChannel(channels []*dbmodel.Channel, errorRates map[int64]float64, ignoreChannel ...int64) (*dbmodel.Channel, error) { if len(channels) == 0 { return nil, ErrChannelsNotFound } @@ -122,8 +164,11 @@ func getRandomChannel(channels []*dbmodel.Channel, ignoreChannel ...int) (*dbmod } var totalWeight int32 - for _, ch := range channels { - totalWeight += ch.GetPriority() + cachedPrioritys := make([]int32, len(channels)) + for i, ch := range channels { + priority := getPriority(ch, errorRates[int64(ch.ID)]) + totalWeight += priority + cachedPrioritys[i] = priority } if totalWeight == 0 { @@ -131,8 +176,8 @@ func getRandomChannel(channels []*dbmodel.Channel, ignoreChannel ...int) (*dbmod } r := rand.Int32N(totalWeight) - for _, ch := range channels { - r -= ch.GetPriority() + for i, ch := range channels { + r -= cachedPrioritys[i] if r < 0 { return ch, nil } @@ -141,18 +186,18 @@ func getRandomChannel(channels []*dbmodel.Channel, ignoreChannel ...int) (*dbmod return channels[rand.IntN(len(channels))], nil } -func getChannelWithFallback(cache *dbmodel.ModelCaches, model string, ignoreChannelIDs ...int) (*dbmodel.Channel, error) { - channel, err := GetRandomChannel(cache, model, ignoreChannelIDs...) +func getChannelWithFallback(cache *dbmodel.ModelCaches, model string, errorRates map[int64]float64, ignoreChannelIDs ...int64) (*dbmodel.Channel, error) { + channel, err := GetRandomChannel(cache, model, errorRates, ignoreChannelIDs...) if err == nil { return channel, nil } if !errors.Is(err, ErrChannelsExhausted) { return nil, err } - return GetRandomChannel(cache, model) + return GetRandomChannel(cache, model, errorRates) } -func NewRelay(mode int) func(c *gin.Context) { +func NewRelay(mode relaymode.Mode) func(c *gin.Context) { relayController, ok := relayController(mode) if !ok { log.Fatalf("relay mode %d not implemented", mode) @@ -162,13 +207,13 @@ func NewRelay(mode int) func(c *gin.Context) { } } -func relay(c *gin.Context, mode int, relayController RelayController) { +func relay(c *gin.Context, mode relaymode.Mode, relayController RelayController) { log := middleware.GetLogger(c) requestModel := middleware.GetOriginalModel(c) // Get initial channel - channel, ignoreChannelIDs, err := getInitialChannel(c, requestModel, log) - if err != nil || channel == nil { + initialChannel, err := getInitialChannel(c, requestModel, log) + if err != nil || initialChannel == nil || initialChannel.channel == nil { c.JSON(http.StatusServiceUnavailable, gin.H{ "error": &model.Error{ Message: "the upstream load is saturated, please try again later", @@ -180,46 +225,62 @@ func relay(c *gin.Context, mode int, relayController RelayController) { } // First attempt - meta := middleware.NewMetaByContext(c, channel, requestModel, mode) + meta := middleware.NewMetaByContext(c, initialChannel.channel, requestModel, mode) bizErr, retry := RelayHelper(meta, c, relayController) if handleRelayResult(c, bizErr, retry) { return } // Setup retry state - retryState := initRetryState(channel, bizErr, ignoreChannelIDs) + retryState := initRetryState(initialChannel.channel, + bizErr, + initialChannel.ignoreChannelIDs, + initialChannel.errorRates, + ) // Retry loop retryLoop(c, mode, requestModel, retryState, relayController, log) } type retryState struct { - retryTimes int64 - lastCanContinueChannel *dbmodel.Channel - ignoreChannelIDs []int - exhausted bool - bizErr *model.ErrorWithStatusCode + retryTimes int + lastHasPermissionChannel *dbmodel.Channel + ignoreChannelIDs []int64 + errorRates map[int64]float64 + exhausted bool + bizErr *model.ErrorWithStatusCode + startTime time.Time } -func getInitialChannel(c *gin.Context, requestModel string, log *log.Entry) (*dbmodel.Channel, []int, error) { - ids, err := monitor.GetBannedChannels(c.Request.Context(), requestModel) +type initialChannel struct { + channel *dbmodel.Channel + ignoreChannelIDs []int64 + errorRates map[int64]float64 +} + +func getInitialChannel(c *gin.Context, requestModel string, log *log.Entry) (*initialChannel, error) { + ids, err := monitor.GetBannedChannelsWithModel(c.Request.Context(), requestModel) if err != nil { log.Errorf("get %s auto banned channels failed: %+v", requestModel, err) } log.Debugf("%s model banned channels: %+v", requestModel, ids) - ignoreChannelIDs := make([]int, 0, len(ids)) - for _, id := range ids { - ignoreChannelIDs = append(ignoreChannelIDs, int(id)) + errorRates, err := monitor.GetModelChannelErrorRate(c.Request.Context(), requestModel) + if err != nil { + log.Errorf("get channel model error rates failed: %+v", err) } mc := middleware.GetModelCaches(c) - channel, err := getChannelWithFallback(mc, requestModel, ignoreChannelIDs...) + channel, err := getChannelWithFallback(mc, requestModel, errorRates, ids...) if err != nil { - return nil, nil, err + return nil, err } - return channel, ignoreChannelIDs, nil + return &initialChannel{ + channel: channel, + ignoreChannelIDs: ids, + errorRates: errorRates, + }, nil } func handleRelayResult(c *gin.Context, bizErr *model.ErrorWithStatusCode, retry bool) bool { @@ -234,43 +295,58 @@ func handleRelayResult(c *gin.Context, bizErr *model.ErrorWithStatusCode, retry return false } -func initRetryState(channel *dbmodel.Channel, bizErr *model.ErrorWithStatusCode, ignoreChannelIDs []int) *retryState { +func initRetryState(channel *dbmodel.Channel, bizErr *model.ErrorWithStatusCode, ignoreChannelIDs []int64, errorRates map[int64]float64) *retryState { state := &retryState{ - retryTimes: config.GetRetryTimes(), + retryTimes: int(config.GetRetryTimes()), ignoreChannelIDs: ignoreChannelIDs, + errorRates: errorRates, bizErr: bizErr, + startTime: time.Now(), } - if !channelCanContinue(bizErr.StatusCode) { - state.ignoreChannelIDs = append(state.ignoreChannelIDs, channel.ID) + if !channelHasPermission(bizErr.StatusCode) { + state.ignoreChannelIDs = append(state.ignoreChannelIDs, int64(channel.ID)) } else { - state.lastCanContinueChannel = channel + state.lastHasPermissionChannel = channel } return state } -func retryLoop(c *gin.Context, mode int, requestModel string, state *retryState, relayController RelayController, log *log.Entry) { +func retryLoop(c *gin.Context, mode relaymode.Mode, requestModel string, state *retryState, relayController RelayController, log *log.Entry) { mc := middleware.GetModelCaches(c) - for i := 0; i < int(state.retryTimes); i++ { + for i := 0; i < state.retryTimes; i++ { + ctxErr := c.Request.Context().Err() + if ctxErr != nil { + log.Warnf("retry loop context error: %+v", ctxErr) + break + } + newChannel, err := getRetryChannel(mc, requestModel, state) if err != nil { break } + log.Data["retry"] = strconv.Itoa(i + 1) + log.Warnf("using channel %s (type: %d, id: %d) to retry (remain times %d)", newChannel.Name, newChannel.Type, newChannel.ID, - state.retryTimes-int64(i), + state.retryTimes-i, ) if !prepareRetry(c, state.bizErr.StatusCode) { break } - meta := middleware.NewMetaByContext(c, newChannel, requestModel, mode) + meta := middleware.NewMetaByContext(c, + newChannel, + requestModel, + mode, + meta.WithRetryTimes(i+1), + ) bizErr, retry := RelayHelper(meta, c, relayController) done := handleRetryResult(bizErr, retry, newChannel, state) @@ -285,18 +361,18 @@ func retryLoop(c *gin.Context, mode int, requestModel string, state *retryState, } } -func getRetryChannel(mc *dbmodel.ModelCaches, requestModel string, state *retryState) (*dbmodel.Channel, error) { +func getRetryChannel(mc *dbmodel.ModelCaches, model string, state *retryState) (*dbmodel.Channel, error) { if state.exhausted { - return state.lastCanContinueChannel, nil + return state.lastHasPermissionChannel, nil } - newChannel, err := GetRandomChannel(mc, requestModel, state.ignoreChannelIDs...) + newChannel, err := GetRandomChannel(mc, model, state.errorRates, state.ignoreChannelIDs...) if err != nil { - if !errors.Is(err, ErrChannelsExhausted) || state.lastCanContinueChannel == nil { + if !errors.Is(err, ErrChannelsExhausted) || state.lastHasPermissionChannel == nil { return nil, err } state.exhausted = true - return state.lastCanContinueChannel, nil + return state.lastHasPermissionChannel, nil } return newChannel, nil @@ -320,49 +396,39 @@ func prepareRetry(c *gin.Context, statusCode int) bool { func handleRetryResult(bizErr *model.ErrorWithStatusCode, retry bool, newChannel *dbmodel.Channel, state *retryState) (done bool) { state.bizErr = bizErr - if bizErr == nil || !retry { + if !retry || bizErr == nil { return true } if state.exhausted { - if !channelCanContinue(bizErr.StatusCode) { + if !channelHasPermission(bizErr.StatusCode) { return true } } else { - if !channelCanContinue(bizErr.StatusCode) { - state.ignoreChannelIDs = append(state.ignoreChannelIDs, newChannel.ID) + if !channelHasPermission(bizErr.StatusCode) { + state.ignoreChannelIDs = append(state.ignoreChannelIDs, int64(newChannel.ID)) state.retryTimes++ } else { - state.lastCanContinueChannel = newChannel + state.lastHasPermissionChannel = newChannel } } return false } -var shouldRetryStatusCodesMap = map[int]struct{}{ - http.StatusTooManyRequests: {}, +func shouldRetry(_ *gin.Context, statusCode int) bool { + return statusCode != http.StatusBadRequest +} + +var channelNoPermissionStatusCodesMap = map[int]struct{}{ http.StatusUnauthorized: {}, http.StatusPaymentRequired: {}, - http.StatusRequestTimeout: {}, - http.StatusGatewayTimeout: {}, http.StatusForbidden: {}, } -func shouldRetry(_ *gin.Context, statusCode int) bool { - _, ok := shouldRetryStatusCodesMap[statusCode] - return ok -} - -var channelCanContinueStatusCodesMap = map[int]struct{}{ - http.StatusTooManyRequests: {}, - http.StatusRequestTimeout: {}, - http.StatusGatewayTimeout: {}, -} - -func channelCanContinue(statusCode int) bool { - _, ok := channelCanContinueStatusCodesMap[statusCode] - return ok +func channelHasPermission(statusCode int) bool { + _, ok := channelNoPermissionStatusCodesMap[statusCode] + return !ok } func shouldDelay(statusCode int) bool { @@ -379,7 +445,6 @@ func RelayNotImplemented(c *gin.Context) { "error": &model.Error{ Message: "API not implemented", Type: middleware.ErrorTypeAIPROXY, - Param: "", Code: "api_not_implemented", }, }) diff --git a/service/aiproxy/main.go b/service/aiproxy/main.go index ef0735f45..372cb70aa 100644 --- a/service/aiproxy/main.go +++ b/service/aiproxy/main.go @@ -3,6 +3,7 @@ package main import ( "context" "errors" + "flag" "fmt" stdlog "log" "net/http" @@ -20,6 +21,7 @@ import ( "github.com/labring/sealos/service/aiproxy/common/balance" "github.com/labring/sealos/service/aiproxy/common/config" "github.com/labring/sealos/service/aiproxy/common/consume" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/labring/sealos/service/aiproxy/controller" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/model" @@ -27,10 +29,16 @@ import ( log "github.com/sirupsen/logrus" ) +var port int + +func init() { + flag.IntVar(&port, "port", 3000, "http server port") +} + func initializeServices() error { setLog(log.StandardLogger()) - common.Init() + initializeNotifier() if err := initializeBalance(); err != nil { return err @@ -54,6 +62,14 @@ func initializeBalance() error { return balance.InitSealos(sealosJwtKey, os.Getenv("SEALOS_ACCOUNT_URL")) } +func initializeNotifier() { + feishuWh := os.Getenv("NOTIFY_FEISHU_WEBHOOK") + if feishuWh != "" { + notify.SetDefaultNotifier(notify.NewFeishuNotify(feishuWh)) + log.Info("NOTIFY_FEISHU_WEBHOOK is set, notifier will be use feishu") + } +} + var logCallerIgnoreFuncs = map[string]struct{}{ "github.com/labring/sealos/service/aiproxy/middleware.logColor": {}, } @@ -123,13 +139,13 @@ func setupHTTPServer() (*http.Server, *gin.Engine) { Use(middleware.RequestID, middleware.CORS()) router.SetRouter(server) - port := os.Getenv("PORT") - if port == "" { - port = strconv.Itoa(*common.Port) + p := os.Getenv("PORT") + if p == "" { + p = strconv.Itoa(port) } return &http.Server{ - Addr: ":" + port, + Addr: ":" + p, ReadHeaderTimeout: 10 * time.Second, Handler: server, }, server @@ -137,7 +153,7 @@ func setupHTTPServer() (*http.Server, *gin.Engine) { func autoTestBannedModels(ctx context.Context) { log.Info("auto test banned models start") - ticker := time.NewTicker(time.Second * 15) + ticker := time.NewTicker(time.Second * 30) defer ticker.Stop() for { @@ -152,7 +168,8 @@ func autoTestBannedModels(ctx context.Context) { func cleanLog(ctx context.Context) { log.Info("clean log start") - ticker := time.NewTicker(time.Minute * 15) + // the interval should not be too large to avoid cleaning too much at once + ticker := time.NewTicker(time.Second * 15) defer ticker.Stop() for { @@ -160,15 +177,17 @@ func cleanLog(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: - err := model.CleanLog() + err := model.CleanLog(1000) if err != nil { - log.Errorf("clean log failed: %s", err) + notify.ErrorThrottle("cleanLog", time.Minute, "clean log failed", err.Error()) } } } } func main() { + flag.Parse() + if err := initializeServices(); err != nil { log.Fatal("failed to initialize services: " + err.Error()) } @@ -197,6 +216,11 @@ func main() { go autoTestBannedModels(ctx) go cleanLog(ctx) + go controller.UpdateChannelsBalance(time.Minute * 10) + + batchProcessorCtx, batchProcessorCancel := context.WithCancel(context.Background()) + wg.Add(1) + go model.StartBatchProcessor(batchProcessorCtx, &wg) <-ctx.Done() @@ -214,8 +238,13 @@ func main() { log.Info("shutting down consumer...") consume.Wait() + batchProcessorCancel() + log.Info("shutting down sync services...") wg.Wait() + log.Info("shutting down batch processor...") + model.ProcessBatchUpdates() + log.Info("server exiting") } diff --git a/service/aiproxy/middleware/auth.go b/service/aiproxy/middleware/auth.go index f7c161fd0..6502088ed 100644 --- a/service/aiproxy/middleware/auth.go +++ b/service/aiproxy/middleware/auth.go @@ -11,6 +11,7 @@ import ( "github.com/labring/sealos/service/aiproxy/common/network" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/meta" + "github.com/labring/sealos/service/aiproxy/relay/relaymode" "github.com/sirupsen/logrus" ) @@ -61,14 +62,17 @@ func TokenAuth(c *gin.Context) { var token *model.TokenCache var useInternalToken bool - if config.GetInternalToken() != "" && config.GetInternalToken() == key || config.AdminKey != "" && config.AdminKey == key { + if config.AdminKey != "" && config.AdminKey == key || + config.GetInternalToken() != "" && config.GetInternalToken() == key { token = &model.TokenCache{} useInternalToken = true } else { var err error token, err = model.ValidateAndGetToken(key) if err != nil { - abortLogWithMessage(c, http.StatusUnauthorized, err.Error()) + abortLogWithMessage(c, http.StatusUnauthorized, err.Error(), &errorField{ + Code: "invalid_token", + }) return } } @@ -174,8 +178,8 @@ func SetLogFieldsFromMeta(m *meta.Meta, fields logrus.Fields) { SetLogChannelFields(fields, m.Channel) } -func SetLogModeField(fields logrus.Fields, mode int) { - fields["mode"] = mode +func SetLogModeField(fields logrus.Fields, mode relaymode.Mode) { + fields["mode"] = mode.String() } func SetLogIsChannelTestField(fields logrus.Fields, isChannelTest bool) { diff --git a/service/aiproxy/middleware/distributor.go b/service/aiproxy/middleware/distributor.go index b5e8742bc..f6ab3daed 100644 --- a/service/aiproxy/middleware/distributor.go +++ b/service/aiproxy/middleware/distributor.go @@ -17,6 +17,7 @@ import ( "github.com/labring/sealos/service/aiproxy/common/config" "github.com/labring/sealos/service/aiproxy/common/consume" "github.com/labring/sealos/service/aiproxy/common/ctxkey" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/labring/sealos/service/aiproxy/common/rpmlimit" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/meta" @@ -100,7 +101,7 @@ func checkGroupModelRPMAndTPM(c *gin.Context, group *model.GroupCache, mc *model return ErrRequestRateLimitExceeded } } else if common.RedisEnabled { - _, err := rpmlimit.PushRequest(c.Request.Context(), group.ID, mc.Model, time.Minute) + _, _, err := rpmlimit.PushRequest(c.Request.Context(), group.ID, mc.Model, 1, time.Minute) if err != nil { log.Errorf("push request error: %s", err.Error()) } @@ -126,7 +127,16 @@ type GroupBalanceConsumer struct { Consumer balance.PostGroupConsumer } -func checkGroupBalance(c *gin.Context, group *model.GroupCache) bool { +func GetGroupBalanceConsumer(c *gin.Context, group *model.GroupCache) (*GroupBalanceConsumer, error) { + gbcI, ok := c.Get(ctxkey.GroupBalance) + if ok { + groupBalanceConsumer, ok := gbcI.(*GroupBalanceConsumer) + if !ok { + return nil, errors.New("internal error: group balance consumer unavailable") + } + return groupBalanceConsumer, nil + } + var groupBalance float64 var consumer balance.PostGroupConsumer @@ -135,37 +145,50 @@ func checkGroupBalance(c *gin.Context, group *model.GroupCache) bool { } else { log := GetLogger(c) var err error - groupBalance, consumer, err = balance.Default.GetGroupRemainBalance(c.Request.Context(), *group) + groupBalance, consumer, err = balance.GetGroupRemainBalance(c.Request.Context(), *group) if err != nil { - if errors.Is(err, balance.ErrNoRealNameUsedAmountLimit) { - abortLogWithMessage(c, http.StatusForbidden, balance.ErrNoRealNameUsedAmountLimit.Error()) - return false - } - log.Errorf("get group (%s) balance error: %v", group.ID, err) - abortWithMessage(c, http.StatusInternalServerError, fmt.Sprintf("get group (%s) balance error", group.ID)) - return false + return nil, err } log.Data["balance"] = strconv.FormatFloat(groupBalance, 'f', -1, 64) } - if groupBalance <= 0 { - abortLogWithMessage(c, http.StatusForbidden, fmt.Sprintf("group (%s) balance not enough", group.ID)) + gbc := &GroupBalanceConsumer{GroupBalance: groupBalance, Consumer: consumer} + c.Set(ctxkey.GroupBalance, gbc) + return gbc, nil +} + +func checkGroupBalance(c *gin.Context, group *model.GroupCache) bool { + gbc, err := GetGroupBalanceConsumer(c, group) + if err != nil { + if errors.Is(err, balance.ErrNoRealNameUsedAmountLimit) { + abortLogWithMessage(c, http.StatusForbidden, err.Error(), &errorField{ + Code: "no_real_name_used_amount_limit", + }) + return false + } + notify.ErrorThrottle("balance", time.Minute, fmt.Sprintf("get group (%s) balance error", group.ID), err.Error()) + abortWithMessage(c, http.StatusInternalServerError, fmt.Sprintf("get group (%s) balance error", group.ID), &errorField{ + Code: "get_group_balance_error", + }) + return false + } + + if gbc.GroupBalance <= 0 { + abortLogWithMessage(c, http.StatusForbidden, fmt.Sprintf("group (%s) balance not enough", group.ID), &errorField{ + Code: "group_balance_not_enough", + }) return false } - c.Set(ctxkey.GroupBalance, &GroupBalanceConsumer{ - GroupBalance: groupBalance, - Consumer: consumer, - }) return true } -func NewDistribute(mode int) gin.HandlerFunc { +func NewDistribute(mode relaymode.Mode) gin.HandlerFunc { return func(c *gin.Context) { distribute(c, mode) } } -func distribute(c *gin.Context, mode int) { +func distribute(c *gin.Context, mode relaymode.Mode) { if config.GetDisableServe() { abortLogWithMessage(c, http.StatusServiceUnavailable, "service is under maintenance") return @@ -181,11 +204,17 @@ func distribute(c *gin.Context, mode int) { requestModel, err := getRequestModel(c, mode) if err != nil { - abortLogWithMessage(c, http.StatusBadRequest, err.Error()) + abortLogWithMessage(c, http.StatusInternalServerError, err.Error(), &errorField{ + Type: "invalid_request_error", + Code: "get_request_model_error", + }) return } if requestModel == "" { - abortLogWithMessage(c, http.StatusBadRequest, "no model provided") + abortLogWithMessage(c, http.StatusBadRequest, "no model provided", &errorField{ + Type: "invalid_request_error", + Code: "no_model_provided", + }) return } @@ -193,25 +222,20 @@ func distribute(c *gin.Context, mode int) { SetLogModelFields(log.Data, requestModel) - mc, ok := GetModelCaches(c).ModelConfig.GetModelConfig(requestModel) - if !ok { - abortLogWithMessage(c, http.StatusServiceUnavailable, requestModel+" is not available") - return - } - - c.Set(ctxkey.ModelConfig, mc) - token := GetToken(c) - - if len(token.Models) == 0 || !slices.Contains(token.Models, requestModel) { + mc, ok := GetModelCaches(c).ModelConfig.GetModelConfig(requestModel) + if !ok || len(token.Models) == 0 || !slices.Contains(token.Models, requestModel) { abortLogWithMessage(c, - http.StatusForbidden, - fmt.Sprintf("token (%s[%d]) has no permission to use model: %s", - token.Name, token.ID, requestModel, - ), + http.StatusNotFound, + fmt.Sprintf("The model `%s` does not exist or you do not have access to it.", requestModel), + &errorField{ + Type: "invalid_request_error", + Code: "model_not_found", + }, ) return } + c.Set(ctxkey.ModelConfig, mc) if err := checkGroupModelRPMAndTPM(c, group, mc); err != nil { errMsg := err.Error() @@ -224,9 +248,13 @@ func distribute(c *gin.Context, mode int) { 0, errMsg, c.ClientIP(), + 0, nil, ) - abortLogWithMessage(c, http.StatusTooManyRequests, errMsg) + abortLogWithMessage(c, http.StatusTooManyRequests, errMsg, &errorField{ + Type: "invalid_request_error", + Code: "request_rate_limit_exceeded", + }) return } @@ -241,20 +269,30 @@ func GetModelConfig(c *gin.Context) *model.ModelConfig { return c.MustGet(ctxkey.ModelConfig).(*model.ModelConfig) } -func NewMetaByContext(c *gin.Context, channel *model.Channel, modelName string, mode int) *meta.Meta { +func NewMetaByContext(c *gin.Context, + channel *model.Channel, + modelName string, + mode relaymode.Mode, + opts ...meta.Option, +) *meta.Meta { requestID := GetRequestID(c) group := GetGroup(c) token := GetToken(c) + opts = append( + opts, + meta.WithRequestID(requestID), + meta.WithGroup(group), + meta.WithToken(token), + meta.WithEndpoint(c.Request.URL.Path), + ) + return meta.NewMeta( channel, mode, modelName, GetModelConfig(c), - meta.WithRequestID(requestID), - meta.WithGroup(group), - meta.WithToken(token), - meta.WithEndpoint(c.Request.URL.Path), + opts..., ) } @@ -262,7 +300,7 @@ type ModelRequest struct { Model string `form:"model" json:"model"` } -func getRequestModel(c *gin.Context, mode int) (string, error) { +func getRequestModel(c *gin.Context, mode relaymode.Mode) (string, error) { path := c.Request.URL.Path switch { case mode == relaymode.ParsePdf: diff --git a/service/aiproxy/middleware/utils.go b/service/aiproxy/middleware/utils.go index c7264561e..fa733999e 100644 --- a/service/aiproxy/middleware/utils.go +++ b/service/aiproxy/middleware/utils.go @@ -15,16 +15,30 @@ func MessageWithRequestID(c *gin.Context, message string) string { return fmt.Sprintf("%s (aiproxy: %s)", message, GetRequestID(c)) } -func abortLogWithMessage(c *gin.Context, statusCode int, message string) { +func abortLogWithMessage(c *gin.Context, statusCode int, message string, fields ...*errorField) { GetLogger(c).Error(message) - abortWithMessage(c, statusCode, message) + abortWithMessage(c, statusCode, message, fields...) } -func abortWithMessage(c *gin.Context, statusCode int, message string) { +type errorField struct { + Type string `json:"type"` + Code any `json:"code"` +} + +func abortWithMessage(c *gin.Context, statusCode int, message string, fields ...*errorField) { + typeName := ErrorTypeAIPROXY + var code any = nil + if len(fields) > 0 { + if fields[0].Type != "" { + typeName = fields[0].Type + } + code = fields[0].Code + } c.JSON(statusCode, gin.H{ "error": &model.Error{ Message: MessageWithRequestID(c, message), - Type: ErrorTypeAIPROXY, + Type: typeName, + Code: code, }, }) c.Abort() diff --git a/service/aiproxy/model/cache.go b/service/aiproxy/model/cache.go index e207f5958..45855e482 100644 --- a/service/aiproxy/model/cache.go +++ b/service/aiproxy/model/cache.go @@ -16,6 +16,7 @@ import ( "github.com/labring/sealos/service/aiproxy/common" "github.com/labring/sealos/service/aiproxy/common/config" "github.com/labring/sealos/service/aiproxy/common/conv" + "github.com/labring/sealos/service/aiproxy/common/notify" "github.com/maruel/natural" "github.com/redis/go-redis/v9" log "github.com/sirupsen/logrus" @@ -406,14 +407,13 @@ func CacheGetGroupModelTPM(id string, model string) (int64, error) { return tpm, nil } -//nolint:revive type ModelConfigCache interface { GetModelConfig(model string) (*ModelConfig, bool) } // read-only cache // -//nolint:revive + type ModelCaches struct { ModelConfig ModelConfigCache EnabledModel2channels map[string][]*Channel @@ -482,7 +482,7 @@ func InitModelConfigAndChannelCache() error { func LoadEnabledChannels() ([]*Channel, error) { var channels []*Channel - err := DB.Where("status = ? or status = ?", ChannelStatusEnabled, ChannelStatusFail).Find(&channels).Error + err := DB.Where("status = ?", ChannelStatusEnabled).Find(&channels).Error if err != nil { return nil, err } @@ -699,8 +699,7 @@ func SyncModelConfigAndChannelCache(ctx context.Context, wg *sync.WaitGroup, fre case <-ticker.C: err := InitModelConfigAndChannelCache() if err != nil { - log.Error("failed to sync channels: " + err.Error()) - continue + notify.ErrorThrottle("syncModelChannel", time.Minute, "failed to sync channels", err.Error()) } } } diff --git a/service/aiproxy/model/channel.go b/service/aiproxy/model/channel.go index d2ad31ceb..7e4ca01e1 100644 --- a/service/aiproxy/model/channel.go +++ b/service/aiproxy/model/channel.go @@ -1,6 +1,7 @@ package model import ( + "context" "fmt" "slices" "strings" @@ -9,6 +10,8 @@ import ( "github.com/bytedance/sonic" "github.com/labring/sealos/service/aiproxy/common" "github.com/labring/sealos/service/aiproxy/common/config" + "github.com/labring/sealos/service/aiproxy/monitor" + "github.com/labring/sealos/service/aiproxy/relay/relaymode" "gorm.io/gorm" "gorm.io/gorm/clause" ) @@ -18,10 +21,8 @@ const ( ) const ( - ChannelStatusUnknown = 0 - ChannelStatusEnabled = 1 // don't use 0, 0 is the default value! - ChannelStatusDisabled = 2 // also don't use 0 - ChannelStatusFail = 3 + ChannelStatusUnknown = 0 + ChannelStatusEnabled = 1 ) type ChannelConfig struct { @@ -29,31 +30,40 @@ type ChannelConfig struct { } type Channel struct { - CreatedAt time.Time `gorm:"index" json:"created_at"` - LastTestErrorAt time.Time `json:"last_test_error_at"` - ChannelTests []*ChannelTest `gorm:"foreignKey:ChannelID;references:ID" json:"channel_tests,omitempty"` - BalanceUpdatedAt time.Time `json:"balance_updated_at"` - ModelMapping map[string]string `gorm:"serializer:fastjson;type:text" json:"model_mapping"` - Key string `gorm:"type:text;index" json:"key"` - Name string `gorm:"index" json:"name"` - BaseURL string `gorm:"index" json:"base_url"` - Models []string `gorm:"serializer:fastjson;type:text" json:"models"` - Balance float64 `json:"balance"` - ID int `gorm:"primaryKey" json:"id"` - UsedAmount float64 `gorm:"index" json:"used_amount"` - RequestCount int `gorm:"index" json:"request_count"` - Status int `gorm:"default:1;index" json:"status"` - Type int `gorm:"default:0;index" json:"type"` - Priority int32 `json:"priority"` - Config *ChannelConfig `gorm:"serializer:fastjson;type:text" json:"config,omitempty"` + CreatedAt time.Time `gorm:"index" json:"created_at"` + LastTestErrorAt time.Time `json:"last_test_error_at"` + ChannelTests []*ChannelTest `gorm:"foreignKey:ChannelID;references:ID" json:"channel_tests,omitempty"` + BalanceUpdatedAt time.Time `json:"balance_updated_at"` + ModelMapping map[string]string `gorm:"serializer:fastjson;type:text" json:"model_mapping"` + Key string `gorm:"type:text;index" json:"key"` + Name string `gorm:"index" json:"name"` + BaseURL string `gorm:"index" json:"base_url"` + Models []string `gorm:"serializer:fastjson;type:text" json:"models"` + Balance float64 `json:"balance"` + ID int `gorm:"primaryKey" json:"id"` + UsedAmount float64 `gorm:"index" json:"used_amount"` + RequestCount int `gorm:"index" json:"request_count"` + Status int `gorm:"default:1;index" json:"status"` + Type int `gorm:"default:0;index" json:"type"` + Priority int32 `json:"priority"` + EnabledAutoBalanceCheck bool `json:"enabled_auto_balance_check"` + BalanceThreshold float64 `json:"balance_threshold"` + Config *ChannelConfig `gorm:"serializer:fastjson;type:text" json:"config,omitempty"` } func (c *Channel) BeforeDelete(tx *gorm.DB) (err error) { return tx.Model(&ChannelTest{}).Where("channel_id = ?", c.ID).Delete(&ChannelTest{}).Error } +func (c *Channel) GetBalanceThreshold() float64 { + if c.BalanceThreshold < 0 { + return 0 + } + return c.BalanceThreshold +} + const ( - DefaultPriority = 100 + DefaultPriority = 10 ) func (c *Channel) GetPriority() int32 { @@ -269,7 +279,12 @@ func GetChannelByID(id int) (*Channel, error) { return &channel, HandleNotFound(err, ErrChannelNotFound) } -func BatchInsertChannels(channels []*Channel) error { +func BatchInsertChannels(channels []*Channel) (err error) { + defer func() { + if err == nil { + _ = InitModelConfigAndChannelCache() + } + }() for _, channel := range channels { if err := CheckModelConfigExist(channel.Models); err != nil { return err @@ -280,13 +295,29 @@ func BatchInsertChannels(channels []*Channel) error { }) } -func UpdateChannel(channel *Channel) error { +func UpdateChannel(channel *Channel) (err error) { + defer func() { + if err == nil { + _ = InitModelConfigAndChannelCache() + _ = monitor.ClearChannelAllModelErrors(context.Background(), channel.ID) + } + }() if err := CheckModelConfigExist(channel.Models); err != nil { return err } result := DB. Model(channel). - Select("model_mapping", "key", "name", "base_url", "models", "type", "priority", "config"). + Select( + "model_mapping", + "key", + "name", + "base_url", + "models", + "type", + "priority", + "config", + "enabled_auto_balance_check", + "balance_threshold"). Clauses(clause.Returning{}). Where("id = ?", channel.ID). Updates(channel) @@ -298,7 +329,7 @@ func ClearLastTestErrorAt(id int) error { return HandleUpdateResult(result, ErrChannelNotFound) } -func (c *Channel) UpdateModelTest(testAt time.Time, model, actualModel string, mode int, took float64, success bool, response string, code int) (*ChannelTest, error) { +func (c *Channel) UpdateModelTest(testAt time.Time, model, actualModel string, mode relaymode.Mode, took float64, success bool, response string, code int) (*ChannelTest, error) { var ct *ChannelTest err := DB.Transaction(func(tx *gorm.DB) error { if !success { @@ -318,7 +349,7 @@ func (c *Channel) UpdateModelTest(testAt time.Time, model, actualModel string, m ChannelName: c.Name, Model: model, ActualModel: actualModel, - Mode: mode, + Mode: int(mode), TestAt: testAt, Took: took, Success: success, @@ -345,12 +376,26 @@ func (c *Channel) UpdateBalance(balance float64) error { return HandleUpdateResult(result, ErrChannelNotFound) } -func DeleteChannelByID(id int) error { +func DeleteChannelByID(id int) (err error) { + defer func() { + if err == nil { + _ = InitModelConfigAndChannelCache() + _ = monitor.ClearChannelAllModelErrors(context.Background(), id) + } + }() result := DB.Delete(&Channel{ID: id}) return HandleUpdateResult(result, ErrChannelNotFound) } -func DeleteChannelsByIDs(ids []int) error { +func DeleteChannelsByIDs(ids []int) (err error) { + defer func() { + if err == nil { + _ = InitModelConfigAndChannelCache() + for _, id := range ids { + _ = monitor.ClearChannelAllModelErrors(context.Background(), id) + } + } + }() return DB.Transaction(func(tx *gorm.DB) error { return tx. Where("id IN (?)", ids). @@ -375,13 +420,3 @@ func UpdateChannelUsedAmount(id int, amount float64, requestCount int) error { }) return HandleUpdateResult(result, ErrChannelNotFound) } - -func DeleteDisabledChannel() error { - result := DB.Where("status = ?", ChannelStatusDisabled).Delete(&Channel{}) - return HandleUpdateResult(result, ErrChannelNotFound) -} - -func DeleteFailChannel() error { - result := DB.Where("status = ?", ChannelStatusFail).Delete(&Channel{}) - return HandleUpdateResult(result, ErrChannelNotFound) -} diff --git a/service/aiproxy/model/configkey.go b/service/aiproxy/model/configkey.go index c826cdf43..2cd4757b4 100644 --- a/service/aiproxy/model/configkey.go +++ b/service/aiproxy/model/configkey.go @@ -2,7 +2,6 @@ package model import "reflect" -//nolint:revive type ModelConfigKey string const ( @@ -15,7 +14,6 @@ const ( ModelConfigSupportVoicesKey ModelConfigKey = "support_voices" ) -//nolint:revive type ModelConfigOption func(config map[ModelConfigKey]any) func WithModelConfigMaxContextTokens(maxContextTokens int) ModelConfigOption { diff --git a/service/aiproxy/model/group.go b/service/aiproxy/model/group.go index c040a98ff..b983374c5 100644 --- a/service/aiproxy/model/group.go +++ b/service/aiproxy/model/group.go @@ -17,8 +17,8 @@ const ( ) const ( - GroupStatusEnabled = 1 // don't use 0, 0 is the default value! - GroupStatusDisabled = 2 // also don't use 0 + GroupStatusEnabled = 1 + GroupStatusDisabled = 2 GroupStatusInternal = 3 ) @@ -39,7 +39,6 @@ func (g *Group) BeforeDelete(tx *gorm.DB) (err error) { return tx.Model(&Token{}).Where("group_id = ?", g.ID).Delete(&Token{}).Error } -//nolint:goconst func getGroupOrder(order string) string { prefix, suffix, _ := strings.Cut(order, "-") switch prefix { @@ -151,7 +150,7 @@ func UpdateGroup(id string, group *Group) (err error) { } func UpdateGroupUsedAmountAndRequestCount(id string, amount float64, count int) (err error) { - group := &Group{ID: id} + group := &Group{} defer func() { if amount > 0 && err == nil { if err := CacheUpdateGroupUsedAmountOnlyIncrease(group.ID, group.UsedAmount); err != nil { diff --git a/service/aiproxy/model/log.go b/service/aiproxy/model/log.go index 4c279612e..e987a8beb 100644 --- a/service/aiproxy/model/log.go +++ b/service/aiproxy/model/log.go @@ -39,29 +39,30 @@ func (d *RequestDetail) BeforeSave(_ *gorm.DB) (err error) { } type Log struct { - RequestDetail *RequestDetail `gorm:"foreignKey:LogID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;" json:"request_detail,omitempty"` - RequestAt time.Time `gorm:"index" json:"request_at"` + RequestDetail *RequestDetail `gorm:"foreignKey:LogID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;" json:"request_detail,omitempty"` + RequestAt time.Time `gorm:"index" json:"request_at"` TimestampTruncByDay int64 `json:"timestamp_trunc_by_day"` TimestampTruncByHour int64 `json:"timestamp_trunc_by_hour"` - CreatedAt time.Time `json:"created_at"` + CreatedAt time.Time `gorm:"autoCreateTime;index" json:"created_at"` TokenName string `json:"token_name,omitempty"` Endpoint string `json:"endpoint"` - Content string `gorm:"type:text" json:"content,omitempty"` - GroupID string `gorm:"index" json:"group,omitempty"` - Model string `gorm:"index" json:"model"` - RequestID string `gorm:"index" json:"request_id"` - Price float64 `json:"price"` - ID int `gorm:"primaryKey" json:"id"` - CompletionPrice float64 `json:"completion_price"` - TokenID int `gorm:"index" json:"token_id,omitempty"` - UsedAmount float64 `json:"used_amount"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - ChannelID int `gorm:"index" json:"channel"` - Code int `gorm:"index" json:"code"` - Mode int `json:"mode"` - IP string `gorm:"index" json:"ip"` + Content string `gorm:"type:text" json:"content,omitempty"` + GroupID string `gorm:"index" json:"group,omitempty"` + Model string `gorm:"index" json:"model"` + RequestID string `gorm:"index" json:"request_id"` + Price float64 `json:"price,omitempty"` + ID int `gorm:"primaryKey" json:"id"` + CompletionPrice float64 `json:"completion_price,omitempty"` + TokenID int `gorm:"index" json:"token_id,omitempty"` + UsedAmount float64 `json:"used_amount,omitempty"` + PromptTokens int `json:"prompt_tokens,omitempty"` + CompletionTokens int `json:"completion_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` + ChannelID int `gorm:"index" json:"channel,omitempty"` + Code int `gorm:"index" json:"code,omitempty"` + Mode int `json:"mode,omitempty"` + IP string `gorm:"index" json:"ip,omitempty"` + RetryTimes int `json:"retry_times,omitempty"` } func CreateLogIndexes(db *gorm.DB) error { @@ -203,37 +204,49 @@ func GetGroupLogDetail(logID int, group string) (*RequestDetail, error) { return &detail, nil } -func CleanLog() error { - err := cleanLog() +const defaultCleanLogBatchSize = 1000 + +func CleanLog(batchSize int) error { + err := cleanLog(batchSize) if err != nil { return err } - return cleanLogDetail() + return cleanLogDetail(batchSize) } -func cleanLog() error { +func cleanLog(batchSize int) error { logStorageHours := config.GetLogStorageHours() if logStorageHours <= 0 { return nil } + if batchSize <= 0 { + batchSize = defaultCleanLogBatchSize + } return LogDB. + Session(&gorm.Session{SkipDefaultTransaction: true}). Where( "created_at < ?", time.Now().Add(-time.Duration(logStorageHours)*time.Hour), ). + Limit(batchSize). Delete(&Log{}).Error } -func cleanLogDetail() error { +func cleanLogDetail(batchSize int) error { detailStorageHours := config.GetLogDetailStorageHours() if detailStorageHours <= 0 { return nil } + if batchSize <= 0 { + batchSize = defaultCleanLogBatchSize + } return LogDB. + Session(&gorm.Session{SkipDefaultTransaction: true}). Where( "created_at < ?", time.Now().Add(-time.Duration(detailStorageHours)*time.Hour), ). + Limit(batchSize). Delete(&RequestDetail{}).Error } @@ -255,6 +268,7 @@ func RecordConsumeLog( content string, mode int, ip string, + retryTimes int, requestDetail *RequestDetail, ) error { log := &Log{ @@ -277,6 +291,7 @@ func RecordConsumeLog( ChannelID: channelID, Endpoint: endpoint, Content: content, + RetryTimes: retryTimes, RequestDetail: requestDetail, } return LogDB.Create(log).Error diff --git a/service/aiproxy/model/modelconfig.go b/service/aiproxy/model/modelconfig.go index 56a22255a..e4a95a7db 100644 --- a/service/aiproxy/model/modelconfig.go +++ b/service/aiproxy/model/modelconfig.go @@ -7,6 +7,7 @@ import ( "github.com/bytedance/sonic" "github.com/labring/sealos/service/aiproxy/common" + "github.com/labring/sealos/service/aiproxy/relay/relaymode" "gorm.io/gorm" ) @@ -15,7 +16,6 @@ const ( PriceUnit = 1000 ) -//nolint:revive type ModelConfig struct { CreatedAt time.Time `gorm:"index;autoCreateTime" json:"created_at"` UpdatedAt time.Time `gorm:"index;autoUpdateTime" json:"updated_at"` @@ -24,7 +24,7 @@ type ModelConfig struct { Model string `gorm:"primaryKey" json:"model"` Owner ModelOwner `gorm:"type:varchar(255);index" json:"owner"` ImageMaxBatchSize int `json:"image_batch_size,omitempty"` - Type int `json:"type"` // relaymode/define.go + Type relaymode.Mode `json:"type"` // relaymode/define.go InputPrice float64 `json:"input_price,omitempty"` OutputPrice float64 `json:"output_price,omitempty"` RPM int64 `json:"rpm,omitempty"` @@ -159,11 +159,21 @@ func SearchModelConfigs(keyword string, page int, perPage int, model string, own return configs, total, err } -func SaveModelConfig(config *ModelConfig) error { +func SaveModelConfig(config *ModelConfig) (err error) { + defer func() { + if err == nil { + _ = InitModelConfigAndChannelCache() + } + }() return DB.Save(config).Error } -func SaveModelConfigs(configs []*ModelConfig) error { +func SaveModelConfigs(configs []*ModelConfig) (err error) { + defer func() { + if err == nil { + _ = InitModelConfigAndChannelCache() + } + }() return DB.Transaction(func(tx *gorm.DB) error { for _, config := range configs { if err := tx.Save(config).Error; err != nil { diff --git a/service/aiproxy/model/option.go b/service/aiproxy/model/option.go index d5cc62fac..2868f14f5 100644 --- a/service/aiproxy/model/option.go +++ b/service/aiproxy/model/option.go @@ -13,6 +13,7 @@ import ( "github.com/bytedance/sonic" "github.com/labring/sealos/service/aiproxy/common/config" "github.com/labring/sealos/service/aiproxy/common/conv" + "github.com/labring/sealos/service/aiproxy/common/notify" log "github.com/sirupsen/logrus" ) @@ -89,6 +90,7 @@ func initOptionMap() error { } optionMap["GroupConsumeLevelRatio"] = conv.BytesToString(groupConsumeLevelRatioJSON) optionMap["InternalToken"] = config.GetInternalToken() + optionMap["NotifyNote"] = config.GetNotifyNote() optionKeys = make([]string, 0, len(optionMap)) for key := range optionMap { @@ -141,7 +143,7 @@ func SyncOptions(ctx context.Context, wg *sync.WaitGroup, frequency time.Duratio return case <-ticker.C: if err := loadOptionsFromDatabase(false); err != nil { - log.Error("failed to sync options from database: " + err.Error()) + notify.ErrorThrottle("syncOptions", time.Minute, "failed to sync options", err.Error()) } } } @@ -329,6 +331,8 @@ func updateOption(key string, value string, isInit bool) (err error) { newGroupRpmRatioMap[consumeLevel] = v } config.SetGroupConsumeLevelRatio(newGroupRpmRatioMap) + case "NotifyNote": + config.SetNotifyNote(value) default: return ErrUnknownOptionKey } diff --git a/service/aiproxy/model/owner.go b/service/aiproxy/model/owner.go index de465a3fa..a44c1a1f5 100644 --- a/service/aiproxy/model/owner.go +++ b/service/aiproxy/model/owner.go @@ -1,6 +1,5 @@ package model -//nolint:revive type ModelOwner string const ( diff --git a/service/aiproxy/model/token.go b/service/aiproxy/model/token.go index 9f281e001..6db1ec12f 100644 --- a/service/aiproxy/model/token.go +++ b/service/aiproxy/model/token.go @@ -18,8 +18,8 @@ const ( ) const ( - TokenStatusEnabled = 1 // don't use 0, 0 is the default value! - TokenStatusDisabled = 2 // also don't use 0 + TokenStatusEnabled = 1 + TokenStatusDisabled = 2 ) type Token struct { @@ -38,7 +38,6 @@ type Token struct { RequestCount int `gorm:"index" json:"request_count"` } -//nolint:goconst func getTokenOrder(order string) string { prefix, suffix, _ := strings.Cut(order, "-") switch prefix { @@ -429,7 +428,7 @@ func UpdateGroupToken(id int, group string, token *Token) (err error) { } func UpdateTokenUsedAmount(id int, amount float64, requestCount int) (err error) { - token := &Token{ID: id} + token := &Token{} defer func() { if amount > 0 && err == nil && token.Quota > 0 { if err := CacheUpdateTokenUsedAmountOnlyIncrease(token.Key, token.UsedAmount); err != nil { diff --git a/service/aiproxy/model/utils.go b/service/aiproxy/model/utils.go index 49293ae15..7f6b6eeef 100644 --- a/service/aiproxy/model/utils.go +++ b/service/aiproxy/model/utils.go @@ -1,21 +1,23 @@ package model import ( + "context" "database/sql/driver" "errors" "fmt" "strconv" "strings" + "sync" "time" + "github.com/labring/sealos/service/aiproxy/common/notify" + "github.com/shopspring/decimal" "gorm.io/gorm" "gorm.io/gorm/clause" ) -type NotFoundError string - -func (e NotFoundError) Error() string { - return string(e) + " not found" +func NotFoundError(errMsg ...string) error { + return fmt.Errorf("%s %w", strings.Join(errMsg, " "), gorm.ErrRecordNotFound) } func HandleNotFound(err error, errMsg ...string) error { @@ -42,6 +44,115 @@ func OnConflictDoNothing() *gorm.DB { }) } +func IgnoreNotFound(err error) error { + if err != nil && errors.Is(err, gorm.ErrRecordNotFound) { + return nil + } + return err +} + +type BatchUpdateData struct { + Groups map[string]*GroupUpdate + Tokens map[int]*TokenUpdate + Channels map[int]*ChannelUpdate + sync.Mutex +} + +type GroupUpdate struct { + Amount float64 + Count int +} + +type TokenUpdate struct { + Amount float64 + Count int +} + +type ChannelUpdate struct { + Amount float64 + Count int +} + +var batchData BatchUpdateData + +func init() { + batchData = BatchUpdateData{ + Groups: make(map[string]*GroupUpdate), + Tokens: make(map[int]*TokenUpdate), + Channels: make(map[int]*ChannelUpdate), + } +} + +func StartBatchProcessor(ctx context.Context, wg *sync.WaitGroup) { + defer wg.Done() + + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + ProcessBatchUpdates() + return + case <-ticker.C: + ProcessBatchUpdates() + } + } +} + +func ProcessBatchUpdates() { + batchData.Lock() + defer batchData.Unlock() + + if len(batchData.Groups) > 0 { + for groupID, data := range batchData.Groups { + err := UpdateGroupUsedAmountAndRequestCount(groupID, data.Amount, data.Count) + if IgnoreNotFound(err) != nil { + notify.ErrorThrottle( + "batchUpdateGroupUsedAmountAndRequestCount", + time.Minute, + "failed to batch update group", + err.Error(), + ) + } else { + delete(batchData.Groups, groupID) + } + } + } + + if len(batchData.Tokens) > 0 { + for tokenID, data := range batchData.Tokens { + err := UpdateTokenUsedAmount(tokenID, data.Amount, data.Count) + if IgnoreNotFound(err) != nil { + notify.ErrorThrottle( + "batchUpdateTokenUsedAmount", + time.Minute, + "failed to batch update token", + err.Error(), + ) + } else { + delete(batchData.Tokens, tokenID) + } + } + } + + if len(batchData.Channels) > 0 { + for channelID, data := range batchData.Channels { + err := UpdateChannelUsedAmount(channelID, data.Amount, data.Count) + if IgnoreNotFound(err) != nil { + notify.ErrorThrottle( + "batchUpdateChannelUsedAmount", + time.Minute, + "failed to batch update channel", + err.Error(), + ) + } else { + delete(batchData.Channels, channelID) + } + } + } +} + func BatchRecordConsume( requestID string, requestAt time.Time, @@ -60,9 +171,9 @@ func BatchRecordConsume( content string, mode int, ip string, + retryTimes int, requestDetail *RequestDetail, ) error { - errs := []error{} err := RecordConsumeLog( requestID, requestAt, @@ -81,33 +192,55 @@ func BatchRecordConsume( content, mode, ip, + retryTimes, requestDetail, ) - if err != nil { - errs = append(errs, fmt.Errorf("failed to record log: %w", err)) - } + + amountDecimal := decimal.NewFromFloat(amount) + + batchData.Lock() + defer batchData.Unlock() + if group != "" { - err = UpdateGroupUsedAmountAndRequestCount(group, amount, 1) - if err != nil { - errs = append(errs, fmt.Errorf("failed to update group used amount and request count: %w", err)) + if _, ok := batchData.Groups[group]; !ok { + batchData.Groups[group] = &GroupUpdate{} } + + if amount > 0 { + batchData.Groups[group].Amount = amountDecimal. + Add(decimal.NewFromFloat(batchData.Groups[group].Amount)). + InexactFloat64() + } + batchData.Groups[group].Count += 1 } + if tokenID > 0 { - err = UpdateTokenUsedAmount(tokenID, amount, 1) - if err != nil { - errs = append(errs, fmt.Errorf("failed to update token used amount: %w", err)) + if _, ok := batchData.Tokens[tokenID]; !ok { + batchData.Tokens[tokenID] = &TokenUpdate{} } + + if amount > 0 { + batchData.Tokens[tokenID].Amount = amountDecimal. + Add(decimal.NewFromFloat(batchData.Tokens[tokenID].Amount)). + InexactFloat64() + } + batchData.Tokens[tokenID].Count += 1 } + if channelID > 0 { - err = UpdateChannelUsedAmount(channelID, amount, 1) - if err != nil { - errs = append(errs, fmt.Errorf("failed to update channel used amount: %w", err)) + if _, ok := batchData.Channels[channelID]; !ok { + batchData.Channels[channelID] = &ChannelUpdate{} } + + if amount > 0 { + batchData.Channels[channelID].Amount = amountDecimal. + Add(decimal.NewFromFloat(batchData.Channels[channelID].Amount)). + InexactFloat64() + } + batchData.Channels[channelID].Count += 1 } - if len(errs) == 0 { - return nil - } - return errors.Join(errs...) + + return err } type EmptyNullString string diff --git a/service/aiproxy/monitor/memmodel.go b/service/aiproxy/monitor/memmodel.go new file mode 100644 index 000000000..aaa2357f2 --- /dev/null +++ b/service/aiproxy/monitor/memmodel.go @@ -0,0 +1,337 @@ +package monitor + +import ( + "context" + "sync" + "time" + + "github.com/labring/sealos/service/aiproxy/common/config" +) + +var memModelMonitor *MemModelMonitor + +func init() { + memModelMonitor = NewMemModelMonitor() +} + +const ( + timeWindow = 10 * time.Second + maxSliceCount = 12 + banDuration = 5 * time.Minute + minRequestCount = 20 + cleanupInterval = time.Minute +) + +type MemModelMonitor struct { + mu sync.RWMutex + models map[string]*ModelData +} + +type ModelData struct { + channels map[int64]*ChannelStats + totalStats *TimeWindowStats +} + +type ChannelStats struct { + timeWindows *TimeWindowStats + bannedUntil time.Time +} + +type TimeWindowStats struct { + slices []*timeSlice + mu sync.Mutex +} + +type timeSlice struct { + windowStart time.Time + requests int + errors int +} + +func NewTimeWindowStats() *TimeWindowStats { + return &TimeWindowStats{ + slices: make([]*timeSlice, 0, maxSliceCount), + } +} + +func NewMemModelMonitor() *MemModelMonitor { + mm := &MemModelMonitor{ + models: make(map[string]*ModelData), + } + + go mm.periodicCleanup() + + return mm +} + +func (m *MemModelMonitor) periodicCleanup() { + ticker := time.NewTicker(cleanupInterval) + defer ticker.Stop() + + for range ticker.C { + m.cleanupExpiredData() + } +} + +func (m *MemModelMonitor) cleanupExpiredData() { + m.mu.Lock() + defer m.mu.Unlock() + + now := time.Now() + + for modelName, modelData := range m.models { + for channelID, channelStats := range modelData.channels { + hasValidSlices := channelStats.timeWindows.HasValidSlices() + if !hasValidSlices && !channelStats.bannedUntil.After(now) { + delete(modelData.channels, channelID) + } + } + + hasValidSlices := modelData.totalStats.HasValidSlices() + if !hasValidSlices && len(modelData.channels) == 0 { + delete(m.models, modelName) + } + } +} + +func (m *MemModelMonitor) AddRequest(model string, channelID int64, isError, tryBan bool) (beyondThreshold, banExecution bool) { + m.mu.Lock() + defer m.mu.Unlock() + + now := time.Now() + + var modelData *ModelData + var exists bool + if modelData, exists = m.models[model]; !exists { + modelData = &ModelData{ + channels: make(map[int64]*ChannelStats), + totalStats: NewTimeWindowStats(), + } + m.models[model] = modelData + } + + var channel *ChannelStats + if channel, exists = modelData.channels[channelID]; !exists { + channel = &ChannelStats{ + timeWindows: NewTimeWindowStats(), + } + modelData.channels[channelID] = channel + } + + modelData.totalStats.AddRequest(now, isError) + channel.timeWindows.AddRequest(now, isError) + + return m.checkAndBan(now, channel, tryBan) +} + +func (m *MemModelMonitor) checkAndBan(now time.Time, channel *ChannelStats, tryBan bool) (beyondThreshold, banExecution bool) { + canBan := config.GetEnableModelErrorAutoBan() + if tryBan && canBan { + if channel.bannedUntil.After(now) { + return false, false + } + channel.bannedUntil = now.Add(banDuration) + return false, true + } + + req, err := channel.timeWindows.GetStats(maxSliceCount) + if req < minRequestCount { + return false, false + } + + if float64(err)/float64(req) >= config.GetModelErrorAutoBanRate() { + if !canBan || channel.bannedUntil.After(now) { + return true, false + } + channel.bannedUntil = now.Add(banDuration) + return false, true + } + return false, false +} + +func getErrorRateFromStats(stats *TimeWindowStats) float64 { + req, err := stats.GetStats(maxSliceCount) + if req < minRequestCount { + return 0 + } + return float64(err) / float64(req) +} + +func (m *MemModelMonitor) GetModelsErrorRate(ctx context.Context) (map[string]float64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string]float64) + for model, data := range m.models { + result[model] = getErrorRateFromStats(data.totalStats) + } + return result, nil +} + +func (m *MemModelMonitor) GetModelChannelErrorRate(ctx context.Context, model string) (map[int64]float64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[int64]float64) + if data, exists := m.models[model]; exists { + for channelID, channel := range data.channels { + result[channelID] = getErrorRateFromStats(channel.timeWindows) + } + } + return result, nil +} + +func (m *MemModelMonitor) GetChannelModelErrorRates(ctx context.Context, channelID int64) (map[string]float64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string]float64) + for model, data := range m.models { + if channel, exists := data.channels[channelID]; exists { + result[model] = getErrorRateFromStats(channel.timeWindows) + } + } + return result, nil +} + +func (m *MemModelMonitor) GetAllChannelModelErrorRates(ctx context.Context) (map[int64]map[string]float64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[int64]map[string]float64) + for model, data := range m.models { + for channelID, channel := range data.channels { + if _, exists := result[channelID]; !exists { + result[channelID] = make(map[string]float64) + } + result[channelID][model] = getErrorRateFromStats(channel.timeWindows) + } + } + return result, nil +} + +func (m *MemModelMonitor) GetBannedChannelsWithModel(ctx context.Context, model string) ([]int64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + var banned []int64 + if data, exists := m.models[model]; exists { + now := time.Now() + for channelID, channel := range data.channels { + if channel.bannedUntil.After(now) { + banned = append(banned, channelID) + } else { + channel.bannedUntil = time.Time{} + } + } + } + return banned, nil +} + +func (m *MemModelMonitor) GetAllBannedModelChannels(ctx context.Context) (map[string][]int64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string][]int64) + now := time.Now() + + for model, data := range m.models { + for channelID, channel := range data.channels { + if channel.bannedUntil.After(now) { + if _, exists := result[model]; !exists { + result[model] = []int64{} + } + result[model] = append(result[model], channelID) + } else { + channel.bannedUntil = time.Time{} + } + } + } + return result, nil +} + +func (m *MemModelMonitor) ClearChannelModelErrors(ctx context.Context, model string, channelID int) error { + m.mu.Lock() + defer m.mu.Unlock() + + if data, exists := m.models[model]; exists { + delete(data.channels, int64(channelID)) + } + return nil +} + +func (m *MemModelMonitor) ClearChannelAllModelErrors(ctx context.Context, channelID int) error { + m.mu.Lock() + defer m.mu.Unlock() + + for _, data := range m.models { + delete(data.channels, int64(channelID)) + } + return nil +} + +func (m *MemModelMonitor) ClearAllModelErrors(ctx context.Context) error { + m.mu.Lock() + defer m.mu.Unlock() + + m.models = make(map[string]*ModelData) + return nil +} + +func (t *TimeWindowStats) cleanupLocked(callback func(slice *timeSlice)) { + cutoff := time.Now().Add(-timeWindow * time.Duration(maxSliceCount)) + validSlices := t.slices[:0] + for _, s := range t.slices { + if s.windowStart.After(cutoff) || s.windowStart.Equal(cutoff) { + validSlices = append(validSlices, s) + if callback != nil { + callback(s) + } + } + } + t.slices = validSlices +} + +func (t *TimeWindowStats) AddRequest(now time.Time, isError bool) { + t.mu.Lock() + defer t.mu.Unlock() + + t.cleanupLocked(nil) + + currentWindow := now.Truncate(timeWindow) + var slice *timeSlice + for i := range t.slices { + if t.slices[i].windowStart.Equal(currentWindow) { + slice = t.slices[i] + break + } + } + if slice == nil { + slice = &timeSlice{windowStart: currentWindow} + t.slices = append(t.slices, slice) + } + + slice.requests++ + if isError { + slice.errors++ + } +} + +func (t *TimeWindowStats) GetStats(maxSlice int) (totalReq, totalErr int) { + t.mu.Lock() + defer t.mu.Unlock() + + t.cleanupLocked(func(slice *timeSlice) { + totalReq += slice.requests + totalErr += slice.errors + }) + return +} + +func (t *TimeWindowStats) HasValidSlices() bool { + t.mu.Lock() + defer t.mu.Unlock() + + t.cleanupLocked(nil) + return len(t.slices) > 0 +} diff --git a/service/aiproxy/monitor/model.go b/service/aiproxy/monitor/model.go index 19018e73c..767655bbb 100644 --- a/service/aiproxy/monitor/model.go +++ b/service/aiproxy/monitor/model.go @@ -10,7 +10,6 @@ import ( "github.com/labring/sealos/service/aiproxy/common" "github.com/labring/sealos/service/aiproxy/common/config" "github.com/redis/go-redis/v9" - log "github.com/sirupsen/logrus" ) // Redis key prefixes and patterns @@ -25,117 +24,30 @@ const ( // Redis scripts var ( addRequestScript = redis.NewScript(addRequestLuaScript) - getChannelModelErrorRateScript = redis.NewScript(getChannelModelErrorRateLuaScript) - getModelErrorRateScript = redis.NewScript(getModelErrorRateLuaScript) - getBannedChannelsScript = redis.NewScript(getBannedChannelsLuaScript) + getErrorRateScript = redis.NewScript(getErrorRateLuaScript) clearChannelModelErrorsScript = redis.NewScript(clearChannelModelErrorsLuaScript) clearChannelAllModelErrorsScript = redis.NewScript(clearChannelAllModelErrorsLuaScript) clearAllModelErrorsScript = redis.NewScript(clearAllModelErrorsLuaScript) ) -func buildStatsKey(model string, channelID interface{}) string { - return fmt.Sprintf("%s%s%s%v%s", modelKeyPrefix, model, channelKeyPart, channelID, statsKeySuffix) -} - // GetModelErrorRate gets error rate for a specific model across all channels func GetModelsErrorRate(ctx context.Context) (map[string]float64, error) { if !common.RedisEnabled { - return map[string]float64{}, nil + return memModelMonitor.GetModelsErrorRate(ctx) } result := make(map[string]float64) pattern := modelKeyPrefix + "*" + modelTotalStatsSuffix - iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() - for iter.Next(ctx) { - key := iter.Val() - parts := strings.Split(key, ":") - if len(parts) != 3 || parts[2] != "total_stats" { - continue - } - model := parts[1] - - rate, err := getModelErrorRateScript.Run( - ctx, - common.RDB, - []string{key}, - time.Now().UnixMilli(), - ).Float64() - if err != nil { - return nil, err - } - - result[model] = rate - } - - if err := iter.Err(); err != nil { - return nil, err - } - - return result, nil -} - -func canAutoBan() int { - if common.RedisEnabled && config.GetEnableModelErrorAutoBan() { - return 1 - } - return 0 -} - -// AddRequest adds a request record and checks if channel should be banned -func AddRequest(ctx context.Context, model string, channelID int64, isError bool) error { - if !common.RedisEnabled { - return nil - } - - errorFlag := 0 - if isError { - errorFlag = 1 - } - - now := time.Now().UnixMilli() - val, err := addRequestScript.Run( - ctx, - common.RDB, - []string{model}, - channelID, - errorFlag, - now, - config.GetModelErrorAutoBanRate(), - time.Second.Milliseconds()*15, - canAutoBan(), - ).Int64() - if err != nil { - return err - } - - log.Debugf("add request result: %d", val) - if val == 1 { - log.Errorf("channel %d model %s is banned", channelID, model) - } - return nil -} - -// GetChannelModelErrorRates gets error rates for a specific channel -func GetChannelModelErrorRates(ctx context.Context, channelID int64) (map[string]float64, error) { - if !common.RedisEnabled { - return map[string]float64{}, nil - } - - result := make(map[string]float64) - pattern := buildStatsKey("*", channelID) now := time.Now().UnixMilli() iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() for iter.Next(ctx) { key := iter.Val() - parts := strings.Split(key, ":") - if len(parts) != 5 || parts[4] != "stats" { - continue - } - model := parts[1] + model := strings.TrimPrefix(key, modelKeyPrefix) + model = strings.TrimSuffix(model, modelTotalStatsSuffix) - rate, err := getChannelModelErrorRateScript.Run( + rate, err := getErrorRateScript.Run( ctx, common.RDB, []string{key}, @@ -155,22 +67,178 @@ func GetChannelModelErrorRates(ctx context.Context, channelID int64) (map[string return result, nil } -// GetBannedChannels gets banned channels for a specific model -func GetBannedChannels(ctx context.Context, model string) ([]int64, error) { - if !common.RedisEnabled || !config.GetEnableModelErrorAutoBan() { - return []int64{}, nil +func canBan() int { + if config.GetEnableModelErrorAutoBan() { + return 1 } - result, err := getBannedChannelsScript.Run(ctx, common.RDB, []string{model}).Int64Slice() + return 0 +} + +// AddRequest adds a request record and checks if channel should be banned +func AddRequest(ctx context.Context, model string, channelID int64, isError, tryBan bool) (beyondThreshold bool, banExecution bool, err error) { + if !common.RedisEnabled { + beyondThreshold, banExecution = memModelMonitor.AddRequest(model, channelID, isError, tryBan) + return beyondThreshold, banExecution, nil + } + + errorFlag := 0 + if isError { + errorFlag = 1 + } else { + tryBan = false + } + + now := time.Now().UnixMilli() + val, err := addRequestScript.Run( + ctx, + common.RDB, + []string{model}, + channelID, + errorFlag, + now, + config.GetModelErrorAutoBanRate(), + canBan(), + tryBan, + ).Int64() if err != nil { + return false, false, err + } + return val == 3, val == 1, nil +} + +func buildStatsKey(model string, channelID string) string { + return fmt.Sprintf("%s%s%s%v%s", modelKeyPrefix, model, channelKeyPart, channelID, statsKeySuffix) +} + +func getModelChannelID(key string) (string, int64, bool) { + content := strings.TrimPrefix(key, modelKeyPrefix) + content = strings.TrimSuffix(content, statsKeySuffix) + model, channelIDStr, ok := strings.Cut(content, channelKeyPart) + if !ok { + return "", 0, false + } + channelID, err := strconv.ParseInt(channelIDStr, 10, 64) + if err != nil { + return "", 0, false + } + return model, channelID, true +} + +// GetChannelModelErrorRates gets error rates for a specific channel +func GetChannelModelErrorRates(ctx context.Context, channelID int64) (map[string]float64, error) { + if !common.RedisEnabled { + return memModelMonitor.GetChannelModelErrorRates(ctx, channelID) + } + + result := make(map[string]float64) + pattern := buildStatsKey("*", strconv.FormatInt(channelID, 10)) + now := time.Now().UnixMilli() + + iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() + for iter.Next(ctx) { + key := iter.Val() + + model, _, ok := getModelChannelID(key) + if !ok { + continue + } + + rate, err := getErrorRateScript.Run( + ctx, + common.RDB, + []string{key}, + now, + ).Float64() + if err != nil { + return nil, err + } + + result[model] = rate + } + + if err := iter.Err(); err != nil { return nil, err } + + return result, nil +} + +func GetModelChannelErrorRate(ctx context.Context, model string) (map[int64]float64, error) { + if !common.RedisEnabled { + return memModelMonitor.GetModelChannelErrorRate(ctx, model) + } + + result := make(map[int64]float64) + pattern := buildStatsKey(model, "*") + now := time.Now().UnixMilli() + + iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() + for iter.Next(ctx) { + key := iter.Val() + + _, channelID, ok := getModelChannelID(key) + if !ok { + continue + } + + rate, err := getErrorRateScript.Run( + ctx, + common.RDB, + []string{key}, + now, + ).Float64() + if err != nil { + return nil, err + } + + result[channelID] = rate + } + + if err := iter.Err(); err != nil { + return nil, err + } + + return result, nil +} + +// GetBannedChannelsWithModel gets banned channels for a specific model +func GetBannedChannelsWithModel(ctx context.Context, model string) ([]int64, error) { + if !config.GetEnableModelErrorAutoBan() { + return []int64{}, nil + } + + if !common.RedisEnabled { + return memModelMonitor.GetBannedChannelsWithModel(ctx, model) + } + + result := []int64{} + prefix := modelKeyPrefix + model + channelKeyPart + pattern := prefix + "*" + bannedKeySuffix + iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() + + for iter.Next(ctx) { + key := iter.Val() + channelIDStr := strings.TrimSuffix(strings.TrimPrefix(key, prefix), bannedKeySuffix) + + channelID, err := strconv.ParseInt(channelIDStr, 10, 64) + if err != nil { + continue + } + + result = append(result, channelID) + } + + if err := iter.Err(); err != nil { + return nil, err + } + return result, nil } // ClearChannelModelErrors clears errors for a specific channel and model func ClearChannelModelErrors(ctx context.Context, model string, channelID int) error { if !common.RedisEnabled { - return nil + return memModelMonitor.ClearChannelModelErrors(ctx, model, channelID) } return clearChannelModelErrorsScript.Run( ctx, @@ -183,7 +251,7 @@ func ClearChannelModelErrors(ctx context.Context, model string, channelID int) e // ClearChannelAllModelErrors clears all errors for a specific channel func ClearChannelAllModelErrors(ctx context.Context, channelID int) error { if !common.RedisEnabled { - return nil + return memModelMonitor.ClearChannelAllModelErrors(ctx, channelID) } return clearChannelAllModelErrorsScript.Run( ctx, @@ -196,33 +264,44 @@ func ClearChannelAllModelErrors(ctx context.Context, channelID int) error { // ClearAllModelErrors clears all error records func ClearAllModelErrors(ctx context.Context) error { if !common.RedisEnabled { - return nil + return memModelMonitor.ClearAllModelErrors(ctx) } return clearAllModelErrorsScript.Run(ctx, common.RDB, []string{}).Err() } -// GetAllBannedChannels gets all banned channels for all models -func GetAllBannedChannels(ctx context.Context) (map[string][]int64, error) { - if !common.RedisEnabled || !config.GetEnableModelErrorAutoBan() { +// GetAllBannedModelChannels gets all banned channels for all models +func GetAllBannedModelChannels(ctx context.Context) (map[string][]int64, error) { + if !config.GetEnableModelErrorAutoBan() { return map[string][]int64{}, nil } + if !common.RedisEnabled { + return memModelMonitor.GetAllBannedModelChannels(ctx) + } + result := make(map[string][]int64) - iter := common.RDB.Scan(ctx, 0, modelKeyPrefix+"*"+bannedKeySuffix, 0).Iterator() + pattern := modelKeyPrefix + "*" + channelKeyPart + "*" + bannedKeySuffix + iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() for iter.Next(ctx) { key := iter.Val() - model := strings.Split(key, ":")[1] + parts := strings.TrimPrefix(key, modelKeyPrefix) + parts = strings.TrimSuffix(parts, bannedKeySuffix) - channels, err := getBannedChannelsScript.Run( - ctx, - common.RDB, - []string{model}, - ).Int64Slice() - if err != nil { - return nil, err + model, channelIDStr, ok := strings.Cut(parts, channelKeyPart) + if !ok { + continue } - result[model] = channels + + channelID, err := strconv.ParseInt(channelIDStr, 10, 64) + if err != nil { + continue + } + + if _, exists := result[model]; !exists { + result[model] = []int64{} + } + result[model] = append(result[model], channelID) } if err := iter.Err(); err != nil { @@ -235,28 +314,23 @@ func GetAllBannedChannels(ctx context.Context) (map[string][]int64, error) { // GetAllChannelModelErrorRates gets error rates for all channels and models func GetAllChannelModelErrorRates(ctx context.Context) (map[int64]map[string]float64, error) { if !common.RedisEnabled { - return map[int64]map[string]float64{}, nil + return memModelMonitor.GetAllChannelModelErrorRates(ctx) } result := make(map[int64]map[string]float64) - pattern := modelKeyPrefix + "*" + channelKeyPart + "*" + statsKeySuffix + pattern := buildStatsKey("*", "*") now := time.Now().UnixMilli() iter := common.RDB.Scan(ctx, 0, pattern, 0).Iterator() for iter.Next(ctx) { key := iter.Val() - parts := strings.Split(key, ":") - if len(parts) != 5 || parts[4] != "stats" { + + model, channelID, ok := getModelChannelID(key) + if !ok { continue } - model := parts[1] - channelID, err := strconv.ParseInt(parts[3], 10, 64) - if err != nil { - continue - } - - rate, err := getChannelModelErrorRateScript.Run( + rate, err := getErrorRateScript.Run( ctx, common.RDB, []string{key}, @@ -287,18 +361,16 @@ local channel_id = ARGV[1] local is_error = tonumber(ARGV[2]) local now_ts = tonumber(ARGV[3]) local max_error_rate = tonumber(ARGV[4]) -local statsExpiry = tonumber(ARGV[5]) -local can_auto_ban = tonumber(ARGV[6]) +local can_ban = tonumber(ARGV[5]) +local try_ban = tonumber(ARGV[6]) -local banned_key = "model:" .. model .. ":banned" +local banned_key = "model:" .. model .. ":channel:" .. channel_id .. ":banned" local stats_key = "model:" .. model .. ":channel:" .. channel_id .. ":stats" local model_stats_key = "model:" .. model .. ":total_stats" -local maxSliceCount = 6 -local current_slice = math.floor(now_ts / 1000) - -if redis.call("SISMEMBER", banned_key, channel_id) == 1 then - return 2 -end +local maxSliceCount = 12 +local statsExpiry = maxSliceCount * 10 * 1000 +local banExpiry = 5 * 60 * 1000 +local current_slice = math.floor(now_ts / 10 / 1000) local function parse_req_err(value) if not value then return 0, 0 end @@ -306,102 +378,72 @@ local function parse_req_err(value) return tonumber(r) or 0, tonumber(e) or 0 end -local function update_channel_stats() - local req, err = parse_req_err(redis.call("HGET", stats_key, current_slice)) +local function update_stats(key) + local req, err = parse_req_err(redis.call("HGET", key, current_slice)) req = req + 1 err = err + (is_error == 1 and 1 or 0) - redis.call("HSET", stats_key, current_slice, req .. ":" .. err) - redis.call("PEXPIRE", stats_key, statsExpiry) + redis.call("HSET", key, current_slice, req .. ":" .. err) + redis.call("PEXPIRE", key, statsExpiry) return req, err end -local function update_model_stats() - local req, err = parse_req_err(redis.call("HGET", model_stats_key, current_slice)) - req = req + 1 - err = err + (is_error == 1 and 1 or 0) - redis.call("HSET", model_stats_key, current_slice, req .. ":" .. err) - redis.call("PEXPIRE", model_stats_key, statsExpiry) - return req, err -end - -update_channel_stats() -update_model_stats() - -if is_error == 0 or can_auto_ban == 0 then - return 0 -end - -local function check_channel_error() - local total_req, total_err = 0, 0 - local min_valid_slice = current_slice - maxSliceCount - - local all_slices = redis.call("HGETALL", stats_key) - local to_delete = {} - +local function get_clean_req_err(key) + local total_req, total_err = 0, 0 + local min_valid_slice = current_slice - maxSliceCount + local all_slices = redis.call("HGETALL", key) for i = 1, #all_slices, 2 do local slice = tonumber(all_slices[i]) if slice < min_valid_slice then - table.insert(to_delete, all_slices[i]) - else - local req, err = parse_req_err(all_slices[i+1]) - total_req = total_req + req - total_err = total_err + err - end + redis.call("HDEL", key, all_slices[i]) + else + local req, err = parse_req_err(all_slices[i+1]) + total_req = total_req + req + total_err = total_err + err + end end - - if #to_delete > 0 then - redis.call("HDEL", stats_key, unpack(to_delete)) - end - - if total_req >= 10 and (total_err / total_req) >= max_error_rate then - redis.call("SADD", banned_key, channel_id) - redis.call("DEL", stats_key) - return true - end - return false + return total_req, total_err end -if check_channel_error() then - return 1 +update_stats(stats_key) +update_stats(model_stats_key) + +local function check_channel_error() + local already_banned = redis.call("EXISTS", banned_key) == 1 + + if try_ban == 1 and can_ban == 1 then + if already_banned then + return 2 + end + redis.call("SET", banned_key, 1) + redis.call("PEXPIRE", banned_key, banExpiry) + return 1 + end + + local total_req, total_err = get_clean_req_err(stats_key) + if total_req < 20 then + return 0 + end + + if (total_err / total_req) < max_error_rate then + return 0 + else + if can_ban == 0 or already_banned then + return 3 + end + redis.call("SET", banned_key, 1) + redis.call("PEXPIRE", banned_key, banExpiry) + return 1 + end end -return 0 + +return check_channel_error() ` - getModelErrorRateLuaScript = ` -local model_stats_key = KEYS[1] -local now_ts = tonumber(ARGV[1]) -local maxSliceCount = 6 -local current_slice = math.floor(now_ts / 1000) -local min_valid_slice = current_slice - maxSliceCount - -local function parse_req_err(value) - if not value then return 0, 0 end - local r, e = value:match("^(%d+):(%d+)$") - return tonumber(r) or 0, tonumber(e) or 0 -end - -local total_req, total_err = 0, 0 -local all_slices = redis.call("HGETALL", model_stats_key) - -for i = 1, #all_slices, 2 do - local slice = tonumber(all_slices[i]) - if slice >= min_valid_slice then - local req, err = parse_req_err(all_slices[i+1]) - total_req = total_req + req - total_err = total_err + err - end -end - -if total_req == 0 then return 0 end -return total_err / total_req -` - - getChannelModelErrorRateLuaScript = ` + getErrorRateLuaScript = ` local stats_key = KEYS[1] local now_ts = tonumber(ARGV[1]) -local maxSliceCount = 6 -local current_slice = math.floor(now_ts / 1000) -local min_valid_slice = current_slice - maxSliceCount +local maxSliceCount = 12 +local current_slice = math.floor(now_ts / 10 / 1000) local function parse_req_err(value) if not value then return 0, 0 end @@ -409,50 +451,52 @@ local function parse_req_err(value) return tonumber(r) or 0, tonumber(e) or 0 end -local total_req, total_err = 0, 0 -local all_slices = redis.call("HGETALL", stats_key) - -for i = 1, #all_slices, 2 do - local slice = tonumber(all_slices[i]) - if slice >= min_valid_slice then - local req, err = parse_req_err(all_slices[i+1]) - total_req = total_req + req - total_err = total_err + err +local function get_clean_req_err(key) + local total_req, total_err = 0, 0 + local min_valid_slice = current_slice - maxSliceCount + local all_slices = redis.call("HGETALL", key) + for i = 1, #all_slices, 2 do + local slice = tonumber(all_slices[i]) + if slice < min_valid_slice then + redis.call("HDEL", key, all_slices[i]) + else + local req, err = parse_req_err(all_slices[i+1]) + total_req = total_req + req + total_err = total_err + err + end end + return total_req, total_err end -if total_req == 0 then return 0 end +local total_req, total_err = get_clean_req_err(stats_key) +if total_req < 20 then return 0 end return string.format("%.2f", total_err / total_req) -` - - getBannedChannelsLuaScript = ` -local model = KEYS[1] -return redis.call("SMEMBERS", "model:" .. model .. ":banned") ` clearChannelModelErrorsLuaScript = ` local model = KEYS[1] local channel_id = ARGV[1] local stats_key = "model:" .. model .. ":channel:" .. channel_id .. ":stats" -local banned_key = "model:" .. model .. ":banned" +local banned_key = "model:" .. model .. ":channel:" .. channel_id .. ":banned" redis.call("DEL", stats_key) -redis.call("SREM", banned_key, channel_id) +redis.call("DEL", banned_key) return redis.status_reply("ok") ` clearChannelAllModelErrorsLuaScript = ` -local channel_id = ARGV[1] -local pattern = "model:*:channel:" .. channel_id .. ":stats" -local keys = redis.call("KEYS", pattern) - -for _, key in ipairs(keys) do - redis.call("DEL", key) - local model = string.match(key, "model:(.*):channel:") - if model then - redis.call("SREM", "model:"..model..":banned", channel_id) - end +local function del_keys(pattern) + local keys = redis.call("KEYS", pattern) + if #keys > 0 then redis.call("DEL", unpack(keys)) end end + +local channel_id = ARGV[1] +local stats_pattern = "model:*:channel:" .. channel_id .. ":stats" +local banned_pattern = "model:*:channel:" .. channel_id .. ":banned" + +del_keys(stats_pattern) +del_keys(banned_pattern) + return redis.status_reply("ok") ` @@ -463,7 +507,7 @@ local function del_keys(pattern) end del_keys("model:*:channel:*:stats") -del_keys("model:*:banned") +del_keys("model:*:channel:*:banned") return redis.status_reply("ok") ` diff --git a/service/aiproxy/relay/adaptor/ali/adaptor.go b/service/aiproxy/relay/adaptor/ali/adaptor.go index 81191c772..2a4398a2a 100644 --- a/service/aiproxy/relay/adaptor/ali/adaptor.go +++ b/service/aiproxy/relay/adaptor/ali/adaptor.go @@ -2,7 +2,7 @@ package ali import ( "bytes" - "errors" + "fmt" "io" "net/http" "strings" @@ -45,7 +45,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.Rerank: return u + "/api/v1/services/rerank/text-rerank/text-rerank", nil default: - return "", errors.New("unsupported mode") + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -69,7 +69,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht case relaymode.AudioTranscription: return ConvertSTTRequest(meta, req) default: - return "", nil, nil, errors.New("unsupported convert request mode") + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -114,7 +114,7 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons case relaymode.AudioTranscription: usage, err = STTDoResponse(meta, c, resp) default: - return nil, openai.ErrorWrapperWithMessage("unsupported response mode", "unsupported_mode", http.StatusBadRequest) + return nil, openai.ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %s", meta.Mode), "unsupported_mode", http.StatusBadRequest) } return } diff --git a/service/aiproxy/relay/adaptor/ali/embeddings.go b/service/aiproxy/relay/adaptor/ali/embeddings.go index b5d9fda3e..d1fc26859 100644 --- a/service/aiproxy/relay/adaptor/ali/embeddings.go +++ b/service/aiproxy/relay/adaptor/ali/embeddings.go @@ -65,7 +65,7 @@ func embeddingResponse2OpenAI(meta *meta.Meta, response *EmbeddingResponse) *ope for i, embedding := range response.Output.Embeddings { openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, &openai.EmbeddingResponseItem{ - Object: `embedding`, + Object: "embedding", Index: i, Embedding: embedding.Embedding, }) diff --git a/service/aiproxy/relay/adaptor/ali/image.go b/service/aiproxy/relay/adaptor/ali/image.go index d57333828..b36bb441e 100644 --- a/service/aiproxy/relay/adaptor/ali/image.go +++ b/service/aiproxy/relay/adaptor/ali/image.go @@ -48,6 +48,12 @@ func ConvertImageRequest(meta *meta.Meta, req *http.Request) (string, http.Heade } func ImageHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + + defer resp.Body.Close() + log := middleware.GetLogger(c) responseFormat := meta.MustGet(MetaResponseFormat).(string) @@ -57,10 +63,6 @@ func ImageHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model. if err != nil { return nil, openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) } - err = resp.Body.Close() - if err != nil { - return nil, openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError) - } err = sonic.Unmarshal(responseBody, &aliTaskResponse) if err != nil { return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) @@ -81,7 +83,6 @@ func ImageHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model. Error: model.Error{ Message: aliResponse.Output.Message, Type: "ali_error", - Param: "", Code: aliResponse.Output.Code, }, StatusCode: resp.StatusCode, diff --git a/service/aiproxy/relay/adaptor/ali/rerank.go b/service/aiproxy/relay/adaptor/ali/rerank.go index 47cec3510..28b2c0404 100644 --- a/service/aiproxy/relay/adaptor/ali/rerank.go +++ b/service/aiproxy/relay/adaptor/ali/rerank.go @@ -56,6 +56,10 @@ func ConvertRerankRequest(meta *meta.Meta, req *http.Request) (string, http.Head } func RerankHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/anthropic/adaptor.go b/service/aiproxy/relay/adaptor/anthropic/adaptor.go index 975e9d26f..96dc394e9 100644 --- a/service/aiproxy/relay/adaptor/anthropic/adaptor.go +++ b/service/aiproxy/relay/adaptor/anthropic/adaptor.go @@ -67,9 +67,9 @@ func (a *Adaptor) DoRequest(_ *meta.Meta, _ *gin.Context, req *http.Request) (*h func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Response) (usage *relaymodel.Usage, err *relaymodel.ErrorWithStatusCode) { if utils.IsStreamResponse(resp) { - err, usage = StreamHandler(meta, c, resp) + usage, err = StreamHandler(meta, c, resp) } else { - err, usage = Handler(meta, c, resp) + usage, err = Handler(meta, c, resp) } return } diff --git a/service/aiproxy/relay/adaptor/anthropic/main.go b/service/aiproxy/relay/adaptor/anthropic/main.go index 787dbbdf7..742c31492 100644 --- a/service/aiproxy/relay/adaptor/anthropic/main.go +++ b/service/aiproxy/relay/adaptor/anthropic/main.go @@ -104,7 +104,10 @@ func ConvertRequest(meta *meta.Meta, req *http.Request) (*Request, error) { if claudeRequest.Thinking != nil { if claudeRequest.Thinking.BudgetTokens == 0 || claudeRequest.Thinking.BudgetTokens >= claudeRequest.MaxTokens { - claudeRequest.Thinking.BudgetTokens = claudeRequest.MaxTokens / 3 * 2 + claudeRequest.Thinking.BudgetTokens = claudeRequest.MaxTokens / 2 + } + if claudeRequest.Thinking.BudgetTokens < 1024 { + claudeRequest.Thinking.BudgetTokens = 1024 } claudeRequest.Temperature = nil } @@ -316,7 +319,11 @@ func ResponseClaude2OpenAI(meta *meta.Meta, claudeResponse *Response) *openai.Te return &fullTextResponse } -func StreamHandler(m *meta.Meta, c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { +func StreamHandler(m *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) @@ -413,35 +420,28 @@ func StreamHandler(m *meta.Meta, c *gin.Context, resp *http.Response) (*model.Er render.Done(c) - return nil, &usage + return &usage, nil } -func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { +func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() var claudeResponse Response err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&claudeResponse) if err != nil { - return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil - } - if claudeResponse.Error.Type != "" { - return &model.ErrorWithStatusCode{ - Error: model.Error{ - Message: claudeResponse.Error.Message, - Type: claudeResponse.Error.Type, - Param: "", - Code: claudeResponse.Error.Type, - }, - StatusCode: resp.StatusCode, - }, nil + return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } fullTextResponse := ResponseClaude2OpenAI(meta, &claudeResponse) jsonResponse, err := sonic.Marshal(fullTextResponse) if err != nil { - return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError) } c.Writer.Header().Set("Content-Type", "application/json") c.Writer.WriteHeader(resp.StatusCode) _, _ = c.Writer.Write(jsonResponse) - return nil, &fullTextResponse.Usage + return &fullTextResponse.Usage, nil } diff --git a/service/aiproxy/relay/adaptor/aws/claude/main.go b/service/aiproxy/relay/adaptor/aws/claude/main.go index 7812700d8..1e5db012b 100644 --- a/service/aiproxy/relay/adaptor/aws/claude/main.go +++ b/service/aiproxy/relay/adaptor/aws/claude/main.go @@ -6,11 +6,10 @@ import ( "net/http" "time" - "github.com/bytedance/sonic" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" + "github.com/bytedance/sonic" "github.com/gin-gonic/gin" "github.com/jinzhu/copier" "github.com/labring/sealos/service/aiproxy/common/render" diff --git a/service/aiproxy/relay/adaptor/aws/llama3/main.go b/service/aiproxy/relay/adaptor/aws/llama3/main.go index e6cc97426..7b62f9bf6 100644 --- a/service/aiproxy/relay/adaptor/aws/llama3/main.go +++ b/service/aiproxy/relay/adaptor/aws/llama3/main.go @@ -8,11 +8,10 @@ import ( "text/template" "time" - "github.com/bytedance/sonic" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" + "github.com/bytedance/sonic" "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common/random" "github.com/labring/sealos/service/aiproxy/common/render" diff --git a/service/aiproxy/relay/adaptor/azure/constants.go b/service/aiproxy/relay/adaptor/azure/constants.go index 1f0ac7c38..86a2cf76d 100644 --- a/service/aiproxy/relay/adaptor/azure/constants.go +++ b/service/aiproxy/relay/adaptor/azure/constants.go @@ -39,8 +39,12 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.ChatCompletions: // https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api return fmt.Sprintf("%s/openai/deployments/%s/chat/completions?api-version=%s", meta.Channel.BaseURL, model, apiVersion), nil + case relaymode.Completions: + return fmt.Sprintf("%s/openai/deployments/%s/completions?api-version=%s", meta.Channel.BaseURL, model, apiVersion), nil + case relaymode.Embeddings: + return fmt.Sprintf("%s/openai/deployments/%s/embeddings?api-version=%s", meta.Channel.BaseURL, model, apiVersion), nil default: - return "", fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } } diff --git a/service/aiproxy/relay/adaptor/azure/key.go b/service/aiproxy/relay/adaptor/azure/key.go index 7d4596b82..dde766f7e 100644 --- a/service/aiproxy/relay/adaptor/azure/key.go +++ b/service/aiproxy/relay/adaptor/azure/key.go @@ -21,10 +21,12 @@ func (a *Adaptor) KeyHelp() string { return "key or key|api-version" } +const defaultAPIVersion = "2024-02-01" + func getTokenAndAPIVersion(key string) (string, string, error) { split := strings.Split(key, "|") if len(split) == 1 { - return key, "", nil + return key, defaultAPIVersion, nil } if len(split) != 2 { return "", "", errors.New("invalid key format") diff --git a/service/aiproxy/relay/adaptor/baidu/adaptor.go b/service/aiproxy/relay/adaptor/baidu/adaptor.go index f2db94be4..b3a607bc7 100644 --- a/service/aiproxy/relay/adaptor/baidu/adaptor.go +++ b/service/aiproxy/relay/adaptor/baidu/adaptor.go @@ -7,14 +7,13 @@ import ( "net/http" "strings" + "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" + relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" "github.com/labring/sealos/service/aiproxy/relay/relaymode" "github.com/labring/sealos/service/aiproxy/relay/utils" - - "github.com/gin-gonic/gin" - relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" ) type Adaptor struct{} @@ -99,8 +98,10 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht return openai.ConvertRequest(meta, req) case relaymode.ImagesGenerations: return openai.ConvertRequest(meta, req) - default: + case relaymode.ChatCompletions: return ConvertRequest(meta, req) + default: + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -116,12 +117,14 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons usage, err = RerankHandler(meta, c, resp) case relaymode.ImagesGenerations: usage, err = ImageHandler(meta, c, resp) - default: + case relaymode.ChatCompletions: if utils.IsStreamResponse(resp) { err, usage = StreamHandler(meta, c, resp) } else { usage, err = Handler(meta, c, resp) } + default: + return nil, openai.ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %s", meta.Mode), "unsupported_mode", http.StatusBadRequest) } return } diff --git a/service/aiproxy/relay/adaptor/baidu/embeddings.go b/service/aiproxy/relay/adaptor/baidu/embeddings.go index aa5d877b6..c9a9cf685 100644 --- a/service/aiproxy/relay/adaptor/baidu/embeddings.go +++ b/service/aiproxy/relay/adaptor/baidu/embeddings.go @@ -3,7 +3,6 @@ package baidu import ( "io" "net/http" - "strconv" "github.com/bytedance/sonic" "github.com/gin-gonic/gin" @@ -32,8 +31,8 @@ func EmbeddingsHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*r if err != nil { return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } - if baiduResponse.Error != nil && baiduResponse.Error.ErrorCode != 0 { - return &baiduResponse.Usage, openai.ErrorWrapperWithMessage(baiduResponse.Error.ErrorMsg, "baidu_error_"+strconv.Itoa(baiduResponse.Error.ErrorCode), http.StatusInternalServerError) + if baiduResponse.Error != nil && baiduResponse.ErrorCode != 0 { + return &baiduResponse.Usage, ErrorHandler(baiduResponse.Error) } respMap := make(map[string]any) diff --git a/service/aiproxy/relay/adaptor/baidu/error.go b/service/aiproxy/relay/adaptor/baidu/error.go new file mode 100644 index 000000000..c311dafa2 --- /dev/null +++ b/service/aiproxy/relay/adaptor/baidu/error.go @@ -0,0 +1,48 @@ +package baidu + +import ( + "net/http" + "strconv" + + "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" + relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" +) + +// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/tlmyncueh + +func ErrorHandler(baiduError *Error) *relaymodel.ErrorWithStatusCode { + switch baiduError.ErrorCode { + case 13, 14, 100, 110: + return openai.ErrorWrapperWithMessage( + baiduError.ErrorMsg, + "upstream_"+strconv.Itoa(baiduError.ErrorCode), + http.StatusUnauthorized, + ) + case 17, 19, 111: + return openai.ErrorWrapperWithMessage( + baiduError.ErrorMsg, + "upstream_"+strconv.Itoa(baiduError.ErrorCode), + http.StatusForbidden, + ) + case 336001, 336002, 336003, + 336005, 336006, 336007, + 336008, 336103, 336104, + 336106, 336118, 336122, + 336123, 336221, 337006, + 337008, 337009: + return openai.ErrorWrapperWithMessage( + baiduError.ErrorMsg, + "upstream_"+strconv.Itoa(baiduError.ErrorCode), + http.StatusBadRequest, + ) + case 4, 18, 336117, 336501, 336502, + 336503, 336504, 336505, + 336507: + return openai.ErrorWrapperWithMessage( + baiduError.ErrorMsg, + "upstream_"+strconv.Itoa(baiduError.ErrorCode), + http.StatusTooManyRequests, + ) + } + return openai.ErrorWrapperWithMessage(baiduError.ErrorMsg, "upstream_"+strconv.Itoa(baiduError.ErrorCode), http.StatusInternalServerError) +} diff --git a/service/aiproxy/relay/adaptor/baidu/image.go b/service/aiproxy/relay/adaptor/baidu/image.go index 051959f26..8a717314d 100644 --- a/service/aiproxy/relay/adaptor/baidu/image.go +++ b/service/aiproxy/relay/adaptor/baidu/image.go @@ -3,7 +3,6 @@ package baidu import ( "io" "net/http" - "strconv" "github.com/bytedance/sonic" "github.com/gin-gonic/gin" @@ -45,7 +44,7 @@ func ImageHandler(_ *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usa } if imageResponse.Error != nil && imageResponse.Error.ErrorMsg != "" { - return usage, openai.ErrorWrapperWithMessage(imageResponse.Error.ErrorMsg, "baidu_error_"+strconv.Itoa(imageResponse.Error.ErrorCode), http.StatusBadRequest) + return usage, ErrorHandler(imageResponse.Error) } openaiResponse := ToOpenAIImageResponse(&imageResponse) diff --git a/service/aiproxy/relay/adaptor/baidu/main.go b/service/aiproxy/relay/adaptor/baidu/main.go index 30a4198e0..3d7839c2c 100644 --- a/service/aiproxy/relay/adaptor/baidu/main.go +++ b/service/aiproxy/relay/adaptor/baidu/main.go @@ -5,15 +5,13 @@ import ( "bytes" "io" "net/http" - "strconv" "github.com/bytedance/sonic" + "github.com/gin-gonic/gin" + "github.com/labring/sealos/service/aiproxy/common" "github.com/labring/sealos/service/aiproxy/common/conv" "github.com/labring/sealos/service/aiproxy/common/render" "github.com/labring/sealos/service/aiproxy/middleware" - - "github.com/gin-gonic/gin" - "github.com/labring/sealos/service/aiproxy/common" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/constant" "github.com/labring/sealos/service/aiproxy/relay/meta" @@ -179,7 +177,7 @@ func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } if baiduResponse.Error != nil && baiduResponse.Error.ErrorCode != 0 { - return nil, openai.ErrorWrapperWithMessage(baiduResponse.Error.ErrorMsg, "baidu_error_"+strconv.Itoa(baiduResponse.Error.ErrorCode), http.StatusInternalServerError) + return nil, ErrorHandler(baiduResponse.Error) } fullTextResponse := responseBaidu2OpenAI(&baiduResponse) fullTextResponse.Model = meta.OriginModel diff --git a/service/aiproxy/relay/adaptor/baidu/rerank.go b/service/aiproxy/relay/adaptor/baidu/rerank.go index de4a87acc..0e52433a2 100644 --- a/service/aiproxy/relay/adaptor/baidu/rerank.go +++ b/service/aiproxy/relay/adaptor/baidu/rerank.go @@ -3,7 +3,6 @@ package baidu import ( "io" "net/http" - "strconv" "github.com/bytedance/sonic" "github.com/gin-gonic/gin" @@ -33,7 +32,7 @@ func RerankHandler(_ *meta.Meta, c *gin.Context, resp *http.Response) (*model.Us return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } if reRankResp.Error != nil && reRankResp.Error.ErrorCode != 0 { - return nil, openai.ErrorWrapperWithMessage(reRankResp.Error.ErrorMsg, "baidu_error_"+strconv.Itoa(reRankResp.Error.ErrorCode), http.StatusInternalServerError) + return nil, ErrorHandler(reRankResp.Error) } respMap := make(map[string]any) err = sonic.Unmarshal(respBody, &respMap) diff --git a/service/aiproxy/relay/adaptor/baiduv2/adaptor.go b/service/aiproxy/relay/adaptor/baiduv2/adaptor.go index 59b549817..c8273f321 100644 --- a/service/aiproxy/relay/adaptor/baiduv2/adaptor.go +++ b/service/aiproxy/relay/adaptor/baiduv2/adaptor.go @@ -44,7 +44,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.ChatCompletions: return meta.Channel.BaseURL + "/chat/completions", nil default: - return "", fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -68,7 +68,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht } return openai.ConvertRequest(meta, req) default: - return "", nil, nil, fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -82,7 +82,7 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons return openai.DoResponse(meta, c, resp) default: return nil, openai.ErrorWrapperWithMessage( - fmt.Sprintf("unsupported mode: %d", meta.Mode), + fmt.Sprintf("unsupported mode: %s", meta.Mode), nil, http.StatusBadRequest, ) diff --git a/service/aiproxy/relay/adaptor/cohere/adaptor.go b/service/aiproxy/relay/adaptor/cohere/adaptor.go index b5b45a7c4..5b8234f21 100644 --- a/service/aiproxy/relay/adaptor/cohere/adaptor.go +++ b/service/aiproxy/relay/adaptor/cohere/adaptor.go @@ -60,9 +60,9 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons usage, err = openai.RerankHandler(meta, c, resp) default: if utils.IsStreamResponse(resp) { - err, usage = StreamHandler(c, resp) + usage, err = StreamHandler(c, resp) } else { - err, usage = Handler(c, resp, meta.InputTokens, meta.ActualModel) + usage, err = Handler(c, resp, meta.InputTokens, meta.ActualModel) } } return diff --git a/service/aiproxy/relay/adaptor/cohere/main.go b/service/aiproxy/relay/adaptor/cohere/main.go index dd0a01745..79250f7f9 100644 --- a/service/aiproxy/relay/adaptor/cohere/main.go +++ b/service/aiproxy/relay/adaptor/cohere/main.go @@ -131,7 +131,11 @@ func ResponseCohere2OpenAI(cohereResponse *Response) *openai.TextResponse { return &fullTextResponse } -func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { +func StreamHandler(c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) @@ -177,27 +181,23 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC render.Done(c) - return nil, &usage + return &usage, nil } -func Handler(c *gin.Context, resp *http.Response, _ int, modelName string) (*model.ErrorWithStatusCode, *model.Usage) { +func Handler(c *gin.Context, resp *http.Response, _ int, modelName string) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() var cohereResponse Response err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&cohereResponse) if err != nil { - return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } if cohereResponse.ResponseID == "" { - return &model.ErrorWithStatusCode{ - Error: model.Error{ - Message: cohereResponse.Message, - Type: cohereResponse.Message, - Param: "", - Code: resp.StatusCode, - }, - StatusCode: resp.StatusCode, - }, nil + return nil, openai.ErrorWrapperWithMessage(cohereResponse.Message, resp.StatusCode, resp.StatusCode) } fullTextResponse := ResponseCohere2OpenAI(&cohereResponse) fullTextResponse.Model = modelName @@ -209,10 +209,10 @@ func Handler(c *gin.Context, resp *http.Response, _ int, modelName string) (*mod fullTextResponse.Usage = usage jsonResponse, err := sonic.Marshal(fullTextResponse) if err != nil { - return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError) } c.Writer.Header().Set("Content-Type", "application/json") c.Writer.WriteHeader(resp.StatusCode) _, _ = c.Writer.Write(jsonResponse) - return nil, &usage + return &usage, nil } diff --git a/service/aiproxy/relay/adaptor/coze/adaptor.go b/service/aiproxy/relay/adaptor/coze/adaptor.go index 3fce72f2d..b4854f37d 100644 --- a/service/aiproxy/relay/adaptor/coze/adaptor.go +++ b/service/aiproxy/relay/adaptor/coze/adaptor.go @@ -10,7 +10,6 @@ import ( "github.com/bytedance/sonic" "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" - "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" "github.com/labring/sealos/service/aiproxy/relay/relaymode" @@ -80,19 +79,11 @@ func (a *Adaptor) DoRequest(_ *meta.Meta, _ *gin.Context, req *http.Request) (*h } func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Response) (usage *relaymodel.Usage, err *relaymodel.ErrorWithStatusCode) { - var responseText *string if utils.IsStreamResponse(resp) { - err, responseText = StreamHandler(meta, c, resp) + usage, err = StreamHandler(meta, c, resp) } else { - err, responseText = Handler(meta, c, resp) + usage, err = Handler(meta, c, resp) } - if responseText != nil { - usage = openai.ResponseText2Usage(*responseText, meta.ActualModel, meta.InputTokens) - } else { - usage = &relaymodel.Usage{} - } - usage.PromptTokens = meta.InputTokens - usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens return } diff --git a/service/aiproxy/relay/adaptor/coze/main.go b/service/aiproxy/relay/adaptor/coze/main.go index 2ae4d25de..c356d179e 100644 --- a/service/aiproxy/relay/adaptor/coze/main.go +++ b/service/aiproxy/relay/adaptor/coze/main.go @@ -86,7 +86,11 @@ func ResponseCoze2OpenAI(cozeResponse *Response) *openai.TextResponse { return &fullTextResponse } -func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *string) { +func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) @@ -136,10 +140,14 @@ func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model render.Done(c) - return nil, &responseText + return openai.ResponseText2Usage(responseText, meta.ActualModel, meta.InputTokens), nil } -func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *string) { +func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) @@ -147,22 +155,16 @@ func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Error var cozeResponse Response err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&cozeResponse) if err != nil { - return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } if cozeResponse.Code != 0 { - return &model.ErrorWithStatusCode{ - Error: model.Error{ - Message: cozeResponse.Msg, - Code: cozeResponse.Code, - }, - StatusCode: resp.StatusCode, - }, nil + return nil, openai.ErrorWrapperWithMessage(cozeResponse.Msg, cozeResponse.Code, resp.StatusCode) } fullTextResponse := ResponseCoze2OpenAI(&cozeResponse) fullTextResponse.Model = meta.OriginModel jsonResponse, err := sonic.Marshal(fullTextResponse) if err != nil { - return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError) } c.Writer.Header().Set("Content-Type", "application/json") c.Writer.WriteHeader(resp.StatusCode) @@ -174,5 +176,5 @@ func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Error if len(fullTextResponse.Choices) > 0 { responseText = fullTextResponse.Choices[0].Message.StringContent() } - return nil, &responseText + return openai.ResponseText2Usage(responseText, meta.ActualModel, meta.InputTokens), nil } diff --git a/service/aiproxy/relay/adaptor/doc2x/adaptor.go b/service/aiproxy/relay/adaptor/doc2x/adaptor.go index 22998c703..b0e9712ef 100644 --- a/service/aiproxy/relay/adaptor/doc2x/adaptor.go +++ b/service/aiproxy/relay/adaptor/doc2x/adaptor.go @@ -30,7 +30,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.ParsePdf: return meta.Channel.BaseURL + "/api/v2/parse/pdf", nil default: - return "", fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -39,7 +39,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht case relaymode.ParsePdf: return ConvertParsePdfRequest(meta, req) default: - return "", nil, nil, fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -52,7 +52,7 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons case relaymode.ParsePdf: return HandleParsePdfResponse(meta, c, resp) default: - return nil, openai.ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %d", meta.Mode), "unsupported_mode", http.StatusBadRequest) + return nil, openai.ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %s", meta.Mode), "unsupported_mode", http.StatusBadRequest) } } diff --git a/service/aiproxy/relay/adaptor/doubao/main.go b/service/aiproxy/relay/adaptor/doubao/main.go index 7ab31fd75..b1e8a3a50 100644 --- a/service/aiproxy/relay/adaptor/doubao/main.go +++ b/service/aiproxy/relay/adaptor/doubao/main.go @@ -10,6 +10,7 @@ import ( "github.com/bytedance/sonic" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" @@ -84,3 +85,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht func (a *Adaptor) GetChannelName() string { return "doubao" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/adaptor/doubaoaudio/main.go b/service/aiproxy/relay/adaptor/doubaoaudio/main.go index f653a821a..96e31e338 100644 --- a/service/aiproxy/relay/adaptor/doubaoaudio/main.go +++ b/service/aiproxy/relay/adaptor/doubaoaudio/main.go @@ -19,7 +19,7 @@ func GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.AudioSpeech: return u + "/api/v1/tts/ws_binary", nil default: - return "", fmt.Errorf("unsupported relay mode %d for doubao", meta.Mode) + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -44,7 +44,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht case relaymode.AudioSpeech: return ConvertTTSRequest(meta, req) default: - return "", nil, nil, fmt.Errorf("unsupported relay mode %d for doubao", meta.Mode) + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -58,7 +58,7 @@ func (a *Adaptor) SetupRequestHeader(meta *meta.Meta, _ *gin.Context, req *http. req.Header.Set("Authorization", "Bearer;"+token) return nil default: - return fmt.Errorf("unsupported relay mode %d for doubao", meta.Mode) + return fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -67,7 +67,7 @@ func (a *Adaptor) DoRequest(meta *meta.Meta, _ *gin.Context, req *http.Request) case relaymode.AudioSpeech: return TTSDoRequest(meta, req) default: - return nil, fmt.Errorf("unsupported relay mode %d for doubao", meta.Mode) + return nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -77,7 +77,7 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons return TTSDoResponse(meta, c, resp) default: return nil, openai.ErrorWrapperWithMessage( - fmt.Sprintf("unsupported relay mode %d for doubao", meta.Mode), + fmt.Sprintf("unsupported mode: %s", meta.Mode), nil, http.StatusBadRequest, ) diff --git a/service/aiproxy/relay/adaptor/gemini/adaptor.go b/service/aiproxy/relay/adaptor/gemini/adaptor.go index 18ca3b398..88027a9ce 100644 --- a/service/aiproxy/relay/adaptor/gemini/adaptor.go +++ b/service/aiproxy/relay/adaptor/gemini/adaptor.go @@ -1,7 +1,6 @@ package gemini import ( - "errors" "fmt" "io" "net/http" @@ -64,7 +63,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, req *http.Request) (string, ht case relaymode.ChatCompletions: return ConvertRequest(meta, req) default: - return "", nil, nil, errors.New("unsupported mode") + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -83,7 +82,7 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons usage, err = Handler(meta, c, resp) } default: - return nil, openai.ErrorWrapperWithMessage("unsupported mode", "unsupported_mode", http.StatusBadRequest) + return nil, openai.ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %s", meta.Mode), "unsupported_mode", http.StatusBadRequest) } return } diff --git a/service/aiproxy/relay/adaptor/gemini/embeddings.go b/service/aiproxy/relay/adaptor/gemini/embeddings.go index 7859d67be..65cc0f6c2 100644 --- a/service/aiproxy/relay/adaptor/gemini/embeddings.go +++ b/service/aiproxy/relay/adaptor/gemini/embeddings.go @@ -47,6 +47,10 @@ func ConvertEmbeddingRequest(meta *meta.Meta, req *http.Request) (string, http.H } func EmbeddingHandler(c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() var geminiEmbeddingResponse EmbeddingResponse diff --git a/service/aiproxy/relay/adaptor/gemini/main.go b/service/aiproxy/relay/adaptor/gemini/main.go index a49bff13d..830d1d277 100644 --- a/service/aiproxy/relay/adaptor/gemini/main.go +++ b/service/aiproxy/relay/adaptor/gemini/main.go @@ -451,6 +451,10 @@ func streamResponseGeminiChat2OpenAI(meta *meta.Meta, geminiResponse *ChatRespon func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + log := middleware.GetLogger(c) responseText := strings.Builder{} @@ -500,6 +504,10 @@ func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model } func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + defer resp.Body.Close() var geminiResponse ChatResponse diff --git a/service/aiproxy/relay/adaptor/interface.go b/service/aiproxy/relay/adaptor/interface.go index 4d8457e04..32904f403 100644 --- a/service/aiproxy/relay/adaptor/interface.go +++ b/service/aiproxy/relay/adaptor/interface.go @@ -1,6 +1,7 @@ package adaptor import ( + "errors" "io" "net/http" @@ -21,6 +22,8 @@ type Adaptor interface { GetModelList() []*model.ModelConfig } +var ErrGetBalanceNotImplemented = errors.New("get balance not implemented") + type Balancer interface { GetBalance(channel *model.Channel) (float64, error) } diff --git a/service/aiproxy/relay/adaptor/lingyiwanwu/adaptor.go b/service/aiproxy/relay/adaptor/lingyiwanwu/adaptor.go index 3f529b461..cf599fbeb 100644 --- a/service/aiproxy/relay/adaptor/lingyiwanwu/adaptor.go +++ b/service/aiproxy/relay/adaptor/lingyiwanwu/adaptor.go @@ -2,6 +2,7 @@ package lingyiwanwu import ( "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" ) @@ -22,3 +23,7 @@ func (a *Adaptor) GetModelList() []*model.ModelConfig { func (a *Adaptor) GetChannelName() string { return "lingyiwanwu" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/adaptor/minimax/adaptor.go b/service/aiproxy/relay/adaptor/minimax/adaptor.go index 54b8a4080..e4d109777 100644 --- a/service/aiproxy/relay/adaptor/minimax/adaptor.go +++ b/service/aiproxy/relay/adaptor/minimax/adaptor.go @@ -7,6 +7,7 @@ import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" @@ -77,3 +78,7 @@ func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Respons func (a *Adaptor) GetChannelName() string { return "minimax" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/adaptor/minimax/tts.go b/service/aiproxy/relay/adaptor/minimax/tts.go index da53a5e1c..ef0e655f2 100644 --- a/service/aiproxy/relay/adaptor/minimax/tts.go +++ b/service/aiproxy/relay/adaptor/minimax/tts.go @@ -106,6 +106,10 @@ type TTSResponse struct { } func TTSHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, openai.ErrorHanlder(resp) + } + if !strings.Contains(resp.Header.Get("Content-Type"), "application/json") && meta.GetBool("stream") { return ttsStreamHandler(meta, c, resp) } diff --git a/service/aiproxy/relay/adaptor/moonshot/balance.go b/service/aiproxy/relay/adaptor/moonshot/balance.go index 50a854d95..7dbc0e078 100644 --- a/service/aiproxy/relay/adaptor/moonshot/balance.go +++ b/service/aiproxy/relay/adaptor/moonshot/balance.go @@ -17,7 +17,7 @@ func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { if u == "" { u = baseURL } - url := u + "/v1/users/me/balance" + url := u + "/users/me/balance" req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil) if err != nil { return 0, err diff --git a/service/aiproxy/relay/adaptor/ollama/adaptor.go b/service/aiproxy/relay/adaptor/ollama/adaptor.go index 554f34de4..ddc3a90d9 100644 --- a/service/aiproxy/relay/adaptor/ollama/adaptor.go +++ b/service/aiproxy/relay/adaptor/ollama/adaptor.go @@ -8,6 +8,7 @@ import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" "github.com/labring/sealos/service/aiproxy/relay/relaymode" @@ -31,7 +32,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.ChatCompletions: return u + "/api/chat", nil default: - return "", fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -50,7 +51,7 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, request *http.Request) (string case relaymode.ChatCompletions: return ConvertRequest(meta, request) default: - return "", nil, nil, fmt.Errorf("unsupported mode: %d", meta.Mode) + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -61,13 +62,15 @@ func (a *Adaptor) DoRequest(_ *meta.Meta, _ *gin.Context, req *http.Request) (*h func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Response) (usage *relaymodel.Usage, err *relaymodel.ErrorWithStatusCode) { switch meta.Mode { case relaymode.Embeddings: - err, usage = EmbeddingHandler(meta, c, resp) - default: + usage, err = EmbeddingHandler(meta, c, resp) + case relaymode.ChatCompletions: if utils.IsStreamResponse(resp) { - err, usage = StreamHandler(meta, c, resp) + usage, err = StreamHandler(meta, c, resp) } else { - err, usage = Handler(meta, c, resp) + usage, err = Handler(meta, c, resp) } + default: + return nil, openai.ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %s", meta.Mode), "unsupported_mode", http.StatusBadRequest) } return } diff --git a/service/aiproxy/relay/adaptor/ollama/error.go b/service/aiproxy/relay/adaptor/ollama/error.go new file mode 100644 index 000000000..8e7795f2e --- /dev/null +++ b/service/aiproxy/relay/adaptor/ollama/error.go @@ -0,0 +1,24 @@ +package ollama + +import ( + "net/http" + + "github.com/bytedance/sonic" + "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" + relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" +) + +type errorResponse struct { + Error string `json:"error"` +} + +func ErrorHandler(resp *http.Response) *relaymodel.ErrorWithStatusCode { + defer resp.Body.Close() + + var er errorResponse + err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&er) + if err != nil { + return openai.ErrorWrapperWithMessage("decode response error: "+err.Error(), nil, http.StatusInternalServerError) + } + return openai.ErrorWrapperWithMessage(er.Error, nil, http.StatusInternalServerError) +} diff --git a/service/aiproxy/relay/adaptor/ollama/main.go b/service/aiproxy/relay/adaptor/ollama/main.go index de38fb3cd..c7af1011f 100644 --- a/service/aiproxy/relay/adaptor/ollama/main.go +++ b/service/aiproxy/relay/adaptor/ollama/main.go @@ -5,16 +5,15 @@ import ( "bytes" "io" "net/http" - "strings" "time" "github.com/bytedance/sonic" "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common" - "github.com/labring/sealos/service/aiproxy/common/conv" "github.com/labring/sealos/service/aiproxy/common/image" "github.com/labring/sealos/service/aiproxy/common/random" "github.com/labring/sealos/service/aiproxy/common/render" + "github.com/labring/sealos/service/aiproxy/common/splitter" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/constant" @@ -29,10 +28,9 @@ func ConvertRequest(meta *meta.Meta, req *http.Request) (string, http.Header, io if err != nil { return "", nil, nil, err } - request.Model = meta.ActualModel ollamaRequest := ChatRequest{ - Model: request.Model, + Model: meta.ActualModel, Options: &Options{ Seed: int(request.Seed), Temperature: request.Temperature, @@ -127,45 +125,45 @@ func streamResponseOllama2OpenAI(meta *meta.Meta, ollamaResponse *ChatResponse) return &response } -func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) { +func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHandler(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) - var usage relaymodel.Usage + var usage *relaymodel.Usage scanner := bufio.NewScanner(resp.Body) - scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { - if atEOF && len(data) == 0 { - return 0, nil, nil - } - if i := strings.Index(conv.BytesToString(data), "}\n"); i >= 0 { - return i + 2, data[0 : i+1], nil - } - if atEOF { - return len(data), data, nil - } - return 0, nil, nil - }) common.SetEventStreamHeaders(c) + var thinkSplitter *splitter.Splitter + if meta.ChannelConfig.SplitThink { + thinkSplitter = splitter.NewThinkSplitter() + } + for scanner.Scan() { - data := scanner.Text() - if strings.HasPrefix(data, "}") { - data = strings.TrimPrefix(data, "}") + "}" - } + data := scanner.Bytes() var ollamaResponse ChatResponse - err := sonic.Unmarshal(conv.StringToBytes(data), &ollamaResponse) + err := sonic.Unmarshal(data, &ollamaResponse) if err != nil { log.Error("error unmarshalling stream response: " + err.Error()) continue } response := streamResponseOllama2OpenAI(meta, &ollamaResponse) - if response.Usage != nil { - usage = *response.Usage + usage = response.Usage + } + + if meta.ChannelConfig.SplitThink { + openai.StreamSplitThinkModeld(response, thinkSplitter, func(data *openai.ChatCompletionsStreamResponse) { + _ = render.ObjectData(c, data) + }) + continue } _ = render.ObjectData(c, response) @@ -177,7 +175,7 @@ func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relay render.Done(c) - return nil, &usage + return usage, nil } func ConvertEmbeddingRequest(meta *meta.Meta, req *http.Request) (string, http.Header, io.Reader, error) { @@ -203,52 +201,47 @@ func ConvertEmbeddingRequest(meta *meta.Meta, req *http.Request) (string, http.H return http.MethodPost, nil, bytes.NewReader(data), nil } -func EmbeddingHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) { +func EmbeddingHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHandler(resp) + } + defer resp.Body.Close() var ollamaResponse EmbeddingResponse err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&ollamaResponse) if err != nil { - return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } if ollamaResponse.Error != "" { - return &relaymodel.ErrorWithStatusCode{ - Error: relaymodel.Error{ - Message: ollamaResponse.Error, - Type: "ollama_error", - Param: "", - Code: "ollama_error", - }, - StatusCode: resp.StatusCode, - }, nil + return nil, openai.ErrorWrapperWithMessage(ollamaResponse.Error, openai.ErrorTypeUpstream, resp.StatusCode) } fullTextResponse := embeddingResponseOllama2OpenAI(meta, &ollamaResponse) jsonResponse, err := sonic.Marshal(fullTextResponse) if err != nil { - return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError) } c.Writer.Header().Set("Content-Type", "application/json") c.Writer.WriteHeader(resp.StatusCode) _, _ = c.Writer.Write(jsonResponse) - return nil, &fullTextResponse.Usage + return &fullTextResponse.Usage, nil } func embeddingResponseOllama2OpenAI(meta *meta.Meta, response *EmbeddingResponse) *openai.EmbeddingResponse { openAIEmbeddingResponse := openai.EmbeddingResponse{ Object: "list", - Data: make([]*openai.EmbeddingResponseItem, 0, 1), + Data: make([]*openai.EmbeddingResponseItem, 0, len(response.Embeddings)), Model: meta.OriginModel, Usage: relaymodel.Usage{ PromptTokens: response.PromptEvalCount, TotalTokens: response.PromptEvalCount, }, } - for i, embedding := range response.Embeddings { openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, &openai.EmbeddingResponseItem{ - Object: `embedding`, + Object: "embedding", Index: i, Embedding: embedding, }) @@ -256,32 +249,33 @@ func embeddingResponseOllama2OpenAI(meta *meta.Meta, response *EmbeddingResponse return &openAIEmbeddingResponse } -func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) { +func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHandler(resp) + } + defer resp.Body.Close() var ollamaResponse ChatResponse err := sonic.ConfigDefault.NewDecoder(resp.Body).Decode(&ollamaResponse) if err != nil { - return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) } if ollamaResponse.Error != "" { - return &relaymodel.ErrorWithStatusCode{ - Error: relaymodel.Error{ - Message: ollamaResponse.Error, - Type: "ollama_error", - Param: "", - Code: "ollama_error", - }, - StatusCode: resp.StatusCode, - }, nil + return nil, openai.ErrorWrapperWithMessage(ollamaResponse.Error, openai.ErrorTypeUpstream, resp.StatusCode) } fullTextResponse := responseOllama2OpenAI(meta, &ollamaResponse) + + if meta.ChannelConfig.SplitThink { + openai.SplitThinkModeld(fullTextResponse) + } + jsonResponse, err := sonic.Marshal(fullTextResponse) if err != nil { - return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + return nil, openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError) } c.Writer.Header().Set("Content-Type", "application/json") c.Writer.WriteHeader(resp.StatusCode) _, _ = c.Writer.Write(jsonResponse) - return nil, &fullTextResponse.Usage + return &fullTextResponse.Usage, nil } diff --git a/service/aiproxy/relay/adaptor/openai/adaptor.go b/service/aiproxy/relay/adaptor/openai/adaptor.go index ed1a64639..c2baae3ba 100644 --- a/service/aiproxy/relay/adaptor/openai/adaptor.go +++ b/service/aiproxy/relay/adaptor/openai/adaptor.go @@ -3,6 +3,7 @@ package openai import ( "bytes" "errors" + "fmt" "io" "net/http" @@ -53,7 +54,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { case relaymode.Rerank: path = "/rerank" default: - return "", errors.New("unsupported mode") + return "", fmt.Errorf("unsupported mode: %s", meta.Mode) } return u + path, nil @@ -89,7 +90,7 @@ func ConvertRequest(meta *meta.Meta, req *http.Request) (string, http.Header, io case relaymode.Rerank: return ConvertRerankRequest(meta, req) default: - return "", nil, nil, errors.New("unsupported convert request mode") + return "", nil, nil, fmt.Errorf("unsupported mode: %s", meta.Mode) } } @@ -114,7 +115,7 @@ func DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Response) (usage *re usage, err = Handler(meta, c, resp) } default: - return nil, ErrorWrapperWithMessage("unsupported response mode", "unsupported_mode", http.StatusBadRequest) + return nil, ErrorWrapperWithMessage(fmt.Sprintf("unsupported mode: %s", meta.Mode), "unsupported_mode", http.StatusBadRequest) } return } diff --git a/service/aiproxy/relay/utils/error.go b/service/aiproxy/relay/adaptor/openai/error.go similarity index 79% rename from service/aiproxy/relay/utils/error.go rename to service/aiproxy/relay/adaptor/openai/error.go index f31a7e2f3..f57709d75 100644 --- a/service/aiproxy/relay/utils/error.go +++ b/service/aiproxy/relay/adaptor/openai/error.go @@ -1,4 +1,4 @@ -package utils +package openai import ( "fmt" @@ -8,7 +8,7 @@ import ( "github.com/bytedance/sonic" "github.com/labring/sealos/service/aiproxy/common/conv" - "github.com/labring/sealos/service/aiproxy/relay/meta" + "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/relay/model" ) @@ -54,26 +54,14 @@ func (e GeneralErrorResponse) ToMessage() string { } const ( + ErrorTypeAIProxy = middleware.ErrorTypeAIPROXY ErrorTypeUpstream = "upstream_error" ErrorCodeBadResponse = "bad_response" ) -func RelayErrorHandler(_ *meta.Meta, resp *http.Response) *model.ErrorWithStatusCode { - if resp == nil { - return &model.ErrorWithStatusCode{ - StatusCode: 500, - Error: model.Error{ - Message: "resp is nil", - Type: ErrorTypeUpstream, - Code: ErrorCodeBadResponse, - }, - } - } - return RelayDefaultErrorHanlder(resp) -} - -func RelayDefaultErrorHanlder(resp *http.Response) *model.ErrorWithStatusCode { +func ErrorHanlder(resp *http.Response) *model.ErrorWithStatusCode { defer resp.Body.Close() + respBody, err := io.ReadAll(resp.Body) if err != nil { return &model.ErrorWithStatusCode{ @@ -89,10 +77,9 @@ func RelayDefaultErrorHanlder(resp *http.Response) *model.ErrorWithStatusCode { ErrorWithStatusCode := &model.ErrorWithStatusCode{ StatusCode: resp.StatusCode, Error: model.Error{ - Message: "", - Type: ErrorTypeUpstream, - Code: ErrorCodeBadResponse, - Param: strconv.Itoa(resp.StatusCode), + Type: ErrorTypeUpstream, + Code: ErrorCodeBadResponse, + Param: strconv.Itoa(resp.StatusCode), }, } diff --git a/service/aiproxy/relay/adaptor/openai/image.go b/service/aiproxy/relay/adaptor/openai/image.go index 6ba311740..3292de249 100644 --- a/service/aiproxy/relay/adaptor/openai/image.go +++ b/service/aiproxy/relay/adaptor/openai/image.go @@ -41,6 +41,10 @@ func ConvertImageRequest(meta *meta.Meta, req *http.Request) (string, http.Heade } func ImageHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/openai/main.go b/service/aiproxy/relay/adaptor/openai/main.go index 41fccd4b1..bac32fae1 100644 --- a/service/aiproxy/relay/adaptor/openai/main.go +++ b/service/aiproxy/relay/adaptor/openai/main.go @@ -93,6 +93,10 @@ func GetUsageAndChoicesResponseFromNode(node *ast.Node) (*UsageAndChoicesRespons } func StreamHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) @@ -241,6 +245,8 @@ func StreamSplitThink(data map[string]any, thinkSplitter *splitter.Splitter, ren } think, remaining := thinkSplitter.Process(conv.StringToBytes(content)) if len(think) == 0 && len(remaining) == 0 { + delta["content"] = "" + delete(delta, "reasoning_content") renderCallback(data) return } @@ -256,6 +262,38 @@ func StreamSplitThink(data map[string]any, thinkSplitter *splitter.Splitter, ren } } +func StreamSplitThinkModeld(data *ChatCompletionsStreamResponse, thinkSplitter *splitter.Splitter, renderCallback func(data *ChatCompletionsStreamResponse)) { + choices := data.Choices + // only support one choice + if len(data.Choices) != 1 { + renderCallback(data) + return + } + choice := choices[0] + content, ok := choice.Delta.Content.(string) + if !ok { + renderCallback(data) + return + } + think, remaining := thinkSplitter.Process(conv.StringToBytes(content)) + if len(think) == 0 && len(remaining) == 0 { + choice.Delta.Content = "" + choice.Delta.ReasoningContent = "" + renderCallback(data) + return + } + if len(think) > 0 { + choice.Delta.Content = "" + choice.Delta.ReasoningContent = conv.BytesToString(think) + renderCallback(data) + } + if len(remaining) > 0 { + choice.Delta.Content = conv.BytesToString(remaining) + choice.Delta.ReasoningContent = "" + renderCallback(data) + } +} + func SplitThink(data map[string]any) { choices, ok := data["choices"].([]any) if !ok { @@ -266,17 +304,30 @@ func SplitThink(data map[string]any) { if !ok { continue } - delta, ok := choiceMap["delta"].(map[string]any) + message, ok := choiceMap["message"].(map[string]any) if !ok { continue } - content, ok := delta["content"].(string) + content, ok := message["content"].(string) if !ok { continue } think, remaining := splitter.NewThinkSplitter().Process(conv.StringToBytes(content)) - delta["reasoning_content"] = conv.BytesToString(think) - delta["content"] = conv.BytesToString(remaining) + message["reasoning_content"] = conv.BytesToString(think) + message["content"] = conv.BytesToString(remaining) + } +} + +func SplitThinkModeld(data *TextResponse) { + choices := data.Choices + for _, choice := range choices { + content, ok := choice.Message.Content.(string) + if !ok { + continue + } + think, remaining := splitter.NewThinkSplitter().Process(conv.StringToBytes(content)) + choice.Message.ReasoningContent = conv.BytesToString(think) + choice.Message.Content = conv.BytesToString(remaining) } } @@ -325,6 +376,10 @@ func GetSlimTextResponseFromNode(node *ast.Node) (*SlimTextResponse, error) { } func Handler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/openai/moderations.go b/service/aiproxy/relay/adaptor/openai/moderations.go index bef5844d7..29fe98933 100644 --- a/service/aiproxy/relay/adaptor/openai/moderations.go +++ b/service/aiproxy/relay/adaptor/openai/moderations.go @@ -12,6 +12,10 @@ import ( ) func ModerationsHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/openai/rerank.go b/service/aiproxy/relay/adaptor/openai/rerank.go index 8f400535d..1bdd644c4 100644 --- a/service/aiproxy/relay/adaptor/openai/rerank.go +++ b/service/aiproxy/relay/adaptor/openai/rerank.go @@ -33,6 +33,10 @@ func ConvertRerankRequest(meta *meta.Meta, req *http.Request) (string, http.Head } func RerankHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/openai/stt.go b/service/aiproxy/relay/adaptor/openai/stt.go index 00d090bea..3c779386b 100644 --- a/service/aiproxy/relay/adaptor/openai/stt.go +++ b/service/aiproxy/relay/adaptor/openai/stt.go @@ -77,6 +77,10 @@ func ConvertSTTRequest(meta *meta.Meta, request *http.Request) (string, http.Hea } func STTHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*model.Usage, *model.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/openai/tts.go b/service/aiproxy/relay/adaptor/openai/tts.go index b0a946574..54264d877 100644 --- a/service/aiproxy/relay/adaptor/openai/tts.go +++ b/service/aiproxy/relay/adaptor/openai/tts.go @@ -55,6 +55,10 @@ func ConvertTTSRequest(meta *meta.Meta, req *http.Request, defaultVoice string) } func TTSHandler(meta *meta.Meta, c *gin.Context, resp *http.Response) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { + if resp.StatusCode != http.StatusOK { + return nil, ErrorHanlder(resp) + } + defer resp.Body.Close() log := middleware.GetLogger(c) diff --git a/service/aiproxy/relay/adaptor/siliconflow/balance.go b/service/aiproxy/relay/adaptor/siliconflow/balance.go index 08fe2b64c..e588f5185 100644 --- a/service/aiproxy/relay/adaptor/siliconflow/balance.go +++ b/service/aiproxy/relay/adaptor/siliconflow/balance.go @@ -18,7 +18,7 @@ func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { if u == "" { u = baseURL } - url := u + "/v1/user/info" + url := u + "/user/info" req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil) if err != nil { return 0, err diff --git a/service/aiproxy/relay/adaptor/stepfun/adaptor.go b/service/aiproxy/relay/adaptor/stepfun/adaptor.go index 619f9e475..d0847ed9b 100644 --- a/service/aiproxy/relay/adaptor/stepfun/adaptor.go +++ b/service/aiproxy/relay/adaptor/stepfun/adaptor.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" "github.com/labring/sealos/service/aiproxy/relay/relaymode" @@ -36,3 +37,7 @@ func (a *Adaptor) GetModelList() []*model.ModelConfig { func (a *Adaptor) GetChannelName() string { return "stepfun" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/adaptor/tencent/adaptor.go b/service/aiproxy/relay/adaptor/tencent/adaptor.go index 4b3109ed0..b1860a8cc 100644 --- a/service/aiproxy/relay/adaptor/tencent/adaptor.go +++ b/service/aiproxy/relay/adaptor/tencent/adaptor.go @@ -2,6 +2,7 @@ package tencent import ( "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" ) @@ -24,3 +25,7 @@ func (a *Adaptor) GetModelList() []*model.ModelConfig { func (a *Adaptor) GetChannelName() string { return "tencent" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/adaptor/vertexai/claude/adapter.go b/service/aiproxy/relay/adaptor/vertexai/claude/adapter.go index 7b9ad5188..6fef4659a 100644 --- a/service/aiproxy/relay/adaptor/vertexai/claude/adapter.go +++ b/service/aiproxy/relay/adaptor/vertexai/claude/adapter.go @@ -9,12 +9,11 @@ import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/adaptor/anthropic" + "github.com/labring/sealos/service/aiproxy/relay/meta" + relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" "github.com/labring/sealos/service/aiproxy/relay/relaymode" "github.com/labring/sealos/service/aiproxy/relay/utils" "github.com/pkg/errors" - - "github.com/labring/sealos/service/aiproxy/relay/meta" - relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" ) var ModelList = []*model.ModelConfig{ @@ -85,9 +84,9 @@ func (a *Adaptor) ConvertRequest(meta *meta.Meta, request *http.Request) (string func (a *Adaptor) DoResponse(meta *meta.Meta, c *gin.Context, resp *http.Response) (usage *relaymodel.Usage, err *relaymodel.ErrorWithStatusCode) { if utils.IsStreamResponse(resp) { - err, usage = anthropic.StreamHandler(meta, c, resp) + usage, err = anthropic.StreamHandler(meta, c, resp) } else { - err, usage = anthropic.Handler(meta, c, resp) + usage, err = anthropic.Handler(meta, c, resp) } return } diff --git a/service/aiproxy/relay/adaptor/vertexai/gemini/adapter.go b/service/aiproxy/relay/adaptor/vertexai/gemini/adapter.go index e2e6ccf49..88957903a 100644 --- a/service/aiproxy/relay/adaptor/vertexai/gemini/adapter.go +++ b/service/aiproxy/relay/adaptor/vertexai/gemini/adapter.go @@ -7,11 +7,10 @@ import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/adaptor/gemini" - "github.com/labring/sealos/service/aiproxy/relay/relaymode" - "github.com/labring/sealos/service/aiproxy/relay/utils" - "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" + "github.com/labring/sealos/service/aiproxy/relay/relaymode" + "github.com/labring/sealos/service/aiproxy/relay/utils" ) var ModelList = []*model.ModelConfig{ diff --git a/service/aiproxy/relay/adaptor/vertexai/token.go b/service/aiproxy/relay/adaptor/vertexai/token.go index 7199fb3ce..63cd7478d 100644 --- a/service/aiproxy/relay/adaptor/vertexai/token.go +++ b/service/aiproxy/relay/adaptor/vertexai/token.go @@ -58,6 +58,6 @@ func getToken(ctx context.Context, channelID int, adcJSON string) (string, error } _ = resp - Cache.Set(cacheKey, resp.AccessToken, cache.DefaultExpiration) - return resp.AccessToken, nil + Cache.Set(cacheKey, resp.GetAccessToken(), cache.DefaultExpiration) + return resp.GetAccessToken(), nil } diff --git a/service/aiproxy/relay/adaptor/xunfei/adaptor.go b/service/aiproxy/relay/adaptor/xunfei/adaptor.go index d5e59771c..2ac141317 100644 --- a/service/aiproxy/relay/adaptor/xunfei/adaptor.go +++ b/service/aiproxy/relay/adaptor/xunfei/adaptor.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" ) @@ -43,3 +44,7 @@ func (a *Adaptor) GetModelList() []*model.ModelConfig { func (a *Adaptor) GetChannelName() string { return "xunfei" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/adaptor/zhipu/adaptor.go b/service/aiproxy/relay/adaptor/zhipu/adaptor.go index 7d01f8138..2fbb82ea5 100644 --- a/service/aiproxy/relay/adaptor/zhipu/adaptor.go +++ b/service/aiproxy/relay/adaptor/zhipu/adaptor.go @@ -5,6 +5,7 @@ import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/adaptor" "github.com/labring/sealos/service/aiproxy/relay/adaptor/openai" "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" @@ -38,3 +39,7 @@ func (a *Adaptor) GetModelList() []*model.ModelConfig { func (a *Adaptor) GetChannelName() string { return "zhipu" } + +func (a *Adaptor) GetBalance(channel *model.Channel) (float64, error) { + return 0, adaptor.ErrGetBalanceNotImplemented +} diff --git a/service/aiproxy/relay/controller/consume.go b/service/aiproxy/relay/controller/consume.go index f2bfb86ce..94fa56b47 100644 --- a/service/aiproxy/relay/controller/consume.go +++ b/service/aiproxy/relay/controller/consume.go @@ -3,7 +3,6 @@ package controller import ( "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common/balance" - "github.com/labring/sealos/service/aiproxy/common/ctxkey" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/model" "github.com/labring/sealos/service/aiproxy/relay/meta" @@ -47,11 +46,9 @@ func getGroupBalance(ctx *gin.Context, meta *meta.Meta) (float64, balance.PostGr return 0, nil, nil } - groupBalance, ok := ctx.Get(ctxkey.GroupBalance) - if !ok { - return balance.Default.GetGroupRemainBalance(ctx.Request.Context(), *meta.Group) + gbc, err := middleware.GetGroupBalanceConsumer(ctx, meta.Group) + if err != nil { + return 0, nil, err } - - groupBalanceConsumer := groupBalance.(*middleware.GroupBalanceConsumer) - return groupBalanceConsumer.GroupBalance, groupBalanceConsumer.Consumer, nil + return gbc.GroupBalance, gbc.Consumer, nil } diff --git a/service/aiproxy/relay/controller/dohelper.go b/service/aiproxy/relay/controller/dohelper.go index 62c4c739f..26b306c0a 100644 --- a/service/aiproxy/relay/controller/dohelper.go +++ b/service/aiproxy/relay/controller/dohelper.go @@ -20,17 +20,9 @@ import ( "github.com/labring/sealos/service/aiproxy/relay/meta" relaymodel "github.com/labring/sealos/service/aiproxy/relay/model" "github.com/labring/sealos/service/aiproxy/relay/relaymode" - "github.com/labring/sealos/service/aiproxy/relay/utils" log "github.com/sirupsen/logrus" ) -func isErrorHappened(resp *http.Response) bool { - if resp == nil { - return false - } - return resp.StatusCode != http.StatusOK -} - const ( // 0.5MB maxBufferSize = 512 * 1024 @@ -101,14 +93,14 @@ func DoHelper( } // 3. Handle error response - if isErrorHappened(resp) { - relayErr := utils.RelayErrorHandler(meta, resp) + if resp == nil { + relayErr := openai.ErrorWrapperWithMessage("response is nil", openai.ErrorCodeBadResponse, http.StatusInternalServerError) detail.ResponseBody = relayErr.JSONOrEmpty() return nil, &detail, relayErr } // 4. Handle success response - usage, relayErr := handleSuccessResponse(a, c, meta, resp, &detail) + usage, relayErr := handleResponse(a, c, meta, resp, &detail) if relayErr != nil { return nil, &detail, relayErr } @@ -154,7 +146,7 @@ func prepareAndDoRequest(a adaptor.Adaptor, c *gin.Context, meta *meta.Meta) (*h log.Debugf("request url: %s %s", method, fullRequestURL) ctx := context.Background() - if timeout := config.GetTimeoutWithModelType()[meta.Mode]; timeout > 0 { + if timeout := config.GetTimeoutWithModelType()[int(meta.Mode)]; timeout > 0 { // donot use c.Request.Context() because it will be canceled by the client // which will cause the usage of non-streaming requests to be unable to be recorded var cancel context.CancelFunc @@ -198,12 +190,18 @@ func doRequest(a adaptor.Adaptor, c *gin.Context, meta *meta.Meta, req *http.Req if errors.Is(err, context.DeadlineExceeded) { return nil, openai.ErrorWrapperWithMessage("do request failed: request timeout", "request_timeout", http.StatusGatewayTimeout) } + if errors.Is(err, io.EOF) { + return nil, openai.ErrorWrapperWithMessage("do request failed: "+err.Error(), "request_failed", http.StatusServiceUnavailable) + } + if errors.Is(err, io.ErrUnexpectedEOF) { + return nil, openai.ErrorWrapperWithMessage("do request failed: "+err.Error(), "request_failed", http.StatusInternalServerError) + } return nil, openai.ErrorWrapperWithMessage("do request failed: "+err.Error(), "request_failed", http.StatusBadRequest) } return resp, nil } -func handleSuccessResponse(a adaptor.Adaptor, c *gin.Context, meta *meta.Meta, resp *http.Response, detail *model.RequestDetail) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { +func handleResponse(a adaptor.Adaptor, c *gin.Context, meta *meta.Meta, resp *http.Response, detail *model.RequestDetail) (*relaymodel.Usage, *relaymodel.ErrorWithStatusCode) { buf := getBuffer() defer putBuffer(buf) @@ -216,10 +214,15 @@ func handleSuccessResponse(a adaptor.Adaptor, c *gin.Context, meta *meta.Meta, r c.Writer = rw c.Header("Content-Type", resp.Header.Get("Content-Type")) + usage, relayErr := a.DoResponse(meta, c, resp) - // copy body buffer - // do not use bytes conv - detail.ResponseBody = rw.body.String() + if relayErr != nil { + detail.ResponseBody = relayErr.JSONOrEmpty() + } else { + // copy body buffer + // do not use bytes conv + detail.ResponseBody = rw.body.String() + } return usage, relayErr } diff --git a/service/aiproxy/relay/controller/handle.go b/service/aiproxy/relay/controller/handle.go index e3313ac49..268c83c2e 100644 --- a/service/aiproxy/relay/controller/handle.go +++ b/service/aiproxy/relay/controller/handle.go @@ -41,6 +41,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa 0, errMsg, c.ClientIP(), + meta.RetryTimes, nil, ) return openai.ErrorWrapperWithMessage( @@ -71,6 +72,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa 0, err.Error(), c.ClientIP(), + meta.RetryTimes, detail, ) return openai.ErrorWrapper(err, "invalid_request", http.StatusBadRequest) @@ -92,12 +94,12 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa logDetail = detail log.Errorf( "handle failed: %+v\nrequest detail:\n%s\nresponse detail:\n%s", - respErr.Error, + respErr, logDetail.RequestBody, logDetail.ResponseBody, ) } else { - log.Errorf("handle failed: %+v", respErr.Error) + log.Errorf("handle failed: %+v", respErr) } consume.AsyncConsume( @@ -109,6 +111,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa preCheckReq.OutputPrice, respErr.Error.JSONOrEmpty(), c.ClientIP(), + meta.RetryTimes, detail, ) return respErr @@ -133,6 +136,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa preCheckReq.OutputPrice, "", c.ClientIP(), + meta.RetryTimes, detail, ) diff --git a/service/aiproxy/relay/meta/meta.go b/service/aiproxy/relay/meta/meta.go index f0417aa84..438206276 100644 --- a/service/aiproxy/relay/meta/meta.go +++ b/service/aiproxy/relay/meta/meta.go @@ -5,6 +5,7 @@ import ( "time" "github.com/labring/sealos/service/aiproxy/model" + "github.com/labring/sealos/service/aiproxy/relay/relaymode" ) type ChannelMeta struct { @@ -28,9 +29,10 @@ type Meta struct { RequestID string OriginModel string ActualModel string - Mode int + Mode relaymode.Mode InputTokens int IsChannelTest bool + RetryTimes int } type Option func(meta *Meta) @@ -71,9 +73,15 @@ func WithToken(token *model.TokenCache) Option { } } +func WithRetryTimes(retryTimes int) Option { + return func(meta *Meta) { + meta.RetryTimes = retryTimes + } +} + func NewMeta( channel *model.Channel, - mode int, + mode relaymode.Mode, modelName string, modelConfig *model.ModelConfig, opts ...Option, @@ -151,7 +159,6 @@ func (m *Meta) GetBool(key string) bool { return b } -//nolint:unparam func GetMappedModelName(modelName string, mapping map[string]string) (string, bool) { if len(modelName) == 0 { return modelName, false diff --git a/service/aiproxy/relay/model/misc.go b/service/aiproxy/relay/model/misc.go index d86211049..5330b103b 100644 --- a/service/aiproxy/relay/model/misc.go +++ b/service/aiproxy/relay/model/misc.go @@ -34,17 +34,17 @@ func (e *Error) JSONOrEmpty() string { } type ErrorWithStatusCode struct { - Error Error `json:"error"` + Error Error `json:"error,omitempty"` StatusCode int `json:"-"` } func (e *ErrorWithStatusCode) JSONOrEmpty() string { - if e.Error.IsEmpty() { + if e.StatusCode == 0 && e.Error.IsEmpty() { return "" } - jsonBuf, err := sonic.Marshal(e) + jsonBuf, err := sonic.MarshalString(e) if err != nil { return "" } - return conv.BytesToString(jsonBuf) + return jsonBuf } diff --git a/service/aiproxy/relay/relaymode/define.go b/service/aiproxy/relay/relaymode/define.go index ecaabc618..931ab0675 100644 --- a/service/aiproxy/relay/relaymode/define.go +++ b/service/aiproxy/relay/relaymode/define.go @@ -1,7 +1,42 @@ package relaymode +import "fmt" + +type Mode int + +func (m Mode) String() string { + switch m { + case Unknown: + return "Unknown" + case ChatCompletions: + return "ChatCompletions" + case Completions: + return "Completions" + case Embeddings: + return "Embeddings" + case Moderations: + return "Moderations" + case ImagesGenerations: + return "ImagesGenerations" + case Edits: + return "Edits" + case AudioSpeech: + return "AudioSpeech" + case AudioTranscription: + return "AudioTranscription" + case AudioTranslation: + return "AudioTranslation" + case Rerank: + return "Rerank" + case ParsePdf: + return "ParsePdf" + default: + return fmt.Sprintf("Mode(%d)", m) + } +} + const ( - Unknown = iota + Unknown Mode = iota ChatCompletions Completions Embeddings diff --git a/service/aiproxy/relay/utils/testreq.go b/service/aiproxy/relay/utils/testreq.go index 55bfb36b2..33f8e83a8 100644 --- a/service/aiproxy/relay/utils/testreq.go +++ b/service/aiproxy/relay/utils/testreq.go @@ -4,7 +4,6 @@ import ( "bytes" "fmt" "io" - "strconv" "github.com/bytedance/sonic" "github.com/labring/sealos/service/aiproxy/model" @@ -24,7 +23,7 @@ func NewErrUnsupportedModelType(modelType string) *UnsupportedModelTypeError { return &UnsupportedModelTypeError{ModelType: modelType} } -func BuildRequest(modelConfig *model.ModelConfig) (io.Reader, int, error) { +func BuildRequest(modelConfig *model.ModelConfig) (io.Reader, relaymode.Mode, error) { switch modelConfig.Type { case relaymode.ChatCompletions: body, err := BuildChatCompletionRequest(modelConfig.Model) @@ -73,7 +72,7 @@ func BuildRequest(modelConfig *model.ModelConfig) (io.Reader, int, error) { case relaymode.ParsePdf: return nil, relaymode.Unknown, NewErrUnsupportedModelType("parse pdf") default: - return nil, relaymode.Unknown, NewErrUnsupportedModelType(strconv.Itoa(modelConfig.Type)) + return nil, relaymode.Unknown, NewErrUnsupportedModelType(modelConfig.Type.String()) } } diff --git a/service/aiproxy/relay/utils/utils.go b/service/aiproxy/relay/utils/utils.go index ae307bf00..74b3d1570 100644 --- a/service/aiproxy/relay/utils/utils.go +++ b/service/aiproxy/relay/utils/utils.go @@ -64,5 +64,10 @@ func DoRequest(req *http.Request) (*http.Response, error) { } func IsStreamResponse(resp *http.Response) bool { - return strings.Contains(resp.Header.Get("Content-Type"), "event-stream") + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + return false + } + return strings.Contains(contentType, "event-stream") || + strings.Contains(contentType, "x-ndjson") } diff --git a/service/aiproxy/router/api.go b/service/aiproxy/router/api.go index 819f91a0b..b554a036e 100644 --- a/service/aiproxy/router/api.go +++ b/service/aiproxy/router/api.go @@ -2,11 +2,10 @@ package router import ( "github.com/gin-contrib/gzip" + "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/common/env" "github.com/labring/sealos/service/aiproxy/controller" "github.com/labring/sealos/service/aiproxy/middleware" - - "github.com/gin-gonic/gin" ) func SetAPIRouter(router *gin.Engine) { @@ -168,6 +167,7 @@ func SetAPIRouter(router *gin.Engine) { monitorRoute.DELETE("/:id", controller.ClearChannelAllModelErrors) monitorRoute.DELETE("/:id/:model", controller.ClearChannelModelErrors) monitorRoute.GET("/models", controller.GetModelsErrorRate) + monitorRoute.GET("/banned_channels", controller.GetAllBannedModelChannels) } } } diff --git a/service/aiproxy/router/relay.go b/service/aiproxy/router/relay.go index a1bcddb27..3f03f9233 100644 --- a/service/aiproxy/router/relay.go +++ b/service/aiproxy/router/relay.go @@ -1,11 +1,10 @@ package router import ( + "github.com/gin-gonic/gin" "github.com/labring/sealos/service/aiproxy/controller" "github.com/labring/sealos/service/aiproxy/middleware" "github.com/labring/sealos/service/aiproxy/relay/relaymode" - - "github.com/gin-gonic/gin" ) func SetRelayRouter(router *gin.Engine) {