优化ai初始化逻辑

This commit is contained in:
2025-12-02 16:34:14 +08:00
parent da2c296c05
commit 6764684fe7
7 changed files with 56 additions and 11 deletions

View File

@@ -128,6 +128,7 @@ alarm_notification:
# AI 服务配置
ai:
model: "gemini" # 不指定就是不用AI
gemini:
api_key: "YOUR_GEMINI_API_KEY" # 替换为你的 Gemini API Key
model_name: "gemini-2.5-flash" # Gemini 模型名称,例如 "gemini-pro"

View File

@@ -106,6 +106,7 @@ alarm_notification:
# AI 服务配置
ai:
model: Gemini
gemini:
api_key: "AIzaSyAJdXUmoN07LIswDac6YxPeRnvXlR73OO8" # 替换为你的 Gemini API Key
model_name: "gemini-2.0-flash" # Gemini 模型名称,例如 "gemini-pro"

View File

@@ -35,7 +35,7 @@ type Infrastructure struct {
storage database.Storage
repos *Repositories
lora *LoraComponents
aiManager infra_ai.AI
ai infra_ai.AI
tokenGenerator token.Generator
}
@@ -56,7 +56,7 @@ func initInfrastructure(ctx context.Context, cfg *config.Config) (*Infrastructur
tokenGenerator := token.NewTokenGenerator([]byte(cfg.App.JWTSecret))
// 初始化 AI
aiManager, err := infra_ai.NewGeminiAI(logs.AddCompName(ctx, "GeminiAI"), &cfg.AI)
ai, err := initAI(ctx, cfg.AI)
if err != nil {
return nil, fmt.Errorf("初始化 AI 管理器失败: %w", err)
}
@@ -65,7 +65,7 @@ func initInfrastructure(ctx context.Context, cfg *config.Config) (*Infrastructur
storage: storage,
repos: repos,
lora: lora,
aiManager: aiManager,
ai: ai,
tokenGenerator: tokenGenerator,
}, nil
}
@@ -248,7 +248,7 @@ func initDomainServices(ctx context.Context, cfg *config.Config, infra *Infrastr
recipeCoreService,
recipeGenerateManager,
infra.repos.recipeRepo,
infra.aiManager,
infra.ai,
)
return &DomainServices{
@@ -519,3 +519,12 @@ func initStorage(ctx context.Context, cfg config.DatabaseConfig) (database.Stora
logs.GetLogger(ctx).Info("数据库初始化完成。")
return storage, nil
}
func initAI(ctx context.Context, cfg config.AIConfig) (infra_ai.AI, error) {
switch cfg.Model {
case models.AI_MODEL_GEMINI:
return infra_ai.NewGeminiAI(ctx, cfg.Gemini)
default:
return infra_ai.NewNoneAI(ctx), nil
}
}

View File

@@ -16,25 +16,25 @@ import (
// geminiImpl 是 Gemini AI 服务的实现。
type geminiImpl struct {
client *genai.GenerativeModel
cfg *config.Gemini
cfg config.Gemini
}
// NewGeminiAI 创建一个新的 geminiImpl 实例。
func NewGeminiAI(ctx context.Context, cfg *config.AIConfig) (AI, error) {
func NewGeminiAI(ctx context.Context, cfg config.Gemini) (AI, error) {
// 检查 API Key 是否存在
if cfg.Gemini.APIKey == "" {
if cfg.APIKey == "" {
return nil, fmt.Errorf("Gemini API Key 未配置")
}
// 创建 Gemini 客户端
genaiClient, err := genai.NewClient(ctx, option.WithAPIKey(cfg.Gemini.APIKey))
genaiClient, err := genai.NewClient(ctx, option.WithAPIKey(cfg.APIKey))
if err != nil {
return nil, fmt.Errorf("创建 Gemini 客户端失败: %w", err)
}
return &geminiImpl{
client: genaiClient.GenerativeModel(cfg.Gemini.ModelName),
cfg: &cfg.Gemini,
client: genaiClient.GenerativeModel(cfg.ModelName),
cfg: cfg,
}, nil
}

View File

@@ -0,0 +1,31 @@
package ai
import (
"context"
"errors"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
)
var NoneAIError = errors.New("当前没有配置AI, 暂不支持此功能")
type NoneAI struct {
ctx context.Context
}
func NewNoneAI(ctx context.Context) AI {
return &NoneAI{
ctx: ctx,
}
}
func (n *NoneAI) GenerateReview(ctx context.Context, prompt string) (string, error) {
logger := logs.TraceLogger(ctx, n.ctx, "GenerateReview")
logger.Warnf("当前没有配置AI, 无法处理AI请求, 消息: %s", prompt)
return "", NoneAIError
}
func (n *NoneAI) AIModel() models.AIModel {
return models.AI_MODEL_NONE
}

View File

@@ -7,6 +7,7 @@ import (
"fmt"
"os"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
"gopkg.in/yaml.v2"
)
@@ -236,7 +237,8 @@ type AlarmNotificationConfig struct {
// AIConfig AI 服务配置
type AIConfig struct {
Gemini Gemini `yaml:"gemini"`
Model models.AIModel `yaml:"model"`
Gemini Gemini `yaml:"gemini"`
}
// Gemini 代表 Gemini AI 服务的配置

View File

@@ -15,6 +15,7 @@ import (
type AIModel string
const (
AI_MODEL_NONE AIModel = "None"
AI_MODEL_GEMINI AIModel = "Gemini"
)