feat: sync deployed v1.0.0-rc.38-custom.67

model/group_metadata.go

@@ -0,0 +1,74 @@
+package model
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"unicode/utf8"
+
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+type GroupMetadata struct {
+	GroupName string `gorm:"primaryKey;size:191"`
+	Value     string `gorm:"type:text;not null"`
+}
+type GroupMetadataValue struct {
+	Tags               []string `json:"tags"`
+	HideRecentDegraded bool     `json:"hide_recent_degraded"`
+}
+
+func GetGroupMetadata(ctx context.Context) (map[string]GroupMetadataValue, error) {
+	var rows []GroupMetadata
+	if err := DB.WithContext(ctx).Find(&rows).Error; err != nil {
+		return nil, err
+	}
+	result := map[string]GroupMetadataValue{}
+	for _, row := range rows {
+		var value GroupMetadataValue
+		if err := json.Unmarshal([]byte(row.Value), &value); err != nil {
+			return nil, err
+		}
+		result[row.GroupName] = value
+	}
+	return result, nil
+}
+
+// Updates only submitted groups; absent groups keep their existing metadata.
+func SaveGroupMetadata(ctx context.Context, groups map[string]GroupMetadataValue) error {
+	rows := []GroupMetadata{}
+	for group, value := range groups {
+		if !ratio_setting.ContainsGroupRatio(group) || group == "auto" {
+			return fmt.Errorf("invalid group: %s", group)
+		}
+		tags, seen := []string{}, map[string]bool{}
+		for _, tag := range value.Tags {
+			tag = strings.TrimSpace(tag)
+			if tag == "" || seen[tag] {
+				continue
+			}
+			if utf8.RuneCountInString(tag) > 40 || len(tags) >= 20 {
+				return fmt.Errorf("too many or too long group tags")
+			}
+			seen[tag] = true
+			tags = append(tags, tag)
+		}
+		value.Tags = tags
+		body, err := json.Marshal(value)
+		if err != nil {
+			return err
+		}
+		rows = append(rows, GroupMetadata{GroupName: group, Value: string(body)})
+	}
+	return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		for _, row := range rows {
+			if err := tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "group_name"}}, DoUpdates: clause.AssignmentColumns([]string{"value"})}).Create(&row).Error; err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}

model/group_model_cost.go

@@ -0,0 +1,86 @@
+package model
+
+import (
+	"context"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+)
+
+type GroupModelCost struct {
+	USDPerMillion *float64 `json:"usd_per_million"`
+	Requests      int64    `json:"requests"`
+	Tokens        float64  `json:"tokens"`
+	quota         float64
+}
+
+// Consumption logs contain the settled charge, not the pre-consumed quota.
+// Attribute by the group actually billed, never by today's channel membership.
+func GetGroupModelCosts(ctx context.Context, start, end time.Time) (map[string]map[string]GroupModelCost, error) {
+	result := map[string]map[string]GroupModelCost{}
+	rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).Select("group", "model_name", "quota", "prompt_tokens", "completion_tokens", "other").
+		Where("type = ? AND token_id > 0 AND created_at >= ? AND created_at < ?", LogTypeConsume, start.Unix(), end.Unix()).Rows()
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var group, name, raw string
+		var quota, prompt, completion int64
+		if err := rows.Scan(&group, &name, &quota, &prompt, &completion, &raw); err != nil {
+			return nil, err
+		}
+		if group == "" || name == "" || quota < 0 {
+			continue
+		}
+		tokens, valid := consumptionTokenTotal(prompt, completion, raw)
+		if !valid || tokens <= 0 {
+			continue
+		}
+		if result[group] == nil {
+			result[group] = map[string]GroupModelCost{}
+		}
+		value := result[group][name]
+		value.Requests++
+		value.Tokens += tokens
+		value.quota += float64(quota)
+		result[group][name] = value
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	for group, models := range result {
+		for name, value := range models {
+			if common.QuotaPerUnit > 0 && value.Tokens > 0 {
+				rate := value.quota / common.QuotaPerUnit / value.Tokens * 1_000_000
+				value.USDPerMillion = &rate
+			}
+			result[group][name] = value
+		}
+	}
+	return result, nil
+}
+
+func consumptionTokenTotal(prompt, completion int64, raw string) (float64, bool) {
+	var usage struct {
+		Total     int64  `json:"input_tokens_total"`
+		Cached    int64  `json:"cache_tokens"`
+		Created   int64  `json:"cache_creation_tokens"`
+		Created5m int64  `json:"cache_creation_tokens_5m"`
+		Created1h int64  `json:"cache_creation_tokens_1h"`
+		Written   int64  `json:"cache_write_tokens"`
+		Semantic  string `json:"usage_semantic"`
+		Claude    bool   `json:"claude"`
+	}
+	if raw != "" && common.UnmarshalJsonStr(raw, &usage) != nil {
+		return 0, false
+	}
+	input := float64(max(prompt, 0))
+	if usage.Total > 0 {
+		input = float64(usage.Total)
+	} else if usage.Semantic == "anthropic" || (usage.Semantic == "" && usage.Claude) {
+		created := max(float64(usage.Written), float64(usage.Created), float64(max(usage.Created5m, 0))+float64(max(usage.Created1h, 0)), 0)
+		input += float64(max(usage.Cached, 0)) + created
+	}
+	return input + float64(max(completion, 0)), true
+}

model/group_status_test.go

@@ -3,6 +3,7 @@ package model
 import (
 	"context"
 	"fmt"
+	"github.com/QuantumNous/new-api/common"
 	"os"
 	"testing"
 	"time"
@@ -16,6 +17,63 @@ import (
 	"gorm.io/gorm/schema"
 )
 
+func TestGroupModelCostsDatabaseMatrix(t *testing.T) {
+	for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
+		t.Run(dialect, func(t *testing.T) {
+			var driver gorm.Dialector
+			switch dialect {
+			case "sqlite":
+				driver = sqlite.Open(":memory:")
+			case "mysql":
+				if os.Getenv("TEST_MYSQL_DSN") == "" {
+					t.Skip("no MySQL test database")
+				}
+				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
+			case "postgres":
+				if os.Getenv("TEST_POSTGRES_DSN") == "" {
+					t.Skip("no PostgreSQL test database")
+				}
+				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
+			}
+			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("cost_%d_", time.Now().UnixNano())}})
+			require.NoError(t, err)
+			oldDB, oldUnit := LOG_DB, common.QuotaPerUnit
+			LOG_DB = db
+			common.QuotaPerUnit = 500000
+			t.Cleanup(func() {
+				LOG_DB = oldDB
+				common.QuotaPerUnit = oldUnit
+				require.NoError(t, db.Migrator().DropTable(&Log{}))
+				sqlDB, e := db.DB()
+				require.NoError(t, e)
+				require.NoError(t, sqlDB.Close())
+			})
+			require.NoError(t, db.AutoMigrate(&Log{}))
+			logs := []Log{
+				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 50, PromptTokens: 100, CompletionTokens: 100, Other: `{"cache_tokens":80}`},
+				{Type: 2, TokenId: 2, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 150, PromptTokens: 200, CompletionTokens: 100, Other: `{"usage_semantic":"anthropic","cache_tokens":100,"cache_creation_tokens":400}`},
+				{Type: 2, TokenId: 3, CreatedAt: 150, Group: "b", ModelName: "m", Quota: 1000, PromptTokens: 200, CompletionTokens: 100, Other: `{"input_tokens_total":900}`},
+				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "free", Quota: 0, PromptTokens: 100},
+				{Type: 2, TokenId: 0, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 2, TokenId: 1, CreatedAt: 99, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 2, TokenId: 1, CreatedAt: 200, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 5, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "no-tokens", Quota: 500},
+			}
+			require.NoError(t, db.Create(&logs).Error)
+			result, err := GetGroupModelCosts(context.Background(), time.Unix(100, 0), time.Unix(200, 0))
+			require.NoError(t, err)
+			require.NotNil(t, result["a"]["m"].USDPerMillion)
+			assert.InDelta(t, 0.4, *result["a"]["m"].USDPerMillion, 1e-9)
+			assert.Equal(t, int64(2), result["a"]["m"].Requests)
+			assert.Equal(t, float64(1000), result["a"]["m"].Tokens)
+			assert.InDelta(t, 2, *result["b"]["m"].USDPerMillion, 1e-9)
+			assert.Zero(t, *result["a"]["free"].USDPerMillion)
+			assert.NotContains(t, result["a"], "no-tokens")
+		})
+	}
+}
+
 func TestGroupCacheStatsDatabaseMatrix(t *testing.T) {
 	for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
 		t.Run(dialect, func(t *testing.T) {

model/main.go

@@ -352,6 +352,8 @@ func migrateDB() error {
 		&Token{},
 		&RelayAttempt{},
 		&User{},
+		&PersonalGroupRatio{},
+		&GroupMetadata{},
 		&UserSession{},
 		&AuthFlow{},
 		&ExternalIdentityClaim{},

model/personal_group_ratio.go

@@ -0,0 +1,87 @@
+package model
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"math"
+
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
+	"github.com/QuantumNous/new-api/types"
+	"github.com/gin-gonic/gin"
+	"gorm.io/gorm/clause"
+)
+
+// Stored independently from User and UserSetting to prevent profile mass assignment.
+type PersonalGroupRatio struct {
+	UserID int    `gorm:"primaryKey;autoIncrement:false"`
+	Ratios string `gorm:"type:text;not null"`
+}
+
+const personalRatioContextKey = "personal_group_ratios"
+
+func GetPersonalGroupRatios(ctx context.Context, userID int) (map[string]float64, error) {
+	var rows []PersonalGroupRatio
+	if err := DB.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
+		return nil, err
+	}
+	result := map[string]float64{}
+	if len(rows) != 0 {
+		if err := json.Unmarshal([]byte(rows[0].Ratios), &result); err != nil {
+			return nil, err
+		}
+	}
+	return result, nil
+}
+
+func SavePersonalGroupRatios(ctx context.Context, userID int, ratios map[string]float64) error {
+	for group, ratio := range ratios {
+		if !ratio_setting.ContainsGroupRatio(group) || group == "auto" || ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
+			return fmt.Errorf("invalid group or ratio: %s", group)
+		}
+	}
+	if ratios == nil {
+		ratios = map[string]float64{}
+	}
+	body, err := json.Marshal(ratios)
+	if err != nil {
+		return err
+	}
+	row := PersonalGroupRatio{UserID: userID, Ratios: string(body)}
+	if err = DB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "user_id"}}, DoUpdates: clause.AssignmentColumns([]string{"ratios"})}).Create(&row).Error; err != nil {
+		return err
+	}
+	// Pricing is loaded from the database at request admission, never from a stale
+	// process/Redis cache. The request snapshot survives retries and price updates.
+	return invalidateUserCache(userID)
+}
+
+func LoadPersonalGroupRatios(c *gin.Context, userID int) error {
+	ratios, err := GetPersonalGroupRatios(c.Request.Context(), userID)
+	if err == nil {
+		c.Set(personalRatioContextKey, ratios)
+	}
+	return err
+}
+
+func EffectiveGroupRatio(c *gin.Context, userGroup, group string) types.GroupRatioInfo {
+	result := types.GroupRatioInfo{GroupRatio: ratio_setting.GetGroupRatio(group), GroupSpecialRatio: -1}
+	if ratio, ok := ratio_setting.GetGroupGroupRatio(userGroup, group); ok {
+		result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
+	}
+	if value, ok := c.Get(personalRatioContextKey); ok {
+		if ratios, ok := value.(map[string]float64); ok {
+			if ratio, ok := ratios[group]; ok {
+				result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
+			}
+		}
+	}
+	return result
+}
+
+func HasPersonalGroupRatio(c *gin.Context, group string) bool {
+	value, _ := c.Get(personalRatioContextKey)
+	ratios, _ := value.(map[string]float64)
+	_, ok := ratios[group]
+	return ok
+}

model/personal_group_ratio_test.go

@@ -0,0 +1,96 @@
+package model
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"math"
+	"net/http/httptest"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
+	"github.com/gin-gonic/gin"
+	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/require"
+	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+	"gorm.io/gorm/schema"
+)
+
+func TestPersonalGroupConfigurationDatabaseMatrix(t *testing.T) {
+	for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
+		t.Run(dialect, func(t *testing.T) {
+			var driver gorm.Dialector
+			switch dialect {
+			case "sqlite":
+				driver = sqlite.Open(":memory:")
+			case "mysql":
+				if os.Getenv("TEST_MYSQL_DSN") == "" {
+					t.Skip("DSN unset")
+				}
+				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
+			case "postgres":
+				if os.Getenv("TEST_POSTGRES_DSN") == "" {
+					t.Skip("DSN unset")
+				}
+				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
+			}
+			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("personal_%d_", time.Now().UnixNano())}})
+			require.NoError(t, err)
+			old, redis := DB, common.RedisEnabled
+			DB, common.RedisEnabled = db, false
+			ratios, special := ratio_setting.GroupRatio2JSONString(), ratio_setting.GroupGroupRatio2JSONString()
+			rows := []any{&PersonalGroupRatio{}, &GroupMetadata{}}
+			t.Cleanup(func() {
+				for _, row := range rows {
+					require.NoError(t, db.Migrator().DropTable(row))
+				}
+				DB, common.RedisEnabled = old, redis
+				require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(ratios))
+				require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(special))
+				conn, _ := db.DB()
+				_ = conn.Close()
+			})
+			require.NoError(t, db.AutoMigrate(rows...))
+			require.NoError(t, db.AutoMigrate(rows...))
+			require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString("{\"default\":1,\"vip\":2}"))
+			require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString("{\"default\":{\"vip\":0.7}}"))
+			makeCtx := func() *gin.Context {
+				c, _ := gin.CreateTestContext(httptest.NewRecorder())
+				c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+				require.NoError(t, LoadPersonalGroupRatios(c, 42))
+				return c
+			}
+			require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
+			require.Equal(t, 1.0, EffectiveGroupRatio(makeCtx(), "default", "default").GroupRatio)
+			require.NoError(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{"vip": 0.2, "default": 0}))
+			inFlight := makeCtx()
+			require.Equal(t, 0.2, EffectiveGroupRatio(inFlight, "default", "vip").GroupRatio)
+			require.Equal(t, 0.0, EffectiveGroupRatio(inFlight, "default", "default").GroupRatio)
+			require.True(t, HasPersonalGroupRatio(inFlight, "vip"))
+			for _, bad := range []float64{-1, math.NaN(), math.Inf(1)} {
+				require.Error(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{"vip": bad}))
+			}
+			require.NoError(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{}))
+			require.Equal(t, 0.2, EffectiveGroupRatio(inFlight, "default", "vip").GroupRatio)
+			require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
+			require.False(t, HasPersonalGroupRatio(makeCtx(), "vip"))
+			require.NoError(t, SaveGroupMetadata(context.Background(), map[string]GroupMetadataValue{"vip": {Tags: []string{"官Key", "官Key", " 打野 ", "不稳定"}, HideRecentDegraded: true}}))
+			metadata, err := GetGroupMetadata(context.Background())
+			require.NoError(t, err)
+			require.Equal(t, []string{"官Key", "打野", "不稳定"}, metadata["vip"].Tags)
+			require.True(t, metadata["vip"].HideRecentDegraded)
+			require.NoError(t, db.AutoMigrate(rows...))
+			metadata, err = GetGroupMetadata(context.Background())
+			require.NoError(t, err)
+			require.Len(t, metadata["vip"].Tags, 3)
+			var u User
+			require.NoError(t, json.Unmarshal([]byte("{\"personal_group_ratios\":{\"vip\":0},\"ratios\":{\"vip\":0}}"), &u))
+			require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
+		})
+	}
+}

model/token_routing.go

@@ -16,19 +16,26 @@ type ModelRoutingRule struct {
 }
 
 type TokenRoutingPolicy struct {
-	Version    int                `json:"version"`
-	Mode       string             `json:"mode"`
-	Groups     []TokenGroupEntry  `json:"groups"`
-	ModelRules []ModelRoutingRule `json:"model_rules"`
+	SkipDegradedNonAstra bool               `json:"skip_degraded_non_astra,omitempty"`
+	AutoPreference       string             `json:"auto_preference,omitempty"`
+	AllowDegraded        bool               `json:"allow_degraded,omitempty"`
+	SkipDegradedGroups   bool               `json:"skip_degraded_groups,omitempty"`
+	Version              int                `json:"version"`
+	Mode                 string             `json:"mode"`
+	Groups               []TokenGroupEntry  `json:"groups"`
+	ModelRules           []ModelRoutingRule `json:"model_rules"`
 }
 
 func (p *TokenRoutingPolicy) Validate(maxGroups int, allowed func(string) bool) error {
-	if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user") {
+	if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user" && p.Mode != "smart") {
 		return fmt.Errorf("invalid routing policy version or mode")
 	}
 	if p.Mode == "custom" && len(p.Groups) == 0 {
 		return fmt.Errorf("a group pool must contain at least one group")
 	}
+	if p.Mode == "smart" && (len(p.ModelRules) != 0 || (p.AutoPreference != "" && p.AutoPreference != "cost" && p.AutoPreference != "speed" && p.AutoPreference != "stable")) {
+		return fmt.Errorf("invalid automatic routing settings")
+	}
 	if p.Mode != "custom" && len(p.Groups) != 0 {
 		return fmt.Errorf("inherited pools cannot contain custom groups")
 	}