From 32cda4486cbaa034cec3e4c8166d183b63f7bc46 Mon Sep 17 00:00:00 2001 From: weichen Date: Mon, 21 Sep 2026 08:29:10 +0000 Subject: [PATCH] feat(mineru): add MinerU precision parse client Implement signed batch upload, poll, zip extract, and real page anchor rebuild from content_list.json. Co-authored-by: Cursor --- internal/mineru/client.go | 426 +++++++++++++++++++++++++++++++++ internal/mineru/client_test.go | 137 +++++++++++ internal/mineru/extract.go | 153 ++++++++++++ internal/mineru/live_test.go | 52 ++++ 4 files changed, 768 insertions(+) create mode 100644 internal/mineru/client.go create mode 100644 internal/mineru/client_test.go create mode 100644 internal/mineru/extract.go create mode 100644 internal/mineru/live_test.go diff --git a/internal/mineru/client.go b/internal/mineru/client.go new file mode 100644 index 00000000..60b3bac3 --- /dev/null +++ b/internal/mineru/client.go @@ -0,0 +1,426 @@ +package mineru + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "path/filepath" + "strings" + "sync" + "time" +) + +const ( + defaultEndpoint = "https://mineru.net" + defaultModelVersion = "vlm" + defaultLanguage = "ch" + pollInterval = 3 * time.Second + submitRatePerMin = 45 // stay under documented 50/min limit +) + +// Client calls MinerU precision parse API (token-auth batch upload). +type Client struct { + Endpoint string + APIKey string + ModelVersion string + Language string + HTTP *http.Client + + submitMu sync.Mutex + submitLog []time.Time +} + +// Config wires Client from application settings. +type Config struct { + Endpoint string + APIKey string + ModelVersion string + Language string + HTTPTimeout time.Duration +} + +// NewClient builds a MinerU client. Empty APIKey means Disabled. +func NewClient(cfg Config) *Client { + endpoint := strings.TrimRight(strings.TrimSpace(cfg.Endpoint), "/") + if endpoint == "" { + endpoint = defaultEndpoint + } + model := strings.TrimSpace(cfg.ModelVersion) + if model == "" { + model = defaultModelVersion + } + lang := strings.TrimSpace(cfg.Language) + if lang == "" { + lang = defaultLanguage + } + timeout := cfg.HTTPTimeout + if timeout <= 0 { + timeout = 120 * time.Second + } + return &Client{ + Endpoint: endpoint, + APIKey: strings.TrimSpace(cfg.APIKey), + ModelVersion: model, + Language: lang, + HTTP: &http.Client{Timeout: timeout}, + } +} + +// Enabled reports whether the client has credentials. +func (c *Client) Enabled() bool { + return c != nil && c.APIKey != "" +} + +// ParseResult is the extracted markdown and optional content_list JSON. +type ParseResult struct { + BatchID string + Markdown string + ContentList json.RawMessage + TotalPages int +} + +// ParseFile uploads local bytes via signed URL batch API and waits for markdown. +func (c *Client) ParseFile(ctx context.Context, fileName string, data []byte) (*ParseResult, error) { + if !c.Enabled() { + return nil, fmt.Errorf("mineru: api key not configured") + } + if len(data) == 0 { + return nil, fmt.Errorf("mineru: empty file") + } + name := filepath.Base(fileName) + if name == "" || name == "." { + name = "upload.bin" + } + + if err := c.waitSubmitSlot(ctx); err != nil { + return nil, err + } + + batchID, uploadURL, err := c.applyUploadURL(ctx, name) + if err != nil { + return nil, err + } + if err := c.putFile(ctx, uploadURL, data); err != nil { + return nil, err + } + zipURL, err := c.pollBatch(ctx, batchID, name) + if err != nil { + return nil, err + } + zipBytes, err := c.download(ctx, zipURL) + if err != nil { + return nil, err + } + md, contentList, pages, err := ExtractFromZip(zipBytes) + if err != nil { + return nil, err + } + return &ParseResult{ + BatchID: batchID, + Markdown: md, + ContentList: contentList, + TotalPages: pages, + }, nil +} + +func (c *Client) waitSubmitSlot(ctx context.Context) error { + c.submitMu.Lock() + defer c.submitMu.Unlock() + for { + now := time.Now() + cutoff := now.Add(-time.Minute) + kept := c.submitLog[:0] + for _, t := range c.submitLog { + if t.After(cutoff) { + kept = append(kept, t) + } + } + c.submitLog = kept + if len(c.submitLog) < submitRatePerMin { + c.submitLog = append(c.submitLog, now) + return nil + } + wait := time.Minute - now.Sub(c.submitLog[0]) + 50*time.Millisecond + c.submitMu.Unlock() + timer := time.NewTimer(wait) + select { + case <-ctx.Done(): + timer.Stop() + c.submitMu.Lock() + return ctx.Err() + case <-timer.C: + } + c.submitMu.Lock() + } +} + +type apiEnvelope struct { + Code any `json:"code"` + Msg string `json:"msg"` + TraceID string `json:"trace_id"` + Data json.RawMessage `json:"data"` +} + +func (c *Client) applyUploadURL(ctx context.Context, fileName string) (batchID, uploadURL string, err error) { + body := map[string]any{ + "files": []map[string]any{ + {"name": fileName}, + }, + "model_version": c.ModelVersion, + "language": c.Language, + "enable_table": true, + "enable_formula": true, + } + raw, err := c.postJSON(ctx, c.Endpoint+"/api/v4/file-urls/batch", body) + if err != nil { + return "", "", err + } + var env apiEnvelope + if err := json.Unmarshal(raw, &env); err != nil { + return "", "", fmt.Errorf("mineru: decode apply url: %w", err) + } + if !codeOK(env.Code) { + return "", "", classifyAPIError(env.Code, env.Msg) + } + var data struct { + BatchID string `json:"batch_id"` + FileURLs []string `json:"file_urls"` + } + if err := json.Unmarshal(env.Data, &data); err != nil { + return "", "", fmt.Errorf("mineru: decode apply data: %w", err) + } + if data.BatchID == "" || len(data.FileURLs) == 0 { + return "", "", fmt.Errorf("mineru: empty batch_id or file_urls") + } + return data.BatchID, data.FileURLs[0], nil +} + +func (c *Client) putFile(ctx context.Context, uploadURL string, data []byte) error { + req, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, bytes.NewReader(data)) + if err != nil { + return err + } + // MinerU docs: do not set Content-Type on signed PUT. + resp, err := c.HTTP.Do(req) + if err != nil { + return fmt.Errorf("mineru: upload put: %w", err) + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return fmt.Errorf("mineru: upload put http %s", resp.Status) + } + return nil +} + +func (c *Client) pollBatch(ctx context.Context, batchID, fileName string) (zipURL string, err error) { + url := fmt.Sprintf("%s/api/v4/extract-results/batch/%s", c.Endpoint, batchID) + ticker := time.NewTicker(pollInterval) + defer ticker.Stop() + for { + raw, err := c.getJSON(ctx, url) + if err != nil { + return "", err + } + var env apiEnvelope + if err := json.Unmarshal(raw, &env); err != nil { + return "", fmt.Errorf("mineru: decode poll: %w", err) + } + if !codeOK(env.Code) { + return "", classifyAPIError(env.Code, env.Msg) + } + var data struct { + BatchID string `json:"batch_id"` + ExtractResult []struct { + FileName string `json:"file_name"` + State string `json:"state"` + FullZipURL string `json:"full_zip_url"` + ErrMsg string `json:"err_msg"` + } `json:"extract_result"` + } + if err := json.Unmarshal(env.Data, &data); err != nil { + return "", fmt.Errorf("mineru: decode poll data: %w", err) + } + item := pickResult(data.ExtractResult, fileName) + if item == nil { + // still waiting for file to appear + } else { + switch item.State { + case "done": + if item.FullZipURL == "" { + return "", fmt.Errorf("mineru: done without zip url") + } + return item.FullZipURL, nil + case "failed": + return "", fmt.Errorf("mineru: parse failed: %s", strings.TrimSpace(item.ErrMsg)) + } + } + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-ticker.C: + } + } +} + +func pickResult(items []struct { + FileName string `json:"file_name"` + State string `json:"state"` + FullZipURL string `json:"full_zip_url"` + ErrMsg string `json:"err_msg"` +}, fileName string) *struct { + FileName string `json:"file_name"` + State string `json:"state"` + FullZipURL string `json:"full_zip_url"` + ErrMsg string `json:"err_msg"` +} { + if len(items) == 0 { + return nil + } + base := strings.ToLower(filepath.Base(fileName)) + for i := range items { + if strings.ToLower(filepath.Base(items[i].FileName)) == base { + return &items[i] + } + } + return &items[0] +} + +func (c *Client) download(ctx context.Context, url string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + resp, err := c.HTTP.Do(req) + if err != nil { + return nil, fmt.Errorf("mineru: download zip: %w", err) + } + defer resp.Body.Close() + raw, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20)) + if err != nil { + return nil, err + } + if resp.StatusCode >= 300 { + return nil, fmt.Errorf("mineru: download zip http %s", resp.Status) + } + return raw, nil +} + +func (c *Client) postJSON(ctx context.Context, url string, body any) ([]byte, error) { + b, _ := json.Marshal(body) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.APIKey) + resp, err := c.HTTP.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + if resp.StatusCode == http.StatusTooManyRequests { + return nil, &RetryableError{Msg: "mineru: rate limited (429)"} + } + if resp.StatusCode >= 300 { + return nil, fmt.Errorf("mineru: http %s: %s", resp.Status, truncate(string(raw), 300)) + } + return raw, nil +} + +func (c *Client) getJSON(ctx context.Context, url string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+c.APIKey) + resp, err := c.HTTP.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + if resp.StatusCode == http.StatusTooManyRequests { + return nil, &RetryableError{Msg: "mineru: rate limited (429)"} + } + if resp.StatusCode >= 300 { + return nil, fmt.Errorf("mineru: http %s: %s", resp.Status, truncate(string(raw), 300)) + } + return raw, nil +} + +func codeOK(code any) bool { + switch v := code.(type) { + case float64: + return v == 0 + case int: + return v == 0 + case string: + return v == "0" || strings.EqualFold(v, "ok") + case nil: + return true + default: + return fmt.Sprint(v) == "0" + } +} + +// RetryableError marks transient MinerU failures (queue full / rate limit). +type RetryableError struct { + Msg string +} + +func (e *RetryableError) Error() string { return e.Msg } + +// PermanentError marks non-retryable failures (auth / quota / bad file). +type PermanentError struct { + Code string + Msg string +} + +func (e *PermanentError) Error() string { + if e.Code == "" { + return e.Msg + } + return fmt.Sprintf("%s: %s", e.Code, e.Msg) +} + +func classifyAPIError(code any, msg string) error { + codeStr := strings.TrimSpace(fmt.Sprint(code)) + msg = strings.TrimSpace(msg) + switch codeStr { + case "A0202", "A0211": + return &PermanentError{Code: codeStr, Msg: msg} + case "-60018", "-60019": + return &PermanentError{Code: codeStr, Msg: msg} + case "-60009": + return &RetryableError{Msg: fmt.Sprintf("%s: %s", codeStr, msg)} + default: + if strings.Contains(strings.ToLower(msg), "queue") { + return &RetryableError{Msg: msg} + } + return &PermanentError{Code: codeStr, Msg: msg} + } +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} + +// SupportsExt reports whether MinerU precision API accepts the file type. +func SupportsExt(fileName string) bool { + ext := strings.ToLower(filepath.Ext(fileName)) + switch ext { + case ".pdf", ".doc", ".docx", ".ppt", ".pptx", ".xls", ".xlsx", + ".png", ".jpg", ".jpeg", ".jp2", ".webp", ".gif", ".bmp", ".html": + return true + default: + return false + } +} diff --git a/internal/mineru/client_test.go b/internal/mineru/client_test.go new file mode 100644 index 00000000..c46796f4 --- /dev/null +++ b/internal/mineru/client_test.go @@ -0,0 +1,137 @@ +package mineru_test + +import ( + "archive/zip" + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "chat-agent-system/internal/mineru" +) + +func TestMarkdownFromContentList_PageAnchors(t *testing.T) { + raw, _ := json.Marshal([]map[string]any{ + {"type": "text", "text": "第一页内容", "page_idx": 0, "text_level": 0}, + {"type": "title", "text": "第二章", "page_idx": 1, "text_level": 2}, + {"type": "text", "text": "第二页正文", "page_idx": 1, "text_level": 0}, + }) + md, pages, err := mineru.MarkdownFromContentList(raw) + if err != nil { + t.Fatal(err) + } + if pages != 2 { + t.Fatalf("pages=%d", pages) + } + if !strings.Contains(md, "") || !strings.Contains(md, "") { + t.Fatalf("missing page anchors:\n%s", md) + } + if !strings.Contains(md, "## 第二章") { + t.Fatalf("missing heading:\n%s", md) + } +} + +func TestExtractFromZip_PrefersContentList(t *testing.T) { + var buf bytes.Buffer + zw := zip.NewWriter(&buf) + w, _ := zw.Create("full.md") + _, _ = w.Write([]byte("# stale full\n")) + list, _ := json.Marshal([]map[string]any{ + {"type": "text", "text": "zip内容", "page_idx": 0}, + }) + w2, _ := zw.Create("demo_content_list.json") + _, _ = w2.Write(list) + _ = zw.Close() + + md, _, pages, err := mineru.ExtractFromZip(buf.Bytes()) + if err != nil { + t.Fatal(err) + } + if pages != 1 || !strings.Contains(md, "zip内容") { + t.Fatalf("md=%q pages=%d", md, pages) + } + if strings.Contains(md, "stale full") { + t.Fatal("should prefer content_list over full.md") + } +} + +func TestClient_ParseFile_MockServer(t *testing.T) { + var pollCount atomic.Int32 + var zipBytes bytes.Buffer + zw := zip.NewWriter(&zipBytes) + list, _ := json.Marshal([]map[string]any{ + {"type": "text", "text": "hello mineru", "page_idx": 0}, + }) + w, _ := zw.Create("full.md") + _, _ = w.Write([]byte("# fallback\n")) + w2, _ := zw.Create("x_content_list.json") + _, _ = w2.Write(list) + _ = zw.Close() + + var baseURL string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/upload" && r.Method == http.MethodPut: + w.WriteHeader(200) + case r.URL.Path == "/result.zip": + _, _ = w.Write(zipBytes.Bytes()) + case r.URL.Path == "/api/v4/file-urls/batch": + _ = json.NewEncoder(w).Encode(map[string]any{ + "code": 0, + "data": map[string]any{ + "batch_id": "batch-1", + "file_urls": []string{baseURL + "/upload"}, + }, + }) + case strings.HasPrefix(r.URL.Path, "/api/v4/extract-results/batch/"): + n := pollCount.Add(1) + state := "running" + zipURL := "" + if n >= 2 { + state = "done" + zipURL = baseURL + "/result.zip" + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "code": 0, + "data": map[string]any{ + "extract_result": []map[string]any{{ + "file_name": "doc.pdf", + "state": state, + "full_zip_url": zipURL, + }}, + }, + }) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + baseURL = srv.URL + + client := mineru.NewClient(mineru.Config{ + Endpoint: srv.URL, + APIKey: "test-token", + HTTPTimeout: 30 * time.Second, + }) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + res, err := client.ParseFile(ctx, "doc.pdf", []byte("%PDF-1.4 mock")) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(res.Markdown, "hello mineru") { + t.Fatalf("markdown=%q", res.Markdown) + } +} + +func TestSupportsExt(t *testing.T) { + if !mineru.SupportsExt("a.PDF") || mineru.SupportsExt("a.txt") { + t.Fatal("ext filter") + } +} diff --git a/internal/mineru/extract.go b/internal/mineru/extract.go new file mode 100644 index 00000000..794fddbf --- /dev/null +++ b/internal/mineru/extract.go @@ -0,0 +1,153 @@ +package mineru + +import ( + "archive/zip" + "bytes" + "encoding/json" + "fmt" + "io" + "path/filepath" + "strings" + "unicode/utf8" +) + +// contentListItem is a subset of MinerU content_list.json entries. +type contentListItem struct { + Type string `json:"type"` + Text string `json:"text"` + PageIdx int `json:"page_idx"` + TextLevel int `json:"text_level"` + TableBody string `json:"table_body"` + ImgCaption json.RawMessage `json:"img_caption"` + TableCaption json.RawMessage `json:"table_caption"` +} + +// ExtractFromZip reads full.md and content_list.json from a MinerU result zip. +// Prefer rebuilding markdown from content_list so page anchors use real page_idx. +func ExtractFromZip(zipBytes []byte) (markdown string, contentList json.RawMessage, totalPages int, err error) { + r, err := zip.NewReader(bytes.NewReader(zipBytes), int64(len(zipBytes))) + if err != nil { + return "", nil, 0, fmt.Errorf("mineru: open zip: %w", err) + } + + var fullMD []byte + var listRaw []byte + for _, f := range r.File { + base := strings.ToLower(filepath.Base(f.Name)) + rc, err := f.Open() + if err != nil { + continue + } + data, readErr := io.ReadAll(io.LimitReader(rc, 32<<20)) + _ = rc.Close() + if readErr != nil { + continue + } + switch { + case base == "full.md": + fullMD = data + case strings.HasSuffix(base, "_content_list.json") || base == "content_list.json": + listRaw = data + } + } + + if len(listRaw) > 0 { + md, pages, buildErr := MarkdownFromContentList(listRaw) + if buildErr == nil && strings.TrimSpace(md) != "" { + return md, json.RawMessage(listRaw), pages, nil + } + } + if len(fullMD) == 0 { + return "", nil, 0, fmt.Errorf("mineru: zip missing full.md and usable content_list") + } + md := strings.TrimSpace(string(fullMD)) + pages := countPageAnchors(md) + if pages == 0 { + pages = 1 + } + return md, json.RawMessage(listRaw), pages, nil +} + +// MarkdownFromContentList rebuilds markdown with anchors from page_idx. +// page_idx in MinerU is 0-based; anchors are 1-based for agent citations. +func MarkdownFromContentList(raw []byte) (string, int, error) { + var items []contentListItem + if err := json.Unmarshal(raw, &items); err != nil { + return "", 0, fmt.Errorf("mineru: decode content_list: %w", err) + } + if len(items) == 0 { + return "", 0, fmt.Errorf("mineru: empty content_list") + } + + var sb strings.Builder + currentPage := -1 + maxPage := 0 + for _, it := range items { + page := it.PageIdx + 1 + if page < 1 { + page = 1 + } + if page > maxPage { + maxPage = page + } + if page != currentPage { + if sb.Len() > 0 { + sb.WriteString("\n") + } + sb.WriteString(fmt.Sprintf("\n", page)) + currentPage = page + } + text := strings.TrimSpace(it.Text) + if text == "" && it.TableBody != "" { + text = strings.TrimSpace(it.TableBody) + } + if text == "" { + continue + } + switch strings.ToLower(it.Type) { + case "title", "text": + if it.TextLevel == 1 { + sb.WriteString("# ") + } else if it.TextLevel == 2 { + sb.WriteString("## ") + } else if it.TextLevel >= 3 { + sb.WriteString("### ") + } + sb.WriteString(text) + sb.WriteString("\n\n") + case "table": + sb.WriteString(text) + sb.WriteString("\n\n") + default: + sb.WriteString(text) + sb.WriteString("\n\n") + } + } + out := strings.TrimSpace(sb.String()) + if out == "" { + return "", 0, fmt.Errorf("mineru: content_list produced empty markdown") + } + if maxPage == 0 { + maxPage = 1 + } + return out, maxPage, nil +} + +func countPageAnchors(md string) int { + n := 0 + for _, line := range strings.Split(md, "\n") { + trim := strings.TrimSpace(line) + if strings.HasPrefix(trim, "") { + n++ + } + } + return n +} + +// EnsureUTF8 returns s if valid UTF-8, otherwise strips invalid sequences. +func EnsureUTF8(s string) string { + if utf8.ValidString(s) { + return s + } + return strings.ToValidUTF8(s, "") +} diff --git a/internal/mineru/live_test.go b/internal/mineru/live_test.go new file mode 100644 index 00000000..a036fda2 --- /dev/null +++ b/internal/mineru/live_test.go @@ -0,0 +1,52 @@ +package mineru_test + +import ( + "context" + "os" + "strings" + "testing" + "time" + + "chat-agent-system/internal/mineru" +) + +// Live smoke against real MinerU API. Requires MINERU_LIVE=1 and MINERU_API_KEY. +func TestLive_ParseTinyPDF(t *testing.T) { + if os.Getenv("MINERU_LIVE") != "1" { + t.Skip("set MINERU_LIVE=1 to run") + } + key := strings.TrimSpace(os.Getenv("MINERU_API_KEY")) + if key == "" { + t.Skip("MINERU_API_KEY empty") + } + // Minimal valid-ish PDF bytes (enough for API to accept as PDF upload). + pdf := []byte("%PDF-1.4\n1 0 obj<<>>endobj\ntrailer<<>>\n%%EOF\n") + client := mineru.NewClient(mineru.Config{ + Endpoint: getenv("MINERU_ENDPOINT_URL", "https://mineru.net"), + APIKey: key, + HTTPTimeout: 3 * time.Minute, + }) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) + defer cancel() + res, err := client.ParseFile(ctx, "live-smoke.pdf", pdf) + if err != nil { + msg := err.Error() + t.Logf("parse result err: %v", err) + if strings.Contains(msg, "A0202") || strings.Contains(msg, "A0211") || strings.Contains(strings.ToLower(msg), "token") { + t.Fatalf("mineru auth failed (refresh MINERU_API_KEY): %v", err) + } + // Synthetic tiny PDF may be rejected by parse engine after successful auth/upload. + return + } + if strings.TrimSpace(res.Markdown) == "" && res.TotalPages == 0 { + t.Fatal("empty parse result") + } + t.Logf("pages=%d md_len=%d batch=%s", res.TotalPages, len(res.Markdown), res.BatchID) +} + +func getenv(k, def string) string { + if v := strings.TrimSpace(os.Getenv(k)); v != "" { + return v + } + return def +}