266 lines
7.3 KiB
Go
266 lines
7.3 KiB
Go
package model
|
||
|
||
import (
|
||
"database/sql/driver"
|
||
"encoding/json"
|
||
"errors"
|
||
"time"
|
||
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
// JSONB 用于存储JSONB数据
|
||
type JSONB map[string]interface{}
|
||
|
||
func (j JSONB) Value() (driver.Value, error) {
|
||
if j == nil {
|
||
return nil, nil
|
||
}
|
||
return json.Marshal(j)
|
||
}
|
||
|
||
func (j *JSONB) Scan(value interface{}) error {
|
||
if value == nil {
|
||
*j = nil
|
||
return nil
|
||
}
|
||
bytes, ok := value.([]byte)
|
||
if !ok {
|
||
return errors.New("type assertion to []byte failed")
|
||
}
|
||
return json.Unmarshal(bytes, j)
|
||
}
|
||
|
||
type Account struct {
|
||
ID int64 `gorm:"primaryKey" json:"id"`
|
||
Name string `gorm:"size:100;not null" json:"name"`
|
||
Platform string `gorm:"size:50;not null" json:"platform"` // anthropic/openai/gemini
|
||
Type string `gorm:"size:20;not null" json:"type"` // oauth/apikey
|
||
Credentials JSONB `gorm:"type:jsonb;default:'{}'" json:"credentials"` // 凭证(加密存储)
|
||
Extra JSONB `gorm:"type:jsonb;default:'{}'" json:"extra"` // 扩展信息
|
||
ProxyID *int64 `gorm:"index" json:"proxy_id"`
|
||
Concurrency int `gorm:"default:3;not null" json:"concurrency"`
|
||
Priority int `gorm:"default:50;not null" json:"priority"` // 1-100，越小越高
|
||
Status string `gorm:"size:20;default:active;not null" json:"status"` // active/disabled/error
|
||
ErrorMessage string `gorm:"type:text" json:"error_message"`
|
||
LastUsedAt *time.Time `gorm:"index" json:"last_used_at"`
|
||
CreatedAt time.Time `gorm:"not null" json:"created_at"`
|
||
UpdatedAt time.Time `gorm:"not null" json:"updated_at"`
|
||
DeletedAt gorm.DeletedAt `gorm:"index" json:"-"`
|
||
|
||
// 调度控制
|
||
Schedulable bool `gorm:"default:true;not null" json:"schedulable"`
|
||
|
||
// 限流状态 (429)
|
||
RateLimitedAt *time.Time `gorm:"index" json:"rate_limited_at"`
|
||
RateLimitResetAt *time.Time `gorm:"index" json:"rate_limit_reset_at"`
|
||
|
||
// 过载状态 (529)
|
||
OverloadUntil *time.Time `gorm:"index" json:"overload_until"`
|
||
|
||
// 5小时时间窗口
|
||
SessionWindowStart *time.Time `json:"session_window_start"`
|
||
SessionWindowEnd *time.Time `json:"session_window_end"`
|
||
SessionWindowStatus string `gorm:"size:20" json:"session_window_status"` // allowed/allowed_warning/rejected
|
||
|
||
// 关联
|
||
Proxy *Proxy `gorm:"foreignKey:ProxyID" json:"proxy,omitempty"`
|
||
AccountGroups []AccountGroup `gorm:"foreignKey:AccountID" json:"account_groups,omitempty"`
|
||
|
||
// 虚拟字段 (不存储到数据库)
|
||
GroupIDs []int64 `gorm:"-" json:"group_ids,omitempty"`
|
||
}
|
||
|
||
func (Account) TableName() string {
|
||
return "accounts"
|
||
}
|
||
|
||
// IsActive 检查是否激活
|
||
func (a *Account) IsActive() bool {
|
||
return a.Status == "active"
|
||
}
|
||
|
||
// IsSchedulable 检查账号是否可调度
|
||
func (a *Account) IsSchedulable() bool {
|
||
if !a.IsActive() || !a.Schedulable {
|
||
return false
|
||
}
|
||
now := time.Now()
|
||
if a.OverloadUntil != nil && now.Before(*a.OverloadUntil) {
|
||
return false
|
||
}
|
||
if a.RateLimitResetAt != nil && now.Before(*a.RateLimitResetAt) {
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
|
||
// IsRateLimited 检查是否处于限流状态
|
||
func (a *Account) IsRateLimited() bool {
|
||
if a.RateLimitResetAt == nil {
|
||
return false
|
||
}
|
||
return time.Now().Before(*a.RateLimitResetAt)
|
||
}
|
||
|
||
// IsOverloaded 检查是否处于过载状态
|
||
func (a *Account) IsOverloaded() bool {
|
||
if a.OverloadUntil == nil {
|
||
return false
|
||
}
|
||
return time.Now().Before(*a.OverloadUntil)
|
||
}
|
||
|
||
// IsOAuth 检查是否为OAuth类型账号（包括oauth和setup-token）
|
||
func (a *Account) IsOAuth() bool {
|
||
return a.Type == AccountTypeOAuth || a.Type == AccountTypeSetupToken
|
||
}
|
||
|
||
// CanGetUsage 检查账号是否可以获取usage信息（只有oauth类型可以，setup-token没有profile权限）
|
||
func (a *Account) CanGetUsage() bool {
|
||
return a.Type == AccountTypeOAuth
|
||
}
|
||
|
||
// GetCredential 获取凭证字段
|
||
func (a *Account) GetCredential(key string) string {
|
||
if a.Credentials == nil {
|
||
return ""
|
||
}
|
||
if v, ok := a.Credentials[key]; ok {
|
||
if s, ok := v.(string); ok {
|
||
return s
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// GetModelMapping 获取模型映射配置
|
||
// 返回格式: map[请求模型名]实际模型名
|
||
func (a *Account) GetModelMapping() map[string]string {
|
||
if a.Credentials == nil {
|
||
return nil
|
||
}
|
||
raw, ok := a.Credentials["model_mapping"]
|
||
if !ok || raw == nil {
|
||
return nil
|
||
}
|
||
// 处理map[string]interface{}类型
|
||
if m, ok := raw.(map[string]interface{}); ok {
|
||
result := make(map[string]string)
|
||
for k, v := range m {
|
||
if s, ok := v.(string); ok {
|
||
result[k] = s
|
||
}
|
||
}
|
||
if len(result) > 0 {
|
||
return result
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// IsModelSupported 检查请求的模型是否被该账号支持
|
||
// 如果没有设置模型映射，则支持所有模型
|
||
func (a *Account) IsModelSupported(requestedModel string) bool {
|
||
mapping := a.GetModelMapping()
|
||
if mapping == nil || len(mapping) == 0 {
|
||
return true // 没有映射配置，支持所有模型
|
||
}
|
||
_, exists := mapping[requestedModel]
|
||
return exists
|
||
}
|
||
|
||
// GetMappedModel 获取映射后的实际模型名
|
||
// 如果没有映射，返回原始模型名
|
||
func (a *Account) GetMappedModel(requestedModel string) string {
|
||
mapping := a.GetModelMapping()
|
||
if mapping == nil || len(mapping) == 0 {
|
||
return requestedModel
|
||
}
|
||
if mappedModel, exists := mapping[requestedModel]; exists {
|
||
return mappedModel
|
||
}
|
||
return requestedModel
|
||
}
|
||
|
||
// GetBaseURL 获取API基础URL（用于apikey类型账号）
|
||
func (a *Account) GetBaseURL() string {
|
||
if a.Type != AccountTypeApiKey {
|
||
return ""
|
||
}
|
||
baseURL := a.GetCredential("base_url")
|
||
if baseURL == "" {
|
||
return "https://api.anthropic.com" // 默认URL
|
||
}
|
||
return baseURL
|
||
}
|
||
|
||
// GetExtraString 从Extra字段获取字符串值
|
||
func (a *Account) GetExtraString(key string) string {
|
||
if a.Extra == nil {
|
||
return ""
|
||
}
|
||
if v, ok := a.Extra[key]; ok {
|
||
if s, ok := v.(string); ok {
|
||
return s
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// IsCustomErrorCodesEnabled 检查是否启用自定义错误码功能（仅适用于 apikey 类型）
|
||
func (a *Account) IsCustomErrorCodesEnabled() bool {
|
||
if a.Type != AccountTypeApiKey || a.Credentials == nil {
|
||
return false
|
||
}
|
||
if v, ok := a.Credentials["custom_error_codes_enabled"]; ok {
|
||
if enabled, ok := v.(bool); ok {
|
||
return enabled
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// GetCustomErrorCodes 获取自定义错误码列表
|
||
func (a *Account) GetCustomErrorCodes() []int {
|
||
if a.Credentials == nil {
|
||
return nil
|
||
}
|
||
raw, ok := a.Credentials["custom_error_codes"]
|
||
if !ok || raw == nil {
|
||
return nil
|
||
}
|
||
// 处理 []interface{} 类型（JSON反序列化后的格式）
|
||
if arr, ok := raw.([]interface{}); ok {
|
||
result := make([]int, 0, len(arr))
|
||
for _, v := range arr {
|
||
// JSON 数字默认解析为 float64
|
||
if f, ok := v.(float64); ok {
|
||
result = append(result, int(f))
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ShouldHandleErrorCode 检查指定错误码是否应该被处理（停止调度/标记限流等）
|
||
// 如果未启用自定义错误码或列表为空，返回 true（使用默认策略）
|
||
// 如果启用且列表非空，只有在列表中的错误码才返回 true
|
||
func (a *Account) ShouldHandleErrorCode(statusCode int) bool {
|
||
if !a.IsCustomErrorCodesEnabled() {
|
||
return true // 未启用，使用默认策略
|
||
}
|
||
codes := a.GetCustomErrorCodes()
|
||
if len(codes) == 0 {
|
||
return true // 启用但列表为空，fallback到默认策略
|
||
}
|
||
// 检查是否在自定义列表中
|
||
for _, code := range codes {
|
||
if code == statusCode {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|