Agent-System/internal/mineru/extract.go

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 <!-- page N --> 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("<!-- page %d -->\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, "<!-- page ") && strings.HasSuffix(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, "")
}