LokiStack/backend/services/ai_config.go

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
}