llmgw/internal/kb/extract/helpers.go

package extract

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

// Helpers are programs on this machine that read what Go here cannot. Any
// may be "", and the formats that need it are then skipped with a reason.
type Helpers struct {
	// Office is LibreOffice's soffice: the binary Office formats and others.
	Office string
	// Tesseract reads images and scanned pages.
	Tesseract string
	// PdfToText and PdfToPPM are Poppler's: a second PDF reader, and pages
	// rendered for Tesseract.
	PdfToText string
	PdfToPPM  string
}

// helperTimeout bounds one helper run; a conversion that hangs is a file
// skipped, not a sync stopped.
const helperTimeout = 3 * time.Minute

// FindHelpers looks for the helpers on PATH and where their installers put
// them on each system.
func FindHelpers() Helpers {
	return Helpers{
		Office:    find("soffice", officePlaces()...),
		Tesseract: find("tesseract", tesseractPlaces()...),
		PdfToText: find("pdftotext"),
		PdfToPPM:  find("pdftoppm"),
	}
}

// Describe says which helpers were found, and what each missing one would
// read.
func (h Helpers) Describe() []string {
	line := func(have, name, reads string) string {
		if have != "" {
			return fmt.Sprintf("%-11s %s", name, have)
		}
		return fmt.Sprintf("%-11s not found — %s are skipped", name, reads)
	}
	return []string{
		line(h.Office, "LibreOffice", "Word 97–2003, Excel 97–2003, PowerPoint 97–2003, Pages, Numbers, Keynote and similar files"),
		line(h.Tesseract, "Tesseract", "images and scanned PDFs"),
		line(h.PdfToText, "pdftotext", "PDFs the built-in reader cannot parse"),
		line(h.PdfToPPM, "pdftoppm", "scanned PDFs (Tesseract needs their pages as images)"),
	}
}

func find(name string, places ...string) string {
	exe := name
	if runtime.GOOS == "windows" {
		exe += ".exe"
	}
	if p, err := exec.LookPath(exe); err == nil {
		return p
	}
	if runtime.GOOS != "windows" {
		if p, err := exec.LookPath(name); err == nil {
			return p
		}
	}
	for _, p := range places {
		if matches, _ := filepath.Glob(p); len(matches) > 0 {
			sort.Strings(matches)
			return matches[len(matches)-1]
		}
	}
	return ""
}

func officePlaces() []string {
	switch runtime.GOOS {
	case "windows":
		return []string{
			`C:\Program Files\LibreOffice*\program\soffice.exe`,
			`C:\Program Files (x86)\LibreOffice*\program\soffice.exe`,
		}
	case "darwin":
		return []string{"/Applications/LibreOffice.app/Contents/MacOS/soffice"}
	}
	return []string{"/usr/bin/libreoffice", "/usr/lib/libreoffice/program/soffice", "/opt/libreoffice*/program/soffice",
		"/snap/bin/libreoffice"}
}

func tesseractPlaces() []string {
	switch runtime.GOOS {
	case "windows":
		return []string{`C:\Program Files\Tesseract-OCR\tesseract.exe`, `C:\Program Files (x86)\Tesseract-OCR\tesseract.exe`}
	case "darwin":
		return []string{"/opt/homebrew/bin/tesseract", "/usr/local/bin/tesseract"}
	}
	return nil
}

// workDir is a temporary directory for one helper run, removed after.
func workDir() (string, func(), error) {
	d, err := os.MkdirTemp("", "llmgw-kb-")
	if err != nil {
		return "", nil, err
	}
	return d, func() { _ = os.RemoveAll(d) }, nil
}

// writeInput puts a document where a helper can read it, under a safe name
// that keeps its extension.
func writeInput(dir, name string, data []byte) (string, error) {
	p := filepath.Join(dir, "input."+Ext(name))
	return p, os.WriteFile(p, data, 0o600)
}

func run(ctx context.Context, dir string, bin string, args ...string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, helperTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, bin, args...)
	cmd.Dir = dir
	out, err := cmd.Output()
	if ctx.Err() != nil {
		return nil, fmt.Errorf("%s took over %s", filepath.Base(bin), helperTimeout)
	}
	if err != nil {
		var ee *exec.ExitError
		if errors.As(err, &ee) && len(ee.Stderr) > 0 {
			return nil, fmt.Errorf("%s: %v: %s", filepath.Base(bin), err, lastLine(ee.Stderr))
		}
		return nil, fmt.Errorf("%s: %w", filepath.Base(bin), err)
	}
	return out, nil
}

func lastLine(b []byte) string {
	lines := strings.Split(strings.TrimSpace(string(b)), "\n")
	return strings.TrimSpace(lines[len(lines)-1])
}

// fileURL is a path as LibreOffice takes its profile directory: a file URL,
// with a drive letter on Windows.
func fileURL(p string) string {
	p = filepath.ToSlash(p)
	if !strings.HasPrefix(p, "/") {
		p = "/" + p
	}
	return "file://" + p
}

var spreadsheets = map[string]bool{"xls": true, "xlt": true, "xlsb": true, "numbers": true, "sxc": true}
var presentations = map[string]bool{"ppt": true, "pps": true, "pot": true, "key": true, "sxi": true}

func viaOffice(ctx context.Context, name string, data []byte, o Options) (string, error) {
	if o.Helpers.Office == "" {
		return "", fmt.Errorf("%w: a .%s file needs LibreOffice, which is not installed here", ErrUnsupported, Ext(name))
	}
	return o.Helpers.officeToText(ctx, name, data, o)
}

// officeToText has LibreOffice convert a document to the matching Office
// Open XML format, then reads that: headings, sheets and slides survive,
// which a conversion to plain text would lose.
func (h Helpers) officeToText(ctx context.Context, name string, data []byte, _ Options) (string, error) {
	dir, done, err := workDir()
	if err != nil {
		return "", err
	}
	defer done()
	in, err := writeInput(dir, name, data)
	if err != nil {
		return "", err
	}
	ext := Ext(name)
	to, read := "docx", docxText
	switch {
	case spreadsheets[ext]:
		to, read = "xlsx", xlsxText
	case presentations[ext]:
		to, read = "pptx", pptxText
	}
	// A profile of its own: two conversions at once with the user's
	// profile wait on each other's lock, or fail.
	profile := fileURL(filepath.Join(dir, "profile"))
	if _, err := run(ctx, dir, h.Office, "-env:UserInstallation="+profile, "--headless", "--norestore",
		"--convert-to", to, "--outdir", dir, in); err != nil {
		return "", err
	}
	converted, err := os.ReadFile(filepath.Join(dir, "input."+to))
	if err != nil {
		return "", fmt.Errorf("LibreOffice did not convert it: %w", err)
	}
	return read(converted)
}

func viaOCR(ctx context.Context, name string, data []byte, o Options) (string, error) {
	if o.Helpers.Tesseract == "" {
		return "", fmt.Errorf("%w: an image needs Tesseract to be read, which is not installed here", ErrUnsupported)
	}
	dir, done, err := workDir()
	if err != nil {
		return "", err
	}
	defer done()
	in, err := writeInput(dir, name, data)
	if err != nil {
		return "", err
	}
	return ocr(ctx, dir, o, in)
}

func ocr(ctx context.Context, dir string, o Options, file string) (string, error) {
	langs := o.OCRLanguages
	if langs == "" {
		langs = "tur+eng"
	}
	out, err := run(ctx, dir, o.Helpers.Tesseract, file, "stdout", "-l", langs)
	if err != nil && langs != "eng" {
		// A Tesseract without the Turkish data still reads English.
		out, err = run(ctx, dir, o.Helpers.Tesseract, file, "stdout", "-l", "eng")
	}
	return string(out), err
}

// maxOCRPages bounds how many pages of one scanned PDF are read.
const maxOCRPages = 500

// pdfWithHelpers reads a PDF the built-in reader could not: with pdftotext,
// and when that finds no text either, page by page through Tesseract.
func pdfWithHelpers(ctx context.Context, name string, data []byte, o Options) (string, error) {
	h := o.Helpers
	if h.PdfToText == "" && (h.Tesseract == "" || h.PdfToPPM == "") {
		return "", fmt.Errorf("%w: no text layer the built-in reader can use; Poppler (pdftotext) or "+
			"Tesseract with pdftoppm would read it", ErrNoText)
	}
	dir, done, err := workDir()
	if err != nil {
		return "", err
	}
	defer done()
	in, err := writeInput(dir, name, data)
	if err != nil {
		return "", err
	}
	if h.PdfToText != "" {
		if out, err := run(ctx, dir, h.PdfToText, "-enc", "UTF-8", in, "-"); err == nil && strings.TrimSpace(string(out)) != "" {
			return string(out), nil
		}
	}
	if h.Tesseract == "" || h.PdfToPPM == "" {
		return "", fmt.Errorf("%w: a scanned PDF needs Tesseract and pdftoppm", ErrNoText)
	}
	if _, err := run(ctx, dir, h.PdfToPPM, "-r", "300", "-gray", "-png", "-l", fmt.Sprint(maxOCRPages), in,
		filepath.Join(dir, "page")); err != nil {
		return "", err
	}
	pages, _ := filepath.Glob(filepath.Join(dir, "page-*.png"))
	// pdftoppm pads page numbers to the width of the last one, so the
	// names sort in page order.
	sort.Strings(pages)
	var all strings.Builder
	for _, p := range pages {
		t, err := ocr(ctx, dir, o, p)
		if err != nil {
			return "", err
		}
		all.WriteString(t)
		all.WriteString("\n\n")
	}
	return all.String(), nil
}