package services import ( "encoding/json" "fmt" "strings" "time" "lokistack/models" "lokistack/pkg/crypto" "lokistack/repository" ) const aiConfigSettingKey = "ai_config" // 成本相关的默认值(宁可保守:额度小、扫描少,用户觉得不够再调) const ( defaultAIMaxTokens = 1500 defaultAITimeoutSec = 120 defaultAITemperature = 0.1 // 分析场景要稳定复现,不要发挥 defaultAIDailyTokenCap = 500000 defaultAIMaxScanBytes = 2 << 30 // 2GiB ) // GetAIConfig 读配置(返回真实 APIKey,仅供内部调用)。 func GetAIConfig() models.AIConfig { cfg := models.AIConfig{ Provider: "openai-compatible", Temperature: defaultAITemperature, MaxTokens: defaultAIMaxTokens, TimeoutSec: defaultAITimeoutSec, DailyTokenCap: defaultAIDailyTokenCap, MaxScanBytes: defaultAIMaxScanBytes, } db := repository.GetDB() if db == nil { return cfg } db.RLock() raw := "" for _, s := range db.Settings { if s.Key == aiConfigSettingKey { raw = s.Value break } } db.RUnlock() if raw == "" { return cfg } var stored models.AIConfig if err := json.Unmarshal([]byte(raw), &stored); err != nil { return cfg } if stored.APIKey != "" { stored.APIKey = crypto.DecryptField(stored.APIKey) } if stored.Provider == "" { stored.Provider = "openai-compatible" } if stored.MaxTokens <= 0 { stored.MaxTokens = defaultAIMaxTokens } if stored.TimeoutSec <= 0 { stored.TimeoutSec = defaultAITimeoutSec } if stored.MaxScanBytes <= 0 { stored.MaxScanBytes = defaultAIMaxScanBytes } return stored } // MaskedAIConfig 返回给前端的版本:APIKey 只留头尾。 func MaskedAIConfig() models.AIConfig { cfg := GetAIConfig() cfg.APIKey = maskKey(cfg.APIKey) return cfg } func maskKey(key string) string { if key == "" { return "" } if len(key) <= 8 { return "***" } return key[:3] + "***" + key[len(key)-4:] } // SaveAIConfig 保存配置。 // // apiKey 传空 = 沿用旧值(前端拿到的是掩码,不能把掩码存回去); // 传 "***" 开头的掩码也视为沿用,避免用户只是改模型名却把 key 清空。 func SaveAIConfig(in models.AIConfig, apiKey string) error { old := GetAIConfig() key := strings.TrimSpace(apiKey) switch { case key == "": key = old.APIKey case strings.Contains(key, "***"): key = old.APIKey } if in.Enabled && (strings.TrimSpace(in.BaseURL) == "" || strings.TrimSpace(in.Model) == "") { return fmt.Errorf("启用 AI 分析需要填写 Base URL 与模型名") } if in.Enabled && key == "" { return fmt.Errorf("启用 AI 分析需要填写 API Key(内网无鉴权的服务可填任意占位值)") } if in.BaseURL != "" && !strings.HasPrefix(in.BaseURL, "http") { return fmt.Errorf("Base URL 必须以 http:// 或 https:// 开头(例如 https://api.deepseek.com/v1)") } enc, err := crypto.EncryptField(key) if err != nil { return fmt.Errorf("加密 API Key 失败: %w", err) } saved := models.AIConfig{ Enabled: in.Enabled, Provider: "openai-compatible", BaseURL: strings.TrimSuffix(strings.TrimSpace(in.BaseURL), "/"), APIKey: enc, Model: strings.TrimSpace(in.Model), Temperature: in.Temperature, MaxTokens: in.MaxTokens, TimeoutSec: in.TimeoutSec, DailyTokenCap: in.DailyTokenCap, MaxScanBytes: in.MaxScanBytes, RedactExtra: in.RedactExtra, } if saved.MaxTokens <= 0 { saved.MaxTokens = defaultAIMaxTokens } if saved.TimeoutSec <= 0 { saved.TimeoutSec = defaultAITimeoutSec } if saved.MaxScanBytes <= 0 { saved.MaxScanBytes = defaultAIMaxScanBytes } data, err := json.Marshal(saved) if err != nil { return err } db := repository.GetDB() if db == nil { return fmt.Errorf("数据库未初始化") } newID := db.NextID("settings") db.Lock() defer db.Unlock() for i := range db.Settings { if db.Settings[i].Key == aiConfigSettingKey { db.Settings[i].Value = string(data) db.RequestSave() return nil } } db.Settings = append(db.Settings, models.SystemSetting{ID: newID, Key: aiConfigSettingKey, Value: string(data)}) db.RequestSave() return nil } // truncateToDay 取当天 00:00(本地时区),用于"今日已用 token"统计。 func truncateToDay() time.Time { now := time.Now() return time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) } // AITokenUsedToday 某用户今日已消耗的 token 总量(每日配额闸门用)。 func AITokenUsedToday(userID uint) int { db := repository.GetDB() if db == nil { return 0 } today := truncateToDay() db.RLock() defer db.RUnlock() total := 0 for _, r := range db.AIAnalysisRuns { if r.UserID == userID && !r.CreatedAt.Before(today) { total += r.PromptTokens + r.CompletionTokens } } return total }