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
|
||
}
|