506 lines
14 KiB
Go
506 lines
14 KiB
Go
package chat
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"mime/multipart"
|
||
"net/http"
|
||
"net/textproto"
|
||
"regexp"
|
||
"sort"
|
||
"strings"
|
||
"time"
|
||
"unicode/utf8"
|
||
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
const (
|
||
maxQrCodeOCRFileSize = 10 * 1024 * 1024
|
||
|
||
ocrJobRedisKeyPrefix = "ocr:job:"
|
||
ocrJobRedisTTL = 10 * time.Minute
|
||
|
||
ocrJobStatusSubmitted = "submitted"
|
||
ocrJobStatusRunning = "running"
|
||
ocrJobStatusDone = "done"
|
||
ocrJobStatusFailed = "failed"
|
||
)
|
||
|
||
var (
|
||
ErrQrCodeOCRNotConfigured = errors.New("二维码 OCR Token 未配置")
|
||
ErrQrCodeOCRInvalidFile = errors.New("无效的二维码图片")
|
||
ErrQrCodeOCRUnavailable = errors.New("二维码 OCR 服务不可用")
|
||
ErrQrCodeOCRJobNotFound = errors.New("OCR 任务不存在或已过期")
|
||
ErrQrCodeOCRRedisDisabled = errors.New("OCR 异步模式需要 Redis")
|
||
|
||
ocrHTTPClient = &http.Client{Timeout: 30 * time.Second}
|
||
)
|
||
|
||
type QrCodeOCRResult struct {
|
||
GroupName string `json:"group_name"`
|
||
Candidates []string `json:"candidates"`
|
||
RawText string `json:"raw_text"`
|
||
}
|
||
|
||
type OCRJobRecord struct {
|
||
PaddleJobID string `json:"paddle_job_id"`
|
||
Status string `json:"status"`
|
||
Result *QrCodeOCRResult `json:"result,omitempty"`
|
||
Error string `json:"error,omitempty"`
|
||
}
|
||
|
||
type paddleOCRJobResponse struct {
|
||
Data struct {
|
||
JobID string `json:"jobId"`
|
||
} `json:"data"`
|
||
Message string `json:"message"`
|
||
Error string `json:"error"`
|
||
}
|
||
|
||
type paddleOCRJobStatusResponse struct {
|
||
Data struct {
|
||
State string `json:"state"`
|
||
ErrorMsg string `json:"errorMsg"`
|
||
ResultURL struct {
|
||
JSONURL string `json:"jsonUrl"`
|
||
} `json:"resultUrl"`
|
||
} `json:"data"`
|
||
Message string `json:"message"`
|
||
Error string `json:"error"`
|
||
}
|
||
|
||
func readAndValidateOCRFile(reader io.Reader, contentType string) ([]byte, string, error) {
|
||
data, err := io.ReadAll(io.LimitReader(reader, maxQrCodeOCRFileSize+1))
|
||
if err != nil || len(data) == 0 || len(data) > maxQrCodeOCRFileSize {
|
||
return nil, "", ErrQrCodeOCRInvalidFile
|
||
}
|
||
if contentType == "" {
|
||
contentType = http.DetectContentType(data)
|
||
}
|
||
if !strings.HasPrefix(contentType, "image/") {
|
||
return nil, "", ErrQrCodeOCRInvalidFile
|
||
}
|
||
return data, contentType, nil
|
||
}
|
||
|
||
func ocrJobKey(jobID string) string {
|
||
return ocrJobRedisKeyPrefix + jobID
|
||
}
|
||
|
||
func (r *Repository) saveOCRJobRecord(ctx context.Context, jobID string, record *OCRJobRecord) error {
|
||
if r.redis == nil {
|
||
return ErrQrCodeOCRRedisDisabled
|
||
}
|
||
raw, err := json.Marshal(record)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return r.redis.Set(ctx, ocrJobKey(jobID), string(raw), ocrJobRedisTTL).Err()
|
||
}
|
||
|
||
func (r *Repository) getOCRJobRecord(ctx context.Context, jobID string) (*OCRJobRecord, error) {
|
||
if r.redis == nil {
|
||
return nil, ErrQrCodeOCRRedisDisabled
|
||
}
|
||
raw, err := r.redis.Get(ctx, ocrJobKey(jobID)).Result()
|
||
if err != nil {
|
||
return nil, ErrQrCodeOCRJobNotFound
|
||
}
|
||
var record OCRJobRecord
|
||
if err := json.Unmarshal([]byte(raw), &record); err != nil {
|
||
return nil, err
|
||
}
|
||
return &record, nil
|
||
}
|
||
|
||
// RecognizeQrCodeGroupName 同步调用 PaddleOCR 识别二维码图片中的群名（保留向后兼容）。
|
||
func (r *Repository) RecognizeQrCodeGroupName(ctx context.Context, filename, contentType string, reader io.Reader) (*QrCodeOCRResult, error) {
|
||
config, err := r.GetQrCodeOCRConfig(ctx)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if strings.TrimSpace(config.Token) == "" {
|
||
return nil, ErrQrCodeOCRNotConfigured
|
||
}
|
||
data, contentType, err := readAndValidateOCRFile(reader, contentType)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
paddleJobID, err := submitPaddleOCRJob(ctx, config, filename, contentType, data)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
jsonURL, err := waitPaddleOCRJob(ctx, config, paddleJobID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
rawText, err := fetchPaddleOCRText(ctx, jsonURL)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
candidates := parseQrCodeGroupNameCandidates(rawText)
|
||
groupName := ""
|
||
if len(candidates) > 0 {
|
||
groupName = candidates[0]
|
||
}
|
||
return &QrCodeOCRResult{
|
||
GroupName: groupName,
|
||
Candidates: candidates,
|
||
RawText: rawText,
|
||
}, nil
|
||
}
|
||
|
||
// SubmitOCRJob 提交 OCR 任务到 PaddleOCR，立即返回任务 ID 供前端轮询。
|
||
func (r *Repository) SubmitOCRJob(ctx context.Context, filename, contentType string, reader io.Reader) (string, error) {
|
||
if r.redis == nil {
|
||
return "", ErrQrCodeOCRRedisDisabled
|
||
}
|
||
config, err := r.GetQrCodeOCRConfig(ctx)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if strings.TrimSpace(config.Token) == "" {
|
||
return "", ErrQrCodeOCRNotConfigured
|
||
}
|
||
data, contentType, err := readAndValidateOCRFile(reader, contentType)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
paddleJobID, err := submitPaddleOCRJob(ctx, config, filename, contentType, data)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
jobID := uuid.NewString()
|
||
record := &OCRJobRecord{
|
||
PaddleJobID: paddleJobID,
|
||
Status: ocrJobStatusSubmitted,
|
||
}
|
||
if err := r.saveOCRJobRecord(ctx, jobID, record); err != nil {
|
||
return "", err
|
||
}
|
||
return jobID, nil
|
||
}
|
||
|
||
// PollOCRJobResult 轮询 OCR 任务状态，每次调用向 PaddleOCR 查询一次并更新缓存。
|
||
func (r *Repository) PollOCRJobResult(ctx context.Context, jobID string) (*OCRJobRecord, error) {
|
||
record, err := r.getOCRJobRecord(ctx, jobID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if record.Status == ocrJobStatusDone || record.Status == ocrJobStatusFailed {
|
||
return record, nil
|
||
}
|
||
|
||
config, err := r.GetQrCodeOCRConfig(ctx)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
statusURL := strings.TrimRight(config.JobURL, "/") + "/" + record.PaddleJobID
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, statusURL, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
req.Header.Set("Authorization", "bearer "+config.Token)
|
||
|
||
resp, err := ocrHTTPClient.Do(req)
|
||
if err != nil {
|
||
// 网络错误不更新 Redis，保留当前状态让前端重试
|
||
return record, nil
|
||
}
|
||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
|
||
resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
return record, nil
|
||
}
|
||
|
||
var statusPayload paddleOCRJobStatusResponse
|
||
if err := json.Unmarshal(respBody, &statusPayload); err != nil {
|
||
return record, nil
|
||
}
|
||
|
||
switch statusPayload.Data.State {
|
||
case "pending", "running":
|
||
record.Status = ocrJobStatusRunning
|
||
_ = r.saveOCRJobRecord(ctx, jobID, record)
|
||
return record, nil
|
||
|
||
case "done":
|
||
if statusPayload.Data.ResultURL.JSONURL == "" {
|
||
record.Status = ocrJobStatusFailed
|
||
record.Error = "未返回识别结果地址"
|
||
_ = r.saveOCRJobRecord(ctx, jobID, record)
|
||
return record, nil
|
||
}
|
||
rawText, err := fetchPaddleOCRText(ctx, statusPayload.Data.ResultURL.JSONURL)
|
||
if err != nil {
|
||
record.Status = ocrJobStatusFailed
|
||
record.Error = err.Error()
|
||
_ = r.saveOCRJobRecord(ctx, jobID, record)
|
||
return record, nil
|
||
}
|
||
candidates := parseQrCodeGroupNameCandidates(rawText)
|
||
groupName := ""
|
||
if len(candidates) > 0 {
|
||
groupName = candidates[0]
|
||
}
|
||
record.Status = ocrJobStatusDone
|
||
record.Result = &QrCodeOCRResult{
|
||
GroupName: groupName,
|
||
Candidates: candidates,
|
||
RawText: rawText,
|
||
}
|
||
_ = r.saveOCRJobRecord(ctx, jobID, record)
|
||
return record, nil
|
||
|
||
case "failed":
|
||
record.Status = ocrJobStatusFailed
|
||
if statusPayload.Data.ErrorMsg != "" {
|
||
record.Error = statusPayload.Data.ErrorMsg
|
||
} else {
|
||
record.Error = "PaddleOCR 识别失败"
|
||
}
|
||
_ = r.saveOCRJobRecord(ctx, jobID, record)
|
||
return record, nil
|
||
}
|
||
|
||
return record, nil
|
||
}
|
||
|
||
func submitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, filename, contentType string, data []byte) (string, error) {
|
||
body := &bytes.Buffer{}
|
||
writer := multipart.NewWriter(body)
|
||
_ = writer.WriteField("model", config.Model)
|
||
optionalPayload, _ := json.Marshal(map[string]bool{
|
||
"useDocOrientationClassify": false,
|
||
"useDocUnwarping": false,
|
||
"useChartRecognition": false,
|
||
})
|
||
_ = writer.WriteField("optionalPayload", string(optionalPayload))
|
||
partHeader := make(textproto.MIMEHeader)
|
||
partHeader.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, escapeQuotes(filename)))
|
||
partHeader.Set("Content-Type", contentType)
|
||
fileWriter, err := writer.CreatePart(partHeader)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
if _, err := fileWriter.Write(data); err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
if err := writer.Close(); err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, config.JobURL, body)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
req.Header.Set("Authorization", "bearer "+config.Token)
|
||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||
|
||
resp, err := ocrHTTPClient.Do(req)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
defer resp.Body.Close()
|
||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
|
||
if resp.StatusCode != http.StatusOK {
|
||
return "", fmt.Errorf("%w: status=%d body=%s", ErrQrCodeOCRUnavailable, resp.StatusCode, string(respBody))
|
||
}
|
||
var payload paddleOCRJobResponse
|
||
if err := json.Unmarshal(respBody, &payload); err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
if payload.Data.JobID == "" {
|
||
return "", fmt.Errorf("%w: 未返回 jobId", ErrQrCodeOCRUnavailable)
|
||
}
|
||
return payload.Data.JobID, nil
|
||
}
|
||
|
||
func waitPaddleOCRJob(ctx context.Context, config *QrCodeOCRConfig, jobID string) (string, error) {
|
||
for i := 0; i < 12; i++ {
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(config.JobURL, "/")+"/"+jobID, nil)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
req.Header.Set("Authorization", "bearer "+config.Token)
|
||
resp, err := ocrHTTPClient.Do(req)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
|
||
_ = resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
return "", fmt.Errorf("%w: status=%d body=%s", ErrQrCodeOCRUnavailable, resp.StatusCode, string(respBody))
|
||
}
|
||
var payload paddleOCRJobStatusResponse
|
||
if err := json.Unmarshal(respBody, &payload); err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
switch payload.Data.State {
|
||
case "done":
|
||
if payload.Data.ResultURL.JSONURL == "" {
|
||
return "", fmt.Errorf("%w: 未返回识别结果地址", ErrQrCodeOCRUnavailable)
|
||
}
|
||
return payload.Data.ResultURL.JSONURL, nil
|
||
case "failed":
|
||
if payload.Data.ErrorMsg != "" {
|
||
return "", fmt.Errorf("%w: %s", ErrQrCodeOCRUnavailable, payload.Data.ErrorMsg)
|
||
}
|
||
return "", ErrQrCodeOCRUnavailable
|
||
}
|
||
select {
|
||
case <-ctx.Done():
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, ctx.Err())
|
||
case <-time.After(2 * time.Second):
|
||
}
|
||
}
|
||
return "", fmt.Errorf("%w: 识别超时", ErrQrCodeOCRUnavailable)
|
||
}
|
||
|
||
func fetchPaddleOCRText(ctx context.Context, jsonURL string) (string, error) {
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, jsonURL, nil)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
resp, err := ocrHTTPClient.Do(req)
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
|
||
return "", fmt.Errorf("%w: status=%d body=%s", ErrQrCodeOCRUnavailable, resp.StatusCode, string(respBody))
|
||
}
|
||
body, err := io.ReadAll(io.LimitReader(resp.Body, 10*1024*1024))
|
||
if err != nil {
|
||
return "", fmt.Errorf("%w: %v", ErrQrCodeOCRUnavailable, err)
|
||
}
|
||
return collectPaddleOCRJSONLText(string(body)), nil
|
||
}
|
||
|
||
func collectPaddleOCRJSONLText(raw string) string {
|
||
var texts []string
|
||
for _, line := range strings.Split(raw, "\n") {
|
||
line = strings.TrimSpace(line)
|
||
if line == "" {
|
||
continue
|
||
}
|
||
var value any
|
||
if err := json.Unmarshal([]byte(line), &value); err != nil {
|
||
texts = append(texts, line)
|
||
continue
|
||
}
|
||
collectPaddleOCRText(value, &texts)
|
||
}
|
||
return strings.Join(texts, "\n")
|
||
}
|
||
|
||
func collectPaddleOCRText(value any, texts *[]string) {
|
||
switch v := value.(type) {
|
||
case string:
|
||
if strings.TrimSpace(v) != "" {
|
||
*texts = append(*texts, v)
|
||
}
|
||
case []any:
|
||
for _, item := range v {
|
||
collectPaddleOCRText(item, texts)
|
||
}
|
||
case map[string]any:
|
||
if markdown, ok := v["markdown"].(map[string]any); ok {
|
||
if text, ok := markdown["text"].(string); ok && strings.TrimSpace(text) != "" {
|
||
*texts = append(*texts, text)
|
||
}
|
||
}
|
||
for key, item := range v {
|
||
if key == "images" || key == "outputImages" {
|
||
continue
|
||
}
|
||
collectPaddleOCRText(item, texts)
|
||
}
|
||
}
|
||
}
|
||
|
||
func parseQrCodeGroupNameCandidates(rawText string) []string {
|
||
excluded := []*regexp.Regexp{
|
||
regexp.MustCompile(`使用.*微信`),
|
||
regexp.MustCompile(`企业微信`),
|
||
regexp.MustCompile(`扫码`),
|
||
regexp.MustCompile(`加入`),
|
||
regexp.MustCompile(`二维码`),
|
||
regexp.MustCompile(`有效`),
|
||
regexp.MustCompile(`天内`),
|
||
regexp.MustCompile(`月\d+日`),
|
||
regexp.MustCompile(`前`),
|
||
regexp.MustCompile(`^\d+$`),
|
||
}
|
||
scores := map[string]int{}
|
||
for _, line := range strings.Split(rawText, "\n") {
|
||
line = normalizeQrCodeOCRLine(line)
|
||
if line == "" || utf8.RuneCountInString(line) < 3 || utf8.RuneCountInString(line) > 64 {
|
||
continue
|
||
}
|
||
if !regexp.MustCompile(`[\p{Han}]`).MatchString(line) {
|
||
continue
|
||
}
|
||
skip := false
|
||
for _, pattern := range excluded {
|
||
if pattern.MatchString(line) {
|
||
skip = true
|
||
break
|
||
}
|
||
}
|
||
if skip {
|
||
continue
|
||
}
|
||
scores[line] = scoreQrCodeGroupNameCandidate(line)
|
||
}
|
||
candidates := make([]string, 0, len(scores))
|
||
for candidate := range scores {
|
||
candidates = append(candidates, candidate)
|
||
}
|
||
sort.SliceStable(candidates, func(i, j int) bool {
|
||
return scores[candidates[i]] > scores[candidates[j]]
|
||
})
|
||
return candidates
|
||
}
|
||
|
||
func normalizeQrCodeOCRLine(line string) string {
|
||
replacer := strings.NewReplacer(" ", "", "\t", "", "|", "", "｜", "", "#", "", "*", "", "`", "", ">", "")
|
||
line = replacer.Replace(line)
|
||
allowed := regexp.MustCompile(`[^\p{Han}a-zA-Z0-9_（）()《》\-—·]+`)
|
||
return allowed.ReplaceAllString(line, "")
|
||
}
|
||
|
||
func scoreQrCodeGroupNameCandidate(line string) int {
|
||
length := utf8.RuneCountInString(line)
|
||
score := length
|
||
if score > 20 {
|
||
score = 20
|
||
}
|
||
if strings.Contains(line, "群") {
|
||
score += 100
|
||
}
|
||
if strings.Contains(line, "-") || strings.Contains(line, "—") {
|
||
score += 30
|
||
}
|
||
if length >= 6 && length <= 24 {
|
||
score += 20
|
||
}
|
||
return score
|
||
}
|
||
|
||
func escapeQuotes(value string) string {
|
||
return strings.NewReplacer("\\", "\\\\", `"`, "\\\"").Replace(value)
|
||
}
|