Agent-System/internal/mineru/agent.go

package mineru

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

// parseAgentFile uses MinerU Agent lightweight API (no token; IP rate-limited).
// Limits: ≤10MB, ≤20 pages. Output is markdown only (CDN URL).
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"}
	}
	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,
	}
	raw, err := c.postJSONNoAuth(ctx, c.Endpoint+"/api/v1/agent/parse/file", body)
	if err != nil {
		return nil, err
	}
	var env apiEnvelope
	if err := json.Unmarshal(raw, &env); err != nil {
		return nil, fmt.Errorf("mineru agent: decode apply: %w", err)
	}
	if !codeOK(env.Code) {
		return nil, 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 nil, fmt.Errorf("mineru agent: decode apply data: %w", err)
	}
	if dataOut.TaskID == "" || dataOut.FileURL == "" {
		return nil, fmt.Errorf("mineru agent: empty task_id or file_url")
	}
	if err := c.putFile(ctx, dataOut.FileURL, data); err != nil {
		return nil, err
	}
	mdURL, err := c.pollAgent(ctx, dataOut.TaskID)
	if err != nil {
		return nil, err
	}
	mdBytes, err := c.download(ctx, mdURL)
	if err != nil {
		return nil, err
	}
	md := EnsureUTF8(strings.TrimSpace(string(mdBytes)))
	if md == "" {
		return nil, fmt.Errorf("mineru agent: empty markdown")
	}
	pages := countPageAnchors(md)
	if pages == 0 {
		// Agent markdown often lacks page markers; keep as single logical page.
		md = "<!-- page 1 -->\n" + md
		pages = 1
	}
	return &ParseResult{
		BatchID:    dataOut.TaskID,
		Markdown:   md,
		TotalPages: pages,
	}, 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")
}