sift/internal/pdfdoc/pdfdoc.go

// Package pdfdoc renders an extracted document for an agent: the whole text with
// page markers when it fits, otherwise a map (title, outline with page numbers,
// leading pages and how to ask for more), a page range, or grep results with page
// numbers. It works on a *pdftext.Doc with no I/O; every length is in runes, which
// is what bounds an agent's context regardless of the byte width of a code point.
package pdfdoc

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"unicode/utf8"

	"git.wntrmute.dev/kyle/sift/internal/pdftext"
)

// The errors Render returns for a malformed Pages or Grep option.
var (
	ErrBadPages = errors.New("invalid pages")
	ErrBadGrep  = errors.New("invalid grep")
)

// maxPages bounds a range request so a caller cannot ask for the whole of a
// million-page document at once; the cap below still applies on top of it.
const maxPages = 100

// Options selects how Render draws from a document.
type Options struct {
	Pages  string // "" = no range; "a-b", "a" or "a-" (1-based, inclusive)
	Grep   string // "" = no grep; an RE2 pattern
	OutCap int    // characters (runes) the content may use
}

// Rendered is what an agent sees. When Map is true Content is the map of a
// document longer than OutCap; a grep answer sets GrepHits and leaves First and
// Last at zero.
type Rendered struct {
	Content     string
	Map         bool   // Content is the map of a document longer than OutCap
	First, Last int    // pages whose text is in Content; 0, 0 for a grep answer
	GrepHits    int    // matches found (more may exist than are shown)
	Note        string // "" or how the answer was cut and how to continue
}

// Marker is the header printed before a page's text.
func Marker(page int) string {
	return fmt.Sprintf("--- page %d ---", page)
}

// Full renders every extracted page with its marker, trimmed of surrounding
// whitespace. It is the unit the map and page ranges are built from.
func Full(doc *pdftext.Doc) string {
	parts := make([]string, 0, len(doc.Text))
	for i := range doc.Text {
		parts = append(parts, Marker(i+1)+"\n\n"+strings.TrimSpace(doc.Text[i]))
	}
	return strings.Join(parts, "\n\n")
}

// Render dispatches on the option that is set: grep, a page range, or the whole
// document.
func Render(doc *pdftext.Doc, opt Options) (Rendered, error) {
	switch {
	case opt.Grep != "":
		return renderGrep(doc, opt)
	case opt.Pages != "":
		return renderPages(doc, opt)
	default:
		return renderWhole(doc, opt)
	}
}

// renderWhole shows the whole document when it fits OutCap, otherwise a map that
// fits.
func renderWhole(doc *pdftext.Doc, opt Options) (Rendered, error) {
	full := Full(doc)
	if utf8.RuneCountInString(full) <= opt.OutCap {
		return Rendered{Content: full, First: 1, Last: len(doc.Text)}, nil
	}
	return renderMap(doc, opt), nil
}

// header is the fixed top of a map: the title line and the pages/token line,
// which notes how many pages were extracted when fewer than the document has.
func header(doc *pdftext.Doc) string {
	title := doc.Title
	if title == "" {
		title = "Untitled PDF"
	}
	pages := fmt.Sprintf("%d pages · ~%d tokens", doc.Pages, tokenEstimate(doc))
	if len(doc.Text) < doc.Pages {
		pages += fmt.Sprintf(" · pages 1-%d extracted", len(doc.Text))
	}
	return "# " + title + "\n\n" + pages + "\n\n"
}

// tokenEstimate roughs a document's size in tokens as a quarter of its runes.
func tokenEstimate(doc *pdftext.Doc) int {
	return utf8.RuneCountInString(Full(doc)) / 4
}

// instruction is the fixed line that tells the agent how to ask for more.
const instruction = "This document is longer than one answer: ask for pages=\"a-b\" or grep=\"pattern\".\n\n"

// renderMap assembles the header, a cuttable outline, the instruction, and whole
// leading pages that fit in a quarter of the cap.
func renderMap(doc *pdftext.Doc, opt Options) Rendered {
	h := header(doc)
	outlineLines := outlineItems(doc.Outline)

	// The outline is the only flexible part: cut it rather than the header or the
	// instruction when the document's outline would not fit.
	pagesBudget := opt.OutCap / 4
	outlineBudget := opt.OutCap - utf8.RuneCountInString(h) - utf8.RuneCountInString(instruction) - pagesBudget
	if outlineBudget < 0 {
		outlineBudget = 0
	}
	var kept []string
	running := 0
	for _, line := range outlineLines {
		line += "\n"
		if running+utf8.RuneCountInString(line) > outlineBudget {
			break
		}
		running += utf8.RuneCountInString(line)
		kept = append(kept, strings.TrimRight(line, "\n"))
	}
	outlineSection := ""
	if len(kept) > 0 {
		outlineSection = "## Outline\n" + strings.Join(kept, "\n") + "\n\n"
	}

	pages := leadingPages(doc, opt.OutCap-utf8.RuneCountInString(h)-utf8.RuneCountInString(outlineSection)-utf8.RuneCountInString(instruction), pagesBudget)
	content := h + outlineSection + instruction + pages.text
	if utf8.RuneCountInString(content) > opt.OutCap {
		content = truncateRunes(content, opt.OutCap)
	}
	return Rendered{Content: content, Map: true, First: 1, Last: pages.last}
}

// outlineItems renders each outline item as "- <indent><title> (p. <page>)", two
// spaces of indent per level below the top.
func outlineItems(items []pdftext.OutlineItem) []string {
	lines := make([]string, 0, len(items))
	for _, it := range items {
		indent := strings.Repeat(" ", max(0, (it.Level-1)*2))
		lines = append(lines, indent+"- "+it.Title+" (p. "+strconv.Itoa(it.Page)+")")
	}
	return lines
}

// pagesResult is a run of whole pages plus the last one shown.
type pagesResult struct {
	text string
	last int
}

// leadingPages renders whole pages from page 1 into budget runes, at least one,
// cutting the first to fit if it is alone larger than the budget.
func leadingPages(doc *pdftext.Doc, budget, pagesBudget int) pagesResult {
	if budget < 0 {
		budget = 0
	}
	if pagesBudget > budget {
		pagesBudget = budget
	}
	extracted := len(doc.Text)
	var parts []string
	running := 0
	last := 0
	for p := 1; p <= extracted; p++ {
		part := Marker(p) + "\n\n" + strings.TrimSpace(doc.Text[p-1])
		if p == 1 && utf8.RuneCountInString(part) > pagesBudget {
			parts = append(parts, truncateRunes(part, pagesBudget))
			return pagesResult{text: strings.Join(parts, "\n\n"), last: 1}
		}
		if running+utf8.RuneCountInString(part) > pagesBudget {
			break
		}
		running += utf8.RuneCountInString(part)
		parts = append(parts, part)
		last = p
	}
	return pagesResult{text: strings.Join(parts, "\n\n"), last: last}
}

// renderPages renders a page range as whole pages, capped at maxPages pages and
// at OutCap runes, noting how to ask for the rest.
func renderPages(doc *pdftext.Doc, opt Options) (Rendered, error) {
	start, end, err := parsePages(opt.Pages, len(doc.Text))
	if err != nil {
		return Rendered{}, err
	}

	limit := end
	if end-start+1 > maxPages {
		limit = start + maxPages - 1
	}

	var parts []string
	running := 0
	last := start - 1
	note := ""
	for p := start; p <= limit; p++ {
		part := Marker(p) + "\n\n" + strings.TrimSpace(doc.Text[p-1])
		if p == start && utf8.RuneCountInString(part) > opt.OutCap {
			parts = append(parts, truncateRunes(part, opt.OutCap))
			note = fmt.Sprintf("page %d cut to fit the answer", p)
			last = p
			break
		}
		if running+utf8.RuneCountInString(part) > opt.OutCap {
			break
		}
		running += utf8.RuneCountInString(part)
		parts = append(parts, part)
		last = p
	}
	if last < end {
		note = fmt.Sprintf("cut at page %d; ask for pages=\"%d-%d\"", last, last+1, end)
	}
	return Rendered{Content: strings.Join(parts, "\n\n"), First: start, Last: last, Note: note}, nil
}

// parsePages parses a "a-b", "a" or "a-" spec against the number of extracted
// pages. It returns a range error wrapping ErrBadPages when the numbers are out
// of range, and a format error wrapping ErrBadPages otherwise.
func parsePages(spec string, extracted int) (start, end int, err error) {
	s := strings.TrimSpace(spec)
	start, end, endToExtracted, ok := splitSpec(s)
	if !ok {
		return 0, 0, fmt.Errorf("%w: invalid pages", ErrBadPages)
	}
	if endToExtracted {
		end = extracted
	}
	if start < 1 || end < start || start > extracted {
		return 0, 0, fmt.Errorf("%w: pages out of range (1-%d)", ErrBadPages, extracted)
	}
	if end > extracted {
		end = extracted
	}
	return start, end, nil
}

// splitSpec turns a spec into its numbers. It reports ok=false for a malformed
// spec so the caller returns a format error. A bare "a" is a single page; "a-"
// spans to the last extracted page.
func splitSpec(s string) (start, end int, endToExtracted, ok bool) {
	if !strings.Contains(s, "-") {
		n, err := strconv.Atoi(s)
		if err != nil {
			return 0, 0, false, false
		}
		return n, n, false, true
	}
	parts := strings.Split(s, "-")
	if len(parts) != 2 || parts[0] == "" {
		return 0, 0, false, false
	}
	a, err := strconv.Atoi(strings.TrimSpace(parts[0]))
	if err != nil {
		return 0, 0, false, false
	}
	if strings.TrimSpace(parts[1]) == "" {
		return a, 0, true, true
	}
	b, err := strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return 0, 0, false, false
	}
	return a, b, false, true
}

// truncateRunes returns the first n runes of s, never splitting a multi-byte one.
func truncateRunes(s string, n int) string {
	if n <= 0 {
		return ""
	}
	if utf8.RuneCountInString(s) <= n {
		return s
	}
	return string([]rune(s)[:n])
}