151 lines
6.7 KiB
Go
151 lines
6.7 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"github.com/QuantumNous/new-api/common"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"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) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
dsn := os.Getenv("TEST_MYSQL_DSN")
|
|
if dsn == "" {
|
|
t.Skip("TEST_MYSQL_DSN is not configured")
|
|
}
|
|
driver = mysql.Open(dsn)
|
|
case "postgres":
|
|
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
|
if dsn == "" {
|
|
t.Skip("TEST_POSTGRES_DSN is not configured")
|
|
}
|
|
driver = postgres.Open(dsn)
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("gs_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
originalDB, originalLogDB := DB, LOG_DB
|
|
DB, LOG_DB = db, db
|
|
t.Cleanup(func() { DB, LOG_DB = originalDB, originalLogDB })
|
|
channelTable := db.NamingStrategy.TableName("channels")
|
|
logTable := db.NamingStrategy.TableName("logs")
|
|
type sample struct {
|
|
Id int
|
|
ChannelId int
|
|
PromptTokens int64
|
|
Other string
|
|
Type int
|
|
TokenId int
|
|
CreatedAt int64
|
|
}
|
|
require.NoError(t, db.Table(channelTable).AutoMigrate(&StatusChannel{}))
|
|
require.NoError(t, db.Table(logTable).AutoMigrate(&sample{}))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, db.Migrator().DropTable(channelTable, logTable))
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
require.NoError(t, sqlDB.Close())
|
|
})
|
|
require.NoError(t, db.Table(channelTable).Create(&[]StatusChannel{
|
|
{Id: 1, Group: "renamed", Status: 1}, {Id: 2, Group: "automatic-down", Status: 3},
|
|
{Id: 3, Group: "manual-off", Status: 2}, {Id: 4, Group: "unused", Status: 1},
|
|
}).Error)
|
|
require.NoError(t, db.Table(logTable).Create(&[]sample{
|
|
{ChannelId: 1, PromptTokens: 100, Other: `{"cache_tokens":80}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 300, Other: `{"cache_tokens":0}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 20, Other: `{"cache_tokens":50,"usage_semantic":"anthropic","cache_creation_tokens":30}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 10, Other: `{"cache_tokens":40,"input_tokens_total":100,"claude":true}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 2, PromptTokens: 50, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 0, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 5, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 99},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 200},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `invalid`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 0, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 3, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
}).Error)
|
|
channels, err := GetStatusChannels(context.Background())
|
|
require.NoError(t, err)
|
|
require.Len(t, channels, 3)
|
|
stats, err := GetGroupCacheStats(context.Background(), channels, time.Unix(100, 0), time.Unix(200, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, GroupCacheStats{Requests: 4, CachedRequests: 3, InputTokens: 600, CachedTokens: 170}, stats["renamed"])
|
|
assert.Equal(t, GroupCacheStats{Requests: 1, InputTokens: 50}, stats["automatic-down"])
|
|
assert.Equal(t, GroupCacheStats{}, stats["unused"])
|
|
assert.NotContains(t, stats, "manual-off")
|
|
})
|
|
}
|
|
}
|