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) }