118 lines
3.7 KiB
Go
118 lines
3.7 KiB
Go
package initdb
|
|
|
|
import (
|
|
_ "embed"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
"senlinai-agent/backend/internal/models"
|
|
)
|
|
|
|
//go:embed experts_catalog.json
|
|
var expertCatalogJSON []byte
|
|
|
|
type expertCatalog struct {
|
|
Source string `json:"source"`
|
|
Reference string `json:"reference"`
|
|
License string `json:"license"`
|
|
Experts []expertCatalogItem `json:"experts"`
|
|
}
|
|
|
|
type expertCatalogItem struct {
|
|
Slug string `json:"slug"`
|
|
Category string `json:"category"`
|
|
CategoryName string `json:"categoryName"`
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
Emoji string `json:"emoji"`
|
|
Color string `json:"color"`
|
|
SystemPrompt string `json:"systemPrompt"`
|
|
}
|
|
|
|
// InitExpert inserts categories and expert items that are missing from the catalog tables.
|
|
func InitExpert(database *gorm.DB) error {
|
|
var catalog expertCatalog
|
|
if err := json.Unmarshal(expertCatalogJSON, &catalog); err != nil {
|
|
return fmt.Errorf("decode embedded AI expert catalog: %w", err)
|
|
}
|
|
|
|
return database.Transaction(func(tx *gorm.DB) error {
|
|
categories, err := initExpertCategories(tx, catalog.Experts)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return initExpertItems(tx, catalog, categories)
|
|
})
|
|
}
|
|
|
|
func initExpertCategories(tx *gorm.DB, experts []expertCatalogItem) (map[string]models.SenlinAgentAIExpertCategory, error) {
|
|
categoryNames := make(map[string]string)
|
|
categoryOrder := make([]string, 0)
|
|
for _, expert := range experts {
|
|
if _, exists := categoryNames[expert.Category]; exists {
|
|
continue
|
|
}
|
|
categoryNames[expert.Category] = expert.CategoryName
|
|
categoryOrder = append(categoryOrder, expert.Category)
|
|
}
|
|
|
|
records := make([]models.SenlinAgentAIExpertCategory, 0, len(categoryOrder))
|
|
for _, slug := range categoryOrder {
|
|
records = append(records, models.SenlinAgentAIExpertCategory{
|
|
Slug: slug,
|
|
Name: categoryNames[slug],
|
|
})
|
|
}
|
|
if len(records) > 0 {
|
|
if err := tx.Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: "slug"}},
|
|
DoNothing: true,
|
|
}).Create(&records).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
var stored []models.SenlinAgentAIExpertCategory
|
|
if err := tx.Where("slug IN ?", categoryOrder).Find(&stored).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
categories := make(map[string]models.SenlinAgentAIExpertCategory, len(stored))
|
|
for _, category := range stored {
|
|
categories[category.Slug] = category
|
|
}
|
|
if len(categories) != len(categoryOrder) {
|
|
return nil, fmt.Errorf("initialize AI expert categories: expected %d, found %d", len(categoryOrder), len(categories))
|
|
}
|
|
return categories, nil
|
|
}
|
|
|
|
func initExpertItems(tx *gorm.DB, catalog expertCatalog, categories map[string]models.SenlinAgentAIExpertCategory) error {
|
|
records := make([]models.SenlinAgentAIExpertItem, 0, len(catalog.Experts))
|
|
for _, expert := range catalog.Experts {
|
|
category, exists := categories[expert.Category]
|
|
if !exists {
|
|
return fmt.Errorf("initialize AI expert item %q: category %q not found", expert.Slug, expert.Category)
|
|
}
|
|
records = append(records, models.SenlinAgentAIExpertItem{
|
|
Slug: expert.Slug, CategoryID: category.ID, CategoryIdentity: category.Identity,
|
|
Category: expert.Category, CategoryName: expert.CategoryName,
|
|
Name: expert.Name, Description: expert.Description, Emoji: expert.Emoji,
|
|
Color: expert.Color, SystemPrompt: expert.SystemPrompt, Source: catalog.Source,
|
|
Reference: catalog.Reference, SourceLicense: catalog.License, Enabled: true,
|
|
})
|
|
}
|
|
if len(records) == 0 {
|
|
return nil
|
|
}
|
|
return tx.Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: "slug"}},
|
|
DoUpdates: clause.AssignmentColumns([]string{
|
|
"category_id",
|
|
"category_identity",
|
|
"reference",
|
|
}),
|
|
}).CreateInBatches(&records, 25).Error
|
|
}
|