242 lines
6.5 KiB
Go
242 lines
6.5 KiB
Go
|
package service
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"bytes"
|
|||
|
|
"context"
|
|||
|
|
"encoding/base64"
|
|||
|
|
"encoding/json"
|
|||
|
|
"fmt"
|
|||
|
|
"io"
|
|||
|
|
"log/slog"
|
|||
|
|
"net/http"
|
|||
|
|
"os"
|
|||
|
|
"regexp"
|
|||
|
|
"strings"
|
|||
|
|
"sync"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"schedule_converter/client"
|
|||
|
|
"schedule_converter/config"
|
|||
|
|
"schedule_converter/models"
|
|||
|
|
apperr "schedule_converter/pkg/errors"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
var (
|
|||
|
|
jsonBlockRegex = regexp.MustCompile("(?s)```json\\s*(\\[[\\s\\S]*?\\])\\s*```")
|
|||
|
|
codeBlockRegex = regexp.MustCompile("(?s)```\\s*(\\[[\\s\\S]*?\\])\\s*```")
|
|||
|
|
jsonArrayRegex = regexp.MustCompile(`(\[[\s\S]*\])`)
|
|||
|
|
regexPool = sync.Pool{
|
|||
|
|
New: func() any {
|
|||
|
|
return nil
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type RecognizerService struct {
|
|||
|
|
client *http.Client
|
|||
|
|
maxRetry int
|
|||
|
|
retryWait time.Duration
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func NewRecognizerService() *RecognizerService {
|
|||
|
|
return &RecognizerService{
|
|||
|
|
client: client.RecognizerClient,
|
|||
|
|
maxRetry: 3,
|
|||
|
|
retryWait: 5 * time.Second,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *RecognizerService) RecognizeFromFile(ctx context.Context, imagePath string) ([]models.Course, error) {
|
|||
|
|
log := slog.With("image_path", imagePath)
|
|||
|
|
log.Info("正在读取图片")
|
|||
|
|
|
|||
|
|
imageData, err := os.ReadFile(imagePath)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, apperr.New("recognize_file", err, "读取图片失败")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
log.Info("图片读取完成", "size", len(imageData))
|
|||
|
|
return r.RecognizeFromBytes(ctx, imageData)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *RecognizerService) RecognizeFromBytes(ctx context.Context, imageData []byte) ([]models.Course, error) {
|
|||
|
|
log := slog.With("size", len(imageData))
|
|||
|
|
log.Info("正在处理图片数据")
|
|||
|
|
|
|||
|
|
imageBase64 := base64.StdEncoding.EncodeToString(imageData)
|
|||
|
|
return r.callAPI(ctx, imageBase64)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *RecognizerService) RecognizeFromBase64(ctx context.Context, imageBase64 string) ([]models.Course, error) {
|
|||
|
|
slog.Info("正在处理base64图片数据", "length", len(imageBase64))
|
|||
|
|
return r.callAPI(ctx, imageBase64)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *RecognizerService) callAPI(ctx context.Context, imageBase64 string) ([]models.Course, error) {
|
|||
|
|
var lastErr error
|
|||
|
|
for retry := range r.maxRetry {
|
|||
|
|
select {
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
return nil, ctx.Err()
|
|||
|
|
default:
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if retry > 0 {
|
|||
|
|
slog.Info("重试中", "attempt", retry+1, "max_retry", r.maxRetry)
|
|||
|
|
time.Sleep(r.retryWait)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
courses, err := r.callAPISingle(ctx, imageBase64)
|
|||
|
|
if err == nil {
|
|||
|
|
return courses, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
lastErr = err
|
|||
|
|
slog.Warn("API调用失败", "attempt", retry+1, "error", err)
|
|||
|
|
|
|||
|
|
if !strings.Contains(err.Error(), "504") {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return nil, apperr.New("recognize_api", lastErr, "已达到最大重试次数")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *RecognizerService) callAPISingle(ctx context.Context, imageBase64 string) ([]models.Course, error) {
|
|||
|
|
cfg := config.Get()
|
|||
|
|
if cfg.Recognizer.APIKey == "" {
|
|||
|
|
return nil, apperr.ErrRecognitionFailed
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
prompt := `请解析以下课表图片，提取所有课程信息并以JSON数组格式返回。
|
|||
|
|
|
|||
|
|
JSON格式要求：
|
|||
|
|
[{"name":"课程名","teacher":"老师名","room":"教室","schedule":[{"day":"星期X","section":[节次数组],"weeks":[周数数组]}]}]
|
|||
|
|
|
|||
|
|
注意事项：
|
|||
|
|
1. day使用"星期一"到"星期日"
|
|||
|
|
2. section节次：上午1-2节为[1,2]，上午3-4节为[3,4]，下午5-6节为[5,6]，下午7-8节为[7,8]，晚上9-10节为[9,10]，晚上11-12节为[11,12]
|
|||
|
|
3. weeks周数：格式如[1,2,3,4,5,6,7,8]或[9,10,11,12,13,14,15,16]
|
|||
|
|
4. 只返回JSON数组，不要任何其他文字。`
|
|||
|
|
|
|||
|
|
falseVal := false
|
|||
|
|
zero := 0
|
|||
|
|
apiReq := models.APIRequest{
|
|||
|
|
Model: cfg.Recognizer.Model,
|
|||
|
|
MaxTokens: 4000,
|
|||
|
|
EnableThinking: &falseVal,
|
|||
|
|
ThinkingBudget: &zero,
|
|||
|
|
Messages: []models.Message{
|
|||
|
|
{
|
|||
|
|
Role: "user",
|
|||
|
|
Content: []models.ContentItem{
|
|||
|
|
{Type: "text", Text: prompt},
|
|||
|
|
{
|
|||
|
|
Type: "image_url",
|
|||
|
|
ImageURL: &models.ImageURL{
|
|||
|
|
URL: "data:image/png;base64," + imageBase64,
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
jsonData, err := json.Marshal(apiReq)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, apperr.New("marshal_request", err, "")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
req, err := http.NewRequestWithContext(ctx, "POST", cfg.Recognizer.APIURL, bytes.NewBuffer(jsonData))
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, apperr.New("create_request", err, "")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
req.Header = http.Header{
|
|||
|
|
"Content-Type": []string{"application/json"},
|
|||
|
|
"Authorization": []string{"Bearer " + cfg.Recognizer.APIKey},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
slog.Info("调用识别API")
|
|||
|
|
resp, err := r.client.Do(req)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, apperr.New("api_call", err, "")
|
|||
|
|
}
|
|||
|
|
defer resp.Body.Close()
|
|||
|
|
|
|||
|
|
if resp.StatusCode != 200 {
|
|||
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|||
|
|
return nil, fmt.Errorf("API调用失败: %d, %s", resp.StatusCode, string(respBody))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
var apiResp models.APIResponse
|
|||
|
|
if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
|
|||
|
|
return nil, apperr.New("decode_response", err, "")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if len(apiResp.Choices) == 0 {
|
|||
|
|
return nil, apperr.New("empty_response", apperr.ErrRecognitionFailed, "API返回为空")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
resultText := extractContent(apiResp.Choices[0].Message.Content)
|
|||
|
|
slog.Debug("API响应接收完成", "response_length", len(resultText))
|
|||
|
|
|
|||
|
|
jsonMatch := extractJSON(resultText)
|
|||
|
|
if jsonMatch == "" {
|
|||
|
|
return nil, apperr.New("extract_json", apperr.ErrRecognitionFailed,
|
|||
|
|
fmt.Sprintf("未找到JSON数据，原始响应: %s", resultText))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
var courses []models.Course
|
|||
|
|
if err := json.Unmarshal([]byte(jsonMatch), &courses); err != nil {
|
|||
|
|
return nil, apperr.New("parse_json", err, fmt.Sprintf("数据: %s", jsonMatch))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
slog.Info("识别完成", "course_count", len(courses))
|
|||
|
|
return courses, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func extractContent(content any) string {
|
|||
|
|
switch c := content.(type) {
|
|||
|
|
case string:
|
|||
|
|
return c
|
|||
|
|
case []any:
|
|||
|
|
for _, item := range c {
|
|||
|
|
if itemMap, ok := item.(map[string]any); ok {
|
|||
|
|
if text, ok := itemMap["text"].(string); ok {
|
|||
|
|
return text
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func extractJSON(text string) string {
|
|||
|
|
if match := jsonBlockRegex.FindStringSubmatch(text); len(match) > 1 {
|
|||
|
|
slog.Debug("从json代码块中提取JSON")
|
|||
|
|
return match[1]
|
|||
|
|
}
|
|||
|
|
if match := codeBlockRegex.FindStringSubmatch(text); len(match) > 1 {
|
|||
|
|
slog.Debug("从代码块中提取JSON")
|
|||
|
|
return match[1]
|
|||
|
|
}
|
|||
|
|
if match := jsonArrayRegex.FindString(text); match != "" {
|
|||
|
|
slog.Debug("从文本中提取JSON")
|
|||
|
|
return match
|
|||
|
|
}
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func RecognizeScheduleFromFile(imagePath string) ([]models.Course, error) {
|
|||
|
|
return NewRecognizerService().RecognizeFromFile(context.Background(), imagePath)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func RecognizeScheduleFromBytes(imageData []byte) ([]models.Course, error) {
|
|||
|
|
return NewRecognizerService().RecognizeFromBytes(context.Background(), imageData)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func RecognizeScheduleFromBase64(imageBase64 string) ([]models.Course, error) {
|
|||
|
|
return NewRecognizerService().RecognizeFromBase64(context.Background(), imageBase64)
|
|||
|
|
}
|