llmgw/internal/kb/extract/misc.go

package extract

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"path"
	"regexp"
	"strings"

	"github.com/ledongthuc/pdf"
)

// epubText reads an e-book: its chapters in reading order (the spine).
func epubText(b []byte) (string, error) {
	zr, err := openZip(b)
	if err != nil {
		return "", err
	}
	container, _ := member(zr, "META-INF/container.xml")
	if container == nil {
		return "", errors.New("not an EPUB: no META-INF/container.xml")
	}
	var opf string
	_ = walkXML(container, func(e xmlEvent) {
		if e.start && e.name == "rootfile" && opf == "" {
			opf = e.get("full-path")
		}
	})
	pkg, _ := member(zr, opf)
	if pkg == nil {
		return "", fmt.Errorf("not an EPUB: no package document %q", opf)
	}
	dir := path.Dir(opf)
	items := map[string]string{}
	var spine []string
	_ = walkXML(pkg, func(e xmlEvent) {
		switch {
		case e.start && e.name == "item":
			items[e.get("id")] = path.Clean(path.Join(dir, e.get("href")))
		case e.start && e.name == "itemref":
			spine = append(spine, e.get("idref"))
		}
	})
	var all strings.Builder
	for _, id := range spine {
		page, _ := member(zr, items[id])
		if page == nil {
			continue
		}
		t, err := htmlText(page)
		if err != nil {
			continue
		}
		all.WriteString(t)
		all.WriteString("\n\n")
	}
	return all.String(), nil
}

// notebookText reads a Jupyter notebook: Markdown cells as they are, code
// cells fenced; outputs left out, since they are mostly numbers and images.
func notebookText(b []byte) (string, error) {
	var nb struct {
		Cells []struct {
			CellType string          `json:"cell_type"`
			Source   json.RawMessage `json:"source"`
		} `json:"cells"`
		Metadata struct {
			Kernel struct {
				Language string `json:"language"`
			} `json:"kernelspec"`
		} `json:"metadata"`
	}
	if err := json.Unmarshal(b, &nb); err != nil {
		return "", err
	}
	var out strings.Builder
	for _, c := range nb.Cells {
		src := jsonText(c.Source)
		if strings.TrimSpace(src) == "" {
			continue
		}
		switch c.CellType {
		case "code":
			out.WriteString("```" + nb.Metadata.Kernel.Language + "\n" + strings.TrimRight(src, "\n") + "\n```\n\n")
		default:
			out.WriteString(strings.TrimRight(src, "\n") + "\n\n")
		}
	}
	return out.String(), nil
}

// jsonText is a notebook source: a string, or a list of lines.
func jsonText(raw json.RawMessage) string {
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return s
	}
	var lines []string
	if json.Unmarshal(raw, &lines) == nil {
		return strings.Join(lines, "")
	}
	return ""
}

var (
	subtitleTiming = regexp.MustCompile(`^\d{1,2}:\d{2}(:\d{2})?[.,]\d{3}\s+-->`)
	subtitleIndex  = regexp.MustCompile(`^\d+$`)
	subtitleTag    = regexp.MustCompile(`<[^>]+>`)
)

// subtitleText reads SRT and WebVTT: the spoken lines, without numbers,
// timings and styling.
func subtitleText(b []byte) (string, error) {
	var out strings.Builder
	last := ""
	for _, line := range strings.Split(decodeText(b), "\n") {
		line = strings.TrimSpace(line)
		switch {
		case line == "", line == "WEBVTT", strings.HasPrefix(line, "NOTE"),
			subtitleIndex.MatchString(line), subtitleTiming.MatchString(line):
			continue
		}
		line = strings.TrimSpace(subtitleTag.ReplaceAllString(line, ""))
		// Captions repeat a line across cues as it scrolls.
		if line == "" || line == last {
			continue
		}
		last = line
		out.WriteString(line + "\n")
	}
	return out.String(), nil
}

// pdfText reads a PDF's text layer, page by page, rows in reading order.
// A scanned PDF has none; pdfWithHelpers is then tried.
func pdfText(b []byte) (text string, err error) {
	// The reader panics on some malformed files rather than failing; a bad
	// PDF is one skipped file, not a stopped sync.
	defer func() {
		if r := recover(); r != nil {
			text, err = "", fmt.Errorf("the PDF could not be parsed: %v", r)
		}
	}()
	r, err := pdf.NewReader(bytes.NewReader(b), int64(len(b)))
	if err != nil {
		return "", err
	}
	var o out
	for i := 1; i <= r.NumPage(); i++ {
		p := r.Page(i)
		if p.V.IsNull() {
			continue
		}
		rows, err := p.GetTextByRow()
		if err != nil {
			continue
		}
		for _, row := range rows {
			var line strings.Builder
			for _, w := range row.Content {
				line.WriteString(w.S)
			}
			o.text(line.String())
			o.newline()
		}
		o.para()
	}
	return o.String(), nil
}