package aiagent import ( "code.tczkiot.com/wlw/ai-agent/cmd/testdata/seedlang" "code.tczkiot.com/wlw/ai-agent/cmd/testdata/seeds" "code.tczkiot.com/wlw/ai-agent/internal/models" "code.tczkiot.com/wlw/ai-agent/internal/pkg/enums" "code.tczkiot.com/wlw/ai-agent/internal/pkg/utils" "code.tczkiot.com/wlw/ai-agent/internal/repositories" "fmt" "strings" "time" "github.com/mlogclub/simple/sqls" ) type InitResult struct { Created int Updated int } // Init 初始化 AI Agent 测试数据 // 依赖于 AI Config 和 Knowledge Base 已初始化 func Init(lang seedlang.Language) (*InitResult, error) { result := &InitResult{} aiConfigID, err := getDefaultAIConfigID() if err != nil { return result, fmt.Errorf("get default ai config id failed: %w", err) } if aiConfigID == 0 { return result, fmt.Errorf("no default ai config found, please init ai config first") } knowledgeIDs, err := getDefaultKnowledgeIDs() if err != nil { return result, fmt.Errorf("get default knowledge ids failed: %w", err) } defaultTeamIDs := getDefaultTeamIDs() agentSeeds := seeds.AIAgentSeeds(lang) seedItems := buildModels(lang, aiConfigID, knowledgeIDs, defaultTeamIDs) for index, item := range seedItems { itemCopy := item if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error { existing := repositories.AIAgentRepository.Take(ctx.Tx, "name = ?", itemCopy.Name) if existing == nil { for _, legacyName := range agentSeeds[index].LegacyNames { legacyName = strings.TrimSpace(legacyName) if legacyName == "" { continue } existing = repositories.AIAgentRepository.Take(ctx.Tx, "name = ?", legacyName) if existing != nil { break } } } if existing != nil { if err := ctx.Tx.Model(existing).Updates(seedUpdateColumns(itemCopy)).Error; err != nil { return err } result.Updated++ } else { if err := ctx.Tx.Create(&itemCopy).Error; err != nil { return err } result.Created++ } return nil }); err != nil { return nil, fmt.Errorf("upsert ai agent failed: %w", err) } } return result, nil } func buildModels(lang seedlang.Language, aiConfigID int64, knowledgeIDs []int64, defaultTeamIDs string) []models.AIAgent { now := time.Now() seedItems := seeds.AIAgentSeeds(lang) items := make([]models.AIAgent, 0, len(seedItems)) for _, seed := range seedItems { items = append(items, models.AIAgent{ Name: seed.Name, Description: seed.Description, Status: enums.StatusOk, AIConfigID: aiConfigID, ServiceMode: seed.ServiceMode, SystemPrompt: seed.SystemPrompt, WelcomeMessage: seed.WelcomeMessage, ReplyTimeoutSeconds: seed.ReplyTimeoutSeconds, TeamIDs: defaultTeamIDs, HandoffMode: seed.HandoffMode, FallbackMode: seed.FallbackMode, FallbackMessage: seed.FallbackMessage, KnowledgeIDs: utils.JoinInt64s(knowledgeIDs), SortNo: seed.SortNo, AuditFields: models.AuditFields{ CreatedAt: now, CreateUserID: 0, CreateUserName: "System", UpdatedAt: now, UpdateUserID: 0, UpdateUserName: "System", }, }) } return items } func seedUpdateColumns(item models.AIAgent) map[string]any { return map[string]any{ "name": item.Name, "description": item.Description, "status": item.Status, "ai_config_id": item.AIConfigID, "service_mode": item.ServiceMode, "system_prompt": item.SystemPrompt, "welcome_message": item.WelcomeMessage, "reply_timeout_seconds": item.ReplyTimeoutSeconds, "team_ids": item.TeamIDs, "handoff_mode": item.HandoffMode, "fallback_mode": item.FallbackMode, "fallback_message": item.FallbackMessage, "knowledge_ids": item.KnowledgeIDs, "sort_no": item.SortNo, "updated_at": item.UpdatedAt, "update_user_id": item.UpdateUserID, "update_user_name": item.UpdateUserName, } } func getDefaultAIConfigID() (int64, error) { aiConfig := repositories.AIConfigRepository.Take( sqls.DB(), "model_type = ? AND status = ?", string(enums.AIModelTypeLLM), enums.StatusOk, ) if aiConfig == nil { return 0, nil } return aiConfig.ID, nil } func getDefaultKnowledgeIDs() ([]int64, error) { knowledges := repositories.KnowledgeBaseRepository.Find( sqls.DB(), sqls.NewCnd().Where("status = ?", enums.StatusOk), ) ids := make([]int64, 0, len(knowledges)) for _, knowledge := range knowledges { ids = append(ids, knowledge.ID) } return ids, nil } func getDefaultTeamIDs() string { teams := repositories.AgentTeamRepository.Find( sqls.DB(), sqls.NewCnd().Where("status = ?", enums.StatusOk), ) teamIDs := make([]int64, 0, len(teams)) for _, team := range teams { teamIDs = append(teamIDs, team.ID) } return utils.JoinInt64s(teamIDs) }