Agent-System/internal/mineru/agent.go

package mineru

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

var chunkAnchorRe = regexp.MustCompile(`(?m)^<!--\s*page\s+(\d+)\s*-->\s*$`)

// bumpPageAnchors rewrites chunk-local `<!-- page N -->` anchors to global
// numbering by adding offset (start-1) to each page number.
func BumpPageAnchors(md string, offset int) string {
	if offset <= 0 {
		return md
	}
	return chunkAnchorRe.ReplaceAllStringFunc(md, func(s string) string {
		m := chunkAnchorRe.FindStringSubmatch(s)
		if len(m) == 2 {
			var n int
			if _, err := fmt.Sscanf(m[1], "%d", &n); err == nil {
				return fmt.Sprintf("<!-- page %d -->", n+offset)
			}
		}
		return s
	})
}

// agentMaxPagesPerRequest is the documented Agent lightweight API page cap.
const agentMaxPagesPerRequest = 20

// agentMaxChunks bounds chunked parsing (20 pages × 40 = 800 pages max).
const agentMaxChunks = 40

// parseAgentFile uses MinerU Agent lightweight API (no token; IP rate-limited).
// Limits: ≤10MB, ≤20 pages per request. Output is markdown only (CDN URL).
// When the whole file exceeds the page cap (-30003), fall back to page_range
// chunking so documents without a valid precision token still parse.
func (c *Client) parseAgentFile(ctx context.Context, fileName string, data []byte) (*ParseResult, error) {
	if len(data) > 10<<20 {
		return nil, &PermanentError{Code: "-30001", Msg: "file exceeds agent lightweight 10MB limit"}
	}
	md, pages, taskID, err := c.parseAgentRange(ctx, fileName, data, "")
	if err == nil {
		if pages == 0 {
			// Agent markdown lacks page markers; keep as single logical page.
			md = "<!-- page 1 -->\n" + md
			pages = 1
		}
		return &ParseResult{BatchID: taskID, Markdown: md, TotalPages: pages, EngineMode: "agent"}, nil
	}
	if !IsPageLimitError(err) {
		return nil, err
	}

	// Whole file exceeded the 20-page cap: parse in page_range windows.
	var sb strings.Builder
	totalPages := 0
	firstTaskID := ""
	for i := 0; i < agentMaxChunks; i++ {
		start := 1 + i*agentMaxPagesPerRequest
		end := start + agentMaxPagesPerRequest - 1
		rng := fmt.Sprintf("%d-%d", start, end)
		md, _, tid, cerr := c.parseAgentRange(ctx, fileName, data, rng)
		if cerr != nil {
			if IsPageLimitError(cerr) || i > 0 {
				// Beyond document end (after at least one success) or re-limit: stop.
				break
			}
			return nil, cerr
		}
		if firstTaskID == "" {
			firstTaskID = tid
		}
		if countPageAnchors(md) > 0 {
			// Chunk-local anchors: renumber to global pages.
			sb.WriteString("\n")
			sb.WriteString(BumpPageAnchors(md, start-1))
			sb.WriteString("\n")
		} else {
			sb.WriteString(fmt.Sprintf("\n<!-- page %d -->\n", start))
			sb.WriteString(md)
			sb.WriteString("\n")
		}
		totalPages = end
	}
	out := strings.TrimSpace(sb.String())
	if out == "" {
		return nil, fmt.Errorf("mineru agent: chunked parse produced no markdown")
	}
	return &ParseResult{
		BatchID:    firstTaskID,
		Markdown:   out,
		TotalPages: totalPages,
		EngineMode: "agent_chunked",
	}, nil
}

// isPageLimitError reports the Agent API "-30003 file page count exceeds
// API limit" condition (both submit-time and poll-time shapes).
func IsPageLimitError(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "-30003") ||
		strings.Contains(msg, "page count exceeds")
}

// parseAgentRange parses fileName/data via Agent API with optional page_range
// ("" = whole file). Returns markdown, anchor count and task id.
func (c *Client) parseAgentRange(ctx context.Context, fileName string, data []byte, pageRange string) (md string, pages int, taskID string, err error) {
	name := filepath.Base(fileName)
	if name == "" || name == "." {
		name = "upload.bin"
	}

	body := map[string]any{
		"file_name":      name,
		"language":       c.Language,
		"enable_table":   true,
		"enable_formula": true,
		"is_ocr":         false,
	}
	if pageRange != "" {
		body["page_range"] = pageRange
	}
	raw, err := c.postJSONNoAuth(ctx, c.Endpoint+"/api/v1/agent/parse/file", body)
	if err != nil {
		return "", 0, "", err
	}
	var env apiEnvelope
	if err := json.Unmarshal(raw, &env); err != nil {
		return "", 0, "", fmt.Errorf("mineru agent: decode apply: %w", err)
	}
	if !codeOK(env.Code) {
		return "", 0, "", classifyAPIError(env.Code, env.Msg)
	}
	var dataOut struct {
		TaskID  string `json:"task_id"`
		FileURL string `json:"file_url"`
	}
	if err := json.Unmarshal(env.Data, &dataOut); err != nil {
		return "", 0, "", fmt.Errorf("mineru agent: decode apply data: %w", err)
	}
	if dataOut.TaskID == "" || dataOut.FileURL == "" {
		return "", 0, "", fmt.Errorf("mineru agent: empty task_id or file_url")
	}
	if err := c.putFile(ctx, dataOut.FileURL, data); err != nil {
		return "", 0, dataOut.TaskID, err
	}
	mdURL, err := c.pollAgent(ctx, dataOut.TaskID)
	if err != nil {
		return "", 0, dataOut.TaskID, err
	}
	mdBytes, err := c.download(ctx, mdURL)
	if err != nil {
		return "", 0, dataOut.TaskID, err
	}
	md = EnsureUTF8(strings.TrimSpace(string(mdBytes)))
	if md == "" {
		return "", 0, dataOut.TaskID, fmt.Errorf("mineru agent: empty markdown")
	}
	pages = countPageAnchors(md)
	return md, pages, dataOut.TaskID, nil
}

func (c *Client) pollAgent(ctx context.Context, taskID string) (markdownURL string, err error) {
	url := fmt.Sprintf("%s/api/v1/agent/parse/%s", c.Endpoint, taskID)
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()
	for {
		raw, err := c.getJSONNoAuth(ctx, url)
		if err != nil {
			return "", err
		}
		var env apiEnvelope
		if err := json.Unmarshal(raw, &env); err != nil {
			return "", fmt.Errorf("mineru agent: decode poll: %w", err)
		}
		if !codeOK(env.Code) {
			return "", classifyAPIError(env.Code, env.Msg)
		}
		var data struct {
			TaskID      string `json:"task_id"`
			State       string `json:"state"`
			MarkdownURL string `json:"markdown_url"`
			ErrMsg      string `json:"err_msg"`
			ErrCode     int    `json:"err_code"`
		}
		if err := json.Unmarshal(env.Data, &data); err != nil {
			return "", fmt.Errorf("mineru agent: decode poll data: %w", err)
		}
		switch data.State {
		case "done":
			if data.MarkdownURL == "" {
				return "", fmt.Errorf("mineru agent: done without markdown_url")
			}
			return data.MarkdownURL, nil
		case "failed":
			return "", fmt.Errorf("mineru agent: failed (%d): %s", data.ErrCode, data.ErrMsg)
		}
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-ticker.C:
		}
	}
}

func (c *Client) postJSONNoAuth(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")
	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 agent: rate limited (429)"}
	}
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("mineru agent: http %s: %s", resp.Status, truncate(string(raw), 300))
	}
	return raw, nil
}

func (c *Client) getJSONNoAuth(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, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode == http.StatusTooManyRequests {
		return nil, &RetryableError{Msg: "mineru agent: rate limited (429)"}
	}
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("mineru agent: http %s: %s", resp.Status, truncate(string(raw), 300))
	}
	return raw, nil
}

func isPrecisionAuthFailure(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "A0202") ||
		strings.Contains(msg, "A0211") ||
		strings.Contains(strings.ToLower(msg), "token") ||
		strings.Contains(msg, "401")
}