实现ai

internal/infra/ai/ai.go

@@ -0,0 +1,19 @@
+package ai
+
+import (
+	"context"
+
+	"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
+)
+
+// AI 定义了通用的 AI 管理接口。
+// 它可以用于处理各种 AI 相关的任务，例如文本生成、内容审核等。
+type AI interface {
+	// GenerateReview 根据提供的文本内容生成评论。
+	// prompt: 用于生成评论的输入文本。
+	// 返回生成的评论字符串和可能发生的错误。
+	GenerateReview(ctx context.Context, prompt string) (string, error)
+
+	// AIModel 返回当前使用的 AI 模型。
+	AIModel() models.AIModel
+}

internal/infra/ai/gemini.go

@@ -0,0 +1,73 @@
+package ai
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"git.huangwc.com/pig/pig-farm-controller/internal/infra/config"
+	"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"
+	"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
+
+	"github.com/google/generative-ai-go/genai"
+	"google.golang.org/api/option"
+)
+
+// geminiImpl 是 Gemini AI 服务的实现。
+type geminiImpl struct {
+	client *genai.GenerativeModel
+	cfg    *config.Gemini
+}
+
+// NewGeminiAI 创建一个新的 geminiImpl 实例。
+func NewGeminiAI(ctx context.Context, cfg *config.AIConfig) (AI, error) {
+	// 检查 API Key 是否存在
+	if cfg.Gemini.APIKey == "" {
+		return nil, fmt.Errorf("Gemini API Key 未配置")
+	}
+
+	// 创建 Gemini 客户端
+	genaiClient, err := genai.NewClient(ctx, option.WithAPIKey(cfg.Gemini.APIKey))
+	if err != nil {
+		return nil, fmt.Errorf("创建 Gemini 客户端失败: %w", err)
+	}
+
+	return &geminiImpl{
+		client: genaiClient.GenerativeModel(cfg.Gemini.ModelName),
+		cfg:    &cfg.Gemini,
+	}, nil
+}
+
+// GenerateReview 根据提供的文本内容生成评论。
+func (g *geminiImpl) GenerateReview(ctx context.Context, prompt string) (string, error) {
+	serviceCtx, logger := logs.Trace(ctx, context.Background(), "GenerateReview")
+	logger.Debugf("开始调用 Gemini 生成评论，prompt: %s", prompt)
+
+	timeoutCtx, cancel := context.WithTimeout(serviceCtx, time.Duration(g.cfg.Timeout)*time.Second)
+	defer cancel()
+
+	resp, err := g.client.GenerateContent(timeoutCtx, genai.Text(prompt))
+	if err != nil {
+		logger.Errorf("调用 Gemini API 失败: %v", err)
+		return "", fmt.Errorf("调用 Gemini API 失败: %w", err)
+	}
+
+	if resp == nil || len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
+		logger.Warn("Gemini API 返回空内容或无候选评论")
+		return "", fmt.Errorf("Gemini API 返回空内容或无候选评论")
+	}
+
+	var review string
+	for _, part := range resp.Candidates[0].Content.Parts {
+		if txt, ok := part.(genai.Text); ok {
+			review += string(txt)
+		}
+	}
+
+	logger.Debugf("成功从 Gemini 生成评论: %s", review)
+	return review, nil
+}
+
+func (g *geminiImpl) AIModel() models.AIModel {
+	return models.AI_MODEL_GEMINI
+}