ai-agent/cmd/testdata/aiconfig/init.go

package aiconfig

import (
	"fmt"
	"os"
	"strings"
	"time"

	"code.tczkiot.com/wlw/ai-agent/internal/models"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/constants"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"
	"code.tczkiot.com/wlw/ai-agent/internal/repositories"

	"github.com/mlogclub/simple/sqls"
	"gopkg.in/yaml.v3"
	"gorm.io/gorm"
)

const (
	defaultSeedFilePath = "cmd/testdata/aiconfig/ai_config.yaml"
	seedFileEnvKey      = "TESTDATA_AI_CONFIG_FILE"
)

type SeedItem struct {
	Name             string `yaml:"name"`
	Provider         string `yaml:"provider"`
	BaseURL          string `yaml:"baseUrl"`
	APIKey           string `yaml:"apiKey"`
	ModelType        string `yaml:"modelType"`
	ModelName        string `yaml:"modelName"`
	Dimension        int    `yaml:"dimension"`
	MaxContextTokens int    `yaml:"maxContextTokens"`
	MaxOutputTokens  int    `yaml:"maxOutputTokens"`
	TimeoutMS        int    `yaml:"timeoutMs"`
	MaxRetryCount    int    `yaml:"maxRetryCount"`
	RPMLimit         int    `yaml:"rpmLimit"`
	TPMLimit         int    `yaml:"tpmLimit"`
	SortNo           int    `yaml:"sortNo"`
	Remark           string `yaml:"remark"`
}

type InitResult struct {
	FilePath string
	Skipped  bool
	Created  int
	Updated  int
}

type seedPayload struct {
	Items []SeedItem `yaml:"items"`
}

func Init() (*InitResult, error) {
	return InitFromFile(resolveSeedFilePath())
}

func InitFromFile(filePath string) (*InitResult, error) {
	result := &InitResult{FilePath: filePath}

	payload, err := loadSeedPayload(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			result.Skipped = true
			return result, nil
		}
		return nil, err
	}
	if len(payload.Items) == 0 {
		result.Skipped = true
		return result, nil
	}

	err = sqls.WithTransaction(func(ctx *sqls.TxContext) error {
		for _, item := range payload.Items {
			created, upsertErr := upsertAIConfig(ctx.Tx, item)
			if upsertErr != nil {
				return upsertErr
			}
			if created {
				result.Created++
			} else {
				result.Updated++
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	return result, nil
}

func loadSeedPayload(filePath string) (*seedPayload, error) {
	content, err := os.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	payload := &seedPayload{}
	if err := yaml.Unmarshal(content, payload); err != nil {
		return nil, fmt.Errorf("parse ai config seed file failed: %w", err)
	}
	return payload, nil
}

func resolveSeedFilePath() string {
	if filePath := strings.TrimSpace(os.Getenv(seedFileEnvKey)); filePath != "" {
		return filePath
	}
	return defaultSeedFilePath
}

func upsertAIConfig(db *gorm.DB, item SeedItem) (bool, error) {
	model, err := buildAIConfigModel(item)
	if err != nil {
		return false, err
	}

	now := time.Now()
	existing := repositories.AIConfigRepository.FindOne(db, sqls.NewCnd().Eq("name", model.Name).Eq("model_type", model.ModelType))
	if existing == nil {
		model.AuditFields = models.AuditFields{
			CreatedAt:      now,
			CreateUserID:   constants.SystemAuditUserID,
			CreateUserName: constants.SystemAuditUserName,
			UpdatedAt:      now,
			UpdateUserID:   constants.SystemAuditUserID,
			UpdateUserName: constants.SystemAuditUserName,
		}
		if err := repositories.AIConfigRepository.Create(db, model); err != nil {
			return false, err
		}
		return true, nil
	}

	err = repositories.AIConfigRepository.Updates(db, existing.ID, map[string]any{
		"provider":           model.Provider,
		"base_url":           model.BaseURL,
		"api_key":            model.APIKey,
		"model_name":         model.ModelName,
		"dimension":          model.Dimension,
		"max_context_tokens": model.MaxContextTokens,
		"max_output_tokens":  model.MaxOutputTokens,
		"timeout_ms":         model.TimeoutMS,
		"max_retry_count":    model.MaxRetryCount,
		"rpm_limit":          model.RPMLimit,
		"tpm_limit":          model.TPMLimit,
		"status":             model.Status,
		"sort_no":            model.SortNo,
		"remark":             model.Remark,
		"update_user_id":     constants.SystemAuditUserID,
		"update_user_name":   constants.SystemAuditUserName,
		"updated_at":         now,
	})
	if err != nil {
		return false, err
	}
	return false, nil
}

func buildAIConfigModel(item SeedItem) (*models.AIConfig, error) {
	name := strings.TrimSpace(item.Name)
	if name == "" {
		return nil, fmt.Errorf("ai config seed name is required")
	}

	provider, err := parseProvider(item.Provider)
	if err != nil {
		return nil, fmt.Errorf("ai config %s: %w", name, err)
	}
	modelType, err := parseModelType(item.ModelType)
	if err != nil {
		return nil, fmt.Errorf("ai config %s: %w", name, err)
	}

	baseURL := strings.TrimSpace(item.BaseURL)
	if baseURL == "" {
		return nil, fmt.Errorf("ai config %s: baseUrl is required", name)
	}
	modelName := strings.TrimSpace(item.ModelName)
	if modelName == "" {
		return nil, fmt.Errorf("ai config %s: modelName is required", name)
	}

	return &models.AIConfig{
		Name:             name,
		Provider:         provider,
		BaseURL:          baseURL,
		APIKey:           strings.TrimSpace(item.APIKey),
		ModelType:        modelType,
		ModelName:        modelName,
		Dimension:        normalizeNonNegative(item.Dimension),
		MaxContextTokens: normalizeNonNegative(item.MaxContextTokens),
		MaxOutputTokens:  normalizeNonNegative(item.MaxOutputTokens),
		TimeoutMS:        defaultIfNotPositive(item.TimeoutMS, 30000),
		MaxRetryCount:    normalizeNonNegative(item.MaxRetryCount),
		RPMLimit:         normalizeNonNegative(item.RPMLimit),
		TPMLimit:         normalizeNonNegative(item.TPMLimit),
		SortNo:           normalizeNonNegative(item.SortNo),
		Status:           enums.StatusOk,
		Remark:           strings.TrimSpace(item.Remark),
	}, nil
}

func parseProvider(raw string) (enums.AIProvider, error) {
	provider := enums.AIProvider(strings.TrimSpace(raw))
	switch provider {
	case enums.AIProviderOpenAI:
		return provider, nil
	default:
		return "", fmt.Errorf("unsupported provider: %s", raw)
	}
}

func parseModelType(raw string) (enums.AIModelType, error) {
	modelType := enums.AIModelType(strings.TrimSpace(raw))
	switch modelType {
	case enums.AIModelTypeLLM, enums.AIModelTypeEmbedding, enums.AIModelTypeRerank:
		return modelType, nil
	default:
		return "", fmt.Errorf("unsupported modelType: %s", raw)
	}
}

func normalizeNonNegative(value int) int {
	if value < 0 {
		return 0
	}
	return value
}

func defaultIfNotPositive(value int, def int) int {
	if value > 0 {
		return value
	}
	return def
}