llmgw/internal/kb/extract/html.go

package extract

import (
	"bytes"
	"io"
	"mime"
	"mime/multipart"
	"net/mail"
	"strings"

	"golang.org/x/net/html"
	"golang.org/x/net/html/charset"
)

// out builds Markdown-shaped text: paragraphs, headings, list items, table
// rows, with whitespace collapsed inside them.
type out struct {
	b strings.Builder
	// line is the text of the line being built, whitespace collapsed.
	line   strings.Builder
	prefix string
}

// text adds running text; runs of whitespace become one space, and none is
// added where the source had none ("<b>Mer</b>haba" is one word).
func (o *out) text(s string) {
	if s == "" {
		return
	}
	r := []rune(s)
	lead, trail := isSpace(r[0]), isSpace(r[len(r)-1])
	for i, f := range strings.FieldsFunc(s, isSpace) {
		if i > 0 || lead {
			o.space()
		}
		o.line.WriteString(f)
	}
	if trail {
		o.space()
	}
}

// space adds one space unless the line is empty or already ends in one.
func (o *out) space() {
	l := o.line.String()
	if l != "" && !strings.HasSuffix(l, " ") && !strings.HasSuffix(l, "\t") {
		o.line.WriteByte(' ')
	}
}

// raw adds text as it is, for preformatted blocks.
func (o *out) raw(s string) { o.line.WriteString(s) }

// tab separates cells.
func (o *out) tab() {
	o.trimLine()
	o.line.WriteString("\t")
}

func (o *out) trimLine() {
	s := strings.TrimRight(o.line.String(), " ")
	o.line.Reset()
	o.line.WriteString(s)
}

// newline ends the line being built, if it has anything.
func (o *out) newline() {
	s := strings.TrimSpace(o.line.String())
	o.line.Reset()
	if s == "" {
		return
	}
	o.b.WriteString(o.prefix)
	o.b.WriteString(s)
	o.b.WriteByte('\n')
	o.prefix = ""
}

// para ends a paragraph: the line, and a blank line after it.
func (o *out) para() {
	o.newline()
	if o.b.Len() > 0 && !strings.HasSuffix(o.b.String(), "\n\n") {
		o.b.WriteByte('\n')
	}
}

// heading starts a heading of a level, 1 to 6.
func (o *out) heading(level int) {
	o.para()
	level = max(1, min(level, 6))
	o.prefix = strings.Repeat("#", level) + " "
}

// item starts a list item.
func (o *out) item() {
	o.newline()
	o.prefix = "- "
}

func (o *out) String() string {
	o.newline()
	return o.b.String()
}

func isSpace(r rune) bool {
	return r == ' ' || r == '\t' || r == '\n' || r == '\r' || r == '\f' || r == ' '
}

// htmlSkip are elements whose content is not text a reader sees.
var htmlSkip = map[string]bool{
	"script": true, "style": true, "noscript": true, "template": true, "svg": true,
	"head": true, "nav": true, "iframe": true, "object": true, "canvas": true, "button": true, "select": true,
}

// htmlBlock are elements that end a paragraph.
var htmlBlock = map[string]bool{
	"p": true, "div": true, "section": true, "article": true, "main": true, "header": true,
	"footer": true, "aside": true, "blockquote": true, "pre": true, "table": true, "ul": true,
	"ol": true, "dl": true, "figure": true, "figcaption": true, "address": true, "hr": true,
	"form": true, "fieldset": true, "details": true, "summary": true, "body": true, "dt": true, "dd": true,
	"caption": true,
}

// htmlText reads an HTML page, in the charset it declares.
func htmlText(b []byte) (string, error) {
	r, err := charset.NewReader(bytes.NewReader(b), "text/html")
	if err != nil {
		r = bytes.NewReader(b)
	}
	return htmlFrom(r)
}

// htmlFrom reads HTML already decoded to UTF-8.
func htmlFrom(r io.Reader) (string, error) {
	z := html.NewTokenizer(r)
	var o out
	skip, pre := 0, 0
	for {
		tt := z.Next()
		switch tt {
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return o.String(), nil
			}
			return o.String(), z.Err()
		case html.TextToken:
			if skip > 0 {
				continue
			}
			if pre > 0 {
				o.raw(string(z.Text()))
			} else {
				o.text(string(z.Text()))
			}
		case html.StartTagToken, html.SelfClosingTagToken, html.EndTagToken:
			name, _ := z.TagName()
			tag := string(name)
			end := tt == html.EndTagToken
			if htmlSkip[tag] {
				if tt == html.StartTagToken {
					skip++
				} else if end && skip > 0 {
					skip--
				}
				continue
			}
			if skip > 0 {
				continue
			}
			switch {
			case len(tag) == 2 && tag[0] == 'h' && tag[1] >= '1' && tag[1] <= '6':
				if end {
					o.para()
				} else {
					o.heading(int(tag[1] - '0'))
				}
			case tag == "title" && !end:
				// Inside head, which is skipped; a title outside it is a
				// heading.
				o.heading(1)
			case tag == "li":
				if !end {
					o.item()
				} else {
					o.newline()
				}
			case tag == "br":
				o.newline()
			case tag == "tr":
				o.newline()
			case tag == "td" || tag == "th":
				if end {
					o.tab()
				}
			case tag == "pre":
				if end {
					pre--
				} else {
					pre++
				}
				o.para()
			case htmlBlock[tag]:
				o.para()
			}
		}
	}
}

// mhtmlText reads a web page saved as one file (MHTML): a MIME message whose
// HTML part is the page.
func mhtmlText(b []byte) (string, error) {
	m, err := mail.ReadMessage(bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	var best string
	err = walkMIME(m.Header.Get("Content-Type"), m.Header.Get("Content-Transfer-Encoding"), m.Body, 0,
		func(ctype, _ string, params map[string]string, body []byte) {
			if best != "" {
				return
			}
			switch ctype {
			case "text/html":
				best, _ = htmlFrom(strings.NewReader(decodeCharset(body, params["charset"])))
			}
		})
	return best, err
}

// walkMIME calls leaf for every non-multipart part of a MIME body, decoded
// from its transfer encoding.
func walkMIME(contentType, cte string, body io.Reader, depth int,
	leaf func(ctype, filename string, params map[string]string, body []byte),
) error {
	if depth > 8 {
		return nil
	}
	ctype, params, err := mime.ParseMediaType(contentType)
	if err != nil {
		ctype, params = "text/plain", map[string]string{}
	}
	if strings.HasPrefix(ctype, "multipart/") {
		mr := multipart.NewReader(body, params["boundary"])
		for {
			p, err := mr.NextRawPart()
			if err == io.EOF {
				return nil
			}
			if err != nil {
				return err
			}
			if err := walkMIME(p.Header.Get("Content-Type"), p.Header.Get("Content-Transfer-Encoding"), p, depth+1,
				func(ct, fn string, pr map[string]string, b []byte) {
					if fn == "" {
						fn = filenameOf(p.Header.Get("Content-Disposition"))
					}
					leaf(ct, fn, pr, b)
				}); err != nil {
				return err
			}
		}
	}
	data, err := io.ReadAll(transferDecoder(cte, body))
	if err != nil {
		return err
	}
	leaf(ctype, params["name"], params, data)
	return nil
}

func filenameOf(disposition string) string {
	_, p, err := mime.ParseMediaType(disposition)
	if err != nil {
		return ""
	}
	return p["filename"]
}