74 lines
2.1 KiB
Go
74 lines
2.1 KiB
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.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
|
|||
|
|
}
|