46 lines
990 B
Go
46 lines
990 B
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"github.com/google/generative-ai-go/genai"
|
|
"google.golang.org/api/option"
|
|
)
|
|
|
|
type GeminiService struct {
|
|
apiKey string
|
|
client *genai.Client
|
|
}
|
|
|
|
func NewGeminiService(apiKey string) (*GeminiService, error) {
|
|
ctx := context.Background()
|
|
client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &GeminiService{
|
|
apiKey: apiKey,
|
|
client: client,
|
|
}, nil
|
|
}
|
|
|
|
func (s *GeminiService) Close() {
|
|
if s.client != nil {
|
|
s.client.Close()
|
|
}
|
|
}
|
|
|
|
func (s *GeminiService) ProcessText(ctx context.Context, prompt string) (string, error) {
|
|
model := s.client.GenerativeModel("gemini-2.0-flash-exp")
|
|
resp, err := model.GenerateContent(ctx, genai.Text(prompt))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if len(resp.Candidates) > 0 && len(resp.Candidates[0].Content.Parts) > 0 {
|
|
if textPart, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
|
|
return string(textPart), nil
|
|
}
|
|
}
|
|
return "", nil
|
|
}
|