geo/server/internal/tenant/app/knowledge_parser.go

package app

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"path/filepath"
	"regexp"
	"strings"
	"time"
	"unicode/utf8"

	pdf "github.com/ledongthuc/pdf"
	"github.com/xuri/excelize/v2"
	"golang.org/x/net/html"
)

const (
	maxKnowledgeFetchBytes = 5 << 20
)

var whitespacePattern = regexp.MustCompile(`[ \t]+`)

func extractKnowledgeTextFromURL(ctx context.Context, rawURL string) (string, error) {
	data, contentType, err := fetchKnowledgeURLSnapshot(ctx, rawURL)
	if err != nil {
		return "", err
	}
	return extractKnowledgeTextFromSnapshot(resolveSnapshotFileName(rawURL, contentType), data)
}

func fetchKnowledgeURLSnapshot(ctx context.Context, rawURL string) ([]byte, string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimSpace(rawURL), nil)
	if err != nil {
		return nil, "", fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("User-Agent", "geo-knowledge-fetcher/1.0")

	client := &http.Client{Timeout: 20 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, "", fmt.Errorf("fetch url: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= http.StatusBadRequest {
		return nil, "", fmt.Errorf("fetch url failed: status=%d", resp.StatusCode)
	}

	data, err := io.ReadAll(io.LimitReader(resp.Body, maxKnowledgeFetchBytes))
	if err != nil {
		return nil, "", fmt.Errorf("read url body: %w", err)
	}

	contentType := strings.ToLower(resp.Header.Get("Content-Type"))
	return data, contentType, nil
}

func resolveSnapshotFileName(rawURL string, contentType string) string {
	if strings.Contains(strings.ToLower(contentType), "text/plain") {
		return rawURL + ".txt"
	}
	return rawURL + ".html"
}

func extractKnowledgeTextFromSnapshot(fileName string, content []byte) (string, error) {
	ext := strings.ToLower(filepath.Ext(fileName))
	if ext != ".html" && ext != ".htm" {
		return extractKnowledgeTextFromFile(fileName, content)
	}

	doc, err := html.Parse(bytes.NewReader(content))
	if err != nil {
		return "", fmt.Errorf("parse html: %w", err)
	}

	var parts []string
	var walk func(*html.Node, bool)
	walk = func(node *html.Node, skip bool) {
		if node == nil {
			return
		}

		nextSkip := skip
		if node.Type == html.ElementNode {
			switch strings.ToLower(node.Data) {
			case "script", "style", "noscript", "svg":
				nextSkip = true
			}
		}

		if node.Type == html.TextNode && !nextSkip {
			text := strings.TrimSpace(node.Data)
			if text != "" {
				parts = append(parts, text)
			}
		}

		for child := node.FirstChild; child != nil; child = child.NextSibling {
			walk(child, nextSkip)
		}
	}
	walk(doc, false)

	return normalizeKnowledgeText(strings.Join(parts, "\n")), nil
}

func extractKnowledgeTextFromFile(fileName string, content []byte) (string, error) {
	ext := strings.ToLower(filepath.Ext(fileName))

	switch ext {
	case ".txt", ".md", ".markdown", ".json", ".csv", ".html", ".htm":
		if !utf8.Valid(content) {
			return "", fmt.Errorf("file %s is not valid UTF-8 text", fileName)
		}
		return normalizeKnowledgeText(string(content)), nil
	case ".docx":
		return extractDOCXText(content)
	case ".xlsx":
		return extractXLSXText(content)
	case ".pdf":
		return extractPDFText(content)
	case ".doc", ".xls":
		return "", fmt.Errorf("file type %s is not supported yet", ext)
	default:
		if utf8.Valid(content) {
			return normalizeKnowledgeText(string(content)), nil
		}
		return "", fmt.Errorf("unsupported file type %s", ext)
	}
}

func extractDOCXText(content []byte) (string, error) {
	reader, err := zip.NewReader(bytes.NewReader(content), int64(len(content)))
	if err != nil {
		return "", fmt.Errorf("open docx: %w", err)
	}

	var parts []string
	for _, file := range reader.File {
		if file.Name != "word/document.xml" {
			continue
		}
		rc, err := file.Open()
		if err != nil {
			return "", fmt.Errorf("open docx entry: %w", err)
		}
		data, err := io.ReadAll(rc)
		rc.Close()
		if err != nil {
			return "", fmt.Errorf("read docx entry: %w", err)
		}

		decoder := xml.NewDecoder(bytes.NewReader(data))
		for {
			token, tokenErr := decoder.Token()
			if tokenErr == io.EOF {
				break
			}
			if tokenErr != nil {
				return "", fmt.Errorf("decode docx xml: %w", tokenErr)
			}
			switch typed := token.(type) {
			case xml.CharData:
				text := strings.TrimSpace(string(typed))
				if text != "" {
					parts = append(parts, text)
				}
			}
		}
	}
	if len(parts) == 0 {
		return "", fmt.Errorf("docx contains no readable text")
	}
	return normalizeKnowledgeText(strings.Join(parts, "\n")), nil
}

func extractXLSXText(content []byte) (string, error) {
	book, err := excelize.OpenReader(bytes.NewReader(content))
	if err != nil {
		return "", fmt.Errorf("open xlsx: %w", err)
	}
	defer func() { _ = book.Close() }()

	lines := make([]string, 0)
	for _, sheet := range book.GetSheetList() {
		rows, err := book.GetRows(sheet)
		if err != nil {
			return "", fmt.Errorf("read xlsx rows: %w", err)
		}
		if len(rows) == 0 {
			continue
		}
		lines = append(lines, sheet)
		for _, row := range rows {
			cells := make([]string, 0, len(row))
			for _, cell := range row {
				cell = strings.TrimSpace(cell)
				if cell != "" {
					cells = append(cells, cell)
				}
			}
			if len(cells) > 0 {
				lines = append(lines, strings.Join(cells, " | "))
			}
		}
	}
	if len(lines) == 0 {
		return "", fmt.Errorf("xlsx contains no readable text")
	}
	return normalizeKnowledgeText(strings.Join(lines, "\n")), nil
}

func extractPDFText(content []byte) (string, error) {
	reader, err := pdf.NewReader(bytes.NewReader(content), int64(len(content)))
	if err != nil {
		return "", fmt.Errorf("open pdf: %w", err)
	}

	var builder strings.Builder
	totalPage := reader.NumPage()
	for pageIndex := 1; pageIndex <= totalPage; pageIndex += 1 {
		page := reader.Page(pageIndex)
		if page.V.IsNull() {
			continue
		}
		text, err := page.GetPlainText(nil)
		if err != nil {
			return "", fmt.Errorf("extract pdf text: %w", err)
		}
		if text != "" {
			builder.WriteString(text)
			builder.WriteString("\n")
		}
	}

	result := normalizeKnowledgeText(builder.String())
	if result == "" {
		return "", fmt.Errorf("pdf contains no readable text")
	}
	return result, nil
}

func normalizeKnowledgeText(input string) string {
	input = strings.ReplaceAll(input, "\u00a0", " ")
	lines := strings.Split(input, "\n")
	cleaned := make([]string, 0, len(lines))
	for _, line := range lines {
		line = whitespacePattern.ReplaceAllString(strings.TrimSpace(line), " ")
		if line != "" {
			cleaned = append(cleaned, line)
		}
	}
	return strings.TrimSpace(strings.Join(cleaned, "\n"))
}

func splitKnowledgeTextIntoChunks(text string, chunkSize, chunkOverlap int) []string {
	text = normalizeKnowledgeText(text)
	if text == "" {
		return nil
	}

	runes := []rune(text)
	if chunkSize <= 0 {
		chunkSize = 900
	}
	if chunkOverlap < 0 {
		chunkOverlap = 0
	}
	if chunkOverlap >= chunkSize {
		chunkOverlap = chunkSize / 4
	}

	chunks := make([]string, 0)
	for start := 0; start < len(runes); {
		end := start + chunkSize
		if end > len(runes) {
			end = len(runes)
		}

		if end < len(runes) {
			for offset := end; offset > start+chunkSize/2; offset -= 1 {
				if runes[offset-1] == '\n' || runes[offset-1] == '。' || runes[offset-1] == '.' {
					end = offset
					break
				}
			}
		}

		chunk := strings.TrimSpace(string(runes[start:end]))
		if chunk != "" {
			chunks = append(chunks, chunk)
		}
		if end >= len(runes) {
			break
		}
		start = end - chunkOverlap
		if start < 0 {
			start = 0
		}
	}

	return chunks
}

func estimateKnowledgeTokenCount(text string) int {
	if text == "" {
		return 0
	}
	return utf8.RuneCountInString(text)
}