refactor: initialize expert catalog from initdb

This commit is contained in:
2026-07-23 14:24:01 +08:00
parent bb90ed1d38
commit 6e7be2f3a4
8 changed files with 206 additions and 124 deletions

View File

@@ -0,0 +1,117 @@
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
}

View File

@@ -0,0 +1,77 @@
package initdb
import (
"fmt"
"testing"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"senlinai-agent/backend/internal/models"
)
func TestNewInitializesUserAndExpertCatalog(t *testing.T) {
database := newExpertTestDatabase(t)
require.NoError(t, New(database))
requireExpertCatalogCounts(t, database)
var user models.SenlinAgentUser
require.NoError(t, database.Where("email = ?", rootUsername).First(&user).Error)
require.Equal(t, rootRole, user.Role)
var expert models.SenlinAgentAIExpertItem
require.NoError(t, database.Preload("ExpertCategory").Where("slug = ?", "academic-geographer").First(&expert).Error)
require.Equal(t, "academic", expert.ExpertCategory.Slug)
require.Equal(t, "学术研究", expert.ExpertCategory.Name)
require.Equal(t, "https://ao.aiolaola.com/experts", expert.Reference)
}
func TestInitExpertOnlyInsertsMissingRecords(t *testing.T) {
database := newExpertTestDatabase(t)
require.NoError(t, InitExpert(database))
var preserved models.SenlinAgentAIExpertItem
require.NoError(t, database.Where("slug = ?", "academic-geographer").First(&preserved).Error)
require.NoError(t, database.Model(&preserved).Updates(map[string]any{
"name": "保留名称",
"category_id": 0,
"category_identity": "",
"reference": "",
}).Error)
require.NoError(t, database.Where("slug = ?", "academic-historian").Delete(&models.SenlinAgentAIExpertItem{}).Error)
require.NoError(t, InitExpert(database))
requireExpertCatalogCounts(t, database)
require.NoError(t, database.Where("slug = ?", "academic-geographer").First(&preserved).Error)
require.Equal(t, "保留名称", preserved.Name)
require.NotZero(t, preserved.CategoryID)
require.NotEmpty(t, preserved.CategoryIdentity)
require.Equal(t, "https://ao.aiolaola.com/experts", preserved.Reference)
var restored models.SenlinAgentAIExpertItem
require.NoError(t, database.Where("slug = ?", "academic-historian").First(&restored).Error)
}
func newExpertTestDatabase(t *testing.T) *gorm.DB {
t.Helper()
database, err := gorm.Open(sqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(
&models.SenlinAgentUser{},
&models.SenlinAgentAIExpertCategory{},
&models.SenlinAgentAIExpertItem{},
))
return database
}
func requireExpertCatalogCounts(t *testing.T, database *gorm.DB) {
t.Helper()
var categoryCount int64
require.NoError(t, database.Model(&models.SenlinAgentAIExpertCategory{}).Count(&categoryCount).Error)
require.Equal(t, int64(19), categoryCount)
var expertCount int64
require.NoError(t, database.Model(&models.SenlinAgentAIExpertItem{}).Count(&expertCount).Error)
require.Equal(t, int64(267), expertCount)
}

File diff suppressed because one or more lines are too long

View File

@@ -4,5 +4,10 @@ import "gorm.io/gorm"
// New initializes the default database records.
func New(database *gorm.DB) error {
return InitUser(database)
return database.Transaction(func(tx *gorm.DB) error {
if err := InitUser(tx); err != nil {
return err
}
return InitExpert(tx)
})
}

View File

@@ -13,7 +13,7 @@ import (
func TestInitUserCreatesRootUserWhenTableIsEmpty(t *testing.T) {
database := newUserTestDatabase(t)
require.NoError(t, New(database))
require.NoError(t, InitUser(database))
var users []models.SenlinAgentUser
require.NoError(t, database.Find(&users).Error)
@@ -34,7 +34,7 @@ func TestInitUserDoesNothingWhenTableIsNotEmpty(t *testing.T) {
}
require.NoError(t, database.Create(&existing).Error)
require.NoError(t, New(database))
require.NoError(t, InitUser(database))
var users []models.SenlinAgentUser
require.NoError(t, database.Find(&users).Error)