326 lines
9.9 KiB
Go
326 lines
9.9 KiB
Go
package main
|
||
|
||
import (
|
||
"context"
|
||
"flag"
|
||
"fmt"
|
||
"os"
|
||
"os/signal"
|
||
"strconv"
|
||
"syscall"
|
||
"time"
|
||
|
||
_ "dubbo.apache.org/dubbo-go/v3/imports"
|
||
_ "dubbo.apache.org/dubbo-go/v3/protocol/triple"
|
||
"dubbo.apache.org/dubbo-go/v3/protocol"
|
||
"dubbo.apache.org/dubbo-go/v3/server"
|
||
"github.com/joho/godotenv"
|
||
"github.com/redis/go-redis/v9"
|
||
"gorm.io/gorm"
|
||
|
||
"github.com/topfans/backend/pkg/database"
|
||
"github.com/topfans/backend/pkg/health"
|
||
"github.com/topfans/backend/pkg/logger"
|
||
"github.com/topfans/backend/services/aiChatService/model"
|
||
"github.com/topfans/backend/services/aiChatService/provider"
|
||
"github.com/topfans/backend/services/aiChatService/repository"
|
||
"github.com/topfans/backend/services/aiChatService/service"
|
||
pbAIChat "github.com/topfans/backend/pkg/proto/ai_chat"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
var (
|
||
port = flag.Int("port", getEnvInt("PORT", 20008), "Dubbo service port")
|
||
dbHost = flag.String("db-host", getEnv("DB_HOST", "localhost"), "Database host")
|
||
dbPort = flag.Int("db-port", getEnvInt("DB_PORT", 5432), "Database port")
|
||
dbUser = flag.String("db-user", getEnv("DB_USER", "postgres"), "Database user")
|
||
dbPassword = flag.String("db-password", getEnv("DB_PASSWORD", ""), "Database password")
|
||
dbName = flag.String("db-name", getEnv("DB_NAME", "top-fans"), "Database name")
|
||
redisHost = flag.String("redis-host", getEnv("REDIS_HOST", "127.0.0.1"), "Redis host")
|
||
redisPort = flag.Int("redis-port", getEnvInt("REDIS_PORT", 6379), "Redis port")
|
||
redisPassword = flag.String("redis-password", getEnv("REDIS_PASSWORD", ""), "Redis password")
|
||
redisDB = flag.Int("redis-db", getEnvInt("REDIS_DB", 0), "Redis db")
|
||
contextTTL = flag.Int("context-ttl", getEnvInt("CONTEXT_TTL", 604800), "Context TTL in seconds (default 7 days)")
|
||
healthHandler *health.Handler
|
||
)
|
||
|
||
func getEnv(key, fallback string) string {
|
||
if v := os.Getenv(key); v != "" {
|
||
return v
|
||
}
|
||
return fallback
|
||
}
|
||
|
||
func getEnvInt(key string, fallback int) int {
|
||
if v := os.Getenv(key); v != "" {
|
||
if n, err := strconv.Atoi(v); err == nil {
|
||
return n
|
||
}
|
||
}
|
||
return fallback
|
||
}
|
||
|
||
func main() {
|
||
godotenv.Load()
|
||
flag.Parse()
|
||
|
||
env := os.Getenv("ENV")
|
||
if env == "" {
|
||
env = "development"
|
||
}
|
||
|
||
if err := logger.Init(logger.Config{
|
||
ServiceName: "ai-chat-service",
|
||
Environment: env,
|
||
LogLevel: os.Getenv("LOG_LEVEL"),
|
||
}); err != nil {
|
||
panic(fmt.Sprintf("Failed to initialize logger: %v", err))
|
||
}
|
||
defer logger.Sync()
|
||
|
||
logger.Logger.Info("Starting AI Chat Service...")
|
||
|
||
// 初始化数据库
|
||
dbConfig := database.Config{
|
||
Host: *dbHost,
|
||
Port: *dbPort,
|
||
User: *dbUser,
|
||
Password: *dbPassword,
|
||
DBName: *dbName,
|
||
SSLMode: "disable",
|
||
TimeZone: "Asia/Shanghai",
|
||
}
|
||
|
||
if err := database.Init(dbConfig); err != nil {
|
||
logger.Logger.Fatal(fmt.Sprintf("Failed to initialize database: %v", err))
|
||
}
|
||
logger.Logger.Info("Database initialized successfully")
|
||
|
||
// 初始化 Redis
|
||
redisClient := redis.NewClient(&redis.Options{
|
||
Addr: fmt.Sprintf("%s:%d", *redisHost, *redisPort),
|
||
Password: *redisPassword,
|
||
DB: *redisDB,
|
||
})
|
||
|
||
// 测试 Redis 连接
|
||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
if err := redisClient.Ping(ctx).Err(); err != nil {
|
||
logger.Logger.Fatal(fmt.Sprintf("Failed to connect to Redis: %v", err))
|
||
}
|
||
cancel()
|
||
logger.Logger.Info("Redis connected successfully")
|
||
|
||
// 启动健康检查 HTTP 服务器
|
||
healthPort := *port + 1000
|
||
healthHandler = health.NewHandler("ai-chat-service", healthPort)
|
||
healthHandler.Start()
|
||
|
||
// 自动迁移数据库表
|
||
if err := autoMigrate(); err != nil {
|
||
logger.Logger.Fatal(fmt.Sprintf("Failed to migrate database: %v", err))
|
||
}
|
||
|
||
// 创建 Repository 层实例
|
||
starRepo := repository.NewPostgreSQLStarRepository(database.GetDB())
|
||
personaRepo := repository.NewPostgreSQLPersonaRepository(database.GetDB(), starRepo)
|
||
configRepo := repository.NewPostgreSQLConfigRepository(database.GetDB())
|
||
shortTermMemoryRepo := repository.NewRedisMemoryRepository(redisClient, *contextTTL)
|
||
longTermMemoryRepo := repository.NewPostgreSQLMemoryRepository(database.GetDB())
|
||
logger.Logger.Info("Repository layer initialized")
|
||
|
||
// 从数据库加载 LLM 配置
|
||
loadCtx := context.Background()
|
||
llmConfigs, err := configRepo.GetByCategory(loadCtx, "llm")
|
||
if err != nil {
|
||
logger.Logger.Warn("Failed to load LLM configs from database, using env defaults", zap.Error(err))
|
||
}
|
||
|
||
// 获取 MiniMax 配置
|
||
miniMaxAPIKey := getEnv("MINIMAX_API_KEY", "")
|
||
miniMaxAPIURL := getEnv("MINIMAX_API_URL", "https://api.minimaxi.com/v1")
|
||
miniMaxModel := getEnv("MINIMAX_MODEL", "M2-her")
|
||
if val, ok := llmConfigs["minimax.api_key"]; ok && val != "" {
|
||
miniMaxAPIKey = val
|
||
}
|
||
if val, ok := llmConfigs["minimax.api_url"]; ok && val != "" {
|
||
miniMaxAPIURL = val
|
||
}
|
||
if val, ok := llmConfigs["minimax.model"]; ok && val != "" {
|
||
miniMaxModel = val
|
||
}
|
||
|
||
// 获取 Qwen 配置
|
||
qwenAPIKey := getEnv("QWEN_API_KEY", "")
|
||
qwenAPIURL := getEnv("QWEN_API_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
|
||
qwenModel := getEnv("QWEN_MODEL", "qwen-plus")
|
||
if val, ok := llmConfigs["qwen.api_key"]; ok && val != "" {
|
||
qwenAPIKey = val
|
||
}
|
||
if val, ok := llmConfigs["qwen.api_url"]; ok && val != "" {
|
||
qwenAPIURL = val
|
||
}
|
||
if val, ok := llmConfigs["qwen.model"]; ok && val != "" {
|
||
qwenModel = val
|
||
}
|
||
|
||
logger.Logger.Info("LLM configs loaded",
|
||
zap.String("minimax_url", miniMaxAPIURL),
|
||
zap.String("minimax_model", miniMaxModel),
|
||
zap.Bool("has_minimax_key", miniMaxAPIKey != ""),
|
||
zap.String("qwen_url", qwenAPIURL),
|
||
zap.String("qwen_model", qwenModel),
|
||
zap.Bool("has_qwen_key", qwenAPIKey != ""),
|
||
)
|
||
|
||
// 获取 Dify 配置（MVP 阶段从环境变量读）
|
||
difyAPIKey := getEnv("DIFY_API_KEY", "")
|
||
difyAPIBase := getEnv("DIFY_API_BASE", "")
|
||
if val, ok := llmConfigs["dify.api_key"]; ok && val != "" {
|
||
difyAPIKey = val
|
||
}
|
||
if val, ok := llmConfigs["dify.api_base"]; ok && val != "" {
|
||
difyAPIBase = val
|
||
}
|
||
logger.Logger.Info("Dify config loaded",
|
||
zap.String("dify_api_base", difyAPIBase),
|
||
zap.Bool("has_dify_key", difyAPIKey != ""),
|
||
)
|
||
|
||
// 创建 AI Provider（Dify 或 LLM）
|
||
var aiProvider service.AIProvider
|
||
if difyAPIKey != "" {
|
||
difyClient := service.NewDifyClient(service.DifyConfig{
|
||
APIBase: difyAPIBase,
|
||
APIKey: difyAPIKey,
|
||
TimeoutSec: 120,
|
||
})
|
||
aiProvider = service.NewDifyAdapter(difyClient)
|
||
logger.Logger.Info("Using Dify as AI provider")
|
||
} else {
|
||
llmService := service.NewLLMService(
|
||
miniMaxAPIURL,
|
||
miniMaxAPIKey,
|
||
miniMaxModel,
|
||
qwenAPIURL,
|
||
qwenAPIKey,
|
||
qwenModel,
|
||
)
|
||
aiProvider = llmService
|
||
logger.Logger.Info("Using LLM (MiniMax) as AI provider")
|
||
}
|
||
|
||
// 创建 Service 层实例
|
||
personaService := service.NewPersonaService(personaRepo)
|
||
memoryService := service.NewMemoryService(shortTermMemoryRepo, longTermMemoryRepo)
|
||
auditService := service.NewAuditService()
|
||
chatService := service.NewChatService(
|
||
aiProvider,
|
||
personaService,
|
||
memoryService,
|
||
auditService,
|
||
)
|
||
logger.Logger.Info("Service layer initialized")
|
||
|
||
// 创建 Provider 层实例
|
||
aiChatProvider := provider.NewAIChatProvider(
|
||
chatService,
|
||
personaService,
|
||
memoryService,
|
||
auditService,
|
||
)
|
||
logger.Logger.Info("Provider layer initialized")
|
||
|
||
// 创建 Dubbo 服务器
|
||
srv, err := server.NewServer(
|
||
server.WithServerProtocol(
|
||
protocol.WithPort(*port),
|
||
protocol.WithTriple(),
|
||
),
|
||
)
|
||
if err != nil {
|
||
logger.Logger.Fatal(fmt.Sprintf("Failed to create Dubbo server: %v", err))
|
||
}
|
||
|
||
// 注册 AI Chat Service
|
||
if err := pbAIChat.RegisterAIChatServiceHandler(srv, aiChatProvider); err != nil {
|
||
logger.Logger.Fatal(fmt.Sprintf("Failed to register AI Chat Service: %v", err))
|
||
}
|
||
|
||
// 启动服务
|
||
if err := srv.Serve(); err != nil {
|
||
logger.Logger.Fatal(fmt.Sprintf("Failed to start AI Chat Service: %v", err))
|
||
}
|
||
|
||
logger.Logger.Info(fmt.Sprintf("AI Chat Service started successfully on port %d", *port))
|
||
|
||
// 等待退出信号
|
||
quit := make(chan os.Signal, 1)
|
||
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
||
<-quit
|
||
|
||
logger.Logger.Info("Shutting down AI Chat Service...")
|
||
|
||
if healthHandler != nil {
|
||
healthHandler.Stop()
|
||
}
|
||
|
||
redisClient.Close()
|
||
}
|
||
|
||
// autoMigrate 自动迁移数据库表
|
||
func autoMigrate() error {
|
||
db := database.GetDB()
|
||
if db == nil {
|
||
return fmt.Errorf("database is not initialized")
|
||
}
|
||
|
||
tables := []interface{}{
|
||
&model.Persona{},
|
||
&model.UserMemory{},
|
||
&model.Config{},
|
||
}
|
||
|
||
for _, table := range tables {
|
||
if err := db.AutoMigrate(table); err != nil {
|
||
// 如果约束删除失败（表已存在该约束），忽略并继续
|
||
logger.Logger.Warn(fmt.Sprintf("Migration warning for %T: %v", table, err))
|
||
}
|
||
}
|
||
|
||
// 手动迁移：将 ai_personas 的唯一索引从 (user_id) 升级为 (user_id, star_id)
|
||
if err := migratePersonaIndex(db); err != nil {
|
||
logger.Logger.Warn(fmt.Sprintf("Persona index migration warning: %v", err))
|
||
}
|
||
|
||
logger.Logger.Info("Database migration completed successfully")
|
||
return nil
|
||
}
|
||
|
||
// migratePersonaIndex 升级 ai_personas 的部分唯一索引并清理旧数据
|
||
func migratePersonaIndex(db *gorm.DB) error {
|
||
// 1. 清理旧的孤儿记录（star_id=0 的默认人设，会被新逻辑重新创建）
|
||
result := db.Exec(`DELETE FROM ai_personas WHERE star_id = 0 AND is_default = true`)
|
||
if result.Error != nil {
|
||
logger.Logger.Warn("Failed to clean orphaned personas", zap.Error(result.Error))
|
||
} else if result.RowsAffected > 0 {
|
||
logger.Logger.Info("Cleaned orphaned personas (star_id=0)",
|
||
zap.Int64("count", result.RowsAffected),
|
||
)
|
||
}
|
||
|
||
// 2. 删除旧索引（仅 (user_id) 维度的部分唯一索引）
|
||
dropOld := `DROP INDEX IF EXISTS idx_ai_personas_user_default`
|
||
if err := db.Exec(dropOld).Error; err != nil {
|
||
logger.Logger.Warn("Failed to drop old persona index (may not exist)", zap.Error(err))
|
||
}
|
||
|
||
// 3. 创建新索引（(user_id, star_id) 维度的部分唯一索引）
|
||
createNew := `CREATE UNIQUE INDEX IF NOT EXISTS idx_ai_personas_user_star_default
|
||
ON ai_personas (user_id, star_id) WHERE is_default = true`
|
||
if err := db.Exec(createNew).Error; err != nil {
|
||
return fmt.Errorf("failed to create new persona index: %w", err)
|
||
}
|
||
|
||
return nil
|
||
} |