pig-farm-controller/internal/infra/ai/gemini.go

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.Gemini) (AI, error) {
	// 检查 API Key 是否存在
	if cfg.APIKey == "" {
		return nil, fmt.Errorf("Gemini API Key 未配置")
	}

	// 创建 Gemini 客户端
	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.ModelName),
		cfg:    cfg,
	}, 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
}