472 lines
14 KiB
Go
472 lines
14 KiB
Go
package kb
|
||
|
||
import (
|
||
"encoding/csv"
|
||
"fmt"
|
||
"io"
|
||
"path/filepath"
|
||
"strings"
|
||
"unicode/utf8"
|
||
|
||
"github.com/xuri/excelize/v2"
|
||
)
|
||
|
||
// ========================================================================
|
||
// Chunker：把长文本切成小段（chunk）
|
||
// ========================================================================
|
||
// 设计目标：
|
||
// - md 文件：按 H1/H2/H3 标题切，标题作为 chunk.title
|
||
// - txt 文件：按段落（双换行）切，无标题
|
||
// - 通用文本：滑动窗口按字符数切（不破坏 utf-8 字符边界）
|
||
// ========================================================================
|
||
|
||
// Chunk 切分结果（还没入库的中间产物）
|
||
type Chunk struct {
|
||
Title string // 分段标题
|
||
Content string // 分段正文
|
||
Meta map[string]any // 附加元数据（可空）
|
||
}
|
||
|
||
// ChunkOptions 切分参数
|
||
type ChunkOptions struct {
|
||
// MaxLen 单个 chunk 的最大字符数（按 rune 计，中文友好）
|
||
// 默认 500（中医知识片段大多 200-500 字）
|
||
MaxLen int
|
||
|
||
// Overlap 滑动窗口重叠字符数（按 rune 计）
|
||
// 默认 50，避免在关键句中间断开
|
||
Overlap int
|
||
}
|
||
|
||
// DefaultChunkOptions 默认切分参数
|
||
func DefaultChunkOptions() ChunkOptions {
|
||
return ChunkOptions{MaxLen: 500, Overlap: 50}
|
||
}
|
||
|
||
// ChunkMarkdown 切分 Markdown 文本
|
||
//
|
||
// 切分规则：
|
||
// 1. 先按 H1/H2/H3 (# / ## / ###) 把文档切成"章节"
|
||
// 2. 每个章节：标题作为 chunk.title，章节正文作为 chunk.content
|
||
// 3. 章节正文仍超 MaxLen 时，按段落滑动窗口再切
|
||
// 4. 没有标题的段落（如文档开头）归到一个 title="" 的 chunk
|
||
func ChunkMarkdown(text string, opt ChunkOptions) []*Chunk {
|
||
if opt.MaxLen <= 0 {
|
||
opt = DefaultChunkOptions()
|
||
}
|
||
lines := strings.Split(text, "\n")
|
||
var chunks []*Chunk
|
||
|
||
curTitle := ""
|
||
curBody := strings.Builder{}
|
||
|
||
// flush 把当前 buffer 的内容打包成一个或多个 chunk
|
||
flush := func() {
|
||
body := strings.TrimSpace(curBody.String())
|
||
if body == "" {
|
||
curBody.Reset()
|
||
return
|
||
}
|
||
// 仍然超过 MaxLen 的，按滑动窗口再切
|
||
for _, piece := range slideWindow(body, opt.MaxLen, opt.Overlap) {
|
||
chunks = append(chunks, &Chunk{
|
||
Title: curTitle,
|
||
Content: piece,
|
||
})
|
||
}
|
||
curBody.Reset()
|
||
}
|
||
|
||
for _, raw := range lines {
|
||
line := strings.TrimRight(raw, "\r")
|
||
trimmed := strings.TrimSpace(line)
|
||
// 命中 H1/H2/H3
|
||
if strings.HasPrefix(trimmed, "# ") ||
|
||
strings.HasPrefix(trimmed, "## ") ||
|
||
strings.HasPrefix(trimmed, "### ") {
|
||
flush()
|
||
curTitle = strings.TrimSpace(strings.TrimLeft(trimmed, "#"))
|
||
continue
|
||
}
|
||
curBody.WriteString(line)
|
||
curBody.WriteString("\n")
|
||
}
|
||
flush()
|
||
return chunks
|
||
}
|
||
|
||
// ChunkPlainText 切分纯文本（txt）
|
||
//
|
||
// 切分规则：
|
||
// 1. 按双换行切段落
|
||
// 2. 段落累计达到 MaxLen 时打包成 chunk
|
||
// 3. 单段落仍超 MaxLen 时，按滑动窗口切
|
||
func ChunkPlainText(text string, opt ChunkOptions) []*Chunk {
|
||
if opt.MaxLen <= 0 {
|
||
opt = DefaultChunkOptions()
|
||
}
|
||
text = strings.ReplaceAll(text, "\r\n", "\n")
|
||
paras := strings.Split(text, "\n\n")
|
||
|
||
var chunks []*Chunk
|
||
curBody := strings.Builder{}
|
||
flush := func() {
|
||
body := strings.TrimSpace(curBody.String())
|
||
if body == "" {
|
||
curBody.Reset()
|
||
return
|
||
}
|
||
for _, piece := range slideWindow(body, opt.MaxLen, opt.Overlap) {
|
||
chunks = append(chunks, &Chunk{Title: "", Content: piece})
|
||
}
|
||
curBody.Reset()
|
||
}
|
||
|
||
for _, p := range paras {
|
||
p = strings.TrimSpace(p)
|
||
if p == "" {
|
||
continue
|
||
}
|
||
// 如果累加后超过 MaxLen，先 flush 再开新段
|
||
if utf8.RuneCountInString(curBody.String())+utf8.RuneCountInString(p) > opt.MaxLen {
|
||
flush()
|
||
}
|
||
curBody.WriteString(p)
|
||
curBody.WriteString("\n\n")
|
||
}
|
||
flush()
|
||
return chunks
|
||
}
|
||
|
||
// slideWindow 滑动窗口切长文本
|
||
//
|
||
// 不破坏 utf-8 字符边界（按 rune 而不是 byte 切）
|
||
// 当文本长度 <= maxLen 时直接返回原文
|
||
func slideWindow(text string, maxLen, overlap int) []string {
|
||
runes := []rune(text)
|
||
if len(runes) <= maxLen {
|
||
return []string{text}
|
||
}
|
||
if overlap < 0 {
|
||
overlap = 0
|
||
}
|
||
if overlap >= maxLen {
|
||
overlap = maxLen / 4 // 防止 overlap 过大导致死循环
|
||
}
|
||
var out []string
|
||
step := maxLen - overlap
|
||
if step <= 0 {
|
||
step = maxLen
|
||
}
|
||
for i := 0; i < len(runes); i += step {
|
||
end := i + maxLen
|
||
if end > len(runes) {
|
||
end = len(runes)
|
||
}
|
||
out = append(out, string(runes[i:end]))
|
||
if end >= len(runes) {
|
||
break
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
// ========================================================================
|
||
// Importer：从文件解析成 []*Chunk
|
||
// ========================================================================
|
||
|
||
// ParsedDoc 文件解析结果（导入流程的中间产物）
|
||
type ParsedDoc struct {
|
||
Title string // 文档标题（默认取文件名不带扩展名）
|
||
SourceType string // 文件类型：xlsx / csv / md / txt / pdf / docx / html / manual
|
||
SourceFile string // 原始文件名（含扩展名）
|
||
RawContent string // 完整原文（用于详情展示）
|
||
Chunks []*Chunk // 切分好的分段
|
||
}
|
||
|
||
// ParseFileFromBytes 从字节流解析文件（不依赖磁盘）
|
||
//
|
||
// 用 Bytes 而非文件路径：HTTP 上传场景拿到的是 multipart.File，
|
||
// 调用方读成 []byte 传进来更通用
|
||
//
|
||
// 支持扩展名（解析产出纯文本后统一走自动分段，入库/检索链路完全复用）：
|
||
// .xlsx/.xls → MaxKB 导出三列格式按行成段；≥4 列的普通表格逐行拼文本再切分
|
||
// .csv → 通用表格，逐行拼文本后滑动窗口切分
|
||
// .md → 按 markdown H1/H2/H3 标题切分
|
||
// .txt → 按段落（双换行）切分
|
||
// .pdf → 逐页抽取文字层后按段落切分（扫描版无文字层会明确报错）
|
||
// .docx → 解析 word/document.xml 抽正文后按段落切分（老版 .doc 不支持）
|
||
// .html/.htm → 抽正文并把 h1~h6 转成 # 标题，走 markdown 章节切分
|
||
//
|
||
// txt/csv/md/html 会自动识别编码（UTF-8 / UTF-16 BOM / GBK）转成 UTF-8 入库
|
||
func ParseFileFromBytes(filename string, data []byte, opt ChunkOptions) (*ParsedDoc, error) {
|
||
ext := strings.ToLower(filepath.Ext(filename))
|
||
title := strings.TrimSuffix(filename, filepath.Ext(filename))
|
||
|
||
pdoc := &ParsedDoc{
|
||
Title: title,
|
||
SourceFile: filename,
|
||
}
|
||
|
||
switch ext {
|
||
case ".xlsx", ".xls":
|
||
pdoc.SourceType = "xlsx"
|
||
chunks, raw, err := parseExcel(data, opt)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pdoc.Chunks = chunks
|
||
pdoc.RawContent = raw
|
||
case ".csv":
|
||
pdoc.SourceType = "csv"
|
||
chunks, raw, err := parseCSV(data, opt)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pdoc.Chunks = chunks
|
||
pdoc.RawContent = raw
|
||
case ".md", ".markdown":
|
||
pdoc.SourceType = "md"
|
||
text := decodeToUTF8(data)
|
||
pdoc.RawContent = text
|
||
pdoc.Chunks = ChunkMarkdown(text, opt)
|
||
case ".txt", "":
|
||
pdoc.SourceType = "txt"
|
||
text := decodeToUTF8(data)
|
||
pdoc.RawContent = text
|
||
pdoc.Chunks = ChunkPlainText(text, opt)
|
||
case ".pdf":
|
||
pdoc.SourceType = "pdf"
|
||
text, err := extractPDFText(data)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pdoc.RawContent = text
|
||
pdoc.Chunks = ChunkPlainText(text, opt)
|
||
case ".docx":
|
||
pdoc.SourceType = "docx"
|
||
text, err := extractDocxText(data)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pdoc.RawContent = text
|
||
pdoc.Chunks = ChunkPlainText(text, opt)
|
||
case ".doc":
|
||
// 老版 .doc 是 OLE 二进制格式，纯 Go 解析成本极高且不可靠，明确拒绝并给出替代方案
|
||
return nil, fmt.Errorf("kb: 不支持老版二进制 .doc，请用 Word 另存为 .docx 后再导入")
|
||
case ".html", ".htm":
|
||
pdoc.SourceType = "html"
|
||
text, err := extractHTMLText(data)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pdoc.RawContent = text
|
||
// HTML 抽取时已把 h1~h6 转成 # 标题，按 markdown 章节切分能保留标题到 chunk.title
|
||
pdoc.Chunks = ChunkMarkdown(text, opt)
|
||
default:
|
||
return nil, fmt.Errorf("kb: 不支持的文件格式 %s（支持 xlsx/xls/csv/md/txt/pdf/docx/html）", ext)
|
||
}
|
||
|
||
// 空内容防御：解析成功但没有任何 chunk 时明确报错，
|
||
// 避免入库一个空文档让用户误以为导入成功
|
||
if len(pdoc.Chunks) == 0 {
|
||
if ext == ".pdf" {
|
||
return nil, fmt.Errorf("kb: PDF 未提取到文本——扫描版/图片型 PDF 没有文字层，请先 OCR 或转成 txt 再导入")
|
||
}
|
||
return nil, fmt.Errorf("kb: 文件中没有可导入的文本内容")
|
||
}
|
||
return pdoc, nil
|
||
}
|
||
|
||
// parseExcel 解析 xlsx，自动识别两种模式
|
||
//
|
||
// 模式一（MaxKB 导出格式，兼容原有行为）：
|
||
// 第 1 列：分段标题（必填）
|
||
// 第 2 列：分段内容（必填）
|
||
// 第 3 列：关联问题列表（可选，分号分隔）
|
||
// 判定条件：首行是表头（含 标题/内容/title/content 字样）或最大列数 ≤ 3，
|
||
// 每行直接成为一个 chunk，不再二次切分
|
||
//
|
||
// 模式二（通用表格）：
|
||
// 列数 ≥ 4 且无 MaxKB 表头的普通业务表格（如 药名|性味|归经|功效），
|
||
// 逐行把单元格用「 | 」拼成一行文本，按滑动窗口切分，遍历所有 sheet
|
||
//
|
||
// 注意：老版二进制 .xls 无法被 excelize 打开，会走到「打开失败」的报错提示
|
||
func parseExcel(data []byte, opt ChunkOptions) ([]*Chunk, string, error) {
|
||
f, err := excelize.OpenReader(bytesReader(data))
|
||
if err != nil {
|
||
return nil, "", fmt.Errorf("kb: 打开 Excel 失败（若为老版 .xls 请另存为 .xlsx）: %w", err)
|
||
}
|
||
defer f.Close()
|
||
|
||
sheets := f.GetSheetList()
|
||
if len(sheets) == 0 {
|
||
return nil, "", fmt.Errorf("kb: xlsx 没有任何 sheet")
|
||
}
|
||
// 先读第一个 sheet 判定模式（MaxKB 导出只有一个 sheet）
|
||
rows, err := f.GetRows(sheets[0])
|
||
if err != nil {
|
||
return nil, "", fmt.Errorf("kb: 读取 xlsx 行失败: %w", err)
|
||
}
|
||
|
||
// 模式判定：无 MaxKB 表头且列数 ≥ 4 → 通用表格模式
|
||
maxCols := 0
|
||
for _, row := range rows {
|
||
if len(row) > maxCols {
|
||
maxCols = len(row)
|
||
}
|
||
}
|
||
hasHeader := len(rows) > 0 && isHeaderRow(rows[0])
|
||
if !hasHeader && maxCols >= 4 {
|
||
return parseExcelGeneric(f, opt)
|
||
}
|
||
|
||
var chunks []*Chunk
|
||
var rawBuilder strings.Builder
|
||
for i, row := range rows {
|
||
// 第一行如果是表头（含"标题"/"内容"等字样）则跳过
|
||
if i == 0 && isHeaderRow(row) {
|
||
continue
|
||
}
|
||
if len(row) == 0 {
|
||
continue
|
||
}
|
||
// 容错：不足 2 列时补空字符串
|
||
title := cellOr(row, 0, "")
|
||
content := cellOr(row, 1, "")
|
||
related := cellOr(row, 2, "")
|
||
|
||
// 至少要有 content（如果只有 title 没 content，跳过）
|
||
if strings.TrimSpace(content) == "" {
|
||
// 如果只有 1 列且非空，把第 1 列当 content（无标题）
|
||
if strings.TrimSpace(title) != "" && len(row) == 1 {
|
||
chunks = append(chunks, &Chunk{Title: "", Content: title})
|
||
rawBuilder.WriteString(title)
|
||
rawBuilder.WriteString("\n")
|
||
}
|
||
continue
|
||
}
|
||
|
||
chunk := &Chunk{Title: title, Content: content}
|
||
// 关联问题塞到 meta 里（V1 仅作展示，V2 也用作 keyword boost）
|
||
if strings.TrimSpace(related) != "" {
|
||
qs := strings.Split(related, ";")
|
||
cleaned := make([]string, 0, len(qs))
|
||
for _, q := range qs {
|
||
if s := strings.TrimSpace(q); s != "" {
|
||
cleaned = append(cleaned, s)
|
||
}
|
||
}
|
||
if len(cleaned) > 0 {
|
||
chunk.Meta = map[string]any{"related_questions": cleaned}
|
||
}
|
||
}
|
||
chunks = append(chunks, chunk)
|
||
|
||
rawBuilder.WriteString(title)
|
||
rawBuilder.WriteString("\n")
|
||
rawBuilder.WriteString(content)
|
||
rawBuilder.WriteString("\n\n")
|
||
}
|
||
return chunks, rawBuilder.String(), nil
|
||
}
|
||
|
||
// parseExcelGeneric 通用表格模式：遍历所有 sheet，逐行拼接文本后滑动窗口切分
|
||
//
|
||
// 每行单元格用「 | 」连接保持列对应关系，chunk.title 用 sheet 名，
|
||
// 检索时 sheet 名也参与 FULLTEXT 匹配（如 sheet 叫「感冒用药」）
|
||
func parseExcelGeneric(f *excelize.File, opt ChunkOptions) ([]*Chunk, string, error) {
|
||
if opt.MaxLen <= 0 {
|
||
opt = DefaultChunkOptions()
|
||
}
|
||
var chunks []*Chunk
|
||
var rawBuilder strings.Builder
|
||
for _, sheet := range f.GetSheetList() {
|
||
rows, err := f.GetRows(sheet)
|
||
if err != nil {
|
||
continue // 单个 sheet 读取失败不影响其他 sheet
|
||
}
|
||
var sheetText strings.Builder
|
||
for _, row := range rows {
|
||
line := strings.TrimSpace(strings.Join(row, " | "))
|
||
// 跳过全空行（只剩分隔符和空白）
|
||
if strings.Trim(line, "| ") == "" {
|
||
continue
|
||
}
|
||
sheetText.WriteString(line)
|
||
sheetText.WriteString("\n")
|
||
}
|
||
text := strings.TrimSpace(sheetText.String())
|
||
if text == "" {
|
||
continue
|
||
}
|
||
rawBuilder.WriteString("【" + sheet + "】\n")
|
||
rawBuilder.WriteString(text)
|
||
rawBuilder.WriteString("\n\n")
|
||
for _, piece := range slideWindow(text, opt.MaxLen, opt.Overlap) {
|
||
chunks = append(chunks, &Chunk{Title: sheet, Content: piece})
|
||
}
|
||
}
|
||
return chunks, rawBuilder.String(), nil
|
||
}
|
||
|
||
// parseCSV 解析 CSV 文件：逐行拼接文本后滑动窗口切分
|
||
//
|
||
// 容错设计：
|
||
// FieldsPerRecord=-1 允许每行列数不同（业务系统导出的 csv 经常不规整）
|
||
// LazyQuotes=true 容忍不规范的引号用法
|
||
// 编码自动识别（Excel 另存的 csv 大概率是 GBK）
|
||
func parseCSV(data []byte, opt ChunkOptions) ([]*Chunk, string, error) {
|
||
if opt.MaxLen <= 0 {
|
||
opt = DefaultChunkOptions()
|
||
}
|
||
r := csv.NewReader(strings.NewReader(decodeToUTF8(data)))
|
||
r.FieldsPerRecord = -1
|
||
r.LazyQuotes = true
|
||
|
||
var lines []string
|
||
for {
|
||
record, err := r.Read()
|
||
if err == io.EOF {
|
||
break
|
||
}
|
||
if err != nil {
|
||
return nil, "", fmt.Errorf("kb: 解析 CSV 失败: %w", err)
|
||
}
|
||
line := strings.TrimSpace(strings.Join(record, " | "))
|
||
if strings.Trim(line, "| ") == "" {
|
||
continue
|
||
}
|
||
lines = append(lines, line)
|
||
}
|
||
text := strings.Join(lines, "\n")
|
||
if strings.TrimSpace(text) == "" {
|
||
return nil, "", nil
|
||
}
|
||
var chunks []*Chunk
|
||
for _, piece := range slideWindow(text, opt.MaxLen, opt.Overlap) {
|
||
chunks = append(chunks, &Chunk{Title: "", Content: piece})
|
||
}
|
||
return chunks, text, nil
|
||
}
|
||
|
||
// isHeaderRow 判断是否表头行（包含"标题"或"内容"字样）
|
||
func isHeaderRow(row []string) bool {
|
||
if len(row) == 0 {
|
||
return false
|
||
}
|
||
joined := strings.ToLower(strings.Join(row, ""))
|
||
return strings.Contains(joined, "标题") ||
|
||
strings.Contains(joined, "title") ||
|
||
strings.Contains(joined, "内容") ||
|
||
strings.Contains(joined, "content")
|
||
}
|
||
|
||
// cellOr 安全取单元格（防止行不满列数导致 index out of range）
|
||
func cellOr(row []string, idx int, def string) string {
|
||
if idx >= len(row) {
|
||
return def
|
||
}
|
||
return row[idx]
|
||
}
|