llmgw/internal/kb/extract/extract.go

// Package extract turns documents into plain text for knowledge bases
// (BACKLOG F2), on the machine that holds them.
//
// The text it produces is Markdown-shaped: headings as "# ", "## ", list
// items as "- ", paragraphs separated by a blank line. The chunker splits on
// exactly those, so a format that knows its headings (DOCX styles, HTML h1–h6,
// ODT outline levels, slides, sheets) says so, and its chunks carry them.
//
// Built-in formats are pure Go, so the client builds for every system without
// a C toolchain. What Go cannot read — Word 97 files, scanned pages — goes to
// a helper when one is on the machine (LibreOffice, Tesseract, Poppler), and
// is otherwise skipped with a sentence saying which helper would read it.
package extract

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"unicode/utf8"
)

// ErrUnsupported is a file this cannot read, with the reason wrapped around
// it ("a .doc needs LibreOffice").
var ErrUnsupported = errors.New("unsupported")

// ErrNoText is a file that was read and held no text: a scanned PDF without
// an OCR helper, an empty sheet.
var ErrNoText = errors.New("no text in it")

// Part is the text of one document. A file is one part; an archive or a mail
// with attachments is several, each named by its path inside, after the
// outer file's name and "!/".
type Part struct {
	// Name is "" for the file itself, or the path inside it.
	Name string
	Text string
	// Err is why this part was skipped, when it was; Text is then empty.
	Err error
}

// Options are the limits and helpers extraction uses.
type Options struct {
	Helpers Helpers
	// MaxBytes is the largest file, or member of an archive, read; 0 is
	// 256 MiB.
	MaxBytes int64
	// OCRLanguages are Tesseract's languages; "" is "tur+eng".
	OCRLanguages string
}

func (o Options) maxBytes() int64 {
	if o.MaxBytes > 0 {
		return o.MaxBytes
	}
	return 256 << 20
}

// File extracts a file on disk.
func File(ctx context.Context, name string, o Options) ([]Part, error) {
	st, err := os.Stat(name)
	if err != nil {
		return nil, err
	}
	if st.Size() > o.maxBytes() {
		return nil, fmt.Errorf("%w: %d MiB is over the %d MiB limit", ErrUnsupported, st.Size()>>20, o.maxBytes()>>20)
	}
	data, err := os.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return Bytes(ctx, filepath.Base(name), data, o)
}

// Bytes extracts a document held in memory; name gives its format.
func Bytes(ctx context.Context, name string, data []byte, o Options) ([]Part, error) {
	return extractDepth(ctx, name, data, o, 0)
}

// maxDepth is how many archives deep members are followed: a zip in a mail
// in a zip, and no further.
const maxDepth = 3

func extractDepth(ctx context.Context, name string, data []byte, o Options, depth int) ([]Part, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	ext := Ext(name)
	if f, ok := containers[ext]; ok {
		if depth >= maxDepth {
			return nil, fmt.Errorf("%w: archives nested more than %d deep", ErrUnsupported, maxDepth)
		}
		return f(ctx, name, data, o, depth)
	}
	text, err := one(ctx, name, ext, data, o)
	if err != nil {
		return nil, err
	}
	text = Clean(text)
	if strings.TrimSpace(text) == "" {
		return nil, ErrNoText
	}
	return []Part{{Text: text}}, nil
}

// one extracts a single document.
func one(ctx context.Context, name, ext string, data []byte, o Options) (string, error) {
	if f, ok := builtin[ext]; ok {
		text, err := f(data)
		if err == nil && strings.TrimSpace(text) != "" {
			return text, nil
		}
		// A PDF the Go reader cannot parse, or one with no text layer, is
		// what the helpers are for.
		if alt, ok := fallback[ext]; ok {
			if t, aerr := alt(ctx, name, data, o); aerr == nil {
				return t, nil
			} else if err == nil {
				err = aerr
			}
		}
		if err == nil {
			err = ErrNoText
		}
		return "", err
	}
	if alt, ok := fallback[ext]; ok {
		return alt(ctx, name, data, o)
	}
	if looksLikeText(data) {
		return decodeText(data), nil
	}
	if o.Helpers.Office != "" {
		return o.Helpers.officeToText(ctx, name, data, o)
	}
	return "", fmt.Errorf("%w: no reader for %q files (LibreOffice, if installed, may read it)", ErrUnsupported, ext)
}

// Ext is a name's extension, lower case and without the dot.
func Ext(name string) string {
	return strings.TrimPrefix(strings.ToLower(path.Ext(strings.ReplaceAll(name, `\`, "/"))), ".")
}

// Supported lists the extensions read without a helper, and those read with
// one, for --help and the docs.
func Supported() (builtIn, withHelper []string) {
	for e := range builtin {
		builtIn = append(builtIn, e)
	}
	for e := range containers {
		builtIn = append(builtIn, e)
	}
	for e := range textExts {
		builtIn = append(builtIn, e)
	}
	for e := range fallback {
		if _, ok := builtin[e]; !ok {
			withHelper = append(withHelper, e)
		}
	}
	sort.Strings(builtIn)
	sort.Strings(withHelper)
	return builtIn, withHelper
}

// Known reports whether a name's extension is one this reads, with or without
// a helper, or one of the text formats.
func Known(name string) bool {
	e := Ext(name)
	_, a := builtin[e]
	_, b := containers[e]
	_, c := fallback[e]
	return a || b || c || textExts[e]
}

var (
	builtin    map[string]func([]byte) (string, error)
	containers map[string]func(context.Context, string, []byte, Options, int) ([]Part, error)
	fallback   map[string]func(context.Context, string, []byte, Options) (string, error)
)

func init() {
	builtin = map[string]func([]byte) (string, error){
		"pdf":  pdfText,
		"html": htmlText, "htm": htmlText, "xhtml": htmlText, "mht": mhtmlText, "mhtml": mhtmlText,
		"docx": docxText, "docm": docxText, "dotx": docxText, "dotm": docxText,
		"xlsx": xlsxText, "xlsm": xlsxText, "xltx": xlsxText, "xltm": xlsxText,
		"pptx": pptxText, "pptm": pptxText, "ppsx": pptxText, "ppsm": pptxText, "potx": pptxText,
		"odt": odfText, "ott": odfText, "ods": odfText, "ots": odfText, "odp": odfText, "otp": odfText,
		"fodt": flatODFText, "fods": flatODFText, "fodp": flatODFText,
		"rtf":   rtfText,
		"epub":  epubText,
		"ipynb": notebookText,
		"srt":   subtitleText, "vtt": subtitleText,
	}
	for e := range textExts {
		builtin[e] = func(b []byte) (string, error) { return decodeText(b), nil }
	}
	containers = map[string]func(context.Context, string, []byte, Options, int) ([]Part, error){
		"zip": zipParts,
		"eml": mailParts,
	}
	fallback = map[string]func(context.Context, string, []byte, Options) (string, error){
		"pdf": pdfWithHelpers,
		"rtf": viaOffice,
	}
	for _, e := range officeOnly {
		fallback[e] = viaOffice
	}
	for _, e := range imageExts {
		fallback[e] = viaOCR
	}
}

// textExts are formats that are text already: read as they are, decoded.
var textExts = map[string]bool{
	"txt": true, "text": true, "md": true, "markdown": true, "mdx": true, "rst": true, "adoc": true,
	"asciidoc": true, "org": true, "tex": true, "log": true, "csv": true, "tsv": true, "json": true,
	"jsonl": true, "ndjson": true, "yaml": true, "yml": true, "toml": true, "ini": true, "cfg": true,
	"conf": true, "properties": true, "env": true, "xml": true, "svg": true, "sql": true,
	"go": true, "py": true, "js": true, "mjs": true, "cjs": true, "ts": true, "tsx": true, "jsx": true,
	"java": true, "kt": true, "kts": true, "scala": true, "c": true, "h": true, "cc": true, "cpp": true,
	"cxx": true, "hpp": true, "hh": true, "cs": true, "fs": true, "vb": true, "rs": true, "swift": true,
	"m": true, "mm": true, "rb": true, "php": true, "pl": true, "pm": true, "lua": true, "r": true,
	"jl": true, "dart": true, "ex": true, "exs": true, "erl": true, "hs": true, "clj": true,
	"sh": true, "bash": true, "zsh": true, "fish": true, "ps1": true, "psm1": true, "bat": true,
	"cmd": true, "groovy": true, "gradle": true, "tf": true, "hcl": true, "proto": true,
	"graphql": true, "gql": true, "css": true, "scss": true, "sass": true, "less": true,
	"vue": true, "svelte": true, "dockerfile": true, "makefile": true, "mk": true, "cmake": true,
	"bib": true, "nfo": true, "diff": true, "patch": true,
}

// officeOnly are formats only LibreOffice reads here: the binary Office
// formats before 2007, and others it opens.
var officeOnly = []string{
	"doc", "dot", "xls", "xlt", "xlsb", "ppt", "pps", "pot", "wpd", "wps", "wri",
	"pages", "numbers", "key", "sxw", "sxc", "sxi", "hwp", "pub", "vsd", "vsdx",
}

// imageExts are read by OCR.
var imageExts = []string{"png", "jpg", "jpeg", "tif", "tiff", "bmp", "gif", "webp", "jp2", "pnm"}

// looksLikeText reports whether bytes are text: no NUL in the first 8 KiB
// (UTF-16 has its BOM checked first), and mostly printable.
func looksLikeText(b []byte) bool {
	if hasUTF16BOM(b) {
		return true
	}
	head := b[:min(len(b), 8192)]
	if bytes.IndexByte(head, 0) >= 0 {
		return false
	}
	if utf8.Valid(head) {
		return true
	}
	// A single-byte code page: allow it when control characters are rare.
	ctl := 0
	for _, c := range head {
		if c < 0x20 && c != '\n' && c != '\r' && c != '\t' && c != '\f' {
			ctl++
		}
	}
	return ctl*100 < len(head)
}

var (
	spaceRun = regexp.MustCompile(`[ \t\x{00A0}]+\n`)
	blankRun = regexp.MustCompile(`\n{3,}`)
)

// Clean normalises extracted text: LF line ends, no trailing spaces, no more
// than one blank line in a row, no control characters but tab and newline.
func Clean(s string) string {
	s = strings.ReplaceAll(s, "\r\n", "\n")
	s = strings.ReplaceAll(s, "\r", "\n")
	s = strings.Map(func(r rune) rune {
		switch {
		case r == '\n' || r == '\t':
			return r
		case r == '\f' || r == '\v':
			return '\n'
		case r < 0x20 || r == 0x7f || r == utf8.RuneError || r == '\uFEFF':
			return -1
		}
		return r
	}, s)
	s = spaceRun.ReplaceAllString(s, "\n")
	s = blankRun.ReplaceAllString(s, "\n\n")
	return strings.TrimSpace(s) + "\n"
}