125 lines
4.0 KiB
Go
125 lines
4.0 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"encoding/json"
|
||
"strings"
|
||
|
||
"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
|
||
}
|
||
|
||
type RateResult struct {
|
||
Rate int `json:"rate"`
|
||
Summary string `json:"summary"`
|
||
DetailedReview string `json:"detailed_review"`
|
||
}
|
||
|
||
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-2.0-flash-exp"),
|
||
apiKey: apiKey,
|
||
}, nil
|
||
}
|
||
|
||
func (s *GeminiService) ProcessText(ctx context.Context, content string, criteria string, writingRequirement string) (*RateResult, error) {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
|
||
if criteria == "" {
|
||
criteria = `你是一名语文老师。你正在给学生的作文打分。根据以下中考作文评分标准，给作文打分。
|
||
## 评分总分值：100分。
|
||
### 88-100分 符合题意；写作目的和对象明确；思考充分，立意深刻，感情真挚；选材精当，内容充实；中心突出，条理清晰；表达准确，语言流畅。
|
||
### 75-87分 符合题意；写作目的和对象较明确；思考较充分，立意清楚，感情真实；选材合理，内容具体；中心明确，有一定条理；表达较准确，语言通畅。
|
||
### 60-74分 符合题意；写作目的和对象较模糊；有一定思考，感情真实；有一定内容；结构基本完整；语言尚通畅。
|
||
### 60分以下 不符合题意；缺乏写作目的和对象；基本没有思考，感情虚假；内容空洞；结构混乱；不成篇。`
|
||
}
|
||
|
||
if writingRequirement == "" {
|
||
writingRequirement = "写一篇不少于600字的作文，体裁不限。"
|
||
}
|
||
|
||
format := `请按照以下JSON格式输出：
|
||
{
|
||
"rate": 分数, // 最多100分制的分数，int类型
|
||
"summary": "总体评价", // 100字以内的总体评价，string类型
|
||
"detailed_review": "详细点评" // 300字以内的详细点评，包含优点和建议，string类型
|
||
}`
|
||
|
||
prompt := "作文要求：\n" + writingRequirement + "\n\n" + "评分标准：\n" + criteria + format + "\n\n" + "\n\n作文内容：\n" + content
|
||
|
||
resp, err := s.model.GenerateContent(ctx, genai.Text(prompt))
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "failed to generate content")
|
||
}
|
||
|
||
if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
|
||
return nil, errors.New("no response from Gemini")
|
||
}
|
||
|
||
if textPart, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
|
||
return s.parseRateResponse(string(textPart))
|
||
}
|
||
|
||
return nil, errors.New("invalid response format from Gemini")
|
||
}
|
||
|
||
func (s *GeminiService) parseRateResponse(response string) (*RateResult, error) {
|
||
//去除所有\n
|
||
response = strings.ReplaceAll(response, "\n", "")
|
||
//去除所有\t
|
||
response = strings.ReplaceAll(response, "\t", "")
|
||
// 去除response中的```json前缀和```后缀
|
||
response = strings.TrimSpace(response)
|
||
response = strings.TrimPrefix(response, "```json")
|
||
response = strings.TrimSuffix(response, "```")
|
||
|
||
// 检查response是否是json格式
|
||
if !strings.HasPrefix(response, "{") {
|
||
return nil, errors.New("response missing opening brace")
|
||
}
|
||
if !strings.HasSuffix(response, "}") {
|
||
return nil, errors.New("response missing closing brace")
|
||
}
|
||
|
||
var result RateResult
|
||
err := json.Unmarshal([]byte(response), &result)
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "failed to parse JSON response")
|
||
}
|
||
|
||
// 合并所有验证条件
|
||
if result.Rate <= 0 || result.Rate > 100 ||
|
||
result.Summary == "" || result.DetailedReview == "" {
|
||
return nil, errors.New("invalid response fields")
|
||
}
|
||
|
||
return &result, nil
|
||
}
|
||
|
||
// 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
|
||
}
|