62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/google/generative-ai-go/genai"
|
|
"google.golang.org/api/option"
|
|
"tencent_ocr/pkg/errors"
|
|
)
|
|
|
|
type GeminiService struct {
|
|
client *genai.Client
|
|
model *genai.GenerativeModel
|
|
apiKey string
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
func NewGeminiService(apiKey string) (*GeminiService, error) {
|
|
ctx := context.Background()
|
|
client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "failed to create Gemini client")
|
|
}
|
|
|
|
return &GeminiService{
|
|
client: client,
|
|
model: client.GenerativeModel("gemini-pro"),
|
|
apiKey: apiKey,
|
|
}, nil
|
|
}
|
|
|
|
func (s *GeminiService) ProcessText(ctx context.Context, text string) (string, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
resp, err := s.model.GenerateContent(ctx, genai.Text(text))
|
|
if err != nil {
|
|
return "", errors.Wrap(err, "failed to generate content")
|
|
}
|
|
|
|
if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
|
|
return "", errors.New("no response from Gemini")
|
|
}
|
|
|
|
if textPart, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
|
|
return string(textPart), nil
|
|
}
|
|
|
|
return "", errors.New("invalid response format from Gemini")
|
|
}
|
|
|
|
// Close implements graceful shutdown
|
|
func (s *GeminiService) Close() error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if s.client != nil {
|
|
s.client.Close()
|
|
}
|
|
return nil
|
|
}
|