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 }