refactor: initialize expert catalog from initdb
This commit is contained in:
117
backend/internal/initdb/expert.go
Normal file
117
backend/internal/initdb/expert.go
Normal 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
|
||||
}
|
||||
77
backend/internal/initdb/expert_test.go
Normal file
77
backend/internal/initdb/expert_test.go
Normal 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)
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"gorm.io/gorm/logger"
|
||||
"senlinai-agent/backend/internal/config"
|
||||
"senlinai-agent/backend/internal/httpx"
|
||||
"senlinai-agent/backend/internal/initdb"
|
||||
"senlinai-agent/backend/internal/models"
|
||||
)
|
||||
|
||||
@@ -333,6 +334,7 @@ func newAIHandlerTestDB(t *testing.T) *gorm.DB {
|
||||
database, err := gorm.Open(sqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, models.AutoMigrate(database))
|
||||
require.NoError(t, initdb.InitExpert(database))
|
||||
models.DBService = database
|
||||
return database
|
||||
}
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
@@ -32,8 +31,6 @@ func runVersionedMigrations(database *gorm.DB) error {
|
||||
migrations := []versionedMigration{
|
||||
{version: 1, name: "normalize_project_identifiers", run: normalizeLegacyProjectIdentifiers},
|
||||
{version: 2, name: "deduplicate_project_tags", run: deduplicateLegacyProjectTags},
|
||||
{version: 3, name: "seed_ai_experts", run: seedAIExperts},
|
||||
{version: 4, name: "normalize_ai_expert_categories", run: normalizeAIExpertCategories},
|
||||
}
|
||||
for _, migration := range migrations {
|
||||
// Connection callbacks can provide an initialized Gorm session. Start each
|
||||
@@ -66,108 +63,6 @@ func runVersionedMigrations(database *gorm.DB) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
//go:embed experts_catalog.json
|
||||
var expertCatalogJSON []byte
|
||||
|
||||
type embeddedExpertCatalog struct {
|
||||
Source string `json:"source"`
|
||||
Reference string `json:"reference"`
|
||||
License string `json:"license"`
|
||||
Experts []embeddedExpertCatalogItem `json:"experts"`
|
||||
}
|
||||
|
||||
type embeddedExpertCatalogItem 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"`
|
||||
}
|
||||
|
||||
func seedAIExperts(tx *gorm.DB) (map[string]int, error) {
|
||||
catalog, categories, err := loadAIExpertCatalog(tx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
records := make([]SenlinAgentAIExpertItem, 0, len(catalog.Experts))
|
||||
for _, expert := range catalog.Experts {
|
||||
category := categories[expert.Category]
|
||||
records = append(records, 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 {
|
||||
if err := tx.CreateInBatches(&records, 25).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return map[string]int{"inserted": len(records)}, nil
|
||||
}
|
||||
|
||||
func normalizeAIExpertCategories(tx *gorm.DB) (map[string]int, error) {
|
||||
catalog, categories, err := loadAIExpertCatalog(tx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
linked := 0
|
||||
for _, expert := range catalog.Experts {
|
||||
category := categories[expert.Category]
|
||||
result := tx.Model(&SenlinAgentAIExpertItem{}).Where("slug = ?", expert.Slug).Updates(map[string]any{
|
||||
"category_id": category.ID,
|
||||
"category_identity": category.Identity,
|
||||
"reference": catalog.Reference,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return nil, result.Error
|
||||
}
|
||||
linked += int(result.RowsAffected)
|
||||
}
|
||||
return map[string]int{"categories": len(categories), "linked": linked}, nil
|
||||
}
|
||||
|
||||
func loadAIExpertCatalog(tx *gorm.DB) (embeddedExpertCatalog, map[string]SenlinAgentAIExpertCategory, error) {
|
||||
var catalog embeddedExpertCatalog
|
||||
if err := json.Unmarshal(expertCatalogJSON, &catalog); err != nil {
|
||||
return catalog, nil, fmt.Errorf("decode embedded AI expert catalog: %w", err)
|
||||
}
|
||||
|
||||
categoryNames := make(map[string]string)
|
||||
categoryOrder := make([]string, 0)
|
||||
for _, expert := range catalog.Experts {
|
||||
if _, exists := categoryNames[expert.Category]; exists {
|
||||
continue
|
||||
}
|
||||
categoryNames[expert.Category] = expert.CategoryName
|
||||
categoryOrder = append(categoryOrder, expert.Category)
|
||||
}
|
||||
|
||||
categories := make(map[string]SenlinAgentAIExpertCategory, len(categoryOrder))
|
||||
for _, slug := range categoryOrder {
|
||||
category := SenlinAgentAIExpertCategory{Slug: slug}
|
||||
if err := tx.Where("slug = ?", slug).
|
||||
Attrs(SenlinAgentAIExpertCategory{Name: categoryNames[slug]}).
|
||||
FirstOrCreate(&category).Error; err != nil {
|
||||
return catalog, nil, err
|
||||
}
|
||||
if category.Name != categoryNames[slug] {
|
||||
category.Name = categoryNames[slug]
|
||||
if err := tx.Model(&category).Update("name", category.Name).Error; err != nil {
|
||||
return catalog, nil, err
|
||||
}
|
||||
}
|
||||
categories[slug] = category
|
||||
}
|
||||
return catalog, categories, nil
|
||||
}
|
||||
|
||||
func normalizeLegacyProjectIdentifiers(tx *gorm.DB) (map[string]int, error) {
|
||||
var projects []SenlinAgentProject
|
||||
if err := tx.Model(&SenlinAgentProject{}).Select("id", "owner_id", "identifier").Order("owner_id asc, id asc").Find(&projects).Error; err != nil {
|
||||
|
||||
@@ -57,28 +57,14 @@ func TestAutoMigrateUpgradesLegacyProjectsAndTagsWithoutLosingAssociations(t *te
|
||||
|
||||
var migrationCount int64
|
||||
require.NoError(t, database.Table("senlin_agent_schema_migrations").Count(&migrationCount).Error)
|
||||
require.Equal(t, int64(4), migrationCount)
|
||||
require.Equal(t, int64(2), migrationCount)
|
||||
var auditDetails []string
|
||||
require.NoError(t, database.Table("senlin_agent_schema_migrations").Order("version asc").Pluck("details", &auditDetails).Error)
|
||||
require.Contains(t, auditDetails[0], `"updated":3`)
|
||||
require.Contains(t, auditDetails[1], `"deduplicated":1`)
|
||||
require.Contains(t, auditDetails[2], `"inserted":267`)
|
||||
require.Contains(t, auditDetails[3], `"categories":19`)
|
||||
require.Contains(t, auditDetails[3], `"linked":267`)
|
||||
var categoryCount int64
|
||||
require.NoError(t, database.Model(&SenlinAgentAIExpertCategory{}).Count(&categoryCount).Error)
|
||||
require.Equal(t, int64(19), categoryCount)
|
||||
var expertCount int64
|
||||
require.NoError(t, database.Model(&SenlinAgentAIExpertItem{}).Count(&expertCount).Error)
|
||||
require.Equal(t, int64(267), expertCount)
|
||||
var expert 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)
|
||||
require.NoError(t, AutoMigrate(database), "versioned migrations must be safe to run again")
|
||||
require.NoError(t, database.Table("senlin_agent_schema_migrations").Count(&migrationCount).Error)
|
||||
require.Equal(t, int64(4), migrationCount)
|
||||
require.Equal(t, int64(2), migrationCount)
|
||||
|
||||
require.Error(t, database.Exec(`INSERT INTO senlin_agent_projects (owner_id, name, identifier) VALUES (7, 'still duplicate', 'DUP')`).Error)
|
||||
require.Error(t, database.Exec(`INSERT INTO senlin_agent_tags (project_id, name) VALUES (1, 'UI')`).Error)
|
||||
|
||||
Reference in New Issue
Block a user