Merge remote-tracking branch 'origin/upstream/main'

common/blacklist/main.go

@@ -0,0 +1,29 @@
+package blacklist
+
+import (
+	"fmt"
+	"sync"
+)
+
+var blackList sync.Map
+
+func init() {
+	blackList = sync.Map{}
+}
+
+func userId2Key(id int) string {
+	return fmt.Sprintf("userid_%d", id)
+}
+
+func BanUser(id int) {
+	blackList.Store(userId2Key(id), true)
+}
+
+func UnbanUser(id int) {
+	blackList.Delete(userId2Key(id))
+}
+
+func IsUserBanned(id int) bool {
+	_, ok := blackList.Load(userId2Key(id))
+	return ok
+}

common/config/config.go

@@ -66,6 +66,7 @@ var EmailDomainWhitelist = []string{
 }
 
 var DebugEnabled = os.Getenv("DEBUG") == "true"
+var DebugSQLEnabled = os.Getenv("DEBUG_SQL") == "true"
 var MemoryCacheEnabled = os.Getenv("MEMORY_CACHE_ENABLED") == "true"
 
 var LogConsumeEnabled = true
@@ -83,6 +84,9 @@ var WeChatServerAddress = ""
 var WeChatServerToken = ""
 var WeChatAccountQRCodeImageURL = ""
 
+var MessagePusherAddress = ""
+var MessagePusherToken = ""
+
 var TurnstileSiteKey = ""
 var TurnstileSecretKey = ""
 
@@ -140,3 +144,9 @@ var (
 )
 
 var RateLimitKeyExpirationDuration = 20 * time.Minute
+
+var EnableMetric = helper.GetOrDefaultEnvBool("ENABLE_METRIC", false)
+var MetricQueueSize = helper.GetOrDefaultEnvInt("METRIC_QUEUE_SIZE", 10)
+var MetricSuccessRateThreshold = helper.GetOrDefaultEnvFloat64("METRIC_SUCCESS_RATE_THRESHOLD", 0.8)
+var MetricSuccessChanSize = helper.GetOrDefaultEnvInt("METRIC_SUCCESS_CHAN_SIZE", 1024)
+var MetricFailChanSize = helper.GetOrDefaultEnvInt("METRIC_FAIL_CHAN_SIZE", 128)

common/constants.go

@@ -15,6 +15,7 @@ const (
 const (
 	UserStatusEnabled  = 1 // don't use 0, 0 is the default value!
 	UserStatusDisabled = 2 // also don't use 0
+	UserStatusDeleted  = 3
 )
 
 const (
@@ -67,6 +68,7 @@ const (
 	ChannelTypeBaichuan
 	ChannelTypeMinimax
 	ChannelTypeMistral
+	ChannelTypeGroq
 
 	ChannelTypeDummy
 )
@@ -101,6 +103,7 @@ var ChannelBaseURLs = []string{
 	"https://api.baichuan-ai.com",               // 26
 	"https://api.minimax.chat",                  // 27
 	"https://api.mistral.ai",                    // 28
+	"https://api.groq.com/openai",               // 29
 }
 
 const (

common/database.go

@@ -4,6 +4,7 @@ import "github.com/songquanpeng/one-api/common/helper"
 
 var UsingSQLite = false
 var UsingPostgreSQL = false
+var UsingMySQL = false
 
 var SQLitePath = "one-api.db"
 var SQLiteBusyTimeout = helper.GetOrDefaultEnvInt("SQLITE_BUSY_TIMEOUT", 3000)

common/helper/helper.go

@@ -195,6 +195,13 @@ func Max(a int, b int) int {
 	}
 }
 
+func GetOrDefaultEnvBool(env string, defaultValue bool) bool {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	return os.Getenv(env) == "true"
+}
+
 func GetOrDefaultEnvInt(env string, defaultValue int) int {
 	if env == "" || os.Getenv(env) == "" {
 		return defaultValue
@@ -207,6 +214,18 @@ func GetOrDefaultEnvInt(env string, defaultValue int) int {
 	return num
 }
 
+func GetOrDefaultEnvFloat64(env string, defaultValue float64) float64 {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	num, err := strconv.ParseFloat(os.Getenv(env), 64)
+	if err != nil {
+		logger.SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %f", env, err.Error(), defaultValue))
+		return defaultValue
+	}
+	return num
+}
+
 func GetOrDefaultEnvString(env string, defaultValue string) string {
 	if env == "" || os.Getenv(env) == "" {
 		return defaultValue

common/logger/logger.go

@@ -19,9 +19,6 @@ const (
 	loggerError = "ERR"
 )
 
-const maxLogCount = 1000000
-
-var logCount int
 var setupLogLock sync.Mutex
 var setupLogWorking bool
 
@@ -96,9 +93,7 @@ func logHelper(ctx context.Context, level string, msg string) {
 	id := ctx.Value(RequestIdKey)
 	now := time.Now()
 	_, _ = fmt.Fprintf(writer, "[%s] %v | %s | %s \n", level, now.Format("2006/01/02 - 15:04:05"), id, msg)
-	logCount++ // we don't need accurate count, so no lock here
-	if logCount > maxLogCount && !setupLogWorking {
-		logCount = 0
+	if !setupLogWorking {
 		setupLogWorking = true
 		go func() {
 			SetupLogger()

common/message/email.go

@@ -1,4 +1,4 @@
-package common
+package message
 
 import (
 	"crypto/rand"
@@ -12,6 +12,9 @@ import (
 )
 
 func SendEmail(subject string, receiver string, content string) error {
+	if receiver == "" {
+		return fmt.Errorf("receiver is empty")
+	}
 	if config.SMTPFrom == "" { // for compatibility
 		config.SMTPFrom = config.SMTPAccount
 	}

common/message/main.go

@@ -0,0 +1,22 @@
+package message
+
+import (
+	"fmt"
+	"github.com/songquanpeng/one-api/common/config"
+)
+
+const (
+	ByAll           = "all"
+	ByEmail         = "email"
+	ByMessagePusher = "message_pusher"
+)
+
+func Notify(by string, title string, description string, content string) error {
+	if by == ByEmail {
+		return SendEmail(title, config.RootUserEmail, content)
+	}
+	if by == ByMessagePusher {
+		return SendMessage(title, description, content)
+	}
+	return fmt.Errorf("unknown notify method: %s", by)
+}

common/message/message-pusher.go

@@ -0,0 +1,53 @@
+package message
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"github.com/songquanpeng/one-api/common/config"
+	"net/http"
+)
+
+type request struct {
+	Title       string `json:"title"`
+	Description string `json:"description"`
+	Content     string `json:"content"`
+	URL         string `json:"url"`
+	Channel     string `json:"channel"`
+	Token       string `json:"token"`
+}
+
+type response struct {
+	Success bool   `json:"success"`
+	Message string `json:"message"`
+}
+
+func SendMessage(title string, description string, content string) error {
+	if config.MessagePusherAddress == "" {
+		return errors.New("message pusher address is not set")
+	}
+	req := request{
+		Title:       title,
+		Description: description,
+		Content:     content,
+		Token:       config.MessagePusherToken,
+	}
+	data, err := json.Marshal(req)
+	if err != nil {
+		return err
+	}
+	resp, err := http.Post(config.MessagePusherAddress,
+		"application/json", bytes.NewBuffer(data))
+	if err != nil {
+		return err
+	}
+	var res response
+	err = json.NewDecoder(resp.Body).Decode(&res)
+	if err != nil {
+		return err
+	}
+	if !res.Success {
+		return errors.New(res.Message)
+	}
+	return nil
+}

common/model-ratio.go

@@ -125,6 +125,11 @@ var ModelRatio = map[string]float64{
 	"mistral-medium-latest": 2.7 / 1000 * USD,
 	"mistral-large-latest":  8.0 / 1000 * USD,
 	"mistral-embed":         0.1 / 1000 * USD,
+	// https://wow.groq.com/
+	"llama2-70b-4096":    0.7 / 1000 * USD,
+	"llama2-7b-2048":     0.1 / 1000 * USD,
+	"mixtral-8x7b-32768": 0.27 / 1000 * USD,
+	"gemma-7b-it":        0.1 / 1000 * USD,
 }
 
 var CompletionRatio = map[string]float64{}
@@ -143,6 +148,26 @@ func init() {
 	}
 }
 
+func AddNewMissingRatio(oldRatio string) string {
+	newRatio := make(map[string]float64)
+	err := json.Unmarshal([]byte(oldRatio), &newRatio)
+	if err != nil {
+		logger.SysError("error unmarshalling old ratio: " + err.Error())
+		return oldRatio
+	}
+	for k, v := range DefaultModelRatio {
+		if _, ok := newRatio[k]; !ok {
+			newRatio[k] = v
+		}
+	}
+	jsonBytes, err := json.Marshal(newRatio)
+	if err != nil {
+		logger.SysError("error marshalling new ratio: " + err.Error())
+		return oldRatio
+	}
+	return string(jsonBytes)
+}
+
 func ModelRatio2JSONString() string {
 	jsonBytes, err := json.Marshal(ModelRatio)
 	if err != nil {
@@ -212,7 +237,7 @@ func GetCompletionRatio(name string) float64 {
 				return 2
 			}
 		}
-		return 1.333333
+		return 4.0 / 3.0
 	}
 	if strings.HasPrefix(name, "gpt-4") {
 		if strings.HasSuffix(name, "preview") {
@@ -229,5 +254,9 @@ func GetCompletionRatio(name string) float64 {
 	if strings.HasPrefix(name, "mistral-") {
 		return 3
 	}
+	switch name {
+	case "llama2-70b-4096":
+		return 0.8 / 0.7
+	}
 	return 1
 }