package chat import ( "bytes" "context" "encoding/json" "errors" "fmt" "io" "mime/multipart" "net/http" "net/textproto" "regexp" "sort" "strings" "time" "unicode/utf8" ) const ( maxQrCodeOCRFileSize = 10 * 1024 * 1024 qrCodeOCRPollInterval = 2 * time.Second qrCodeOCRMaxPollTimes = 45 ) var ( ErrQrCodeOCRNotConfigured = errors.New("二维码 OCR Token 未配置") ErrQrCodeOCRInvalidFile = errors.New("无效的二维码图片") ErrQrCodeOCRUnavailable = errors.New("二维码 OCR 服务不可用") ) type QrCodeOCRResult struct { GroupName string `json:"group_name"` Candidates []string `json:"candidates"` RawText string `json:"raw_text"` } 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"` } // RecognizeQrCodeGroupName 调用 PaddleOCR 异步 API 识别二维码图片中的群名。 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, 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 } jobID, err := submitPaddleOCRJob(ctx, config, filename, contentType, data) if err != nil { return nil, err } jsonURL, err := waitPaddleOCRJob(ctx, config, jobID) 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 } 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 := http.DefaultClient.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 < qrCodeOCRMaxPollTimes; 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 := http.DefaultClient.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 "", ctx.Err() case <-time.After(qrCodeOCRPollInterval): } } 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 := http.DefaultClient.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) }