WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,268 @@
|
||||
// Package chunker - header_tracker.go implements context-preserving header
|
||||
// tracking for document chunking, ported from docreader/splitter/header_hook.py.
|
||||
//
|
||||
// When a large Markdown table is split across multiple chunks, each chunk after
|
||||
// the first would lose the table header context. The headerTracker detects table
|
||||
// headers and signals the merge logic to prepend them to subsequent chunks.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// headerTrackerHook defines a pattern pair for detecting contextual headers.
|
||||
// When startPattern matches a unit's text, that text becomes an "active header".
|
||||
// The header stays active until endPattern matches a subsequent unit.
|
||||
type headerTrackerHook struct {
|
||||
startPattern *regexp.Regexp
|
||||
endPattern *regexp.Regexp
|
||||
priority int
|
||||
}
|
||||
|
||||
// defaultHeaderHooks returns header tracking hooks matching the Python defaults
|
||||
// in docreader/splitter/header_hook.py.
|
||||
var defaultHeaderHooks = []headerTrackerHook{
|
||||
{
|
||||
// Markdown table: header row + separator row (e.g. "| A | B |\n| --- | --- |\n")
|
||||
startPattern: regexp.MustCompile(`(?si)^\s*(?:\|[^|\n]*)+[\r\n]+\s*(?:\|\s*:?-{3,}:?\s*)+\|?[\r\n]+$`),
|
||||
// Empty/whitespace line or a line that doesn't start with | or whitespace
|
||||
endPattern: regexp.MustCompile(`(?si)^\s*$|^\s*[^|\s].*$`),
|
||||
priority: 15,
|
||||
},
|
||||
}
|
||||
|
||||
// tableRowPattern matches a single Markdown table row: "| cell | cell | ... |\n"
|
||||
var tableRowPattern = regexp.MustCompile(`(?m)^\s*(?:\|[^|\n]*)+\|\s*$`)
|
||||
|
||||
// markdownTableHookPriority matches DEFAULT_CONFIGS / defaultHeaderHooks table hook.
|
||||
const markdownTableHookPriority = 15
|
||||
|
||||
// headerTracker maintains the state of active headers across split units.
|
||||
type headerTracker struct {
|
||||
hooks []headerTrackerHook
|
||||
activeHeaders map[int]string // priority -> header text
|
||||
endedHeaders map[int]bool // priorities that have been ended
|
||||
pendingExtend map[int]bool // headers with empty column names awaiting first data row
|
||||
// pendingTableBreak is set when a table row unit ends with a paragraph break
|
||||
// (the blank line between tables is consumed by \n\n splitting). The header
|
||||
// stays active until the next unit is seen so we can detect a new table.
|
||||
pendingTableBreak bool
|
||||
// headerEndedThisUnit tells mergeUnits to flush before the current unit when a
|
||||
// new table starts (column mismatch or pendingTableBreak + table row).
|
||||
headerEndedThisUnit bool
|
||||
}
|
||||
|
||||
func newHeaderTracker() *headerTracker {
|
||||
return &headerTracker{
|
||||
hooks: defaultHeaderHooks,
|
||||
activeHeaders: make(map[int]string),
|
||||
endedHeaders: make(map[int]bool),
|
||||
pendingExtend: make(map[int]bool),
|
||||
}
|
||||
}
|
||||
|
||||
// update checks split text for header start/end markers and updates internal state.
|
||||
func (ht *headerTracker) update(split string) {
|
||||
ht.headerEndedThisUnit = false
|
||||
|
||||
if ht.pendingTableBreak {
|
||||
ht.pendingTableBreak = false
|
||||
if _, active := ht.activeHeaders[markdownTableHookPriority]; active {
|
||||
if firstTableRowColumnCount(split) > 0 {
|
||||
ht.clearTableHeader()
|
||||
ht.headerEndedThisUnit = true
|
||||
} else {
|
||||
ht.clearTableHeader()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 1. Check for header-end markers among currently active headers
|
||||
for _, hook := range ht.hooks {
|
||||
if _, active := ht.activeHeaders[hook.priority]; active {
|
||||
if hook.endPattern.MatchString(split) {
|
||||
ht.endedHeaders[hook.priority] = true
|
||||
delete(ht.activeHeaders, hook.priority)
|
||||
delete(ht.pendingExtend, hook.priority)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 1b. Paragraph splits consume the blank line between tables. Mark a break
|
||||
// after "| last row |\n\n" and resolve on the next unit; also end when a new
|
||||
// table row has a different column count than the active header.
|
||||
if _, active := ht.activeHeaders[markdownTableHookPriority]; active {
|
||||
if !ht.pendingExtend[markdownTableHookPriority] {
|
||||
if splitEndsWithParagraphBreak(split) {
|
||||
ht.pendingTableBreak = true
|
||||
} else {
|
||||
ht.endTableHeaderOnColumnMismatch(split)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. If a header has an empty column-name row (e.g. "||"), replace it with
|
||||
// a proper Markdown table header using the first data row as column names.
|
||||
//
|
||||
// Before: "||" + "| --- | --- |\n"
|
||||
// After: "| col1 | col2 |\n" + "| --- | --- |\n"
|
||||
for p := range ht.pendingExtend {
|
||||
if _, active := ht.activeHeaders[p]; active && tableRowPattern.MatchString(split) {
|
||||
sep := extractSeparatorLine(ht.activeHeaders[p])
|
||||
ht.activeHeaders[p] = split + sep
|
||||
}
|
||||
delete(ht.pendingExtend, p)
|
||||
}
|
||||
|
||||
// 3. Check for new header-start markers (only for hooks that are neither active nor ended)
|
||||
for _, hook := range ht.hooks {
|
||||
if _, active := ht.activeHeaders[hook.priority]; active {
|
||||
continue
|
||||
}
|
||||
if ht.endedHeaders[hook.priority] {
|
||||
continue
|
||||
}
|
||||
if loc := hook.startPattern.FindString(split); loc != "" {
|
||||
ht.activeHeaders[hook.priority] = loc
|
||||
if isEmptyTableHeaderRow(loc) {
|
||||
ht.pendingExtend[hook.priority] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 4. If all headers ended, clear the ended set so future tables can be tracked
|
||||
if len(ht.activeHeaders) == 0 {
|
||||
for k := range ht.endedHeaders {
|
||||
delete(ht.endedHeaders, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getHeaders returns all active headers concatenated, sorted by priority descending.
|
||||
func (ht *headerTracker) getHeaders() string {
|
||||
if len(ht.activeHeaders) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
type entry struct {
|
||||
priority int
|
||||
text string
|
||||
}
|
||||
entries := make([]entry, 0, len(ht.activeHeaders))
|
||||
for p, t := range ht.activeHeaders {
|
||||
entries = append(entries, entry{p, t})
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool {
|
||||
return entries[i].priority > entries[j].priority
|
||||
})
|
||||
|
||||
parts := make([]string, len(entries))
|
||||
for i, e := range entries {
|
||||
parts[i] = e.text
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
// isEmptyTableHeaderRow checks if the header row (the line before the separator)
|
||||
// contains only pipes and whitespace — meaning the column names are empty.
|
||||
// This is common with MarkItDown and similar converters that produce tables like:
|
||||
//
|
||||
// ||
|
||||
// | --- | --- |
|
||||
// | real column A | real column B |
|
||||
func isEmptyTableHeaderRow(header string) bool {
|
||||
idx := strings.IndexByte(header, '\n')
|
||||
if idx < 0 {
|
||||
return false
|
||||
}
|
||||
row := strings.TrimSpace(header[:idx])
|
||||
for _, r := range row {
|
||||
if r != '|' && r != ' ' && r != '\t' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// extractSeparatorLine returns the separator line (e.g. "| --- | --- |\n") from
|
||||
// a table header string. It looks for the line containing "---".
|
||||
func extractSeparatorLine(header string) string {
|
||||
for _, line := range strings.Split(header, "\n") {
|
||||
if strings.Contains(line, "---") {
|
||||
return line + "\n"
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (ht *headerTracker) clearTableHeader() {
|
||||
ht.endedHeaders[markdownTableHookPriority] = true
|
||||
delete(ht.activeHeaders, markdownTableHookPriority)
|
||||
delete(ht.pendingExtend, markdownTableHookPriority)
|
||||
}
|
||||
|
||||
func (ht *headerTracker) endTableHeaderOnColumnMismatch(split string) {
|
||||
header, ok := ht.activeHeaders[markdownTableHookPriority]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
rowCols := firstTableRowColumnCount(split)
|
||||
headerCols := headerTableColumnCount(header)
|
||||
if rowCols > 0 && headerCols > 0 && rowCols != headerCols {
|
||||
ht.clearTableHeader()
|
||||
ht.headerEndedThisUnit = true
|
||||
}
|
||||
}
|
||||
|
||||
func splitEndsWithParagraphBreak(split string) bool {
|
||||
trimmed := strings.TrimRight(split, " \t\r")
|
||||
return strings.HasSuffix(trimmed, "\n\n") || strings.HasSuffix(trimmed, "\r\n\r\n")
|
||||
}
|
||||
|
||||
func tableRowColumnCount(line string) int {
|
||||
line = strings.TrimSpace(line)
|
||||
if !strings.HasPrefix(line, "|") {
|
||||
return 0
|
||||
}
|
||||
parts := strings.Split(line, "|")
|
||||
if len(parts) > 0 && strings.TrimSpace(parts[0]) == "" {
|
||||
parts = parts[1:]
|
||||
}
|
||||
if len(parts) > 0 && strings.TrimSpace(parts[len(parts)-1]) == "" {
|
||||
parts = parts[:len(parts)-1]
|
||||
}
|
||||
return len(parts)
|
||||
}
|
||||
|
||||
func firstTableRowColumnCount(text string) int {
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line != "" && tableRowPattern.MatchString(line) {
|
||||
return tableRowColumnCount(line)
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func headerTableColumnCount(header string) int {
|
||||
for _, line := range strings.Split(header, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.Contains(line, "---") {
|
||||
continue
|
||||
}
|
||||
if n := tableRowColumnCount(line); n > 0 {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// headerColumnMismatch reports whether the next split unit starts a new table
|
||||
// whose width differs from the active markdown table header.
|
||||
func headerColumnMismatch(headers, nextUnit string) bool {
|
||||
headerCols := headerTableColumnCount(headers)
|
||||
rowCols := firstTableRowColumnCount(nextUnit)
|
||||
return headerCols > 0 && rowCols > 0 && headerCols != rowCols
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
// Package chunker - heading_hierarchy.go tracks Markdown heading nesting so
|
||||
// that the heading-aware splitter (heading_splitter.go) can prepend a
|
||||
// breadcrumb like "# Top > ## Section > ### Subsection" to each chunk.
|
||||
//
|
||||
// Conceptually similar to header_tracker.go (which tracks table headers via
|
||||
// start/end hooks) but a Markdown heading does not have an explicit end —
|
||||
// it is ended by the next heading of equal or shallower depth. The hook
|
||||
// abstraction in header_tracker.go does not fit that pattern, so we model
|
||||
// it as an explicit level-stack.
|
||||
package chunker
|
||||
|
||||
import "strings"
|
||||
|
||||
// HeadingHierarchy maintains a stack of active Markdown headings indexed by
|
||||
// level (1..6). Pushing a level-N heading pops every entry of level >= N
|
||||
// because the previous siblings/descendants are no longer in scope.
|
||||
type HeadingHierarchy struct {
|
||||
// stack[i] holds the heading text for level i+1 (so stack[0] = H1).
|
||||
// Entries beyond the deepest active level are empty strings.
|
||||
stack [6]string
|
||||
depth int // current deepest active level (0 if no active heading)
|
||||
}
|
||||
|
||||
// NewHeadingHierarchy returns an empty hierarchy.
|
||||
func NewHeadingHierarchy() *HeadingHierarchy {
|
||||
return &HeadingHierarchy{}
|
||||
}
|
||||
|
||||
// Observe parses line and updates the hierarchy if line is a Markdown
|
||||
// heading. Returns the (level, headingText) when a heading was recognized,
|
||||
// or (0, "") otherwise. Lines that look like headings inside fenced code
|
||||
// blocks are NOT detected here — callers must avoid feeding code-block
|
||||
// content to Observe (the heading splitter does so).
|
||||
func (h *HeadingHierarchy) Observe(line string) (int, string) {
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m == nil {
|
||||
return 0, ""
|
||||
}
|
||||
level := len(m[1])
|
||||
if level < 1 || level > 6 {
|
||||
return 0, ""
|
||||
}
|
||||
heading := strings.TrimSpace(m[2])
|
||||
// Replace this level and clear deeper ones — siblings/descendants of
|
||||
// the previous heading at this level are no longer in scope.
|
||||
h.stack[level-1] = heading
|
||||
for i := level; i < 6; i++ {
|
||||
h.stack[i] = ""
|
||||
}
|
||||
if level > h.depth {
|
||||
h.depth = level
|
||||
} else {
|
||||
// Recompute depth: it might shrink if we just pushed a shallower heading.
|
||||
h.depth = 0
|
||||
for i := 0; i < 6; i++ {
|
||||
if h.stack[i] != "" {
|
||||
h.depth = i + 1
|
||||
}
|
||||
}
|
||||
}
|
||||
return level, heading
|
||||
}
|
||||
|
||||
// Breadcrumb returns the current heading path joined by " > ", e.g.
|
||||
// "Chapter 1 > Section 2 > Subsection a". Returns "" when no headings
|
||||
// are active.
|
||||
func (h *HeadingHierarchy) Breadcrumb() string {
|
||||
if h.depth == 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, h.depth)
|
||||
for i := 0; i < h.depth; i++ {
|
||||
if h.stack[i] != "" {
|
||||
parts = append(parts, h.stack[i])
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " > ")
|
||||
}
|
||||
|
||||
// BreadcrumbWithHashes returns the path with the original `#` prefixes,
|
||||
// suitable for embedding back into chunk content as a context header.
|
||||
// Example: "# Chapter 1\n## Section 2\n### Subsection a"
|
||||
func (h *HeadingHierarchy) BreadcrumbWithHashes() string {
|
||||
if h.depth == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
for i := 0; i < h.depth; i++ {
|
||||
if h.stack[i] == "" {
|
||||
continue
|
||||
}
|
||||
if sb.Len() > 0 {
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
sb.WriteString(strings.Repeat("#", i+1))
|
||||
sb.WriteByte(' ')
|
||||
sb.WriteString(h.stack[i])
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// Depth returns the current deepest active heading level.
|
||||
func (h *HeadingHierarchy) Depth() int { return h.depth }
|
||||
|
||||
// Reset clears all state.
|
||||
func (h *HeadingHierarchy) Reset() {
|
||||
for i := range h.stack {
|
||||
h.stack[i] = ""
|
||||
}
|
||||
h.depth = 0
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestHeadingHierarchy_LinearNesting(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Chapter 1")
|
||||
h.Observe("## Section 1.1")
|
||||
h.Observe("### Subsection 1.1.1")
|
||||
if got := h.Breadcrumb(); got != "Chapter 1 > Section 1.1 > Subsection 1.1.1" {
|
||||
t.Errorf("breadcrumb mismatch: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_PopsDeeperOnSiblingHeading(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Chapter 1")
|
||||
h.Observe("## Section 1.1")
|
||||
h.Observe("### Subsection 1.1.1")
|
||||
h.Observe("## Section 1.2") // pops the H3
|
||||
if got := h.Breadcrumb(); got != "Chapter 1 > Section 1.2" {
|
||||
t.Errorf("breadcrumb after sibling: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_PopsAllOnNewTopLevel(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Chapter 1")
|
||||
h.Observe("## Section A")
|
||||
h.Observe("### Sub")
|
||||
h.Observe("# Chapter 2")
|
||||
if got := h.Breadcrumb(); got != "Chapter 2" {
|
||||
t.Errorf("breadcrumb after new H1: %q", got)
|
||||
}
|
||||
if h.Depth() != 1 {
|
||||
t.Errorf("depth should be 1, got %d", h.Depth())
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_NonHeadingIgnored(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# Title")
|
||||
level, txt := h.Observe("just a paragraph")
|
||||
if level != 0 || txt != "" {
|
||||
t.Errorf("non-heading should not register: %d %q", level, txt)
|
||||
}
|
||||
if got := h.Breadcrumb(); got != "Title" {
|
||||
t.Errorf("breadcrumb after non-heading: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_BreadcrumbWithHashes(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# A")
|
||||
h.Observe("## B")
|
||||
got := h.BreadcrumbWithHashes()
|
||||
want := "# A\n## B"
|
||||
if got != want {
|
||||
t.Errorf("hashes breadcrumb:\n got: %q\n want: %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_EmptyState(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
if h.Breadcrumb() != "" || h.BreadcrumbWithHashes() != "" || h.Depth() != 0 {
|
||||
t.Error("empty hierarchy should have no breadcrumb / depth=0")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_Reset(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
h.Observe("# A")
|
||||
h.Observe("## B")
|
||||
h.Reset()
|
||||
if h.Breadcrumb() != "" || h.Depth() != 0 {
|
||||
t.Error("Reset did not clear state")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingHierarchy_SkipLevels(t *testing.T) {
|
||||
h := NewHeadingHierarchy()
|
||||
// Document jumps from H1 directly to H3, skipping H2.
|
||||
h.Observe("# Top")
|
||||
h.Observe("### Deep")
|
||||
if got := h.Breadcrumb(); got != "Top > Deep" {
|
||||
t.Errorf("skipped-level breadcrumb: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,342 @@
|
||||
// Package chunker - heading_splitter.go implements Tier 1: Markdown
|
||||
// heading-aware chunking. Documents with proper heading structure are split
|
||||
// at heading boundaries and each chunk is prefixed with a breadcrumb of
|
||||
// active heading context (e.g. "# Chapter 1\n## Section 1.2").
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// init wires this implementation into the strategy resolver.
|
||||
func init() {
|
||||
splitByHeadings = splitByHeadingsImpl
|
||||
}
|
||||
|
||||
// headingBoundary marks where a section starts. The first boundary is at
|
||||
// rune offset 0 (covers any preamble before the first heading), subsequent
|
||||
// boundaries sit at headings whose level is <= primaryLevel.
|
||||
type headingBoundary struct {
|
||||
runeStart int
|
||||
line string // raw heading line, "" when this is the leading boundary
|
||||
}
|
||||
|
||||
// splitByHeadingsImpl is the Tier-1 implementation. It falls through to the
|
||||
// legacy splitter when the document has no usable heading structure or when
|
||||
// the heading split would produce a single section anyway.
|
||||
//
|
||||
// profile may be nil; we compute one on demand. When the strategy resolver
|
||||
// already ran the profiler (auto strategy), the same profile is threaded
|
||||
// through here so we don't re-scan the entire document.
|
||||
func splitByHeadingsImpl(text string, cfg SplitterConfig, profile *DocProfile) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
if profile == nil {
|
||||
profile = ProfileDocument(text)
|
||||
}
|
||||
primaryLevel := profile.DominantHeadingLevel()
|
||||
if primaryLevel == 0 {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
bounds := findHeadingBoundaries(text, primaryLevel)
|
||||
if len(bounds) <= 1 {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
runes := []rune(text)
|
||||
hierarchy := NewHeadingHierarchy()
|
||||
|
||||
// Pre-walk every heading (not just primary-level) so the hierarchy
|
||||
// reflects the full nesting context for each section's start. We only
|
||||
// snapshot the breadcrumb at section boundaries; deeper sub-headings
|
||||
// inside a section update the hierarchy but do not change the chunk's
|
||||
// breadcrumb (chunks within a section share one breadcrumb).
|
||||
var out []Chunk
|
||||
seq := 0
|
||||
|
||||
for i, b := range bounds {
|
||||
endRune := len(runes)
|
||||
if i+1 < len(bounds) {
|
||||
endRune = bounds[i+1].runeStart
|
||||
}
|
||||
if b.line != "" {
|
||||
hierarchy.Observe(b.line)
|
||||
}
|
||||
// Catch sub-headings that occur between this primary boundary and
|
||||
// the next so the hierarchy stays in sync for subsequent sections.
|
||||
// We intentionally do this after observing the section header so
|
||||
// the breadcrumb reflects the section-leading heading.
|
||||
breadcrumb := hierarchy.BreadcrumbWithHashes()
|
||||
sectionStart := *hierarchy
|
||||
observeSubHeadings(runes[b.runeStart:endRune], primaryLevel, hierarchy)
|
||||
|
||||
sectionRunes := runes[b.runeStart:endRune]
|
||||
sectionContent := string(sectionRunes)
|
||||
secLen := len(sectionRunes)
|
||||
if secLen == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
bcLen := utf8.RuneCountInString(breadcrumb)
|
||||
// Single-chunk section: emit as-is, breadcrumb tracked separately.
|
||||
// The breadcrumb is delivered via Chunk.ContextHeader (not Content)
|
||||
// to preserve End-Start == len(Content) invariants relied on by
|
||||
// document reconstruction (knowledge.go:2278+).
|
||||
if bcLen+2+secLen <= cfg.ChunkSize {
|
||||
out = append(out, Chunk{
|
||||
Content: sectionContent,
|
||||
ContextHeader: breadcrumb,
|
||||
Seq: seq,
|
||||
Start: b.runeStart,
|
||||
End: endRune,
|
||||
})
|
||||
seq++
|
||||
continue
|
||||
}
|
||||
|
||||
// Section too large: defer to the legacy splitter for inner
|
||||
// segmentation. We do NOT shrink the inner ChunkSize budget here
|
||||
// because the breadcrumb no longer counts against Content size.
|
||||
// Each sub-chunk gets a breadcrumb reflecting the deepest heading
|
||||
// active at its start, so deep `###`/`####` sub-headings inside a
|
||||
// long section aren't collapsed to the section-level header.
|
||||
subBreadcrumbs := sectionBreadcrumbs(sectionRunes, primaryLevel, sectionStart)
|
||||
subChunks := SplitText(sectionContent, cfg)
|
||||
for _, sub := range subChunks {
|
||||
out = append(out, Chunk{
|
||||
Content: sub.Content,
|
||||
ContextHeader: breadcrumbAtOffset(subBreadcrumbs, sub.Start, breadcrumb),
|
||||
Seq: seq,
|
||||
Start: b.runeStart + sub.Start,
|
||||
End: b.runeStart + sub.End,
|
||||
})
|
||||
seq++
|
||||
}
|
||||
}
|
||||
|
||||
return coalesceTinyChunks(out, cfg.ChunkSize)
|
||||
}
|
||||
|
||||
// coalesceTinyChunks merges adjacent small chunks under their shared heading
|
||||
// context so that documents whose primary sections are mostly short (FAQs,
|
||||
// install logs, change-lists) don't trip the validator's "too many tiny
|
||||
// chunks" rule and fall through all the way to legacy. The merged breadcrumb
|
||||
// is the line-prefix shared by both inputs; the original sub-headings remain
|
||||
// visible because heading_splitter includes the heading line in each
|
||||
// section's Content.
|
||||
//
|
||||
// Safety:
|
||||
// - We only merge when cur.End == next.Start. That preserves the
|
||||
// End-Start == len([]rune(Content)) invariant that document
|
||||
// reconstruction relies on, and naturally skips legacy sub-chunks (which
|
||||
// may overlap due to ChunkOverlap).
|
||||
// - We stop accumulating once the running chunk reaches the merge target
|
||||
// (≈ ChunkSize/2) so we don't aggressively pack chunks beyond what the
|
||||
// validator considers comfortable.
|
||||
func coalesceTinyChunks(in []Chunk, chunkSize int) []Chunk {
|
||||
if len(in) <= 1 || chunkSize <= 0 {
|
||||
return in
|
||||
}
|
||||
target := chunkSize / 2
|
||||
if target < 200 {
|
||||
target = 200
|
||||
}
|
||||
|
||||
out := make([]Chunk, 0, len(in))
|
||||
cur := in[0]
|
||||
curLen := utf8.RuneCountInString(cur.Content)
|
||||
|
||||
for i := 1; i < len(in); i++ {
|
||||
next := in[i]
|
||||
nextLen := utf8.RuneCountInString(next.Content)
|
||||
sharedHeader := commonHeadingPrefix(cur.ContextHeader, next.ContextHeader)
|
||||
// Adjacent + still-small + would not blow the size budget → merge.
|
||||
if sharedHeader != "" && cur.End == next.Start && curLen < target && curLen+nextLen <= chunkSize {
|
||||
cur.Content += next.Content
|
||||
cur.ContextHeader = sharedHeader
|
||||
cur.End = next.End
|
||||
curLen += nextLen
|
||||
continue
|
||||
}
|
||||
out = append(out, cur)
|
||||
cur = next
|
||||
curLen = nextLen
|
||||
}
|
||||
out = append(out, cur)
|
||||
|
||||
// Re-sequence — downstream code (knowledge.go) expects Seq to be a dense
|
||||
// 0..N-1 range over the returned slice.
|
||||
for i := range out {
|
||||
out[i].Seq = i
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// commonHeadingPrefix returns the longest line-aligned prefix shared by two
|
||||
// breadcrumb strings. Heading hierarchies are emitted as
|
||||
// "# Top\n## Section\n### Sub", so a line-by-line comparison is sufficient
|
||||
// and avoids partial-line truncation that would corrupt the breadcrumb.
|
||||
func commonHeadingPrefix(a, b string) string {
|
||||
if a == b {
|
||||
return a
|
||||
}
|
||||
la := strings.Split(a, "\n")
|
||||
lb := strings.Split(b, "\n")
|
||||
n := len(la)
|
||||
if len(lb) < n {
|
||||
n = len(lb)
|
||||
}
|
||||
common := 0
|
||||
for i := 0; i < n; i++ {
|
||||
if la[i] != lb[i] {
|
||||
break
|
||||
}
|
||||
common = i + 1
|
||||
}
|
||||
if common == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(la[:common], "\n")
|
||||
}
|
||||
|
||||
// findHeadingBoundaries returns one boundary at offset 0 plus one per
|
||||
// Markdown heading at level <= primaryLevel that sits outside fenced code
|
||||
// blocks. Heading detection is line-oriented — a heading must occupy a
|
||||
// whole line to be recognized.
|
||||
func findHeadingBoundaries(text string, primaryLevel int) []headingBoundary {
|
||||
runes := []rune(text)
|
||||
bounds := []headingBoundary{{runeStart: 0}}
|
||||
if len(runes) == 0 {
|
||||
return bounds
|
||||
}
|
||||
|
||||
pos := 0
|
||||
inFence := false
|
||||
lines := strings.Split(text, "\n")
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++ // newline
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !inFence {
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m != nil {
|
||||
level := len(m[1])
|
||||
if level >= 1 && level <= primaryLevel && pos > 0 {
|
||||
bounds = append(bounds, headingBoundary{
|
||||
runeStart: pos,
|
||||
line: line,
|
||||
})
|
||||
}
|
||||
if level >= 1 && level <= primaryLevel && pos == 0 {
|
||||
// First line is a heading — replace the leading boundary
|
||||
bounds[0].line = line
|
||||
}
|
||||
}
|
||||
}
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++ // account for the \n that strings.Split removed
|
||||
}
|
||||
}
|
||||
return bounds
|
||||
}
|
||||
|
||||
// observeSubHeadings walks the section's lines and feeds every Markdown
|
||||
// heading deeper than primaryLevel into the hierarchy. This keeps the
|
||||
// hierarchy state correct so the breadcrumb at the next primary section
|
||||
// reflects the truly active stack.
|
||||
func observeSubHeadings(runes []rune, primaryLevel int, h *HeadingHierarchy) {
|
||||
if len(runes) == 0 {
|
||||
return
|
||||
}
|
||||
text := string(runes)
|
||||
inFence := false
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
continue
|
||||
}
|
||||
if inFence {
|
||||
continue
|
||||
}
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
level := len(m[1])
|
||||
if level > primaryLevel {
|
||||
h.Observe(line)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sectionBreadcrumb pairs a rune offset within a section with the breadcrumb
|
||||
// in effect from that offset onward.
|
||||
type sectionBreadcrumb struct {
|
||||
runeStart int
|
||||
breadcrumb string
|
||||
}
|
||||
|
||||
// sectionBreadcrumbs walks a section's deeper sub-headings (level >
|
||||
// primaryLevel) and records, for each, the rune offset where it takes effect
|
||||
// and the resulting breadcrumb. seed is the hierarchy state at the section's
|
||||
// start (already including the section heading and its ancestors). The
|
||||
// returned slice is ordered by runeStart and always begins with the seed
|
||||
// breadcrumb at offset 0, so a sub-chunk sitting far below a deep heading
|
||||
// still resolves to that heading's path rather than the section header.
|
||||
func sectionBreadcrumbs(sectionRunes []rune, primaryLevel int, seed HeadingHierarchy) []sectionBreadcrumb {
|
||||
h := seed
|
||||
result := []sectionBreadcrumb{{runeStart: 0, breadcrumb: h.BreadcrumbWithHashes()}}
|
||||
pos := 0
|
||||
inFence := false
|
||||
lines := strings.Split(string(sectionRunes), "\n")
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !inFence {
|
||||
if m := MarkdownHeadingPattern.FindStringSubmatch(line); m != nil && len(m[1]) > primaryLevel {
|
||||
h.Observe(line)
|
||||
result = append(result, sectionBreadcrumb{
|
||||
runeStart: pos,
|
||||
breadcrumb: h.BreadcrumbWithHashes(),
|
||||
})
|
||||
}
|
||||
}
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// breadcrumbAtOffset returns the breadcrumb in effect at the given rune offset
|
||||
// — the last entry whose runeStart <= offset. fallback covers the (unreachable
|
||||
// in practice) empty-slice case.
|
||||
func breadcrumbAtOffset(bcs []sectionBreadcrumb, offset int, fallback string) string {
|
||||
bc := fallback
|
||||
for _, e := range bcs {
|
||||
if e.runeStart > offset {
|
||||
break
|
||||
}
|
||||
bc = e.breadcrumb
|
||||
}
|
||||
return bc
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitByHeadings_BasicSections(t *testing.T) {
|
||||
// Each section is intentionally larger than the merge-target (≈
|
||||
// ChunkSize/2) so the post-split coalesce pass leaves them as distinct
|
||||
// chunks. We're testing per-section emission + breadcrumb here, not
|
||||
// merging.
|
||||
body := strings.Repeat("Lorem ipsum dolor sit amet consectetur adipiscing elit. ", 4)
|
||||
doc := "# Top\n" + body + "\n\n## Section A\n" + body + "\n\n## Section B\n" + body + "\n\n## Section C\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("expected ≥3 chunks (one per section), got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Breadcrumb is delivered via ContextHeader, not Content.
|
||||
for i, c := range chunks {
|
||||
if !strings.Contains(c.ContextHeader, "# Top") {
|
||||
t.Errorf("chunk %d missing H1 in ContextHeader:\n%q", i, c.ContextHeader)
|
||||
}
|
||||
// EmbeddingContent merges header + content for the embedder.
|
||||
if !strings.Contains(c.EmbeddingContent(), "# Top") {
|
||||
t.Errorf("chunk %d EmbeddingContent missing H1", i)
|
||||
}
|
||||
}
|
||||
|
||||
found := false
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.Content, "## Section B") && strings.Contains(c.Content, "Lorem ipsum") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("no chunk contains Section B with its body")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_FallsThroughForUnstructuredDoc(t *testing.T) {
|
||||
doc := "Just a plain paragraph without any headings at all in this text."
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
// no headings → falls through to SplitText, which keeps the whole thing
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("expected fallthrough single chunk, got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_LargeSectionRecursesIntoLegacy(t *testing.T) {
|
||||
body := strings.Repeat("This is a long sentence repeated many times. ", 50)
|
||||
doc := "# Top\n## Big\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Separators: []string{". "}}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("large section should be sub-split, got %d chunks", len(chunks))
|
||||
}
|
||||
// Every sub-chunk should carry the breadcrumb via ContextHeader.
|
||||
for i, c := range chunks {
|
||||
if !strings.Contains(c.ContextHeader, "# Top") {
|
||||
t.Errorf("sub-chunk %d missing H1 in ContextHeader", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_BreadcrumbReflectsLatestPath(t *testing.T) {
|
||||
// Sized so each section stays its own chunk after the tiny-section
|
||||
// coalesce pass — we're verifying breadcrumb assignment per section,
|
||||
// not the merge behavior.
|
||||
body := strings.Repeat("Lorem ipsum dolor sit amet consectetur adipiscing elit. ", 4)
|
||||
doc := "# Chapter 1\n" + body + "\n\n## Section A\n" + body + "\n\n## Section B\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("expected ≥3 chunks, got %d", len(chunks))
|
||||
}
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.Content, "text B") {
|
||||
if strings.Contains(c.ContextHeader, "## Section A") {
|
||||
t.Errorf("Section B chunk should not include Section A in breadcrumb:\n%s", c.ContextHeader)
|
||||
}
|
||||
if !strings.Contains(c.ContextHeader, "## Section B") {
|
||||
t.Errorf("Section B chunk should include its own heading in breadcrumb:\n%s", c.ContextHeader)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_IgnoresHeadingsInsideCodeFence(t *testing.T) {
|
||||
doc := "# Real\n\n```\n# Fake heading inside code\n```\n\nbody"
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.ContextHeader, "# Real") || strings.Contains(c.Content, "# Real") {
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Error("expected real H1 breadcrumb on some chunk")
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_PreservesPositionRelativeToOriginal(t *testing.T) {
|
||||
doc := "# Top\nintro\n\n## A\nbody A\n\n## B\nbody B"
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
for i, c := range chunks {
|
||||
if c.Start < 0 {
|
||||
t.Errorf("chunk %d has negative Start", i)
|
||||
}
|
||||
if c.End < c.Start {
|
||||
t.Errorf("chunk %d End < Start", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_PositionInvariant ensures End-Start == len(Content)
|
||||
// and runes[Start:End] == Content for every emitted chunk. This invariant
|
||||
// is required by knowledge.go:2278+ document reconstruction logic.
|
||||
func TestSplitByHeadings_PositionInvariant(t *testing.T) {
|
||||
doc := `# Top
|
||||
intro paragraph here.
|
||||
|
||||
## Section A
|
||||
content of A here, several sentences.
|
||||
|
||||
## Section B
|
||||
content of B here.
|
||||
|
||||
## Section C
|
||||
content of C here.`
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
docRunes := []rune(doc)
|
||||
for i, c := range chunks {
|
||||
contentRuneLen := len([]rune(c.Content))
|
||||
span := c.End - c.Start
|
||||
if span != contentRuneLen {
|
||||
t.Errorf("chunk %d: span(%d) != content_runes(%d)\nContent:\n%q", i, span, contentRuneLen, c.Content)
|
||||
}
|
||||
if c.Start >= 0 && c.End <= len(docRunes) {
|
||||
if string(docRunes[c.Start:c.End]) != c.Content {
|
||||
t.Errorf("chunk %d: runes[Start:End] != Content", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_CoalescesTinyAdjacentSections covers the FAQ /
|
||||
// install-log case where a parent heading hosts many short sub-sections.
|
||||
// Without merging, each `##` becomes its own <50-char chunk and the
|
||||
// validator rejects the tier with "too many tiny chunks". After merging,
|
||||
// they collapse into a small number of properly-sized chunks while still
|
||||
// surfacing the shared parent breadcrumb.
|
||||
func TestSplitByHeadings_CoalescesTinyAdjacentSections(t *testing.T) {
|
||||
doc := `# Install Log
|
||||
|
||||
## Docker镜像
|
||||
使用 daocloud 部署 v0.3.1。
|
||||
|
||||
## 前端老版本
|
||||
浏览器缓存了旧前端资源。
|
||||
|
||||
## 登录报错
|
||||
ERROR: column missing.
|
||||
|
||||
## 解析失败
|
||||
embedding 表缺列。`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected at least one chunk")
|
||||
}
|
||||
if len(chunks) >= 5 {
|
||||
t.Errorf("expected coalesce to produce <5 chunks, got %d", len(chunks))
|
||||
}
|
||||
// Every merged chunk must carry the shared parent in its breadcrumb so
|
||||
// retrieval can still answer "what document is this from".
|
||||
for i, c := range chunks {
|
||||
if !strings.Contains(c.ContextHeader, "# Install Log") {
|
||||
t.Errorf("chunk %d missing parent H1 in breadcrumb: %q", i, c.ContextHeader)
|
||||
}
|
||||
}
|
||||
// All four sub-section headings must remain visible somewhere in the
|
||||
// merged content (heading_splitter keeps the heading line as part of
|
||||
// each section's Content).
|
||||
for _, h := range []string{"## Docker镜像", "## 前端老版本", "## 登录报错", "## 解析失败"} {
|
||||
seen := false
|
||||
for _, c := range chunks {
|
||||
if strings.Contains(c.Content, h) {
|
||||
seen = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !seen {
|
||||
t.Errorf("merged chunks should still contain heading %q somewhere", h)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeadings_DoesNotCoalesceDistinctTopLevelHeadings(t *testing.T) {
|
||||
doc := `# Intro
|
||||
short intro.
|
||||
|
||||
# Usage
|
||||
short usage.
|
||||
|
||||
# FAQ
|
||||
short faq.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
if len(chunks) != 3 {
|
||||
t.Fatalf("expected one chunk per top-level heading, got %d:\n%v", len(chunks), chunks)
|
||||
}
|
||||
for i, heading := range []string{"# Intro", "# Usage", "# FAQ"} {
|
||||
if !strings.Contains(chunks[i].Content, heading) {
|
||||
t.Errorf("chunk %d should contain heading %q, got:\n%s", i, heading, chunks[i].Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_CoalescePreservesPositionInvariant guards the
|
||||
// End-Start == len([]rune(Content)) invariant after merging. Adjacent
|
||||
// chunks (cur.End == next.Start) must concatenate cleanly; the merge must
|
||||
// refuse to combine non-adjacent chunks (e.g. legacy sub-chunks from an
|
||||
// oversized section that overlap).
|
||||
func TestSplitByHeadings_CoalescePreservesPositionInvariant(t *testing.T) {
|
||||
doc := `# Top
|
||||
|
||||
## A
|
||||
short A.
|
||||
|
||||
## B
|
||||
short B.
|
||||
|
||||
## C
|
||||
short C.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
docRunes := []rune(doc)
|
||||
for i, c := range chunks {
|
||||
contentRuneLen := len([]rune(c.Content))
|
||||
if c.End-c.Start != contentRuneLen {
|
||||
t.Errorf("chunk %d: End-Start(%d) != content_runes(%d) after merge",
|
||||
i, c.End-c.Start, contentRuneLen)
|
||||
}
|
||||
if c.Start >= 0 && c.End <= len(docRunes) {
|
||||
if string(docRunes[c.Start:c.End]) != c.Content {
|
||||
t.Errorf("chunk %d: source[Start:End] != Content after merge", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_CoalesceRespectsChunkSize ensures the merge target
|
||||
// stays within the ChunkSize budget — the validator caps oversize chunks
|
||||
// at 2x and we should never approach that line via merging.
|
||||
func TestSplitByHeadings_CoalesceRespectsChunkSize(t *testing.T) {
|
||||
const sections = 30
|
||||
var sb strings.Builder
|
||||
sb.WriteString("# Doc\n")
|
||||
for i := 0; i < sections; i++ {
|
||||
sb.WriteString("\n## Section ")
|
||||
sb.WriteString(strings.Repeat("X", 1)) // unique-ish heading
|
||||
sb.WriteString("\nshort body line.\n")
|
||||
}
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(sb.String(), cfg, nil)
|
||||
for i, c := range chunks {
|
||||
if l := len([]rune(c.Content)); l > cfg.ChunkSize {
|
||||
t.Errorf("chunk %d exceeds ChunkSize: %d > %d", i, l, cfg.ChunkSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCommonHeadingPrefix exercises the breadcrumb-prefix helper directly.
|
||||
func TestCommonHeadingPrefix(t *testing.T) {
|
||||
cases := []struct {
|
||||
a, b, want string
|
||||
}{
|
||||
{"# Top\n## A", "# Top\n## B", "# Top"},
|
||||
{"# Top", "# Top", "# Top"},
|
||||
{"# X", "# Y", ""},
|
||||
{"# Top\n## A\n### x", "# Top\n## A\n### y", "# Top\n## A"},
|
||||
{"", "# Top", ""},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
got := commonHeadingPrefix(tc.a, tc.b)
|
||||
if got != tc.want {
|
||||
t.Errorf("commonHeadingPrefix(%q, %q) = %q, want %q", tc.a, tc.b, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_NoBreadcrumbDuplication ensures the section's own
|
||||
// heading line does not appear twice in the chunk content (once as part of
|
||||
// the breadcrumb, once as the section's first line).
|
||||
func TestSplitByHeadings_NoBreadcrumbDuplication(t *testing.T) {
|
||||
doc := `# Chapter 1
|
||||
intro.
|
||||
|
||||
## Section A
|
||||
body A.
|
||||
|
||||
## Section B
|
||||
body B.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
for i, c := range chunks {
|
||||
// Count occurrences of "## Section A" / "## Section B"
|
||||
for _, heading := range []string{"## Section A", "## Section B"} {
|
||||
n := strings.Count(c.Content, heading)
|
||||
if n > 1 {
|
||||
t.Errorf("chunk %d contains %q %d times — duplicated by breadcrumb prepend:\n%s",
|
||||
i, heading, n, c.Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitByHeadings_DeepSubHeadingInLargeSection covers issue #1674: when a
|
||||
// shallow heading level repeats and therefore becomes the dominant split
|
||||
// backbone, a long section split into sub-chunks must still carry the deep
|
||||
// `###`/`####` heading that defines a sub-chunk's meaning, not just the
|
||||
// section-level header. The marker line sits far below its deep heading, in a
|
||||
// later sub-chunk that has no heading of its own.
|
||||
func TestSplitByHeadings_DeepSubHeadingInLargeSection(t *testing.T) {
|
||||
filler := strings.Repeat("clause body sentence that pads the section out. ", 20)
|
||||
doc := "# Standard XYZ\n" +
|
||||
"## Preface\n" + filler + "\n\n" +
|
||||
"## 5 Classification\n" + filler + "\n\n" +
|
||||
"### 5.9 Grade Nine\n" + filler + "\n\n" +
|
||||
"#### 5.9.2 Clause Series\n" + filler + "\n\n" +
|
||||
"the item users search for is MARKER_ITEM_23 graded here.\n\n" +
|
||||
"## Appendix A\n" + filler + "\n\n" +
|
||||
"## Appendix B\n" + filler + "\n\n" +
|
||||
"## Appendix C\n" + filler
|
||||
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0, Separators: []string{". "}}
|
||||
chunks := splitByHeadingsImpl(doc, cfg, nil)
|
||||
|
||||
var marker *Chunk
|
||||
for i := range chunks {
|
||||
if strings.Contains(chunks[i].Content, "MARKER_ITEM_23") {
|
||||
marker = &chunks[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if marker == nil {
|
||||
t.Fatal("no chunk contains the marker line")
|
||||
}
|
||||
if !strings.Contains(marker.ContextHeader, "#### 5.9.2 Clause Series") {
|
||||
t.Errorf("marker chunk lost its deep sub-heading in the breadcrumb:\nContextHeader=%q", marker.ContextHeader)
|
||||
}
|
||||
if !strings.Contains(marker.ContextHeader, "## 5 Classification") {
|
||||
t.Errorf("marker chunk lost its section heading in the breadcrumb:\nContextHeader=%q", marker.ContextHeader)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
// Package chunker - heuristic_splitter.go implements Tier 2: boundary-driven
|
||||
// chunking for documents that lack proper Markdown headings but contain
|
||||
// recognizable structural cues (page breaks, numbered sections, multilingual
|
||||
// chapter markers, visual separators, all-caps section titles, page footers).
|
||||
//
|
||||
// The algorithm finds all candidate boundary positions, then performs greedy
|
||||
// bin-packing — accumulating blocks between boundaries into chunks until the
|
||||
// next block would exceed cfg.ChunkSize. Blocks larger than ChunkSize are
|
||||
// recursively delegated to the legacy splitter for inner segmentation.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func init() {
|
||||
splitByHeuristics = splitByHeuristicsImpl
|
||||
}
|
||||
|
||||
// boundary marks a candidate split point in the document.
|
||||
type boundary struct {
|
||||
runeStart int // rune offset where the next chunk should start
|
||||
priority int
|
||||
}
|
||||
|
||||
// splitByHeuristicsImpl is the Tier-2 implementation. Falls through to the
|
||||
// legacy splitter when no heuristic boundaries are found.
|
||||
//
|
||||
// profile is currently unused (this tier scans for boundaries directly) but
|
||||
// is accepted to keep the splitByHeadings / splitByHeuristics signatures
|
||||
// uniform — see strategy.runTier.
|
||||
func splitByHeuristicsImpl(text string, cfg SplitterConfig, _ *DocProfile) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
runes := []rune(text)
|
||||
totalRunes := len(runes)
|
||||
if totalRunes <= cfg.ChunkSize {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
bounds := findHeuristicBoundaries(text, cfg.Languages)
|
||||
// Drop any boundary that falls strictly inside a protected region (table,
|
||||
// fenced code block, LaTeX block, etc.) — splitting there would cut
|
||||
// through atomic content. Boundaries on a span edge are kept since they
|
||||
// align with the protected region start/end.
|
||||
if prot := protectedSpansRune(text, protectedSpans(text)); len(prot) > 0 {
|
||||
bounds = dropBoundsInsideSpans(bounds, prot)
|
||||
}
|
||||
if len(bounds) == 0 {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// Append a sentinel at end-of-document so the bin-packer can flush.
|
||||
bounds = append(bounds, boundary{runeStart: totalRunes})
|
||||
// Always start with a boundary at offset 0 if not already there.
|
||||
if bounds[0].runeStart != 0 {
|
||||
bounds = append([]boundary{{runeStart: 0}}, bounds...)
|
||||
}
|
||||
|
||||
// Greedy bin-packing.
|
||||
var out []Chunk
|
||||
seq := 0
|
||||
chunkStart := bounds[0].runeStart
|
||||
curEnd := chunkStart
|
||||
minChunkSize := cfg.ChunkSize / 4
|
||||
if minChunkSize < 50 {
|
||||
minChunkSize = 50
|
||||
}
|
||||
|
||||
for i := 1; i < len(bounds); i++ {
|
||||
nextEnd := bounds[i].runeStart
|
||||
blockLen := nextEnd - curEnd
|
||||
|
||||
if blockLen > cfg.ChunkSize {
|
||||
// The block between the previous and this boundary is itself too
|
||||
// large to fit in any chunk. Flush current accumulation, then
|
||||
// recursively chunk the oversize block via the legacy splitter.
|
||||
if curEnd-chunkStart > 0 {
|
||||
out = appendChunk(out, runes, chunkStart, curEnd, &seq)
|
||||
chunkStart = curEnd
|
||||
}
|
||||
out = appendOversizeBlock(out, runes, curEnd, nextEnd, cfg, &seq)
|
||||
curEnd = nextEnd
|
||||
chunkStart = nextEnd
|
||||
continue
|
||||
}
|
||||
|
||||
// Would adding this block exceed the budget?
|
||||
accumulated := nextEnd - chunkStart
|
||||
if accumulated > cfg.ChunkSize && curEnd-chunkStart >= minChunkSize {
|
||||
// Flush accumulated content as a chunk, restart at curEnd.
|
||||
out = appendChunk(out, runes, chunkStart, curEnd, &seq)
|
||||
// Snap overlap start to the nearest semantic boundary or line
|
||||
// break instead of slicing mid-line / mid-word.
|
||||
chunkStart = applyOverlapAligned(runes, curEnd, cfg.ChunkOverlap, bounds)
|
||||
}
|
||||
curEnd = nextEnd
|
||||
}
|
||||
|
||||
// Flush remaining content.
|
||||
if curEnd > chunkStart {
|
||||
out = appendChunk(out, runes, chunkStart, curEnd, &seq)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// findHeuristicBoundaries scans text and returns boundary positions in
|
||||
// ascending order. Lower-priority duplicates at the same offset are dropped.
|
||||
func findHeuristicBoundaries(text string, langs []string) []boundary {
|
||||
var bounds []boundary
|
||||
|
||||
// Form feeds — strongest single-character boundary.
|
||||
for _, idx := range allRuneIndices(text, "\f") {
|
||||
bounds = append(bounds, boundary{runeStart: idx, priority: PrioFormFeed})
|
||||
}
|
||||
|
||||
// Per-line patterns walk the text once, line by line.
|
||||
lines := strings.Split(text, "\n")
|
||||
chapterPatterns := ChapterPatternsForLangs(langs)
|
||||
pos := 0
|
||||
inFence := false
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
} else if !inFence {
|
||||
runeStart := pos
|
||||
added := false
|
||||
for _, pat := range chapterPatterns {
|
||||
if pat.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioChapterMarker})
|
||||
added = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !added && NumberedSectionPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioNumberedHead})
|
||||
added = true
|
||||
}
|
||||
if !added && AllCapsHeadingPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioAllCapsHeading})
|
||||
added = true
|
||||
}
|
||||
if !added && VisualSeparatorPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioVisualSep})
|
||||
added = true
|
||||
}
|
||||
if !added && PageFooterPattern.MatchString(line) {
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioPageFooter})
|
||||
}
|
||||
}
|
||||
pos += utf8.RuneCountInString(line)
|
||||
if i < len(lines)-1 {
|
||||
pos++ // \n
|
||||
}
|
||||
}
|
||||
|
||||
// Excessive blank blocks (\n{3,}). Match at the *start* of the run so we
|
||||
// drop into the next paragraph cleanly.
|
||||
for _, idx := range ExcessiveBlanksPattern.FindAllStringIndex(text, -1) {
|
||||
runeStart := utf8.RuneCountInString(text[:idx[1]])
|
||||
bounds = append(bounds, boundary{runeStart: runeStart, priority: PrioBlankBlock})
|
||||
}
|
||||
|
||||
if len(bounds) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Sort by position; drop near-duplicate offsets keeping the highest priority.
|
||||
sort.Slice(bounds, func(i, j int) bool {
|
||||
if bounds[i].runeStart != bounds[j].runeStart {
|
||||
return bounds[i].runeStart < bounds[j].runeStart
|
||||
}
|
||||
return bounds[i].priority > bounds[j].priority
|
||||
})
|
||||
deduped := bounds[:0]
|
||||
prev := -1
|
||||
for _, b := range bounds {
|
||||
if b.runeStart != prev {
|
||||
deduped = append(deduped, b)
|
||||
prev = b.runeStart
|
||||
}
|
||||
}
|
||||
return deduped
|
||||
}
|
||||
|
||||
// dropBoundsInsideSpans returns bounds with entries that fall strictly
|
||||
// inside any of the (rune-offset) protected spans removed. Bounds at a
|
||||
// span's start or end are kept — they align with the span edge and don't
|
||||
// split protected content. spans must be sorted by start.
|
||||
func dropBoundsInsideSpans(bounds []boundary, spans []span) []boundary {
|
||||
if len(spans) == 0 {
|
||||
return bounds
|
||||
}
|
||||
out := bounds[:0]
|
||||
boundLoop:
|
||||
for _, b := range bounds {
|
||||
for _, s := range spans {
|
||||
if s.start >= b.runeStart {
|
||||
break // remaining spans start at or after b — can't contain b
|
||||
}
|
||||
if b.runeStart < s.end {
|
||||
continue boundLoop
|
||||
}
|
||||
}
|
||||
out = append(out, b)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// allRuneIndices returns every rune offset where needle starts in text.
|
||||
// Only used for single-rune needles like form-feed.
|
||||
func allRuneIndices(text, needle string) []int {
|
||||
var out []int
|
||||
if needle == "" {
|
||||
return out
|
||||
}
|
||||
pos := 0
|
||||
for _, r := range text {
|
||||
if string(r) == needle {
|
||||
out = append(out, pos)
|
||||
}
|
||||
pos++
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// appendChunk slices runes[start:end] into a Chunk and appends it to out.
|
||||
// Pure-whitespace slices are skipped (boundary clustering can occasionally
|
||||
// produce them). The Content stored is the raw slice — Start/End rune
|
||||
// offsets must match utf8.RuneCountInString(Content) for downstream
|
||||
// reconstruction code; whitespace stripping for embedding happens in
|
||||
// Chunk.EmbeddingContent.
|
||||
func appendChunk(out []Chunk, runes []rune, start, end int, seq *int) []Chunk {
|
||||
if end <= start {
|
||||
return out
|
||||
}
|
||||
raw := string(runes[start:end])
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
return out
|
||||
}
|
||||
c := Chunk{Content: raw, Seq: *seq, Start: start, End: end}
|
||||
*seq++
|
||||
return append(out, c)
|
||||
}
|
||||
|
||||
// appendOversizeBlock recursively chunks a region that is itself larger than
|
||||
// cfg.ChunkSize, using the legacy splitter so internal length budgets and
|
||||
// protected patterns are still respected.
|
||||
func appendOversizeBlock(out []Chunk, runes []rune, start, end int, cfg SplitterConfig, seq *int) []Chunk {
|
||||
if end <= start {
|
||||
return out
|
||||
}
|
||||
subText := string(runes[start:end])
|
||||
subs := SplitText(subText, cfg)
|
||||
for _, s := range subs {
|
||||
out = append(out, Chunk{
|
||||
Content: s.Content,
|
||||
Seq: *seq,
|
||||
Start: start + s.Start,
|
||||
End: start + s.End,
|
||||
})
|
||||
*seq++
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// applyOverlapAligned returns the rune offset where the next chunk should
|
||||
// start. The target is `curEnd - overlap`, but we snap to the nearest
|
||||
// preceding boundary (within 2x overlap) or, failing that, the previous
|
||||
// newline so chunks don't begin mid-line / mid-word. Falls back to the raw
|
||||
// target only if neither option is available.
|
||||
//
|
||||
// curEnd itself is always a boundary (the bin-packer flushes at boundary
|
||||
// positions), so we exclude it from the search — picking it would yield
|
||||
// zero overlap, defeating the purpose of this function.
|
||||
func applyOverlapAligned(runes []rune, curEnd, overlap int, bounds []boundary) int {
|
||||
if overlap <= 0 {
|
||||
return curEnd
|
||||
}
|
||||
target := curEnd - overlap
|
||||
if target < 0 {
|
||||
target = 0
|
||||
}
|
||||
// Allowed search window: [curEnd - 2*overlap, curEnd)
|
||||
windowStart := curEnd - 2*overlap
|
||||
if windowStart < 0 {
|
||||
windowStart = 0
|
||||
}
|
||||
|
||||
// Prefer a semantic boundary strictly inside the window.
|
||||
bestBound := -1
|
||||
for _, b := range bounds {
|
||||
if b.runeStart >= windowStart && b.runeStart < curEnd && b.runeStart > bestBound {
|
||||
bestBound = b.runeStart
|
||||
}
|
||||
}
|
||||
if bestBound >= 0 {
|
||||
return bestBound
|
||||
}
|
||||
|
||||
// Fallback: scan backwards from `target` to the previous newline, but
|
||||
// not past windowStart so we keep the overlap roughly the right size.
|
||||
for i := target; i > windowStart && i < len(runes); i-- {
|
||||
if runes[i] == '\n' {
|
||||
return i + 1
|
||||
}
|
||||
}
|
||||
return target
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitByHeuristics_FormFeedBoundary(t *testing.T) {
|
||||
doc := strings.Repeat("page one body text. ", 30) + "\f" + strings.Repeat("page two body. ", 30)
|
||||
cfg := SplitterConfig{ChunkSize: 400, ChunkOverlap: 20, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("form feed should produce ≥2 chunks, got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_NumberedSections(t *testing.T) {
|
||||
body := strings.Repeat("body sentence. ", 8)
|
||||
doc := "1. Introduction\n" + body + "\n\n2. Methods\n" + body + "\n\n3. Results\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("numbered sections should split: got %d chunks", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_GermanChapterMarkers(t *testing.T) {
|
||||
body := strings.Repeat("Beispieltext. ", 10)
|
||||
doc := "Kapitel 1: Einführung\n" + body + "\n\nKapitel 2: Hauptteil\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("German chapter markers should split: got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_ChineseChapterMarkers(t *testing.T) {
|
||||
body := strings.Repeat("内容内容内容。", 60)
|
||||
doc := "第一章 引言\n" + body + "\n\n第二章 方法\n" + body
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{"。"}, Languages: []string{LangChinese}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("Chinese chapter markers should split: got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_FallsThroughForUnstructuredDoc(t *testing.T) {
|
||||
doc := strings.Repeat("plain prose without structure. ", 5)
|
||||
cfg := SplitterConfig{ChunkSize: 1000, ChunkOverlap: 20}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("unstructured short doc should be one chunk, got %d", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_OversizeBlockRecursesIntoLegacy(t *testing.T) {
|
||||
huge := strings.Repeat("This is a long sentence. ", 200) // ~5000 chars
|
||||
doc := "1. Intro\n" + huge
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 50, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 5 {
|
||||
t.Errorf("oversize block should produce many sub-chunks, got %d", len(chunks))
|
||||
}
|
||||
// No single chunk should massively exceed the budget.
|
||||
for i, c := range chunks {
|
||||
if len([]rune(c.Content)) > 2*cfg.ChunkSize {
|
||||
t.Errorf("chunk %d exceeds 2x size: %d runes", i, len([]rune(c.Content)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_BoundariesAreOrdered(t *testing.T) {
|
||||
doc := "Kapitel 1: A\nbody\n\n---\n\n2. Section B\nbody\n\nPage 3 of 10\n\n第三章 C\nbody"
|
||||
bounds := findHeuristicBoundaries(doc, nil)
|
||||
if len(bounds) < 2 {
|
||||
t.Fatalf("expected multiple boundaries, got %d", len(bounds))
|
||||
}
|
||||
for i := 1; i < len(bounds); i++ {
|
||||
if bounds[i].runeStart < bounds[i-1].runeStart {
|
||||
t.Errorf("bounds not sorted: %d before %d", bounds[i].runeStart, bounds[i-1].runeStart)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitByHeuristics_EmptyText(t *testing.T) {
|
||||
if got := splitByHeuristicsImpl("", DefaultConfig(), nil); got != nil {
|
||||
t.Errorf("empty doc should be nil, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: applyOverlapAligned previously included curEnd itself in its
|
||||
// boundary search, and curEnd is always one of the bounds (the bin-packer
|
||||
// flushes at boundary positions). The function therefore always returned
|
||||
// curEnd, producing zero overlap regardless of cfg.ChunkOverlap.
|
||||
func TestSplitByHeuristics_OverlapActuallyOverlaps(t *testing.T) {
|
||||
// Build many small numbered sections so the bin-packer flushes mid-doc
|
||||
// with at least one earlier boundary inside the overlap window.
|
||||
var sb strings.Builder
|
||||
for i := 1; i <= 12; i++ {
|
||||
sb.WriteString("\n\n")
|
||||
sb.WriteByte(byte('0' + i%10))
|
||||
sb.WriteString(". ")
|
||||
sb.WriteString(strings.Repeat("alpha beta gamma. ", 4)) // ~72 chars / section
|
||||
}
|
||||
doc := sb.String()
|
||||
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 80, Separators: []string{". "}}
|
||||
chunks := splitByHeuristicsImpl(doc, cfg, nil)
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("need >=2 chunks to test overlap, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// At least one consecutive chunk pair must share a non-trivial suffix /
|
||||
// prefix. We don't require *every* pair to overlap (oversize blocks
|
||||
// short-circuit through legacy and reset chunkStart), but at least one
|
||||
// regular flush boundary should produce real overlap.
|
||||
saw := false
|
||||
for i := 1; i < len(chunks); i++ {
|
||||
prev := strings.TrimSpace(chunks[i-1].Content)
|
||||
cur := strings.TrimSpace(chunks[i].Content)
|
||||
// Walk back the longest suffix of prev that prefixes cur.
|
||||
match := 0
|
||||
maxScan := len(prev)
|
||||
if len(cur) < maxScan {
|
||||
maxScan = len(cur)
|
||||
}
|
||||
for n := 1; n <= maxScan; n++ {
|
||||
if strings.HasPrefix(cur, prev[len(prev)-n:]) {
|
||||
match = n
|
||||
}
|
||||
}
|
||||
if match >= 20 {
|
||||
saw = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !saw {
|
||||
t.Fatalf("expected at least one chunk pair to overlap by >=20 chars, none did. chunk sizes: %v",
|
||||
chunkLengths(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// Heuristic boundaries that fall inside protected regions (LaTeX block,
|
||||
// table, link, etc.) must be dropped so the bin-packer doesn't break
|
||||
// atomic content. Without the protected-span filter, a numbered-section
|
||||
// looking line inside a $$...$$ math block would be picked as a boundary.
|
||||
func TestSplitByHeuristics_DropsBoundariesInsideProtectedSpans(t *testing.T) {
|
||||
body := strings.Repeat("filler. ", 30)
|
||||
// LaTeX block whose middle line matches NumberedSectionPattern. The
|
||||
// filter should drop that boundary so the math block stays intact.
|
||||
doc := body + "\n\n$$\nx = 1\n1. equation step one\ny = 2\n$$\n\n" + body
|
||||
|
||||
bounds := findHeuristicBoundaries(doc, nil)
|
||||
prot := protectedSpansRune(doc, protectedSpans(doc))
|
||||
if len(prot) == 0 {
|
||||
t.Fatalf("expected protected spans for doc, got none")
|
||||
}
|
||||
filtered := dropBoundsInsideSpans(bounds, prot)
|
||||
for _, b := range filtered {
|
||||
for _, s := range prot {
|
||||
if b.runeStart > s.start && b.runeStart < s.end {
|
||||
t.Errorf("boundary %d still inside protected span [%d,%d)", b.runeStart, s.start, s.end)
|
||||
}
|
||||
}
|
||||
}
|
||||
// And it should actually have removed at least one boundary.
|
||||
if len(filtered) >= len(bounds) {
|
||||
t.Errorf("filter removed nothing: before=%d after=%d", len(bounds), len(filtered))
|
||||
}
|
||||
}
|
||||
|
||||
func chunkLengths(chunks []Chunk) []int {
|
||||
out := make([]int, len(chunks))
|
||||
for i, c := range chunks {
|
||||
out[i] = len([]rune(c.Content))
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Package chunker - patterns.go is the source of truth for multilingual
|
||||
// regex patterns used by the heading-aware and heuristic splitters.
|
||||
//
|
||||
// Patterns are grouped by purpose (chapter markers, numbering, separators)
|
||||
// and tagged with a priority that the heuristic splitter uses to rank
|
||||
// candidate chunk boundaries.
|
||||
package chunker
|
||||
|
||||
import "regexp"
|
||||
|
||||
// BoundaryPriority levels for heuristic chunk boundaries. Higher = stronger.
|
||||
const (
|
||||
PrioFormFeed = 100
|
||||
PrioNumberedHead = 90
|
||||
PrioChapterMarker = 85
|
||||
PrioAllCapsHeading = 70
|
||||
PrioVisualSep = 60
|
||||
PrioPageFooter = 50
|
||||
PrioBlankBlock = 40
|
||||
)
|
||||
|
||||
// MarkdownHeadingPattern matches an ATX-style Markdown heading at line start.
|
||||
// Capture groups: (1) hashes, (2) heading text.
|
||||
var MarkdownHeadingPattern = regexp.MustCompile(`(?m)^(#{1,6})\s+(.+?)\s*#*\s*$`)
|
||||
|
||||
// FormFeedPattern matches the form-feed control character used by some PDF
|
||||
// converters as a page break marker.
|
||||
var FormFeedPattern = regexp.MustCompile(`\f`)
|
||||
|
||||
// NumberedSectionPattern matches lines starting with numeric or roman numbering
|
||||
// followed by a non-empty title, e.g. "1. Intro", "2.3 Methods", "IV. Results",
|
||||
// "2.2.1 用户与权限". The trailing dot after a multi-level numeral is optional
|
||||
// because many technical documents write "1.1 Foo" without a closing dot.
|
||||
var NumberedSectionPattern = regexp.MustCompile(`(?m)^[ \t]*(?:\d+(?:\.\d+){1,3}\.?|(?:\d+|[IVX]{1,5})\.)[ \t]+\S.{0,200}$`)
|
||||
|
||||
// AllCapsHeadingPattern matches short all-caps lines (likely section titles
|
||||
// rendered without Markdown headings). It requires at least 4 letters and
|
||||
// up to ~10 words. Trailing colons are tolerated.
|
||||
var AllCapsHeadingPattern = regexp.MustCompile(`(?m)^[ \t]*([A-ZÄÖÜ][A-ZÄÖÜ \-]{3,80}):?\s*$`)
|
||||
|
||||
// VisualSeparatorPattern matches horizontal rules / divider lines used as
|
||||
// section separators in plain text or pre-Markdown documents.
|
||||
var VisualSeparatorPattern = regexp.MustCompile(`(?m)^[ \t]*(?:-{3,}|={3,}|\*{3,}|_{3,})[ \t]*$`)
|
||||
|
||||
// ExcessiveBlanksPattern matches three or more consecutive newlines, which
|
||||
// usually denote a hard section break.
|
||||
var ExcessiveBlanksPattern = regexp.MustCompile(`\n{3,}`)
|
||||
|
||||
// PageFooterPattern matches typical "Seite X von Y" / "Page X of Y" lines.
|
||||
var PageFooterPattern = regexp.MustCompile(`(?mi)^[ \t]*(?:Seite|Page|页码?)\s+\d+(?:\s*(?:von|of|/)\s*\d+)?[ \t]*$`)
|
||||
|
||||
// GermanChapterPattern matches German chapter / section markers.
|
||||
var GermanChapterPattern = regexp.MustCompile(`(?m)^[ \t]*(?:Kapitel|Abschnitt|Teil)\s+(?:[0-9]+|[IVX]{1,5})[\.: ].{0,200}$`)
|
||||
|
||||
// EnglishChapterPattern matches English chapter / section markers.
|
||||
var EnglishChapterPattern = regexp.MustCompile(`(?m)^[ \t]*(?:Chapter|Section|Part)\s+(?:[0-9]+|[IVX]{1,5})[\.: ].{0,200}$`)
|
||||
|
||||
// ChineseChapterPattern matches CJK chapter / section markers like 第一章,
|
||||
// 第3节, 第 1 章 (whitespace between 第 / numeral / unit is tolerated).
|
||||
var ChineseChapterPattern = regexp.MustCompile(`(?m)^[ \t]*第[ \t]*[一二三四五六七八九十百千零〇0-9]+[ \t]*(?:章|节|節|部分|篇)[ \t]?.{0,200}$`)
|
||||
|
||||
// SentenceSeparators returns sentence-level separators tuned for the language.
|
||||
// Used for fine-grained sub-splitting when a section is still too large.
|
||||
func SentenceSeparators(lang string) []string {
|
||||
switch lang {
|
||||
case LangChinese:
|
||||
return []string{"。", "!", "?", ";", "\n"}
|
||||
case LangGerman, LangEnglish:
|
||||
return []string{". ", "! ", "? ", "; ", "\n"}
|
||||
default:
|
||||
return []string{"。", "!", "?", ";", ". ", "! ", "? ", "; ", "\n"}
|
||||
}
|
||||
}
|
||||
|
||||
// ChapterPatternsForLangs returns the chapter-marker regexes that apply for
|
||||
// the given language hints. An empty / unknown list returns all of them so
|
||||
// that auto-detected documents still match.
|
||||
func ChapterPatternsForLangs(langs []string) []*regexp.Regexp {
|
||||
if len(langs) == 0 {
|
||||
return []*regexp.Regexp{GermanChapterPattern, EnglishChapterPattern, ChineseChapterPattern}
|
||||
}
|
||||
var out []*regexp.Regexp
|
||||
for _, l := range langs {
|
||||
switch l {
|
||||
case LangGerman:
|
||||
out = append(out, GermanChapterPattern)
|
||||
case LangEnglish:
|
||||
out = append(out, EnglishChapterPattern)
|
||||
case LangChinese:
|
||||
out = append(out, ChineseChapterPattern)
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
out = []*regexp.Regexp{GermanChapterPattern, EnglishChapterPattern, ChineseChapterPattern}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestMarkdownHeadingPattern_BasicLevels(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"# Heading 1", true},
|
||||
{"## Heading 2", true},
|
||||
{"###### Heading 6", true},
|
||||
{"####### Too many", false},
|
||||
{"#NoSpace", false},
|
||||
{" # Indented", false},
|
||||
{"plain text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := MarkdownHeadingPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("MarkdownHeadingPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNumberedSectionPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"1. Introduction", true},
|
||||
{"2.3 Methodology", true}, // multi-level numerals no longer require trailing dot
|
||||
{"2.3. Methodology", true},
|
||||
{"2.2.1 用户与权限", true}, // three-level numbering, Chinese title, no trailing dot
|
||||
{"3.2.1 单机 Docker Compose", true},
|
||||
{"IV. Results", true},
|
||||
{"1.Introduction", false}, // requires whitespace
|
||||
{"1.", false}, // requires title
|
||||
{"1.1", false}, // requires title after numeral
|
||||
{"1 NoDotSingleLevel", false}, // single-level numerals still need trailing dot
|
||||
{"1.2.3.4.5 TooDeep", false}, // more than 3 sub-levels not accepted
|
||||
{"plain text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := NumberedSectionPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("NumberedSectionPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGermanChapterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"Kapitel 1: Einführung", true},
|
||||
{"Abschnitt 2.3 Methodik", true},
|
||||
{"Abschnitt 3 Methodik", true},
|
||||
{"Teil II Ergebnisse", true},
|
||||
{"chapter 1", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := GermanChapterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("GermanChapterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnglishChapterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"Chapter 1: Intro", true},
|
||||
{"Section 5 Methods", true},
|
||||
{"Part IV Results", true},
|
||||
{"Kapitel 1", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := EnglishChapterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("EnglishChapterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChineseChapterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"第一章 引言", true},
|
||||
{"第3节 方法论", true},
|
||||
{"第二部分 结果", true},
|
||||
{"第 1 章 引言", true}, // space between 第 / numeral / unit
|
||||
{"第 一 章 引言", true},
|
||||
{"第1章 引言", true},
|
||||
{"Chapter 1", false},
|
||||
{"第 章 空数字", false}, // missing numeral
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := ChineseChapterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("ChineseChapterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVisualSeparatorPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"---", true},
|
||||
{"========", true},
|
||||
{"***", true},
|
||||
{"____", true},
|
||||
{"--", false}, // needs at least 3
|
||||
{"-- text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := VisualSeparatorPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("VisualSeparatorPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPageFooterPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"Seite 3 von 24", true},
|
||||
{"Page 5 of 12", true},
|
||||
{"page 7", true},
|
||||
{"Seite 9", true},
|
||||
{"页 3", true},
|
||||
{"页码 3 / 12", true},
|
||||
{"Some text", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := PageFooterPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("PageFooterPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllCapsHeadingPattern(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
match bool
|
||||
}{
|
||||
{"INTRODUCTION", true},
|
||||
{"METHODS AND MATERIALS", true},
|
||||
{"Mixed Case Heading", false},
|
||||
{"ABC", false}, // <4 chars
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := AllCapsHeadingPattern.MatchString(c.in)
|
||||
if got != c.match {
|
||||
t.Errorf("AllCapsHeadingPattern(%q): got %v want %v", c.in, got, c.match)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSentenceSeparators(t *testing.T) {
|
||||
if got := SentenceSeparators(LangChinese); got[0] != "。" {
|
||||
t.Errorf("Chinese should start with 。, got %v", got)
|
||||
}
|
||||
if got := SentenceSeparators(LangEnglish); got[0] != ". " {
|
||||
t.Errorf("English should start with '. ', got %v", got)
|
||||
}
|
||||
if got := SentenceSeparators("xx"); len(got) < 5 {
|
||||
t.Errorf("Unknown lang should return mixed (>=5 separators), got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChapterPatternsForLangs(t *testing.T) {
|
||||
if got := ChapterPatternsForLangs(nil); len(got) != 3 {
|
||||
t.Errorf("nil langs should return all 3, got %d", len(got))
|
||||
}
|
||||
if got := ChapterPatternsForLangs([]string{LangGerman}); len(got) != 1 {
|
||||
t.Errorf("only DE requested should return 1, got %d", len(got))
|
||||
}
|
||||
if got := ChapterPatternsForLangs([]string{LangGerman, LangChinese}); len(got) != 2 {
|
||||
t.Errorf("DE+ZH requested should return 2, got %d", len(got))
|
||||
}
|
||||
if got := ChapterPatternsForLangs([]string{"xx"}); len(got) != 3 {
|
||||
t.Errorf("unknown lang should fall back to all 3, got %d", len(got))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
// Package chunker - profiler.go scans a document once to gather structure
|
||||
// indicators that drive strategy selection (heading-aware vs. heuristic vs.
|
||||
// recursive). Profiling is cheap (a few regex passes plus rune counting)
|
||||
// and runs before any chunking decision is made.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// DocProfile holds the document-level signals used to choose a chunking tier.
|
||||
//
|
||||
// The JSON shape (snake_case via struct tags) is part of the public preview
|
||||
// endpoint API — keep field names stable. Internal callers should use the Go
|
||||
// field names; only the preview handler relies on the wire format.
|
||||
type DocProfile struct {
|
||||
TotalChars int `json:"total_chars"`
|
||||
TotalLines int `json:"total_lines"`
|
||||
AvgLineLen float64 `json:"avg_line_len"`
|
||||
StdLineLen float64 `json:"std_line_len"`
|
||||
|
||||
// Markdown structure
|
||||
MdHeadingCounts map[int]int `json:"md_heading_counts"` // level (1..6) → count
|
||||
MdHeadingTotal int `json:"md_heading_total"`
|
||||
|
||||
// Heuristic indicators
|
||||
NumberedSectionCount int `json:"numbered_section_count"`
|
||||
AllCapsShortLineCount int `json:"all_caps_short_line_count"`
|
||||
BlankParagraphBreaks int `json:"blank_paragraph_breaks"`
|
||||
FormFeedCount int `json:"form_feed_count"`
|
||||
VisualSepCount int `json:"visual_sep_count"`
|
||||
GermanChapterCount int `json:"german_chapter_count"`
|
||||
EnglishChapterCount int `json:"english_chapter_count"`
|
||||
ChineseChapterCount int `json:"chinese_chapter_count"`
|
||||
RepeatedFooterCount int `json:"repeated_footer_count"`
|
||||
|
||||
// Content characteristics
|
||||
HasTables bool `json:"has_tables"`
|
||||
HasCode bool `json:"has_code"`
|
||||
CodeRatio float64 `json:"code_ratio"`
|
||||
|
||||
// Detected language hints (best-effort)
|
||||
DetectedLangs []string `json:"detected_langs"`
|
||||
}
|
||||
|
||||
// HeadingDensity returns the share of lines that are Markdown headings.
|
||||
func (p *DocProfile) HeadingDensity() float64 {
|
||||
if p.TotalLines == 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(p.MdHeadingTotal) / float64(p.TotalLines)
|
||||
}
|
||||
|
||||
// DominantHeadingLevel returns the heading level (1..6) that should drive
|
||||
// section splitting. Preference order:
|
||||
// 1. The lowest level (closest to root) that has at least 3 occurrences —
|
||||
// a "real" structural backbone of the document.
|
||||
// 2. Otherwise the deepest level present at least once — gives finer-grained
|
||||
// boundaries for small documents that just have an H1 + a few H2s.
|
||||
//
|
||||
// Returns 0 when no Markdown headings exist.
|
||||
func (p *DocProfile) DominantHeadingLevel() int {
|
||||
if p.MdHeadingTotal == 0 {
|
||||
return 0
|
||||
}
|
||||
for level := 1; level <= 6; level++ {
|
||||
if p.MdHeadingCounts[level] >= 3 {
|
||||
return level
|
||||
}
|
||||
}
|
||||
for level := 6; level >= 1; level-- {
|
||||
if p.MdHeadingCounts[level] > 0 {
|
||||
return level
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// HeuristicMarkerTotal sums the non-Markdown structural markers.
|
||||
func (p *DocProfile) HeuristicMarkerTotal() int {
|
||||
return p.NumberedSectionCount +
|
||||
p.GermanChapterCount + p.EnglishChapterCount + p.ChineseChapterCount +
|
||||
p.AllCapsShortLineCount + p.VisualSepCount + p.FormFeedCount
|
||||
}
|
||||
|
||||
// ProfileDocument runs a single pass over text and returns its profile.
|
||||
func ProfileDocument(text string) *DocProfile {
|
||||
p := &DocProfile{
|
||||
MdHeadingCounts: make(map[int]int),
|
||||
}
|
||||
if text == "" {
|
||||
return p
|
||||
}
|
||||
|
||||
p.TotalChars = len([]rune(text))
|
||||
p.FormFeedCount = strings.Count(text, "\f")
|
||||
|
||||
lines := strings.Split(text, "\n")
|
||||
p.TotalLines = len(lines)
|
||||
|
||||
// First pass: per-line markers and length stats
|
||||
var lengths []float64
|
||||
inFence := false
|
||||
codeChars := 0
|
||||
for _, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
|
||||
// Toggle fenced-code state. We use a 3-backtick prefix detector here
|
||||
// rather than a full regex so we don't have to fight with the
|
||||
// protected-pattern logic later.
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
inFence = !inFence
|
||||
p.HasCode = true
|
||||
continue
|
||||
}
|
||||
if inFence {
|
||||
codeChars += len([]rune(line))
|
||||
continue
|
||||
}
|
||||
|
||||
runeLen := len([]rune(line))
|
||||
lengths = append(lengths, float64(runeLen))
|
||||
|
||||
if matchHeading(line, &p.MdHeadingCounts) {
|
||||
p.MdHeadingTotal++
|
||||
continue
|
||||
}
|
||||
if NumberedSectionPattern.MatchString(line) {
|
||||
p.NumberedSectionCount++
|
||||
}
|
||||
if GermanChapterPattern.MatchString(line) {
|
||||
p.GermanChapterCount++
|
||||
}
|
||||
if EnglishChapterPattern.MatchString(line) {
|
||||
p.EnglishChapterCount++
|
||||
}
|
||||
if ChineseChapterPattern.MatchString(line) {
|
||||
p.ChineseChapterCount++
|
||||
}
|
||||
if AllCapsHeadingPattern.MatchString(line) {
|
||||
p.AllCapsShortLineCount++
|
||||
}
|
||||
if VisualSeparatorPattern.MatchString(line) {
|
||||
p.VisualSepCount++
|
||||
}
|
||||
if PageFooterPattern.MatchString(line) {
|
||||
p.RepeatedFooterCount++
|
||||
}
|
||||
if strings.HasPrefix(trimmed, "|") && strings.HasSuffix(trimmed, "|") {
|
||||
p.HasTables = true
|
||||
}
|
||||
}
|
||||
|
||||
if len(lengths) > 0 {
|
||||
var sum float64
|
||||
for _, l := range lengths {
|
||||
sum += l
|
||||
}
|
||||
p.AvgLineLen = sum / float64(len(lengths))
|
||||
var variance float64
|
||||
for _, l := range lengths {
|
||||
d := l - p.AvgLineLen
|
||||
variance += d * d
|
||||
}
|
||||
variance /= float64(len(lengths))
|
||||
p.StdLineLen = math.Sqrt(variance)
|
||||
}
|
||||
|
||||
if p.TotalChars > 0 {
|
||||
p.CodeRatio = float64(codeChars) / float64(p.TotalChars)
|
||||
}
|
||||
|
||||
p.BlankParagraphBreaks = strings.Count(text, "\n\n\n")
|
||||
|
||||
// Sample a slice of the document for language detection — avoids paying
|
||||
// O(N) scan cost on huge inputs while still giving a stable signal.
|
||||
sample := text
|
||||
if len(sample) > 4096 {
|
||||
sample = sample[:4096]
|
||||
}
|
||||
lang := DetectLanguage(sample)
|
||||
p.DetectedLangs = []string{lang}
|
||||
if lang == LangMixed {
|
||||
// Provide all three for downstream pattern selection.
|
||||
p.DetectedLangs = []string{LangEnglish, LangGerman, LangChinese}
|
||||
}
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
// matchHeading checks whether line is an ATX heading and increments the
|
||||
// appropriate level counter when so. Returns true on match.
|
||||
func matchHeading(line string, counts *map[int]int) bool {
|
||||
m := MarkdownHeadingPattern.FindStringSubmatch(line)
|
||||
if m == nil {
|
||||
return false
|
||||
}
|
||||
level := len(m[1])
|
||||
if level < 1 || level > 6 {
|
||||
return false
|
||||
}
|
||||
(*counts)[level]++
|
||||
return true
|
||||
}
|
||||
|
||||
// StrategyTier identifies which chunking implementation should run.
|
||||
type StrategyTier string
|
||||
|
||||
const (
|
||||
TierHeading StrategyTier = "heading"
|
||||
TierHeuristic StrategyTier = "heuristic"
|
||||
TierLegacy StrategyTier = "legacy"
|
||||
)
|
||||
|
||||
// SelectStrategy returns the ordered tier chain to attempt for this document.
|
||||
// The first tier is the primary choice; subsequent tiers are fallbacks if
|
||||
// validation rejects the previous output. The "legacy" tier is appended as
|
||||
// a final safety net so callers always receive at least one chunk-set.
|
||||
func SelectStrategy(p *DocProfile) []StrategyTier {
|
||||
if p == nil {
|
||||
return []StrategyTier{TierLegacy}
|
||||
}
|
||||
var chain []StrategyTier
|
||||
|
||||
// Tier 1 candidate: Markdown heading-aware
|
||||
if p.MdHeadingTotal >= 3 && p.HeadingDensity() > 0.005 && p.DominantHeadingLevel() > 0 {
|
||||
chain = append(chain, TierHeading)
|
||||
}
|
||||
|
||||
// Tier 2 candidate: heuristic boundary detection
|
||||
if p.HeuristicMarkerTotal() >= 5 || p.FormFeedCount > 0 ||
|
||||
p.GermanChapterCount+p.EnglishChapterCount+p.ChineseChapterCount > 0 {
|
||||
chain = append(chain, TierHeuristic)
|
||||
}
|
||||
|
||||
// Legacy is the ultimate fallback: always returns chunks even when
|
||||
// validation fails, so callers never get an empty result.
|
||||
chain = append(chain, TierLegacy)
|
||||
return chain
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestProfileDocument_Empty(t *testing.T) {
|
||||
p := ProfileDocument("")
|
||||
if p.TotalChars != 0 || p.TotalLines != 0 {
|
||||
t.Errorf("empty doc should have zero stats, got %+v", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_MarkdownHeadings(t *testing.T) {
|
||||
doc := `# Title
|
||||
Some intro text here.
|
||||
|
||||
## Section 1
|
||||
Body of section 1.
|
||||
|
||||
## Section 2
|
||||
Body of section 2.
|
||||
|
||||
### Subsection 2.1
|
||||
Detail.
|
||||
|
||||
## Section 3
|
||||
More body.`
|
||||
p := ProfileDocument(doc)
|
||||
if p.MdHeadingCounts[1] != 1 {
|
||||
t.Errorf("expected 1 H1, got %d", p.MdHeadingCounts[1])
|
||||
}
|
||||
if p.MdHeadingCounts[2] != 3 {
|
||||
t.Errorf("expected 3 H2, got %d", p.MdHeadingCounts[2])
|
||||
}
|
||||
if p.MdHeadingCounts[3] != 1 {
|
||||
t.Errorf("expected 1 H3, got %d", p.MdHeadingCounts[3])
|
||||
}
|
||||
if p.MdHeadingTotal != 5 {
|
||||
t.Errorf("expected 5 headings total, got %d", p.MdHeadingTotal)
|
||||
}
|
||||
if p.DominantHeadingLevel() != 2 {
|
||||
t.Errorf("dominant level should be 2 (≥3 occurrences), got %d", p.DominantHeadingLevel())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_DominantLevelFallback(t *testing.T) {
|
||||
// No level reaches 3 occurrences — should fall back to most frequent.
|
||||
doc := "# Single H1\n## H2 a\n## H2 b\n"
|
||||
p := ProfileDocument(doc)
|
||||
if p.DominantHeadingLevel() != 2 {
|
||||
t.Errorf("expected fallback to level 2 (most frequent), got %d", p.DominantHeadingLevel())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_NumberedSections(t *testing.T) {
|
||||
doc := `1. Introduction
|
||||
text
|
||||
|
||||
2. Methodology
|
||||
text
|
||||
|
||||
3. Results
|
||||
text`
|
||||
p := ProfileDocument(doc)
|
||||
if p.NumberedSectionCount < 3 {
|
||||
t.Errorf("expected ≥3 numbered sections, got %d", p.NumberedSectionCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_GermanChapters(t *testing.T) {
|
||||
doc := "Kapitel 1: Einführung\n\nText\n\nKapitel 2: Hauptteil\n\nText"
|
||||
p := ProfileDocument(doc)
|
||||
if p.GermanChapterCount != 2 {
|
||||
t.Errorf("expected 2 German chapters, got %d", p.GermanChapterCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_ChineseChapters(t *testing.T) {
|
||||
doc := "第一章 引言\n\n内容\n\n第二章 方法\n\n内容"
|
||||
p := ProfileDocument(doc)
|
||||
if p.ChineseChapterCount != 2 {
|
||||
t.Errorf("expected 2 Chinese chapters, got %d", p.ChineseChapterCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_FormFeed(t *testing.T) {
|
||||
doc := "page 1 content\f\npage 2 content\f\npage 3 content"
|
||||
p := ProfileDocument(doc)
|
||||
if p.FormFeedCount != 2 {
|
||||
t.Errorf("expected 2 form feeds, got %d", p.FormFeedCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_DetectsCodeBlock(t *testing.T) {
|
||||
doc := "Some prose.\n\n```go\nfunc main() {}\n```\n\nMore prose."
|
||||
p := ProfileDocument(doc)
|
||||
if !p.HasCode {
|
||||
t.Error("expected HasCode=true for fenced block")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_DetectsTable(t *testing.T) {
|
||||
doc := "Intro.\n\n| col a | col b |\n| --- | --- |\n| 1 | 2 |\n"
|
||||
p := ProfileDocument(doc)
|
||||
if !p.HasTables {
|
||||
t.Error("expected HasTables=true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProfileDocument_LineStatistics(t *testing.T) {
|
||||
doc := "short\nthis is a longer line of text\nanother line here"
|
||||
p := ProfileDocument(doc)
|
||||
if p.TotalLines != 3 {
|
||||
t.Errorf("expected 3 lines, got %d", p.TotalLines)
|
||||
}
|
||||
if p.AvgLineLen <= 0 {
|
||||
t.Error("expected positive avg line len")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_HeadingDoc(t *testing.T) {
|
||||
doc := "# A\nbody\n## B\nbody\n## C\nbody\n## D\nbody"
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
if chain[0] != TierHeading {
|
||||
t.Errorf("expected heading tier first, got %v", chain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_HeuristicDoc(t *testing.T) {
|
||||
doc := strings.Repeat("Kapitel 1: Foo\nbody body body\n\n", 1) +
|
||||
strings.Repeat("Kapitel 2: Bar\nbody body body\n\n", 1)
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
// no markdown headings → heuristic must come first (heading tier skipped)
|
||||
if chain[0] != TierHeuristic {
|
||||
t.Errorf("expected heuristic tier first, got %v", chain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_PlainDoc(t *testing.T) {
|
||||
doc := "just a paragraph of plain text without any structure indicators at all here"
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
if chain[0] != TierLegacy {
|
||||
t.Errorf("expected legacy tier first for unstructured doc, got %v", chain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectStrategy_AlwaysFallsBackToLegacy(t *testing.T) {
|
||||
for _, doc := range []string{"", "simple", "# H1\nbody"} {
|
||||
p := ProfileDocument(doc)
|
||||
chain := SelectStrategy(p)
|
||||
if chain[len(chain)-1] != TierLegacy {
|
||||
t.Errorf("chain must end with legacy, got %v for doc=%q", chain, doc)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,914 @@
|
||||
// Package chunker implements text splitting for document chunking.
|
||||
//
|
||||
// Ported from the Python docreader/splitter/splitter.py recursive text splitter.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser"
|
||||
)
|
||||
|
||||
// Chunk represents a piece of split text with position tracking.
|
||||
//
|
||||
// Content holds exactly the text from the original document between Start
|
||||
// and End (rune offsets), so End-Start == utf8.RuneCountInString(Content).
|
||||
// This invariant is relied on by document-reconstruction code paths
|
||||
// (knowledge.go:2278+ for summary generation, UI highlighting, etc.).
|
||||
//
|
||||
// ContextHeader is a separately-tracked context string (e.g. a Markdown
|
||||
// heading breadcrumb) that should be prepended at embedding/retrieval time
|
||||
// but is NOT part of Content. Keeping the two apart preserves the
|
||||
// position invariant while still letting embedding pipelines see the
|
||||
// section context.
|
||||
type Chunk struct {
|
||||
Content string
|
||||
ContextHeader string
|
||||
Seq int
|
||||
Start int
|
||||
End int
|
||||
}
|
||||
|
||||
// EmbeddingContent returns the text that should be fed to the embedding
|
||||
// model — the ContextHeader prepended (when set) plus the chunk content.
|
||||
// Use this where Content alone would lose semantic context (Tier-1 chunks).
|
||||
//
|
||||
// Content is returned verbatim from the source document (the End-Start
|
||||
// rune-count invariant requires that), but for embedding we trim the
|
||||
// surrounding whitespace so leading/trailing newlines from boundary slices
|
||||
// don't dilute the embedded vector or waste tokens. Inner whitespace is
|
||||
// preserved.
|
||||
func (c Chunk) EmbeddingContent() string {
|
||||
body := strings.TrimSpace(c.Content)
|
||||
if c.ContextHeader == "" {
|
||||
return body
|
||||
}
|
||||
return c.ContextHeader + "\n\n" + body
|
||||
}
|
||||
|
||||
// ImageRef is an image reference found within a chunk's content.
|
||||
type ImageRef struct {
|
||||
OriginalRef string
|
||||
AltText string
|
||||
Start int // offset within the chunk content
|
||||
End int
|
||||
}
|
||||
|
||||
// SplitterConfig configures the text splitter. Strategy and TokenLimit are
|
||||
// honored by the strategy entry point in strategy.go; the legacy SplitText
|
||||
// path uses only ChunkSize/Overlap/Separators.
|
||||
type SplitterConfig struct {
|
||||
ChunkSize int
|
||||
ChunkOverlap int
|
||||
Separators []string
|
||||
|
||||
// Strategy selects an adaptive tier. Empty = legacy (backwards-compatible).
|
||||
// See strategy.go for valid values.
|
||||
Strategy string
|
||||
// TokenLimit caps chunk size in approximate tokens. 0 = use ChunkSize chars.
|
||||
TokenLimit int
|
||||
// Languages hints multilingual heuristic patterns. Empty = auto-detect.
|
||||
Languages []string
|
||||
}
|
||||
|
||||
// Default chunk sizing constants. Single source of truth for the entire
|
||||
// chunker package and (via knowledge.go::buildSplitterConfig) the
|
||||
// knowledge service. The frontend KnowledgeBaseEditorModal mirrors these
|
||||
// numbers in its initial form state — keep them in sync if you change
|
||||
// either value here.
|
||||
//
|
||||
// DefaultChunkSize = 512 chars: ~100–130 English tokens / ~300 Chinese
|
||||
// tokens. Validated as a strong baseline by the Vecta Feb-2026 benchmark
|
||||
// across 50 academic papers. Use 200–400 for FAQ-style atomic content,
|
||||
// 1000–2000 for narrative / argumentative documents.
|
||||
//
|
||||
// DefaultChunkOverlap = 80 chars (≈15% of DefaultChunkSize): community-
|
||||
// recommended sweet spot between recall (an answer split across a
|
||||
// boundary needs overlap to be retrievable) and storage cost. Use 0 for
|
||||
// strictly atomic data (FAQ, JSON records), 150–200 for long narratives
|
||||
// where reasoning crosses chunks.
|
||||
//
|
||||
// MIGRATION NOTE: Prior versions had three different overlap defaults
|
||||
// (Go DefaultConfig: 64, knowledge.go buildSplitterConfig: 50, Python
|
||||
// docreader: 100). All consolidated to 80 here.
|
||||
//
|
||||
// Existing knowledge bases that stored ChunkOverlap=0 in the DB pick
|
||||
// this 80 up on next re-index; their previously-indexed embeddings will
|
||||
// not match new ones bit-for-bit. Recall stays similar but search
|
||||
// ranking can shift slightly. To freeze the old behavior on a per-KB
|
||||
// basis, explicitly set ChunkingConfig.ChunkOverlap to 64 before
|
||||
// re-indexing.
|
||||
const (
|
||||
DefaultChunkSize = 512
|
||||
DefaultChunkOverlap = 80
|
||||
)
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
func DefaultConfig() SplitterConfig {
|
||||
return SplitterConfig{
|
||||
ChunkSize: DefaultChunkSize,
|
||||
ChunkOverlap: DefaultChunkOverlap,
|
||||
Separators: []string{"\n\n", "\n", "。"},
|
||||
}
|
||||
}
|
||||
|
||||
// protectedPatterns are regex patterns for content that must not be split.
|
||||
var protectedPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?s)\$\$.*?\$\$`), // LaTeX block math
|
||||
regexp.MustCompile(`!\[[^\]]*\]\([^)]+\)`), // Markdown images
|
||||
regexp.MustCompile(`\[[^\]]*\]\([^)]+\)`), // Markdown links
|
||||
regexp.MustCompile("(?m)[ ]*(?:\\|[^|\\n]*)+\\|[\\r\\n]+\\s*(?:\\|\\s*:?-{3,}:?\\s*)+\\|[\\r\\n]+"), // Table header+separator
|
||||
regexp.MustCompile("(?m)[ ]*(?:\\|[^|\\n]*)+\\|[\\r\\n]+"), // Table rows
|
||||
regexp.MustCompile("(?s)```(?:\\w+)?[\\r\\n].*?```"), // Fenced code blocks
|
||||
regexp.MustCompile("`[^`\\r\\n]+`"), // Markdown inline code
|
||||
}
|
||||
|
||||
type span struct {
|
||||
start, end int
|
||||
}
|
||||
|
||||
// protectedSpansRune converts byte-offset protected spans to rune offsets
|
||||
// in a single forward pass over text. Used by callers that work in rune
|
||||
// space (e.g. the heuristic splitter) to avoid choosing chunk boundaries
|
||||
// that cut through protected content. byteSpans must be sorted by start
|
||||
// (protectedSpans guarantees this).
|
||||
func protectedSpansRune(text string, byteSpans []span) []span {
|
||||
if len(byteSpans) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]span, 0, len(byteSpans))
|
||||
runeIdx := 0
|
||||
byteIdx := 0
|
||||
for _, s := range byteSpans {
|
||||
for byteIdx < s.start && byteIdx < len(text) {
|
||||
_, size := utf8.DecodeRuneInString(text[byteIdx:])
|
||||
byteIdx += size
|
||||
runeIdx++
|
||||
}
|
||||
startRune := runeIdx
|
||||
for byteIdx < s.end && byteIdx < len(text) {
|
||||
_, size := utf8.DecodeRuneInString(text[byteIdx:])
|
||||
byteIdx += size
|
||||
runeIdx++
|
||||
}
|
||||
out = append(out, span{start: startRune, end: runeIdx})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// protectedSpans finds all non-overlapping protected regions in text.
|
||||
func protectedSpans(text string) []span {
|
||||
type match struct {
|
||||
start, end int
|
||||
}
|
||||
var all []match
|
||||
for _, pat := range protectedPatterns {
|
||||
locs := pat.FindAllStringIndex(text, -1)
|
||||
for _, loc := range locs {
|
||||
if loc[1]-loc[0] > 0 {
|
||||
all = append(all, match{loc[0], loc[1]})
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(all) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Sort by start, then by length descending
|
||||
for i := 1; i < len(all); i++ {
|
||||
for j := i; j > 0; j-- {
|
||||
if all[j].start < all[j-1].start ||
|
||||
(all[j].start == all[j-1].start && (all[j].end-all[j].start) > (all[j-1].end-all[j-1].start)) {
|
||||
all[j], all[j-1] = all[j-1], all[j]
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove overlaps
|
||||
var result []span
|
||||
lastEnd := 0
|
||||
for _, m := range all {
|
||||
if m.start >= lastEnd {
|
||||
result = append(result, span{m.start, m.end})
|
||||
lastEnd = m.end
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// splitUnit is a piece of text with its original position.
|
||||
type splitUnit struct {
|
||||
text string
|
||||
start, end int
|
||||
}
|
||||
|
||||
// splitBySeparators splits text by separators in priority order, recursively
|
||||
// applying the next separator to any piece that is still larger than
|
||||
// chunkSize. Mirrors the recursive priority semantics of the Python
|
||||
// reference splitter (docreader/splitter/splitter.py:_split): if `\n\n`
|
||||
// produces a piece that's still too big, `\n` (and subsequent separators)
|
||||
// are applied within that piece — not to the whole text.
|
||||
//
|
||||
// chunkSize == 0 disables the recursion guard; callers that don't care
|
||||
// about size budget (e.g. a final mergeUnits-style pass) pass 0.
|
||||
func splitBySeparators(text string, separators []string, chunkSize int) []string {
|
||||
if text == "" || len(separators) == 0 {
|
||||
return []string{text}
|
||||
}
|
||||
if chunkSize > 0 && runeLen(text) <= chunkSize {
|
||||
return []string{text}
|
||||
}
|
||||
|
||||
for i, sep := range separators {
|
||||
if sep == "" {
|
||||
continue
|
||||
}
|
||||
re := regexp.MustCompile("(" + regexp.QuoteMeta(sep) + ")")
|
||||
splits := re.Split(text, -1)
|
||||
matches := re.FindAllString(text, -1)
|
||||
if len(matches) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var pieces []string
|
||||
for j, s := range splits {
|
||||
if s != "" {
|
||||
pieces = append(pieces, s)
|
||||
}
|
||||
if j < len(matches) && matches[j] != "" {
|
||||
pieces = append(pieces, matches[j])
|
||||
}
|
||||
}
|
||||
if len(pieces) <= 1 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Recursively split any piece that is still too large with the
|
||||
// remaining (lower-priority) separators.
|
||||
var out []string
|
||||
remaining := separators[i+1:]
|
||||
for _, p := range pieces {
|
||||
if chunkSize > 0 && runeLen(p) > chunkSize && len(remaining) > 0 {
|
||||
out = append(out, splitBySeparators(p, remaining, chunkSize)...)
|
||||
} else {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
return []string{text}
|
||||
}
|
||||
|
||||
// runeLen returns the number of runes in s.
|
||||
func runeLen(s string) int {
|
||||
return utf8.RuneCountInString(s)
|
||||
}
|
||||
|
||||
// SplitText splits text into chunks with overlap, respecting protected patterns.
|
||||
func SplitText(text string, cfg SplitterConfig) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
chunkSize := cfg.ChunkSize
|
||||
chunkOverlap := cfg.ChunkOverlap
|
||||
separators := cfg.Separators
|
||||
|
||||
if chunkSize <= 0 {
|
||||
chunkSize = 512
|
||||
}
|
||||
if chunkOverlap < 0 {
|
||||
chunkOverlap = 0
|
||||
}
|
||||
|
||||
// Step 1: Find protected spans
|
||||
protected := protectedSpans(text)
|
||||
|
||||
// Step 2: Split non-protected regions by separators, keep protected as atomic units.
|
||||
// chunkSize is forwarded so splitBySeparators can recursively apply lower-priority
|
||||
// separators to oversize pieces (Python-parity recursive split).
|
||||
units := buildUnitsWithProtection(text, protected, separators, chunkSize)
|
||||
|
||||
// Step 3: Merge units into chunks with overlap
|
||||
return mergeUnits(units, chunkSize, chunkOverlap)
|
||||
}
|
||||
|
||||
// buildUnitsWithProtection splits text into units, preserving protected spans as atomic.
|
||||
// Start/End positions in the returned units are rune offsets (not byte offsets),
|
||||
// because downstream merge logic indexes content via []rune slicing.
|
||||
// If a protected span exceeds maxProtectedSize, it will be forcibly split to prevent
|
||||
// creating chunks that are too large for downstream processing (e.g., embedding APIs).
|
||||
// chunkSize is forwarded to splitBySeparators so recursive splitting can keep pieces
|
||||
// under the budget when one separator alone leaves a piece oversize.
|
||||
func buildUnitsWithProtection(text string, protected []span, separators []string, chunkSize int) []splitUnit {
|
||||
const maxProtectedSize = 7500 // Maximum size for a protected unit (留余量给标题等)
|
||||
|
||||
var units []splitUnit
|
||||
bytePos := 0
|
||||
runePos := 0
|
||||
|
||||
for _, p := range protected {
|
||||
if p.start > bytePos {
|
||||
pre := text[bytePos:p.start]
|
||||
parts := splitBySeparators(pre, separators, chunkSize)
|
||||
runeOffset := runePos
|
||||
for _, part := range parts {
|
||||
partRuneLen := runeLen(part)
|
||||
units = append(units, splitUnit{
|
||||
text: part,
|
||||
start: runeOffset,
|
||||
end: runeOffset + partRuneLen,
|
||||
})
|
||||
runeOffset += partRuneLen
|
||||
}
|
||||
runePos += runeLen(pre)
|
||||
}
|
||||
|
||||
protText := text[p.start:p.end]
|
||||
protRuneLen := runeLen(protText)
|
||||
|
||||
// If protected content is too large, forcibly split it
|
||||
if protRuneLen > maxProtectedSize {
|
||||
// Split into smaller chunks at line breaks or spaces
|
||||
runes := []rune(protText)
|
||||
offset := 0
|
||||
for offset < len(runes) {
|
||||
chunkEnd := offset + maxProtectedSize
|
||||
if chunkEnd > len(runes) {
|
||||
chunkEnd = len(runes)
|
||||
} else {
|
||||
// Try to break at a newline or space
|
||||
for i := chunkEnd - 1; i > offset && i > chunkEnd-200; i-- {
|
||||
if runes[i] == '\n' || runes[i] == ' ' {
|
||||
chunkEnd = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
chunkText := string(runes[offset:chunkEnd])
|
||||
chunkLen := chunkEnd - offset
|
||||
units = append(units, splitUnit{
|
||||
text: chunkText,
|
||||
start: runePos + offset,
|
||||
end: runePos + offset + chunkLen,
|
||||
})
|
||||
offset = chunkEnd
|
||||
}
|
||||
} else {
|
||||
// Normal case: keep protected content as a single unit
|
||||
units = append(units, splitUnit{
|
||||
text: protText,
|
||||
start: runePos,
|
||||
end: runePos + protRuneLen,
|
||||
})
|
||||
}
|
||||
runePos += protRuneLen
|
||||
bytePos = p.end
|
||||
}
|
||||
|
||||
if bytePos < len(text) {
|
||||
remaining := text[bytePos:]
|
||||
parts := splitBySeparators(remaining, separators, chunkSize)
|
||||
runeOffset := runePos
|
||||
for _, part := range parts {
|
||||
partRuneLen := runeLen(part)
|
||||
units = append(units, splitUnit{
|
||||
text: part,
|
||||
start: runeOffset,
|
||||
end: runeOffset + partRuneLen,
|
||||
})
|
||||
runeOffset += partRuneLen
|
||||
}
|
||||
}
|
||||
|
||||
return units
|
||||
}
|
||||
|
||||
// mergeUnits combines split units into chunks with overlap tracking.
|
||||
// Enforces an absolute maximum chunk size to prevent exceeding downstream limits (e.g., embedding APIs).
|
||||
// Active contextual headers (e.g., Markdown table headers) are prepended to new
|
||||
// chunks so that every chunk carries its own header context.
|
||||
func mergeUnits(units []splitUnit, chunkSize, chunkOverlap int) []Chunk {
|
||||
if len(units) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
const absoluteMaxSize = 7500
|
||||
|
||||
ht := newHeaderTracker()
|
||||
|
||||
var chunks []Chunk
|
||||
var current []splitUnit
|
||||
curLen := 0
|
||||
|
||||
for _, u := range units {
|
||||
uLen := runeLen(u.text)
|
||||
|
||||
// If this single unit exceeds absolute max, force split it further
|
||||
if uLen > absoluteMaxSize {
|
||||
// Flush current chunk if any
|
||||
if len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
current = nil
|
||||
curLen = 0
|
||||
}
|
||||
|
||||
// Update header state even for oversized units
|
||||
ht.update(u.text)
|
||||
|
||||
// Split this oversized unit into smaller chunks
|
||||
runes := []rune(u.text)
|
||||
offset := 0
|
||||
for offset < len(runes) {
|
||||
chunkEnd := offset + absoluteMaxSize
|
||||
if chunkEnd > len(runes) {
|
||||
chunkEnd = len(runes)
|
||||
} else {
|
||||
for i := chunkEnd - 1; i > offset && i > chunkEnd-200; i-- {
|
||||
if runes[i] == '\n' || runes[i] == ' ' {
|
||||
chunkEnd = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
chunkText := string(runes[offset:chunkEnd])
|
||||
chunks = append(chunks, Chunk{
|
||||
Content: chunkText,
|
||||
Seq: len(chunks),
|
||||
Start: u.start + offset,
|
||||
End: u.start + chunkEnd,
|
||||
})
|
||||
offset = chunkEnd
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Update header tracking
|
||||
ht.update(u.text)
|
||||
// Flush at table boundary so the next table is not merged into a chunk
|
||||
// that still carries the previous table's prepended header context.
|
||||
if ht.headerEndedThisUnit && len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
current = nil
|
||||
curLen = 0
|
||||
}
|
||||
headers := ht.getHeaders()
|
||||
headersLen := runeLen(headers)
|
||||
if headersLen > chunkSize {
|
||||
headers = ""
|
||||
headersLen = 0
|
||||
}
|
||||
|
||||
// If adding this unit (plus reserving space for headers in a potential
|
||||
// next chunk) would exceed chunk size, flush the current chunk.
|
||||
if curLen+uLen+headersLen > chunkSize && len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
|
||||
// Keep overlap from the end of current
|
||||
current, curLen = computeOverlap(current, chunkOverlap, chunkSize, uLen)
|
||||
|
||||
// Shrink overlap further if needed to fit headers + next unit
|
||||
if headers != "" && headersLen+uLen <= chunkSize {
|
||||
for len(current) > 0 && curLen+uLen+headersLen > chunkSize {
|
||||
curLen -= runeLen(current[0].text)
|
||||
current = current[1:]
|
||||
}
|
||||
|
||||
// Prepend headers if the column-name context is not already present
|
||||
// in the overlap or the next unit being added.
|
||||
overlapText := unitsText(current)
|
||||
if !headerAlreadyPresent(headers, overlapText, u.text) &&
|
||||
!headerColumnMismatch(headers, u.text) {
|
||||
startPos := u.start
|
||||
if len(current) > 0 {
|
||||
startPos = current[0].start
|
||||
}
|
||||
hUnit := splitUnit{text: headers, start: startPos, end: startPos}
|
||||
current = append([]splitUnit{hUnit}, current...)
|
||||
curLen += headersLen
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if adding this unit would exceed absolute max
|
||||
if curLen+uLen > absoluteMaxSize {
|
||||
if len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
current = nil
|
||||
curLen = 0
|
||||
}
|
||||
}
|
||||
|
||||
current = append(current, u)
|
||||
curLen += uLen
|
||||
}
|
||||
|
||||
// Flush remaining
|
||||
if len(current) > 0 {
|
||||
chunks = append(chunks, buildChunk(current, len(chunks)))
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// unitsText concatenates the text of all units.
|
||||
func unitsText(units []splitUnit) string {
|
||||
var sb strings.Builder
|
||||
for _, u := range units {
|
||||
sb.WriteString(u.text)
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// headerAlreadyPresent returns true if the column-name row from the header
|
||||
// is already present in the overlap or the next unit, preventing duplication.
|
||||
func headerAlreadyPresent(headers, overlapText, unitText string) bool {
|
||||
// Fast path: full header already in overlap or unit
|
||||
if strings.Contains(overlapText, headers) || strings.Contains(unitText, headers) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Extract the column-name row (first meaningful non-separator line).
|
||||
// For a rewritten header like "| col1 | col2 |\n| --- | --- |\n",
|
||||
// the first line is the column names.
|
||||
colRow := headerColumnRow(headers)
|
||||
if colRow == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
return strings.Contains(overlapText, colRow) || strings.Contains(unitText, colRow)
|
||||
}
|
||||
|
||||
// headerColumnRow extracts the column-name line from a header string.
|
||||
// Returns empty string if the header has no meaningful column names.
|
||||
func headerColumnRow(header string) string {
|
||||
for _, line := range strings.Split(header, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.Contains(line, "---") {
|
||||
continue
|
||||
}
|
||||
// Skip lines that are only pipes/whitespace (empty header rows)
|
||||
onlyPipes := true
|
||||
for _, r := range line {
|
||||
if r != '|' && r != ' ' && r != '\t' {
|
||||
onlyPipes = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if !onlyPipes {
|
||||
return line
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func buildChunk(units []splitUnit, seq int) Chunk {
|
||||
var sb strings.Builder
|
||||
for _, u := range units {
|
||||
sb.WriteString(u.text)
|
||||
}
|
||||
return Chunk{
|
||||
Content: sb.String(),
|
||||
Seq: seq,
|
||||
Start: units[0].start,
|
||||
End: units[len(units)-1].end,
|
||||
}
|
||||
}
|
||||
|
||||
// computeOverlap returns the semantic suffix to keep for overlap and its total
|
||||
// rune length.
|
||||
//
|
||||
// The configured overlap is a hard upper bound, not a raw character slice.
|
||||
// Boundary detection may inspect up to four additional runes immediately
|
||||
// before that tail window so a separator cut by the window boundary remains
|
||||
// visible. A candidate is eligible only when its last rune is at position -1
|
||||
// relative to the original window, or later, so retained content never exceeds
|
||||
// the overlap limit. Eligible boundaries use this priority:
|
||||
//
|
||||
// 1. paragraph break (\n\n / \r\n\r\n)
|
||||
// 2. line break (\n / \r\n)
|
||||
// 3. sentence end (。?!, ". ", "? ", "! ")
|
||||
//
|
||||
// Priority wins first; for boundaries with the same priority, the earliest
|
||||
// one in the window wins so the useful overlap is as large as possible. The
|
||||
// next chunk starts immediately after the selected separator. If the window
|
||||
// has no valid semantic boundary, no overlap is retained. This keeps overlap
|
||||
// deterministic and bounded while allowing a boundary inside a large split
|
||||
// unit (the previous implementation could only retain whole units and often
|
||||
// collapsed to zero overlap for ordinary paragraphs).
|
||||
const semanticOverlapLookbehind = 4 // rune length of the longest separator: \r\n\r\n
|
||||
|
||||
func computeOverlap(current []splitUnit, chunkOverlap, chunkSize, nextLen int) ([]splitUnit, int) {
|
||||
if chunkOverlap <= 0 {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
// Overlap is part of the next chunk's total size. When the incoming unit
|
||||
// is already close to the chunk budget, shrink the search window rather
|
||||
// than creating an oversized chunk.
|
||||
maxOverlap := chunkOverlap
|
||||
if remaining := chunkSize - nextLen; remaining < maxOverlap {
|
||||
maxOverlap = remaining
|
||||
}
|
||||
if maxOverlap <= 0 {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
window := semanticOverlapWindow(current, maxOverlap+semanticOverlapLookbehind)
|
||||
if len(window) == 0 {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
windowText := unitsText(window)
|
||||
// Coordinates before originalWindowStart are the lookbehind region. With
|
||||
// end-exclusive offsets, requiring boundaryEnd >= originalWindowStart is
|
||||
// equivalent to requiring the separator's final rune to be at -1 or later.
|
||||
originalWindowStart := runeLen(windowText) - maxOverlap
|
||||
if originalWindowStart < 0 {
|
||||
originalWindowStart = 0
|
||||
}
|
||||
boundaryEnd, ok := findSemanticOverlapBoundaryEndingAtOrAfter(windowText, originalWindowStart)
|
||||
if !ok {
|
||||
return nil, 0
|
||||
}
|
||||
|
||||
overlap := trimUnitsPrefix(window, boundaryEnd)
|
||||
overlapLen := 0
|
||||
for _, u := range overlap {
|
||||
overlapLen += runeLen(u.text)
|
||||
}
|
||||
if overlapLen <= 0 || overlapLen > maxOverlap || strings.TrimSpace(unitsText(overlap)) == "" {
|
||||
return nil, 0
|
||||
}
|
||||
return overlap, overlapLen
|
||||
}
|
||||
|
||||
// semanticOverlapWindow returns at most maxLen source-backed runes from the
|
||||
// tail of current. It may slice inside the first retained splitUnit so a
|
||||
// semantic boundary inside a large paragraph remains discoverable. Synthetic
|
||||
// zero-width units (for example repeated table headers) form a hard barrier:
|
||||
// crossing them would break the Start/End-to-Content source invariant.
|
||||
func semanticOverlapWindow(current []splitUnit, maxLen int) []splitUnit {
|
||||
if maxLen <= 0 || len(current) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
remaining := maxLen
|
||||
reversed := make([]splitUnit, 0, len(current))
|
||||
for i := len(current) - 1; i >= 0 && remaining > 0; i-- {
|
||||
u := current[i]
|
||||
uLen := runeLen(u.text)
|
||||
if uLen == 0 {
|
||||
continue
|
||||
}
|
||||
// Header markers contain generated text but occupy no source range.
|
||||
// Do not include or cross them when deriving an overlap suffix.
|
||||
if u.start == u.end || u.end-u.start != uLen {
|
||||
break
|
||||
}
|
||||
|
||||
if uLen <= remaining {
|
||||
reversed = append(reversed, u)
|
||||
remaining -= uLen
|
||||
continue
|
||||
}
|
||||
|
||||
runes := []rune(u.text)
|
||||
start := uLen - remaining
|
||||
reversed = append(reversed, splitUnit{
|
||||
text: string(runes[start:]),
|
||||
start: u.start + start,
|
||||
end: u.end,
|
||||
})
|
||||
remaining = 0
|
||||
}
|
||||
|
||||
if len(reversed) == 0 {
|
||||
return nil
|
||||
}
|
||||
window := make([]splitUnit, len(reversed))
|
||||
for i := range reversed {
|
||||
window[len(reversed)-1-i] = reversed[i]
|
||||
}
|
||||
return window
|
||||
}
|
||||
|
||||
type semanticOverlapBoundary struct {
|
||||
start int
|
||||
end int
|
||||
priority int
|
||||
}
|
||||
|
||||
// findSemanticOverlapBoundary returns the rune offset immediately after the
|
||||
// selected separator. Boundaries inside protected Markdown/code/math regions
|
||||
// are ignored, and a boundary is invalid when only whitespace follows it.
|
||||
func findSemanticOverlapBoundary(text string) (int, bool) {
|
||||
return findSemanticOverlapBoundaryEndingAtOrAfter(text, 0)
|
||||
}
|
||||
|
||||
// findSemanticOverlapBoundaryEndingAtOrAfter applies the normal priority and
|
||||
// earliest-position rules only to candidates whose end-exclusive rune offset
|
||||
// is at least minEnd. Filtering before comparison prevents an earlier but
|
||||
// ineligible lookbehind separator from hiding a later eligible boundary.
|
||||
func findSemanticOverlapBoundaryEndingAtOrAfter(text string, minEnd int) (int, bool) {
|
||||
runes := []rune(text)
|
||||
if len(runes) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
if minEnd < 0 {
|
||||
minEnd = 0
|
||||
}
|
||||
if minEnd > len(runes) {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
protected := protectedSpansRune(text, protectedSpans(text))
|
||||
insideProtected := func(pos int) bool {
|
||||
for _, p := range protected {
|
||||
if pos < p.start {
|
||||
return false
|
||||
}
|
||||
if pos >= p.start && pos < p.end {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
hasMeaningfulTail := func(end int) bool {
|
||||
return end >= 0 && end < len(runes) && strings.TrimSpace(string(runes[end:])) != ""
|
||||
}
|
||||
|
||||
var best semanticOverlapBoundary
|
||||
found := false
|
||||
consider := func(start, end, priority int) {
|
||||
if start < 0 || end <= start || end < minEnd || end > len(runes) ||
|
||||
insideProtected(start) || !hasMeaningfulTail(end) {
|
||||
return
|
||||
}
|
||||
candidate := semanticOverlapBoundary{start: start, end: end, priority: priority}
|
||||
if !found || candidate.priority < best.priority ||
|
||||
(candidate.priority == best.priority && candidate.start < best.start) {
|
||||
best = candidate
|
||||
found = true
|
||||
}
|
||||
}
|
||||
|
||||
// Mark paragraph-break runes so their component newlines are not also
|
||||
// emitted as lower-priority line-break candidates.
|
||||
paragraphRune := make([]bool, len(runes))
|
||||
for i := 0; i < len(runes); i++ {
|
||||
switch {
|
||||
case i+3 < len(runes) && runes[i] == '\r' && runes[i+1] == '\n' &&
|
||||
runes[i+2] == '\r' && runes[i+3] == '\n':
|
||||
consider(i, i+4, 1)
|
||||
for j := i; j < i+4; j++ {
|
||||
paragraphRune[j] = true
|
||||
}
|
||||
i += 3
|
||||
case i+1 < len(runes) && runes[i] == '\n' && runes[i+1] == '\n':
|
||||
consider(i, i+2, 1)
|
||||
paragraphRune[i], paragraphRune[i+1] = true, true
|
||||
i++
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < len(runes); i++ {
|
||||
if paragraphRune[i] {
|
||||
continue
|
||||
}
|
||||
if runes[i] == '\r' && i+1 < len(runes) && runes[i+1] == '\n' && !paragraphRune[i+1] {
|
||||
consider(i, i+2, 2)
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if runes[i] == '\n' {
|
||||
consider(i, i+1, 2)
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < len(runes); i++ {
|
||||
switch runes[i] {
|
||||
case '。', '?', '!':
|
||||
consider(i, i+1, 3)
|
||||
case '.', '?', '!':
|
||||
if i+1 < len(runes) && runes[i+1] == ' ' {
|
||||
consider(i, i+2, 3)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !found {
|
||||
return 0, false
|
||||
}
|
||||
return best.end, true
|
||||
}
|
||||
|
||||
// trimUnitsPrefix removes prefixLen source runes while preserving the
|
||||
// remaining units' original positions. prefixLen is relative to unitsText.
|
||||
func trimUnitsPrefix(units []splitUnit, prefixLen int) []splitUnit {
|
||||
if prefixLen <= 0 {
|
||||
out := make([]splitUnit, len(units))
|
||||
copy(out, units)
|
||||
return out
|
||||
}
|
||||
|
||||
remaining := prefixLen
|
||||
out := make([]splitUnit, 0, len(units))
|
||||
for _, u := range units {
|
||||
uLen := runeLen(u.text)
|
||||
if remaining >= uLen {
|
||||
remaining -= uLen
|
||||
continue
|
||||
}
|
||||
if remaining > 0 {
|
||||
runes := []rune(u.text)
|
||||
u.text = string(runes[remaining:])
|
||||
u.start += remaining
|
||||
remaining = 0
|
||||
}
|
||||
out = append(out, u)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ParentChildResult holds the two-level chunking output.
|
||||
// Parent chunks provide context (large window), child chunks are used for
|
||||
// embedding/retrieval (small window). Each child carries its ParentIndex so
|
||||
// the caller can wire up ParentChunkID after DB insertion.
|
||||
type ParentChildResult struct {
|
||||
Parents []Chunk
|
||||
Children []ChildChunk
|
||||
}
|
||||
|
||||
// ChildChunk extends Chunk with a reference to its parent.
|
||||
type ChildChunk struct {
|
||||
Chunk
|
||||
ParentIndex int // index into ParentChildResult.Parents
|
||||
}
|
||||
|
||||
// SplitTextParentChild performs two-level chunking:
|
||||
// 1. Split text into large parent chunks (parentCfg).
|
||||
// 2. Split each parent into smaller child chunks (childCfg) for embedding.
|
||||
//
|
||||
// The child Seq is globally unique across the entire document.
|
||||
func SplitTextParentChild(text string, parentCfg, childCfg SplitterConfig) ParentChildResult {
|
||||
parents := SplitText(text, parentCfg)
|
||||
if len(parents) == 0 {
|
||||
return ParentChildResult{}
|
||||
}
|
||||
|
||||
var newParents []Chunk
|
||||
var children []ChildChunk
|
||||
childSeq := 0
|
||||
for _, parent := range parents {
|
||||
subs := SplitText(parent.Content, childCfg)
|
||||
|
||||
parentIndex := -1
|
||||
if len(subs) > 1 || (len(subs) == 1 && subs[0].Content != parent.Content) {
|
||||
parentIndex = len(newParents)
|
||||
newParents = append(newParents, parent)
|
||||
}
|
||||
|
||||
for _, sub := range subs {
|
||||
// Adjust offsets: sub positions are relative to parent content,
|
||||
// shift to document-level offsets.
|
||||
// Use additive shift (not Content-length based) so that chunks with
|
||||
// prepended context headers keep correct positional tracking.
|
||||
sub.Seq = childSeq
|
||||
sub.Start += parent.Start
|
||||
sub.End += parent.Start
|
||||
children = append(children, ChildChunk{
|
||||
Chunk: sub,
|
||||
ParentIndex: parentIndex,
|
||||
})
|
||||
childSeq++
|
||||
}
|
||||
}
|
||||
return ParentChildResult{Parents: newParents, Children: children}
|
||||
}
|
||||
|
||||
// ExtractImageRefs extracts markdown image references from text.
|
||||
// The URL group supports one level of balanced parentheses so that URLs
|
||||
// like https://example.com/item_(abc)/123 are captured in full.
|
||||
var imageRefPattern = regexp.MustCompile(`!\[([^\]]*)\]\(([^()\s]*(?:\([^)]*\)[^()\s]*)*)\)`)
|
||||
|
||||
func ExtractImageRefs(text string) []ImageRef {
|
||||
text = docparser.UnwrapLinkedImages(text)
|
||||
matches := imageRefPattern.FindAllStringSubmatchIndex(text, -1)
|
||||
var refs []ImageRef
|
||||
for _, m := range matches {
|
||||
refs = append(refs, ImageRef{
|
||||
OriginalRef: text[m[4]:m[5]], // group 2: URL
|
||||
AltText: text[m[2]:m[3]], // group 1: alt text
|
||||
Start: m[0],
|
||||
End: m[1],
|
||||
})
|
||||
}
|
||||
return refs
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,363 @@
|
||||
// Package chunker - strategy.go is the public entry point for adaptive
|
||||
// chunking. Callers invoke Split / SplitParentChild instead of the legacy
|
||||
// SplitText / SplitTextParentChild functions; the strategy resolver picks
|
||||
// a tier based on document profile and the SplitterConfig.Strategy hint.
|
||||
//
|
||||
// The legacy entry points still exist in splitter.go for backwards
|
||||
// compatibility — strategy.go simply layers a tier-selector on top.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// Strategy values for SplitterConfig.Strategy.
|
||||
const (
|
||||
StrategyAuto = "auto"
|
||||
StrategyHeading = "heading"
|
||||
StrategyHeuristic = "heuristic"
|
||||
StrategyRecursive = "recursive"
|
||||
StrategyLegacy = "legacy"
|
||||
)
|
||||
|
||||
// Split chunks text using the strategy configured in cfg. When cfg.Strategy
|
||||
// is empty or "auto" the document profiler picks the tier. The function
|
||||
// always returns a non-nil result: on tier failure the chain falls through
|
||||
// to the legacy splitter, which is the original Tier 3 implementation.
|
||||
//
|
||||
// Hot path: avoids the Diagnostics struct allocation that
|
||||
// SplitWithDiagnostics performs (matters in SplitParentChild where
|
||||
// Split is called once per parent).
|
||||
func Split(text string, cfg SplitterConfig) []Chunk {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
cfg = ensureDefaults(cfg)
|
||||
chain, profile := resolveChainWithProfile(text, cfg)
|
||||
totalChars := len([]rune(text))
|
||||
|
||||
var lastOut []Chunk
|
||||
for i, tier := range chain {
|
||||
out := runTier(tier, text, cfg, profile)
|
||||
if v := ValidateChunks(out, totalChars, cfg.ChunkSize); v.OK {
|
||||
return out
|
||||
} else {
|
||||
logger.Debugf(context.Background(), "chunker: tier %s rejected: %s", tier, v.Reason)
|
||||
}
|
||||
if tier == TierLegacy && i == len(chain)-1 {
|
||||
lastOut = out
|
||||
}
|
||||
}
|
||||
if lastOut != nil {
|
||||
return lastOut
|
||||
}
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// TierRejection records why a tier was rejected by the validator and the
|
||||
// chain advanced to the next tier. Surfaced by SplitWithDiagnostics for
|
||||
// the debug/preview endpoint.
|
||||
type TierRejection struct {
|
||||
Tier StrategyTier `json:"tier"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
// Diagnostics captures which tier produced the returned chunks plus the
|
||||
// chain that was attempted, any rejected tiers along the way, and the
|
||||
// document profile that drove tier selection.
|
||||
//
|
||||
// Useful for surfacing in a debug UI; not produced by the normal Split
|
||||
// path. The JSON shape is part of the public preview-endpoint API — keep
|
||||
// field names stable.
|
||||
type Diagnostics struct {
|
||||
SelectedTier StrategyTier `json:"selected_tier"`
|
||||
TierChain []StrategyTier `json:"tier_chain"`
|
||||
Rejected []TierRejection `json:"rejected"`
|
||||
// Profile is set when the auto strategy resolved the chain via the
|
||||
// document profiler. nil when an explicit Strategy bypassed profiling.
|
||||
Profile *DocProfile `json:"profile,omitempty"`
|
||||
}
|
||||
|
||||
// SplitWithDiagnostics is the same as Split but also returns the
|
||||
// diagnostic trace (selected tier, full chain, rejection reasons,
|
||||
// profile when available). Use this for the chunker preview endpoint
|
||||
// where the caller wants to know which tier won and why others lost.
|
||||
func SplitWithDiagnostics(text string, cfg SplitterConfig) ([]Chunk, *Diagnostics) {
|
||||
// Default selected tier to legacy so an empty diag never carries the
|
||||
// zero string — that would render as a blank tag in the debug UI.
|
||||
diag := &Diagnostics{SelectedTier: TierLegacy}
|
||||
if text == "" {
|
||||
return nil, diag
|
||||
}
|
||||
cfg = ensureDefaults(cfg)
|
||||
chain, profile := resolveChainWithProfile(text, cfg)
|
||||
diag.TierChain = chain
|
||||
diag.Profile = profile
|
||||
totalChars := len([]rune(text))
|
||||
|
||||
var lastOut []Chunk
|
||||
var lastTier StrategyTier
|
||||
for i, tier := range chain {
|
||||
out := runTier(tier, text, cfg, profile)
|
||||
v := ValidateChunks(out, totalChars, cfg.ChunkSize)
|
||||
if v.OK {
|
||||
diag.SelectedTier = tier
|
||||
return out, diag
|
||||
}
|
||||
diag.Rejected = append(diag.Rejected, TierRejection{Tier: tier, Reason: v.Reason})
|
||||
logger.Debugf(context.Background(), "chunker: tier %s rejected: %s", tier, v.Reason)
|
||||
if tier == TierLegacy && i == len(chain)-1 {
|
||||
lastOut = out
|
||||
lastTier = tier
|
||||
}
|
||||
}
|
||||
if lastOut != nil {
|
||||
diag.SelectedTier = lastTier
|
||||
return lastOut, diag
|
||||
}
|
||||
// Defensive last-ditch fallback.
|
||||
return SplitText(text, cfg), diag
|
||||
}
|
||||
|
||||
// SplitParentChild is the strategy-aware analog of SplitTextParentChild.
|
||||
// It runs the tier selector for parent splitting, then re-splits each
|
||||
// parent into children with the small-chunk config.
|
||||
//
|
||||
// Child splitting honours childCfg.Strategy. If it is empty/auto and a
|
||||
// parent has its own internal structure (sub-headings, numbered sub-
|
||||
// sections), the appropriate tier picks it up so child chunks carry a
|
||||
// finer-grained breadcrumb than the parent's. Re-profiling each parent
|
||||
// is bounded by O(sum(parent_size)) ≈ O(N) total, which is the same
|
||||
// order as the original parent profiling pass.
|
||||
func SplitParentChild(text string, parentCfg, childCfg SplitterConfig) ParentChildResult {
|
||||
result, _ := splitParentChild(text, parentCfg, childCfg, false)
|
||||
return result
|
||||
}
|
||||
|
||||
// SplitParentChildWithDiagnostics is the diagnostics-aware form of
|
||||
// SplitParentChild. The result is identical to SplitParentChild; diagnostics
|
||||
// describe the strategy selected for the parent split over the full document.
|
||||
// It is intended for debug surfaces rather than the ingestion hot path.
|
||||
func SplitParentChildWithDiagnostics(text string, parentCfg, childCfg SplitterConfig) (ParentChildResult, *Diagnostics) {
|
||||
return splitParentChild(text, parentCfg, childCfg, true)
|
||||
}
|
||||
|
||||
func splitParentChild(text string, parentCfg, childCfg SplitterConfig, withDiagnostics bool) (ParentChildResult, *Diagnostics) {
|
||||
parentCfg = ensureDefaults(parentCfg)
|
||||
childCfg = ensureDefaults(childCfg)
|
||||
|
||||
var (
|
||||
parents []Chunk
|
||||
diag *Diagnostics
|
||||
)
|
||||
if withDiagnostics {
|
||||
parents, diag = SplitWithDiagnostics(text, parentCfg)
|
||||
} else {
|
||||
parents = Split(text, parentCfg)
|
||||
}
|
||||
if len(parents) == 0 {
|
||||
return ParentChildResult{}, diag
|
||||
}
|
||||
|
||||
var newParents []Chunk
|
||||
var children []ChildChunk
|
||||
childSeq := 0
|
||||
for _, parent := range parents {
|
||||
subs := Split(parent.Content, childCfg)
|
||||
|
||||
parentIndex := -1
|
||||
if len(subs) > 1 || (len(subs) == 1 && subs[0].Content != parent.Content) {
|
||||
parentIndex = len(newParents)
|
||||
newParents = append(newParents, parent)
|
||||
}
|
||||
for _, sub := range subs {
|
||||
sub.Seq = childSeq
|
||||
sub.Start += parent.Start
|
||||
sub.End += parent.Start
|
||||
sub.ContextHeader = mergeBreadcrumbs(parent.ContextHeader, sub.ContextHeader)
|
||||
children = append(children, ChildChunk{Chunk: sub, ParentIndex: parentIndex})
|
||||
childSeq++
|
||||
}
|
||||
}
|
||||
return ParentChildResult{Parents: newParents, Children: children}, diag
|
||||
}
|
||||
|
||||
// NormalizeSplitterConfig applies the same base splitter defaults used by
|
||||
// knowledge ingestion before parent-child derivation or single-pass splitting.
|
||||
func NormalizeSplitterConfig(cfg SplitterConfig) SplitterConfig {
|
||||
if cfg.ChunkSize <= 0 {
|
||||
cfg.ChunkSize = DefaultChunkSize
|
||||
}
|
||||
if cfg.ChunkOverlap <= 0 {
|
||||
cfg.ChunkOverlap = DefaultChunkOverlap
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
cfg.Separators = []string{"\n\n", "\n", "。"}
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// NormalizeLineEndings canonicalizes text before chunking. Browser textareas
|
||||
// normalize pasted CRLF text to LF, while uploaded files preserve their
|
||||
// original line endings; without this, the same document produces different
|
||||
// character counts and chunk boundaries.
|
||||
func NormalizeLineEndings(text string) string {
|
||||
if !strings.Contains(text, "\r") {
|
||||
return text
|
||||
}
|
||||
text = strings.ReplaceAll(text, "\r\n", "\n")
|
||||
return strings.ReplaceAll(text, "\r", "\n")
|
||||
}
|
||||
|
||||
// DeriveParentChildConfigs produces the exact parent and child splitter
|
||||
// configurations used by knowledge ingestion. Keeping this here lets preview
|
||||
// and ingestion remain in lockstep as the parent-child defaults evolve.
|
||||
// Languages are copied to both levels so parent and child splitters use the same boundary rules.
|
||||
// TokenLimit is copied only to children because parents keep the configured context window.
|
||||
func DeriveParentChildConfigs(base SplitterConfig, parentSize, childSize int) (parent, child SplitterConfig) {
|
||||
if parentSize <= 0 {
|
||||
parentSize = 4096
|
||||
}
|
||||
if childSize <= 0 {
|
||||
childSize = 384
|
||||
}
|
||||
parent = SplitterConfig{
|
||||
ChunkSize: parentSize,
|
||||
ChunkOverlap: base.ChunkOverlap,
|
||||
Separators: base.Separators,
|
||||
Strategy: base.Strategy,
|
||||
Languages: base.Languages,
|
||||
}
|
||||
child = SplitterConfig{
|
||||
ChunkSize: childSize,
|
||||
ChunkOverlap: childSize / 5,
|
||||
Separators: base.Separators,
|
||||
Strategy: base.Strategy,
|
||||
TokenLimit: base.TokenLimit,
|
||||
Languages: base.Languages,
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// mergeBreadcrumbs combines the parent and child heading breadcrumbs into a
|
||||
// single ContextHeader. When the child re-runs heading detection on parent
|
||||
// content, its first breadcrumb line typically duplicates the parent's last
|
||||
// line (the parent's leading heading sits at the top of the child's input);
|
||||
// drop that duplicate so the embedding context isn't redundant.
|
||||
func mergeBreadcrumbs(parent, child string) string {
|
||||
if parent == "" {
|
||||
return child
|
||||
}
|
||||
if child == "" {
|
||||
return parent
|
||||
}
|
||||
parentLines := strings.Split(parent, "\n")
|
||||
childLines := strings.Split(child, "\n")
|
||||
if len(parentLines) > 0 && len(childLines) > 0 &&
|
||||
strings.TrimSpace(parentLines[len(parentLines)-1]) == strings.TrimSpace(childLines[0]) {
|
||||
childLines = childLines[1:]
|
||||
}
|
||||
if len(childLines) == 0 {
|
||||
return parent
|
||||
}
|
||||
return parent + "\n" + strings.Join(childLines, "\n")
|
||||
}
|
||||
|
||||
// resolveChainWithProfile returns the strategy chain to attempt and, when
|
||||
// the chain was selected by the profiler (auto strategy), the DocProfile
|
||||
// that drove the selection. Profile is nil for explicit non-auto strategies
|
||||
// so callers don't pay for an unused profiling pass.
|
||||
func resolveChainWithProfile(text string, cfg SplitterConfig) ([]StrategyTier, *DocProfile) {
|
||||
switch cfg.Strategy {
|
||||
case StrategyHeading:
|
||||
return []StrategyTier{TierHeading, TierLegacy}, nil
|
||||
case StrategyHeuristic:
|
||||
return []StrategyTier{TierHeuristic, TierLegacy}, nil
|
||||
case StrategyRecursive:
|
||||
// "recursive" is a public-API alias for "legacy": both invoke
|
||||
// SplitText. Kept for backwards compatibility with stored configs.
|
||||
return []StrategyTier{TierLegacy}, nil
|
||||
case StrategyLegacy, "":
|
||||
// Empty == legacy preserves backwards compatibility with stored
|
||||
// ChunkingConfig rows that pre-date the Strategy field.
|
||||
return []StrategyTier{TierLegacy}, nil
|
||||
case StrategyAuto:
|
||||
fallthrough
|
||||
default:
|
||||
profile := ProfileDocument(text)
|
||||
return SelectStrategy(profile), profile
|
||||
}
|
||||
}
|
||||
|
||||
// runTier dispatches the splitter implementation for the given tier.
|
||||
// splitByHeadings / splitByHeuristics are package-level vars overridden
|
||||
// from heading_splitter.go / heuristic_splitter.go via init(); legacy
|
||||
// runs SplitText. The default branch is defensive for future
|
||||
// StrategyTier additions.
|
||||
//
|
||||
// profile may be nil when the caller did not run the document profiler
|
||||
// (explicit non-auto strategies skip profiling); splitters that need a
|
||||
// profile compute one on demand.
|
||||
func runTier(tier StrategyTier, text string, cfg SplitterConfig, profile *DocProfile) []Chunk {
|
||||
switch tier {
|
||||
case TierHeading:
|
||||
return splitByHeadings(text, cfg, profile)
|
||||
case TierHeuristic:
|
||||
return splitByHeuristics(text, cfg, profile)
|
||||
case TierLegacy:
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// ensureDefaults fills in zero-value config fields with sane defaults.
|
||||
// Mirrors buildSplitterConfig in internal/application/service/knowledge.go
|
||||
// so direct callers of this package get the same numbers.
|
||||
//
|
||||
// When cfg.TokenLimit is set, ChunkSize is clamped to the character budget
|
||||
// that fits within that token limit (with a 10% safety factor). This makes
|
||||
// chunks safe for embedding APIs that have hard token caps.
|
||||
func ensureDefaults(cfg SplitterConfig) SplitterConfig {
|
||||
if cfg.ChunkSize <= 0 {
|
||||
cfg.ChunkSize = DefaultChunkSize
|
||||
}
|
||||
if cfg.ChunkOverlap <= 0 {
|
||||
cfg.ChunkOverlap = DefaultChunkOverlap
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
cfg.Separators = []string{"\n\n", "\n", "。"}
|
||||
}
|
||||
if cfg.TokenLimit > 0 {
|
||||
lang := LangMixed
|
||||
if len(cfg.Languages) > 0 {
|
||||
lang = cfg.Languages[0]
|
||||
}
|
||||
charBudget := CharsForTokenLimit(cfg.TokenLimit, lang)
|
||||
if charBudget > 0 && (cfg.ChunkSize == 0 || charBudget < cfg.ChunkSize) {
|
||||
cfg.ChunkSize = charBudget
|
||||
}
|
||||
}
|
||||
// Guard against pathological overlap configurations: if Overlap exceeds
|
||||
// half of ChunkSize, almost every chunk is duplicate content. Cap it at
|
||||
// ChunkSize/2 so Overlap stays a useful smoothing band rather than a
|
||||
// near-clone of the previous chunk.
|
||||
if cfg.ChunkOverlap > cfg.ChunkSize/2 && cfg.ChunkSize > 0 {
|
||||
cfg.ChunkOverlap = cfg.ChunkSize / 2
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// splitByHeadings is overridden by heading_splitter.go. The default no-op
|
||||
// fallback ignores the profile (profile may be nil; the override computes
|
||||
// one on demand when so).
|
||||
var splitByHeadings = func(text string, cfg SplitterConfig, _ *DocProfile) []Chunk {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
|
||||
// splitByHeuristics is overridden by heuristic_splitter.go. profile may be nil.
|
||||
var splitByHeuristics = func(text string, cfg SplitterConfig, _ *DocProfile) []Chunk {
|
||||
return SplitText(text, cfg)
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplitWithDiagnostics_LegacyStrategy_ReportsLegacyTier(t *testing.T) {
|
||||
// Splittable input so the validator accepts the legacy output cleanly.
|
||||
text := strings.Repeat("Hello world.\n\nNext paragraph here.\n\n", 50)
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Separators: []string{"\n\n", "\n"}, Strategy: StrategyLegacy}
|
||||
chunks, diag := SplitWithDiagnostics(text, cfg)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
if diag.SelectedTier != TierLegacy {
|
||||
t.Errorf("expected SelectedTier=legacy, got %s", diag.SelectedTier)
|
||||
}
|
||||
if len(diag.TierChain) != 1 || diag.TierChain[0] != TierLegacy {
|
||||
t.Errorf("expected single-tier chain [legacy], got %v", diag.TierChain)
|
||||
}
|
||||
if len(diag.Rejected) != 0 {
|
||||
t.Errorf("expected no rejections for splittable input, got %v", diag.Rejected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitWithDiagnostics_AutoOnHeadingDoc_PicksHeading(t *testing.T) {
|
||||
doc := strings.Repeat("# Top\nintro paragraph here.\n\n## Section A\nbody A here.\n\n## Section B\nbody B here.\n\n## Section C\nbody C here.\n\n", 1)
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Strategy: StrategyAuto}
|
||||
_, diag := SplitWithDiagnostics(doc, cfg)
|
||||
if len(diag.TierChain) == 0 {
|
||||
t.Fatal("expected non-empty tier chain")
|
||||
}
|
||||
// Heading tier should be tried first for this doc.
|
||||
if diag.TierChain[0] != TierHeading {
|
||||
t.Errorf("expected heading tier first, got chain %v", diag.TierChain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitWithDiagnostics_EmptyText(t *testing.T) {
|
||||
chunks, diag := SplitWithDiagnostics("", DefaultConfig())
|
||||
if chunks != nil {
|
||||
t.Errorf("expected nil chunks for empty text, got %v", chunks)
|
||||
}
|
||||
if diag == nil {
|
||||
t.Fatal("diag must never be nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplit_AndDiagnostics_AgreeOnChunks ensures Split (no diagnostics)
|
||||
// and SplitWithDiagnostics produce the same chunk set for a given input.
|
||||
// They run independent loops as of the post-audit refactor — this test
|
||||
// is the regression wall against them drifting.
|
||||
func TestSplit_AndDiagnostics_AgreeOnChunks(t *testing.T) {
|
||||
text := "para one.\n\npara two.\n\npara three."
|
||||
cfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 10}
|
||||
a := Split(text, cfg)
|
||||
b, diag := SplitWithDiagnostics(text, cfg)
|
||||
if len(a) != len(b) {
|
||||
t.Fatalf("chunk count disagrees: Split=%d Diagnostics=%d", len(a), len(b))
|
||||
}
|
||||
for i := range a {
|
||||
if a[i].Content != b[i].Content || a[i].Start != b[i].Start || a[i].End != b[i].End {
|
||||
t.Errorf("chunk %d differs:\n Split: %+v\n Diag : %+v", i, a[i], b[i])
|
||||
}
|
||||
}
|
||||
if diag == nil {
|
||||
t.Error("diagnostics must not be nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitWithDiagnostics_ProfileSetForAuto verifies that auto-strategy
|
||||
// returns the DocProfile that drove tier selection — required by the
|
||||
// preview endpoint to avoid double-profiling.
|
||||
func TestSplitWithDiagnostics_ProfileSetForAuto(t *testing.T) {
|
||||
doc := "# Top\nintro.\n\n## A\nbody A.\n\n## B\nbody B."
|
||||
_, diag := SplitWithDiagnostics(doc, SplitterConfig{ChunkSize: 200, Strategy: StrategyAuto})
|
||||
if diag.Profile == nil {
|
||||
t.Fatal("auto strategy must populate diag.Profile")
|
||||
}
|
||||
if diag.Profile.MdHeadingTotal == 0 {
|
||||
t.Errorf("profile should have detected headings, got %+v", diag.Profile)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplitWithDiagnostics_ProfileNilForExplicit verifies the inverse:
|
||||
// explicit strategies bypass profiling and leave Profile nil so the
|
||||
// preview handler knows to materialize one if it needs stats.
|
||||
func TestSplitWithDiagnostics_ProfileNilForExplicit(t *testing.T) {
|
||||
for _, strat := range []string{StrategyHeading, StrategyHeuristic, StrategyRecursive, StrategyLegacy} {
|
||||
t.Run(strat, func(t *testing.T) {
|
||||
_, diag := SplitWithDiagnostics("plain text", SplitterConfig{ChunkSize: 200, Strategy: strat})
|
||||
if diag.Profile != nil {
|
||||
t.Errorf("strategy %q should leave Profile nil, got %+v", strat, diag.Profile)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitParentChildWithDiagnostics_MatchesSplitParentChild(t *testing.T) {
|
||||
text := strings.Repeat("## Record\n"+strings.Repeat("A sufficiently long entry body. ", 10)+"\n\n", 12)
|
||||
parentCfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Separators: []string{"\n\n", "\n"}, Strategy: StrategyHeading}
|
||||
childCfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 20, Separators: []string{"\n\n", "\n"}, Strategy: StrategyHeading}
|
||||
|
||||
want := SplitParentChild(text, parentCfg, childCfg)
|
||||
got, diag := SplitParentChildWithDiagnostics(text, parentCfg, childCfg)
|
||||
|
||||
if len(got.Parents) != len(want.Parents) || len(got.Children) != len(want.Children) {
|
||||
t.Fatalf("parent-child result differs: got %d parents/%d children, want %d parents/%d children",
|
||||
len(got.Parents), len(got.Children), len(want.Parents), len(want.Children))
|
||||
}
|
||||
for i := range want.Children {
|
||||
if got.Children[i] != want.Children[i] {
|
||||
t.Errorf("child %d differs:\n got: %+v\n want: %+v", i, got.Children[i], want.Children[i])
|
||||
}
|
||||
}
|
||||
if diag == nil || diag.SelectedTier == "" {
|
||||
t.Fatal("expected diagnostics for parent split")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSplit_EmptyText(t *testing.T) {
|
||||
if got := Split("", DefaultConfig()); got != nil {
|
||||
t.Errorf("empty text should return nil, got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_LegacyStrategy_MatchesSplitText(t *testing.T) {
|
||||
text := strings.Repeat("Hello world.\n\n", 30)
|
||||
cfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 20, Separators: []string{"\n\n"}, Strategy: StrategyLegacy}
|
||||
a := Split(text, cfg)
|
||||
b := SplitText(text, cfg)
|
||||
if len(a) != len(b) {
|
||||
t.Errorf("legacy strategy should match SplitText: got %d vs %d chunks", len(a), len(b))
|
||||
}
|
||||
for i := range a {
|
||||
if a[i].Content != b[i].Content {
|
||||
t.Errorf("chunk %d differs", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_EmptyStrategyEqualsLegacy(t *testing.T) {
|
||||
text := strings.Repeat("Sentence one. Sentence two.\n", 20)
|
||||
cfg := SplitterConfig{ChunkSize: 80, ChunkOverlap: 10}
|
||||
a := Split(text, cfg)
|
||||
cfg.Strategy = StrategyLegacy
|
||||
b := Split(text, cfg)
|
||||
if len(a) != len(b) {
|
||||
t.Errorf("empty Strategy should equal legacy: %d vs %d", len(a), len(b))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_AutoStrategy_PicksHeadingForMarkdownDoc(t *testing.T) {
|
||||
doc := strings.Repeat("# A\nbody\n## B\nbody\n## C\nbody\n## D\nbody\n", 1)
|
||||
cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Strategy: StrategyAuto}
|
||||
// Until heading_splitter is wired, this falls through to SplitText —
|
||||
// just assert we get a valid result.
|
||||
chunks := Split(doc, cfg)
|
||||
if len(chunks) == 0 {
|
||||
t.Error("auto strategy should produce chunks")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit_HeadingStrategyKeepsDistinctTopLevelHeadings(t *testing.T) {
|
||||
doc := `# Intro
|
||||
short intro.
|
||||
|
||||
# Usage
|
||||
short usage.
|
||||
|
||||
# FAQ
|
||||
short faq.`
|
||||
cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0, Strategy: StrategyHeading}
|
||||
chunks := Split(doc, cfg)
|
||||
if len(chunks) != 3 {
|
||||
t.Fatalf("expected one chunk per top-level heading, got %d:\n%v", len(chunks), chunks)
|
||||
}
|
||||
for i, heading := range []string{"# Intro", "# Usage", "# FAQ"} {
|
||||
if !strings.Contains(chunks[i].Content, heading) {
|
||||
t.Errorf("chunk %d should contain heading %q, got:\n%s", i, heading, chunks[i].Content)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSplit_PreservesPositionInvariantAcrossTiers ensures every chunk's
|
||||
// (Start, End, Content) triple stays consistent — End-Start must equal the
|
||||
// rune length of Content, and runes[Start:End] must equal Content. This is
|
||||
// the contract that knowledge.go:2278+ relies on for document reconstruction
|
||||
// during summary generation.
|
||||
func TestSplit_PreservesPositionInvariantAcrossTiers(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"heading-tier": "# Top\nintro paragraph here.\n\n## Section A\nbody A here.\n\n## Section B\nbody B here.\n\n## Section C\nbody C.",
|
||||
"heuristic-tier": strings.Repeat("Kapitel 1: Einleitung\n", 1) + strings.Repeat("Beispieltext. ", 50) +
|
||||
"\n\n" + strings.Repeat("Kapitel 2: Hauptteil\n", 1) + strings.Repeat("Mehr Text. ", 50),
|
||||
"recursive-tier": strings.Repeat("plain prose without structure. ", 100),
|
||||
}
|
||||
cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 30, Separators: []string{"\n\n", "\n", "。", ". "}, Strategy: StrategyAuto}
|
||||
|
||||
for name, doc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
runes := []rune(doc)
|
||||
chunks := Split(doc, cfg)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
for i, c := range chunks {
|
||||
contentRuneLen := len([]rune(c.Content))
|
||||
spanLen := c.End - c.Start
|
||||
if spanLen != contentRuneLen {
|
||||
t.Errorf("chunk %d: End(%d)-Start(%d)=%d but Content has %d runes:\n%q",
|
||||
i, c.End, c.Start, spanLen, contentRuneLen, c.Content)
|
||||
}
|
||||
if c.Start < 0 || c.End > len(runes) {
|
||||
t.Errorf("chunk %d: position out of range Start=%d End=%d totalRunes=%d",
|
||||
i, c.Start, c.End, len(runes))
|
||||
}
|
||||
if c.Start >= 0 && c.End <= len(runes) {
|
||||
sliced := string(runes[c.Start:c.End])
|
||||
if sliced != c.Content {
|
||||
t.Errorf("chunk %d: runes[Start:End] differs from Content", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Children should pick up sub-heading context when childCfg uses auto
|
||||
// strategy on a heading-rich document. Previously child splitting was
|
||||
// pinned to the recursive tier, so sub-chunks all shared the parent's
|
||||
// top-level breadcrumb regardless of inner H2/H3 structure.
|
||||
func TestSplitParentChild_AutoStrategyEnrichesChildBreadcrumbs(t *testing.T) {
|
||||
body := strings.Repeat("Lorem ipsum dolor sit amet. ", 40)
|
||||
doc := "# Chapter\n" + body +
|
||||
"\n\n## Section A\n" + body +
|
||||
"\n\n## Section B\n" + body +
|
||||
"\n\n## Section C\n" + body
|
||||
seps := []string{"\n\n", "\n", ". "}
|
||||
parentCfg := SplitterConfig{ChunkSize: 800, ChunkOverlap: 80, Strategy: StrategyAuto, Separators: seps}
|
||||
childCfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 20, Strategy: StrategyAuto, Separators: seps}
|
||||
res := SplitParentChild(doc, parentCfg, childCfg)
|
||||
if len(res.Children) == 0 {
|
||||
t.Fatal("expected children")
|
||||
}
|
||||
|
||||
// At least one child should carry a breadcrumb that includes a
|
||||
// sub-heading (## Section …) — that's only possible if the child
|
||||
// splitter re-detected headings inside the parent.
|
||||
saw := false
|
||||
for _, c := range res.Children {
|
||||
if strings.Contains(c.ContextHeader, "## Section") {
|
||||
saw = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !saw {
|
||||
t.Errorf("no child carries a sub-heading breadcrumb. samples:\n %q\n %q",
|
||||
firstHeader(res.Children), lastHeader(res.Children))
|
||||
}
|
||||
|
||||
// And no breadcrumb should contain the same line twice in a row
|
||||
// (mergeBreadcrumbs deduplicates the parent/child seam).
|
||||
for i, c := range res.Children {
|
||||
lines := strings.Split(c.ContextHeader, "\n")
|
||||
for j := 1; j < len(lines); j++ {
|
||||
if strings.TrimSpace(lines[j]) != "" && lines[j] == lines[j-1] {
|
||||
t.Errorf("child[%d] has duplicate breadcrumb lines: %q", i, c.ContextHeader)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func firstHeader(cs []ChildChunk) string {
|
||||
if len(cs) == 0 {
|
||||
return ""
|
||||
}
|
||||
return cs[0].ContextHeader
|
||||
}
|
||||
|
||||
func lastHeader(cs []ChildChunk) string {
|
||||
if len(cs) == 0 {
|
||||
return ""
|
||||
}
|
||||
return cs[len(cs)-1].ContextHeader
|
||||
}
|
||||
|
||||
func TestMergeBreadcrumbs(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
parent, child string
|
||||
want string
|
||||
}{
|
||||
{"empty parent", "", "## Sub", "## Sub"},
|
||||
{"empty child", "# Top", "", "# Top"},
|
||||
{"both empty", "", "", ""},
|
||||
{"disjoint", "# Top", "## Other", "# Top\n## Other"},
|
||||
{"duplicate seam", "# Top\n## A", "## A\n### A1", "# Top\n## A\n### A1"},
|
||||
{"deep duplicate", "# Top", "# Top", "# Top"},
|
||||
{"only whitespace differs", "# Top\n## A", " ## A \n### A1", "# Top\n## A\n### A1"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := mergeBreadcrumbs(tc.parent, tc.child)
|
||||
if got != tc.want {
|
||||
t.Errorf("mergeBreadcrumbs(%q, %q) = %q, want %q", tc.parent, tc.child, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitParentChild_LegacyStrategy(t *testing.T) {
|
||||
text := strings.Repeat("This is a sentence. Another one.\n\n", 50)
|
||||
parentCfg := SplitterConfig{ChunkSize: 400, ChunkOverlap: 40, Strategy: StrategyLegacy}
|
||||
childCfg := SplitterConfig{ChunkSize: 100, ChunkOverlap: 20, Strategy: StrategyLegacy}
|
||||
res := SplitParentChild(text, parentCfg, childCfg)
|
||||
if len(res.Children) == 0 {
|
||||
t.Fatal("expected children chunks")
|
||||
}
|
||||
for i, c := range res.Children {
|
||||
if c.ParentIndex >= 0 && c.ParentIndex >= len(res.Parents) {
|
||||
t.Errorf("child[%d] has invalid ParentIndex %d (parents=%d)", i, c.ParentIndex, len(res.Parents))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureDefaults(t *testing.T) {
|
||||
cfg := ensureDefaults(SplitterConfig{})
|
||||
if cfg.ChunkSize != DefaultChunkSize {
|
||||
t.Errorf("expected default ChunkSize %d, got %d", DefaultChunkSize, cfg.ChunkSize)
|
||||
}
|
||||
if cfg.ChunkOverlap != DefaultChunkOverlap {
|
||||
t.Errorf("expected default ChunkOverlap %d, got %d", DefaultChunkOverlap, cfg.ChunkOverlap)
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
t.Error("expected default separators")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeSplitterConfig_MatchesIngestionDefaults(t *testing.T) {
|
||||
cfg := NormalizeSplitterConfig(SplitterConfig{})
|
||||
if cfg.ChunkSize != DefaultChunkSize {
|
||||
t.Errorf("chunk size: got %d want %d", cfg.ChunkSize, DefaultChunkSize)
|
||||
}
|
||||
if cfg.ChunkOverlap != DefaultChunkOverlap {
|
||||
t.Errorf("chunk overlap: got %d want %d", cfg.ChunkOverlap, DefaultChunkOverlap)
|
||||
}
|
||||
if len(cfg.Separators) != 3 {
|
||||
t.Fatalf("separators: got %v", cfg.Separators)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeLineEndings(t *testing.T) {
|
||||
input := "first\r\nsecond\rthird\nfourth"
|
||||
if got, want := NormalizeLineEndings(input), "first\nsecond\nthird\nfourth"; got != want {
|
||||
t.Errorf("NormalizeLineEndings() = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeriveParentChildConfigs_DefaultSizes(t *testing.T) {
|
||||
base := SplitterConfig{
|
||||
ChunkSize: 1000,
|
||||
ChunkOverlap: 100,
|
||||
Separators: []string{"\n\n", "\n"},
|
||||
Strategy: StrategyHeading,
|
||||
}
|
||||
parent, child := DeriveParentChildConfigs(base, 0, 0)
|
||||
if parent.ChunkSize != 4096 || child.ChunkSize != 384 {
|
||||
t.Fatalf("default sizes: parent=%d child=%d", parent.ChunkSize, child.ChunkSize)
|
||||
}
|
||||
if parent.Strategy != StrategyHeading || child.Strategy != StrategyHeading {
|
||||
t.Fatalf("strategy not propagated: parent=%q child=%q", parent.Strategy, child.Strategy)
|
||||
}
|
||||
if child.ChunkOverlap != 384/5 {
|
||||
t.Fatalf("child overlap: got %d want %d", child.ChunkOverlap, 384/5)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeriveParentChildConfigs_AppliesEmbeddingTokenLimitOnlyToChildren(t *testing.T) {
|
||||
const (
|
||||
parentSize = 4096
|
||||
childSize = 1024
|
||||
tokenLimit = 200
|
||||
)
|
||||
base := SplitterConfig{
|
||||
ChunkSize: DefaultChunkSize,
|
||||
ChunkOverlap: DefaultChunkOverlap,
|
||||
Separators: []string{"。"},
|
||||
Strategy: StrategyLegacy,
|
||||
TokenLimit: tokenLimit,
|
||||
Languages: []string{LangChinese},
|
||||
}
|
||||
|
||||
parent, child := DeriveParentChildConfigs(base, parentSize, childSize)
|
||||
if parent.TokenLimit != 0 {
|
||||
t.Fatalf("parent token limit: got %d want 0", parent.TokenLimit)
|
||||
}
|
||||
if len(parent.Languages) != 1 || parent.Languages[0] != LangChinese {
|
||||
t.Fatalf("parent languages: got %v want [%s]", parent.Languages, LangChinese)
|
||||
}
|
||||
if child.TokenLimit != tokenLimit {
|
||||
t.Fatalf("child token limit: got %d want %d", child.TokenLimit, tokenLimit)
|
||||
}
|
||||
if len(child.Languages) != 1 || child.Languages[0] != LangChinese {
|
||||
t.Fatalf("child languages: got %v want [%s]", child.Languages, LangChinese)
|
||||
}
|
||||
|
||||
text := strings.Repeat("这是用于验证中文分块预算的句子。", 100)
|
||||
result := SplitParentChild(text, parent, child)
|
||||
if len(result.Parents) != 1 {
|
||||
t.Fatalf("parents: got %d want 1", len(result.Parents))
|
||||
}
|
||||
if got := len([]rune(result.Parents[0].Content)); got != len([]rune(text)) {
|
||||
t.Fatalf("parent length: got %d want %d", got, len([]rune(text)))
|
||||
}
|
||||
budget := CharsForTokenLimit(tokenLimit, LangChinese)
|
||||
for i, c := range result.Children {
|
||||
if got := len([]rune(c.Content)); got > budget {
|
||||
t.Fatalf("child[%d] length: got %d want <= %d", i, got, budget)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_Empty(t *testing.T) {
|
||||
if v := ValidateChunks(nil, 1000, 500); v.OK {
|
||||
t.Error("nil chunks should be invalid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_SingleChunkLargeDoc(t *testing.T) {
|
||||
c := []Chunk{{Content: strings.Repeat("a", 5000)}}
|
||||
if v := ValidateChunks(c, 5000, 500); v.OK {
|
||||
t.Error("single 10x-too-large chunk should be invalid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_AcceptsReasonableOutput(t *testing.T) {
|
||||
chunks := []Chunk{
|
||||
{Content: strings.Repeat("a", 480)},
|
||||
{Content: strings.Repeat("b", 510)},
|
||||
{Content: strings.Repeat("c", 460)},
|
||||
}
|
||||
if v := ValidateChunks(chunks, 1500, 512); !v.OK {
|
||||
t.Errorf("reasonable chunks should validate, got: %s", v.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_RejectsOversized(t *testing.T) {
|
||||
chunks := []Chunk{
|
||||
{Content: strings.Repeat("a", 100)},
|
||||
{Content: strings.Repeat("b", 5000)}, // > 2x chunkSize
|
||||
}
|
||||
if v := ValidateChunks(chunks, 5100, 1000); v.OK {
|
||||
t.Error("chunk >2x size should be invalid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChunks_TolerantTinyTail(t *testing.T) {
|
||||
chunks := []Chunk{
|
||||
{Content: strings.Repeat("a", 480)},
|
||||
{Content: strings.Repeat("b", 510)},
|
||||
{Content: "tail"},
|
||||
}
|
||||
if v := ValidateChunks(chunks, 994, 512); !v.OK {
|
||||
t.Errorf("tiny last chunk should be tolerated, got: %s", v.Reason)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestEnsureDefaults_TokenLimitClampsChunkSize(t *testing.T) {
|
||||
cfg := SplitterConfig{
|
||||
ChunkSize: 10000, // huge
|
||||
TokenLimit: 100,
|
||||
Languages: []string{LangEnglish},
|
||||
}
|
||||
out := ensureDefaults(cfg)
|
||||
// 100 tokens * 4 chars/token * 0.9 ≈ 360 chars
|
||||
if out.ChunkSize >= 1000 {
|
||||
t.Errorf("expected ChunkSize clamped by TokenLimit, got %d", out.ChunkSize)
|
||||
}
|
||||
if out.ChunkOverlap >= out.ChunkSize {
|
||||
t.Errorf("overlap should be smaller than clamped chunk size: overlap=%d size=%d", out.ChunkOverlap, out.ChunkSize)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureDefaults_TokenLimitChineseTighter(t *testing.T) {
|
||||
cfgEN := ensureDefaults(SplitterConfig{TokenLimit: 200, Languages: []string{LangEnglish}})
|
||||
cfgZH := ensureDefaults(SplitterConfig{TokenLimit: 200, Languages: []string{LangChinese}})
|
||||
if cfgZH.ChunkSize >= cfgEN.ChunkSize {
|
||||
t.Errorf("Chinese char budget should be tighter than English: zh=%d en=%d", cfgZH.ChunkSize, cfgEN.ChunkSize)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureDefaults_NoTokenLimitKeepsChunkSize(t *testing.T) {
|
||||
cfg := SplitterConfig{ChunkSize: 800}
|
||||
out := ensureDefaults(cfg)
|
||||
if out.ChunkSize != 800 {
|
||||
t.Errorf("ChunkSize should stay 800, got %d", out.ChunkSize)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
// Package chunker - tokens.go provides language-aware token count approximation.
|
||||
//
|
||||
// We avoid pulling in a tokenizer dependency (e.g. tiktoken) and instead use
|
||||
// per-language chars-per-token ratios derived from common embedding model
|
||||
// vocabularies. The numbers are conservative — they tend to slightly
|
||||
// over-estimate token counts so that chunks stay safely under model limits.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Language identifiers used by the token estimator and the heuristic splitter.
|
||||
const (
|
||||
LangEnglish = "en"
|
||||
LangGerman = "de"
|
||||
LangChinese = "zh"
|
||||
LangMixed = "mixed"
|
||||
)
|
||||
|
||||
// charsPerToken holds approximate chars/token ratios per language.
|
||||
// Numbers err on the conservative side so estimates over-shoot a little.
|
||||
var charsPerToken = map[string]float64{
|
||||
LangEnglish: 4.0,
|
||||
LangGerman: 4.5,
|
||||
LangChinese: 1.7,
|
||||
LangMixed: 3.0,
|
||||
}
|
||||
|
||||
// ApproxTokenCount returns a conservative token estimate for s in the given
|
||||
// language. An empty or unknown lang falls back to "mixed".
|
||||
func ApproxTokenCount(s string, lang string) int {
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
return ApproxTokenCountFromRuneLen(utf8.RuneCountInString(s), lang)
|
||||
}
|
||||
|
||||
// ApproxTokenCountFromRuneLen is the allocation-free variant of
|
||||
// ApproxTokenCount when the caller has already computed the rune length.
|
||||
// Use this in hot loops where the same content's rune count would
|
||||
// otherwise be recomputed multiple times (e.g. preview endpoint emitting
|
||||
// per-chunk stats).
|
||||
func ApproxTokenCountFromRuneLen(runeLen int, lang string) int {
|
||||
if runeLen <= 0 {
|
||||
return 0
|
||||
}
|
||||
ratio, ok := charsPerToken[lang]
|
||||
if !ok {
|
||||
ratio = charsPerToken[LangMixed]
|
||||
}
|
||||
approx := float64(runeLen) / ratio
|
||||
if approx < 1 {
|
||||
return 1
|
||||
}
|
||||
return int(approx + 0.5)
|
||||
}
|
||||
|
||||
// DetectLanguage returns a coarse language label by counting CJK runes vs.
|
||||
// Latin runes. The result is one of LangChinese, LangGerman, LangEnglish or
|
||||
// LangMixed. Detection is cheap and meant only for heuristic dispatch — it
|
||||
// is NOT a replacement for proper language identification.
|
||||
func DetectLanguage(s string) string {
|
||||
if s == "" {
|
||||
return LangMixed
|
||||
}
|
||||
var cjk, latin, umlaut int
|
||||
for _, r := range s {
|
||||
switch {
|
||||
case unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hangul, r) || unicode.Is(unicode.Hiragana, r) || unicode.Is(unicode.Katakana, r):
|
||||
cjk++
|
||||
case isGermanUmlaut(r):
|
||||
umlaut++
|
||||
latin++
|
||||
case (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z'):
|
||||
latin++
|
||||
}
|
||||
}
|
||||
total := cjk + latin
|
||||
if total == 0 {
|
||||
return LangMixed
|
||||
}
|
||||
cjkRatio := float64(cjk) / float64(total)
|
||||
latinRatio := float64(latin) / float64(total)
|
||||
// Mixed: meaningful presence of both scripts (>=15% each).
|
||||
if cjkRatio >= 0.15 && latinRatio >= 0.15 {
|
||||
return LangMixed
|
||||
}
|
||||
if cjkRatio > 0.3 {
|
||||
return LangChinese
|
||||
}
|
||||
if umlaut > 0 || hasGermanWords(s) {
|
||||
return LangGerman
|
||||
}
|
||||
return LangEnglish
|
||||
}
|
||||
|
||||
func isGermanUmlaut(r rune) bool {
|
||||
switch r {
|
||||
case 'ä', 'ö', 'ü', 'Ä', 'Ö', 'Ü', 'ß':
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// hasGermanWords does a tiny stop-word check to bias towards "de" when the
|
||||
// text uses common German function words. Cheap heuristic — false positives
|
||||
// on borrowed terms are acceptable.
|
||||
func hasGermanWords(s string) bool {
|
||||
const sample = 512
|
||||
if len(s) > sample {
|
||||
s = s[:sample]
|
||||
}
|
||||
for _, w := range []string{" der ", " die ", " das ", " und ", " ist ", " nicht ", " mit ", " auf "} {
|
||||
if containsLower(s, w) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func containsLower(haystack, needle string) bool {
|
||||
if len(haystack) < len(needle) {
|
||||
return false
|
||||
}
|
||||
for i := 0; i+len(needle) <= len(haystack); i++ {
|
||||
match := true
|
||||
for j := 0; j < len(needle); j++ {
|
||||
h := haystack[i+j]
|
||||
if h >= 'A' && h <= 'Z' {
|
||||
h += 'a' - 'A'
|
||||
}
|
||||
if h != needle[j] {
|
||||
match = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if match {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// CharsForTokenLimit converts a token limit into an approximate character
|
||||
// budget for a given language. Used to size chunks so they fit within an
|
||||
// embedding model's max-token window with a small safety margin.
|
||||
func CharsForTokenLimit(tokens int, lang string) int {
|
||||
if tokens <= 0 {
|
||||
return 0
|
||||
}
|
||||
ratio, ok := charsPerToken[lang]
|
||||
if !ok {
|
||||
ratio = charsPerToken[LangMixed]
|
||||
}
|
||||
// 0.9 safety factor so we under-shoot the model limit instead of overshooting.
|
||||
return int(float64(tokens) * ratio * 0.9)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package chunker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestApproxTokenCount_English(t *testing.T) {
|
||||
got := ApproxTokenCount("The quick brown fox jumps over the lazy dog.", LangEnglish)
|
||||
// 44 chars / 4 ≈ 11 tokens
|
||||
if got < 9 || got > 13 {
|
||||
t.Errorf("English token estimate out of range: got %d, want 9..13", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApproxTokenCount_Chinese(t *testing.T) {
|
||||
got := ApproxTokenCount("这是一段中文测试内容用于检验分词估算", LangChinese)
|
||||
// 18 runes / 1.7 ≈ 10
|
||||
if got < 9 || got > 12 {
|
||||
t.Errorf("Chinese token estimate out of range: got %d, want 9..12", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApproxTokenCount_Empty(t *testing.T) {
|
||||
if got := ApproxTokenCount("", LangEnglish); got != 0 {
|
||||
t.Errorf("empty string should return 0 tokens, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApproxTokenCount_UnknownLang(t *testing.T) {
|
||||
got := ApproxTokenCount("Hello world hello world", "xx")
|
||||
if got <= 0 {
|
||||
t.Errorf("unknown lang should fall back to mixed, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_English(t *testing.T) {
|
||||
if got := DetectLanguage("The quick brown fox jumps over the lazy dog."); got != LangEnglish {
|
||||
t.Errorf("expected English, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_German(t *testing.T) {
|
||||
if got := DetectLanguage("Der schnelle braune Fuchs springt über den faulen Hund."); got != LangGerman {
|
||||
t.Errorf("expected German, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_GermanByStopwords(t *testing.T) {
|
||||
// No umlauts but plenty of German function words.
|
||||
if got := DetectLanguage("Das ist ein Test und nicht mit Umlauten."); got != LangGerman {
|
||||
t.Errorf("expected German via stopwords, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_Chinese(t *testing.T) {
|
||||
if got := DetectLanguage("这是一段中文测试内容"); got != LangChinese {
|
||||
t.Errorf("expected Chinese, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectLanguage_Mixed(t *testing.T) {
|
||||
got := DetectLanguage("This 这是 mixed 测试 content with 多语言 inside")
|
||||
if got != LangMixed {
|
||||
t.Errorf("expected Mixed, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCharsForTokenLimit_AppliesSafetyMargin(t *testing.T) {
|
||||
got := CharsForTokenLimit(1000, LangEnglish)
|
||||
// 1000 * 4 * 0.9 = 3600
|
||||
if got < 3500 || got > 3700 {
|
||||
t.Errorf("char budget for 1000 EN tokens out of range: got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCharsForTokenLimit_ZeroTokens(t *testing.T) {
|
||||
if got := CharsForTokenLimit(0, LangEnglish); got != 0 {
|
||||
t.Errorf("zero tokens should give zero chars, got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
// Package chunker - validator.go inspects a tier's output and decides whether
|
||||
// it is good enough to ship or whether the strategy chain should fall through
|
||||
// to the next tier. The validator is intentionally permissive: a single
|
||||
// "obviously broken" output is rejected, but plausible-looking variation is
|
||||
// accepted so we don't oscillate between tiers.
|
||||
package chunker
|
||||
|
||||
import "math"
|
||||
|
||||
// ValidationResult captures the verdict and reason for a chunk-set.
|
||||
type ValidationResult struct {
|
||||
OK bool
|
||||
Reason string
|
||||
}
|
||||
|
||||
// ValidateChunks checks whether the given chunks form a usable result for a
|
||||
// document of totalChars characters with a target chunkSize. Returns OK=true
|
||||
// when no broken-output indicator triggers.
|
||||
func ValidateChunks(chunks []Chunk, totalChars, chunkSize int) ValidationResult {
|
||||
if len(chunks) == 0 {
|
||||
return ValidationResult{Reason: "no chunks produced"}
|
||||
}
|
||||
|
||||
// A single chunk for a document much larger than chunkSize means the
|
||||
// strategy did not actually split — fail so the next tier runs.
|
||||
if len(chunks) == 1 && totalChars > 2*chunkSize {
|
||||
return ValidationResult{Reason: "single chunk for large document"}
|
||||
}
|
||||
|
||||
// Compute size statistics.
|
||||
var sum, sumSq float64
|
||||
maxLen, minLen := 0, math.MaxInt32
|
||||
for _, c := range chunks {
|
||||
l := len([]rune(c.Content))
|
||||
sum += float64(l)
|
||||
sumSq += float64(l * l)
|
||||
if l > maxLen {
|
||||
maxLen = l
|
||||
}
|
||||
if l < minLen {
|
||||
minLen = l
|
||||
}
|
||||
}
|
||||
avg := sum / float64(len(chunks))
|
||||
|
||||
// All but the last chunk should carry meaningful content. We allow the
|
||||
// last chunk to be tiny because tail residue is normal.
|
||||
tinyCount := 0
|
||||
for i, c := range chunks {
|
||||
if i == len(chunks)-1 {
|
||||
continue
|
||||
}
|
||||
if len([]rune(c.Content)) < 50 {
|
||||
tinyCount++
|
||||
}
|
||||
}
|
||||
if tinyCount > len(chunks)/4 && tinyCount > 2 {
|
||||
return ValidationResult{Reason: "too many tiny chunks"}
|
||||
}
|
||||
|
||||
// Reject when no chunk reached at least 25% of the target — the splitter
|
||||
// is fragmenting too aggressively to be useful.
|
||||
if maxLen < chunkSize/4 && totalChars > chunkSize {
|
||||
return ValidationResult{Reason: "all chunks far below target size"}
|
||||
}
|
||||
|
||||
// Sanity check on absolute upper bound. Anything past 2x chunkSize is a
|
||||
// red flag — the splitter ignored its size budget.
|
||||
if maxLen > 2*chunkSize && chunkSize > 0 {
|
||||
return ValidationResult{Reason: "chunk exceeds 2x target size"}
|
||||
}
|
||||
|
||||
_ = avg
|
||||
return ValidationResult{OK: true}
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
// Package anydoc converts office documents to Markdown inside the Go process,
|
||||
// without the Python docreader service.
|
||||
//
|
||||
// The conversion itself is the anydoc Rust library, linked as a static archive
|
||||
// through cgo (see third_party/anydoc-go). That archive needs a Rust toolchain
|
||||
// to produce, so it is opt-in: only builds tagged `anydoc` link it. Every other
|
||||
// build compiles the stub in backend_stub.go, where Available reports false and
|
||||
// the engine registry hides the engine with a reason the UI can show.
|
||||
//
|
||||
// The exported surface is deliberately narrow — bytes in, Markdown and embedded
|
||||
// images out — so the backend can be replaced (a WASI runtime, a future
|
||||
// upstream Go module) without touching the parser code that calls it.
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ImageDir is the markdown path prefix for extracted images. The image
|
||||
// resolver matches references by this path and swaps them for storage URLs.
|
||||
const ImageDir = "images/"
|
||||
|
||||
// Result is one converted document.
|
||||
type Result struct {
|
||||
// Markdown is GitHub-Flavored Markdown for the whole document.
|
||||
// When WithAssets is set, embedded images appear in place as
|
||||
// `` via anydoc's official serializer.
|
||||
Markdown string
|
||||
// Assets are the images embedded in the document, in document order.
|
||||
// Always empty for PDF, which anydoc renders straight to Markdown
|
||||
// without a document model.
|
||||
Assets []Asset
|
||||
// AssetsError explains why Assets is empty when images were asked for.
|
||||
// The document-model parse is what both extracts images and places them
|
||||
// in the markdown; when it fails, Markdown still comes from the text-only
|
||||
// renderer so the conversion succeeds and the caller decides whether to
|
||||
// log or ignore the loss.
|
||||
AssetsError error
|
||||
}
|
||||
|
||||
// Asset is one image embedded in a document.
|
||||
type Asset struct {
|
||||
// ID is the document-model asset index, as referenced by in-place image
|
||||
// sources. It is stable for a single conversion.
|
||||
ID uint64
|
||||
// Name is a generated, extension-carrying file name ("image-1.png"):
|
||||
// embedded assets have no name of their own inside the container.
|
||||
Name string
|
||||
// MediaType is the IANA media type the container declared.
|
||||
MediaType string
|
||||
// Data is the raw image bytes.
|
||||
Data []byte
|
||||
// Alt is the image's alternative text as the document author wrote it,
|
||||
// empty when the document carries none.
|
||||
Alt string
|
||||
// Section is the text of the nearest heading above the image, empty when
|
||||
// the image sits before the first heading.
|
||||
Section string
|
||||
}
|
||||
|
||||
// Options tunes a single conversion.
|
||||
type Options struct {
|
||||
// Format names the parser explicitly ("docx", "csv", ...). Empty means
|
||||
// the format is detected from the content, which every format except
|
||||
// CSV carries a signature for.
|
||||
Format string
|
||||
// WithAssets extracts embedded images and renders them in place in the
|
||||
// markdown using anydoc's official GFM serializer (asset images are
|
||||
// rewritten to `images/image-N.ext` links first). Ignored for PDF.
|
||||
WithAssets bool
|
||||
}
|
||||
|
||||
// ErrUnavailable is returned by Convert when the binding is not linked into
|
||||
// this build. Callers that can fall back to another engine should check for it
|
||||
// with errors.Is.
|
||||
var ErrUnavailable = errors.New("anydoc: binding not built into this binary")
|
||||
|
||||
// supportedExtensions are the file types anydoc converts, mapped to the format
|
||||
// name its parser selector uses. The list is static so that engine metadata
|
||||
// (file types shown in the UI) is identical whether or not the binding is
|
||||
// linked in.
|
||||
var supportedExtensions = map[string]string{
|
||||
"doc": "doc",
|
||||
"docx": "docx",
|
||||
"docm": "docx",
|
||||
"odt": "odt",
|
||||
"rtf": "rtf",
|
||||
"ppt": "ppt",
|
||||
"pptx": "pptx",
|
||||
"pptm": "pptx",
|
||||
"odp": "odp",
|
||||
"xls": "xlsx",
|
||||
"xlsx": "xlsx",
|
||||
"xlsm": "xlsx",
|
||||
"ods": "ods",
|
||||
"epub": "epub",
|
||||
"csv": "csv",
|
||||
"pdf": "pdf",
|
||||
}
|
||||
|
||||
// SupportedFileTypes returns the extensions anydoc can convert, sorted so the
|
||||
// engine list is stable across restarts.
|
||||
func SupportedFileTypes() []string {
|
||||
types := make([]string, 0, len(supportedExtensions))
|
||||
for ext := range supportedExtensions {
|
||||
types = append(types, ext)
|
||||
}
|
||||
slices.Sort(types)
|
||||
return types
|
||||
}
|
||||
|
||||
// FormatForFile resolves the anydoc format name for a file type or file name.
|
||||
// ok is false for anything anydoc does not convert.
|
||||
func FormatForFile(fileType, fileName string) (format string, ok bool) {
|
||||
ext := normalizeExt(fileType)
|
||||
if ext == "" {
|
||||
ext = normalizeExt(filepath.Ext(fileName))
|
||||
}
|
||||
format, ok = supportedExtensions[ext]
|
||||
return format, ok
|
||||
}
|
||||
|
||||
// Supports reports whether anydoc converts this file type, regardless of
|
||||
// whether the binding is linked into this build.
|
||||
func Supports(fileType, fileName string) bool {
|
||||
_, ok := FormatForFile(fileType, fileName)
|
||||
return ok
|
||||
}
|
||||
|
||||
// Available reports whether conversions can actually run in this binary.
|
||||
func Available() bool { return backendAvailable() }
|
||||
|
||||
// UnavailableReason explains, for the UI, why Available is false. It returns
|
||||
// "" when the binding is available.
|
||||
func UnavailableReason() string { return backendUnavailableReason() }
|
||||
|
||||
// Version reports the anydoc release the linked binding was built from, or ""
|
||||
// when no binding is linked.
|
||||
func Version() string { return backendVersion() }
|
||||
|
||||
// Convert turns document bytes into Markdown. It returns ErrUnavailable when
|
||||
// the binding is not linked into this build.
|
||||
func Convert(data []byte, opts Options) (*Result, error) {
|
||||
if len(data) == 0 {
|
||||
return nil, fmt.Errorf("anydoc: empty document")
|
||||
}
|
||||
if !backendAvailable() {
|
||||
return nil, fmt.Errorf("%w: %s", ErrUnavailable, backendUnavailableReason())
|
||||
}
|
||||
// PDF has no document model, so asset extraction would only produce an
|
||||
// error the caller has to special-case. Drop the request instead.
|
||||
if opts.Format == "pdf" {
|
||||
opts.WithAssets = false
|
||||
}
|
||||
return backendConvert(data, opts)
|
||||
}
|
||||
|
||||
// PDFNeedsOCR reports whether a Convert error means the PDF has no usable
|
||||
// text layer. Scanned exam papers fail this way; the caller should fall back
|
||||
// to an engine that can rasterize pages for OCR.
|
||||
func PDFNeedsOCR(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(err.Error())
|
||||
return strings.Contains(msg, "ocr is required") ||
|
||||
strings.Contains(msg, "no extractable text")
|
||||
}
|
||||
|
||||
func normalizeExt(s string) string {
|
||||
return strings.ToLower(strings.TrimPrefix(strings.TrimSpace(s), "."))
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFormatForFileResolvesTypeAndName(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
fileType string
|
||||
fileName string
|
||||
wantFormat string
|
||||
wantOK bool
|
||||
}{
|
||||
{name: "file type", fileType: "docx", wantFormat: "docx", wantOK: true},
|
||||
{name: "dotted file type", fileType: ".DOCX", wantFormat: "docx", wantOK: true},
|
||||
{name: "macro format maps to base format", fileType: "pptm", wantFormat: "pptx", wantOK: true},
|
||||
{name: "legacy excel maps to xlsx parser", fileType: "xls", wantFormat: "xlsx", wantOK: true},
|
||||
{name: "falls back to file name", fileName: "report.odt", wantFormat: "odt", wantOK: true},
|
||||
{name: "unsupported type", fileType: "png", wantOK: false},
|
||||
{name: "nothing to go on", wantOK: false},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
format, ok := FormatForFile(tc.fileType, tc.fileName)
|
||||
if ok != tc.wantOK {
|
||||
t.Fatalf("FormatForFile(%q, %q) ok = %v, want %v", tc.fileType, tc.fileName, ok, tc.wantOK)
|
||||
}
|
||||
if format != tc.wantFormat {
|
||||
t.Fatalf("FormatForFile(%q, %q) = %q, want %q", tc.fileType, tc.fileName, format, tc.wantFormat)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The engine list must look the same whether or not the converter is linked
|
||||
// in, so that the UI can show the file types with a "not built in" reason.
|
||||
func TestSupportedFileTypesIsStableAndSorted(t *testing.T) {
|
||||
types := SupportedFileTypes()
|
||||
if !slices.IsSorted(types) {
|
||||
t.Fatalf("SupportedFileTypes is not sorted: %v", types)
|
||||
}
|
||||
for _, want := range []string{"docx", "pptx", "xlsx", "epub", "pdf", "csv"} {
|
||||
if !slices.Contains(types, want) {
|
||||
t.Errorf("SupportedFileTypes does not include %q: %v", want, types)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFNeedsOCR(t *testing.T) {
|
||||
if PDFNeedsOCR(nil) {
|
||||
t.Fatal("nil error was treated as needing OCR")
|
||||
}
|
||||
if !PDFNeedsOCR(errors.New("PDF has no extractable text (Scanned, 5 pages): OCR is required")) {
|
||||
t.Error("scanned-PDF error was not detected")
|
||||
}
|
||||
if PDFNeedsOCR(errors.New("not a readable zip archive")) {
|
||||
t.Error("unrelated error was treated as needing OCR")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertRejectsEmptyInput(t *testing.T) {
|
||||
if _, err := Convert(nil, Options{Format: "docx"}); err == nil {
|
||||
t.Fatal("Convert(nil) succeeded, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// Without the build tag there is no converter, and callers must be able to
|
||||
// detect that so they can fall back to another engine.
|
||||
func TestConvertReportsUnavailableBackend(t *testing.T) {
|
||||
if Available() {
|
||||
t.Skip("converter is linked into this build")
|
||||
}
|
||||
_, err := Convert([]byte("id,name\n1,a\n"), Options{Format: "csv"})
|
||||
if !errors.Is(err, ErrUnavailable) {
|
||||
t.Fatalf("Convert error = %v, want ErrUnavailable", err)
|
||||
}
|
||||
if UnavailableReason() == "" {
|
||||
t.Error("UnavailableReason is empty while the converter is unavailable")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
//go:build anydoc && cgo
|
||||
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"mime"
|
||||
"strings"
|
||||
|
||||
upstream "github.com/firecrawl/anydoc/go"
|
||||
)
|
||||
|
||||
// This build links the anydoc Rust archive through the upstream cgo bindings
|
||||
// (vendored under third_party/anydoc-go until they are published as a module).
|
||||
// Conversion runs in-process: no subprocess, no HTTP hop, no Python.
|
||||
|
||||
func backendAvailable() bool { return true }
|
||||
|
||||
func backendUnavailableReason() string { return "" }
|
||||
|
||||
func backendVersion() string { return upstream.Version }
|
||||
|
||||
func backendConvert(data []byte, opts Options) (*Result, error) {
|
||||
format, err := upstreamFormat(opts.Format)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if !opts.WithAssets {
|
||||
markdown, err := upstream.ToMarkdownBytes(data, format)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc: markdown conversion failed: %w", err)
|
||||
}
|
||||
return &Result{Markdown: markdown}, nil
|
||||
}
|
||||
|
||||
// Official GFM serializer after rewriting Asset images to External URLs
|
||||
// (`images/image-N.ext`). ToDocument still supplies the image bytes.
|
||||
document, err := upstream.ToDocument(data, format)
|
||||
if err != nil {
|
||||
markdown, mdErr := upstream.ToMarkdownBytes(data, format)
|
||||
if mdErr != nil {
|
||||
return nil, fmt.Errorf("anydoc: markdown conversion failed: %w", mdErr)
|
||||
}
|
||||
return &Result{
|
||||
Markdown: markdown,
|
||||
AssetsError: fmt.Errorf("anydoc: image extraction failed: %w", err),
|
||||
}, nil
|
||||
}
|
||||
markdown, err := upstream.ToMarkdownWithAssetLinks(data, format)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc: markdown conversion failed: %w", err)
|
||||
}
|
||||
return &Result{Markdown: markdown, Assets: collectAssets(document)}, nil
|
||||
}
|
||||
|
||||
// upstreamFormat maps our format name onto the binding's constant. An empty
|
||||
// name means "detect from content".
|
||||
func upstreamFormat(name string) (*upstream.Format, error) {
|
||||
if name == "" {
|
||||
return nil, nil
|
||||
}
|
||||
format, ok := upstream.FormatFromExtension(name)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("anydoc: unsupported format %q", name)
|
||||
}
|
||||
return &format, nil
|
||||
}
|
||||
|
||||
// collectAssets names each embedded image after its position in the document,
|
||||
// because containers store assets without a usable file name of their own, and
|
||||
// attaches the alt text and section the image appeared under.
|
||||
func collectAssets(document *upstream.Document) []Asset {
|
||||
if document == nil || len(document.Assets) == 0 {
|
||||
return nil
|
||||
}
|
||||
placements := placeAssets(document.Blocks)
|
||||
assets := make([]Asset, 0, len(document.Assets))
|
||||
for i, asset := range document.Assets {
|
||||
if len(asset.Data) == 0 {
|
||||
continue
|
||||
}
|
||||
placement := placements[uint64(i)]
|
||||
assets = append(assets, Asset{
|
||||
ID: uint64(i),
|
||||
Name: fmt.Sprintf("image-%d%s", i+1, extensionFor(asset.MediaType)),
|
||||
MediaType: asset.MediaType,
|
||||
Data: asset.Data,
|
||||
Alt: placement.alt,
|
||||
Section: placement.section,
|
||||
})
|
||||
}
|
||||
return assets
|
||||
}
|
||||
|
||||
// placement is where in the document an asset was referenced from.
|
||||
type placement struct {
|
||||
alt string
|
||||
section string
|
||||
}
|
||||
|
||||
// placeAssets walks the document body and records, per asset, the alt text of
|
||||
// the image that referenced it and the heading it sat under. The walk
|
||||
// recurses into lists, tables, and quotes, since an image can be nested
|
||||
// anywhere a paragraph can.
|
||||
func placeAssets(blocks []upstream.Block) map[uint64]placement {
|
||||
placements := make(map[uint64]placement)
|
||||
section := ""
|
||||
walkBlocks(blocks, §ion, placements)
|
||||
return placements
|
||||
}
|
||||
|
||||
func walkBlocks(blocks []upstream.Block, section *string, placements map[uint64]placement) {
|
||||
for _, block := range blocks {
|
||||
switch block.Kind {
|
||||
case "heading":
|
||||
*section = inlineText(block.Content)
|
||||
walkInlines(block.Content, *section, placements)
|
||||
case "paragraph":
|
||||
walkInlines(block.Content, *section, placements)
|
||||
case "block_quote":
|
||||
walkBlocks(block.Blocks, section, placements)
|
||||
case "list":
|
||||
if block.List != nil {
|
||||
for _, item := range block.List.Items {
|
||||
walkBlocks(item.Blocks, section, placements)
|
||||
}
|
||||
}
|
||||
case "table":
|
||||
if block.Table != nil {
|
||||
for _, row := range block.Table.Grid {
|
||||
for _, slot := range row {
|
||||
if slot.Cell != nil {
|
||||
walkBlocks(slot.Cell.Blocks, section, placements)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func walkInlines(inlines []upstream.Inline, section string, placements map[uint64]placement) {
|
||||
for _, inline := range inlines {
|
||||
switch inline.Kind {
|
||||
case "image":
|
||||
if inline.Source == nil || inline.Source.Kind != "asset" || inline.Source.AssetID == nil {
|
||||
continue
|
||||
}
|
||||
alt := ""
|
||||
if inline.Alt != nil {
|
||||
alt = strings.TrimSpace(*inline.Alt)
|
||||
}
|
||||
// First reference wins: the same asset can be placed more than
|
||||
// once, and the earliest placement is the one a reader would
|
||||
// associate it with.
|
||||
if _, seen := placements[*inline.Source.AssetID]; !seen {
|
||||
placements[*inline.Source.AssetID] = placement{alt: alt, section: section}
|
||||
}
|
||||
case "link":
|
||||
walkInlines(inline.Content, section, placements)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// inlineText flattens inline content to plain text, which is all a heading is
|
||||
// needed for here.
|
||||
func inlineText(inlines []upstream.Inline) string {
|
||||
var text strings.Builder
|
||||
for _, inline := range inlines {
|
||||
switch inline.Kind {
|
||||
case "text":
|
||||
if inline.Text != nil {
|
||||
text.WriteString(*inline.Text)
|
||||
}
|
||||
case "link":
|
||||
text.WriteString(inlineText(inline.Content))
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(text.String())
|
||||
}
|
||||
|
||||
func extensionFor(mediaType string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(mediaType)) {
|
||||
case "image/jpeg", "image/jpg":
|
||||
return ".jpg"
|
||||
case "image/png":
|
||||
return ".png"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/bmp":
|
||||
return ".bmp"
|
||||
case "image/tiff":
|
||||
return ".tiff"
|
||||
case "image/svg+xml":
|
||||
return ".svg"
|
||||
}
|
||||
// Anything else (EMF/WMF drawings, unknown object payloads) keeps
|
||||
// whatever extension the media type registry knows, and .bin otherwise,
|
||||
// so downstream storage never writes an extension-less blob.
|
||||
if exts, err := mime.ExtensionsByType(mediaType); err == nil && len(exts) > 0 {
|
||||
return exts[0]
|
||||
}
|
||||
return ".bin"
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
//go:build !anydoc || !cgo
|
||||
|
||||
package anydoc
|
||||
|
||||
// The default build links no converter: the anydoc archive is a Rust build
|
||||
// artifact, and requiring a Rust toolchain to compile WeKnora would be a steep
|
||||
// price for one optional engine. Everything still compiles and the engine
|
||||
// registry simply reports the engine as unavailable.
|
||||
|
||||
const unavailableReason = "anydoc engine not built into this binary " +
|
||||
"(rebuild with `make build-anydoc` / `-tags anydoc`; " +
|
||||
"Docker images need `--build-arg WITH_ANYDOC=1`)"
|
||||
|
||||
func backendAvailable() bool { return false }
|
||||
|
||||
func backendUnavailableReason() string { return unavailableReason }
|
||||
|
||||
func backendVersion() string { return "" }
|
||||
|
||||
func backendConvert(_ []byte, _ Options) (*Result, error) {
|
||||
// Convert checks Available first, so this is only reachable if a caller
|
||||
// bypasses it.
|
||||
return nil, ErrUnavailable
|
||||
}
|
||||
@@ -0,0 +1,379 @@
|
||||
//go:build anydoc && cgo
|
||||
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// These tests exercise the linked Rust converter, so they only build with the
|
||||
// `anydoc` tag. Run them with:
|
||||
//
|
||||
// scripts/build-anydoc-lib.sh && go test -tags anydoc ./internal/infrastructure/docparser/...
|
||||
|
||||
// SupportedFileTypes is a static list, so that the engine can advertise its
|
||||
// file types in builds that link no converter. This is the check that it still
|
||||
// agrees with what the converter actually accepts, in either direction.
|
||||
func TestSupportedFileTypesMatchTheConverter(t *testing.T) {
|
||||
for _, fileType := range SupportedFileTypes() {
|
||||
format, ok := FormatForFile(fileType, "")
|
||||
if !ok {
|
||||
t.Errorf("%q is advertised but maps to no format", fileType)
|
||||
continue
|
||||
}
|
||||
if _, err := upstreamFormat(format); err != nil {
|
||||
t.Errorf("%q maps to format %q, which the converter rejects: %v", fileType, format, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertCSV(t *testing.T) {
|
||||
result, err := Convert([]byte("quarter,widgets\nQ1,12\nQ2,15\n"), Options{Format: "csv"})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "| quarter | widgets |") {
|
||||
t.Fatalf("expected a markdown table header, got:\n%s", result.Markdown)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "| Q1 | 12 |") {
|
||||
t.Fatalf("expected the first data row, got:\n%s", result.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertDocxWithEmbeddedImage(t *testing.T) {
|
||||
document := buildDocx(t)
|
||||
|
||||
result, err := Convert(document, Options{Format: "docx", WithAssets: true})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "# Quarterly report") {
|
||||
t.Fatalf("expected the heading, got:\n%s", result.Markdown)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "Widgets shipped on time.") {
|
||||
t.Fatalf("expected the body paragraph, got:\n%s", result.Markdown)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "Closing remarks.") {
|
||||
t.Fatalf("expected the trailing paragraph, got:\n%s", result.Markdown)
|
||||
}
|
||||
if len(result.Assets) != 1 {
|
||||
t.Fatalf("got %d assets, want 1", len(result.Assets))
|
||||
}
|
||||
asset := result.Assets[0]
|
||||
if asset.Name != "image-1.png" {
|
||||
t.Errorf("asset name = %q, want image-1.png", asset.Name)
|
||||
}
|
||||
if !bytes.Equal(asset.Data, onePixelPNG()) {
|
||||
t.Errorf("asset data does not round-trip the embedded image")
|
||||
}
|
||||
if asset.Alt != "Shipping chart" {
|
||||
t.Errorf("asset alt = %q, want %q", asset.Alt, "Shipping chart")
|
||||
}
|
||||
if asset.Section != "Quarterly report" {
|
||||
t.Errorf("asset section = %q, want %q", asset.Section, "Quarterly report")
|
||||
}
|
||||
imageLink := ""
|
||||
if !strings.Contains(result.Markdown, imageLink) {
|
||||
t.Fatalf("expected the image in place, got:\n%s", result.Markdown)
|
||||
}
|
||||
before := strings.Index(result.Markdown, "Widgets shipped on time.")
|
||||
at := strings.Index(result.Markdown, imageLink)
|
||||
after := strings.Index(result.Markdown, "Closing remarks.")
|
||||
if before < 0 || at < 0 || after < 0 || !(before < at && at < after) {
|
||||
t.Fatalf("image is not between the surrounding paragraphs:\n%s", result.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertTextlessPDFNeedsOCR(t *testing.T) {
|
||||
_, err := Convert(textlessPDF(), Options{Format: "pdf"})
|
||||
if err == nil {
|
||||
t.Fatal("Convert succeeded on a PDF with no text, want an error")
|
||||
}
|
||||
if !PDFNeedsOCR(err) {
|
||||
t.Fatalf("Convert error = %v, want an OCR-required error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The ABI reports error messages through a thread-local slot that a second
|
||||
// call reads, so a goroutine resumed on another OS thread would report an
|
||||
// empty or someone else's message. The binding pins the goroutine for the
|
||||
// pair; this is the regression test for that.
|
||||
func TestErrorDetailSurvivesConcurrency(t *testing.T) {
|
||||
const conversions = 4000
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var empty, crossed atomic.Int64
|
||||
for i := 0; i < conversions; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
// docx and xlsx failures carry distinctive messages, so a detail
|
||||
// naming the other format proves it came from the wrong thread.
|
||||
format, want, other := "docx", "not a readable zip archive", "unreadable workbook"
|
||||
if i%2 == 1 {
|
||||
format, want, other = "xlsx", "unreadable workbook", "not a readable zip archive"
|
||||
}
|
||||
_, err := Convert([]byte(fmt.Sprintf("garbage %d", i)), Options{Format: format})
|
||||
if err == nil {
|
||||
t.Errorf("Convert(garbage) succeeded")
|
||||
return
|
||||
}
|
||||
detail := err.Error()
|
||||
switch {
|
||||
case strings.Contains(detail, other):
|
||||
crossed.Add(1)
|
||||
case !strings.Contains(detail, want):
|
||||
empty.Add(1)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if empty.Load() > 0 || crossed.Load() > 0 {
|
||||
t.Fatalf("error details lost or crossed between goroutines: lost=%d crossed=%d (of %d)",
|
||||
empty.Load(), crossed.Load(), conversions)
|
||||
}
|
||||
}
|
||||
|
||||
// Detection reads the container itself, so a document whose format is not
|
||||
// named still converts.
|
||||
func TestConvertDetectsFormatFromContent(t *testing.T) {
|
||||
result, err := Convert(buildDocx(t), Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "# Quarterly report") {
|
||||
t.Fatalf("expected the heading, got:\n%s", result.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertRejectsGarbage(t *testing.T) {
|
||||
if _, err := Convert([]byte("not a document at all"), Options{Format: "docx"}); err == nil {
|
||||
t.Fatal("Convert succeeded on garbage input, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// PDF has no document model, so WithAssets has to be dropped rather than turn
|
||||
// a working conversion into an error.
|
||||
func TestConvertPDFIgnoresAssetRequest(t *testing.T) {
|
||||
result, err := Convert(minimalPDF(), Options{Format: "pdf", WithAssets: true})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "Shipping summary") {
|
||||
t.Fatalf("expected the page text, got:\n%s", result.Markdown)
|
||||
}
|
||||
if len(result.Assets) != 0 {
|
||||
t.Errorf("got %d assets for a PDF, want 0", len(result.Assets))
|
||||
}
|
||||
}
|
||||
|
||||
// A PDF whose catalog holds a deeply nested array is the shape of
|
||||
// RUSTSEC-2026-0187: lopdf 0.41 recursed per nesting level and killed the
|
||||
// process with an uncatchable stack overflow, which no Go-side or Rust-side
|
||||
// recover can contain. The pinned dependency in third_party/anydoc-go bounds
|
||||
// the recursion, so this must come back as an ordinary error — and if a
|
||||
// dependency bump ever reintroduces it, this test crashes rather than passing
|
||||
// quietly.
|
||||
func TestDeeplyNestedPDFFailsWithoutKillingTheProcess(t *testing.T) {
|
||||
const depth = 50000
|
||||
nested := strings.Repeat("[", depth) + strings.Repeat("]", depth)
|
||||
|
||||
_, err := Convert(nestedPDF(nested), Options{Format: "pdf"})
|
||||
if err == nil {
|
||||
t.Fatal("Convert succeeded on a PDF with no text, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// A show-text operator with no operand made the page-classification scan walk
|
||||
// back over the whole content stream looking for one, so a stream of bare
|
||||
// `] TJ` tokens cost time quadratic in its length: before pdf-inspector 1.14.2
|
||||
// this input took 26 seconds of CPU, and a 2 MB one took nearly two minutes.
|
||||
// Unbounded work is the half of the hostile-PDF problem that guarded() cannot
|
||||
// catch — it never panics, it just holds the core — so the bound is pinned
|
||||
// here. The budget is deliberately far above the ~7ms a bounded lookback needs:
|
||||
// what it has to distinguish is linear from quadratic, not fast from slow.
|
||||
func TestDetectorLookbackStaysLinear(t *testing.T) {
|
||||
const (
|
||||
operators = 200000
|
||||
budget = 15 * time.Second
|
||||
)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
// The result is irrelevant: the page carries one real text run, so
|
||||
// this converts either way. Only how long it takes is under test.
|
||||
_, _ = Convert(unmatchedShowTextPDF(operators), Options{Format: "pdf"})
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(budget):
|
||||
t.Fatalf("converting a PDF with %d unmatched TJ operators took over %s; "+
|
||||
"the detector's operand lookback is no longer bounded", operators, budget)
|
||||
}
|
||||
}
|
||||
|
||||
// buildDocx writes the smallest OOXML package that carries a heading, a
|
||||
// paragraph, and one embedded image.
|
||||
func buildDocx(t *testing.T) []byte {
|
||||
t.Helper()
|
||||
|
||||
parts := map[string]string{
|
||||
"[Content_Types].xml": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
|
||||
<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
|
||||
<Default Extension="xml" ContentType="application/xml"/>
|
||||
<Default Extension="png" ContentType="image/png"/>
|
||||
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
|
||||
<Override PartName="/word/styles.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.styles+xml"/>
|
||||
</Types>`,
|
||||
"_rels/.rels": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
|
||||
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
|
||||
</Relationships>`,
|
||||
"word/styles.xml": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<w:styles xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
|
||||
<w:style w:type="paragraph" w:styleId="Heading1">
|
||||
<w:name w:val="heading 1"/>
|
||||
<w:pPr><w:outlineLvl w:val="0"/></w:pPr>
|
||||
</w:style>
|
||||
</w:styles>`,
|
||||
"word/_rels/document.xml.rels": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
|
||||
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles" Target="styles.xml"/>
|
||||
<Relationship Id="rId10" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/image" Target="media/image1.png"/>
|
||||
</Relationships>`,
|
||||
"word/document.xml": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"
|
||||
xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"
|
||||
xmlns:wp="http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing"
|
||||
xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main"
|
||||
xmlns:pic="http://schemas.openxmlformats.org/drawingml/2006/picture">
|
||||
<w:body>
|
||||
<w:p><w:pPr><w:pStyle w:val="Heading1"/></w:pPr><w:r><w:t>Quarterly report</w:t></w:r></w:p>
|
||||
<w:p><w:r><w:t>Widgets shipped on time.</w:t></w:r></w:p>
|
||||
<w:p><w:r><w:drawing><wp:inline><wp:docPr id="1" name="Chart" descr="Shipping chart"/>
|
||||
<a:graphic><a:graphicData uri="http://schemas.openxmlformats.org/drawingml/2006/picture">
|
||||
<pic:pic><pic:nvPicPr><pic:cNvPr id="1" name="Chart"/><pic:cNvPicPr/></pic:nvPicPr>
|
||||
<pic:blipFill><a:blip r:embed="rId10"/></pic:blipFill>
|
||||
<pic:spPr/></pic:pic>
|
||||
</a:graphicData></a:graphic>
|
||||
</wp:inline></w:drawing></w:r></w:p>
|
||||
<w:p><w:r><w:t>Closing remarks.</w:t></w:r></w:p>
|
||||
</w:body>
|
||||
</w:document>`,
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
archive := zip.NewWriter(&buf)
|
||||
for name, content := range parts {
|
||||
writer, err := archive.Create(name)
|
||||
if err != nil {
|
||||
t.Fatalf("create %s: %v", name, err)
|
||||
}
|
||||
if _, err := writer.Write([]byte(content)); err != nil {
|
||||
t.Fatalf("write %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
image, err := archive.Create("word/media/image1.png")
|
||||
if err != nil {
|
||||
t.Fatalf("create image part: %v", err)
|
||||
}
|
||||
if _, err := image.Write(onePixelPNG()); err != nil {
|
||||
t.Fatalf("write image part: %v", err)
|
||||
}
|
||||
if err := archive.Close(); err != nil {
|
||||
t.Fatalf("close archive: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// onePixelPNG is a 1x1 transparent PNG: the smallest thing a container will
|
||||
// accept as an image part.
|
||||
func onePixelPNG() []byte {
|
||||
return []byte{
|
||||
0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a,
|
||||
0x00, 0x00, 0x00, 0x0d, 'I', 'H', 'D', 'R',
|
||||
0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
|
||||
0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4, 0x89,
|
||||
0x00, 0x00, 0x00, 0x0a, 'I', 'D', 'A', 'T',
|
||||
0x78, 0x9c, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00, 0x01,
|
||||
0x0d, 0x0a, 0x2d, 0xb4,
|
||||
0x00, 0x00, 0x00, 0x00, 'I', 'E', 'N', 'D', 0xae, 0x42, 0x60, 0x82,
|
||||
}
|
||||
}
|
||||
|
||||
// minimalPDF is a one-page PDF with a single text run, written by hand so the
|
||||
// test carries no binary fixture.
|
||||
func minimalPDF() []byte {
|
||||
content := "BT /F1 12 Tf 20 100 Td (Shipping summary) Tj ET\n"
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] /Contents 4 0 R " +
|
||||
"/Resources << /Font << /F1 5 0 R >> >> >>",
|
||||
fmt.Sprintf("<< /Length %d >>\nstream\n%sendstream", len(content), content),
|
||||
"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
|
||||
})
|
||||
}
|
||||
|
||||
// nestedPDF is a one-page PDF carrying the given nested value in its catalog.
|
||||
func nestedPDF(nested string) []byte {
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R /Nested " + nested + " >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] >>",
|
||||
})
|
||||
}
|
||||
|
||||
// unmatchedShowTextPDF is a one-page PDF whose content stream shows text once
|
||||
// and then repeats a TJ operator whose array operand is never opened.
|
||||
func unmatchedShowTextPDF(operators int) []byte {
|
||||
content := "BT /F1 12 Tf 20 100 Td (Shipping summary) Tj ET\n" +
|
||||
strings.Repeat("] TJ\n", operators)
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] /Contents 4 0 R " +
|
||||
"/Resources << /Font << /F1 5 0 R >> >> >>",
|
||||
fmt.Sprintf("<< /Length %d >>\nstream\n%sendstream", len(content), content),
|
||||
"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
|
||||
})
|
||||
}
|
||||
|
||||
// textlessPDF is a one-page PDF with no text stream: the shape of a scanned
|
||||
// exam paper as far as pdf-inspector is concerned.
|
||||
func textlessPDF() []byte {
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] >>",
|
||||
})
|
||||
}
|
||||
|
||||
func writePDF(objects []string) []byte {
|
||||
var pdf bytes.Buffer
|
||||
offsets := make([]int, len(objects))
|
||||
pdf.WriteString("%PDF-1.4\n")
|
||||
for i, object := range objects {
|
||||
offsets[i] = pdf.Len()
|
||||
fmt.Fprintf(&pdf, "%d 0 obj\n%s\nendobj\n", i+1, object)
|
||||
}
|
||||
|
||||
xref := pdf.Len()
|
||||
fmt.Fprintf(&pdf, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1)
|
||||
for _, offset := range offsets {
|
||||
fmt.Fprintf(&pdf, "%010d 00000 n \n", offset)
|
||||
}
|
||||
fmt.Fprintf(&pdf, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref)
|
||||
return pdf.Bytes()
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// AnydocReader converts office documents to markdown in this process, using
|
||||
// the anydoc converter instead of the Python docreader service.
|
||||
//
|
||||
// Embedded images are rendered in place from the document model. PDFs with no
|
||||
// text layer (scanned exam papers) cannot be converted by anydoc; when a
|
||||
// fallback reader is configured (typically the builtin docreader) those files
|
||||
// are handed off so pages can be rasterized for OCR.
|
||||
type AnydocReader struct {
|
||||
// extractImages controls whether embedded images are parsed and placed
|
||||
// in the markdown. Off halves parse time for text-only indexing.
|
||||
extractImages bool
|
||||
// fallback is used when a PDF has no extractable text. Nil means the
|
||||
// conversion error is returned to the caller.
|
||||
fallback interfaces.DocReader
|
||||
}
|
||||
|
||||
// NewAnydocReader builds a reader. The "anydoc_extract_images" override turns
|
||||
// off image extraction. fallback may be nil.
|
||||
func NewAnydocReader(overrides map[string]string, fallback interfaces.DocReader) *AnydocReader {
|
||||
return &AnydocReader{
|
||||
extractImages: !isFalsey(overrides["anydoc_extract_images"]),
|
||||
fallback: fallback,
|
||||
}
|
||||
}
|
||||
|
||||
// Read converts the document carried by the request.
|
||||
func (r *AnydocReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if req.URL != "" && len(req.FileContent) == 0 {
|
||||
return nil, fmt.Errorf("anydoc engine reads uploaded documents, not URLs")
|
||||
}
|
||||
|
||||
format, ok := anydoc.FormatForFile(req.FileType, req.FileName)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("anydoc engine does not support file type %q", fileTypeOf(req))
|
||||
}
|
||||
|
||||
converted, err := anydoc.Convert(req.FileContent, anydoc.Options{
|
||||
Format: format,
|
||||
WithAssets: r.extractImages,
|
||||
})
|
||||
if format == "pdf" && (anydoc.PDFNeedsOCR(err) || pdfMarkdownIsEmpty(converted, err)) {
|
||||
return r.readScannedPDF(ctx, req, err)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc conversion failed for %q: %w", req.FileName, err)
|
||||
}
|
||||
if converted.AssetsError != nil {
|
||||
logger.Warnf(ctx, "[anydoc] %q parsed as text but its images were dropped: %v",
|
||||
req.FileName, converted.AssetsError)
|
||||
}
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: converted.Markdown,
|
||||
ImageRefs: imageRefsFromAssets(converted.Assets),
|
||||
Metadata: map[string]string{
|
||||
"parser": AnydocEngineName,
|
||||
"anydoc_version": anydoc.Version(),
|
||||
"source_format": format,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func pdfMarkdownIsEmpty(converted *anydoc.Result, err error) bool {
|
||||
if err != nil || converted == nil {
|
||||
return false
|
||||
}
|
||||
return strings.TrimSpace(converted.Markdown) == ""
|
||||
}
|
||||
|
||||
func (r *AnydocReader) readScannedPDF(ctx context.Context, req *types.ReadRequest, convertErr error) (*types.ReadResult, error) {
|
||||
if !r.canFallback() {
|
||||
if convertErr != nil {
|
||||
return nil, fmt.Errorf("anydoc conversion failed for %q: %w", req.FileName, convertErr)
|
||||
}
|
||||
return nil, fmt.Errorf("anydoc conversion failed for %q: PDF has no extractable text; OCR is required", req.FileName)
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[anydoc] %q has no text layer, falling back to builtin for scanned-PDF OCR", req.FileName)
|
||||
fallbackReq := *req
|
||||
fallbackReq.ParserEngine = BuiltinEngineName
|
||||
result, err := r.fallback.Read(ctx, &fallbackReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc scanned-PDF fallback failed for %q: %w", req.FileName, err)
|
||||
}
|
||||
if result == nil {
|
||||
return nil, fmt.Errorf("anydoc scanned-PDF fallback returned no result for %q", req.FileName)
|
||||
}
|
||||
if result.Metadata == nil {
|
||||
result.Metadata = map[string]string{}
|
||||
}
|
||||
if result.Metadata["parser"] == "" {
|
||||
result.Metadata["parser"] = BuiltinEngineName
|
||||
}
|
||||
result.Metadata["anydoc_fallback"] = "scanned_pdf"
|
||||
if result.Metadata["image_source_type"] == "" {
|
||||
result.Metadata["image_source_type"] = "scanned_pdf"
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r *AnydocReader) canFallback() bool {
|
||||
if r.fallback == nil {
|
||||
return false
|
||||
}
|
||||
if connected, ok := r.fallback.(interface{ IsConnected() bool }); ok {
|
||||
return connected.IsConnected()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// imageRefsFromAssets turns extracted assets into ImageRefs that match the
|
||||
// in-place markdown links renderMarkdown wrote.
|
||||
func imageRefsFromAssets(assets []anydoc.Asset) []types.ImageRef {
|
||||
if len(assets) == 0 {
|
||||
return nil
|
||||
}
|
||||
refs := make([]types.ImageRef, 0, len(assets))
|
||||
for _, asset := range assets {
|
||||
ref := anydoc.ImageDir + asset.Name
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: asset.Name,
|
||||
OriginalRef: ref,
|
||||
MimeType: asset.MediaType,
|
||||
ImageData: asset.Data,
|
||||
})
|
||||
}
|
||||
return refs
|
||||
}
|
||||
|
||||
// fileTypeOf reports the request's file type, falling back to the file name's
|
||||
// extension the way every reader in this package does.
|
||||
func fileTypeOf(req *types.ReadRequest) string {
|
||||
if ft := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(req.FileType)), "."); ft != "" {
|
||||
return ft
|
||||
}
|
||||
return strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
|
||||
func isFalsey(value string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "false", "0", "no", "off":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestAnydocReaderRejectsUnsupportedFileType(t *testing.T) {
|
||||
_, err := NewAnydocReader(nil, nil).Read(context.Background(), &types.ReadRequest{
|
||||
FileName: "photo.png",
|
||||
FileType: "png",
|
||||
FileContent: []byte("not really a png"),
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Read succeeded for an unsupported file type, want an error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "png") {
|
||||
t.Errorf("error does not name the file type: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderRejectsURLs(t *testing.T) {
|
||||
_, err := NewAnydocReader(nil, nil).Read(context.Background(), &types.ReadRequest{
|
||||
URL: "https://example.com/report.docx",
|
||||
FileType: "docx",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Read succeeded for a URL request, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAnydocReaderHonoursImageExtractionOverride(t *testing.T) {
|
||||
if !NewAnydocReader(nil, nil).extractImages {
|
||||
t.Error("image extraction is off by default, want on")
|
||||
}
|
||||
if NewAnydocReader(map[string]string{"anydoc_extract_images": "false"}, nil).extractImages {
|
||||
t.Error("image extraction stayed on after the override turned it off")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageRefsFromAssetsMatchInPlaceMarkdownPaths(t *testing.T) {
|
||||
refs := imageRefsFromAssets([]anydoc.Asset{
|
||||
{Name: "image-1.png", MediaType: "image/png", Data: []byte("first")},
|
||||
{Name: "image-2.jpg", MediaType: "image/jpeg", Data: []byte("second")},
|
||||
})
|
||||
if len(refs) != 2 {
|
||||
t.Fatalf("got %d image refs, want 2", len(refs))
|
||||
}
|
||||
if refs[0].OriginalRef != "images/image-1.png" {
|
||||
t.Errorf("first OriginalRef = %q, want images/image-1.png", refs[0].OriginalRef)
|
||||
}
|
||||
if refs[1].MimeType != "image/jpeg" {
|
||||
t.Errorf("second mime type = %q, want image/jpeg", refs[1].MimeType)
|
||||
}
|
||||
if len(refs[0].ImageData) == 0 {
|
||||
t.Error("image ref carries no bytes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageRefsFromAssetsEmpty(t *testing.T) {
|
||||
if refs := imageRefsFromAssets(nil); refs != nil {
|
||||
t.Errorf("got %d image refs, want none", len(refs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderFallsBackForScannedPDF(t *testing.T) {
|
||||
fallback := &stubDocReader{result: &types.ReadResult{
|
||||
MarkdownContent: "",
|
||||
ImageRefs: []types.ImageRef{{
|
||||
Filename: "page_1.jpg",
|
||||
OriginalRef: "images/page_1.jpg",
|
||||
ImageData: []byte("jpeg"),
|
||||
}},
|
||||
Metadata: map[string]string{"image_source_type": "scanned_pdf"},
|
||||
}}
|
||||
result, err := NewAnydocReader(nil, fallback).readScannedPDF(context.Background(), &types.ReadRequest{
|
||||
FileName: "exam.pdf",
|
||||
FileType: "pdf",
|
||||
}, errors.New("PDF has no extractable text (Scanned, 5 pages): OCR is required"))
|
||||
if err != nil {
|
||||
t.Fatalf("readScannedPDF: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.MarkdownContent, "images/page_1.jpg") {
|
||||
t.Errorf("fallback markdown missing page image:\n%s", result.MarkdownContent)
|
||||
}
|
||||
if result.Metadata["anydoc_fallback"] != "scanned_pdf" {
|
||||
t.Errorf("anydoc_fallback = %q, want scanned_pdf", result.Metadata["anydoc_fallback"])
|
||||
}
|
||||
if result.Metadata["image_source_type"] != "scanned_pdf" {
|
||||
t.Errorf("image_source_type = %q, want scanned_pdf", result.Metadata["image_source_type"])
|
||||
}
|
||||
if fallback.got == nil || fallback.got.ParserEngine != BuiltinEngineName {
|
||||
t.Errorf("fallback request engine = %v, want builtin", fallback.got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderScannedPDFWithoutFallbackKeepsError(t *testing.T) {
|
||||
_, err := NewAnydocReader(nil, nil).readScannedPDF(context.Background(), &types.ReadRequest{
|
||||
FileName: "exam.pdf",
|
||||
}, errors.New("PDF has no extractable text: OCR is required"))
|
||||
if err == nil {
|
||||
t.Fatal("readScannedPDF succeeded without a fallback, want an error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "OCR is required") {
|
||||
t.Errorf("error does not mention OCR: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderDisconnectedFallbackKeepsError(t *testing.T) {
|
||||
fallback := &disconnectedDocReader{}
|
||||
_, err := NewAnydocReader(nil, fallback).readScannedPDF(context.Background(), &types.ReadRequest{
|
||||
FileName: "exam.pdf",
|
||||
}, errors.New("PDF has no extractable text: OCR is required"))
|
||||
if err == nil {
|
||||
t.Fatal("readScannedPDF succeeded with a disconnected fallback, want an error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "OCR is required") {
|
||||
t.Errorf("error does not mention OCR: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
type stubDocReader struct {
|
||||
got *types.ReadRequest
|
||||
result *types.ReadResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stubDocReader) Read(_ context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
s.got = req
|
||||
return s.result, s.err
|
||||
}
|
||||
|
||||
type disconnectedDocReader struct{}
|
||||
|
||||
func (disconnectedDocReader) Read(context.Context, *types.ReadRequest) (*types.ReadResult, error) {
|
||||
return nil, errors.New("should not be called")
|
||||
}
|
||||
|
||||
func (disconnectedDocReader) IsConnected() bool { return false }
|
||||
@@ -0,0 +1,223 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// simpleFormats lists file extensions that Go can handle without the Python service.
|
||||
var simpleFormats = map[string]bool{
|
||||
"md": true, "markdown": true,
|
||||
"txt": true, "text": true,
|
||||
"csv": true,
|
||||
"json": true,
|
||||
}
|
||||
|
||||
var imageFormats = map[string]bool{
|
||||
"jpg": true, "jpeg": true, "png": true, "gif": true,
|
||||
"bmp": true, "tiff": true, "webp": true,
|
||||
}
|
||||
|
||||
var audioFormats = map[string]bool{
|
||||
"mp3": true, "wav": true, "m4a": true, "flac": true, "ogg": true,
|
||||
}
|
||||
|
||||
func init() {
|
||||
for k := range imageFormats {
|
||||
simpleFormats[k] = true
|
||||
}
|
||||
for k := range audioFormats {
|
||||
simpleFormats[k] = true
|
||||
}
|
||||
}
|
||||
|
||||
// IsSimpleFormat returns true if the file type can be handled by the Go SimpleFormatReader.
|
||||
func IsSimpleFormat(fileType string) bool {
|
||||
return simpleFormats[strings.ToLower(strings.TrimPrefix(fileType, "."))]
|
||||
}
|
||||
|
||||
// SimpleFormatReader handles simple file formats and images directly in Go,
|
||||
// bypassing the Python docreader service.
|
||||
type SimpleFormatReader struct{}
|
||||
|
||||
// Read reads simple format files and returns markdown.
|
||||
func (b *SimpleFormatReader) Read(_ context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
ft := strings.ToLower(strings.TrimPrefix(req.FileType, "."))
|
||||
if ft == "" {
|
||||
ft = strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
|
||||
switch {
|
||||
case ft == "md" || ft == "markdown":
|
||||
return &types.ReadResult{MarkdownContent: string(req.FileContent)}, nil
|
||||
case ft == "txt" || ft == "text":
|
||||
return &types.ReadResult{MarkdownContent: string(req.FileContent)}, nil
|
||||
case ft == "csv":
|
||||
md, err := csvToMarkdown(req.FileContent)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("csv conversion failed: %w", err)
|
||||
}
|
||||
return &types.ReadResult{MarkdownContent: md}, nil
|
||||
case ft == "json":
|
||||
md, err := jsonToMarkdown(req.FileContent)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("json conversion failed: %w", err)
|
||||
}
|
||||
return &types.ReadResult{MarkdownContent: md}, nil
|
||||
case imageFormats[ft]:
|
||||
return imageToResult(req.FileName, req.FileContent), nil
|
||||
case audioFormats[ft]:
|
||||
return audioToResult(req.FileName, req.FileContent), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported simple format: %s", ft)
|
||||
}
|
||||
}
|
||||
|
||||
// imageToResult wraps a standalone image as a markdown image reference with
|
||||
// the raw bytes in ImageRefs, matching Python ImageParser behaviour.
|
||||
func imageToResult(fileName string, data []byte) *types.ReadResult {
|
||||
if fileName == "" {
|
||||
fileName = "image.png"
|
||||
}
|
||||
refPath := "images/" + fileName
|
||||
// Encode spaces so the markdown URL is valid and matches the regex in ResolveAndStore.
|
||||
safeRef := strings.ReplaceAll(refPath, " ", "%20")
|
||||
mime := http.DetectContentType(data)
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: fmt.Sprintf("", fileName, safeRef),
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: fileName,
|
||||
OriginalRef: safeRef,
|
||||
MimeType: mime,
|
||||
ImageData: data,
|
||||
IsOriginal: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// IsImageFormat returns true if the file type is a recognized image format.
|
||||
func IsImageFormat(fileType string) bool {
|
||||
return imageFormats[strings.ToLower(strings.TrimPrefix(fileType, "."))]
|
||||
}
|
||||
|
||||
// IsAudioFormat returns true if the file type is a recognized audio format.
|
||||
func IsAudioFormat(fileType string) bool {
|
||||
return audioFormats[strings.ToLower(strings.TrimPrefix(fileType, "."))]
|
||||
}
|
||||
|
||||
// audioToResult wraps a standalone audio file. The actual transcription is
|
||||
// handled by the ASR model in the knowledge service pipeline. Here we just
|
||||
// return a placeholder markdown with the raw bytes preserved for upstream
|
||||
// processing.
|
||||
func audioToResult(fileName string, data []byte) *types.ReadResult {
|
||||
if fileName == "" {
|
||||
fileName = "audio.mp3"
|
||||
}
|
||||
// Return a placeholder; the knowledge service will replace this with
|
||||
// the ASR transcription result.
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: fmt.Sprintf("[Audio file: %s]", fileName),
|
||||
IsAudio: true,
|
||||
AudioData: data,
|
||||
}
|
||||
}
|
||||
|
||||
// ensureOriginalImageRef checks whether the input file is an image and, if the
|
||||
// returned markdown does not already contain a markdown image reference for it,
|
||||
// prepends one and appends the raw bytes to imageRefs. This guarantees that
|
||||
// when MinerU OCRs a standalone image, the downstream chunks still carry the
|
||||
// original image link for retrieval display.
|
||||
func ensureOriginalImageRef(req *types.ReadRequest, mdContent string, imageRefs []types.ImageRef) (string, []types.ImageRef) {
|
||||
ft := strings.ToLower(strings.TrimPrefix(req.FileType, "."))
|
||||
if ft == "" {
|
||||
ft = strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
if !imageFormats[ft] {
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
if len(req.FileContent) == 0 {
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
|
||||
fileName := req.FileName
|
||||
if fileName == "" {
|
||||
fileName = "image." + ft
|
||||
}
|
||||
refPath := "images/" + fileName
|
||||
|
||||
if strings.Contains(mdContent, refPath) {
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
|
||||
imgLine := fmt.Sprintf("", fileName, refPath)
|
||||
if strings.TrimSpace(mdContent) == "" {
|
||||
mdContent = imgLine
|
||||
} else {
|
||||
mdContent = imgLine + "\n\n" + mdContent
|
||||
}
|
||||
|
||||
mime := http.DetectContentType(req.FileContent)
|
||||
imageRefs = append(imageRefs, types.ImageRef{
|
||||
Filename: fileName,
|
||||
OriginalRef: refPath,
|
||||
MimeType: mime,
|
||||
ImageData: req.FileContent,
|
||||
IsOriginal: true,
|
||||
})
|
||||
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
|
||||
func csvToMarkdown(data []byte) (string, error) {
|
||||
reader := csv.NewReader(strings.NewReader(string(data)))
|
||||
reader.LazyQuotes = true
|
||||
reader.TrimLeadingSpace = true
|
||||
|
||||
records, err := reader.ReadAll()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(records) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
|
||||
// Header row
|
||||
header := records[0]
|
||||
sb.WriteString("| ")
|
||||
sb.WriteString(strings.Join(header, " | "))
|
||||
sb.WriteString(" |\n")
|
||||
|
||||
// Separator
|
||||
sb.WriteString("|")
|
||||
for range header {
|
||||
sb.WriteString(" --- |")
|
||||
}
|
||||
sb.WriteString("\n")
|
||||
|
||||
// Data rows
|
||||
for _, row := range records[1:] {
|
||||
sb.WriteString("| ")
|
||||
// Pad row if shorter than header
|
||||
cells := make([]string, len(header))
|
||||
for i := range cells {
|
||||
if i < len(row) {
|
||||
cells[i] = row[i]
|
||||
}
|
||||
}
|
||||
sb.WriteString(strings.Join(cells, " | "))
|
||||
sb.WriteString(" |\n")
|
||||
}
|
||||
|
||||
return sb.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
// Package docparser turns uploaded files and URLs into Markdown plus the
|
||||
// images they contain, which is the input the chunking and multimodal stages
|
||||
// downstream expect.
|
||||
//
|
||||
// A parse request names a parser engine, and every engine is one entry in the
|
||||
// registry:
|
||||
//
|
||||
// - engines.go declares the catalog — each engine's name, description,
|
||||
// supported file types, availability probe, and the reader that does the
|
||||
// work.
|
||||
// - engine_registry.go turns an engine name into a reader (NewReader) and
|
||||
// merges the local catalog with the engines the Python docreader reports
|
||||
// over its ListEngines RPC (ListAllEngines).
|
||||
//
|
||||
// Readers differ in where the parsing happens: in this process
|
||||
// (SimpleFormatReader, AnydocReader), in the docreader service over gRPC or
|
||||
// HTTP (GRPCDocumentReader, HTTPDocumentReader), or in a remote API (MinerU,
|
||||
// PaddleOCR-VL, WeKnora Cloud). They all return types.ReadResult, so the rest
|
||||
// of the package — image resolution and storage, table normalization — is
|
||||
// shared regardless of which engine ran.
|
||||
package docparser
|
||||
@@ -0,0 +1,151 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// EngineRegistration is what every locally registered parser engine provides:
|
||||
// the metadata the engine list shows, and the reader that does the parsing.
|
||||
// Remote-only engines (e.g. markitdown) live in the Python docreader, are
|
||||
// discovered through its ListEngines RPC, and never register here — the
|
||||
// registry routes them to the docreader client by default.
|
||||
type EngineRegistration interface {
|
||||
Name() string
|
||||
Description() string
|
||||
FileTypes(docreaderConnected bool) []string
|
||||
CheckAvailable(docreaderConnected bool, overrides map[string]string) (available bool, reason string)
|
||||
// NewReader builds the reader for one parse request. Returning an error
|
||||
// means this engine cannot serve the request (missing credentials,
|
||||
// unreachable service); the caller reports it rather than silently
|
||||
// parsing with something else.
|
||||
NewReader(ctx context.Context, deps ReaderDeps) (interfaces.DocReader, error)
|
||||
}
|
||||
|
||||
// ReaderDeps carries everything an engine may need to build its reader but
|
||||
// cannot construct itself: tenant configuration, tenant credentials, and the
|
||||
// shared docreader connection.
|
||||
type ReaderDeps struct {
|
||||
// Overrides holds tenant-level engine configuration (service endpoints,
|
||||
// API keys), as produced by ParserEngineConfig.ToOverridesMap.
|
||||
Overrides map[string]string
|
||||
// Remote is the docreader client. Nil when the service is not connected.
|
||||
Remote interfaces.DocReader
|
||||
// WeKnoraCloudCredentials resolves the tenant's WeKnora Cloud
|
||||
// credentials. It is a function rather than a value because resolving
|
||||
// them can hit the database, which most engines never need. Nil, or a
|
||||
// nil return, means the tenant has not configured them.
|
||||
WeKnoraCloudCredentials func(ctx context.Context) *types.WeKnoraCloudCredentials
|
||||
}
|
||||
|
||||
// localEngines holds all locally registered parser engines, in registration
|
||||
// order — which is also the order the engine list is shown in.
|
||||
var localEngines []EngineRegistration
|
||||
|
||||
// RegisterEngine adds an engine to the local registry. Called from init().
|
||||
func RegisterEngine(e EngineRegistration) {
|
||||
localEngines = append(localEngines, e)
|
||||
}
|
||||
|
||||
// lookupEngine returns the locally registered engine with this name.
|
||||
func lookupEngine(name string) (EngineRegistration, bool) {
|
||||
for _, engine := range localEngines {
|
||||
if engine.Name() == name {
|
||||
return engine, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// NewReader builds the reader for an engine.
|
||||
//
|
||||
// An empty engine name means "no explicit choice": simple formats are handled
|
||||
// in Go and everything else goes to the docreader service. An unknown name is
|
||||
// routed to the docreader too, so engines that only exist in the Python
|
||||
// service keep working without a Go-side registration.
|
||||
func NewReader(
|
||||
ctx context.Context, engine, fileType string, isURL bool, deps ReaderDeps,
|
||||
) (interfaces.DocReader, error) {
|
||||
if registration, ok := lookupEngine(engine); ok {
|
||||
return registration.NewReader(ctx, deps)
|
||||
}
|
||||
if engine == "" && !isURL && IsSimpleFormat(fileType) {
|
||||
return &SimpleFormatReader{}, nil
|
||||
}
|
||||
return remoteReader(deps)
|
||||
}
|
||||
|
||||
// remoteReader returns the docreader client, or an error when the service is
|
||||
// not connected — a nil interface value here would panic at the call site.
|
||||
func remoteReader(deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
if deps.Remote == nil {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
return deps.Remote, nil
|
||||
}
|
||||
|
||||
// ListAllEngines returns the merged engine list: locally registered engines
|
||||
// plus engines discovered from the remote docreader via ListEngines RPC.
|
||||
//
|
||||
// Merge rules:
|
||||
// - Local engines are always included, with Go-side availability checks.
|
||||
// - For a remote engine whose name matches a local one, the remote's
|
||||
// file_types and description take precedence (the remote service is
|
||||
// authoritative for its own capabilities).
|
||||
// - Remote engines not present locally are appended as-is, enabling
|
||||
// auto-discovery of newly added docreader engines without Go changes.
|
||||
func ListAllEngines(
|
||||
docreaderConnected bool, overrides map[string]string, remoteEngines []types.ParserEngineInfo,
|
||||
) []types.ParserEngineInfo {
|
||||
remoteMap := make(map[string]types.ParserEngineInfo, len(remoteEngines))
|
||||
for _, re := range remoteEngines {
|
||||
remoteMap[re.Name] = re
|
||||
}
|
||||
|
||||
seen := make(map[string]bool, len(localEngines))
|
||||
result := make([]types.ParserEngineInfo, 0, len(localEngines)+len(remoteEngines))
|
||||
|
||||
for _, e := range localEngines {
|
||||
name := e.Name()
|
||||
seen[name] = true
|
||||
|
||||
fileTypes := e.FileTypes(docreaderConnected)
|
||||
description := e.Description()
|
||||
|
||||
if re, ok := remoteMap[name]; ok {
|
||||
if len(re.FileTypes) > 0 {
|
||||
fileTypes = re.FileTypes
|
||||
}
|
||||
if re.Description != "" {
|
||||
description = re.Description
|
||||
}
|
||||
}
|
||||
|
||||
available, reason := e.CheckAvailable(docreaderConnected, overrides)
|
||||
result = append(result, types.ParserEngineInfo{
|
||||
Name: name,
|
||||
Description: description,
|
||||
FileTypes: fileTypes,
|
||||
Available: available,
|
||||
UnavailableReason: reason,
|
||||
})
|
||||
}
|
||||
|
||||
for _, re := range remoteEngines {
|
||||
if seen[re.Name] {
|
||||
continue
|
||||
}
|
||||
result = append(result, re)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// errEngineUnavailable reports an engine that is registered but cannot run for
|
||||
// this tenant or this build.
|
||||
func errEngineUnavailable(engine, reason string) error {
|
||||
return fmt.Errorf("parser engine %q is unavailable: %s", engine, reason)
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// stubRemote stands in for the docreader client. NewReader only ever hands it
|
||||
// back, so the embedded interface is never called — and would panic if routing
|
||||
// ever sent a parse here by mistake.
|
||||
type stubRemote struct{ interfaces.DocReader }
|
||||
|
||||
func TestNewReaderRoutesByEngine(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
remote := &stubRemote{}
|
||||
deps := ReaderDeps{Remote: remote}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
engine string
|
||||
fileType string
|
||||
isURL bool
|
||||
want any
|
||||
}{
|
||||
{name: "simple engine", engine: SimpleEngineName, fileType: "md", want: &SimpleFormatReader{}},
|
||||
{name: "builtin engine goes remote", engine: BuiltinEngineName, fileType: "md", want: remote},
|
||||
{name: "unset engine handles simple formats in Go", fileType: "csv", want: &SimpleFormatReader{}},
|
||||
{name: "unset engine sends complex formats to docreader", fileType: "docx", want: remote},
|
||||
{name: "default pptx engine is remote markitdown", engine: "markitdown", fileType: "pptx", want: remote},
|
||||
{name: "URLs always go to docreader", fileType: "md", isURL: true, want: remote},
|
||||
{name: "docreader-only engines fall through", engine: "markitdown", fileType: "docx", want: remote},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
reader, err := NewReader(ctx, tc.engine, tc.fileType, tc.isURL, deps)
|
||||
if err != nil {
|
||||
t.Fatalf("NewReader: %v", err)
|
||||
}
|
||||
if _, isSimple := tc.want.(*SimpleFormatReader); isSimple {
|
||||
if _, ok := reader.(*SimpleFormatReader); !ok {
|
||||
t.Fatalf("reader = %T, want *SimpleFormatReader", reader)
|
||||
}
|
||||
return
|
||||
}
|
||||
if reader != tc.want {
|
||||
t.Fatalf("reader = %T, want the docreader client", reader)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A disconnected docreader has to surface as an error, not as a nil reader the
|
||||
// caller would dereference.
|
||||
func TestNewReaderReportsDisconnectedDocReader(t *testing.T) {
|
||||
if _, err := NewReader(context.Background(), BuiltinEngineName, "docx", false, ReaderDeps{}); err == nil {
|
||||
t.Fatal("NewReader succeeded without a docreader connection, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewReaderRequiresWeKnoraCloudCredentials(t *testing.T) {
|
||||
_, err := NewReader(context.Background(), WeKnoraCloudEngineName, "docx", false, ReaderDeps{
|
||||
WeKnoraCloudCredentials: func(context.Context) *types.WeKnoraCloudCredentials { return nil },
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("NewReader succeeded without credentials, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// The anydoc engine is only linked into builds tagged `anydoc`; everywhere
|
||||
// else it must be listed as unavailable and refuse to build a reader, so a
|
||||
// knowledge base configured for it fails loudly instead of parsing with
|
||||
// something else.
|
||||
func TestAnydocEngineFollowsBuildAvailability(t *testing.T) {
|
||||
reader, err := NewReader(context.Background(), AnydocEngineName, "docx", false, ReaderDeps{})
|
||||
|
||||
if anydoc.Available() {
|
||||
if err != nil {
|
||||
t.Fatalf("NewReader: %v", err)
|
||||
}
|
||||
if _, ok := reader.(*AnydocReader); !ok {
|
||||
t.Fatalf("reader = %T, want *AnydocReader", reader)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("NewReader succeeded without the converter linked in, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAllEnginesIncludesAnydoc(t *testing.T) {
|
||||
for _, engine := range ListAllEngines(true, nil, nil) {
|
||||
if engine.Name != AnydocEngineName {
|
||||
continue
|
||||
}
|
||||
if engine.Available != anydoc.Available() {
|
||||
t.Errorf("anydoc availability = %v, want %v", engine.Available, anydoc.Available())
|
||||
}
|
||||
if !engine.Available && engine.UnavailableReason == "" {
|
||||
t.Error("anydoc is unavailable without a reason to show")
|
||||
}
|
||||
if len(engine.FileTypes) == 0 {
|
||||
t.Error("anydoc lists no file types")
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatal("anydoc engine not found in the engine list")
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestListAllEnginesBuiltinIncludesDocumentFormats(t *testing.T) {
|
||||
engines := ListAllEngines(true, nil, nil)
|
||||
for _, engine := range engines {
|
||||
if engine.Name != "builtin" {
|
||||
continue
|
||||
}
|
||||
if !engine.Available {
|
||||
t.Fatalf("builtin engine is unavailable: %s", engine.UnavailableReason)
|
||||
}
|
||||
|
||||
fileTypes := make(map[string]bool, len(engine.FileTypes))
|
||||
for _, fileType := range engine.FileTypes {
|
||||
fileTypes[fileType] = true
|
||||
}
|
||||
for _, want := range []string{"html", "htm", "xmind", "ppt", "pptx"} {
|
||||
if !fileTypes[want] {
|
||||
t.Errorf("builtin engine file types do not include %q: %v", want, engine.FileTypes)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
t.Fatal("builtin engine not found")
|
||||
}
|
||||
|
||||
func TestDefaultParserEnginePrefersAnydocWhenLinked(t *testing.T) {
|
||||
cases := []string{"pptx", "ppt", "pdf", "docx"}
|
||||
if anydoc.Available() {
|
||||
for _, ft := range cases {
|
||||
if got := types.DefaultParserEngine(ft); got != AnydocEngineName {
|
||||
t.Errorf("DefaultParserEngine(%s) = %q, want anydoc when the binding is linked", ft, got)
|
||||
}
|
||||
}
|
||||
if got := types.DefaultParserEngine("csv"); got != "" {
|
||||
t.Errorf("DefaultParserEngine(csv) = %q, want empty so the Go simple reader stays default", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
if got := types.DefaultParserEngine("pptx"); got != "markitdown" {
|
||||
t.Fatalf("DefaultParserEngine(pptx) = %q, want markitdown when anydoc is unavailable", got)
|
||||
}
|
||||
if got := types.DefaultParserEngine("pdf"); got != "" {
|
||||
t.Fatalf("DefaultParserEngine(pdf) = %q, want empty when anydoc is unavailable", got)
|
||||
}
|
||||
if got := types.DefaultParserEngine("docx"); got != "" {
|
||||
t.Fatalf("DefaultParserEngine(docx) = %q, want empty when anydoc is unavailable", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// Engine names, as stored in knowledge base parser rules and shown in the UI.
|
||||
const (
|
||||
// BuiltinEngineName is the DocReader (Python) parser suite.
|
||||
BuiltinEngineName = "builtin"
|
||||
// SimpleEngineName is Go-native handling of text formats and images.
|
||||
SimpleEngineName = "simple"
|
||||
// AnydocEngineName is the in-process anydoc office-document converter.
|
||||
AnydocEngineName = "anydoc"
|
||||
// WeKnoraCloudEngineName is the hosted WeKnora Cloud document reader.
|
||||
WeKnoraCloudEngineName = "weknoracloud"
|
||||
// MinerUEngineName is a self-hosted MinerU service.
|
||||
MinerUEngineName = "mineru"
|
||||
// MinerUCloudEngineName is the MinerU Cloud API.
|
||||
MinerUCloudEngineName = "mineru_cloud"
|
||||
// PaddleOCRVLEngineName is a self-hosted PaddleOCR-VL pipeline.
|
||||
PaddleOCRVLEngineName = "paddleocr_vl"
|
||||
// PaddleOCRVLCloudEngineName is the PaddleOCR-VL AI Studio cloud API.
|
||||
PaddleOCRVLCloudEngineName = "paddleocr_vl_cloud"
|
||||
)
|
||||
|
||||
func init() {
|
||||
types.SetPreferParserEngine(preferAnydocWhenAvailable)
|
||||
RegisterEngine(&builtinEngine{})
|
||||
RegisterEngine(&simpleEngine{})
|
||||
RegisterEngine(&anydocEngine{})
|
||||
RegisterEngine(&weKnoraCloudEngine{})
|
||||
RegisterEngine(&mineruEngine{})
|
||||
RegisterEngine(&mineruCloudEngine{})
|
||||
RegisterEngine(&paddleOCRVLEngine{})
|
||||
RegisterEngine(&paddleOCRVLCloudEngine{})
|
||||
}
|
||||
|
||||
// preferAnydocWhenAvailable is the type-level default override: when the
|
||||
// anydoc binding is linked and converts this file type, use it instead of
|
||||
// builtin or the markitdown fallback. Simple formats stay on the Go reader.
|
||||
func preferAnydocWhenAvailable(fileType string) string {
|
||||
if IsSimpleFormat(fileType) {
|
||||
return ""
|
||||
}
|
||||
if anydoc.Available() && anydoc.Supports(fileType, "") {
|
||||
return AnydocEngineName
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// builtin — DocReader-backed parser for complex document formats.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type builtinEngine struct{}
|
||||
|
||||
func (e *builtinEngine) Name() string { return BuiltinEngineName }
|
||||
|
||||
func (e *builtinEngine) Description() string { return "DocReader built-in parser engine" }
|
||||
|
||||
func (e *builtinEngine) FileTypes(_ bool) []string {
|
||||
return []string{
|
||||
"docx", "doc", "pdf", "md", "markdown", "xlsx", "xls",
|
||||
"pptx", "ppt", "epub",
|
||||
"html", "htm", "mhtml", "xmind",
|
||||
"jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp",
|
||||
"mp3", "wav", "m4a", "flac", "ogg",
|
||||
}
|
||||
}
|
||||
|
||||
func (e *builtinEngine) CheckAvailable(docreaderConnected bool, _ map[string]string) (bool, string) {
|
||||
if docreaderConnected {
|
||||
return true, ""
|
||||
}
|
||||
return false, "DocReader service not connected"
|
||||
}
|
||||
|
||||
// NewReader returns the docreader client. Selecting "builtin" explicitly means
|
||||
// the docreader is wanted even for formats the simple reader could handle.
|
||||
func (e *builtinEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return remoteReader(deps)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// simple — Go handles md/txt/csv/json natively, no external service needed.
|
||||
// Distinct from docreader's "builtin", which uses Python libraries for complex
|
||||
// formats (docx, pdf, ...).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type simpleEngine struct{}
|
||||
|
||||
func (e *simpleEngine) Name() string { return SimpleEngineName }
|
||||
|
||||
func (e *simpleEngine) Description() string {
|
||||
return "Simple format & image parsing (no external service required)"
|
||||
}
|
||||
|
||||
func (e *simpleEngine) FileTypes(_ bool) []string {
|
||||
return []string{
|
||||
"md", "markdown", "txt", "csv", "json",
|
||||
"jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp",
|
||||
"mp3", "wav", "m4a", "flac", "ogg",
|
||||
}
|
||||
}
|
||||
|
||||
func (e *simpleEngine) CheckAvailable(_ bool, _ map[string]string) (bool, string) {
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func (e *simpleEngine) NewReader(_ context.Context, _ ReaderDeps) (interfaces.DocReader, error) {
|
||||
return &SimpleFormatReader{}, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// anydoc — office documents converted in this process, no external service.
|
||||
// Only present in builds that link the converter (see the anydoc package).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type anydocEngine struct{}
|
||||
|
||||
func (e *anydocEngine) Name() string { return AnydocEngineName }
|
||||
|
||||
func (e *anydocEngine) Description() string {
|
||||
return "anydoc in-process office document converter (no external service required)"
|
||||
}
|
||||
|
||||
func (e *anydocEngine) FileTypes(_ bool) []string { return anydoc.SupportedFileTypes() }
|
||||
|
||||
func (e *anydocEngine) CheckAvailable(_ bool, _ map[string]string) (bool, string) {
|
||||
if anydoc.Available() {
|
||||
return true, ""
|
||||
}
|
||||
return false, anydoc.UnavailableReason()
|
||||
}
|
||||
|
||||
func (e *anydocEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
if !anydoc.Available() {
|
||||
return nil, errEngineUnavailable(AnydocEngineName, anydoc.UnavailableReason())
|
||||
}
|
||||
return NewAnydocReader(deps.Overrides, deps.Remote), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// weknoracloud — Tenant-scoped WeKnoraCloud docreader with signed requests.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type weKnoraCloudEngine struct{}
|
||||
|
||||
func (e *weKnoraCloudEngine) Name() string { return WeKnoraCloudEngineName }
|
||||
|
||||
func (e *weKnoraCloudEngine) Description() string { return "WeKnoraCloud document reader" }
|
||||
|
||||
func (e *weKnoraCloudEngine) FileTypes(_ bool) []string {
|
||||
return []string{"docx", "doc", "pdf", "md", "markdown", "xlsx", "xls", "pptx", "ppt"}
|
||||
}
|
||||
|
||||
func (e *weKnoraCloudEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
if overrides["weknoracloud_app_id"] != "" {
|
||||
return true, ""
|
||||
}
|
||||
return false, "WeKnora Cloud credentials not configured. Go to Settings → WeKnora Cloud to set up."
|
||||
}
|
||||
|
||||
func (e *weKnoraCloudEngine) NewReader(
|
||||
ctx context.Context, deps ReaderDeps,
|
||||
) (interfaces.DocReader, error) {
|
||||
if deps.WeKnoraCloudCredentials == nil {
|
||||
return nil, errEngineUnavailable(WeKnoraCloudEngineName, "no credential resolver configured")
|
||||
}
|
||||
creds := deps.WeKnoraCloudCredentials(ctx)
|
||||
if creds == nil {
|
||||
return nil, errEngineUnavailable(WeKnoraCloudEngineName, "tenant credentials not configured")
|
||||
}
|
||||
return NewWeKnoraCloudSignedDocumentReader(creds.AppID, creds.AppSecret)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mineru — Go-native, calls self-hosted MinerU API directly
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type mineruEngine struct{}
|
||||
|
||||
func (e *mineruEngine) Name() string { return MinerUEngineName }
|
||||
|
||||
func (e *mineruEngine) Description() string { return "MinerU self-hosted service" }
|
||||
|
||||
func (e *mineruEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff", "doc", "docx", "ppt", "pptx"}
|
||||
}
|
||||
|
||||
func (e *mineruEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
endpoint := strings.TrimSpace(overrides["mineru_endpoint"])
|
||||
if endpoint == "" {
|
||||
return false, "MinerU service not configured"
|
||||
}
|
||||
return PingMinerU(endpoint)
|
||||
}
|
||||
|
||||
func (e *mineruEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return NewMinerUReader(deps.Overrides), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mineru_cloud — Go-native, calls MinerU Cloud API directly
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type mineruCloudEngine struct{}
|
||||
|
||||
func (e *mineruCloudEngine) Name() string { return MinerUCloudEngineName }
|
||||
|
||||
func (e *mineruCloudEngine) Description() string { return "MinerU Cloud API" }
|
||||
|
||||
func (e *mineruCloudEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff", "doc", "docx", "ppt", "pptx"}
|
||||
}
|
||||
|
||||
func (e *mineruCloudEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
apiKey := strings.TrimSpace(overrides["mineru_api_key"])
|
||||
if apiKey == "" {
|
||||
return false, "MinerU API Key not configured"
|
||||
}
|
||||
return PingMinerUCloud(apiKey)
|
||||
}
|
||||
|
||||
func (e *mineruCloudEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return NewMinerUCloudReader(deps.Overrides), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// paddleocr_vl — Go-native, calls a self-hosted PaddleOCR-VL pipeline service
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type paddleOCRVLEngine struct{}
|
||||
|
||||
func (e *paddleOCRVLEngine) Name() string { return PaddleOCRVLEngineName }
|
||||
|
||||
func (e *paddleOCRVLEngine) Description() string { return "PaddleOCR-VL self-hosted service" }
|
||||
|
||||
func (e *paddleOCRVLEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff"}
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
endpoint := strings.TrimSpace(overrides["paddleocr_vl_endpoint"])
|
||||
if endpoint == "" {
|
||||
return false, "PaddleOCR-VL service not configured"
|
||||
}
|
||||
return PingPaddleOCRVL(endpoint)
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return NewPaddleOCRVLReader(deps.Overrides), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// paddleocr_vl_cloud — Go-native, calls the PaddleOCR-VL AI Studio cloud API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type paddleOCRVLCloudEngine struct{}
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) Name() string { return PaddleOCRVLCloudEngineName }
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) Description() string { return "PaddleOCR-VL Cloud API" }
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff"}
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
token := strings.TrimSpace(overrides["paddleocr_vl_cloud_token"])
|
||||
if token == "" {
|
||||
return false, "PaddleOCR-VL Cloud Token not configured"
|
||||
}
|
||||
return PingPaddleOCRVLCloud(token)
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) NewReader(
|
||||
_ context.Context, deps ReaderDeps,
|
||||
) (interfaces.DocReader, error) {
|
||||
return NewPaddleOCRVLCloudReader(deps.Overrides), nil
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
docclient "github.com/Tencent/WeKnora/docreader/client"
|
||||
"github.com/Tencent/WeKnora/docreader/proto"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/resolver"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
func getMaxMessageSize() int {
|
||||
return int(utils.GetMaxFileSizeMB() * 1024 * 1024)
|
||||
}
|
||||
|
||||
// GRPCDocumentReader implements DocumentReader over gRPC.
|
||||
type GRPCDocumentReader struct {
|
||||
mu sync.RWMutex
|
||||
conn *grpc.ClientConn
|
||||
client proto.DocReaderClient
|
||||
addr string
|
||||
}
|
||||
|
||||
func NewGRPCDocumentReader(addr string) (*GRPCDocumentReader, error) {
|
||||
p := &GRPCDocumentReader{}
|
||||
if addr != "" {
|
||||
if err := p.connect(addr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) connect(addr string) error {
|
||||
authConfig := docclient.LoadAuthConfigFromEnv()
|
||||
opts, err := authConfig.BuildDialOptions(getMaxMessageSize())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to build docreader dial options: %w", err)
|
||||
}
|
||||
if authConfig.TLSEnabled {
|
||||
logger.Infof(context.Background(), "TLS enabled for docreader gRPC client")
|
||||
}
|
||||
if authConfig.AuthToken != "" {
|
||||
logger.Infof(context.Background(),
|
||||
"Token authentication enabled for docreader gRPC client (TLS=%v)",
|
||||
authConfig.TLSEnabled,
|
||||
)
|
||||
}
|
||||
|
||||
resolver.SetDefaultScheme("dns")
|
||||
|
||||
start := time.Now()
|
||||
conn, err := grpc.Dial("dns:///"+addr, opts...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to docreader: %w", err)
|
||||
}
|
||||
logger.Infof(context.Background(), "Connected to docreader in %v", time.Since(start))
|
||||
|
||||
p.conn = conn
|
||||
p.client = proto.NewDocReaderClient(conn)
|
||||
p.addr = addr
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) Reconnect(addr string) error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if p.conn != nil {
|
||||
_ = p.conn.Close()
|
||||
p.conn = nil
|
||||
p.client = nil
|
||||
p.addr = ""
|
||||
}
|
||||
return p.connect(addr)
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) IsConnected() bool {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
return p.conn != nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) Close() error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.conn != nil {
|
||||
return p.conn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var errNotConnected = fmt.Errorf("docreader service not connected")
|
||||
|
||||
func (p *GRPCDocumentReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
p.mu.RLock()
|
||||
client := p.client
|
||||
p.mu.RUnlock()
|
||||
if client == nil {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
|
||||
protoReq := &proto.ReadRequest{
|
||||
FileContent: req.FileContent,
|
||||
FileName: req.FileName,
|
||||
FileType: req.FileType,
|
||||
Url: req.URL,
|
||||
Title: req.Title,
|
||||
RequestId: req.RequestID,
|
||||
Config: &proto.ReadConfig{
|
||||
ParserEngine: req.ParserEngine,
|
||||
ParserEngineOverrides: req.ParserEngineOverrides,
|
||||
},
|
||||
}
|
||||
|
||||
// Use the streaming RPC so documents with many page images (large scanned
|
||||
// PDFs) are not capped by the unary message-size limit. The meta frame
|
||||
// arrives first, followed by one frame per image.
|
||||
result, err := p.readStream(ctx, client, protoReq)
|
||||
if err != nil {
|
||||
// An older docreader build may not implement ReadStream. Fall back to
|
||||
// the unary Read RPC so a version-skewed deployment still parses
|
||||
// documents (small/medium docs only — the unary path remains capped by
|
||||
// the gRPC message-size limit, which is exactly what streaming avoids).
|
||||
if status.Code(err) == codes.Unimplemented {
|
||||
logger.Warnf(ctx, "docreader ReadStream unimplemented, falling back to unary Read: %v", err)
|
||||
return p.readUnary(ctx, client, protoReq)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// readStream consumes the server-streaming ReadStream RPC: one meta frame
|
||||
// followed by one frame per image. Errors are returned verbatim so the caller
|
||||
// can inspect the gRPC status code (e.g. Unimplemented) for fallback.
|
||||
func (p *GRPCDocumentReader) readStream(
|
||||
ctx context.Context, client proto.DocReaderClient, protoReq *proto.ReadRequest,
|
||||
) (*types.ReadResult, error) {
|
||||
stream, err := client.ReadStream(ctx, protoReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gRPC ReadStream failed: %w", err)
|
||||
}
|
||||
|
||||
result := &types.ReadResult{}
|
||||
gotMeta := false
|
||||
for {
|
||||
frame, recvErr := stream.Recv()
|
||||
if recvErr == io.EOF {
|
||||
break
|
||||
}
|
||||
if recvErr != nil {
|
||||
return nil, fmt.Errorf("gRPC ReadStream recv failed: %w", recvErr)
|
||||
}
|
||||
|
||||
if meta := frame.GetMeta(); meta != nil {
|
||||
gotMeta = true
|
||||
result.MarkdownContent = meta.GetMarkdownContent()
|
||||
result.ImageDirPath = meta.GetImageDirPath()
|
||||
result.Metadata = meta.GetMetadata()
|
||||
result.Error = meta.GetError()
|
||||
if n := meta.GetImageCount(); n > 0 {
|
||||
result.ImageRefs = make([]types.ImageRef, 0, n)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if img := frame.GetImage(); img != nil {
|
||||
result.ImageRefs = append(result.ImageRefs, types.ImageRef{
|
||||
Filename: img.GetFilename(),
|
||||
OriginalRef: img.GetOriginalRef(),
|
||||
MimeType: img.GetMimeType(),
|
||||
StorageKey: img.GetStorageKey(),
|
||||
ImageData: img.GetImageData(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if !gotMeta {
|
||||
return nil, fmt.Errorf("gRPC ReadStream returned no metadata frame")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// readUnary calls the legacy unary Read RPC. Used only as a compatibility
|
||||
// fallback when the connected docreader does not implement ReadStream.
|
||||
func (p *GRPCDocumentReader) readUnary(
|
||||
ctx context.Context, client proto.DocReaderClient, protoReq *proto.ReadRequest,
|
||||
) (*types.ReadResult, error) {
|
||||
resp, err := client.Read(ctx, protoReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gRPC Read failed: %w", err)
|
||||
}
|
||||
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: resp.GetMarkdownContent(),
|
||||
ImageDirPath: resp.GetImageDirPath(),
|
||||
Metadata: resp.GetMetadata(),
|
||||
Error: resp.GetError(),
|
||||
}
|
||||
if refs := resp.GetImageRefs(); len(refs) > 0 {
|
||||
result.ImageRefs = make([]types.ImageRef, 0, len(refs))
|
||||
for _, img := range refs {
|
||||
result.ImageRefs = append(result.ImageRefs, types.ImageRef{
|
||||
Filename: img.GetFilename(),
|
||||
OriginalRef: img.GetOriginalRef(),
|
||||
MimeType: img.GetMimeType(),
|
||||
StorageKey: img.GetStorageKey(),
|
||||
ImageData: img.GetImageData(),
|
||||
})
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) ListEngines(ctx context.Context, overrides map[string]string) ([]types.ParserEngineInfo, error) {
|
||||
p.mu.RLock()
|
||||
client := p.client
|
||||
p.mu.RUnlock()
|
||||
if client == nil {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
|
||||
resp, err := client.ListEngines(ctx, &proto.ListEnginesRequest{ConfigOverrides: overrides})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gRPC ListEngines failed: %w", err)
|
||||
}
|
||||
|
||||
result := make([]types.ParserEngineInfo, 0, len(resp.GetEngines()))
|
||||
for _, e := range resp.GetEngines() {
|
||||
result = append(result, types.ParserEngineInfo{
|
||||
Name: e.GetName(),
|
||||
Description: e.GetDescription(),
|
||||
FileTypes: e.GetFileTypes(),
|
||||
Available: e.GetAvailable(),
|
||||
UnavailableReason: e.GetUnavailableReason(),
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// stringOr returns val (trimmed) if non-empty, otherwise fallback.
|
||||
func stringOr(val, fallback string) string {
|
||||
val = strings.TrimSpace(val)
|
||||
if val == "" {
|
||||
return fallback
|
||||
}
|
||||
return val
|
||||
}
|
||||
|
||||
// parseBoolOr parses a truthy string ("true","1","yes"), returning fallback on empty.
|
||||
func parseBoolOr(val string, fallback bool) bool {
|
||||
val = strings.ToLower(strings.TrimSpace(val))
|
||||
if val == "" {
|
||||
return fallback
|
||||
}
|
||||
return val == "true" || val == "1" || val == "yes"
|
||||
}
|
||||
|
||||
// firstNonEmpty returns the first non-empty string, or "" if all are empty.
|
||||
func firstNonEmpty(vals ...string) string {
|
||||
for _, v := range vals {
|
||||
if v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// sleepCtx sleeps for d but returns early if ctx is cancelled.
|
||||
func sleepCtx(ctx context.Context, d time.Duration) {
|
||||
t := time.NewTimer(d)
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
|
||||
// logResponseStructure recursively logs the structure of an API response,
|
||||
// truncating large string values. label identifies the subsystem (e.g. "MinerU").
|
||||
func logResponseStructure(label string, obj interface{}, prefix string) {
|
||||
switch v := obj.(type) {
|
||||
case map[string]interface{}:
|
||||
keys := make([]string, 0, len(v))
|
||||
for k := range v {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
logger.Infof(context.Background(), "[%s] %s = {object with %d keys: %s}", label, prefix, len(v), strings.Join(keys, ", "))
|
||||
for _, key := range keys {
|
||||
val := v[key]
|
||||
path := prefix + "." + key
|
||||
switch inner := val.(type) {
|
||||
case map[string]interface{}:
|
||||
logResponseStructure(label, inner, path)
|
||||
case []interface{}:
|
||||
logger.Infof(context.Background(), "[%s] %s = [array with %d items]", label, path, len(inner))
|
||||
if len(inner) > 0 {
|
||||
logger.Infof(context.Background(), "[%s] %s[0] type=%T", label, path, inner[0])
|
||||
if len(inner) <= 3 {
|
||||
for i, item := range inner {
|
||||
logResponseStructure(label, item, fmt.Sprintf("%s[%d]", path, i))
|
||||
}
|
||||
} else {
|
||||
logResponseStructure(label, inner[0], path+"[0]")
|
||||
logger.Infof(context.Background(), "[%s] ... and %d more items in %s", label, len(inner)-1, path)
|
||||
}
|
||||
}
|
||||
case string:
|
||||
if len(inner) > 200 {
|
||||
logger.Infof(context.Background(), "[%s] %s = string(%d chars): %.200s...", label, path, len(inner), inner)
|
||||
} else {
|
||||
logger.Infof(context.Background(), "[%s] %s = %q", label, path, inner)
|
||||
}
|
||||
case float64:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (number)", label, path, inner)
|
||||
case bool:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (bool)", label, path, inner)
|
||||
case nil:
|
||||
logger.Infof(context.Background(), "[%s] %s = null", label, path)
|
||||
default:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (%T)", label, path, val, val)
|
||||
}
|
||||
}
|
||||
case []interface{}:
|
||||
logger.Infof(context.Background(), "[%s] %s = [array with %d items]", label, prefix, len(v))
|
||||
if len(v) > 0 {
|
||||
if len(v) <= 3 {
|
||||
for i, item := range v {
|
||||
logResponseStructure(label, item, fmt.Sprintf("%s[%d]", prefix, i))
|
||||
}
|
||||
} else {
|
||||
logResponseStructure(label, v[0], prefix+"[0]")
|
||||
logger.Infof(context.Background(), "[%s] ... and %d more items in %s", label, len(v)-1, prefix)
|
||||
}
|
||||
}
|
||||
default:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (%T)", label, prefix, v, v)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/converter"
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/plugin/base"
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/plugin/commonmark"
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/plugin/table"
|
||||
)
|
||||
|
||||
var (
|
||||
// htmlTableBlockPattern matches a single (non-nested) <table>...</table>
|
||||
// block, the form OCR/layout engines such as PaddleOCR-VL emit tables in.
|
||||
htmlTableBlockPattern = regexp.MustCompile(`(?is)<table\b[^>]*>.*?</table>`)
|
||||
|
||||
// htmlLayoutAttrPattern matches presentational HTML attributes that carry
|
||||
// no semantic value (text-align styles, CSS classes, sizing). Structural
|
||||
// attributes like rowspan/colspan are intentionally excluded.
|
||||
htmlLayoutAttrPattern = regexp.MustCompile(
|
||||
`(?is)\s+(?:style|class|align|valign|width|height|bgcolor)\s*=\s*(?:"[^"]*"|'[^']*'|[^\s>]+)`,
|
||||
)
|
||||
|
||||
// htmlSpanAttrPattern detects rowspan/colspan, which Markdown tables cannot
|
||||
// represent; such tables keep their HTML form (attributes stripped) instead.
|
||||
htmlSpanAttrPattern = regexp.MustCompile(`(?i)\b(?:row|col)span\b`)
|
||||
|
||||
// markdownTableSeparatorPattern matches the |---|---| delimiter row that a
|
||||
// valid GFM table must contain.
|
||||
markdownTableSeparatorPattern = regexp.MustCompile(`(?m)^\s*\|?\s*:?-+:?\s*(?:\|\s*:?-+:?\s*)+\|?\s*$`)
|
||||
)
|
||||
|
||||
// normalizeHTMLTables rewrites inline HTML <table> blocks embedded in OCR
|
||||
// markdown output. PaddleOCR-VL emits tables as HTML with per-cell text-align
|
||||
// styles, which (1) waste tokens on layout markup and (2) are not recognized
|
||||
// by the chunker's table-protection logic, so large tables get split mid-row.
|
||||
//
|
||||
// Each table block is converted to a GFM Markdown table when possible. Tables
|
||||
// that use rowspan/colspan (which Markdown cannot express) fall back to having
|
||||
// their presentational attributes stripped so they stay intact as HTML.
|
||||
func normalizeHTMLTables(md string) string {
|
||||
if !strings.Contains(strings.ToLower(md), "<table") {
|
||||
return md
|
||||
}
|
||||
|
||||
conv := converter.NewConverter(
|
||||
converter.WithPlugins(
|
||||
base.NewBasePlugin(),
|
||||
commonmark.NewCommonmarkPlugin(),
|
||||
table.NewTablePlugin(),
|
||||
),
|
||||
)
|
||||
|
||||
return htmlTableBlockPattern.ReplaceAllStringFunc(md, func(block string) string {
|
||||
if htmlSpanAttrPattern.MatchString(block) {
|
||||
return stripHTMLLayoutAttrs(block)
|
||||
}
|
||||
converted, err := conv.ConvertString(block)
|
||||
if err != nil {
|
||||
return stripHTMLLayoutAttrs(block)
|
||||
}
|
||||
converted = strings.TrimSpace(converted)
|
||||
if converted == "" || !markdownTableSeparatorPattern.MatchString(converted) {
|
||||
return stripHTMLLayoutAttrs(block)
|
||||
}
|
||||
// Pad with blank lines so the Markdown table is a standalone block that
|
||||
// the chunker recognizes (and protects) as a table.
|
||||
return "\n\n" + converted + "\n\n"
|
||||
})
|
||||
}
|
||||
|
||||
// stripHTMLLayoutAttrs removes presentational attributes from an HTML fragment
|
||||
// while preserving structural attributes (rowspan/colspan) and text content.
|
||||
func stripHTMLLayoutAttrs(html string) string {
|
||||
return htmlLayoutAttrPattern.ReplaceAllString(html, "")
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizeHTMLTables_ConvertsStyledTableToMarkdown(t *testing.T) {
|
||||
// Mirrors PaddleOCR-VL output: a data table where every cell carries a
|
||||
// text-align style that wastes tokens.
|
||||
input := `# 报告
|
||||
|
||||
<table><tr><td style="text-align:center;">指标</td><td style="text-align:center;">数值</td></tr>` +
|
||||
`<tr><td style="text-align:center;">营收</td><td style="text-align:right;">10亿</td></tr>` +
|
||||
`<tr><td style="text-align:center;">利润</td><td style="text-align:right;">2.3亿</td></tr></table>
|
||||
|
||||
结尾。`
|
||||
|
||||
got := normalizeHTMLTables(input)
|
||||
|
||||
if strings.Contains(got, "<table") {
|
||||
t.Fatalf("expected HTML table to be converted away, got:\n%s", got)
|
||||
}
|
||||
if strings.Contains(got, "text-align") {
|
||||
t.Fatalf("expected style attributes removed, got:\n%s", got)
|
||||
}
|
||||
if !markdownTableSeparatorPattern.MatchString(got) {
|
||||
t.Fatalf("expected a Markdown table separator row, got:\n%s", got)
|
||||
}
|
||||
for _, want := range []string{"指标", "数值", "营收", "10亿", "利润", "2.3亿"} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Fatalf("expected converted table to retain %q, got:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(got, "# 报告") || !strings.Contains(got, "结尾。") {
|
||||
t.Fatalf("expected surrounding markdown to be preserved, got:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeHTMLTables_StripsAttrsOnSpanTables(t *testing.T) {
|
||||
// rowspan/colspan cannot be expressed in Markdown, so the table stays HTML
|
||||
// but its presentational attributes are stripped.
|
||||
input := `<table><tr><td colspan="2" style="text-align:center;" class="hdr">合计</td></tr>` +
|
||||
`<tr><td style="text-align:left;">A</td><td width="80">B</td></tr></table>`
|
||||
|
||||
got := normalizeHTMLTables(input)
|
||||
|
||||
if !strings.Contains(got, "<table") {
|
||||
t.Fatalf("expected span table to remain HTML, got:\n%s", got)
|
||||
}
|
||||
if !strings.Contains(got, `colspan="2"`) {
|
||||
t.Fatalf("expected colspan to be preserved, got:\n%s", got)
|
||||
}
|
||||
for _, banned := range []string{"text-align", "class=", "width="} {
|
||||
if strings.Contains(got, banned) {
|
||||
t.Fatalf("expected %q to be stripped, got:\n%s", banned, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeHTMLTables_NoTableUnchanged(t *testing.T) {
|
||||
input := "# 标题\n\n普通段落,没有表格。\n\n| a | b |\n| --- | --- |\n| 1 | 2 |"
|
||||
if got := normalizeHTMLTables(input); got != input {
|
||||
t.Fatalf("expected content without HTML tables to be unchanged, got:\n%s", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
PathRead = "/read"
|
||||
PathListEngines = "/list-engines"
|
||||
)
|
||||
|
||||
// --- JSON DTOs ---
|
||||
|
||||
type httpReadConfig struct {
|
||||
ParserEngine string `json:"parser_engine,omitempty"`
|
||||
ParserEngineOverrides map[string]string `json:"parser_engine_overrides,omitempty"`
|
||||
}
|
||||
|
||||
type httpReadRequest struct {
|
||||
FileContent string `json:"file_content,omitempty"` // base64
|
||||
FileName string `json:"file_name,omitempty"`
|
||||
FileType string `json:"file_type,omitempty"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Config *httpReadConfig `json:"config,omitempty"`
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
}
|
||||
|
||||
type httpImageRef struct {
|
||||
Filename string `json:"filename"`
|
||||
OriginalRef string `json:"original_ref"`
|
||||
MimeType string `json:"mime_type"`
|
||||
StorageKey string `json:"storage_key,omitempty"`
|
||||
ImageData []byte `json:"image_data,omitempty"`
|
||||
}
|
||||
|
||||
type httpReadResponse struct {
|
||||
MarkdownContent string `json:"markdown_content"`
|
||||
ImageRefs []httpImageRef `json:"image_refs,omitempty"`
|
||||
ImageDirPath string `json:"image_dir_path,omitempty"`
|
||||
Metadata map[string]string `json:"metadata,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// HTTPDocumentReader implements DocumentReader over HTTP/JSON.
|
||||
type HTTPDocumentReader struct {
|
||||
mu sync.RWMutex
|
||||
baseURL string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func NewHTTPDocumentReader(baseURL string) (*HTTPDocumentReader, error) {
|
||||
baseURL = strings.TrimSuffix(strings.TrimSpace(baseURL), "/")
|
||||
if baseURL != "" {
|
||||
if err := secutils.ValidateURLForSSRF(baseURL); err != nil {
|
||||
return nil, fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
}
|
||||
clientCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
|
||||
clientCfg.Timeout = 5 * time.Minute
|
||||
p := &HTTPDocumentReader{
|
||||
baseURL: baseURL,
|
||||
client: secutils.NewSSRFSafeHTTPClient(clientCfg),
|
||||
}
|
||||
if p.baseURL != "" {
|
||||
logger.Infof(context.Background(), "INFO: HTTP docreader base URL: %s", p.baseURL)
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) base() string {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
return p.baseURL
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) Reconnect(addr string) error {
|
||||
addr = strings.TrimSuffix(strings.TrimSpace(addr), "/")
|
||||
if addr != "" {
|
||||
if err := secutils.ValidateURLForSSRF(addr); err != nil {
|
||||
return fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.baseURL = addr
|
||||
logger.Infof(context.Background(), "INFO: HTTP docreader base URL set to %s", p.baseURL)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) IsConnected() bool {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
return p.baseURL != ""
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) Close() error { return nil }
|
||||
|
||||
type httpListEnginesRequest struct {
|
||||
ConfigOverrides map[string]string `json:"config_overrides,omitempty"`
|
||||
}
|
||||
|
||||
type httpParserEngineInfo struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
FileTypes []string `json:"file_types"`
|
||||
Available bool `json:"available"`
|
||||
UnavailableReason string `json:"unavailable_reason,omitempty"`
|
||||
}
|
||||
|
||||
type httpListEnginesResponse struct {
|
||||
Engines []httpParserEngineInfo `json:"engines"`
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) ListEngines(ctx context.Context, overrides map[string]string) ([]types.ParserEngineInfo, error) {
|
||||
base := p.base()
|
||||
if base == "" {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
if err := secutils.ValidateURLForSSRF(base); err != nil {
|
||||
return nil, fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
|
||||
body := httpListEnginesRequest{ConfigOverrides: overrides}
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http marshal list-engines request: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, base+PathListEngines, bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http new request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http list-engines failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBytes, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("http list-engines status %d: %s", resp.StatusCode, string(respBytes))
|
||||
}
|
||||
|
||||
var out httpListEnginesResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, fmt.Errorf("http decode list-engines response: %w", err)
|
||||
}
|
||||
|
||||
result := make([]types.ParserEngineInfo, 0, len(out.Engines))
|
||||
for _, e := range out.Engines {
|
||||
result = append(result, types.ParserEngineInfo{
|
||||
Name: e.Name,
|
||||
Description: e.Description,
|
||||
FileTypes: e.FileTypes,
|
||||
Available: e.Available,
|
||||
UnavailableReason: e.UnavailableReason,
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func fromHTTPReadResponse(resp *httpReadResponse) *types.ReadResult {
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: resp.MarkdownContent,
|
||||
ImageDirPath: resp.ImageDirPath,
|
||||
Metadata: resp.Metadata,
|
||||
Error: resp.Error,
|
||||
}
|
||||
for _, ref := range resp.ImageRefs {
|
||||
result.ImageRefs = append(result.ImageRefs, types.ImageRef{
|
||||
Filename: ref.Filename,
|
||||
OriginalRef: ref.OriginalRef,
|
||||
MimeType: ref.MimeType,
|
||||
StorageKey: ref.StorageKey,
|
||||
ImageData: ref.ImageData,
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
base := p.base()
|
||||
if base == "" {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
if err := secutils.ValidateURLForSSRF(base); err != nil {
|
||||
return nil, fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
|
||||
body := httpReadRequest{
|
||||
FileName: req.FileName,
|
||||
FileType: req.FileType,
|
||||
URL: req.URL,
|
||||
Title: req.Title,
|
||||
RequestID: req.RequestID,
|
||||
Config: &httpReadConfig{
|
||||
ParserEngine: req.ParserEngine,
|
||||
ParserEngineOverrides: req.ParserEngineOverrides,
|
||||
},
|
||||
}
|
||||
if len(req.FileContent) > 0 {
|
||||
body.FileContent = base64.StdEncoding.EncodeToString(req.FileContent)
|
||||
}
|
||||
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http marshal read request: %w", err)
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, base+PathRead, bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http new request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.ContentLength = int64(len(jsonBody))
|
||||
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http read failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("http read status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
var out httpReadResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, fmt.Errorf("http decode read response: %w", err)
|
||||
}
|
||||
return fromHTTPReadResponse(&out), nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,80 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestResolveAndStoreRelativeHTMLImages(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: `<table><tr><td><img src="images/profile.jpg" alt="profile"/></td></tr></table>`,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "profile.jpg",
|
||||
OriginalRef: "images/profile.jpg",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 stored image but got %d", len(imgs))
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes to be called once but got %d", len(svc.saved))
|
||||
}
|
||||
if strings.Contains(out, `src="images/profile.jpg"`) {
|
||||
t.Fatalf("expected relative html img src to be replaced, got: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `src="local://test/`) {
|
||||
t.Fatalf("expected stored file url in html img src, got: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveAndStoreDedupsSameImageRefVariants ensures that when the same
|
||||
// MinerU image is referenced under multiple path forms (e.g. "images/foo.png"
|
||||
// and "./images/foo.png"), the bytes are uploaded to object storage only once.
|
||||
func TestResolveAndStoreDedupsSameImageRefVariants(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: strings.Join([]string{
|
||||
"",
|
||||
`<img src="./images/foo.png"/>`,
|
||||
}, "\n"),
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "foo.png",
|
||||
OriginalRef: "images/foo.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
{
|
||||
Filename: "foo.png",
|
||||
OriginalRef: "./images/foo.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
_, _, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes to be called once for duplicate refs, got %d", len(svc.saved))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,600 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/png"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// createTestPNG generates a minimal PNG image with the given dimensions.
|
||||
func createTestPNG(w, h int) []byte {
|
||||
img := image.NewRGBA(image.Rect(0, 0, w, h))
|
||||
for y := 0; y < h; y++ {
|
||||
for x := 0; x < w; x++ {
|
||||
img.Set(x, y, color.RGBA{R: 128, G: 128, B: 128, A: 255})
|
||||
}
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
_ = png.Encode(&buf, img)
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestIsIconImage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
expect bool
|
||||
}{
|
||||
{
|
||||
name: "tiny bytes (< 512B)",
|
||||
data: make([]byte, 256),
|
||||
expect: true,
|
||||
},
|
||||
{
|
||||
name: "small icon 32x32",
|
||||
data: createTestPNG(32, 32),
|
||||
expect: true,
|
||||
},
|
||||
{
|
||||
name: "small icon 48x48",
|
||||
data: createTestPNG(48, 48),
|
||||
expect: true,
|
||||
},
|
||||
{
|
||||
name: "borderline 64x64",
|
||||
data: createTestPNG(64, 64),
|
||||
expect: false,
|
||||
},
|
||||
{
|
||||
name: "normal image 200x150",
|
||||
data: createTestPNG(200, 150),
|
||||
expect: false,
|
||||
},
|
||||
{
|
||||
name: "wide but short 200x30",
|
||||
data: createTestPNG(200, 30),
|
||||
expect: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := isIconImage(tt.data)
|
||||
if got != tt.expect {
|
||||
t.Errorf("isIconImage() = %v, want %v (data len=%d)", got, tt.expect, len(tt.data))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// captureSaveBytes implements interfaces.FileService for tests; only SaveBytes records data.
|
||||
type captureSaveBytes struct {
|
||||
saved [][]byte
|
||||
urls []string
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) CheckConnectivity(context.Context) error { return nil }
|
||||
|
||||
func (c *captureSaveBytes) SaveFile(context.Context, *multipart.FileHeader, uint64, string) (string, error) {
|
||||
return "", fmt.Errorf("not implemented")
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) SaveBytes(_ context.Context, data []byte, _ uint64, fileName string, _ bool) (string, error) {
|
||||
c.saved = append(c.saved, append([]byte(nil), data...))
|
||||
u := "local://test/" + fileName
|
||||
c.urls = append(c.urls, u)
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) GetFile(context.Context, string) (io.ReadCloser, error) {
|
||||
return nil, fmt.Errorf("not implemented")
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) GetFileURL(context.Context, string) (string, error) { return "", nil }
|
||||
|
||||
func (c *captureSaveBytes) DeleteFile(context.Context, string) error { return nil }
|
||||
|
||||
func (c *captureSaveBytes) CopyFile(context.Context, string, uint64, string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
var _ interfaces.FileService = (*captureSaveBytes)(nil)
|
||||
|
||||
func TestResolveDataURIImages(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := "pre  post"
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("got %d images", len(imgs))
|
||||
}
|
||||
if len(svc.saved) != 1 || !bytes.Equal(svc.saved[0], png) {
|
||||
t.Fatal("SaveBytes payload mismatch")
|
||||
}
|
||||
if !strings.Contains(out, "local://test/") || strings.Contains(out, "data:image") {
|
||||
t.Fatalf("markdown: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveDataURIImages_CaseInsensitive(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := ""
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || !strings.Contains(out, "local://test/") {
|
||||
t.Fatalf("imgs=%d out=%q", len(imgs), out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveDataURIImages_AltTextWithBracket(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := `pre ![C:\Users\Frank\RichOle\YAJ]7D)SW_W6.png](data:image/png;base64,` + b64 + `) post`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image but got %d; alt text with ] was not handled", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 content still present in output: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveDataURIImages_XEmfMimeType(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := ``
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image for x-emf mime type but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "data:image") {
|
||||
t.Fatalf("data URI still present for x-emf: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveHTMLDataURIImages(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := `text <img src="data:image/png;base64,` + b64 + `" alt="test"> more`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveHTMLDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image from HTML img tag but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "<img") || strings.Contains(out, "base64") {
|
||||
t.Fatalf("HTML img tag or base64 still present: %s", out[:min(200, len(out))])
|
||||
}
|
||||
if !strings.Contains(out, " {
|
||||
t.Fatalf("expected markdown image replacement, got: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveBareBase64Content_BareDataURI(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := `some text data:image/png;base64,` + b64 + ` more text`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveBareBase64Content(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image from bare data URI but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 content still present: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveBareBase64Content_InsideBrokenMarkdownRef(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
// Simulate a broken markdown ref that the primary regex already handled
|
||||
// If primary regex handles it, this won't reach the bare handler
|
||||
// But test the bare handler's behavior when prev char is '('
|
||||
md := `text (data:image/png;base64,` + b64 + `) more`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveBareBase64Content(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 still present: %s", out[:min(200, len(out))])
|
||||
}
|
||||
// When preceded by '(', should replace with just URL (not wrapped in ![image]())
|
||||
if strings.Contains(out, "![image]") {
|
||||
t.Fatalf("should not wrap in ![image]() when inside parentheses: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStore_MultipleFormats(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
// Markdown with: standard data URI, alt text with ], x-emf mime type
|
||||
md := strings.Join([]string{
|
||||
`# Document`,
|
||||
``,
|
||||
`![C:\path\file]name.png](data:image/png;base64,` + b64 + `)`,
|
||||
``,
|
||||
}, "\n\n")
|
||||
|
||||
result := &types.ReadResult{MarkdownContent: md}
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 3 {
|
||||
t.Fatalf("expected 3 images but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 content still in output after ResolveAndStore")
|
||||
}
|
||||
if strings.Contains(out, "data:image") {
|
||||
t.Fatalf("data URI still in output after ResolveAndStore")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithTitle(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image but got %d", len(imgs))
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes to be called once but got %d", len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, ` "图片")`) {
|
||||
t.Fatalf("markdown image title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
if strings.Contains(out, "images/test.png") {
|
||||
t.Fatalf("original image path was not replaced: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithSingleQuotedTitle(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, ` '图片')`) {
|
||||
t.Fatalf("markdown image title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithSpacedFilename(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "第 1 页.png",
|
||||
OriginalRef: "images/第 1 页.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` {
|
||||
t.Fatalf("spaced filename path was not replaced: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageTitleContainingRightParen(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ` 图片")`,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, ` "阶段 1) 图片")`) {
|
||||
t.Fatalf("right-paren title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithMultilineTitle(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: "",
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, " || !strings.Contains(out, "\n \"图片说明\")") {
|
||||
t.Fatalf("multiline title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithAngleDestination(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: `.png> "阶段 1) 图片")`,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "第 1 页 (测试).png",
|
||||
OriginalRef: "images/第 1 页 (测试).png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, `> "阶段 1) 图片")`) {
|
||||
t.Fatalf("angle destination wrapper or title was not preserved: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreDuplicateReferencesStoreOnce(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: strings.Join([]string{
|
||||
``,
|
||||
``,
|
||||
}, "\n\n"),
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected duplicate references to save once but got %d", len(svc.saved))
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected one stored image record but got %d", len(imgs))
|
||||
}
|
||||
if strings.Count(out, "local://test/") != 2 {
|
||||
t.Fatalf("expected both markdown references to be replaced: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"第一处"`) || !strings.Contains(out, `"第二处"`) {
|
||||
t.Fatalf("duplicate reference titles were not preserved: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreUnknownReferenceUnchanged(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 0 || len(imgs) != 0 {
|
||||
t.Fatalf("unknown reference should not be saved, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if out != result.MarkdownContent {
|
||||
t.Fatalf("unknown reference changed: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreSharedMHTMLContract(t *testing.T) {
|
||||
type contractImage struct {
|
||||
Filename string `json:"filename"`
|
||||
OriginalRef string `json:"original_ref"`
|
||||
MimeType string `json:"mime_type"`
|
||||
ImageDataBase64 string `json:"image_data_base64"`
|
||||
}
|
||||
var contract struct {
|
||||
MarkdownContent string `json:"markdown_content"`
|
||||
Images []contractImage `json:"images"`
|
||||
}
|
||||
|
||||
contractPath := filepath.Join("..", "..", "..", "testdata", "mhtml", "titled-image-contract.json")
|
||||
raw, err := os.ReadFile(contractPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := json.Unmarshal(raw, &contract); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(contract.Images) != 1 {
|
||||
t.Fatalf("expected one contract image but got %d", len(contract.Images))
|
||||
}
|
||||
|
||||
imageContract := contract.Images[0]
|
||||
imageData, err := base64.StdEncoding.DecodeString(imageContract.ImageDataBase64)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: contract.MarkdownContent,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: imageContract.Filename,
|
||||
OriginalRef: imageContract.OriginalRef,
|
||||
MimeType: imageContract.MimeType,
|
||||
ImageData: imageData,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes once but got %d", len(svc.saved))
|
||||
}
|
||||
if !bytes.Equal(svc.saved[0], imageData) {
|
||||
t.Fatal("stored image bytes do not match contract fixture")
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected one stored image record but got %d", len(imgs))
|
||||
}
|
||||
if imgs[0].OriginalRef != imageContract.OriginalRef || imgs[0].MimeType != imageContract.MimeType {
|
||||
t.Fatalf("unexpected stored image metadata: %#v", imgs[0])
|
||||
}
|
||||
if strings.Contains(out, imageContract.OriginalRef) {
|
||||
t.Fatalf("original image path remained in output: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, ` {
|
||||
t.Fatalf("stored URL missing from markdown: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"阶段 1) 图片"`) {
|
||||
t.Fatalf("image title was not preserved: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "| 赛季制建立 | BP、Rank |\n\n {
|
||||
t.Fatalf("table-to-image block boundary changed: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"阶段 1) 图片")`+"\n\n高机动性身法与独特枪械反馈") {
|
||||
t.Fatalf("image-to-body block boundary changed: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "alpha \nbeta") {
|
||||
t.Fatalf("markdown hard break changed: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "```\nline1\n\n\nline2\n```") {
|
||||
t.Fatalf("fenced code blank lines changed: %s", out)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// defaultJSONChunkSize is the target chunk size in bytes for JSON semantic
|
||||
// splitting. Approximates 512 tokens (1 token ≈ 3-4 bytes for mixed content).
|
||||
const defaultJSONChunkSize = 1536
|
||||
|
||||
// minJSONChunkSize is the minimum chunk size. A new chunk is only started
|
||||
// when the current chunk has reached at least this size.
|
||||
var minJSONChunkSize = defaultJSONChunkSize - 200
|
||||
|
||||
// jsonToMarkdown converts raw JSON bytes into markdown text
|
||||
//
|
||||
// Key properties:
|
||||
// - Every output chunk is a **valid JSON object** (not a fragment).
|
||||
// - Nested paths from root to leaf are **fully preserved** in each chunk.
|
||||
// - Arrays are converted to index-keyed dicts so the algorithm is uniform.
|
||||
// - Small objects that fit within maxChunkSize are kept intact (not split).
|
||||
// - The output is a series of fenced ```json code blocks separated by \n\n,
|
||||
// which the downstream text chunker can split at block boundaries.
|
||||
func jsonToMarkdown(data []byte) (string, error) {
|
||||
data = trimBOM(data)
|
||||
if len(data) == 0 {
|
||||
return "", fmt.Errorf("empty JSON content")
|
||||
}
|
||||
if !json.Valid(data) {
|
||||
return "", fmt.Errorf("invalid JSON content")
|
||||
}
|
||||
|
||||
var parsed interface{}
|
||||
if err := json.Unmarshal(data, &parsed); err != nil {
|
||||
return "", fmt.Errorf("failed to parse JSON: %w", err)
|
||||
}
|
||||
|
||||
// Normalize: convert top-level arrays to index-keyed dicts
|
||||
normalized := listToDictPreprocess(parsed)
|
||||
|
||||
// If the whole thing fits in one chunk, just format it
|
||||
wholeSize := jsonSize(normalized)
|
||||
if wholeSize <= defaultJSONChunkSize {
|
||||
formatted := formatValue(normalized)
|
||||
return wrapCodeBlock(formatted), nil
|
||||
}
|
||||
|
||||
// Recursive split
|
||||
chunks := recursiveJSONSplit(normalized, nil, nil)
|
||||
|
||||
// Convert each chunk dict to a fenced code block
|
||||
blocks := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
if len(chunk) == 0 {
|
||||
continue
|
||||
}
|
||||
blocks = append(blocks, wrapCodeBlock(formatValue(chunk)))
|
||||
}
|
||||
|
||||
if len(blocks) == 0 {
|
||||
return wrapCodeBlock(formatValue(normalized)), nil
|
||||
}
|
||||
return strings.Join(blocks, "\n\n"), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// RecursiveJsonSplitter core algorithm
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// recursiveJSONSplit splits a JSON dict into a list of JSON dicts,
|
||||
// each fitting within defaultJSONChunkSize. It preserves the full nested
|
||||
// path from root to each leaf by using setNestedDict.
|
||||
//
|
||||
// This is a Go port of LangChain's RecursiveJsonSplitter._json_split.
|
||||
func recursiveJSONSplit(
|
||||
data interface{},
|
||||
currentPath []string,
|
||||
chunks []map[string]interface{},
|
||||
) []map[string]interface{} {
|
||||
if chunks == nil {
|
||||
chunks = []map[string]interface{}{{}}
|
||||
}
|
||||
|
||||
dict, ok := data.(map[string]interface{})
|
||||
if !ok {
|
||||
// Scalar or already-processed value: place it at the current path
|
||||
if len(currentPath) > 0 && len(chunks) > 0 {
|
||||
setNestedDict(chunks[len(chunks)-1], currentPath, data)
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
// Sort keys for deterministic output
|
||||
keys := sortedKeys(dict)
|
||||
|
||||
for _, key := range keys {
|
||||
value := dict[key]
|
||||
newPath := append(append([]string{}, currentPath...), key)
|
||||
|
||||
// Measure sizes
|
||||
chunkSize := jsonSize(chunks[len(chunks)-1])
|
||||
itemSize := jsonSize(map[string]interface{}{key: value})
|
||||
remaining := defaultJSONChunkSize - chunkSize
|
||||
|
||||
if itemSize <= remaining {
|
||||
// Item fits in the current chunk — add it preserving the path
|
||||
setNestedDict(chunks[len(chunks)-1], newPath, value)
|
||||
} else {
|
||||
// Item doesn't fit
|
||||
if chunkSize >= minJSONChunkSize {
|
||||
// Current chunk is big enough, start a new one
|
||||
chunks = append(chunks, map[string]interface{}{})
|
||||
}
|
||||
|
||||
// Check if the value itself is a dict/list that can be recursed into
|
||||
normalized := listToDictPreprocess(value)
|
||||
if subDict, isDict := normalized.(map[string]interface{}); isDict && canSplitDict(subDict) {
|
||||
// Recurse into the sub-object
|
||||
chunks = recursiveJSONSplit(subDict, newPath, chunks)
|
||||
} else {
|
||||
// Cannot split further (scalar or single-key dict) — place as-is
|
||||
setNestedDict(chunks[len(chunks)-1], newPath, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// setNestedDict sets a value in a nested dict structure, creating
|
||||
// intermediate dicts as needed. This preserves the full JSON path.
|
||||
//
|
||||
// Example: setNestedDict(d, ["config","db","host"], "localhost")
|
||||
// produces: {"config": {"db": {"host": "localhost"}}}
|
||||
func setNestedDict(d map[string]interface{}, path []string, value interface{}) {
|
||||
if len(path) == 0 {
|
||||
return
|
||||
}
|
||||
current := d
|
||||
for _, key := range path[:len(path)-1] {
|
||||
next, ok := current[key]
|
||||
if !ok {
|
||||
next = map[string]interface{}{}
|
||||
current[key] = next
|
||||
}
|
||||
if nextDict, ok := next.(map[string]interface{}); ok {
|
||||
current = nextDict
|
||||
} else {
|
||||
// Path conflict (existing value is not a dict) — overwrite
|
||||
newDict := map[string]interface{}{}
|
||||
current[key] = newDict
|
||||
current = newDict
|
||||
}
|
||||
}
|
||||
current[path[len(path)-1]] = value
|
||||
}
|
||||
|
||||
// listToDictPreprocess recursively converts JSON arrays to index-keyed
|
||||
// dicts so the splitter can treat everything uniformly.
|
||||
//
|
||||
// Example: ["a","b","c"] → {"0":"a", "1":"b", "2":"c"}
|
||||
func listToDictPreprocess(data interface{}) interface{} {
|
||||
switch v := data.(type) {
|
||||
case map[string]interface{}:
|
||||
result := make(map[string]interface{}, len(v))
|
||||
for k, val := range v {
|
||||
result[k] = listToDictPreprocess(val)
|
||||
}
|
||||
return result
|
||||
case []interface{}:
|
||||
result := make(map[string]interface{}, len(v))
|
||||
for i, item := range v {
|
||||
result[fmt.Sprintf("%d", i)] = listToDictPreprocess(item)
|
||||
}
|
||||
return result
|
||||
default:
|
||||
return data
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// jsonSize returns the serialized JSON byte length of a value.
|
||||
func jsonSize(v interface{}) int {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return len(b)
|
||||
}
|
||||
|
||||
// formatValue formats a JSON value with indentation.
|
||||
func formatValue(v interface{}) string {
|
||||
b, err := json.MarshalIndent(v, "", " ")
|
||||
if err != nil {
|
||||
b, _ = json.Marshal(v)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// indentJSON formats raw JSON bytes with indentation.
|
||||
func indentJSON(data []byte) (string, error) {
|
||||
var buf bytes.Buffer
|
||||
if err := json.Indent(&buf, data, "", " "); err != nil {
|
||||
return string(data), err
|
||||
}
|
||||
return buf.String(), nil
|
||||
}
|
||||
|
||||
// wrapCodeBlock wraps content in a fenced JSON code block.
|
||||
func wrapCodeBlock(content string) string {
|
||||
return "```json\n" + content + "\n```"
|
||||
}
|
||||
|
||||
// trimBOM removes a UTF-8 BOM prefix if present.
|
||||
func trimBOM(data []byte) []byte {
|
||||
if len(data) >= 3 && data[0] == 0xEF && data[1] == 0xBB && data[2] == 0xBF {
|
||||
return data[3:]
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
// canSplitDict returns true if a dict can be meaningfully split.
|
||||
// A dict with multiple keys can be split by distributing keys across chunks.
|
||||
// A dict with a single key can still be split if its value is a splittable dict.
|
||||
func canSplitDict(d map[string]interface{}) bool {
|
||||
if len(d) > 1 {
|
||||
return true
|
||||
}
|
||||
if len(d) == 1 {
|
||||
for _, v := range d {
|
||||
if sub, ok := v.(map[string]interface{}); ok && len(sub) > 1 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// sortedKeys returns the keys of a map in sorted order.
|
||||
// When all keys are numeric strings (from array-to-dict conversion),
|
||||
// sorts numerically so "2" comes before "10".
|
||||
func sortedKeys(m map[string]interface{}) []string {
|
||||
keys := make([]string, 0, len(m))
|
||||
allNumeric := true
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
if allNumeric {
|
||||
if _, err := strconv.Atoi(k); err != nil {
|
||||
allNumeric = false
|
||||
}
|
||||
}
|
||||
}
|
||||
if allNumeric {
|
||||
sort.Slice(keys, func(i, j int) bool {
|
||||
ni, _ := strconv.Atoi(keys[i])
|
||||
nj, _ := strconv.Atoi(keys[j])
|
||||
return ni < nj
|
||||
})
|
||||
} else {
|
||||
sort.Strings(keys)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,705 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Helper: visual test reporter
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func visualReport(t *testing.T, label string, input string, result string, err error, checks []checkResult) {
|
||||
t.Helper()
|
||||
|
||||
const width = 72
|
||||
bar := strings.Repeat("─", width)
|
||||
thickBar := strings.Repeat("━", width)
|
||||
|
||||
t.Logf("\n%s", thickBar)
|
||||
t.Logf("📋 TEST: %s", label)
|
||||
t.Logf("%s", bar)
|
||||
|
||||
// ── Input ──
|
||||
inputPreview := input
|
||||
if len(inputPreview) > 200 {
|
||||
inputPreview = inputPreview[:200] + fmt.Sprintf("... (%d bytes total)", len(input))
|
||||
}
|
||||
t.Logf("📥 INPUT (%d bytes):", len(input))
|
||||
for _, line := range strings.Split(inputPreview, "\n") {
|
||||
t.Logf(" %s", line)
|
||||
}
|
||||
t.Logf("%s", bar)
|
||||
|
||||
// ── Output / Error ──
|
||||
if err != nil {
|
||||
t.Logf("❌ ERROR: %v", err)
|
||||
} else {
|
||||
blockCount := strings.Count(result, "```json")
|
||||
totalLen := len(result)
|
||||
t.Logf("📤 OUTPUT: %d code block(s), %d bytes total", blockCount, totalLen)
|
||||
t.Logf("%s", bar)
|
||||
|
||||
blocks := splitBlocks(result)
|
||||
for i, blk := range blocks {
|
||||
preview := blk
|
||||
lines := strings.Split(preview, "\n")
|
||||
if len(lines) > 10 {
|
||||
preview = strings.Join(lines[:4], "\n") +
|
||||
fmt.Sprintf("\n ... (%d lines omitted) ...\n", len(lines)-8) +
|
||||
strings.Join(lines[len(lines)-4:], "\n")
|
||||
}
|
||||
t.Logf(" 📦 Block #%d (%d bytes):", i+1, len(blk))
|
||||
for _, line := range strings.Split(preview, "\n") {
|
||||
t.Logf(" │ %s", line)
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Logf("%s", bar)
|
||||
|
||||
// ── Checks ──
|
||||
allPass := true
|
||||
for _, c := range checks {
|
||||
icon := "✅"
|
||||
if !c.pass {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
}
|
||||
t.Logf("%s %s", icon, c.desc)
|
||||
if !c.pass {
|
||||
t.Errorf("FAIL: %s", c.desc)
|
||||
}
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("%s\n", thickBar)
|
||||
}
|
||||
|
||||
// splitBlocks extracts the content between ```json and ``` fences.
|
||||
func splitBlocks(result string) []string {
|
||||
var blocks []string
|
||||
rest := result
|
||||
for {
|
||||
start := strings.Index(rest, "```json\n")
|
||||
if start < 0 {
|
||||
break
|
||||
}
|
||||
rest = rest[start+len("```json\n"):]
|
||||
end := strings.Index(rest, "\n```")
|
||||
if end < 0 {
|
||||
blocks = append(blocks, rest)
|
||||
break
|
||||
}
|
||||
blocks = append(blocks, rest[:end])
|
||||
rest = rest[end+len("\n```"):]
|
||||
}
|
||||
return blocks
|
||||
}
|
||||
|
||||
type checkResult struct {
|
||||
desc string
|
||||
pass bool
|
||||
}
|
||||
|
||||
func check(desc string, pass bool) checkResult { return checkResult{desc, pass} }
|
||||
|
||||
// blockIsValidJSON checks that each code block inside the result is valid JSON.
|
||||
func allBlocksValidJSON(result string) bool {
|
||||
for _, blk := range splitBlocks(result) {
|
||||
if !json.Valid([]byte(blk)) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 1: Small inputs — kept intact, NOT split
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_SmallObject(t *testing.T) {
|
||||
input := `{"name": "test", "version": "1.0"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Small Object → single block, kept intact", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains key 'name'", strings.Contains(result, `"name"`)),
|
||||
check("contains key 'version'", strings.Contains(result, `"version"`)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
check("both keys in SAME block (not split apart)",
|
||||
strings.Contains(result, `"name"`) && strings.Contains(result, `"version"`) && blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SmallArray(t *testing.T) {
|
||||
input := `[{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Small Array → single block, kept intact", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_EmptyObject(t *testing.T) {
|
||||
input := `{}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Empty Object {}", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains '{}'", strings.Contains(result, "{}")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_EmptyArray(t *testing.T) {
|
||||
input := `[]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Empty Array []", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains '{}'", strings.Contains(result, "{}")), // [] → {} after list-to-dict
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveString(t *testing.T) {
|
||||
input := `"hello world"`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive String", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains 'hello world'", strings.Contains(result, "hello world")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveNumber(t *testing.T) {
|
||||
input := `42`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive Number", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains '42'", strings.Contains(result, "42")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveBoolean(t *testing.T) {
|
||||
input := `true`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive Boolean", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains 'true'", strings.Contains(result, "true")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveNull(t *testing.T) {
|
||||
input := `null`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive Null", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains 'null'", strings.Contains(result, "null")),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 2: Large data — smart recursive splitting
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_LargeObject(t *testing.T) {
|
||||
obj := make(map[string]interface{})
|
||||
for i := 0; i < 50; i++ {
|
||||
obj[fmt.Sprintf("key_%02d", i)] = strings.Repeat("value", 50)
|
||||
}
|
||||
data, _ := json.Marshal(obj)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Large Object (50 keys) → multiple blocks", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_LargeArray(t *testing.T) {
|
||||
arr := make([]interface{}, 100)
|
||||
for i := range arr {
|
||||
arr[i] = map[string]interface{}{
|
||||
"id": i,
|
||||
"name": strings.Repeat("name", 20),
|
||||
"description": strings.Repeat("desc", 30),
|
||||
}
|
||||
}
|
||||
data, _ := json.Marshal(arr)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Large Array (100 elements) → multiple blocks", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_RecursiveSplitLargeKey(t *testing.T) {
|
||||
inner := make(map[string]interface{})
|
||||
for i := 0; i < 30; i++ {
|
||||
inner[fmt.Sprintf("sub_%02d", i)] = strings.Repeat("val", 50)
|
||||
}
|
||||
obj := map[string]interface{}{"bigkey": inner}
|
||||
data, _ := json.Marshal(obj)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Recursive Split: 1 big key with 30 sub-keys", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
check("path preserved: chunks contain 'bigkey'", strings.Contains(result, "bigkey")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PathPreservation(t *testing.T) {
|
||||
// Core test: verify that nested paths are preserved in each chunk
|
||||
inner := make(map[string]interface{})
|
||||
for i := 0; i < 40; i++ {
|
||||
inner[fmt.Sprintf("field_%02d", i)] = strings.Repeat("data", 60)
|
||||
}
|
||||
obj := map[string]interface{}{
|
||||
"config": map[string]interface{}{
|
||||
"database": inner,
|
||||
},
|
||||
}
|
||||
data, _ := json.Marshal(obj)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := splitBlocks(result)
|
||||
// Every block should contain the full path "config" → "database"
|
||||
allHavePath := true
|
||||
for _, blk := range blocks {
|
||||
if !strings.Contains(blk, `"config"`) || !strings.Contains(blk, `"database"`) {
|
||||
allHavePath = false
|
||||
break
|
||||
}
|
||||
}
|
||||
visualReport(t, "Path Preservation: config.database.* across chunks", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", len(blocks)), len(blocks) >= 2),
|
||||
check("ALL blocks preserve path: config → database", allHavePath),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_MixedArrayElements(t *testing.T) {
|
||||
arr := []interface{}{
|
||||
map[string]interface{}{"id": 1},
|
||||
map[string]interface{}{"id": 2},
|
||||
map[string]interface{}{"id": 3, "data": strings.Repeat("x", 2000)},
|
||||
map[string]interface{}{"id": 4},
|
||||
}
|
||||
data, _ := json.Marshal(arr)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Mixed Array: 3 small + 1 large element", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("at least 2 blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_DeepNested(t *testing.T) {
|
||||
input := `{"l1": {"l2": {"l3": {"data": "` + strings.Repeat("x", 2000) + `"}}}}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Deep Nested (3 levels, large leaf)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("block count ≥ 1 (got %d)", blocks), blocks >= 1),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 3: Error Handling
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_InvalidJSON(t *testing.T) {
|
||||
input := `{invalid json}`
|
||||
_, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Invalid JSON → error", input, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
check("error mentions 'invalid JSON'",
|
||||
err != nil && strings.Contains(err.Error(), "invalid JSON")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_EmptyInput(t *testing.T) {
|
||||
input := ""
|
||||
_, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Empty Input → error", input, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
check("error mentions 'empty'",
|
||||
err != nil && strings.Contains(err.Error(), "empty")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_WhitespaceOnly(t *testing.T) {
|
||||
input := " \n\t "
|
||||
_, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Whitespace-Only Input → error", input, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 4: Edge Cases & Encoding
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_BOM(t *testing.T) {
|
||||
raw := append([]byte{0xEF, 0xBB, 0xBF}, []byte(`{"key": "value"}`)...)
|
||||
result, err := jsonToMarkdown(raw)
|
||||
visualReport(t, "UTF-8 BOM prefix → stripped", string(raw), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains key after BOM removal", strings.Contains(result, `"key"`)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_UnicodeContent(t *testing.T) {
|
||||
input := `{"名称": "WeKnora 知识库", "描述": "支持中文 JSON 🎉", "emoji": "🚀"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Unicode / Chinese / Emoji", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains Chinese", strings.Contains(result, "知识库")),
|
||||
check("contains emoji", strings.Contains(result, "🚀")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SpecialCharsInStrings(t *testing.T) {
|
||||
input := `{"html": "<div class=\"test\">hello</div>", "path": "C:\\Users\\test", "newline": "line1\nline2"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Special characters (HTML, backslash, newline)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_NullValues(t *testing.T) {
|
||||
input := `{"name": "test", "value": null, "list": [null, 1, null]}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Object with null values", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains 'null'", strings.Contains(result, "null")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_MixedValueTypes(t *testing.T) {
|
||||
input := `{"str": "hello", "num": 42, "float": 3.14, "bool": true, "null_val": null, "arr": [1,2], "obj": {"a":1}}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Mixed value types in one object", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains string", strings.Contains(result, "hello")),
|
||||
check("contains number", strings.Contains(result, "42")),
|
||||
check("contains float", strings.Contains(result, "3.14")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_NestedArrayOfArrays(t *testing.T) {
|
||||
input := `[[1, 2, 3], [4, 5, 6], [7, 8, 9]]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Nested Array of Arrays (matrix)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_LargePrimitiveString(t *testing.T) {
|
||||
bigStr := strings.Repeat("abcdefghij", 300)
|
||||
input := fmt.Sprintf(`"%s"`, bigStr)
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Large Primitive String (3000 chars) → fallback", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("content preserved", strings.Contains(result, "abcdefghij")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SingleKeyObject(t *testing.T) {
|
||||
input := `{"only_key": "only_value"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Single-Key Object", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SingleElementArray(t *testing.T) {
|
||||
input := `[{"id": 1, "name": "solo"}]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Single-Element Array", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 5: Realistic Scenarios
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_RealisticConfig(t *testing.T) {
|
||||
input := `{
|
||||
"app": {"name": "WeKnora", "version": "2.0.0", "debug": false},
|
||||
"database": {"host": "localhost", "port": 5432, "name": "weknora_db", "pool_size": 10},
|
||||
"redis": {"host": "localhost", "port": 6379},
|
||||
"features": {"json_upload": true, "multimodel": true, "graph_extraction": false}
|
||||
}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Realistic: App Config (fits in 1 block)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains 'WeKnora'", strings.Contains(result, "WeKnora")),
|
||||
check("contains 'database'", strings.Contains(result, "database")),
|
||||
check(fmt.Sprintf("single block (fits, got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_RealisticAPIResponse(t *testing.T) {
|
||||
users := make([]interface{}, 40)
|
||||
for i := range users {
|
||||
users[i] = map[string]interface{}{
|
||||
"id": i + 1,
|
||||
"username": fmt.Sprintf("user_%03d", i+1),
|
||||
"email": fmt.Sprintf("user%03d@example.com", i+1),
|
||||
"role": "editor",
|
||||
"active": i%3 != 0,
|
||||
}
|
||||
}
|
||||
apiResp := map[string]interface{}{
|
||||
"status": "ok",
|
||||
"total": 40,
|
||||
"page": 1,
|
||||
"data": users,
|
||||
}
|
||||
data, _ := json.Marshal(apiResp)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Realistic: API Response (40 user records)", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains user data", strings.Contains(result, "user_001")),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_RealisticLargeNestedConfig(t *testing.T) {
|
||||
// A config where one section is huge and must be recursively split,
|
||||
// while other sections are small
|
||||
rules := make(map[string]interface{})
|
||||
for i := 0; i < 30; i++ {
|
||||
rules[fmt.Sprintf("rule_%02d", i)] = map[string]interface{}{
|
||||
"pattern": strings.Repeat("pattern", 10),
|
||||
"action": "allow",
|
||||
"priority": i,
|
||||
}
|
||||
}
|
||||
config := map[string]interface{}{
|
||||
"version": "3.0",
|
||||
"metadata": map[string]interface{}{"author": "admin", "updated": "2026-03-24"},
|
||||
"firewall": map[string]interface{}{
|
||||
"enabled": true,
|
||||
"rules": rules,
|
||||
},
|
||||
}
|
||||
data, _ := json.Marshal(config)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := splitBlocks(result)
|
||||
// Each block should still have the path "firewall" → "rules" where applicable
|
||||
visualReport(t, "Realistic: Nested config with 30 firewall rules", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks needed (got %d)", len(blocks)), len(blocks) >= 2),
|
||||
check("contains 'firewall'", strings.Contains(result, "firewall")),
|
||||
check("contains 'version'", strings.Contains(result, "version")),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 6: Integration — SimpleFormatReader
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestSimpleFormatReader_JSON(t *testing.T) {
|
||||
reader := &SimpleFormatReader{}
|
||||
input := `{"test": "integration", "count": 42}`
|
||||
req := &types.ReadRequest{
|
||||
FileName: "config.json",
|
||||
FileType: "json",
|
||||
FileContent: []byte(input),
|
||||
}
|
||||
|
||||
result, err := reader.Read(context.Background(), req)
|
||||
|
||||
var resultStr string
|
||||
if result != nil {
|
||||
resultStr = result.MarkdownContent
|
||||
}
|
||||
visualReport(t, "SimpleFormatReader.Read() with .json file", input, resultStr, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("result not nil", result != nil),
|
||||
check("markdown contains 'test'",
|
||||
result != nil && strings.Contains(result.MarkdownContent, "test")),
|
||||
check("markdown has code block",
|
||||
result != nil && strings.Contains(result.MarkdownContent, "```json")),
|
||||
check("output is valid JSON inside block",
|
||||
result != nil && allBlocksValidJSON(result.MarkdownContent)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestSimpleFormatReader_JSON_Invalid(t *testing.T) {
|
||||
reader := &SimpleFormatReader{}
|
||||
req := &types.ReadRequest{
|
||||
FileName: "bad.json",
|
||||
FileType: "json",
|
||||
FileContent: []byte(`{not valid json}`),
|
||||
}
|
||||
|
||||
_, err := reader.Read(context.Background(), req)
|
||||
visualReport(t, "SimpleFormatReader with invalid JSON → error", `{not valid json}`, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
check("error mentions 'json'",
|
||||
err != nil && strings.Contains(strings.ToLower(err.Error()), "json")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestIsSimpleFormat_JSON(t *testing.T) {
|
||||
cases := []struct {
|
||||
input string
|
||||
want bool
|
||||
}{
|
||||
{"json", true}, {"JSON", true}, {".json", true}, {"Json", true},
|
||||
{"txt", true}, {"csv", true},
|
||||
{"pdf", false}, {"docx", false},
|
||||
}
|
||||
|
||||
t.Logf("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
||||
t.Logf("📋 TEST: IsSimpleFormat() recognizes JSON")
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
allPass := true
|
||||
for _, tc := range cases {
|
||||
got := IsSimpleFormat(tc.input)
|
||||
icon := "✅"
|
||||
if got != tc.want {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
t.Errorf("IsSimpleFormat(%q) = %v, want %v", tc.input, got, tc.want)
|
||||
}
|
||||
t.Logf("%s IsSimpleFormat(%q) = %v (want %v)", icon, tc.input, got, tc.want)
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n")
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 7: list-to-dict preprocessing
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestListToDictPreprocess(t *testing.T) {
|
||||
// Simulate what json.Unmarshal produces: numbers become float64
|
||||
input := []interface{}{"a", "b", map[string]interface{}{"nested": []interface{}{float64(1), float64(2)}}}
|
||||
result := listToDictPreprocess(input)
|
||||
|
||||
dict, ok := result.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected map, got %T", result)
|
||||
}
|
||||
|
||||
t.Logf("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
||||
t.Logf("📋 TEST: listToDictPreprocess")
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
t.Logf("📥 INPUT: [\"a\", \"b\", {\"nested\": [1, 2]}]")
|
||||
t.Logf("📤 OUTPUT: %s", formatValue(dict))
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
|
||||
checks := []checkResult{
|
||||
check("top level is dict", ok),
|
||||
check(`key "0" = "a"`, dict["0"] == "a"),
|
||||
check(`key "1" = "b"`, dict["1"] == "b"),
|
||||
}
|
||||
// Check nested array also converted
|
||||
if nested, ok2 := dict["2"].(map[string]interface{}); ok2 {
|
||||
if innerNested, ok3 := nested["nested"].(map[string]interface{}); ok3 {
|
||||
checks = append(checks, check(`nested array converted to dict`, innerNested["0"] == float64(1)))
|
||||
} else {
|
||||
checks = append(checks, check(`nested array converted to dict`, false))
|
||||
}
|
||||
}
|
||||
allPass := true
|
||||
for _, c := range checks {
|
||||
icon := "✅"
|
||||
if !c.pass {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
t.Errorf("FAIL: %s", c.desc)
|
||||
}
|
||||
t.Logf("%s %s", icon, c.desc)
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n")
|
||||
}
|
||||
|
||||
func TestSetNestedDict(t *testing.T) {
|
||||
d := map[string]interface{}{}
|
||||
setNestedDict(d, []string{"config", "database", "host"}, "localhost")
|
||||
|
||||
t.Logf("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
||||
t.Logf("📋 TEST: setNestedDict path preservation")
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
t.Logf("📥 path: [\"config\", \"database\", \"host\"]")
|
||||
t.Logf("📥 value: \"localhost\"")
|
||||
t.Logf("📤 result: %s", formatValue(d))
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
|
||||
// Verify the structure
|
||||
config, ok1 := d["config"].(map[string]interface{})
|
||||
db, ok2 := map[string]interface{}{}, false
|
||||
host := ""
|
||||
if ok1 {
|
||||
db, ok2 = config["database"].(map[string]interface{})
|
||||
}
|
||||
if ok2 {
|
||||
host, _ = db["host"].(string)
|
||||
}
|
||||
|
||||
checks := []checkResult{
|
||||
check(`d["config"] is dict`, ok1),
|
||||
check(`d["config"]["database"] is dict`, ok2),
|
||||
check(`d["config"]["database"]["host"] = "localhost"`, host == "localhost"),
|
||||
}
|
||||
allPass := true
|
||||
for _, c := range checks {
|
||||
icon := "✅"
|
||||
if !c.pass {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
t.Errorf("FAIL: %s", c.desc)
|
||||
}
|
||||
t.Logf("%s %s", icon, c.desc)
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n")
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
type markdownImageTargetSpan struct {
|
||||
TargetStart int
|
||||
TargetEnd int
|
||||
}
|
||||
|
||||
func scanMarkdownImageTargets(markdown string) []markdownImageTargetSpan {
|
||||
var spans []markdownImageTargetSpan
|
||||
for i := 0; i+1 < len(markdown); i++ {
|
||||
if markdown[i] != '!' || markdown[i+1] != '[' || isEscaped(markdown, i) {
|
||||
continue
|
||||
}
|
||||
|
||||
altEnd := findMarkdownImageAltEnd(markdown, i+2)
|
||||
if altEnd == -1 {
|
||||
continue
|
||||
}
|
||||
|
||||
targetStart := altEnd + 2
|
||||
targetEnd, ok := findMarkdownImageTargetEnd(markdown, targetStart)
|
||||
if !ok {
|
||||
i = altEnd
|
||||
continue
|
||||
}
|
||||
spans = append(spans, markdownImageTargetSpan{
|
||||
TargetStart: targetStart,
|
||||
TargetEnd: targetEnd,
|
||||
})
|
||||
i = targetEnd
|
||||
}
|
||||
return spans
|
||||
}
|
||||
|
||||
func findMarkdownImageAltEnd(markdown string, start int) int {
|
||||
for i := start; i+1 < len(markdown); i++ {
|
||||
if markdown[i] == ']' && markdown[i+1] == '(' && !isEscaped(markdown, i) {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func findMarkdownImageTargetEnd(markdown string, start int) (int, bool) {
|
||||
parenDepth := 1
|
||||
inAngleDestination := false
|
||||
seenNonSpace := false
|
||||
var inQuote byte
|
||||
|
||||
for i := start; i < len(markdown); i++ {
|
||||
ch := markdown[i]
|
||||
if ch == '\\' {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
|
||||
if !seenNonSpace && !isMarkdownSpace(ch) {
|
||||
seenNonSpace = true
|
||||
if ch == '<' {
|
||||
inAngleDestination = true
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
if inAngleDestination {
|
||||
if ch == '>' {
|
||||
inAngleDestination = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if inQuote != 0 {
|
||||
if ch == inQuote {
|
||||
inQuote = 0
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if (ch == '"' || ch == '\'') && i > start && isMarkdownSpace(markdown[i-1]) {
|
||||
inQuote = ch
|
||||
continue
|
||||
}
|
||||
|
||||
switch ch {
|
||||
case '(':
|
||||
parenDepth++
|
||||
case ')':
|
||||
parenDepth--
|
||||
if parenDepth == 0 {
|
||||
return i, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func splitMarkdownImageTarget(
|
||||
raw string,
|
||||
refMap map[string]types.ImageRef,
|
||||
) (path string, pathStart int, pathEnd int, ok bool) {
|
||||
start, end := trimMarkdownSpaceBounds(raw, 0, len(raw))
|
||||
if start >= end {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
trimmed := raw[start:end]
|
||||
if _, found := refMap[trimmed]; found {
|
||||
return trimmed, start, end, true
|
||||
}
|
||||
|
||||
if raw[start] == '<' {
|
||||
return splitAngleMarkdownImageTarget(raw, start, end, refMap)
|
||||
}
|
||||
|
||||
titleStart, found := parseMarkdownImageTitleSuffix(trimmed)
|
||||
if !found {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
pathEnd = start + titleStart
|
||||
for pathEnd > start && isMarkdownSpace(raw[pathEnd-1]) {
|
||||
pathEnd--
|
||||
}
|
||||
if pathEnd == start {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
path = raw[start:pathEnd]
|
||||
if _, found := refMap[path]; !found {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
return path, start, pathEnd, true
|
||||
}
|
||||
|
||||
func splitAngleMarkdownImageTarget(
|
||||
raw string,
|
||||
start int,
|
||||
end int,
|
||||
refMap map[string]types.ImageRef,
|
||||
) (path string, pathStart int, pathEnd int, ok bool) {
|
||||
closeIdx := -1
|
||||
for i := start + 1; i < end; i++ {
|
||||
if raw[i] == '>' && !isEscaped(raw, i) {
|
||||
closeIdx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if closeIdx == -1 {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
path = raw[start+1 : closeIdx]
|
||||
if _, found := refMap[path]; !found {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
if !isEmptyOrMarkdownImageTitleSuffix(raw[closeIdx+1 : end]) {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
return path, start + 1, closeIdx, true
|
||||
}
|
||||
|
||||
func parseMarkdownImageTitleSuffix(raw string) (titleStart int, ok bool) {
|
||||
_, end := trimMarkdownSpaceBounds(raw, 0, len(raw))
|
||||
if end == 0 {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
switch raw[end-1] {
|
||||
case '"', '\'':
|
||||
quote := raw[end-1]
|
||||
for i := end - 2; i >= 0; i-- {
|
||||
if raw[i] != quote || isEscaped(raw, i) {
|
||||
continue
|
||||
}
|
||||
if i == 0 || !isMarkdownSpace(raw[i-1]) {
|
||||
return 0, false
|
||||
}
|
||||
if markdownTitleHasBlankLine(raw[i+1 : end-1]) {
|
||||
return 0, false
|
||||
}
|
||||
return i, true
|
||||
}
|
||||
case ')':
|
||||
depth := 0
|
||||
for i := end - 2; i >= 0; i-- {
|
||||
if isEscaped(raw, i) {
|
||||
continue
|
||||
}
|
||||
switch raw[i] {
|
||||
case ')':
|
||||
depth++
|
||||
case '(':
|
||||
if depth == 0 {
|
||||
if i == 0 || !isMarkdownSpace(raw[i-1]) {
|
||||
return 0, false
|
||||
}
|
||||
if markdownTitleHasBlankLine(raw[i+1 : end-1]) {
|
||||
return 0, false
|
||||
}
|
||||
return i, true
|
||||
}
|
||||
depth--
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func isEmptyOrMarkdownImageTitleSuffix(raw string) bool {
|
||||
start, end := trimMarkdownSpaceBounds(raw, 0, len(raw))
|
||||
if start == end {
|
||||
return true
|
||||
}
|
||||
titleStart, ok := parseMarkdownImageTitleSuffix(raw[:end])
|
||||
return ok && isAllMarkdownSpace(raw[:titleStart])
|
||||
}
|
||||
|
||||
func markdownTitleHasBlankLine(title string) bool {
|
||||
for _, line := range strings.Split(title, "\n") {
|
||||
if strings.Trim(line, " \t") == "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isAllMarkdownSpace(raw string) bool {
|
||||
for i := 0; i < len(raw); i++ {
|
||||
if !isMarkdownSpace(raw[i]) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func trimMarkdownSpaceBounds(raw string, start int, end int) (int, int) {
|
||||
for start < end && isMarkdownSpace(raw[start]) {
|
||||
start++
|
||||
}
|
||||
for end > start && isMarkdownSpace(raw[end-1]) {
|
||||
end--
|
||||
}
|
||||
return start, end
|
||||
}
|
||||
|
||||
func isMarkdownSpace(b byte) bool {
|
||||
return b == ' ' || b == '\t' || b == '\n' || b == '\r'
|
||||
}
|
||||
|
||||
func isEscaped(s string, pos int) bool {
|
||||
backslashes := 0
|
||||
for i := pos - 1; i >= 0 && s[i] == '\\'; i-- {
|
||||
backslashes++
|
||||
}
|
||||
return backslashes%2 == 1
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestScanMarkdownImageTargets(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
targets []string
|
||||
}{
|
||||
{
|
||||
name: "without title",
|
||||
input: `before  after`,
|
||||
targets: []string{`images/a.png`},
|
||||
},
|
||||
{
|
||||
name: "double quoted title",
|
||||
input: ``,
|
||||
targets: []string{`images/a.png "图"`},
|
||||
},
|
||||
{
|
||||
name: "single quoted title",
|
||||
input: ``,
|
||||
targets: []string{`images/a.png '图'`},
|
||||
},
|
||||
{
|
||||
name: "parenthesized title",
|
||||
input: `)`,
|
||||
targets: []string{`images/a.png (图)`},
|
||||
},
|
||||
{
|
||||
name: "title containing right paren",
|
||||
input: ` 结果")`,
|
||||
targets: []string{`images/a.png "阶段 1) 结果"`},
|
||||
},
|
||||
{
|
||||
name: "title containing both parens",
|
||||
input: ` 结果')`,
|
||||
targets: []string{`images/a.png '阶段 (1) 结果'`},
|
||||
},
|
||||
{
|
||||
name: "escaped quote in title",
|
||||
input: ``,
|
||||
targets: []string{`images/a.png "阶段 \"1\""`},
|
||||
},
|
||||
{
|
||||
name: "multiline title",
|
||||
input: "",
|
||||
targets: []string{"images/a.png\n \"多行 title\""},
|
||||
},
|
||||
{
|
||||
name: "spaced path",
|
||||
input: ``,
|
||||
targets: []string{`images/第 1 页.png`},
|
||||
},
|
||||
{
|
||||
name: "path containing balanced parens",
|
||||
input: `.png "title")`,
|
||||
targets: []string{`images/a_(1).png "title"`},
|
||||
},
|
||||
{
|
||||
name: "angle destination",
|
||||
input: ``,
|
||||
targets: []string{`<images/a b.png> "title"`},
|
||||
},
|
||||
{
|
||||
name: "malformed image is skipped",
|
||||
input: ` after `,
|
||||
targets: []string{`images/b.png`},
|
||||
},
|
||||
{
|
||||
name: "escaped image marker is skipped",
|
||||
input: `\ `,
|
||||
targets: []string{`images/b.png`},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
spans := scanMarkdownImageTargets(tt.input)
|
||||
if len(spans) != len(tt.targets) {
|
||||
t.Fatalf("got %d spans, want %d: %#v", len(spans), len(tt.targets), spans)
|
||||
}
|
||||
for i, span := range spans {
|
||||
got := tt.input[span.TargetStart:span.TargetEnd]
|
||||
if got != tt.targets[i] {
|
||||
t.Fatalf("target %d = %q, want %q", i, got, tt.targets[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitMarkdownImageTarget(t *testing.T) {
|
||||
refMap := map[string]types.ImageRef{
|
||||
"images/a.png": {OriginalRef: "images/a.png"},
|
||||
"images/第 1 页.png": {OriginalRef: "images/第 1 页.png"},
|
||||
"images/a_(1).png": {OriginalRef: "images/a_(1).png"},
|
||||
"images/第 1 页 (测试).gif": {OriginalRef: "images/第 1 页 (测试).gif"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
wantPath string
|
||||
wantText string
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "without title",
|
||||
raw: `images/a.png`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "double quoted title",
|
||||
raw: `images/a.png "图"`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored "图"`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "single quoted title",
|
||||
raw: `images/a.png '图'`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored '图'`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "parenthesized title",
|
||||
raw: `images/a.png (阶段 (1) 结果)`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored (阶段 (1) 结果)`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "multiline title",
|
||||
raw: "images/a.png\n \"多行 title\"",
|
||||
wantPath: `images/a.png`,
|
||||
wantText: "local://stored\n \"多行 title\"",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "path with spaces wins without title",
|
||||
raw: `images/第 1 页.png`,
|
||||
wantPath: `images/第 1 页.png`,
|
||||
wantText: `local://stored`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "path with balanced parens",
|
||||
raw: `images/a_(1).png "title"`,
|
||||
wantPath: `images/a_(1).png`,
|
||||
wantText: `local://stored "title"`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "angle destination replaces only inner path",
|
||||
raw: `<images/第 1 页 (测试).gif> "阶段 1) 图片"`,
|
||||
wantPath: `images/第 1 页 (测试).gif`,
|
||||
wantText: `<local://stored> "阶段 1) 图片"`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "unknown reference",
|
||||
raw: `images/missing.png "title"`,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "blank line in title is invalid",
|
||||
raw: "images/a.png \"line1\n\nline2\"",
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
path, pathStart, pathEnd, ok := splitMarkdownImageTarget(tt.raw, refMap)
|
||||
if ok != tt.wantOK {
|
||||
t.Fatalf("ok = %v, want %v", ok, tt.wantOK)
|
||||
}
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if path != tt.wantPath {
|
||||
t.Fatalf("path = %q, want %q", path, tt.wantPath)
|
||||
}
|
||||
replaced := tt.raw[:pathStart] + "local://stored" + tt.raw[pathEnd:]
|
||||
if replaced != tt.wantText {
|
||||
t.Fatalf("replacement = %q, want %q", replaced, tt.wantText)
|
||||
}
|
||||
if strings.Contains(replaced, tt.wantPath) {
|
||||
t.Fatalf("replacement still contains original path: %q", replaced)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,517 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultPollInterval = 3 * time.Second
|
||||
defaultCloudTimeout = 600 * time.Second
|
||||
defaultBaseURL = "https://mineru.net/api/v4"
|
||||
)
|
||||
|
||||
// MinerUCloudReader calls the MinerU Cloud API (mineru.net) to read/convert documents.
|
||||
// Flow: POST /file-urls/batch → PUT file → poll GET /extract-results/batch/{batch_id}.
|
||||
type MinerUCloudReader struct {
|
||||
apiKey string
|
||||
baseURL string
|
||||
model string
|
||||
formulaEnable bool
|
||||
tableEnable bool
|
||||
ocrEnable bool
|
||||
language string
|
||||
}
|
||||
|
||||
// NewMinerUCloudReader creates a reader from ParserEngineOverrides.
|
||||
func NewMinerUCloudReader(overrides map[string]string) *MinerUCloudReader {
|
||||
return &MinerUCloudReader{
|
||||
apiKey: strings.TrimSpace(overrides["mineru_api_key"]),
|
||||
baseURL: defaultBaseURL,
|
||||
model: stringOr(overrides["mineru_cloud_model"], "pipeline"),
|
||||
formulaEnable: parseBoolOr(overrides["mineru_cloud_enable_formula"], true),
|
||||
tableEnable: parseBoolOr(overrides["mineru_cloud_enable_table"], true),
|
||||
ocrEnable: parseBoolOr(overrides["mineru_cloud_enable_ocr"], true),
|
||||
language: stringOr(overrides["mineru_cloud_language"], "ch"),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.apiKey == "" {
|
||||
return &types.ReadResult{Error: "MinerU Cloud API key is not configured"}, nil
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerUCloud] Parsing file=%s size=%d via %s", req.FileName, len(content), c.baseURL)
|
||||
|
||||
ext := filepath.Ext(req.FileName)
|
||||
if ext == "" && req.FileType != "" {
|
||||
ext = "." + req.FileType
|
||||
}
|
||||
if ext == "" {
|
||||
ext = ".pdf"
|
||||
}
|
||||
fileName := strings.TrimSuffix(req.FileName, ext) + ext
|
||||
if fileName == ext {
|
||||
fileName = "document" + ext
|
||||
}
|
||||
|
||||
batchID, uploadURL, err := c.applyUploadURLs(ctx, fileName, ext)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MinerU Cloud apply upload URLs: %w", err)
|
||||
}
|
||||
|
||||
if err := c.uploadFile(ctx, uploadURL, content); err != nil {
|
||||
return nil, fmt.Errorf("MinerU Cloud file upload: %w", err)
|
||||
}
|
||||
|
||||
mdContent, imageRefs, err := c.pollBatchResult(ctx, batchID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MinerU Cloud poll: %w", err)
|
||||
}
|
||||
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// --- batch upload API ---
|
||||
|
||||
type batchApplyResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data struct {
|
||||
BatchID string `json:"batch_id"`
|
||||
FileURLs []string `json:"file_urls"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) applyUploadURLs(ctx context.Context, fileName, ext string) (string, string, error) {
|
||||
modelVersion := c.model
|
||||
if strings.ToLower(ext) == ".html" {
|
||||
modelVersion = "MinerU-HTML"
|
||||
}
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"files": []map[string]string{{"name": fileName, "data_id": uuid.New().String()}},
|
||||
"model_version": modelVersion,
|
||||
"is_ocr": c.ocrEnable,
|
||||
"enable_formula": c.formulaEnable,
|
||||
"enable_table": c.tableEnable,
|
||||
"language": c.language,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("marshal payload: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/file-urls/batch", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 30 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
return "", "", fmt.Errorf("API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result batchApplyResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", "", fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if result.Code != 0 {
|
||||
return "", "", fmt.Errorf("API error: %s", result.Msg)
|
||||
}
|
||||
if len(result.Data.FileURLs) == 0 {
|
||||
return "", "", fmt.Errorf("API returned no file_urls")
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerUCloud] batch apply ok: batch_id=%s, urls=%d", result.Data.BatchID, len(result.Data.FileURLs))
|
||||
return result.Data.BatchID, result.Data.FileURLs[0], nil
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) uploadFile(ctx context.Context, uploadURL string, content []byte) error {
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, bytes.NewReader(content))
|
||||
if err != nil {
|
||||
return fmt.Errorf("create PUT request: %w", err)
|
||||
}
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 120 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("PUT upload: %w", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 300 {
|
||||
return fmt.Errorf("PUT upload status %d", resp.StatusCode)
|
||||
}
|
||||
logger.Infof(context.Background(), "[MinerUCloud] file uploaded, status=%d", resp.StatusCode)
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- polling ---
|
||||
|
||||
type batchPollResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data struct {
|
||||
ExtractResult json.RawMessage `json:"extract_result"` // can be object or array
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type extractResultItem struct {
|
||||
State string `json:"state"`
|
||||
FileName string `json:"file_name"`
|
||||
Markdown string `json:"markdown"`
|
||||
Content string `json:"content"`
|
||||
Text string `json:"text"`
|
||||
ErrMsg string `json:"err_msg"`
|
||||
Progress struct {
|
||||
ExtractedPages int `json:"extracted_pages"`
|
||||
TotalPages int `json:"total_pages"`
|
||||
} `json:"extract_progress"`
|
||||
FullZipURL string `json:"full_zip_url"`
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) pollBatchResult(ctx context.Context, batchID string) (string, []types.ImageRef, error) {
|
||||
deadline := time.Now().Add(defaultCloudTimeout)
|
||||
pollCount := 0
|
||||
headers := map[string]string{
|
||||
"Authorization": "Bearer " + c.apiKey,
|
||||
}
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
// Bail out promptly on caller cancellation instead of spinning:
|
||||
// fetchBatchStatus fails immediately and sleepCtx returns at once on a
|
||||
// cancelled ctx, so without this guard the loop busy-hammers the cloud
|
||||
// API and floods logs until the deadline.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
pollCount++
|
||||
|
||||
items, err := c.fetchBatchStatus(ctx, batchID, headers)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerUCloud] poll #%d failed: %v", pollCount, err)
|
||||
sleepCtx(ctx, defaultPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
if len(items) == 0 {
|
||||
if pollCount <= 3 || pollCount%10 == 0 {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] poll #%d: extract_result empty, retrying", pollCount)
|
||||
}
|
||||
sleepCtx(ctx, defaultPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
item := items[0]
|
||||
state := strings.ToLower(item.State)
|
||||
|
||||
if pollCount == 1 || pollCount%10 == 0 || state == "done" || state == "failed" {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] poll #%d: file=%s state=%s pages=%d/%d",
|
||||
pollCount, item.FileName, state, item.Progress.ExtractedPages, item.Progress.TotalPages)
|
||||
}
|
||||
|
||||
if state == "failed" {
|
||||
return "", nil, fmt.Errorf("MinerU Cloud task failed: %s", item.ErrMsg)
|
||||
}
|
||||
|
||||
if state == "done" {
|
||||
return c.extractDoneResult(ctx, &item)
|
||||
}
|
||||
|
||||
sleepCtx(ctx, defaultPollInterval)
|
||||
}
|
||||
|
||||
return "", nil, fmt.Errorf("MinerU Cloud task timed out after %d polls", pollCount)
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) fetchBatchStatus(ctx context.Context, batchID string, headers map[string]string) ([]extractResultItem, error) {
|
||||
url := fmt.Sprintf("%s/extract-results/batch/%s", c.baseURL, batchID)
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k, v := range headers {
|
||||
httpReq.Header.Set(k, v)
|
||||
}
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 30 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read poll response body: %w", err)
|
||||
}
|
||||
|
||||
var pollResp batchPollResponse
|
||||
if err := json.Unmarshal(respBody, &pollResp); err != nil {
|
||||
return nil, fmt.Errorf("decode poll response: %w", err)
|
||||
}
|
||||
if pollResp.Code != 0 {
|
||||
return nil, fmt.Errorf("poll error code=%d msg=%s", pollResp.Code, pollResp.Msg)
|
||||
}
|
||||
|
||||
if len(pollResp.Data.ExtractResult) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Dump the raw extract_result JSON for debugging
|
||||
rawExtract := string(pollResp.Data.ExtractResult)
|
||||
if len(rawExtract) > 4000 {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] Raw extract_result (truncated to 4000 chars): %s ...", rawExtract[:4000])
|
||||
} else {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] Raw extract_result: %s", rawExtract)
|
||||
}
|
||||
|
||||
// Pretty-print the structure to reveal all available fields
|
||||
var rawObj interface{}
|
||||
if err := json.Unmarshal(pollResp.Data.ExtractResult, &rawObj); err == nil {
|
||||
logResponseStructure("MinerUCloud", rawObj, "extract_result")
|
||||
}
|
||||
|
||||
// The extract_result can be either a single object or an array
|
||||
var items []extractResultItem
|
||||
if pollResp.Data.ExtractResult[0] == '[' {
|
||||
if err := json.Unmarshal(pollResp.Data.ExtractResult, &items); err != nil {
|
||||
return nil, fmt.Errorf("decode extract_result array: %w", err)
|
||||
}
|
||||
} else {
|
||||
var single extractResultItem
|
||||
if err := json.Unmarshal(pollResp.Data.ExtractResult, &single); err != nil {
|
||||
return nil, fmt.Errorf("decode extract_result object: %w", err)
|
||||
}
|
||||
items = []extractResultItem{single}
|
||||
}
|
||||
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// extractDoneResult extracts markdown and images from a completed batch item.
|
||||
// Prefers inline markdown/content fields; falls back to downloading full_zip_url.
|
||||
func (c *MinerUCloudReader) extractDoneResult(_ context.Context, item *extractResultItem) (string, []types.ImageRef, error) {
|
||||
text := firstNonEmpty(item.Markdown, item.Content, item.Text)
|
||||
if text != "" {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] parsed (inline), length=%d", len(text))
|
||||
return text, nil, nil
|
||||
}
|
||||
|
||||
if item.FullZipURL == "" {
|
||||
return "", nil, fmt.Errorf("MinerU Cloud state=done but no markdown/content or full_zip_url")
|
||||
}
|
||||
|
||||
md, imageRefs, err := downloadAndExtractZip(item.FullZipURL)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("extract zip: %w", err)
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerUCloud] parsed (zip), markdown=%d chars, images=%d", len(md), len(imageRefs))
|
||||
return md, imageRefs, nil
|
||||
}
|
||||
|
||||
// --- ZIP handling ---
|
||||
|
||||
var imgRefPattern = regexp.MustCompile(`!\[[^\]]*\]\(([^)]+)\)`)
|
||||
|
||||
func downloadAndExtractZip(zipURL string) (string, []types.ImageRef, error) {
|
||||
if err := utils.ValidateURLForSSRF(zipURL); err != nil {
|
||||
return "", nil, fmt.Errorf("zip URL blocked by SSRF check: %v", err)
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 120 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Get(zipURL)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("download zip: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", nil, fmt.Errorf("download zip status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
zipData, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read zip body: %w", err)
|
||||
}
|
||||
|
||||
zr, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("open zip: %w", err)
|
||||
}
|
||||
|
||||
// Find .md files
|
||||
var mdFiles []string
|
||||
entries := make(map[string]*zip.File)
|
||||
for _, f := range zr.File {
|
||||
entries[f.Name] = f
|
||||
if strings.HasSuffix(f.Name, ".md") {
|
||||
mdFiles = append(mdFiles, f.Name)
|
||||
}
|
||||
}
|
||||
if len(mdFiles) == 0 {
|
||||
return "", nil, fmt.Errorf("no .md file found in zip")
|
||||
}
|
||||
sort.Slice(mdFiles, func(i, j int) bool {
|
||||
di, dj := strings.Count(mdFiles[i], "/"), strings.Count(mdFiles[j], "/")
|
||||
if di != dj {
|
||||
return di < dj
|
||||
}
|
||||
return mdFiles[i] < mdFiles[j]
|
||||
})
|
||||
|
||||
mdText, err := readZipEntry(entries[mdFiles[0]])
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read md file: %w", err)
|
||||
}
|
||||
|
||||
mdDir := filepath.Dir(mdFiles[0])
|
||||
|
||||
// Extract referenced images
|
||||
var imageRefs []types.ImageRef
|
||||
seen := map[string]bool{}
|
||||
for _, match := range imgRefPattern.FindAllStringSubmatch(mdText, -1) {
|
||||
imgPath := match[1]
|
||||
if strings.HasPrefix(imgPath, "http://") || strings.HasPrefix(imgPath, "https://") || strings.HasPrefix(imgPath, "data:") {
|
||||
continue
|
||||
}
|
||||
if seen[imgPath] {
|
||||
continue
|
||||
}
|
||||
seen[imgPath] = true
|
||||
|
||||
resolved := resolveInZip(imgPath, mdDir, entries)
|
||||
if resolved == nil {
|
||||
logger.Errorf(context.Background(), "[MinerUCloud] image not found in zip: %s", imgPath)
|
||||
continue
|
||||
}
|
||||
|
||||
imgData, err := readZipEntryBytes(resolved)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerUCloud] failed to read zip image %s: %v", imgPath, err)
|
||||
continue
|
||||
}
|
||||
|
||||
ext := strings.ToLower(filepath.Ext(resolved.Name))
|
||||
if ext == "" {
|
||||
ext = ".png"
|
||||
}
|
||||
mimeType := mime.TypeByExtension(ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
imageRefs = append(imageRefs, types.ImageRef{
|
||||
Filename: filepath.Base(resolved.Name),
|
||||
OriginalRef: imgPath,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgData,
|
||||
})
|
||||
}
|
||||
|
||||
return mdText, imageRefs, nil
|
||||
}
|
||||
|
||||
func resolveInZip(imgPath, mdDir string, entries map[string]*zip.File) *zip.File {
|
||||
normalized := strings.ReplaceAll(imgPath, "\\", "/")
|
||||
if f, ok := entries[normalized]; ok {
|
||||
return f
|
||||
}
|
||||
if mdDir != "" && mdDir != "." {
|
||||
joined := mdDir + "/" + normalized
|
||||
if f, ok := entries[joined]; ok {
|
||||
return f
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func readZipEntry(f *zip.File) (string, error) {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rc.Close()
|
||||
data, err := io.ReadAll(rc)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func readZipEntryBytes(f *zip.File) ([]byte, error) {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rc.Close()
|
||||
return io.ReadAll(rc)
|
||||
}
|
||||
|
||||
// PingMinerUCloud checks if the MinerU Cloud API is reachable with the given API key.
|
||||
func PingMinerUCloud(apiKey string) (bool, string) {
|
||||
apiKey = strings.TrimSpace(apiKey)
|
||||
if apiKey == "" {
|
||||
return false, "未配置 MinerU Cloud API Key"
|
||||
}
|
||||
|
||||
targetURL := defaultBaseURL + "/file-urls/batch"
|
||||
payload := []byte(`{"files":[],"model_version":"pipeline"}`)
|
||||
req, err := http.NewRequest(http.MethodPost, targetURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("构建请求失败: %v", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: 10 * time.Second,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("MinerU Cloud 不可达: %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return false, "MinerU Cloud API Key 无效"
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
@@ -0,0 +1,472 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const mineruTimeout = 1000 * time.Second // large docs can take a while
|
||||
|
||||
var (
|
||||
b64DataURIPattern = regexp.MustCompile(`^data:image/(\w+);base64,(.+)$`)
|
||||
minerUFileTypePattern = regexp.MustCompile(`^[a-z0-9]+$`)
|
||||
)
|
||||
|
||||
// MinerUReader calls a self-hosted MinerU API to read/convert documents.
|
||||
type MinerUReader struct {
|
||||
endpoint string
|
||||
backend string // "pipeline", "vlm-*", "hybrid-*"
|
||||
vlmServerURL string // vLLM server URL for vlm-http-client / hybrid-http-client
|
||||
formulaEnable bool
|
||||
tableEnable bool
|
||||
parseMethod string
|
||||
language string
|
||||
}
|
||||
|
||||
// NewMinerUReader creates a reader from ParserEngineOverrides.
|
||||
func NewMinerUReader(overrides map[string]string) *MinerUReader {
|
||||
var legacyOCREnabled *bool
|
||||
if raw, ok := overrides["mineru_enable_ocr"]; ok {
|
||||
value := parseBoolOr(raw, true)
|
||||
legacyOCREnabled = &value
|
||||
}
|
||||
|
||||
c := &MinerUReader{
|
||||
endpoint: strings.TrimRight(overrides["mineru_endpoint"], "/"),
|
||||
backend: stringOr(overrides["mineru_model"], "pipeline"),
|
||||
vlmServerURL: overrides["mineru_vlm_server_url"],
|
||||
formulaEnable: parseBoolOr(overrides["mineru_enable_formula"], true),
|
||||
tableEnable: parseBoolOr(overrides["mineru_enable_table"], true),
|
||||
parseMethod: types.ResolveMinerUParseMethod(overrides["mineru_parse_method"], legacyOCREnabled),
|
||||
language: stringOr(overrides["mineru_language"], "ch"),
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *MinerUReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.endpoint == "" {
|
||||
return &types.ReadResult{Error: "MinerU endpoint is not configured"}, nil
|
||||
}
|
||||
if err := validateMinerUOutboundURL(c.endpoint); err != nil {
|
||||
return &types.ReadResult{Error: err.Error()}, nil
|
||||
}
|
||||
if c.vlmServerURL != "" {
|
||||
if err := validateMinerUOutboundURL(c.vlmServerURL); err != nil {
|
||||
return &types.ReadResult{Error: err.Error()}, nil
|
||||
}
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerU] Parsing file=%s size=%d via %s", req.FileName, len(content), c.endpoint)
|
||||
|
||||
mdContent, imagesB64, err := c.callFileParse(ctx, content, req.FileName, req.FileType)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MinerU file_parse: %w", err)
|
||||
}
|
||||
|
||||
// MinerU already returns markdown with embedded HTML blocks (e.g. <table>, <details>).
|
||||
// Re-running the whole document through html-to-markdown corrupts valid markdown
|
||||
// by escaping headings and image syntax. Only apply narrow compatibility fixes.
|
||||
mdContent = normalizeMinerUMarkdown(mdContent)
|
||||
|
||||
// Process images: decode base64, build ImageRef list, replace refs in markdown
|
||||
imageRefs, mdContent := c.processImages(mdContent, imagesB64)
|
||||
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
logger.Infof(context.Background(), "[MinerU] Parsed successfully, markdown=%d chars, images=%d", len(mdContent), len(imageRefs))
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type mineruFileEntry struct {
|
||||
MDContent string `json:"md_content"`
|
||||
Images map[string]string `json:"images"` // path -> "data:image/png;base64,..." or raw base64
|
||||
}
|
||||
|
||||
func minerUCleanFileType(fileType string) string {
|
||||
cleanType := strings.ToLower(strings.TrimSpace(fileType))
|
||||
cleanType = strings.TrimPrefix(cleanType, ".")
|
||||
if minerUFileTypePattern.MatchString(cleanType) {
|
||||
return cleanType
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func minerUUploadFileName(fileName, fileType string) string {
|
||||
cleanName := strings.TrimSpace(fileName)
|
||||
cleanName = strings.ReplaceAll(cleanName, `\`, "/")
|
||||
cleanName = path.Base(cleanName)
|
||||
cleanName = strings.Map(func(r rune) rune {
|
||||
if r == '\r' || r == '\n' || r == 0 {
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, cleanName)
|
||||
cleanName = strings.TrimSpace(cleanName)
|
||||
if cleanName != "" && cleanName != "." && cleanName != ".." && cleanName != "/" {
|
||||
if filepath.Ext(cleanName) == "" {
|
||||
if cleanType := minerUCleanFileType(fileType); cleanType != "" {
|
||||
return cleanName + "." + cleanType
|
||||
}
|
||||
}
|
||||
return cleanName
|
||||
}
|
||||
|
||||
if cleanType := minerUCleanFileType(fileType); cleanType != "" {
|
||||
return "document." + cleanType
|
||||
}
|
||||
return "document"
|
||||
}
|
||||
|
||||
// minerUResultStem mirrors MinerU's upload.stem: the basename without extension.
|
||||
func minerUResultStem(uploadFileName string) string {
|
||||
stem := strings.TrimSuffix(path.Base(uploadFileName), filepath.Ext(uploadFileName))
|
||||
if stem == "" || stem == "." {
|
||||
return ""
|
||||
}
|
||||
return stem
|
||||
}
|
||||
|
||||
func minerUResultLookupKeys(uploadFileName string) []string {
|
||||
keys := make([]string, 0, 3)
|
||||
if stem := minerUResultStem(uploadFileName); stem != "" {
|
||||
keys = append(keys, stem)
|
||||
}
|
||||
return append(keys, "document", "files")
|
||||
}
|
||||
|
||||
func parseMinerUFileParseResponse(respBody []byte, uploadFileName string) (string, map[string]string, string, error) {
|
||||
var envelope struct {
|
||||
Results map[string]mineruFileEntry `json:"results"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &envelope); err != nil {
|
||||
return "", nil, "", fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if len(envelope.Results) == 0 {
|
||||
return "", nil, "", nil
|
||||
}
|
||||
|
||||
for _, key := range minerUResultLookupKeys(uploadFileName) {
|
||||
if entry, ok := envelope.Results[key]; ok {
|
||||
if entry.MDContent != "" || len(entry.Images) > 0 {
|
||||
return entry.MDContent, entry.Images, key, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for key, entry := range envelope.Results {
|
||||
if entry.MDContent != "" || len(entry.Images) > 0 {
|
||||
return entry.MDContent, entry.Images, key, nil
|
||||
}
|
||||
}
|
||||
return "", nil, "", nil
|
||||
}
|
||||
|
||||
func (c *MinerUReader) callFileParse(
|
||||
ctx context.Context,
|
||||
content []byte,
|
||||
fileName string,
|
||||
fileType string,
|
||||
) (string, map[string]string, error) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
|
||||
// Form fields
|
||||
fields := map[string]string{
|
||||
"return_md": "true",
|
||||
"return_images": "true",
|
||||
"table_enable": fmt.Sprintf("%v", c.tableEnable),
|
||||
"formula_enable": fmt.Sprintf("%v", c.formulaEnable),
|
||||
"parse_method": c.parseMethod,
|
||||
"start_page_id": "0",
|
||||
"end_page_id": "99999",
|
||||
"backend": c.backend,
|
||||
"response_format_zip": "false",
|
||||
"return_middle_json": "false",
|
||||
"return_model_output": "false",
|
||||
"return_content_list": "true",
|
||||
}
|
||||
if c.language != "" {
|
||||
fields["lang_list"] = c.language
|
||||
}
|
||||
if c.vlmServerURL != "" && (strings.HasPrefix(c.backend, "vlm-http-client") || strings.HasPrefix(c.backend, "hybrid-http-client")) {
|
||||
fields["server_url"] = c.vlmServerURL
|
||||
}
|
||||
for k, v := range fields {
|
||||
_ = writer.WriteField(k, v)
|
||||
}
|
||||
|
||||
uploadFileName := minerUUploadFileName(fileName, fileType)
|
||||
|
||||
// File part
|
||||
part, err := writer.CreateFormFile("files", uploadFileName)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("create form file: %w", err)
|
||||
}
|
||||
if _, err := part.Write(content); err != nil {
|
||||
return "", nil, fmt.Errorf("write file content: %w", err)
|
||||
}
|
||||
writer.Close()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint+"/file_parse", &body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: mineruTimeout,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
return "", nil, fmt.Errorf("MinerU API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read response body: %w", err)
|
||||
}
|
||||
|
||||
// Dump raw response for debugging (truncate if too large)
|
||||
rawStr := string(respBody)
|
||||
if len(rawStr) > 4000 {
|
||||
logger.Infof(context.Background(), "[MinerU] Raw response (truncated to 4000 chars): %s ...", rawStr[:4000])
|
||||
} else {
|
||||
logger.Infof(context.Background(), "[MinerU] Raw response: %s", rawStr)
|
||||
}
|
||||
|
||||
// Also pretty-print the top-level structure (without large base64 blobs)
|
||||
var rawMap map[string]interface{}
|
||||
if err := json.Unmarshal(respBody, &rawMap); err == nil {
|
||||
c.logMinerUResponseStructure(rawMap, "")
|
||||
}
|
||||
|
||||
mdContent, imagesB64, resultKey, err := parseMinerUFileParseResponse(respBody, uploadFileName)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
if resultKey != "" {
|
||||
logger.Infof(context.Background(), "[MinerU] Using response path: results.%s", resultKey)
|
||||
return mdContent, imagesB64, nil
|
||||
}
|
||||
|
||||
logger.Errorf(context.Background(), "[MinerU] Response has no markdown/images under results")
|
||||
return "", nil, nil
|
||||
}
|
||||
|
||||
// processImages decodes base64 images from MinerU response and returns ImageRef list.
|
||||
// It also replaces image references in the markdown content.
|
||||
func (c *MinerUReader) processImages(mdContent string, imagesB64 map[string]string) ([]types.ImageRef, string) {
|
||||
var refs []types.ImageRef
|
||||
|
||||
for ipath, b64Str := range imagesB64 {
|
||||
matchedRefs := mineruImageOriginalRefs(mdContent, ipath)
|
||||
if len(matchedRefs) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var imgBytes []byte
|
||||
var ext string
|
||||
|
||||
if m := b64DataURIPattern.FindStringSubmatch(b64Str); len(m) == 3 {
|
||||
ext = m[1]
|
||||
decoded, err := base64.StdEncoding.DecodeString(m[2])
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerU] Failed to decode base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
} else {
|
||||
// raw base64 without data URI prefix
|
||||
decoded, err := base64.StdEncoding.DecodeString(b64Str)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerU] Failed to decode raw base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
ext = strings.TrimPrefix(filepath.Ext(ipath), ".")
|
||||
if ext == "" {
|
||||
ext = "png"
|
||||
}
|
||||
}
|
||||
|
||||
mimeType := mime.TypeByExtension("." + ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
for _, originalRef := range matchedRefs {
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: ipath,
|
||||
OriginalRef: originalRef,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgBytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return refs, mdContent
|
||||
}
|
||||
|
||||
// logMinerUResponseStructure logs the structure of the MinerU API response.
|
||||
func (c *MinerUReader) logMinerUResponseStructure(obj interface{}, prefix string) {
|
||||
logResponseStructure("MinerU", obj, prefix)
|
||||
}
|
||||
|
||||
// validateMinerUOutboundURL rejects MinerU endpoints that would reach private
|
||||
// or otherwise restricted hosts when parsed or probed from the app server.
|
||||
func validateMinerUOutboundURL(rawURL string) error {
|
||||
rawURL = strings.TrimSpace(rawURL)
|
||||
if rawURL == "" {
|
||||
return nil
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(rawURL); err != nil {
|
||||
return fmt.Errorf("MinerU URL blocked by SSRF check: %v", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PingMinerU checks if the self-hosted MinerU service is reachable.
|
||||
func PingMinerU(endpoint string) (bool, string) {
|
||||
endpoint = strings.TrimRight(endpoint, "/")
|
||||
if endpoint == "" {
|
||||
return false, "未配置 MinerU 端点"
|
||||
}
|
||||
if err := validateMinerUOutboundURL(endpoint); err != nil {
|
||||
return false, err.Error()
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: 5 * time.Second,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Get(endpoint + "/docs")
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("MinerU 服务不可达: %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode >= 400 {
|
||||
return false, fmt.Sprintf("MinerU 服务返回状态 %d", resp.StatusCode)
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
|
||||
// escapedImageSyntaxPattern matches markdown image references whose `[` was
|
||||
// over-escaped to `\[` by html-to-markdown. The URL group mirrors the
|
||||
// downstream image-extraction regex so escapes are only stripped for actual
|
||||
// image references.
|
||||
var escapedImageSyntaxPattern = regexp.MustCompile(`!\\\[(.*?)\\?\]\(([^()\n]*(?:\([^)]*\)[^()\n]*)*)\)`)
|
||||
|
||||
// escapedHeadingPattern restores markdown headings that were over-escaped to
|
||||
// \# Heading. We only touch line-leading heading markers to avoid rewriting
|
||||
// body text that intentionally contains escaped # characters.
|
||||
var escapedHeadingPattern = regexp.MustCompile(`(?m)^\\(#{1,6})(\s+)`)
|
||||
|
||||
// unescapeMarkdownImageSyntax restores `` from html-to-markdown's
|
||||
// over-escaped `!\[alt\](url)` form. Without this, the downstream image regex
|
||||
// in ImageResolver fails to match and images are never persisted.
|
||||
func unescapeMarkdownImageSyntax(content string) string {
|
||||
return escapedImageSyntaxPattern.ReplaceAllString(content, "")
|
||||
}
|
||||
|
||||
func normalizeEscapedMarkdownHeadings(content string) string {
|
||||
return escapedHeadingPattern.ReplaceAllString(content, `$1$2`)
|
||||
}
|
||||
|
||||
func normalizeMinerUMarkdown(content string) string {
|
||||
content = unescapeMarkdownImageSyntax(content)
|
||||
content = normalizeEscapedMarkdownHeadings(content)
|
||||
return content
|
||||
}
|
||||
|
||||
func mineruImageOriginalRefs(mdContent, imagePath string) []string {
|
||||
normalizedTarget := normalizeMinerUImagePath(imagePath)
|
||||
if normalizedTarget == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
referenced := extractImageRefsFromContent(mdContent)
|
||||
seen := make(map[string]struct{}, len(referenced))
|
||||
var matched []string
|
||||
for _, refPath := range referenced {
|
||||
if normalizeMinerUImagePath(refPath) != normalizedTarget {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[refPath]; ok {
|
||||
continue
|
||||
}
|
||||
matched = append(matched, refPath)
|
||||
seen[refPath] = struct{}{}
|
||||
}
|
||||
|
||||
return matched
|
||||
}
|
||||
|
||||
// imgMarkdownPatternAllowSpaces matches markdown image syntax while allowing
|
||||
// spaces in the URL group, so that paths like "images/第 1 页.jpg" produced by
|
||||
// MinerU on Chinese documents are still detected as image references.
|
||||
var imgMarkdownPatternAllowSpaces = regexp.MustCompile(
|
||||
`!\[(.*?)\]\(([^()\n]*(?:\([^)]*\)[^()\n]*)*)\)`,
|
||||
)
|
||||
|
||||
func extractImageRefsFromContent(content string) []string {
|
||||
var refs []string
|
||||
|
||||
for _, match := range imgMarkdownPatternAllowSpaces.FindAllStringSubmatch(content, -1) {
|
||||
if len(match) >= 3 {
|
||||
refs = append(refs, strings.TrimSpace(match[2]))
|
||||
}
|
||||
}
|
||||
for _, match := range imgHTMLSrc.FindAllStringSubmatch(content, -1) {
|
||||
if len(match) >= 3 {
|
||||
refs = append(refs, match[2])
|
||||
}
|
||||
}
|
||||
|
||||
return refs
|
||||
}
|
||||
|
||||
func normalizeMinerUImagePath(p string) string {
|
||||
p = strings.TrimSpace(p)
|
||||
if p == "" {
|
||||
return ""
|
||||
}
|
||||
if decoded, err := url.PathUnescape(p); err == nil && decoded != "" {
|
||||
p = decoded
|
||||
}
|
||||
p = strings.TrimPrefix(p, "./")
|
||||
p = strings.TrimPrefix(p, "/")
|
||||
p = strings.TrimPrefix(p, "images/")
|
||||
return p
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateMinerUOutboundURL_RejectsLoopback(t *testing.T) {
|
||||
err := validateMinerUOutboundURL("http://127.0.0.1:8080")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "SSRF")
|
||||
}
|
||||
|
||||
func TestPingMinerU_RejectsPrivateEndpoint(t *testing.T) {
|
||||
ok, msg := PingMinerU("http://127.0.0.1:8080")
|
||||
assert.False(t, ok)
|
||||
assert.Contains(t, msg, "SSRF")
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
func TestNewMinerUReaderResolvesParseMethod(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
overrides map[string]string
|
||||
want string
|
||||
}{
|
||||
{name: "default is auto", overrides: map[string]string{}, want: "auto"},
|
||||
{name: "legacy enabled becomes auto", overrides: map[string]string{"mineru_enable_ocr": "true"}, want: "auto"},
|
||||
{name: "legacy disabled becomes text", overrides: map[string]string{"mineru_enable_ocr": "false"}, want: "txt"},
|
||||
{name: "explicit method wins", overrides: map[string]string{"mineru_parse_method": "ocr", "mineru_enable_ocr": "false"}, want: "ocr"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
reader := NewMinerUReader(tt.overrides)
|
||||
if reader.parseMethod != tt.want {
|
||||
t.Fatalf("parseMethod = %q, want %q", reader.parseMethod, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinerUUploadFileName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
fileName string
|
||||
fileType string
|
||||
want string
|
||||
}{
|
||||
{name: "original basename", fileName: "reports/Illustrator 研报.pdf", fileType: "pdf", want: "Illustrator 研报.pdf"},
|
||||
{
|
||||
name: "strip multipart controls",
|
||||
fileName: "reports/\r\nIllustrator\x00 export.pdf",
|
||||
fileType: "pdf",
|
||||
want: "Illustrator export.pdf",
|
||||
},
|
||||
{name: "append extension from type", fileName: "report", fileType: "pdf", want: "report.pdf"},
|
||||
{name: "fallback from bare type", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "fallback from dotted uppercase type", fileType: ".PDF", want: "document.pdf"},
|
||||
{name: "fallback from path-only name", fileName: "/", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "legacy fallback without metadata", want: "document"},
|
||||
{name: "reject unsafe fallback type", fileType: "../pdf", want: "document"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := minerUUploadFileName(tt.fileName, tt.fileType); got != tt.want {
|
||||
t.Fatalf("minerUUploadFileName() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseMinerUFileParseResponse(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
uploadFileName string
|
||||
response map[string]any
|
||||
wantMarkdown string
|
||||
wantResultKey string
|
||||
}{
|
||||
{
|
||||
name: "match upload stem",
|
||||
uploadFileName: "Illustrator 研报.pdf",
|
||||
response: map[string]any{
|
||||
"results": map[string]any{
|
||||
"Illustrator 研报": map[string]any{"md_content": "from-stem", "images": map[string]any{}},
|
||||
},
|
||||
},
|
||||
wantMarkdown: "from-stem",
|
||||
wantResultKey: "Illustrator 研报",
|
||||
},
|
||||
{
|
||||
name: "legacy document key",
|
||||
uploadFileName: "document",
|
||||
response: map[string]any{
|
||||
"results": map[string]any{
|
||||
"document": map[string]any{"md_content": "legacy", "images": map[string]any{}},
|
||||
},
|
||||
},
|
||||
wantMarkdown: "legacy",
|
||||
wantResultKey: "document",
|
||||
},
|
||||
{
|
||||
name: "files variant",
|
||||
uploadFileName: "document.pdf",
|
||||
response: map[string]any{
|
||||
"results": map[string]any{
|
||||
"files": map[string]any{"md_content": "from-files", "images": map[string]any{}},
|
||||
},
|
||||
},
|
||||
wantMarkdown: "from-files",
|
||||
wantResultKey: "files",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
respBody, err := json.Marshal(tt.response)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal response: %v", err)
|
||||
}
|
||||
gotMarkdown, _, gotKey, err := parseMinerUFileParseResponse(respBody, tt.uploadFileName)
|
||||
if err != nil {
|
||||
t.Fatalf("parseMinerUFileParseResponse() error: %v", err)
|
||||
}
|
||||
if gotMarkdown != tt.wantMarkdown {
|
||||
t.Fatalf("markdown = %q, want %q", gotMarkdown, tt.wantMarkdown)
|
||||
}
|
||||
if gotKey != tt.wantResultKey {
|
||||
t.Fatalf("result key = %q, want %q", gotKey, tt.wantResultKey)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinerUReaderPreservesMultipartFilename(t *testing.T) {
|
||||
utils.SetSSRFWhitelistFromRaw("127.0.0.1")
|
||||
t.Cleanup(func() {
|
||||
utils.SetSSRFWhitelistFromRaw("")
|
||||
})
|
||||
|
||||
const fileContent = "%PDF-1.7 synthetic test content"
|
||||
tests := []struct {
|
||||
name string
|
||||
fileName string
|
||||
fileType string
|
||||
want string
|
||||
}{
|
||||
{name: "original basename", fileName: "reports/Illustrator 研报.pdf", fileType: "pdf", want: "Illustrator 研报.pdf"},
|
||||
{
|
||||
name: "strip multipart controls",
|
||||
fileName: "reports/\r\nIllustrator\x00 export.pdf",
|
||||
fileType: "pdf",
|
||||
want: "Illustrator export.pdf",
|
||||
},
|
||||
{name: "fallback from bare type", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "fallback from dotted uppercase type", fileType: ".PDF", want: "document.pdf"},
|
||||
{name: "fallback from path-only name", fileName: "/", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "legacy fallback without metadata", want: "document"},
|
||||
{name: "reject unsafe fallback type", fileType: "../pdf", want: "document"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
method string
|
||||
path string
|
||||
fileName string
|
||||
content string
|
||||
err error
|
||||
}
|
||||
|
||||
captured := make(chan capturedRequest, 1)
|
||||
uploadName := minerUUploadFileName(tt.fileName, tt.fileType)
|
||||
resultStem := minerUResultStem(uploadName)
|
||||
responseBody, err := json.Marshal(map[string]any{
|
||||
"results": map[string]any{
|
||||
resultStem: map[string]any{
|
||||
"md_content": "ok",
|
||||
"images": map[string]any{},
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal response: %v", err)
|
||||
}
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
got := capturedRequest{method: r.Method, path: r.URL.Path}
|
||||
file, header, err := r.FormFile("files")
|
||||
if err != nil {
|
||||
got.err = err
|
||||
} else {
|
||||
got.fileName = header.Filename
|
||||
content, readErr := io.ReadAll(file)
|
||||
got.content = string(content)
|
||||
got.err = readErr
|
||||
_ = file.Close()
|
||||
}
|
||||
captured <- got
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write(responseBody)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
reader := NewMinerUReader(map[string]string{"mineru_endpoint": server.URL})
|
||||
result, err := reader.Read(context.Background(), &types.ReadRequest{
|
||||
FileContent: []byte(fileContent),
|
||||
FileName: tt.fileName,
|
||||
FileType: tt.fileType,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Read returned error: %v", err)
|
||||
}
|
||||
if result.MarkdownContent != "ok" {
|
||||
t.Fatalf("MarkdownContent = %q, want %q", result.MarkdownContent, "ok")
|
||||
}
|
||||
|
||||
request := <-captured
|
||||
if request.err != nil {
|
||||
t.Fatalf("capture multipart request: %v", request.err)
|
||||
}
|
||||
if request.method != http.MethodPost {
|
||||
t.Errorf("method = %q, want %q", request.method, http.MethodPost)
|
||||
}
|
||||
if request.path != "/file_parse" {
|
||||
t.Errorf("path = %q, want %q", request.path, "/file_parse")
|
||||
}
|
||||
if request.fileName != tt.want {
|
||||
t.Errorf("multipart filename = %q, want %q", request.fileName, tt.want)
|
||||
}
|
||||
if request.content != fileContent {
|
||||
t.Errorf("multipart content = %q, want %q", request.content, fileContent)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeMinerUMarkdownPreservesMarkdownAndHTML(t *testing.T) {
|
||||
input := strings.Join([]string{
|
||||
"# Heading",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
`<details><summary>text_image</summary>caption</details>`,
|
||||
"",
|
||||
`<table><tr><td><img src="images/profile.jpg"/></td></tr></table>`,
|
||||
}, "\n")
|
||||
|
||||
got := normalizeMinerUMarkdown(input)
|
||||
|
||||
if !strings.Contains(got, "# Heading") {
|
||||
t.Fatalf("expected heading to stay intact, got: %q", got)
|
||||
}
|
||||
if strings.Contains(got, `\# Heading`) {
|
||||
t.Fatalf("expected heading to avoid escaped form, got: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "") {
|
||||
t.Fatalf("expected markdown image syntax to stay intact, got: %q", got)
|
||||
}
|
||||
if strings.Contains(got, `!\[](images/cover.jpg)`) {
|
||||
t.Fatalf("expected markdown image syntax to avoid escaped form, got: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, `<details><summary>text_image</summary>caption</details>`) {
|
||||
t.Fatalf("expected details/summary block to be preserved, got: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, `<img src="images/profile.jpg"/>`) {
|
||||
t.Fatalf("expected html img tag to be preserved, got: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessImagesKeepsReferencedVariants(t *testing.T) {
|
||||
reader := &MinerUReader{}
|
||||
mdContent := strings.Join([]string{
|
||||
"",
|
||||
`<img src="./images/profile.jpg"/>`,
|
||||
``,
|
||||
}, "\n")
|
||||
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
images := map[string]string{
|
||||
"cover.jpg": "data:image/png;base64," + b64,
|
||||
"profile.jpg": "data:image/png;base64," + b64,
|
||||
"plain.jpg": "data:image/png;base64," + b64,
|
||||
}
|
||||
|
||||
refs, gotMarkdown := reader.processImages(mdContent, images)
|
||||
|
||||
if gotMarkdown != mdContent {
|
||||
t.Fatalf("processImages should not rewrite markdown content")
|
||||
}
|
||||
if len(refs) != 3 {
|
||||
t.Fatalf("expected 3 image refs, got %d", len(refs))
|
||||
}
|
||||
}
|
||||
|
||||
// TestProcessImagesMatchesPathsWithSpaces guards against a regression where
|
||||
// MinerU image filenames containing spaces (common on Chinese documents,
|
||||
// e.g. "images/第 1 页.jpg") would be silently dropped because the markdown
|
||||
// regex used to extract refs disallowed whitespace inside the URL group.
|
||||
func TestProcessImagesMatchesPathsWithSpaces(t *testing.T) {
|
||||
reader := &MinerUReader{}
|
||||
mdContent := ""
|
||||
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
images := map[string]string{
|
||||
"第 1 页.jpg": "data:image/png;base64," + b64,
|
||||
}
|
||||
|
||||
refs, _ := reader.processImages(mdContent, images)
|
||||
if len(refs) != 1 {
|
||||
t.Fatalf("expected 1 image ref for path with spaces, got %d", len(refs))
|
||||
}
|
||||
if refs[0].OriginalRef != "images/第 1 页.jpg" {
|
||||
t.Fatalf("unexpected OriginalRef: %q", refs[0].OriginalRef)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
paddleOCRVLCloudDefaultBaseURL = "https://paddleocr.aistudio-app.com/api/v2/ocr/jobs"
|
||||
paddleOCRVLCloudDefaultModel = "PaddleOCR-VL-1.6"
|
||||
paddleOCRVLCloudPollInterval = 5 * time.Second
|
||||
paddleOCRVLCloudTimeout = 600 * time.Second
|
||||
)
|
||||
|
||||
// PaddleOCRVLCloudReader calls the PaddleOCR-VL AI Studio cloud API.
|
||||
// Flow: POST /jobs (multipart) → poll GET /jobs/{id} → download result JSONL,
|
||||
// then fetch each referenced image URL.
|
||||
type PaddleOCRVLCloudReader struct {
|
||||
token string
|
||||
baseURL string
|
||||
model string
|
||||
useSeal bool
|
||||
useChart bool
|
||||
}
|
||||
|
||||
// NewPaddleOCRVLCloudReader creates a reader from ParserEngineOverrides.
|
||||
func NewPaddleOCRVLCloudReader(overrides map[string]string) *PaddleOCRVLCloudReader {
|
||||
return &PaddleOCRVLCloudReader{
|
||||
token: strings.TrimSpace(overrides["paddleocr_vl_cloud_token"]),
|
||||
baseURL: strings.TrimRight(stringOr(overrides["paddleocr_vl_cloud_base_url"], paddleOCRVLCloudDefaultBaseURL), "/"),
|
||||
model: stringOr(overrides["paddleocr_vl_cloud_model"], paddleOCRVLCloudDefaultModel),
|
||||
useSeal: parseBoolOr(overrides["paddleocr_vl_cloud_use_seal_recognition"], true),
|
||||
useChart: parseBoolOr(overrides["paddleocr_vl_cloud_use_chart_recognition"], false),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.token == "" {
|
||||
return &types.ReadResult{Error: "PaddleOCR-VL Cloud token is not configured"}, nil
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(c.baseURL); err != nil {
|
||||
return &types.ReadResult{Error: fmt.Sprintf("PaddleOCR-VL Cloud base URL blocked by SSRF policy: %v", err)}, nil
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] Parsing file=%s size=%d model=%s",
|
||||
req.FileName, len(content), c.model)
|
||||
|
||||
jobID, err := c.submitJob(ctx, req, content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL Cloud submit: %w", err)
|
||||
}
|
||||
|
||||
jsonlURL, err := c.pollJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL Cloud poll: %w", err)
|
||||
}
|
||||
|
||||
mdContent, imagesURL, err := c.fetchResults(jsonlURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL Cloud fetch results: %w", err)
|
||||
}
|
||||
|
||||
mdContent = normalizeHTMLTables(mdContent)
|
||||
|
||||
imageRefs := c.downloadImages(mdContent, imagesURL)
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] Parsed successfully, markdown=%d chars, images=%d",
|
||||
len(mdContent), len(imageRefs))
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) optionalPayload() map[string]interface{} {
|
||||
// Shared with the self-hosted engine so both produce identical output.
|
||||
return paddleOCRVLRecognitionParams(c.useSeal, c.useChart)
|
||||
}
|
||||
|
||||
// --- job submit ---
|
||||
|
||||
type paddleOCRVLCloudSubmitResponse struct {
|
||||
Data struct {
|
||||
JobID string `json:"jobId"`
|
||||
} `json:"data"`
|
||||
ErrorCode int `json:"errorCode"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) submitJob(ctx context.Context, req *types.ReadRequest, content []byte) (string, error) {
|
||||
optional, err := json.Marshal(c.optionalPayload())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal optionalPayload: %w", err)
|
||||
}
|
||||
|
||||
fileName := req.FileName
|
||||
if fileName == "" {
|
||||
ext := strings.TrimPrefix(req.FileType, ".")
|
||||
if ext == "" {
|
||||
ext = "pdf"
|
||||
}
|
||||
fileName = "document." + ext
|
||||
}
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", c.model)
|
||||
_ = writer.WriteField("optionalPayload", string(optional))
|
||||
part, err := writer.CreateFormFile("file", filepath.Base(fileName))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create form file: %w", err)
|
||||
}
|
||||
if _, err := part.Write(content); err != nil {
|
||||
return "", fmt.Errorf("write file content: %w", err)
|
||||
}
|
||||
writer.Close()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL, &body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "bearer "+c.token)
|
||||
httpReq.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 60 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result paddleOCRVLCloudSubmitResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return "", fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if result.Data.JobID == "" {
|
||||
return "", fmt.Errorf("API returned no jobId: %s", string(respBody))
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] job submitted: jobId=%s", result.Data.JobID)
|
||||
return result.Data.JobID, nil
|
||||
}
|
||||
|
||||
// --- polling ---
|
||||
|
||||
type paddleOCRVLCloudPollResponse struct {
|
||||
Data struct {
|
||||
State string `json:"state"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
ExtractProgress struct {
|
||||
TotalPages int `json:"totalPages"`
|
||||
ExtractedPages int `json:"extractedPages"`
|
||||
} `json:"extractProgress"`
|
||||
ResultURL struct {
|
||||
JSONURL string `json:"jsonUrl"`
|
||||
} `json:"resultUrl"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) pollJob(ctx context.Context, jobID string) (string, error) {
|
||||
deadline := time.Now().Add(paddleOCRVLCloudTimeout)
|
||||
pollCount := 0
|
||||
url := c.baseURL + "/" + jobID
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
// Bail out promptly when the caller cancels (task cancelled / timed
|
||||
// out) instead of spinning: client.Do would fail immediately and
|
||||
// sleepCtx returns at once on a cancelled ctx, so without this guard
|
||||
// the loop busy-hammers the cloud API and floods logs until deadline.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
pollCount++
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create poll request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "bearer "+c.token)
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 30 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] poll #%d failed: %v", pollCount, err)
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
continue
|
||||
}
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] poll #%d status %d: %s", pollCount, resp.StatusCode, string(respBody))
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
var pollResp paddleOCRVLCloudPollResponse
|
||||
if err := json.Unmarshal(respBody, &pollResp); err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] poll #%d decode error: %v", pollCount, err)
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
state := strings.ToLower(pollResp.Data.State)
|
||||
if pollCount == 1 || pollCount%6 == 0 || state == "done" || state == "failed" {
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] poll #%d: state=%s pages=%d/%d",
|
||||
pollCount, state, pollResp.Data.ExtractProgress.ExtractedPages, pollResp.Data.ExtractProgress.TotalPages)
|
||||
}
|
||||
|
||||
switch state {
|
||||
case "done":
|
||||
if pollResp.Data.ResultURL.JSONURL == "" {
|
||||
return "", fmt.Errorf("state=done but no jsonUrl")
|
||||
}
|
||||
return pollResp.Data.ResultURL.JSONURL, nil
|
||||
case "failed":
|
||||
return "", fmt.Errorf("task failed: %s", pollResp.Data.ErrorMsg)
|
||||
}
|
||||
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("task timed out after %d polls", pollCount)
|
||||
}
|
||||
|
||||
// --- result parsing ---
|
||||
|
||||
type paddleOCRVLCloudResultLine struct {
|
||||
Result struct {
|
||||
LayoutParsingResults []struct {
|
||||
Markdown struct {
|
||||
Text string `json:"text"`
|
||||
Images map[string]string `json:"images"`
|
||||
} `json:"markdown"`
|
||||
} `json:"layoutParsingResults"`
|
||||
} `json:"result"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) fetchResults(jsonlURL string) (string, map[string]string, error) {
|
||||
if err := utils.ValidateURLForSSRF(jsonlURL); err != nil {
|
||||
return "", nil, fmt.Errorf("jsonl URL blocked by SSRF check: %v", err)
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 120 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Get(jsonlURL)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("download jsonl: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", nil, fmt.Errorf("download jsonl status %d", resp.StatusCode)
|
||||
}
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read jsonl body: %w", err)
|
||||
}
|
||||
|
||||
texts := make([]string, 0)
|
||||
images := make(map[string]string)
|
||||
for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var parsed paddleOCRVLCloudResultLine
|
||||
if err := json.Unmarshal([]byte(line), &parsed); err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] skip malformed jsonl line: %v", err)
|
||||
continue
|
||||
}
|
||||
for _, p := range parsed.Result.LayoutParsingResults {
|
||||
if t := strings.TrimSpace(p.Markdown.Text); t != "" {
|
||||
texts = append(texts, p.Markdown.Text)
|
||||
}
|
||||
for path, u := range p.Markdown.Images {
|
||||
if _, ok := images[path]; !ok {
|
||||
images[path] = u
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] fetched %d page(s), images=%d", len(texts), len(images))
|
||||
return strings.Join(texts, "\n\n"), images, nil
|
||||
}
|
||||
|
||||
// downloadImages fetches each referenced image URL and builds ImageRef entries.
|
||||
func (c *PaddleOCRVLCloudReader) downloadImages(mdContent string, imagesURL map[string]string) []types.ImageRef {
|
||||
var refs []types.ImageRef
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 60 * time.Second, MaxRedirects: 5})
|
||||
|
||||
for ipath, u := range imagesURL {
|
||||
matchedRefs := mineruImageOriginalRefs(mdContent, ipath)
|
||||
if len(matchedRefs) == 0 {
|
||||
continue
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(u); err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] image URL blocked %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
resp, err := client.Get(u)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] download image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes, err := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if err != nil || resp.StatusCode != http.StatusOK {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] read image %s status=%d err=%v", ipath, resp.StatusCode, err)
|
||||
continue
|
||||
}
|
||||
|
||||
ext := strings.TrimPrefix(filepath.Ext(ipath), ".")
|
||||
if ext == "" {
|
||||
ext = "png"
|
||||
}
|
||||
mimeType := mime.TypeByExtension("." + ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
for _, originalRef := range matchedRefs {
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: ipath,
|
||||
OriginalRef: originalRef,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgBytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return refs
|
||||
}
|
||||
|
||||
// PingPaddleOCRVLCloud checks whether the cloud token is present (the API has
|
||||
// no lightweight health endpoint, so we only validate configuration here).
|
||||
func PingPaddleOCRVLCloud(token string) (bool, string) {
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return false, "未配置 PaddleOCR-VL Cloud Token"
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
@@ -0,0 +1,292 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const paddleOCRVLTimeout = 1000 * time.Second // large scanned PDFs can take a while
|
||||
|
||||
// PaddleOCRVLReader calls a self-hosted PaddleOCR-VL pipeline service
|
||||
// (the full document-parsing API, not the bare VLM inference server).
|
||||
//
|
||||
// Flow: POST {endpoint}/layout-parsing with base64 file → synchronous JSON
|
||||
// response containing per-page markdown + inline base64 images.
|
||||
type PaddleOCRVLReader struct {
|
||||
endpoint string
|
||||
useSeal bool
|
||||
useChart bool
|
||||
}
|
||||
|
||||
// NewPaddleOCRVLReader creates a reader from ParserEngineOverrides.
|
||||
func NewPaddleOCRVLReader(overrides map[string]string) *PaddleOCRVLReader {
|
||||
return &PaddleOCRVLReader{
|
||||
endpoint: strings.TrimRight(overrides["paddleocr_vl_endpoint"], "/"),
|
||||
useSeal: parseBoolOr(overrides["paddleocr_vl_use_seal_recognition"], true),
|
||||
useChart: parseBoolOr(overrides["paddleocr_vl_use_chart_recognition"], false),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.endpoint == "" {
|
||||
return &types.ReadResult{Error: "PaddleOCR-VL endpoint is not configured"}, nil
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(c.endpoint); err != nil {
|
||||
return &types.ReadResult{Error: fmt.Sprintf("PaddleOCR-VL endpoint blocked by SSRF policy: %v", err)}, nil
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL] Parsing file=%s size=%d via %s",
|
||||
req.FileName, len(content), c.endpoint)
|
||||
|
||||
mdContent, imagesB64, err := c.callLayoutParsing(ctx, req, content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL layout-parsing: %w", err)
|
||||
}
|
||||
|
||||
// PaddleOCR-VL renders tables as styled HTML (per-cell text-align), which
|
||||
// wastes tokens and defeats the chunker's table-protection logic. Convert
|
||||
// them to Markdown tables (or strip layout attributes when conversion is
|
||||
// not possible) before downstream processing.
|
||||
mdContent = normalizeHTMLTables(mdContent)
|
||||
|
||||
imageRefs, mdContent := c.processImages(mdContent, imagesB64)
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL] Parsed successfully, markdown=%d chars, images=%d",
|
||||
len(mdContent), len(imageRefs))
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// paddleOCRVLRecognitionParams returns the recognition / page-restructuring
|
||||
// parameters shared by the self-hosted (/layout-parsing, top-level body) and
|
||||
// cloud (optionalPayload) request bodies. Keeping both identical ensures the
|
||||
// self-hosted engine reproduces the cloud output: cross-page table merging,
|
||||
// multi-level heading reconstruction, header/footer stripping, and the same
|
||||
// sampling / resolution settings used by the AI Studio service.
|
||||
func paddleOCRVLRecognitionParams(useSeal, useChart bool) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"markdownIgnoreLabels": []string{
|
||||
"header", "header_image", "footer", "footer_image",
|
||||
"number", "footnote", "aside_text",
|
||||
},
|
||||
"useDocOrientationClassify": false,
|
||||
"useDocUnwarping": false,
|
||||
"useLayoutDetection": true,
|
||||
"useChartRecognition": useChart,
|
||||
"useSealRecognition": useSeal,
|
||||
"useOcrForImageBlock": false,
|
||||
"mergeTables": true,
|
||||
"relevelTitles": true,
|
||||
"restructurePages": true,
|
||||
"layoutShapeMode": "auto",
|
||||
"promptLabel": "ocr",
|
||||
"layoutNms": true,
|
||||
"repetitionPenalty": 1,
|
||||
"temperature": 0,
|
||||
"topP": 1,
|
||||
"minPixels": 147384,
|
||||
"maxPixels": 2822400,
|
||||
}
|
||||
}
|
||||
|
||||
// fileTypeCode maps a request to the PaddleOCR-VL fileType field:
|
||||
// 0 = PDF, 1 = image (including TIFF).
|
||||
func fileTypeCode(req *types.ReadRequest) int {
|
||||
ft := strings.ToLower(strings.TrimPrefix(req.FileType, "."))
|
||||
if ft == "" {
|
||||
ft = strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
if ft == "pdf" {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
// paddleOCRVLResponse mirrors the relevant fields of the PaddleX serving
|
||||
// /layout-parsing response. The service returns one entry per page.
|
||||
type paddleOCRVLResponse struct {
|
||||
ErrorCode int `json:"errorCode"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
Result struct {
|
||||
LayoutParsingResults []struct {
|
||||
Markdown struct {
|
||||
Text string `json:"text"`
|
||||
Images map[string]string `json:"images"`
|
||||
} `json:"markdown"`
|
||||
} `json:"layoutParsingResults"`
|
||||
} `json:"result"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLReader) callLayoutParsing(
|
||||
ctx context.Context, req *types.ReadRequest, content []byte,
|
||||
) (string, map[string]string, error) {
|
||||
payload := paddleOCRVLRecognitionParams(c.useSeal, c.useChart)
|
||||
payload["file"] = base64.StdEncoding.EncodeToString(content)
|
||||
payload["fileType"] = fileTypeCode(req)
|
||||
payload["visualize"] = false
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("marshal payload: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(
|
||||
ctx, http.MethodPost, c.endpoint+"/layout-parsing", bytes.NewReader(body),
|
||||
)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: paddleOCRVLTimeout,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read response body: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", nil, fmt.Errorf("PaddleOCR-VL API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result paddleOCRVLResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return "", nil, fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if result.ErrorCode != 0 {
|
||||
return "", nil, fmt.Errorf("PaddleOCR-VL error %d: %s", result.ErrorCode, result.ErrorMsg)
|
||||
}
|
||||
|
||||
pages := result.Result.LayoutParsingResults
|
||||
if len(pages) == 0 {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL] response has no layoutParsingResults")
|
||||
return "", nil, nil
|
||||
}
|
||||
|
||||
// Merge per-page markdown and image dicts into one document.
|
||||
texts := make([]string, 0, len(pages))
|
||||
images := make(map[string]string)
|
||||
for _, p := range pages {
|
||||
if t := strings.TrimSpace(p.Markdown.Text); t != "" {
|
||||
texts = append(texts, p.Markdown.Text)
|
||||
}
|
||||
for path, data := range p.Markdown.Images {
|
||||
if _, ok := images[path]; !ok {
|
||||
images[path] = data
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL] parsed %d page(s), images=%d", len(pages), len(images))
|
||||
return strings.Join(texts, "\n\n"), images, nil
|
||||
}
|
||||
|
||||
// processImages decodes the inline base64 images returned by PaddleOCR-VL and
|
||||
// builds ImageRef entries, matching them against references in the markdown.
|
||||
func (c *PaddleOCRVLReader) processImages(
|
||||
mdContent string, imagesB64 map[string]string,
|
||||
) ([]types.ImageRef, string) {
|
||||
var refs []types.ImageRef
|
||||
|
||||
for ipath, b64Str := range imagesB64 {
|
||||
matchedRefs := mineruImageOriginalRefs(mdContent, ipath)
|
||||
if len(matchedRefs) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var imgBytes []byte
|
||||
var ext string
|
||||
if m := b64DataURIPattern.FindStringSubmatch(b64Str); len(m) == 3 {
|
||||
ext = m[1]
|
||||
decoded, err := base64.StdEncoding.DecodeString(m[2])
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL] decode base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
} else {
|
||||
decoded, err := base64.StdEncoding.DecodeString(b64Str)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL] decode raw base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
ext = strings.TrimPrefix(filepath.Ext(ipath), ".")
|
||||
if ext == "" {
|
||||
ext = "png"
|
||||
}
|
||||
}
|
||||
|
||||
mimeType := mime.TypeByExtension("." + ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
for _, originalRef := range matchedRefs {
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: ipath,
|
||||
OriginalRef: originalRef,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgBytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return refs, mdContent
|
||||
}
|
||||
|
||||
// PingPaddleOCRVL checks whether a self-hosted PaddleOCR-VL service is reachable.
|
||||
func PingPaddleOCRVL(endpoint string) (bool, string) {
|
||||
endpoint = strings.TrimRight(endpoint, "/")
|
||||
if endpoint == "" {
|
||||
return false, "未配置 PaddleOCR-VL 端点"
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(endpoint); err != nil {
|
||||
return false, fmt.Sprintf("PaddleOCR-VL 端点未通过 SSRF 校验: %v", err)
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: 5 * time.Second,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
// The pipeline only exposes POST /layout-parsing; an empty GET should still
|
||||
// produce a routed HTTP response (e.g. 404/405) when the service is up.
|
||||
resp, err := client.Get(endpoint + "/layout-parsing")
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("PaddleOCR-VL 服务不可达: %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode >= 500 {
|
||||
return false, fmt.Sprintf("PaddleOCR-VL 服务返回状态 %d", resp.StatusCode)
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// remoteImageServer serves a valid PNG for every request and counts the hits so
|
||||
// that budget and dedup behaviour can be asserted directly.
|
||||
func remoteImageServer(t *testing.T) (*httptest.Server, *int64) {
|
||||
t.Helper()
|
||||
var hits int64
|
||||
png := createTestPNG(200, 200)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt64(&hits, 1)
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
return ts, &hits
|
||||
}
|
||||
|
||||
func allowLocalhost(t *testing.T) {
|
||||
t.Helper()
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
secutils.ResetSSRFWhitelistForTest()
|
||||
}
|
||||
|
||||
func resolve(t *testing.T, markdown string) (string, []StoredImage, *mockFileService) {
|
||||
t.Helper()
|
||||
fSvc := &mockFileService{}
|
||||
updated, images, err := NewImageResolver().ResolveRemoteImages(
|
||||
context.Background(), markdown, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
return updated, images, fSvc
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HTMLTagRewritesSrcAndKeepsTag(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
markdown := fmt.Sprintf(
|
||||
`<p><img class="screenshot" src="%s/a.png" alt="login screen" width="600"></p>`, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL still present: %s", updated)
|
||||
}
|
||||
if !strings.Contains(updated, images[0].ServingURL) {
|
||||
t.Errorf("serving URL missing from output: %s", updated)
|
||||
}
|
||||
// The surrounding markup is the reason these documents use HTML at all.
|
||||
for _, want := range []string{`<img`, `class="screenshot"`, `alt="login screen"`, `width="600"`, `</p>`} {
|
||||
if !strings.Contains(updated, want) {
|
||||
t.Errorf("tag structure lost, %q missing from: %s", want, updated)
|
||||
}
|
||||
}
|
||||
if strings.Contains(updated, "\n\n<img src=\"%s/html.png\">", ts.URL, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 2 {
|
||||
t.Fatalf("expected 2 stored images, got %d", len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("a remote URL survived: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_IndependentBudgetPerSyntax(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// Adding HTML support must not reduce how many Markdown images a document
|
||||
// already gets resolved: each syntax carries its own budget.
|
||||
var b strings.Builder
|
||||
for i := 0; i < maxRemoteImages; i++ {
|
||||
fmt.Fprintf(&b, "\n\n", i, ts.URL, i)
|
||||
}
|
||||
for i := 0; i < 5; i++ {
|
||||
fmt.Fprintf(&b, "<img src=\"%s/html%d.png\">\n\n", ts.URL, i)
|
||||
}
|
||||
|
||||
updated, images, _ := resolve(t, b.String())
|
||||
|
||||
if len(images) != maxRemoteImages+5 {
|
||||
t.Fatalf("expected %d stored images, got %d", maxRemoteImages+5, len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("a remote URL survived: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HTMLSrcNormalization(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
t.Run("uppercase scheme", func(t *testing.T) {
|
||||
upper := strings.Replace(ts.URL, "http://", "HTTP://", 1)
|
||||
updated, images, _ := resolve(t, fmt.Sprintf(`<img SRC="%s/up.png">`, upper))
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected uppercase scheme to be fetched, got %d", len(images))
|
||||
}
|
||||
// The scheme is normalized because every fetcher downstream compares it
|
||||
// byte-for-byte; an image counted as resolved but unfetchable is worse
|
||||
// than one left alone.
|
||||
if strings.Contains(images[0].OriginalRef, "HTTP://") {
|
||||
t.Errorf("scheme was not normalized: %q", images[0].OriginalRef)
|
||||
}
|
||||
if strings.Contains(updated, "HTTP://") {
|
||||
t.Errorf("document still carries the uppercase scheme: %s", updated)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("markdown uppercase scheme is left alone", func(t *testing.T) {
|
||||
// Markdown targets are handed through unchanged, so a document that
|
||||
// upstream leaves untouched stays untouched.
|
||||
upper := strings.Replace(ts.URL, "http://", "HTTP://", 1)
|
||||
markdown := fmt.Sprintf("", upper)
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no markdown behaviour change, got %d image(s)", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown content changed:\n got %q\nwant %q", updated, markdown)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("padded src", func(t *testing.T) {
|
||||
_, images, _ := resolve(t, fmt.Sprintf(`<img src=" %s/pad.png ">`, ts.URL))
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected padded src to be fetched, got %d", len(images))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("entity encoded query", func(t *testing.T) {
|
||||
var got string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
got = r.URL.RawQuery
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(createTestPNG(200, 200))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, images, _ := resolve(t, fmt.Sprintf(`<img src="%s/q.png?a=1&b=2">`, srv.URL))
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
if got != "a=1&b=2" {
|
||||
t.Errorf("entities were not decoded before the request: raw query %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HTMLSrcFormsOutOfScope(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, hits := remoteImageServer(t)
|
||||
|
||||
cases := map[string]string{
|
||||
"unquoted src": fmt.Sprintf(`<img src=%s/unquoted.png>`, ts.URL),
|
||||
"srcset only": fmt.Sprintf(`<img srcset="%s/set.png 2x">`, ts.URL),
|
||||
"relative src": `<img src="./images/local.png">`,
|
||||
"data uri src": `<img src="data:image/png;base64,AAAA">`,
|
||||
"provider src": `<img src="local://images/x.png">`,
|
||||
"non-img tag": fmt.Sprintf(`<video src="%s/v.mp4">`, ts.URL),
|
||||
}
|
||||
|
||||
for name, markdown := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
before := atomic.LoadInt64(hits)
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no stored images, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("content changed:\n got %q\nwant %q", updated, markdown)
|
||||
}
|
||||
if got := atomic.LoadInt64(hits) - before; got != 0 {
|
||||
t.Errorf("expected no fetch, got %d", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HyphenatedAttributeIsNotSrc(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// A lazy-loading tag carries a placeholder in src and the real image in
|
||||
// data-src. Treating data-src as src would capture whichever came first and
|
||||
// leave the other unvisited, so only src counts.
|
||||
markdown := fmt.Sprintf(`<img data-src="/placeholder.svg" src="%s/real.png">`, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected the src image to be resolved, got %d", len(images))
|
||||
}
|
||||
if images[0].OriginalRef != ts.URL+"/real.png" {
|
||||
t.Errorf("expected src to be the fetched reference, got %s", images[0].OriginalRef)
|
||||
}
|
||||
if !strings.Contains(updated, `data-src="/placeholder.svg"`) {
|
||||
t.Errorf("data-src must be left alone: %s", updated)
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("src was not rewritten: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HyphenatedAttributeOnlyIsNotResolved(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, hits := remoteImageServer(t)
|
||||
|
||||
// The other half of the same trade-off, pinned deliberately. Requiring
|
||||
// whitespace before src means a tag carrying only a hyphenated *-src is no
|
||||
// longer treated as an image at all — the previous boundary did resolve it.
|
||||
// Recovering that needs an attribute parser rather than a regex; capturing a
|
||||
// lazy-loading placeholder while the real src goes unvisited is the worse of
|
||||
// the two failures.
|
||||
for _, attr := range []string{"data-src", "ng-src", "data-original-src"} {
|
||||
t.Run(attr, func(t *testing.T) {
|
||||
before := atomic.LoadInt64(hits)
|
||||
markdown := fmt.Sprintf(`<img %s="%s/only.png">`, attr, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no stored images, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("content changed:\n got %q\nwant %q", updated, markdown)
|
||||
}
|
||||
if got := atomic.LoadInt64(hits) - before; got != 0 {
|
||||
t.Errorf("expected no fetch, got %d", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_WhitelistedHostAgreesWithDocument(t *testing.T) {
|
||||
var hits int64
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt64(&hits, 1)
|
||||
if got := r.URL.RawQuery; got != "a=1&b=2" {
|
||||
t.Errorf("expected entities decoded for the request, got raw query %q", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(createTestPNG(200, 200))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
t.Setenv("IMAGE_HOST_KEEP_URL", strings.TrimPrefix(ts.URL, "http://"))
|
||||
secutils.ResetSSRFWhitelistForTest()
|
||||
|
||||
markdown := fmt.Sprintf(`<img src=" %s/q.png?a=1&b=2 " alt="kept">`, ts.URL)
|
||||
|
||||
updated, images, fSvc := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
if atomic.LoadInt64(&hits) != 1 {
|
||||
t.Fatalf("expected the image to be downloaded once, got %d", hits)
|
||||
}
|
||||
if len(fSvc.saved) != 0 {
|
||||
t.Errorf("a whitelisted image must not be uploaded, got %d uploads", len(fSvc.saved))
|
||||
}
|
||||
|
||||
serving := images[0].ServingURL
|
||||
// ServingURL carries two obligations that have to agree: later stages locate
|
||||
// the image by searching the document for it, and the multimodal stage fetches
|
||||
// that same string.
|
||||
if !strings.Contains(updated, serving) {
|
||||
t.Errorf("ServingURL %q is not present in the document %q", serving, updated)
|
||||
}
|
||||
if serving != ts.URL+"/q.png?a=1&b=2" {
|
||||
t.Errorf("ServingURL must be the normalized fetchable URL, got %q", serving)
|
||||
}
|
||||
if strings.TrimSpace(serving) != serving {
|
||||
t.Errorf("ServingURL must not carry padding: %q", serving)
|
||||
}
|
||||
if strings.Contains(serving, "&") {
|
||||
t.Errorf("ServingURL must not carry HTML entities: %q", serving)
|
||||
}
|
||||
if !strings.Contains(updated, `alt="kept"`) {
|
||||
t.Errorf("the rest of the tag must survive: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_CodeExampleIsRewritten(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// Accepted limitation: this resolver does not parse Markdown structure, so a
|
||||
// tag shown as a documentation example inside a fence is rewritten like any
|
||||
// other. Recognising code spans reliably needs a Markdown parser; a
|
||||
// hand-rolled scanner risks the opposite and far worse failure of treating a
|
||||
// document's real screenshots as examples.
|
||||
markdown := fmt.Sprintf("```html\n<img src=\"%s/example.png\">\n```", ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected the example to be resolved, got %d", len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("expected the example to be rewritten: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_CancelledContextStops(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, hits := remoteImageServer(t)
|
||||
|
||||
var b strings.Builder
|
||||
for i := 0; i < 5; i++ {
|
||||
fmt.Fprintf(&b, "<img src=\"%s/c%d.png\">\n\n", ts.URL, i)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
_, images, err := NewImageResolver().ResolveRemoteImages(ctx, b.String(), &mockFileService{}, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no images after cancellation, got %d", len(images))
|
||||
}
|
||||
if got := atomic.LoadInt64(hits); got != 0 {
|
||||
t.Errorf("expected no fetch after cancellation, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreThenResolveRemoteImages(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// Mirrors the real ingest order, where ResolveAndStore runs first and
|
||||
// ResolveRemoteImages then works on its output. This is the only place the
|
||||
// two HTML resolvers can interfere with each other.
|
||||
markdown := fmt.Sprintf(
|
||||
"<img src=\"images/relative.png\">\n\n<img src=\"%s/remote.png\">", ts.URL)
|
||||
|
||||
fSvc := &mockFileService{}
|
||||
resolver := NewImageResolver()
|
||||
|
||||
afterStore, storeImages, err := resolver.ResolveAndStore(context.Background(), &types.ReadResult{
|
||||
MarkdownContent: markdown,
|
||||
ImageRefs: []types.ImageRef{{
|
||||
OriginalRef: "images/relative.png",
|
||||
Filename: "relative.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: createTestPNG(200, 200),
|
||||
}},
|
||||
}, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveAndStore: %v", err)
|
||||
}
|
||||
if len(storeImages) != 1 {
|
||||
t.Fatalf("expected the relative image stored, got %d", len(storeImages))
|
||||
}
|
||||
|
||||
afterRemote, remoteImages, err := resolver.ResolveRemoteImages(
|
||||
context.Background(), afterStore, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveRemoteImages: %v", err)
|
||||
}
|
||||
if len(remoteImages) != 1 {
|
||||
t.Fatalf("expected the remote image stored, got %d", len(remoteImages))
|
||||
}
|
||||
if strings.Contains(afterRemote, ts.URL) {
|
||||
t.Errorf("remote URL survived: %s", afterRemote)
|
||||
}
|
||||
if strings.Contains(afterRemote, "images/relative.png") {
|
||||
t.Errorf("relative reference was not rewritten by the earlier pass: %s", afterRemote)
|
||||
}
|
||||
if n := strings.Count(afterRemote, "<img"); n != 2 {
|
||||
t.Errorf("expected both tags preserved, found %d", n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// mockFileService is a minimal FileService implementation for testing.
|
||||
type mockFileService struct {
|
||||
saved []savedEntry
|
||||
}
|
||||
|
||||
type savedEntry struct {
|
||||
Data []byte
|
||||
TenantID uint64
|
||||
FileName string
|
||||
}
|
||||
|
||||
func (m *mockFileService) CheckConnectivity(ctx context.Context) error { return nil }
|
||||
func (m *mockFileService) SaveFile(ctx context.Context, file *multipart.FileHeader, tenantID uint64, knowledgeID string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
func (m *mockFileService) SaveBytes(ctx context.Context, data []byte, tenantID uint64, fileName string, temp bool) (string, error) {
|
||||
m.saved = append(m.saved, savedEntry{Data: data, TenantID: tenantID, FileName: fileName})
|
||||
return fmt.Sprintf("local://images/%s", fileName), nil
|
||||
}
|
||||
func (m *mockFileService) GetFile(ctx context.Context, filePath string) (io.ReadCloser, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockFileService) GetFileURL(ctx context.Context, filePath string) (string, error) {
|
||||
return filePath, nil
|
||||
}
|
||||
func (m *mockFileService) DeleteFile(ctx context.Context, filePath string) error { return nil }
|
||||
func (m *mockFileService) CopyFile(ctx context.Context, srcPath string, tenantID uint64, knowledgeID string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_NormalDownload(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
// Create a test HTTP server that serves a real PNG image.
|
||||
pngData := createTestPNG(200, 200)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write(pngData)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("# Hello\n\n\n\nSome text", ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
|
||||
// URL should have been replaced.
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL should have been replaced in markdown, got: %s", updated)
|
||||
}
|
||||
if !strings.Contains(updated, "local://images/") {
|
||||
t.Errorf("expected local:// URL in markdown, got: %s", updated)
|
||||
}
|
||||
|
||||
// Verify saved data.
|
||||
if len(fSvc.saved) != 1 {
|
||||
t.Fatalf("expected 1 saved entry, got %d", len(fSvc.saved))
|
||||
}
|
||||
if fSvc.saved[0].TenantID != 42 {
|
||||
t.Errorf("expected tenantID 42, got %d", fSvc.saved[0].TenantID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_SSRFBlocked(t *testing.T) {
|
||||
// URLs pointing to private IPs should be blocked by SSRF check.
|
||||
markdown := "\n\n"
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Both images should be left unchanged (SSRF blocked).
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 stored images (SSRF blocked), got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown should be unchanged when SSRF blocked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_NonImageContentType(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("<html>not an image</html>"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("", ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images for non-image content type, got %d", len(images))
|
||||
}
|
||||
// Original URL should be preserved.
|
||||
if !strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL should be preserved for non-image content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_ProviderSchemeSkipped(t *testing.T) {
|
||||
markdown := "\n"
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images for provider:// URLs, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown should be unchanged for provider:// URLs")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_MultipleImages(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
pngData := createTestPNG(256, 256)
|
||||
callCount := 0
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write(pngData)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("\n\ntext\n\n\n\n",
|
||||
ts.URL, ts.URL, ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 3 {
|
||||
t.Fatalf("expected 3 stored images, got %d", len(images))
|
||||
}
|
||||
if callCount != 3 {
|
||||
t.Errorf("expected 3 HTTP requests, got %d", callCount)
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("all original URLs should have been replaced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_NoImages(t *testing.T) {
|
||||
markdown := "# Just text\n\nNo images here."
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown should be unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_Server404(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("", ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images for 404, got %d", len(images))
|
||||
}
|
||||
// Original URL preserved on failure.
|
||||
if !strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL should be preserved on download failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtFromURLPath(t *testing.T) {
|
||||
tests := []struct {
|
||||
url string
|
||||
expect string
|
||||
}{
|
||||
{"https://example.com/photo.jpg", ".jpg"},
|
||||
{"https://example.com/photo.JPEG", ".jpeg"},
|
||||
{"https://example.com/photo.png?v=2", ""}, // query param — path.Ext won't catch it cleanly but that's ok
|
||||
{"https://example.com/photo.gif", ".gif"},
|
||||
{"https://example.com/photo.webp", ".webp"},
|
||||
{"https://example.com/photo.bmp", ".bmp"},
|
||||
{"https://example.com/photo.svg", ".svg"},
|
||||
{"https://example.com/photo.pdf", ""},
|
||||
{"https://example.com/noext", ""},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.url, func(t *testing.T) {
|
||||
got := extFromURLPath(tt.url)
|
||||
if got != tt.expect {
|
||||
t.Errorf("extFromURLPath(%q) = %q, want %q", tt.url, got, tt.expect)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/models/utils"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
weKnoraCloudReaderBaseURL = "https://weknora.weixin.qq.com/api/v1/doc"
|
||||
)
|
||||
|
||||
// WeKnoraCloudSignedDocumentReader implements the docreader HTTP protocol with WeKnoraCloud signing.
|
||||
type WeKnoraCloudSignedDocumentReader struct {
|
||||
appID string
|
||||
apiKey string
|
||||
client *http.Client
|
||||
initialPollInterval time.Duration
|
||||
maxPollInterval time.Duration
|
||||
pollTimeout time.Duration
|
||||
}
|
||||
|
||||
func NewWeKnoraCloudSignedDocumentReader(appID, apiKey string) (*WeKnoraCloudSignedDocumentReader, error) {
|
||||
if appID == "" {
|
||||
return nil, fmt.Errorf("WeKnoraCloud appID is required")
|
||||
}
|
||||
if apiKey == "" {
|
||||
return nil, fmt.Errorf("WeKnoraCloud apiKey is required")
|
||||
}
|
||||
clientCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
|
||||
clientCfg.Timeout = 500 * time.Minute
|
||||
return &WeKnoraCloudSignedDocumentReader{
|
||||
appID: appID,
|
||||
apiKey: apiKey,
|
||||
initialPollInterval: 500 * time.Millisecond,
|
||||
maxPollInterval: 10 * time.Second,
|
||||
pollTimeout: 20 * time.Minute,
|
||||
client: secutils.NewSSRFSafeHTTPClient(clientCfg),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) Reconnect(addr string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) IsConnected() bool { return true }
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) ListEngines(ctx context.Context, overrides map[string]string) ([]types.ParserEngineInfo, error) {
|
||||
return []types.ParserEngineInfo{{
|
||||
Name: WeKnoraCloudEngineName,
|
||||
Description: "WeKnoraCloud signed docreader",
|
||||
FileTypes: []string{"docx", "doc", "pdf", "md", "markdown", "xlsx", "xls", "pptx", "ppt"},
|
||||
Available: true,
|
||||
}}, nil
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
logger.Infof(ctx, "[WeKnoraCloud] read start file=%q type=%q engine=%q hasURL=%v contentLen=%d requestID=%q",
|
||||
req.FileName, req.FileType, req.ParserEngine, strings.TrimSpace(req.URL) != "", len(req.FileContent), req.RequestID)
|
||||
body := httpReadRequest{
|
||||
FileName: req.FileName,
|
||||
FileType: req.FileType,
|
||||
URL: req.URL,
|
||||
Title: req.Title,
|
||||
RequestID: req.RequestID,
|
||||
Config: &httpReadConfig{
|
||||
ParserEngine: req.ParserEngine,
|
||||
ParserEngineOverrides: req.ParserEngineOverrides,
|
||||
},
|
||||
}
|
||||
if len(req.FileContent) > 0 {
|
||||
body.FileContent = base64.StdEncoding.EncodeToString(req.FileContent)
|
||||
}
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] marshal read request: %v", err)
|
||||
return nil, fmt.Errorf("http marshal read request: %w", err)
|
||||
}
|
||||
httpReq, err := p.newSignedRequest(ctx, http.MethodPost, weKnoraCloudReaderBaseURL+"/reader", jsonBody)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] signed read request: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http read request failed: %v", err)
|
||||
return nil, fmt.Errorf("http read failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http read unexpected status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
return nil, fmt.Errorf("http read status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
var submit weKnoraCloudAsyncSubmitResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&submit); err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] decode read submit response: %v", err)
|
||||
return nil, fmt.Errorf("http decode read submit response: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(submit.TaskID) == "" {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] submit response missing task_id (status=%q message=%q)", submit.Status, submit.Message)
|
||||
return nil, fmt.Errorf("weknoracloud docreader submit response missing task_id")
|
||||
}
|
||||
logger.Infof(ctx, "[WeKnoraCloud] task submitted task_id=%s file=%q type=%q", submit.TaskID, req.FileName, req.FileType)
|
||||
return p.pollTaskResult(ctx, submit.TaskID)
|
||||
}
|
||||
|
||||
type weKnoraCloudAsyncSubmitResponse struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Status string `json:"status"`
|
||||
Message string `json:"message"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
}
|
||||
|
||||
type weKnoraCloudAsyncTaskResponse struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Status string `json:"status"`
|
||||
Message string `json:"message"`
|
||||
Progress float64 `json:"progress"`
|
||||
Result *httpReadResponse `json:"result"`
|
||||
Error string `json:"error"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
UpdatedAt int64 `json:"updated_at"`
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) pollTaskResult(ctx context.Context, taskID string) (*types.ReadResult, error) {
|
||||
pollCtx := ctx
|
||||
if _, ok := ctx.Deadline(); !ok && p.pollTimeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
pollCtx, cancel = context.WithTimeout(ctx, p.pollTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
statusURL := weKnoraCloudReaderBaseURL + "/" + taskID
|
||||
currentInterval := p.initialPollInterval
|
||||
for {
|
||||
httpReq, err := p.newSignedRequest(pollCtx, http.MethodGet, statusURL, nil)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] poll signed request task_id=%s: %v", taskID, err)
|
||||
return nil, err
|
||||
}
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http poll task_id=%s failed: %v", taskID, err)
|
||||
return nil, fmt.Errorf("http poll task failed: %w", err)
|
||||
}
|
||||
var taskResp weKnoraCloudAsyncTaskResponse
|
||||
func() {
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
err = fmt.Errorf("http poll task status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] poll task_id=%s status %d: %s", taskID, resp.StatusCode, string(bodyBytes))
|
||||
return
|
||||
}
|
||||
if decodeErr := json.NewDecoder(resp.Body).Decode(&taskResp); decodeErr != nil {
|
||||
err = fmt.Errorf("http decode task response: %w", decodeErr)
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] poll task_id=%s decode response: %v", taskID, decodeErr)
|
||||
}
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch taskResp.Status {
|
||||
case "completed":
|
||||
if taskResp.Result == nil {
|
||||
logger.Infof(ctx, "[WeKnoraCloud] task_id=%s completed with no result payload", taskID)
|
||||
return &types.ReadResult{}, nil
|
||||
}
|
||||
res := fromHTTPReadResponse(taskResp.Result)
|
||||
if res.Error != "" {
|
||||
logger.Errorf(ctx, "[WeKnoraCloud] task_id=%s completed with result.error: %s", taskID, res.Error)
|
||||
} else {
|
||||
logger.Debugf(ctx, "[WeKnoraCloud] task_id=%s completed ok markdownLen=%d", taskID, len(res.MarkdownContent))
|
||||
}
|
||||
return res, nil
|
||||
case "failed":
|
||||
if taskResp.Error != "" {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s failed: %s", taskID, taskResp.Error)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task failed: %s", taskResp.Error)
|
||||
}
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s failed: %s", taskID, taskResp.Message)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task failed: %s", taskResp.Message)
|
||||
case "cancelled":
|
||||
if taskResp.Error != "" {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s cancelled: %s", taskID, taskResp.Error)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task cancelled: %s", taskResp.Error)
|
||||
}
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s cancelled", taskID)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task cancelled")
|
||||
}
|
||||
if err := pollCtx.Err(); err != nil {
|
||||
logger.Errorf(ctx, "[WeKnoraCloud] poll task_id=%s aborted before sleep: %v", taskID, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Exponential backoff: multiply by 1.5 each time, cap at maxPollInterval
|
||||
select {
|
||||
case <-pollCtx.Done():
|
||||
logger.Errorf(ctx, "[WeKnoraCloud] poll task_id=%s stopped: %v", taskID, pollCtx.Err())
|
||||
return nil, pollCtx.Err()
|
||||
case <-time.After(currentInterval):
|
||||
// Update interval for next iteration
|
||||
nextInterval := time.Duration(float64(currentInterval) * 1.5)
|
||||
if nextInterval > p.maxPollInterval {
|
||||
nextInterval = p.maxPollInterval
|
||||
}
|
||||
currentInterval = nextInterval
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) newSignedRequest(ctx context.Context, method, url string, body []byte) (*http.Request, error) {
|
||||
requestID := uuid.New().String()
|
||||
if len(body) == 0 {
|
||||
body = []byte("{}")
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http new request %s %s: %v", method, url, err)
|
||||
return nil, fmt.Errorf("http new request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.ContentLength = int64(len(body))
|
||||
for k, v := range utils.Sign(p.appID, p.apiKey, requestID, string(body)) {
|
||||
httpReq.Header.Set(k, v)
|
||||
}
|
||||
return httpReq, nil
|
||||
}
|
||||
@@ -0,0 +1,485 @@
|
||||
// Package web_fetch provides a public URL content fetcher with SSRF protection.
|
||||
package web_fetch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/PuerkitoBio/goquery"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/chromedp/chromedp"
|
||||
)
|
||||
|
||||
const (
|
||||
fetchTimeout = 60 * time.Second
|
||||
pipelineFetchTimeout = 15 * time.Second
|
||||
maxBodySize = 100 * 1024
|
||||
)
|
||||
|
||||
// ErrorCode identifies the stage and class of a fetch failure.
|
||||
type ErrorCode string
|
||||
|
||||
const (
|
||||
ErrorInvalidURL ErrorCode = "invalid_url"
|
||||
ErrorDNS ErrorCode = "dns_failed"
|
||||
ErrorTimeout ErrorCode = "connection_timeout"
|
||||
ErrorTLS ErrorCode = "tls_failed"
|
||||
ErrorHTTP403 ErrorCode = "http_403"
|
||||
ErrorHTTP429 ErrorCode = "http_429"
|
||||
ErrorHTTP5xx ErrorCode = "http_5xx"
|
||||
ErrorHTTPStatus ErrorCode = "http_status"
|
||||
ErrorSSRFRejected ErrorCode = "ssrf_rejected"
|
||||
ErrorRedirectRejected ErrorCode = "redirect_rejected"
|
||||
ErrorRead ErrorCode = "read_failed"
|
||||
ErrorHTMLParse ErrorCode = "html_parse_failed"
|
||||
ErrorEmptyContent ErrorCode = "empty_content"
|
||||
ErrorConnection ErrorCode = "connection_failed"
|
||||
)
|
||||
|
||||
// FetchError carries stable, machine-readable failure details.
|
||||
type FetchError struct {
|
||||
Code ErrorCode
|
||||
Retryable bool
|
||||
Err error
|
||||
}
|
||||
|
||||
func (e *FetchError) Error() string {
|
||||
if e == nil {
|
||||
return ""
|
||||
}
|
||||
if e.Err == nil {
|
||||
return string(e.Code)
|
||||
}
|
||||
return e.Err.Error()
|
||||
}
|
||||
|
||||
func (e *FetchError) Unwrap() error {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return e.Err
|
||||
}
|
||||
|
||||
// ErrorDetails returns stable fields suitable for tool responses and logs.
|
||||
func ErrorDetails(err error) (ErrorCode, bool, string) {
|
||||
if err == nil {
|
||||
return "", false, ""
|
||||
}
|
||||
var fetchErr *FetchError
|
||||
if errors.As(err, &fetchErr) {
|
||||
return fetchErr.Code, fetchErr.Retryable, fetchErr.Error()
|
||||
}
|
||||
return ErrorConnection, true, err.Error()
|
||||
}
|
||||
|
||||
// Fetcher fetches and extracts public web pages through an SSRF-safe client.
|
||||
type Fetcher struct {
|
||||
client *http.Client
|
||||
timeout time.Duration
|
||||
maxBodySize int64
|
||||
validateURL func(string) error
|
||||
resolveIPs func(context.Context, string) ([]net.IP, error)
|
||||
dialContext func(context.Context, string, string) (net.Conn, error)
|
||||
renderBrowser func(context.Context, pinnedTarget) (string, error)
|
||||
}
|
||||
|
||||
type pinnedTarget struct {
|
||||
URL *url.URL
|
||||
Host string
|
||||
Port string
|
||||
IP net.IP
|
||||
}
|
||||
|
||||
type httpFetchResult struct {
|
||||
body []byte
|
||||
finalURL string
|
||||
}
|
||||
|
||||
// NewFetcher creates a production fetcher with DNS and redirect SSRF guards.
|
||||
func NewFetcher() *Fetcher {
|
||||
return newFetcher(fetchTimeout, renderWithChromium)
|
||||
}
|
||||
|
||||
// NewPipelineFetcher creates an HTTP-only fetcher for the chat pipeline.
|
||||
// It keeps the pre-refactor 15s timeout and does not launch Chromium.
|
||||
func NewPipelineFetcher() *Fetcher {
|
||||
return newFetcher(pipelineFetchTimeout, nil)
|
||||
}
|
||||
|
||||
func newFetcher(timeout time.Duration, renderBrowser func(context.Context, pinnedTarget) (string, error)) *Fetcher {
|
||||
config := utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: timeout,
|
||||
MaxRedirects: 10,
|
||||
}
|
||||
fetcher := &Fetcher{
|
||||
timeout: timeout,
|
||||
maxBodySize: maxBodySize,
|
||||
validateURL: utils.ValidateURLForSSRF,
|
||||
resolveIPs: lookupPublicDNS,
|
||||
renderBrowser: renderBrowser,
|
||||
}
|
||||
transport := utils.NewSSRFSafeTransport(config)
|
||||
transport.DialContext = fetcher.pinnedDialContext()
|
||||
fetcher.client = utils.NewSSRFSafeHTTPClientWithTransport(config, transport)
|
||||
return fetcher
|
||||
}
|
||||
|
||||
// FetchURLContent preserves the existing package-level API for the chat pipeline.
|
||||
func FetchURLContent(ctx context.Context, rawURL string) (string, error) {
|
||||
return NewPipelineFetcher().Fetch(ctx, rawURL)
|
||||
}
|
||||
|
||||
// Fetch downloads a page and returns clean text content.
|
||||
func (f *Fetcher) Fetch(ctx context.Context, rawURL string) (string, error) {
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return "", newFetchError(ErrorInvalidURL, false, "url is empty")
|
||||
}
|
||||
parsedURL, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return "", newFetchError(ErrorInvalidURL, false, "invalid URL: %v", err)
|
||||
}
|
||||
if (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") || parsedURL.Hostname() == "" {
|
||||
return "", newFetchError(ErrorInvalidURL, false, "invalid URL format")
|
||||
}
|
||||
if err := f.validateURL(rawURL); err != nil {
|
||||
return "", classifyValidationError(err)
|
||||
}
|
||||
|
||||
requestCtx, cancel := context.WithTimeout(ctx, f.timeout)
|
||||
defer cancel()
|
||||
httpResult, httpErr := f.fetchHTTP(requestCtx, rawURL, parsedURL)
|
||||
if httpErr == nil {
|
||||
content, parseErr := htmlToText(string(httpResult.body))
|
||||
requiresBrowser := parseErr == nil && needsBrowserFallback(content, httpResult.body)
|
||||
if parseErr == nil && strings.TrimSpace(content) != "" && !requiresBrowser {
|
||||
logger.Infof(ctx, "[WebFetch] fetched %s → %d chars", rawURL, len(content))
|
||||
return content, nil
|
||||
}
|
||||
if f.renderBrowser != nil {
|
||||
browserURL := firstNonEmpty(httpResult.finalURL, rawURL)
|
||||
if rendered, browserErr := f.fetchWithBrowser(requestCtx, browserURL); browserErr == nil {
|
||||
content, browserParseErr := htmlToText(rendered)
|
||||
if browserParseErr == nil && strings.TrimSpace(content) != "" {
|
||||
logger.Infof(ctx, "[WebFetch] rendered %s → %d chars", rawURL, len(content))
|
||||
return content, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
if parseErr != nil {
|
||||
return "", parseErr
|
||||
}
|
||||
if strings.TrimSpace(content) == "" || requiresBrowser {
|
||||
return "", newFetchError(ErrorEmptyContent, false, "page contains no readable text")
|
||||
}
|
||||
return content, nil
|
||||
}
|
||||
|
||||
if f.renderBrowser != nil && canRenderAfterHTTPError(httpErr) {
|
||||
if rendered, browserErr := f.fetchWithBrowser(requestCtx, rawURL); browserErr == nil {
|
||||
content, browserParseErr := htmlToText(rendered)
|
||||
if browserParseErr == nil && strings.TrimSpace(content) != "" {
|
||||
logger.Infof(ctx, "[WebFetch] rendered %s → %d chars", rawURL, len(content))
|
||||
return content, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", httpErr
|
||||
}
|
||||
|
||||
func (f *Fetcher) fetchHTTP(ctx context.Context, rawURL string, parsedURL *url.URL) (*httpFetchResult, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return nil, newFetchError(ErrorInvalidURL, false, "invalid URL: %v", err)
|
||||
}
|
||||
setBrowserHeaders(req, parsedURL)
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, classifyRequestError(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, classifyHTTPStatus(resp.StatusCode, resp.Status)
|
||||
}
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, f.maxBodySize))
|
||||
if err != nil {
|
||||
return nil, newFetchError(ErrorRead, true, "read failed: %v", err)
|
||||
}
|
||||
finalURL := rawURL
|
||||
if resp.Request != nil && resp.Request.URL != nil {
|
||||
finalURL = resp.Request.URL.String()
|
||||
}
|
||||
return &httpFetchResult{body: body, finalURL: finalURL}, nil
|
||||
}
|
||||
|
||||
func (f *Fetcher) pinnedDialContext() func(context.Context, string, string) (net.Conn, error) {
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid address %s: %w", address, err)
|
||||
}
|
||||
if utils.IsSystemProxy(address) || utils.IsSSRFWhitelisted(host) {
|
||||
return (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, address)
|
||||
}
|
||||
ips, err := f.resolveIPs(ctx, host)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("DNS resolution failed for %s: %w", host, err)
|
||||
}
|
||||
if len(ips) == 0 {
|
||||
return nil, fmt.Errorf("DNS resolution failed: no addresses for %s", host)
|
||||
}
|
||||
for _, ip := range ips {
|
||||
if !utils.IsPublicIP(ip) {
|
||||
return nil, fmt.Errorf("connection blocked: %s resolves to restricted IP %s", host, ip)
|
||||
}
|
||||
}
|
||||
pinnedAddress := net.JoinHostPort(ips[0].String(), port)
|
||||
if f.dialContext != nil {
|
||||
return f.dialContext(ctx, network, pinnedAddress)
|
||||
}
|
||||
return (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, pinnedAddress)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *Fetcher) fetchWithBrowser(ctx context.Context, rawURL string) (string, error) {
|
||||
target, err := f.resolvePinnedTarget(ctx, rawURL)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return f.renderBrowser(ctx, target)
|
||||
}
|
||||
|
||||
func (f *Fetcher) resolvePinnedTarget(ctx context.Context, rawURL string) (pinnedTarget, error) {
|
||||
parsedURL, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return pinnedTarget{}, newFetchError(ErrorInvalidURL, false, "invalid URL: %v", err)
|
||||
}
|
||||
port := parsedURL.Port()
|
||||
if port == "" {
|
||||
port = "443"
|
||||
if parsedURL.Scheme == "http" {
|
||||
port = "80"
|
||||
}
|
||||
}
|
||||
ips, err := f.resolveIPs(ctx, parsedURL.Hostname())
|
||||
if err != nil {
|
||||
return pinnedTarget{}, newFetchError(ErrorDNS, true, "DNS lookup failed for %s: %v", parsedURL.Hostname(), err)
|
||||
}
|
||||
if len(ips) == 0 {
|
||||
return pinnedTarget{}, newFetchError(ErrorDNS, true, "DNS lookup returned no addresses for %s", parsedURL.Hostname())
|
||||
}
|
||||
if !utils.IsSSRFWhitelisted(parsedURL.Hostname()) {
|
||||
for _, ip := range ips {
|
||||
if !utils.IsPublicIP(ip) {
|
||||
return pinnedTarget{}, newFetchError(ErrorSSRFRejected, false, "host resolves to restricted IP %s", ip)
|
||||
}
|
||||
}
|
||||
}
|
||||
return pinnedTarget{URL: parsedURL, Host: parsedURL.Hostname(), Port: port, IP: ips[0]}, nil
|
||||
}
|
||||
|
||||
func lookupPublicDNS(ctx context.Context, host string) ([]net.IP, error) {
|
||||
return net.DefaultResolver.LookupIP(ctx, "ip", host)
|
||||
}
|
||||
|
||||
func renderWithChromium(ctx context.Context, target pinnedTarget) (string, error) {
|
||||
hostRule := fmt.Sprintf("MAP %s %s, MAP * ~NOTFOUND", target.Host, target.IP.String())
|
||||
opts := append(chromedp.DefaultExecAllocatorOptions[:],
|
||||
chromedp.Flag("host-resolver-rules", hostRule),
|
||||
chromedp.Flag("headless", true),
|
||||
chromedp.Flag("disable-setuid-sandbox", true),
|
||||
chromedp.Flag("disable-dev-shm-usage", true),
|
||||
chromedp.Flag("disable-gpu", true),
|
||||
chromedp.Flag("disable-blink-features", "AutomationControlled"),
|
||||
)
|
||||
allocatorCtx, cancelAllocator := chromedp.NewExecAllocator(ctx, opts...)
|
||||
defer cancelAllocator()
|
||||
browserCtx, cancelBrowser := chromedp.NewContext(allocatorCtx)
|
||||
defer cancelBrowser()
|
||||
browserCtx, cancelTimeout := context.WithTimeout(browserCtx, fetchTimeout)
|
||||
defer cancelTimeout()
|
||||
var html string
|
||||
if err := chromedp.Run(browserCtx,
|
||||
chromedp.Navigate(target.URL.String()),
|
||||
chromedp.WaitReady("body", chromedp.ByQuery),
|
||||
chromedp.OuterHTML("html", &html),
|
||||
); err != nil {
|
||||
return "", fmt.Errorf("chromium render failed: %w", err)
|
||||
}
|
||||
return string(limitBytes([]byte(html), maxBodySize)), nil
|
||||
}
|
||||
|
||||
func limitBytes(value []byte, max int64) []byte {
|
||||
if int64(len(value)) <= max {
|
||||
return value
|
||||
}
|
||||
return value[:max]
|
||||
}
|
||||
|
||||
func needsBrowserFallback(content string, html []byte) bool {
|
||||
trimmed := strings.TrimSpace(strings.ToLower(content))
|
||||
if trimmed == "" || strings.Contains(trimmed, "enable javascript") || strings.Contains(trimmed, "loading...") {
|
||||
return true
|
||||
}
|
||||
if len([]rune(trimmed)) >= 200 {
|
||||
return false
|
||||
}
|
||||
lowerHTML := strings.ToLower(string(html))
|
||||
hasAppRoot := strings.Contains(lowerHTML, `id="app"`) || strings.Contains(lowerHTML, `id='app'`) ||
|
||||
strings.Contains(lowerHTML, `id="root"`) || strings.Contains(lowerHTML, `id='root'`)
|
||||
return hasAppRoot && strings.Contains(lowerHTML, "<script")
|
||||
}
|
||||
|
||||
func canRenderAfterHTTPError(err error) bool {
|
||||
code, _, _ := ErrorDetails(err)
|
||||
return code == ErrorHTTP403 || code == ErrorEmptyContent
|
||||
}
|
||||
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, value := range values {
|
||||
if value != "" {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func setBrowserHeaders(req *http.Request, parsedURL *url.URL) {
|
||||
req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36")
|
||||
req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7")
|
||||
req.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6")
|
||||
req.Header.Set("Accept-Encoding", "identity")
|
||||
req.Header.Set("Cache-Control", "no-cache")
|
||||
req.Header.Set("Pragma", "no-cache")
|
||||
req.Header.Set("Sec-Ch-Ua", `"Chromium";v="131", "Not_A Brand";v="24"`)
|
||||
req.Header.Set("Sec-Ch-Ua-Mobile", "?0")
|
||||
req.Header.Set("Sec-Ch-Ua-Platform", `"macOS"`)
|
||||
req.Header.Set("Sec-Fetch-Dest", "document")
|
||||
req.Header.Set("Sec-Fetch-Mode", "navigate")
|
||||
req.Header.Set("Sec-Fetch-Site", "none")
|
||||
req.Header.Set("Sec-Fetch-User", "?1")
|
||||
req.Header.Set("Upgrade-Insecure-Requests", "1")
|
||||
req.Header.Set("Referer", parsedURL.Scheme+"://"+parsedURL.Hostname()+"/")
|
||||
}
|
||||
|
||||
func classifyHTTPStatus(statusCode int, status string) error {
|
||||
switch {
|
||||
case statusCode == http.StatusForbidden:
|
||||
return newFetchError(ErrorHTTP403, false, "HTTP %s", status)
|
||||
case statusCode == http.StatusTooManyRequests:
|
||||
return newFetchError(ErrorHTTP429, true, "HTTP %s", status)
|
||||
case statusCode >= http.StatusInternalServerError:
|
||||
return newFetchError(ErrorHTTP5xx, true, "HTTP %s", status)
|
||||
default:
|
||||
return newFetchError(ErrorHTTPStatus, false, "HTTP %s", status)
|
||||
}
|
||||
}
|
||||
|
||||
func classifyValidationError(err error) error {
|
||||
message := strings.ToLower(err.Error())
|
||||
if strings.Contains(message, "dns resolution failed") || strings.Contains(message, "dns lookup failed") {
|
||||
return newFetchError(ErrorDNS, true, "DNS lookup failed: %v", err)
|
||||
}
|
||||
return newFetchError(ErrorSSRFRejected, false, "URL rejected: %v", err)
|
||||
}
|
||||
|
||||
func classifyRequestError(err error) error {
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return newFetchError(ErrorTimeout, true, "fetch timed out: %v", err)
|
||||
}
|
||||
var dnsErr *net.DNSError
|
||||
if errors.As(err, &dnsErr) {
|
||||
return newFetchError(ErrorDNS, true, "DNS lookup failed: %v", err)
|
||||
}
|
||||
var netErr net.Error
|
||||
if errors.As(err, &netErr) && netErr.Timeout() {
|
||||
return newFetchError(ErrorTimeout, true, "fetch timed out: %v", err)
|
||||
}
|
||||
var unknownAuthority x509.UnknownAuthorityError
|
||||
var certificateInvalid x509.CertificateInvalidError
|
||||
var hostnameError x509.HostnameError
|
||||
var recordHeaderError tls.RecordHeaderError
|
||||
if errors.As(err, &unknownAuthority) || errors.As(err, &certificateInvalid) ||
|
||||
errors.As(err, &hostnameError) || errors.As(err, &recordHeaderError) {
|
||||
return newFetchError(ErrorTLS, false, "TLS validation failed: %v", err)
|
||||
}
|
||||
if errors.Is(err, utils.ErrSSRFRedirectBlocked) {
|
||||
return newFetchError(ErrorRedirectRejected, false, "redirect rejected: %v", err)
|
||||
}
|
||||
var urlErr *url.Error
|
||||
if errors.As(err, &urlErr) {
|
||||
redirectMessage := strings.ToLower(urlErr.Err.Error())
|
||||
if strings.Contains(redirectMessage, "redirect") || strings.Contains(redirectMessage, "stopped after") {
|
||||
return newFetchError(ErrorRedirectRejected, false, "redirect rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
message := strings.ToLower(err.Error())
|
||||
if strings.Contains(message, "dns resolution failed") || strings.Contains(message, "dns lookup failed") {
|
||||
return newFetchError(ErrorDNS, true, "DNS lookup failed: %v", err)
|
||||
}
|
||||
if strings.Contains(message, "connection blocked:") {
|
||||
return newFetchError(ErrorSSRFRejected, false, "URL rejected: %v", err)
|
||||
}
|
||||
if strings.Contains(message, "certificate") || strings.Contains(message, "tls") {
|
||||
return newFetchError(ErrorTLS, false, "TLS validation failed: %v", err)
|
||||
}
|
||||
return newFetchError(ErrorConnection, true, "fetch failed: %v", err)
|
||||
}
|
||||
|
||||
func newFetchError(code ErrorCode, retryable bool, format string, args ...interface{}) error {
|
||||
return &FetchError{Code: code, Retryable: retryable, Err: fmt.Errorf(format, args...)}
|
||||
}
|
||||
|
||||
func htmlToText(html string) (string, error) {
|
||||
doc, err := goquery.NewDocumentFromReader(strings.NewReader(html))
|
||||
if err != nil {
|
||||
fallback := stripTags(html)
|
||||
if fallback == "" {
|
||||
return "", newFetchError(ErrorHTMLParse, false, "HTML parse failed: %v", err)
|
||||
}
|
||||
return fallback, nil
|
||||
}
|
||||
doc.Find("script, style, nav, footer, header, iframe, noscript, svg, img").Remove()
|
||||
|
||||
var builder strings.Builder
|
||||
doc.Find("body").Each(func(_ int, selection *goquery.Selection) {
|
||||
builder.WriteString(selection.Text())
|
||||
})
|
||||
lines := strings.Split(builder.String(), "\n")
|
||||
cleaned := make([]string, 0, len(lines))
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if line != "" {
|
||||
cleaned = append(cleaned, line)
|
||||
}
|
||||
}
|
||||
return strings.Join(cleaned, "\n"), nil
|
||||
}
|
||||
|
||||
func stripTags(html string) string {
|
||||
var builder strings.Builder
|
||||
inTag := false
|
||||
for _, char := range html {
|
||||
switch char {
|
||||
case '<':
|
||||
inTag = true
|
||||
case '>':
|
||||
inTag = false
|
||||
default:
|
||||
if !inTag {
|
||||
builder.WriteRune(char)
|
||||
}
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(builder.String())
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package web_fetch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFetchURLContentBlocksRedirectToLoopback(t *testing.T) {
|
||||
internalHit := make(chan struct{}, 1)
|
||||
internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
internalHit <- struct{}{}
|
||||
fmt.Fprint(w, "<html><body>WEKNORA_INTERNAL_CANARY_178193</body></html>")
|
||||
}))
|
||||
defer internal.Close()
|
||||
|
||||
if _, err := FetchURLContent(context.Background(), internal.URL); err == nil {
|
||||
t.Fatal("direct loopback URL was not rejected")
|
||||
}
|
||||
|
||||
attacker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, internal.URL, http.StatusFound)
|
||||
}))
|
||||
defer attacker.Close()
|
||||
|
||||
got, err := FetchURLContent(context.Background(), attacker.URL+"/entry")
|
||||
if err == nil {
|
||||
if strings.Contains(got, "WEKNORA_INTERNAL_CANARY_178193") {
|
||||
t.Fatalf("redirect followed to loopback, got %q", got)
|
||||
}
|
||||
t.Fatal("expected redirect to loopback to be blocked")
|
||||
}
|
||||
select {
|
||||
case <-internalHit:
|
||||
t.Fatal("internal loopback server was reached via redirect")
|
||||
default:
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
package web_fetch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/x509"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
return function(request)
|
||||
}
|
||||
|
||||
func TestFetcherFetchSuccess(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte("<html><body><main>official specifications</main></body></html>"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
fetcher := newTestFetcher(server.Client())
|
||||
content, err := fetcher.Fetch(context.Background(), server.URL)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, content, "official specifications")
|
||||
}
|
||||
|
||||
func TestFetcherClassifiesHTTPStatus(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
status int
|
||||
code ErrorCode
|
||||
retryable bool
|
||||
}{
|
||||
{name: "forbidden", status: http.StatusForbidden, code: ErrorHTTP403, retryable: false},
|
||||
{name: "rate limited", status: http.StatusTooManyRequests, code: ErrorHTTP429, retryable: true},
|
||||
{name: "server error", status: http.StatusServiceUnavailable, code: ErrorHTTP5xx, retryable: true},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writer.WriteHeader(test.status)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := newTestFetcher(server.Client()).Fetch(context.Background(), server.URL)
|
||||
code, retryable, _ := ErrorDetails(err)
|
||||
assert.Equal(t, test.code, code)
|
||||
assert.Equal(t, test.retryable, retryable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetcherClassifiesNetworkFailures(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
err error
|
||||
code ErrorCode
|
||||
retryable bool
|
||||
}{
|
||||
{name: "dns", err: &net.DNSError{Err: "no such host", Name: "invalid.example"}, code: ErrorDNS, retryable: true},
|
||||
{name: "timeout", err: context.DeadlineExceeded, code: ErrorTimeout, retryable: true},
|
||||
{name: "tls", err: x509.HostnameError{Host: "example.com"}, code: ErrorTLS, retryable: false},
|
||||
{name: "redirect", err: errors.New("redirect blocked by SSRF private address"), code: ErrorRedirectRejected, retryable: false},
|
||||
{name: "dial-time SSRF", err: errors.New("connection blocked: host resolves to restricted IP"), code: ErrorSSRFRejected, retryable: false},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return nil, test.err
|
||||
})}
|
||||
_, err := newTestFetcher(client).Fetch(context.Background(), "https://example.com")
|
||||
code, retryable, _ := ErrorDetails(err)
|
||||
assert.Equal(t, test.code, code)
|
||||
assert.Equal(t, test.retryable, retryable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetcherClassifiesDNSFailureDuringSSRFValidation(t *testing.T) {
|
||||
fetcher := newTestFetcher(&http.Client{})
|
||||
fetcher.validateURL = func(string) error {
|
||||
return errors.New("SSRF validation failed: DNS resolution failed for hostname unavailable.example")
|
||||
}
|
||||
|
||||
_, err := fetcher.Fetch(context.Background(), "https://unavailable.example")
|
||||
code, retryable, _ := ErrorDetails(err)
|
||||
|
||||
assert.Equal(t, ErrorDNS, code)
|
||||
assert.True(t, retryable)
|
||||
}
|
||||
|
||||
func TestFetcherRejectsEmptyContent(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte("<html><body><script>ignored()</script></body></html>"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := newTestFetcher(server.Client()).Fetch(context.Background(), server.URL)
|
||||
code, retryable, _ := ErrorDetails(err)
|
||||
assert.Equal(t, ErrorEmptyContent, code)
|
||||
assert.False(t, retryable)
|
||||
}
|
||||
|
||||
func TestFetcherUsesBrowserFallbackForClientRenderedPage(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte(`<html><body><div id="app">Loading...</div><script>render()</script></body></html>`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
fetcher := newTestFetcher(server.Client())
|
||||
fetcher.resolveIPs = func(context.Context, string) ([]net.IP, error) {
|
||||
return []net.IP{net.ParseIP("93.184.216.34")}, nil
|
||||
}
|
||||
fetcher.renderBrowser = func(context.Context, pinnedTarget) (string, error) {
|
||||
return `<html><body><main>rendered product specifications</main></body></html>`, nil
|
||||
}
|
||||
|
||||
content, err := fetcher.Fetch(context.Background(), server.URL)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, content, "rendered product specifications")
|
||||
}
|
||||
|
||||
func TestNewFetcherKeepsAgentCompatibleTimeout(t *testing.T) {
|
||||
assert.Equal(t, 60*time.Second, NewFetcher().timeout)
|
||||
}
|
||||
|
||||
func TestNewPipelineFetcherUsesHTTPOnlyAndLegacyTimeout(t *testing.T) {
|
||||
fetcher := NewPipelineFetcher()
|
||||
assert.Equal(t, 15*time.Second, fetcher.timeout)
|
||||
assert.Nil(t, fetcher.renderBrowser)
|
||||
}
|
||||
|
||||
func TestFetcherReturnsErrorWhenBrowserFallbackFailsOnSPA(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte(`<html><body><div id="app">Loading...</div><script>render()</script></body></html>`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
fetcher := newTestFetcher(server.Client())
|
||||
fetcher.resolveIPs = func(context.Context, string) ([]net.IP, error) {
|
||||
return []net.IP{net.ParseIP("93.184.216.34")}, nil
|
||||
}
|
||||
fetcher.renderBrowser = func(context.Context, pinnedTarget) (string, error) {
|
||||
return "", errors.New("browser unavailable")
|
||||
}
|
||||
|
||||
_, err := fetcher.Fetch(context.Background(), server.URL)
|
||||
code, retryable, _ := ErrorDetails(err)
|
||||
assert.Equal(t, ErrorEmptyContent, code)
|
||||
assert.False(t, retryable)
|
||||
}
|
||||
|
||||
func TestFetcherClassifiesInvalidAndSSRFURLs(t *testing.T) {
|
||||
fetcher := NewFetcher()
|
||||
_, invalidErr := fetcher.Fetch(context.Background(), "not-a-url")
|
||||
invalidCode, invalidRetryable, _ := ErrorDetails(invalidErr)
|
||||
assert.Equal(t, ErrorInvalidURL, invalidCode)
|
||||
assert.False(t, invalidRetryable)
|
||||
|
||||
_, ssrfErr := fetcher.Fetch(context.Background(), "http://127.0.0.1:1/private")
|
||||
ssrfCode, ssrfRetryable, _ := ErrorDetails(ssrfErr)
|
||||
assert.Equal(t, ErrorSSRFRejected, ssrfCode)
|
||||
assert.False(t, ssrfRetryable)
|
||||
}
|
||||
|
||||
func TestPinnedDialUsesValidatedIPAndPreservesPort(t *testing.T) {
|
||||
var dialedAddress string
|
||||
fetcher := &Fetcher{
|
||||
resolveIPs: func(context.Context, string) ([]net.IP, error) {
|
||||
return []net.IP{net.ParseIP("93.184.216.34")}, nil
|
||||
},
|
||||
dialContext: func(_ context.Context, _, address string) (net.Conn, error) {
|
||||
dialedAddress = address
|
||||
return nil, errors.New("stop dial")
|
||||
},
|
||||
}
|
||||
|
||||
_, err := fetcher.pinnedDialContext()(context.Background(), "tcp", "example.com:443")
|
||||
|
||||
assert.Equal(t, "93.184.216.34:443", dialedAddress)
|
||||
assert.EqualError(t, err, "stop dial")
|
||||
}
|
||||
|
||||
func TestPinnedDialRejectsRebindingToRestrictedIP(t *testing.T) {
|
||||
dialCalled := false
|
||||
fetcher := &Fetcher{
|
||||
resolveIPs: func(context.Context, string) ([]net.IP, error) {
|
||||
return []net.IP{net.ParseIP("93.184.216.34"), net.ParseIP("127.0.0.1")}, nil
|
||||
},
|
||||
dialContext: func(context.Context, string, string) (net.Conn, error) {
|
||||
dialCalled = true
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
|
||||
_, err := fetcher.pinnedDialContext()(context.Background(), "tcp", "example.com:443")
|
||||
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "connection blocked")
|
||||
assert.False(t, dialCalled)
|
||||
}
|
||||
|
||||
func TestFetcherKeepsOriginalHostForTLSAndHTTPRouting(t *testing.T) {
|
||||
var requestURL string
|
||||
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
|
||||
requestURL = request.URL.Host
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Body: http.NoBody,
|
||||
Request: request,
|
||||
}, nil
|
||||
})}
|
||||
fetcher := newTestFetcher(client)
|
||||
fetcher.validateURL = func(string) error { return nil }
|
||||
|
||||
_, err := fetcher.Fetch(context.Background(), "https://example.com/specs")
|
||||
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, "example.com", requestURL)
|
||||
assert.Contains(t, err.Error(), "no readable text")
|
||||
}
|
||||
|
||||
func newTestFetcher(client *http.Client) *Fetcher {
|
||||
return &Fetcher{
|
||||
client: client,
|
||||
timeout: time.Second,
|
||||
maxBodySize: maxBodySize,
|
||||
validateURL: func(string) error { return nil },
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,265 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultBaiduWebSearchURL is the hardcoded Baidu AI Search API URL.
|
||||
// Not configurable by tenants — prevents SSRF.
|
||||
defaultBaiduWebSearchURL = "https://qianfan.baidubce.com/v2/ai_search/web_search"
|
||||
defaultBaiduTimeout = 15 * time.Second
|
||||
defaultBaiduResults = 5
|
||||
maxBaiduResults = 50
|
||||
maxBaiduQueryUnits = 72
|
||||
)
|
||||
|
||||
// BaiduProvider implements web search using Baidu AI Search API.
|
||||
type BaiduProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
}
|
||||
|
||||
// NewBaiduProvider creates a new Baidu web search provider from parameters.
|
||||
func NewBaiduProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if params.APIKey == "" {
|
||||
return nil, fmt.Errorf("API key is required for Baidu provider")
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultBaiduTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create Baidu HTTP client: %w", err)
|
||||
}
|
||||
return &BaiduProvider{
|
||||
client: client,
|
||||
baseURL: defaultBaiduWebSearchURL, // Hardcoded — not tenant-configurable
|
||||
apiKey: params.APIKey,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name.
|
||||
func (p *BaiduProvider) Name() string {
|
||||
return "baidu"
|
||||
}
|
||||
|
||||
// Search performs a web search using Baidu AI Search API.
|
||||
func (p *BaiduProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
preparedQuery := normalizeBaiduQuery(query)
|
||||
if preparedQuery == "" {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
if preparedQuery != strings.TrimSpace(query) {
|
||||
logger.Infof(ctx, "[WebSearch][Baidu] normalized query to satisfy API constraints")
|
||||
}
|
||||
|
||||
if maxResults <= 0 {
|
||||
maxResults = defaultBaiduResults
|
||||
}
|
||||
if maxResults > maxBaiduResults {
|
||||
maxResults = maxBaiduResults
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[WebSearch][Baidu] query=%q maxResults=%d url=%s", preparedQuery, maxResults, p.baseURL)
|
||||
req, err := p.buildRequest(ctx, preparedQuery, maxResults)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results, err := p.doSearch(ctx, req, includeDate)
|
||||
if err != nil {
|
||||
logger.Warnf(ctx, "[WebSearch][Baidu] failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Baidu] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (p *BaiduProvider) buildRequest(ctx context.Context, query string, maxResults int) (*http.Request, error) {
|
||||
requestBody := baiduSearchRequest{
|
||||
Messages: []baiduMessage{
|
||||
{Role: "user", Content: query},
|
||||
},
|
||||
SearchSource: "baidu_search_v2",
|
||||
ResourceTypeFilter: []baiduResourceTypeFilter{
|
||||
{Type: "web", TopK: maxResults},
|
||||
},
|
||||
}
|
||||
|
||||
bodyBytes, err := json.Marshal(requestBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (p *BaiduProvider) doSearch(ctx context.Context, req *http.Request, includeDate bool) ([]*types.WebSearchResult, error) {
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) // 2MB limit
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.Warnf(ctx, "[WebSearch][Baidu] API returned status %d: %s", resp.StatusCode, string(body))
|
||||
return nil, fmt.Errorf("baidu API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var respData baiduSearchResponse
|
||||
if err := json.Unmarshal(body, &respData); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
// Check for API-level error (returned with 200 status)
|
||||
if respData.Code != 0 {
|
||||
return nil, fmt.Errorf("baidu API error (code %d): %s", respData.Code, respData.Message)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, len(respData.References))
|
||||
for _, ref := range respData.References {
|
||||
result := &types.WebSearchResult{
|
||||
Title: ref.Title,
|
||||
URL: ref.URL,
|
||||
Content: ref.Content,
|
||||
Source: "baidu",
|
||||
}
|
||||
|
||||
if includeDate && ref.Date != "" {
|
||||
if t, err := parseBaiduDate(ref.Date); err == nil {
|
||||
result.PublishedAt = &t
|
||||
}
|
||||
}
|
||||
|
||||
results = append(results, result)
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
var baiduDateRe = regexp.MustCompile(`^(\d{4})-(\d{1,2})-(\d{1,2})(?:\s+(\d{1,2}):(\d{2})(?::(\d{2}))?)?`)
|
||||
|
||||
// parseBaiduDate extracts date components via regex, handling variable formats
|
||||
// like "2025-4-24", "2025-04-27 18:02:00", "2025-05-20 11:58" uniformly.
|
||||
func parseBaiduDate(dateStr string) (time.Time, error) {
|
||||
m := baiduDateRe.FindStringSubmatch(dateStr)
|
||||
if m == nil {
|
||||
return time.Time{}, fmt.Errorf("unable to parse date: %s", dateStr)
|
||||
}
|
||||
// Pad to "YYYY-MM-DD HH:MM:SS" and parse once
|
||||
normalized := fmt.Sprintf("%s-%02s-%02s %02s:%02s:%02s",
|
||||
m[1], m[2], m[3],
|
||||
defaultStr(m[4], "00"), defaultStr(m[5], "00"), defaultStr(m[6], "00"))
|
||||
return time.Parse("2006-01-02 15:04:05", normalized)
|
||||
}
|
||||
|
||||
func defaultStr(s, fallback string) string {
|
||||
if s == "" {
|
||||
return fallback
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Baidu documents the query length limit as 72 chars, counting CJK/full-width
|
||||
// runes as 2. Use a conservative width model so mixed-language input stays
|
||||
// under the API limit.
|
||||
func normalizeBaiduQuery(query string) string {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return ""
|
||||
}
|
||||
if baiduQueryUnits(query) <= maxBaiduQueryUnits {
|
||||
return query
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.Grow(len(query))
|
||||
used := 0
|
||||
for _, r := range query {
|
||||
width := baiduQueryUnitWidth(r)
|
||||
if used+width > maxBaiduQueryUnits {
|
||||
break
|
||||
}
|
||||
b.WriteRune(r)
|
||||
used += width
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func baiduQueryUnits(query string) int {
|
||||
units := 0
|
||||
for _, r := range query {
|
||||
units += baiduQueryUnitWidth(r)
|
||||
}
|
||||
return units
|
||||
}
|
||||
|
||||
func baiduQueryUnitWidth(r rune) int {
|
||||
if r <= unicode.MaxASCII {
|
||||
return 1
|
||||
}
|
||||
return 2
|
||||
}
|
||||
|
||||
// --- Request/Response types ---
|
||||
|
||||
type baiduSearchRequest struct {
|
||||
Messages []baiduMessage `json:"messages"`
|
||||
SearchSource string `json:"search_source"`
|
||||
ResourceTypeFilter []baiduResourceTypeFilter `json:"resource_type_filter"`
|
||||
SearchRecencyFilter string `json:"search_recency_filter,omitempty"`
|
||||
}
|
||||
|
||||
type baiduMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type baiduResourceTypeFilter struct {
|
||||
Type string `json:"type"`
|
||||
TopK int `json:"top_k"`
|
||||
}
|
||||
|
||||
type baiduSearchResponse struct {
|
||||
References []baiduReference `json:"references"`
|
||||
RequestID string `json:"request_id"`
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type baiduReference struct {
|
||||
ID int `json:"id"`
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
Date string `json:"date"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultBingSearchURL is the hardcoded Bing search API URL.
|
||||
// Not configurable by tenants — prevents SSRF.
|
||||
defaultBingSearchURL = "https://api.bing.microsoft.com/v7.0/search"
|
||||
)
|
||||
|
||||
var (
|
||||
defaultUserAgentHeader = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36"
|
||||
defaultBingTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
type bingSafeSearch string
|
||||
|
||||
const (
|
||||
bingSafeSearchOff bingSafeSearch = "Off"
|
||||
bingSafeSearchModerate bingSafeSearch = "Moderate"
|
||||
bingSafeSearchStrict bingSafeSearch = "Strict"
|
||||
)
|
||||
|
||||
type bingFreshness string
|
||||
|
||||
const (
|
||||
bingFreshnessDay = "Day"
|
||||
bingFreshnessWeek = "Week"
|
||||
bingFreshnessMonth = "Month"
|
||||
)
|
||||
|
||||
// BingProvider implements web search using Bing Search API
|
||||
type BingProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
}
|
||||
|
||||
// NewBingProvider creates a new Bing provider from parameters (no environment variables).
|
||||
func NewBingProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if params.APIKey == "" {
|
||||
return nil, fmt.Errorf("API key is required for Bing provider")
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultBingTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &BingProvider{
|
||||
client: client,
|
||||
baseURL: defaultBingSearchURL, // Hardcoded — not tenant-configurable
|
||||
apiKey: params.APIKey,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name
|
||||
func (p *BingProvider) Name() string {
|
||||
return "bing"
|
||||
}
|
||||
|
||||
// Search performs a web search using Bing Search API
|
||||
func (p *BingProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if len(query) == 0 {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Bing] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
|
||||
req, err := p.buildParams(ctx, query, maxResults, includeDate)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results, err := p.doSearch(ctx, req)
|
||||
if err != nil {
|
||||
logger.Warnf(ctx, "[WebSearch][Bing] failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Bing] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (p *BingProvider) doSearch(ctx context.Context, req *http.Request) ([]*types.WebSearchResult, error) {
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.Warnf(ctx, "[WebSearch][Bing] API returned status %d: %s", resp.StatusCode, string(body))
|
||||
return nil, fmt.Errorf("bing API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var respData bingSearchResponse
|
||||
if err := json.Unmarshal(body, &respData); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
results := make([]*types.WebSearchResult, 0, len(respData.WebPages.Value))
|
||||
for _, item := range respData.WebPages.Value {
|
||||
results = append(results, &types.WebSearchResult{
|
||||
Title: item.Name,
|
||||
URL: item.URL,
|
||||
Snippet: item.Snippet,
|
||||
Source: "bing",
|
||||
PublishedAt: &item.DateLastCrawled,
|
||||
})
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// bingSearchResponse defines the response structure for Bing search API.
|
||||
type bingSearchResponse struct {
|
||||
Type string `json:"_type"`
|
||||
QueryContext struct {
|
||||
OriginalQuery string `json:"originalQuery"`
|
||||
} `json:"queryContext"`
|
||||
WebPages struct {
|
||||
WebSearchURL string `json:"webSearchUrl"`
|
||||
TotalEstimatedMatches int `json:"totalEstimatedMatches"`
|
||||
Value []struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
URL string `json:"url"`
|
||||
IsFamilyFriendly bool `json:"isFamilyFriendly"`
|
||||
DisplayURL string `json:"displayUrl"`
|
||||
Snippet string `json:"snippet"`
|
||||
DateLastCrawled time.Time `json:"dateLastCrawled"`
|
||||
SearchTags []struct {
|
||||
Name string `json:"name"`
|
||||
Content string `json:"content"`
|
||||
} `json:"searchTags,omitempty"`
|
||||
About []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"about,omitempty"`
|
||||
} `json:"value"`
|
||||
} `json:"webPages"`
|
||||
RelatedSearches struct {
|
||||
ID string `json:"id"`
|
||||
Value []struct {
|
||||
Text string `json:"text"`
|
||||
DisplayText string `json:"displayText"`
|
||||
WebSearchURL string `json:"webSearchUrl"`
|
||||
} `json:"value"`
|
||||
} `json:"relatedSearches"`
|
||||
RankingResponse struct {
|
||||
Mainline struct {
|
||||
Items []struct {
|
||||
AnswerType string `json:"answerType"`
|
||||
ResultIndex int `json:"resultIndex"`
|
||||
Value struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"value"`
|
||||
} `json:"items"`
|
||||
} `json:"mainline"`
|
||||
Sidebar struct {
|
||||
Items []struct {
|
||||
AnswerType string `json:"answerType"`
|
||||
Value struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"value"`
|
||||
} `json:"items"`
|
||||
} `json:"sidebar"`
|
||||
} `json:"rankingResponse"`
|
||||
}
|
||||
|
||||
func (p *BingProvider) buildParams(ctx context.Context, query string, maxResults int, includeDate bool) (*http.Request, error) {
|
||||
params := url.Values{}
|
||||
params.Set("q", query)
|
||||
params.Set("count", strconv.Itoa(maxResults))
|
||||
|
||||
queryURL := fmt.Sprintf("%s?%s", p.baseURL, params.Encode())
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", queryURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("User-Agent", defaultUserAgentHeader)
|
||||
req.Header.Set("Ocp-Apim-Subscription-Key", p.apiKey)
|
||||
return req, nil
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/PuerkitoBio/goquery"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// DuckDuckGoProvider implements web search using DuckDuckGo (HTML first, API fallback)
|
||||
type DuckDuckGoProvider struct {
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// NewDuckDuckGoProvider creates a new DuckDuckGo provider.
|
||||
// DuckDuckGo is free and requires no API key; optional ProxyURL tunnels requests via an HTTP/HTTPS proxy.
|
||||
func NewDuckDuckGoProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
client, err := NewSearchHTTPClient(30*time.Second, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &DuckDuckGoProvider{client: client}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name
|
||||
func (p *DuckDuckGoProvider) Name() string {
|
||||
return "duckduckgo"
|
||||
}
|
||||
|
||||
// Search performs a web search using DuckDuckGo HTML endpoint with API fallback
|
||||
func (p *DuckDuckGoProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
// Try HTML scraping first (more reliable for general results)
|
||||
htmlResults, err := p.searchHTML(ctx, query, maxResults)
|
||||
if err == nil && len(htmlResults) > 0 {
|
||||
return htmlResults, nil
|
||||
}
|
||||
// Fallback to Instant Answer API
|
||||
apiResults, apiErr := p.searchAPI(ctx, query, maxResults)
|
||||
if apiErr == nil && len(apiResults) > 0 {
|
||||
return apiResults, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("duckduckgo HTML search failed: %w", err)
|
||||
}
|
||||
return nil, fmt.Errorf("duckduckgo API search failed: %w", apiErr)
|
||||
}
|
||||
|
||||
// searchHTML performs a web search using DuckDuckGo HTML endpoint
|
||||
func (p *DuckDuckGoProvider) searchHTML(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
baseURL := "https://html.duckduckgo.com/html/"
|
||||
params := url.Values{}
|
||||
params.Set("q", query)
|
||||
params.Set("kl", "cn-zh")
|
||||
|
||||
reqURL := baseURL + "?" + params.Encode()
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set(
|
||||
"User-Agent",
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
|
||||
)
|
||||
|
||||
curlCommand := fmt.Sprintf(
|
||||
"curl -X GET '%s' -H 'User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'",
|
||||
req.URL.String(),
|
||||
)
|
||||
logger.Infof(ctx, "Curl of request: %s", secutils.SanitizeForLog(curlCommand))
|
||||
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to perform request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
return nil, fmt.Errorf("duckduckgo HTML returned status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
doc, err := goquery.NewDocumentFromReader(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse HTML: %w", err)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, maxResults)
|
||||
doc.Find(".web-result").Each(func(i int, s *goquery.Selection) {
|
||||
if len(results) >= maxResults {
|
||||
return
|
||||
}
|
||||
titleNode := s.Find(".result__a")
|
||||
title := strings.TrimSpace(titleNode.Text())
|
||||
var link string
|
||||
if href, exists := titleNode.Attr("href"); exists {
|
||||
link = cleanDDGURL(href)
|
||||
}
|
||||
snippet := strings.TrimSpace(s.Find(".result__snippet").Text())
|
||||
if title != "" && link != "" {
|
||||
results = append(results, &types.WebSearchResult{
|
||||
Title: title,
|
||||
URL: link,
|
||||
Snippet: snippet,
|
||||
Source: "duckduckgo",
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
logger.Infof(ctx, "DuckDuckGo HTML search returned %d results for query: %s", len(results), query)
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// searchAPI performs a web search using DuckDuckGo API endpoint
|
||||
func (p *DuckDuckGoProvider) searchAPI(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
baseURL := "https://api.duckduckgo.com/"
|
||||
params := url.Values{}
|
||||
params.Set("q", query)
|
||||
params.Set("format", "json")
|
||||
params.Set("no_html", "1")
|
||||
params.Set("skip_disambig", "1")
|
||||
|
||||
reqURL := baseURL + "?" + params.Encode()
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("User-Agent", "WeKnora/1.0")
|
||||
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to perform request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("duckduckgo API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var apiResponse struct {
|
||||
AbstractText string `json:"AbstractText"`
|
||||
AbstractURL string `json:"AbstractURL"`
|
||||
Heading string `json:"Heading"`
|
||||
RelatedTopics []struct {
|
||||
FirstURL string `json:"FirstURL"`
|
||||
Text string `json:"Text"`
|
||||
} `json:"RelatedTopics"`
|
||||
Results []struct {
|
||||
FirstURL string `json:"FirstURL"`
|
||||
Text string `json:"Text"`
|
||||
} `json:"Results"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&apiResponse); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode API response: %w", err)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, maxResults)
|
||||
if apiResponse.AbstractText != "" && apiResponse.AbstractURL != "" {
|
||||
results = append(results, &types.WebSearchResult{
|
||||
Title: apiResponse.Heading,
|
||||
URL: apiResponse.AbstractURL,
|
||||
Snippet: apiResponse.AbstractText,
|
||||
Source: "duckduckgo",
|
||||
})
|
||||
}
|
||||
for _, topic := range apiResponse.RelatedTopics {
|
||||
if len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
if topic.Text != "" && topic.FirstURL != "" {
|
||||
results = append(results, &types.WebSearchResult{
|
||||
Title: extractTitle(topic.Text),
|
||||
URL: topic.FirstURL,
|
||||
Snippet: topic.Text,
|
||||
Source: "duckduckgo",
|
||||
})
|
||||
}
|
||||
}
|
||||
for _, r := range apiResponse.Results {
|
||||
if len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
if r.Text != "" && r.FirstURL != "" {
|
||||
results = append(results, &types.WebSearchResult{
|
||||
Title: extractTitle(r.Text),
|
||||
URL: r.FirstURL,
|
||||
Snippet: r.Text,
|
||||
Source: "duckduckgo",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "DuckDuckGo API search returned %d results for query: %s", len(results), query)
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// cleanDDGURL cleans the URL from DuckDuckGo HTML endpoint
|
||||
func cleanDDGURL(urlStr string) string {
|
||||
if strings.HasPrefix(urlStr, "//duckduckgo.com/l/?uddg=") {
|
||||
trimmed := strings.TrimPrefix(urlStr, "//duckduckgo.com/l/?uddg=")
|
||||
if idx := strings.Index(trimmed, "&rut="); idx != -1 {
|
||||
decodedStr, err := url.PathUnescape(trimmed[:idx])
|
||||
if err == nil {
|
||||
return decodedStr
|
||||
}
|
||||
return ""
|
||||
}
|
||||
}
|
||||
if strings.HasPrefix(urlStr, "https://duckduckgo.com/l/?uddg=") {
|
||||
if parsedURL, err := url.Parse(urlStr); err == nil {
|
||||
if uddg := parsedURL.Query().Get("uddg"); uddg != "" {
|
||||
return uddg
|
||||
}
|
||||
}
|
||||
}
|
||||
return urlStr
|
||||
}
|
||||
|
||||
// extractTitle extracts the title from the text
|
||||
func extractTitle(text string) string {
|
||||
lines := strings.Split(text, "\n")
|
||||
if len(lines) > 0 {
|
||||
title := strings.TrimSpace(lines[0])
|
||||
if len(title) > 100 {
|
||||
title = title[:100] + "..."
|
||||
}
|
||||
return title
|
||||
}
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultExaSearchURL = "https://api.exa.ai/search"
|
||||
defaultExaTimeout = 15 * time.Second
|
||||
defaultExaResults = 5
|
||||
maxExaResults = 100
|
||||
maxExaResponseBytes = 2 << 20
|
||||
maxExaContentRunes = 12000
|
||||
)
|
||||
|
||||
// ExaProvider implements web search using Exa's official Search API.
|
||||
type ExaProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
includeText bool
|
||||
}
|
||||
|
||||
// NewExaProvider creates an Exa provider from tenant-specific parameters.
|
||||
func NewExaProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
apiKey := strings.TrimSpace(params.APIKey)
|
||||
if apiKey == "" {
|
||||
return nil, fmt.Errorf("API key is required for Exa provider")
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultExaTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ExaProvider{
|
||||
client: client,
|
||||
baseURL: defaultExaSearchURL,
|
||||
apiKey: apiKey,
|
||||
includeText: parseExaBool(params.ExtraConfig, "include_text"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider type identifier.
|
||||
func (p *ExaProvider) Name() string { return "exa" }
|
||||
|
||||
// Search performs a web search through Exa's official Search API.
|
||||
func (p *ExaProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
if maxResults <= 0 {
|
||||
maxResults = defaultExaResults
|
||||
}
|
||||
if maxResults > maxExaResults {
|
||||
maxResults = maxExaResults
|
||||
}
|
||||
|
||||
bodyBytes, err := json.Marshal(exaSearchRequest{
|
||||
Query: query,
|
||||
NumResults: maxResults,
|
||||
Contents: exaContents{
|
||||
Highlights: true,
|
||||
Text: p.includeText,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal Exa request: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create Exa request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("x-api-key", p.apiKey)
|
||||
|
||||
logger.Infof(ctx, "[WebSearch][Exa] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute Exa request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxExaResponseBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read Exa response: %w", err)
|
||||
}
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
logger.Warnf(ctx, "[WebSearch][Exa] API returned status %d: %s", resp.StatusCode, string(body))
|
||||
return nil, fmt.Errorf("exa API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var data exaSearchResponse
|
||||
if err := json.Unmarshal(body, &data); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal Exa response: %w", err)
|
||||
}
|
||||
if data.Error != "" {
|
||||
return nil, fmt.Errorf("exa API error: %s", data.Error)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, len(data.Results))
|
||||
for _, item := range data.Results {
|
||||
if len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
snippet := strings.TrimSpace(strings.Join(item.Highlights, "\n"))
|
||||
content := truncateExaText(strings.TrimSpace(item.Text), maxExaContentRunes)
|
||||
if snippet == "" {
|
||||
snippet = truncateExaText(content, 500)
|
||||
}
|
||||
result := &types.WebSearchResult{Title: item.Title, URL: item.URL, Snippet: snippet, Content: content, Source: "exa"}
|
||||
if includeDate && item.PublishedDate != "" {
|
||||
if publishedAt, err := time.Parse(time.RFC3339, item.PublishedDate); err == nil {
|
||||
result.PublishedAt = &publishedAt
|
||||
}
|
||||
}
|
||||
results = append(results, result)
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Exa] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func parseExaBool(config map[string]string, key string) bool {
|
||||
v, err := strconv.ParseBool(strings.TrimSpace(config[key]))
|
||||
return err == nil && v
|
||||
}
|
||||
|
||||
func truncateExaText(value string, maxRunes int) string {
|
||||
if maxRunes <= 0 {
|
||||
return ""
|
||||
}
|
||||
if utf8.RuneCountInString(value) <= maxRunes {
|
||||
return value
|
||||
}
|
||||
byteIndex := 0
|
||||
for runeIndex := 0; runeIndex < maxRunes && byteIndex < len(value); runeIndex++ {
|
||||
_, size := utf8.DecodeRuneInString(value[byteIndex:])
|
||||
byteIndex += size
|
||||
}
|
||||
return value[:byteIndex]
|
||||
}
|
||||
|
||||
type exaSearchRequest struct {
|
||||
Query string `json:"query"`
|
||||
NumResults int `json:"numResults"`
|
||||
Contents exaContents `json:"contents"`
|
||||
}
|
||||
|
||||
type exaContents struct {
|
||||
Highlights bool `json:"highlights"`
|
||||
Text bool `json:"text,omitempty"`
|
||||
}
|
||||
|
||||
type exaSearchResponse struct {
|
||||
Results []exaResult `json:"results"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type exaResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
PublishedDate string `json:"publishedDate,omitempty"`
|
||||
Highlights []string `json:"highlights,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestExaProviderSearchMapping(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost || r.URL.Path != "/search" {
|
||||
t.Fatalf("request = %s %s", r.Method, r.URL.Path)
|
||||
}
|
||||
if got := r.Header.Get("x-api-key"); got != "exa-test" {
|
||||
t.Fatalf("x-api-key = %q", got)
|
||||
}
|
||||
var request exaSearchRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if request.Query != "hello" || request.NumResults != 2 ||
|
||||
!request.Contents.Highlights || !request.Contents.Text {
|
||||
t.Fatalf("request body = %+v", request)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(exaSearchResponse{Results: []exaResult{
|
||||
{
|
||||
Title: "One",
|
||||
URL: "https://example.com/1",
|
||||
Highlights: []string{"first", "second"},
|
||||
Text: "body",
|
||||
PublishedDate: "2026-05-01T00:00:00Z",
|
||||
},
|
||||
{Title: "Two", URL: "https://example.com/2"},
|
||||
{Title: "Three", URL: "https://example.com/3"},
|
||||
}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &ExaProvider{client: srv.Client(), baseURL: srv.URL + "/search", apiKey: "exa-test", includeText: true}
|
||||
results, err := p.Search(context.Background(), "hello", 2, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 2 || results[0].Snippet != "first\nsecond" ||
|
||||
results[0].Content != "body" || results[0].Source != "exa" {
|
||||
t.Fatalf("unexpected results: %+v", results)
|
||||
}
|
||||
if results[0].PublishedAt == nil || !results[0].PublishedAt.Equal(time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC)) {
|
||||
t.Fatalf("unexpected date: %v", results[0].PublishedAt)
|
||||
}
|
||||
if results[1].Snippet != "" {
|
||||
t.Fatalf("unexpected empty-result fallback: %+v", results[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExaProviderSupportsConnectionTestProbe(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost || r.URL.Path != "/search" {
|
||||
t.Fatalf("request = %s %s", r.Method, r.URL.Path)
|
||||
}
|
||||
if got := r.Header.Get("x-api-key"); got != "exa-test" {
|
||||
t.Fatalf("x-api-key = %q", got)
|
||||
}
|
||||
var request exaSearchRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if request.Query != "test" || request.NumResults != 1 {
|
||||
t.Fatalf("connection test request = %+v", request)
|
||||
}
|
||||
if !request.Contents.Highlights || request.Contents.Text {
|
||||
t.Fatalf("connection test should request highlights without page text: %+v", request.Contents)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(exaSearchResponse{Results: []exaResult{
|
||||
{Title: "Exa", URL: "https://exa.ai/"},
|
||||
}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &ExaProvider{
|
||||
client: srv.Client(),
|
||||
baseURL: srv.URL + "/search",
|
||||
apiKey: "exa-test",
|
||||
}
|
||||
results, err := p.Search(context.Background(), "test", 1, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("len(results) = %d, want 1", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExaProviderValidationAndStatus(t *testing.T) {
|
||||
if _, err := NewExaProvider(types.WebSearchProviderParameters{}); err == nil {
|
||||
t.Fatal("expected missing API key error")
|
||||
}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = w.Write([]byte(`{"error":"rate limited"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
p := &ExaProvider{client: srv.Client(), baseURL: srv.URL, apiKey: "key"}
|
||||
if _, err := p.Search(context.Background(), "q", 1, false); err == nil {
|
||||
t.Fatal("expected status error")
|
||||
}
|
||||
if _, err := p.Search(context.Background(), " ", 1, false); err == nil {
|
||||
t.Fatal("expected empty query error")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"google.golang.org/api/customsearch/v1"
|
||||
"google.golang.org/api/option"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// GoogleProvider implements web search using Google Custom Search Engine API
|
||||
type GoogleProvider struct {
|
||||
srv *customsearch.Service
|
||||
apiKey string
|
||||
engineID string
|
||||
}
|
||||
|
||||
// NewGoogleProvider creates a new Google provider from parameters (no environment variables).
|
||||
// The API endpoint is the official Google Custom Search endpoint — not tenant-configurable.
|
||||
func NewGoogleProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if params.APIKey == "" {
|
||||
return nil, fmt.Errorf("API key is required for Google provider")
|
||||
}
|
||||
if params.EngineID == "" {
|
||||
return nil, fmt.Errorf("engine ID is required for Google provider")
|
||||
}
|
||||
|
||||
httpClient, err := NewSearchHTTPClient(30*time.Second, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
clientOpts := []option.ClientOption{
|
||||
option.WithAPIKey(params.APIKey),
|
||||
option.WithHTTPClient(httpClient),
|
||||
}
|
||||
srv, err := customsearch.NewService(context.Background(), clientOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &GoogleProvider{
|
||||
srv: srv,
|
||||
apiKey: params.APIKey,
|
||||
engineID: params.EngineID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name
|
||||
func (p *GoogleProvider) Name() string {
|
||||
return "google"
|
||||
}
|
||||
|
||||
// Search performs a web search using Google Custom Search Engine API
|
||||
func (p *GoogleProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if len(query) == 0 {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Google] query=%q maxResults=%d engineID=%s", query, maxResults, p.engineID)
|
||||
cseCall := p.srv.Cse.List().Context(ctx).Cx(p.engineID).Q(query)
|
||||
|
||||
if maxResults > 0 {
|
||||
cseCall = cseCall.Num(int64(maxResults))
|
||||
} else {
|
||||
cseCall = cseCall.Num(5)
|
||||
}
|
||||
cseCall = cseCall.Hl("ch-zh")
|
||||
|
||||
resp, err := cseCall.Do()
|
||||
if err != nil {
|
||||
logger.Warnf(ctx, "[WebSearch][Google] failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
results := make([]*types.WebSearchResult, 0)
|
||||
for _, item := range resp.Items {
|
||||
result := &types.WebSearchResult{
|
||||
Title: item.Title,
|
||||
URL: item.Link,
|
||||
Snippet: item.Snippet,
|
||||
Source: "google",
|
||||
}
|
||||
results = append(results, result)
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Google] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultKeenableBaseURL is the hardcoded Keenable API base URL.
|
||||
// Not configurable by tenants — prevents SSRF.
|
||||
defaultKeenableBaseURL = "https://api.keenable.ai"
|
||||
// keenableTitle is the attribution tag Keenable segments integration traffic by.
|
||||
keenableTitle = "WeKnora"
|
||||
// defaultKeenableResults matches the fallback other providers use when the
|
||||
// caller does not specify a positive maxResults.
|
||||
defaultKeenableResults = 5
|
||||
)
|
||||
|
||||
var (
|
||||
defaultKeenableTimeout = 15 * time.Second
|
||||
)
|
||||
|
||||
// KeenableProvider implements web search using the Keenable Search API.
|
||||
// Keyless by default: with no API key it calls the public endpoint
|
||||
// (rate-limited); a key switches to the authenticated endpoint and lifts the cap.
|
||||
type KeenableProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
}
|
||||
|
||||
// NewKeenableProvider creates a new Keenable provider from parameters.
|
||||
// The API key is optional; Keenable works keyless against the public endpoint.
|
||||
func NewKeenableProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
client, err := NewSearchHTTPClient(defaultKeenableTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &KeenableProvider{
|
||||
client: client,
|
||||
baseURL: defaultKeenableBaseURL,
|
||||
apiKey: params.APIKey,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name
|
||||
func (p *KeenableProvider) Name() string {
|
||||
return "keenable"
|
||||
}
|
||||
|
||||
// Search performs a web search using the Keenable Search API.
|
||||
func (p *KeenableProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if len(query) == 0 {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
if maxResults <= 0 {
|
||||
maxResults = defaultKeenableResults
|
||||
}
|
||||
|
||||
// Keyless by default; a configured key switches to the authenticated path.
|
||||
path := "/v1/search/public"
|
||||
if p.apiKey != "" {
|
||||
path = "/v1/search"
|
||||
}
|
||||
endpoint := p.baseURL + path
|
||||
logger.Infof(ctx, "[WebSearch][Keenable] query=%q maxResults=%d url=%s", query, maxResults, endpoint)
|
||||
|
||||
bodyBytes, err := json.Marshal(keenableSearchRequest{Query: query, Mode: "pro"})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("X-Keenable-Title", keenableTitle)
|
||||
if p.apiKey != "" {
|
||||
req.Header.Set("X-API-Key", p.apiKey)
|
||||
}
|
||||
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
logger.Warnf(ctx, "[WebSearch][Keenable] API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("keenable API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
var respData keenableSearchResponse
|
||||
if err := json.Unmarshal(respBody, &respData); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, len(respData.Results))
|
||||
for _, item := range respData.Results {
|
||||
if maxResults > 0 && len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
// Keenable returns both a short "description" and a longer "snippet"
|
||||
// excerpt. Prefer the short description as the summary snippet and keep
|
||||
// the longer excerpt as Content (used by RAG compression). Fall back to
|
||||
// whichever is present so we never drop the only available text.
|
||||
snippet := item.Description
|
||||
if snippet == "" {
|
||||
snippet = item.Snippet
|
||||
}
|
||||
result := &types.WebSearchResult{
|
||||
Title: item.Title,
|
||||
URL: item.URL,
|
||||
Snippet: snippet,
|
||||
Content: item.Snippet,
|
||||
Source: "keenable",
|
||||
}
|
||||
if includeDate && item.PublishedAt != "" {
|
||||
if t, err := time.Parse(time.RFC3339, item.PublishedAt); err == nil {
|
||||
result.PublishedAt = &t
|
||||
}
|
||||
}
|
||||
results = append(results, result)
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Keenable] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// keenableSearchRequest defines the request body for the Keenable search API.
|
||||
type keenableSearchRequest struct {
|
||||
Query string `json:"query"`
|
||||
Mode string `json:"mode"`
|
||||
}
|
||||
|
||||
// keenableSearchResponse defines the response structure for the Keenable search API.
|
||||
type keenableSearchResponse struct {
|
||||
Query string `json:"query"`
|
||||
Results []struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Description string `json:"description"`
|
||||
Snippet string `json:"snippet,omitempty"`
|
||||
PublishedAt string `json:"published_at,omitempty"`
|
||||
AcquiredAt string `json:"acquired_at,omitempty"`
|
||||
} `json:"results"`
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestKeenableProvider_Search_Keyless(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// No key configured => keyless public endpoint, no X-API-Key.
|
||||
if r.URL.Path != "/v1/search/public" {
|
||||
t.Errorf("path = %q, want /v1/search/public", r.URL.Path)
|
||||
}
|
||||
if got := r.Header.Get("X-API-Key"); got != "" {
|
||||
t.Errorf("X-API-Key = %q, want empty for keyless", got)
|
||||
}
|
||||
if got := r.Header.Get("X-Keenable-Title"); got != "WeKnora" {
|
||||
t.Errorf("X-Keenable-Title = %q, want WeKnora", got)
|
||||
}
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["query"] != "hello" || body["mode"] != "pro" {
|
||||
t.Errorf("body = %v, want query=hello mode=pro", body)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"query": "hello",
|
||||
"results": []map[string]any{
|
||||
{"title": "T1", "url": "https://e/1", "description": "d1", "published_at": "2026-05-01T00:00:00Z"},
|
||||
{"title": "T2", "url": "https://e/2", "description": "d2"},
|
||||
{"title": "T3", "url": "https://e/3", "description": "d3"},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
results, err := p.Search(ctx, "hello", 2, true)
|
||||
if err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("expected 2 results (maxResults trim), got %d", len(results))
|
||||
}
|
||||
if results[0].Source != "keenable" {
|
||||
t.Errorf("source = %q, want keenable", results[0].Source)
|
||||
}
|
||||
if results[0].Snippet != "d1" || results[0].URL != "https://e/1" {
|
||||
t.Errorf("unexpected first result: %+v", results[0])
|
||||
}
|
||||
if results[0].PublishedAt == nil {
|
||||
t.Errorf("expected PublishedAt to be set on the first result")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeenableProvider_Search_SnippetMapping(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"results": []map[string]any{
|
||||
// Has both: description is the short summary, snippet the longer excerpt.
|
||||
{"title": "T1", "url": "https://e/1", "description": "short", "snippet": "long excerpt"},
|
||||
// Only snippet present: it must fall back into Snippet so we don't drop text.
|
||||
{"title": "T2", "url": "https://e/2", "snippet": "only long"},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
results, err := p.Search(ctx, "q", 5, false)
|
||||
if err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("expected 2 results, got %d", len(results))
|
||||
}
|
||||
if results[0].Snippet != "short" || results[0].Content != "long excerpt" {
|
||||
t.Errorf("result[0] snippet/content mapping wrong: %+v", results[0])
|
||||
}
|
||||
if results[1].Snippet != "only long" {
|
||||
t.Errorf("result[1] should fall back to snippet when description is empty: %+v", results[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeenableProvider_Search_Keyed(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// A configured key => authenticated endpoint + X-API-Key.
|
||||
if r.URL.Path != "/v1/search" {
|
||||
t.Errorf("path = %q, want /v1/search", r.URL.Path)
|
||||
}
|
||||
if got := r.Header.Get("X-API-Key"); got != "keen_test" {
|
||||
t.Errorf("X-API-Key = %q, want keen_test", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"results": []map[string]any{}})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL, apiKey: "keen_test"}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if _, err := p.Search(ctx, "hi", 5, false); err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeenableProvider_Search_ErrorStatus(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = w.Write([]byte(`{"message":"rate limited"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
p := &KeenableProvider{client: srv.Client(), baseURL: srv.URL}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if _, err := p.Search(ctx, "x", 5, false); err == nil {
|
||||
t.Fatal("expected an error for a non-200 response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeenableProvider_Name(t *testing.T) {
|
||||
p, err := NewKeenableProvider(types.WebSearchProviderParameters{})
|
||||
if err != nil {
|
||||
t.Fatalf("NewKeenableProvider (keyless): %v", err)
|
||||
}
|
||||
if p.Name() != "keenable" {
|
||||
t.Errorf("Name() = %q, want keenable", p.Name())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultMetasoSearchURL = "https://metaso.cn/api/v1/search"
|
||||
defaultMetasoTimeout = 30 * time.Second
|
||||
defaultMetasoResults = 10
|
||||
maxMetasoResults = 50
|
||||
maxMetasoResponseBytes = 4 << 20
|
||||
defaultMetasoScope = "webpage"
|
||||
)
|
||||
|
||||
var validMetasoScopes = map[string]struct{}{
|
||||
"webpage": {}, "document": {}, "scholar": {},
|
||||
"podcast": {}, "video": {}, "image": {},
|
||||
}
|
||||
|
||||
// MetasoProvider implements web search using the official Metaso AI Search API.
|
||||
type MetasoProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
scope string
|
||||
}
|
||||
|
||||
func NewMetasoProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if err := ValidateMetasoParameters(params); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultMetasoTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &MetasoProvider{
|
||||
client: client, baseURL: defaultMetasoSearchURL,
|
||||
apiKey: strings.TrimSpace(params.APIKey), scope: metasoScope(params.ExtraConfig),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func ValidateMetasoParameters(params types.WebSearchProviderParameters) error {
|
||||
if strings.TrimSpace(params.APIKey) == "" {
|
||||
return fmt.Errorf("API key is required for Metaso provider")
|
||||
}
|
||||
scope := metasoScope(params.ExtraConfig)
|
||||
if _, ok := validMetasoScopes[scope]; !ok {
|
||||
return fmt.Errorf("invalid Metaso search scope: %s", scope)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func metasoScope(extraConfig map[string]string) string {
|
||||
if scope := strings.TrimSpace(extraConfig["scope"]); scope != "" {
|
||||
return scope
|
||||
}
|
||||
return defaultMetasoScope
|
||||
}
|
||||
|
||||
func (p *MetasoProvider) Name() string { return "metaso" }
|
||||
|
||||
func (p *MetasoProvider) Search(ctx context.Context, query string, maxResults int, includeDate bool) ([]*types.WebSearchResult, error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
if maxResults <= 0 {
|
||||
maxResults = defaultMetasoResults
|
||||
}
|
||||
if maxResults > maxMetasoResults {
|
||||
maxResults = maxMetasoResults
|
||||
}
|
||||
|
||||
body, err := json.Marshal(metasoSearchRequest{
|
||||
Query: query, Scope: p.scope, Size: maxResults,
|
||||
IncludeSummary: true, IncludeRawContent: false, ConciseSnippet: true,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal Metaso request: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create Metaso request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
logger.Infof(ctx, "[WebSearch][Metaso] query=%q maxResults=%d scope=%s", query, maxResults, p.scope)
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute Metaso request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
respBody, err := readMetasoResponseBody(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, metasoHTTPError(resp.StatusCode, respBody)
|
||||
}
|
||||
|
||||
var response metasoSearchResponse
|
||||
if err := json.Unmarshal(respBody, &response); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal Metaso response: %w", err)
|
||||
}
|
||||
results := make([]*types.WebSearchResult, 0, len(response.Webpages))
|
||||
for _, item := range response.Webpages {
|
||||
if strings.TrimSpace(item.Title) == "" && strings.TrimSpace(item.Link) == "" {
|
||||
continue
|
||||
}
|
||||
snippet := strings.TrimSpace(item.Summary)
|
||||
if snippet == "" {
|
||||
snippet = strings.TrimSpace(item.Snippet)
|
||||
}
|
||||
result := &types.WebSearchResult{Title: item.Title, URL: item.Link, Snippet: snippet, Content: item.RawContent, Source: "metaso"}
|
||||
if includeDate {
|
||||
if publishedAt, ok := parseMetasoDate(item.Date); ok {
|
||||
result.PublishedAt = &publishedAt
|
||||
}
|
||||
}
|
||||
results = append(results, result)
|
||||
if len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Metaso] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func readMetasoResponseBody(reader io.Reader) ([]byte, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(reader, maxMetasoResponseBytes+1))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read Metaso response: %w", err)
|
||||
}
|
||||
if len(body) > maxMetasoResponseBytes {
|
||||
return nil, fmt.Errorf("Metaso response exceeds %d bytes", maxMetasoResponseBytes)
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func metasoHTTPError(statusCode int, body []byte) error {
|
||||
var apiError struct {
|
||||
Message string `json:"message"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(body, &apiError) == nil {
|
||||
detail := strings.TrimSpace(apiError.Message)
|
||||
if detail == "" {
|
||||
detail = strings.TrimSpace(apiError.Error)
|
||||
}
|
||||
if detail != "" {
|
||||
return fmt.Errorf("Metaso API returned status %d: %s", statusCode, detail)
|
||||
}
|
||||
}
|
||||
detail := strings.TrimSpace(string(body))
|
||||
if len(detail) > 4096 {
|
||||
detail = detail[:4096]
|
||||
}
|
||||
if detail == "" {
|
||||
return fmt.Errorf("Metaso API returned status %d", statusCode)
|
||||
}
|
||||
return fmt.Errorf("Metaso API returned status %d: %s", statusCode, detail)
|
||||
}
|
||||
|
||||
func parseMetasoDate(value string) (time.Time, bool) {
|
||||
for _, layout := range []string{time.RFC3339Nano, "2006-01-02 15:04:05", "2006-01-02"} {
|
||||
if parsed, err := time.Parse(layout, strings.TrimSpace(value)); err == nil {
|
||||
return parsed, true
|
||||
}
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
type metasoSearchRequest struct {
|
||||
Query string `json:"q"`
|
||||
Scope string `json:"scope"`
|
||||
Size int `json:"size"`
|
||||
IncludeSummary bool `json:"includeSummary"`
|
||||
IncludeRawContent bool `json:"includeRawContent"`
|
||||
ConciseSnippet bool `json:"conciseSnippet"`
|
||||
}
|
||||
|
||||
type metasoSearchResponse struct {
|
||||
Webpages []metasoWebpage `json:"webpages"`
|
||||
}
|
||||
|
||||
type metasoWebpage struct {
|
||||
Title string `json:"title"`
|
||||
Link string `json:"link"`
|
||||
Snippet string `json:"snippet"`
|
||||
Summary string `json:"summary"`
|
||||
RawContent string `json:"rawContent"`
|
||||
Date string `json:"date"`
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestMetasoProviderSearch(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Fatalf("method = %s, want POST", r.Method)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer mk-test" {
|
||||
t.Fatalf("Authorization = %q", got)
|
||||
}
|
||||
var request metasoSearchRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if request.Query != "WeKnora" || request.Scope != "scholar" || request.Size != 2 {
|
||||
t.Fatalf("unexpected request: %+v", request)
|
||||
}
|
||||
if !request.IncludeSummary || request.IncludeRawContent || !request.ConciseSnippet {
|
||||
t.Fatalf("unexpected content options: %+v", request)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"webpages":[
|
||||
{"title":"First","link":"https://example.com/1","summary":"Summary","snippet":"Snippet","date":"2026-08-09"},
|
||||
{"title":"Second","link":"https://example.com/2","snippet":"Fallback snippet","date":"invalid"},
|
||||
{"title":"Third","link":"https://example.com/3","snippet":"must be capped"}
|
||||
]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
metaso := &MetasoProvider{client: server.Client(), baseURL: server.URL, apiKey: "mk-test", scope: "scholar"}
|
||||
results, err := metaso.Search(context.Background(), " WeKnora ", 2, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("len(results) = %d, want 2", len(results))
|
||||
}
|
||||
if results[0].Snippet != "Summary" || results[1].Snippet != "Fallback snippet" {
|
||||
t.Fatalf("unexpected snippets: %q, %q", results[0].Snippet, results[1].Snippet)
|
||||
}
|
||||
if results[0].Source != "metaso" || results[0].PublishedAt == nil {
|
||||
t.Fatalf("unexpected first result: %+v", results[0])
|
||||
}
|
||||
if want := time.Date(2026, 8, 9, 0, 0, 0, 0, time.UTC); !results[0].PublishedAt.Equal(want) {
|
||||
t.Fatalf("date = %v, want %v", results[0].PublishedAt, want)
|
||||
}
|
||||
if results[1].PublishedAt != nil {
|
||||
t.Fatalf("invalid date should be ignored: %v", results[1].PublishedAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateMetasoParameters(t *testing.T) {
|
||||
if err := ValidateMetasoParameters(types.WebSearchProviderParameters{}); err == nil {
|
||||
t.Fatal("expected missing API key error")
|
||||
}
|
||||
if err := ValidateMetasoParameters(types.WebSearchProviderParameters{APIKey: "mk-test", ExtraConfig: map[string]string{"scope": "unknown"}}); err == nil {
|
||||
t.Fatal("expected invalid scope error")
|
||||
}
|
||||
if err := ValidateMetasoParameters(types.WebSearchProviderParameters{APIKey: "mk-test"}); err != nil {
|
||||
t.Fatalf("default parameters: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetasoProviderHTTPError(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"message":"invalid API key"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
metaso := &MetasoProvider{client: server.Client(), baseURL: server.URL, apiKey: "bad", scope: defaultMetasoScope}
|
||||
_, err := metaso.Search(context.Background(), "test", 1, false)
|
||||
if err == nil || !strings.Contains(err.Error(), "invalid API key") {
|
||||
t.Fatalf("error = %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultOllamaWebSearchURL is the hardcoded Ollama web search API URL.
|
||||
// Not configurable by tenants — prevents SSRF.
|
||||
defaultOllamaWebSearchURL = "https://ollama.com/api/web_search"
|
||||
defaultOllamaTimeout = 10 * time.Second
|
||||
defaultOllamaResults = 5
|
||||
maxOllamaResults = 10 // Ollama限制最多10个结果
|
||||
)
|
||||
|
||||
// OllamaProvider implements web search using Ollama Cloud API
|
||||
type OllamaProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
}
|
||||
|
||||
// NewOllamaProvider creates a new Ollama web search provider from parameters.
|
||||
func NewOllamaProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if params.APIKey == "" {
|
||||
return nil, fmt.Errorf("API key is required for Ollama provider")
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultOllamaTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create Ollama HTTP client: %w", err)
|
||||
}
|
||||
return &OllamaProvider{
|
||||
client: client,
|
||||
baseURL: defaultOllamaWebSearchURL, // Hardcoded — not tenant-configurable
|
||||
apiKey: params.APIKey,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name
|
||||
func (p *OllamaProvider) Name() string {
|
||||
return "ollama"
|
||||
}
|
||||
|
||||
// Search performs a web search using Ollama Cloud API
|
||||
func (p *OllamaProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if len(query) == 0 {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
|
||||
if maxResults <= 0 {
|
||||
maxResults = defaultOllamaResults
|
||||
}
|
||||
|
||||
// Ollama限制最多10个结果
|
||||
if maxResults > maxOllamaResults {
|
||||
maxResults = maxOllamaResults
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[WebSearch][Ollama] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
|
||||
req, err := p.buildRequest(ctx, query, maxResults)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results, err := p.doSearch(ctx, req)
|
||||
if err != nil {
|
||||
logger.Warnf(ctx, "[WebSearch][Ollama] failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Ollama] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (p *OllamaProvider) buildRequest(ctx context.Context, query string, maxResults int) (*http.Request, error) {
|
||||
requestBody := map[string]interface{}{
|
||||
"query": query,
|
||||
"max_results": maxResults,
|
||||
}
|
||||
|
||||
bodyBytes, err := json.Marshal(requestBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", defaultUserAgentHeader)
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (p *OllamaProvider) doSearch(ctx context.Context, req *http.Request) ([]*types.WebSearchResult, error) {
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.Warnf(ctx, "[WebSearch][Ollama] API returned status %d: %s", resp.StatusCode, string(body))
|
||||
return nil, fmt.Errorf("ollama API returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var respData ollamaSearchResponse
|
||||
if err := json.Unmarshal(body, &respData); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, len(respData.Results))
|
||||
for _, item := range respData.Results {
|
||||
results = append(results, &types.WebSearchResult{
|
||||
Title: item.Title,
|
||||
URL: item.URL,
|
||||
Snippet: item.Snippet,
|
||||
Content: item.Content,
|
||||
Source: "ollama",
|
||||
})
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// ollamaSearchResponse defines the response structure for Ollama web search API.
|
||||
type ollamaSearchResponse struct {
|
||||
Results []ollamaSearchResult `json:"results"`
|
||||
}
|
||||
|
||||
type ollamaSearchResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
Snippet string `json:"snippet"`
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// ValidateProxyURL delegates to utils.ValidateURLForSSRF (only http/https pass that check).
|
||||
func ValidateProxyURL(proxyURL string) error {
|
||||
proxyURL = strings.TrimSpace(proxyURL)
|
||||
if proxyURL == "" {
|
||||
return nil
|
||||
}
|
||||
return utils.ValidateURLForSSRF(proxyURL)
|
||||
}
|
||||
|
||||
// NewSearchHTTPClient builds an http.Client for outbound web search requests.
|
||||
// It uses utils.SSRFSafeDialContext, optional explicit or environment proxy, and
|
||||
// redirect validation consistent with utils.NewSSRFSafeHTTPClient.
|
||||
func NewSearchHTTPClient(timeout time.Duration, proxyURL string) (*http.Client, error) {
|
||||
proxyURL = strings.TrimSpace(proxyURL)
|
||||
def, ok := http.DefaultTransport.(*http.Transport)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("default HTTP transport is not *http.Transport")
|
||||
}
|
||||
t := def.Clone()
|
||||
t.DialContext = utils.SSRFSafeDialContext
|
||||
|
||||
if proxyURL != "" {
|
||||
if err := ValidateProxyURL(proxyURL); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
u, err := url.Parse(proxyURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid proxy_url: %w", err)
|
||||
}
|
||||
if u.Scheme == "" || u.Host == "" {
|
||||
return nil, fmt.Errorf("invalid proxy_url: scheme and host are required")
|
||||
}
|
||||
t.Proxy = http.ProxyURL(u)
|
||||
} else {
|
||||
t.Proxy = http.ProxyFromEnvironment
|
||||
}
|
||||
|
||||
cfg := utils.DefaultSSRFSafeHTTPClientConfig()
|
||||
cfg.Timeout = timeout
|
||||
return utils.NewSSRFSafeHTTPClientWithTransport(cfg, t), nil
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestValidateProxyURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
wantErr bool
|
||||
}{
|
||||
{"empty", "", false},
|
||||
{"whitespace only", " ", false},
|
||||
{"no scheme private IP bare host:port", "192.168.0.1:3128", true},
|
||||
{"localhost http", "http://localhost:3128", true},
|
||||
{"direct IPv4 http", "http://192.168.1.1:3128", true},
|
||||
{"socks5 to IP blocked", "socks5://10.0.0.1:1080", true},
|
||||
{"bad scheme", "ftp://proxy.example:21", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := ValidateProxyURL(tt.raw)
|
||||
if tt.wantErr && err == nil {
|
||||
t.Fatalf("expected error, got nil")
|
||||
}
|
||||
if !tt.wantErr && err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSearchHTTPClient_RejectsUnsafeProxy(t *testing.T) {
|
||||
_, err := NewSearchHTTPClient(0, "http://127.0.0.1:3128")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for disallowed proxy URL (127.0.0.1)")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "SSRF") {
|
||||
t.Fatalf("expected SSRF-related error, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSearchHTTPClient_AcceptsEmptyProxy(t *testing.T) {
|
||||
c, err := NewSearchHTTPClient(0, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if c.Transport == nil {
|
||||
t.Fatal("expected transport")
|
||||
}
|
||||
if c.CheckRedirect == nil {
|
||||
t.Fatal("expected CheckRedirect for SSRF-safe redirects")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSsrfSafeRedirect_BlocksNonHTTPSScheme(t *testing.T) {
|
||||
client, err := NewSearchHTTPClient(0, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req := &http.Request{URL: mustParse(t, "ftp://evil.com/")}
|
||||
// len(via)==0 so first redirect
|
||||
err = client.CheckRedirect(req, nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for ftp redirect")
|
||||
}
|
||||
}
|
||||
|
||||
func mustParse(t *testing.T, s string) *url.URL {
|
||||
t.Helper()
|
||||
u, err := url.Parse(s)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return u
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// ProviderFactory creates a new web search provider instance from parameters.
|
||||
type ProviderFactory func(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error)
|
||||
|
||||
// Registry manages web search provider type registrations.
|
||||
// It maps provider type IDs (e.g., "bing", "google") to their factory functions.
|
||||
// Instances are created on-demand with tenant-specific parameters.
|
||||
type Registry struct {
|
||||
factories map[string]ProviderFactory
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewRegistry creates a new web search provider registry
|
||||
func NewRegistry() *Registry {
|
||||
return &Registry{
|
||||
factories: make(map[string]ProviderFactory),
|
||||
}
|
||||
}
|
||||
|
||||
// Register registers a provider type factory by ID
|
||||
func (r *Registry) Register(id string, factory ProviderFactory) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.factories[id] = factory
|
||||
}
|
||||
|
||||
// CreateProvider creates a provider instance by type with the given parameters.
|
||||
func (r *Registry) CreateProvider(providerType string, params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
r.mu.RLock()
|
||||
factory, ok := r.factories[providerType]
|
||||
r.mu.RUnlock()
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("web search provider type %s not registered", providerType)
|
||||
}
|
||||
return factory(params)
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// defaultSearxngTimeout is sized slightly above the SearXNG image's default
|
||||
// outgoing.max_request_timeout (10s in docker/searxng/settings.yml) so a slow
|
||||
// upstream engine surfaces as a SearXNG-side error instead of a client cancel.
|
||||
const defaultSearxngTimeout = 12 * time.Second
|
||||
|
||||
// SearxngProvider implements web search using a self-hosted SearXNG instance.
|
||||
//
|
||||
// Unlike commercial providers, SearXNG is self-hosted, so the instance URL is
|
||||
// supplied by the tenant via WebSearchProviderParameters.BaseURL. The URL is
|
||||
// validated with utils.ValidateURLForSSRF; private/loopback hosts must be added
|
||||
// to the SSRF_WHITELIST environment variable.
|
||||
type SearxngProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
lastUnresponsive [][]string
|
||||
}
|
||||
|
||||
// ValidateSearxngBaseURL validates a SearXNG instance URL: must be a non-empty,
|
||||
// absolute http(s) URL, and must pass the SSRF whitelist check. Shared between
|
||||
// the service-layer parameter validation and the provider constructor so that
|
||||
// "save" and "use" never disagree.
|
||||
func ValidateSearxngBaseURL(rawURL string) error {
|
||||
base := strings.TrimSpace(rawURL)
|
||||
if base == "" {
|
||||
return fmt.Errorf("base_url is required for SearXNG provider")
|
||||
}
|
||||
parsed, err := url.Parse(base)
|
||||
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
||||
return fmt.Errorf("invalid SearXNG base_url: must be an absolute http(s) URL")
|
||||
}
|
||||
if parsed.Scheme != "http" && parsed.Scheme != "https" {
|
||||
return fmt.Errorf("invalid SearXNG base_url scheme: %s", parsed.Scheme)
|
||||
}
|
||||
if parsed.RawQuery != "" || parsed.Fragment != "" {
|
||||
return fmt.Errorf("invalid SearXNG base_url: must not contain query or fragment")
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(base); err != nil {
|
||||
return fmt.Errorf("invalid SearXNG base_url: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewSearxngProvider builds a SearXNG provider from tenant parameters.
|
||||
func NewSearxngProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
base := strings.TrimSpace(params.BaseURL)
|
||||
if err := ValidateSearxngBaseURL(base); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client, err := NewSearchHTTPClient(defaultSearxngTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &SearxngProvider{
|
||||
client: client,
|
||||
baseURL: strings.TrimRight(base, "/"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name.
|
||||
func (p *SearxngProvider) Name() string { return "searxng" }
|
||||
|
||||
// EmptyResultDiagnostics explains why the most recent search returned no
|
||||
// usable results. Used by the settings "test connection" flow.
|
||||
func (p *SearxngProvider) EmptyResultDiagnostics() string {
|
||||
if detail := formatUnresponsiveEngines(p.lastUnresponsive); detail != "" {
|
||||
return detail + "; check that upstream search engines can reach the internet"
|
||||
}
|
||||
return "verify the instance URL is reachable and JSON format is enabled in settings.yml"
|
||||
}
|
||||
|
||||
// Search performs a metasearch query against the configured SearXNG instance.
|
||||
// SearXNG must have `search.formats: [json]` enabled in settings.yml.
|
||||
func (p *SearxngProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if strings.TrimSpace(query) == "" {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
|
||||
q := url.Values{}
|
||||
q.Set("q", query)
|
||||
q.Set("format", "json")
|
||||
// Use "all" (SearXNG's documented value for "no language filter") instead
|
||||
// of "auto", which is a UI-side default and not a valid /search parameter.
|
||||
// safesearch is intentionally not set here so the value configured in the
|
||||
// instance's settings.yml is honored.
|
||||
q.Set("language", "all")
|
||||
|
||||
reqURL := p.baseURL + "/search?" + q.Encode()
|
||||
logger.Infof(ctx, "[WebSearch][SearXNG] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("User-Agent", "WeKnora/1.0")
|
||||
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
|
||||
return nil, fmt.Errorf("searxng returned status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var data searxngResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
|
||||
p.lastUnresponsive = nil
|
||||
return nil, fmt.Errorf("failed to decode SearXNG response (ensure JSON format is enabled in settings.yml): %w", err)
|
||||
}
|
||||
p.lastUnresponsive = data.UnresponsiveEngines
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, maxResults)
|
||||
for _, r := range data.Results {
|
||||
if len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
if r.URL == "" || r.Title == "" {
|
||||
continue
|
||||
}
|
||||
item := &types.WebSearchResult{
|
||||
Title: r.Title,
|
||||
URL: r.URL,
|
||||
Snippet: r.Content,
|
||||
Source: "searxng",
|
||||
}
|
||||
if includeDate && r.PublishedDate != "" {
|
||||
if t, ok := parseSearxngDate(r.PublishedDate); ok {
|
||||
item.PublishedAt = &t
|
||||
} else {
|
||||
logger.Debugf(ctx, "[WebSearch][SearXNG] unparsable publishedDate=%q", r.PublishedDate)
|
||||
}
|
||||
}
|
||||
results = append(results, item)
|
||||
}
|
||||
if len(results) == 0 && len(data.UnresponsiveEngines) > 0 {
|
||||
logger.Warnf(ctx, "[WebSearch][SearXNG] empty results, unresponsive_engines=%v", data.UnresponsiveEngines)
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][SearXNG] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// searxngDateLayouts covers the formats different SearXNG engines emit for
|
||||
// publishedDate. Order matters only for performance; first match wins.
|
||||
// time.RFC3339 already accepts the nanosecond-precision form, so RFC3339Nano
|
||||
// is intentionally omitted.
|
||||
var searxngDateLayouts = []string{
|
||||
time.RFC3339,
|
||||
"2006-01-02T15:04:05",
|
||||
"2006-01-02 15:04:05",
|
||||
"2006-01-02",
|
||||
time.RFC1123Z,
|
||||
time.RFC1123,
|
||||
}
|
||||
|
||||
func parseSearxngDate(s string) (time.Time, bool) {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return time.Time{}, false
|
||||
}
|
||||
for _, layout := range searxngDateLayouts {
|
||||
if t, err := time.Parse(layout, s); err == nil {
|
||||
return t, true
|
||||
}
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
type searxngResponse struct {
|
||||
Results []struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
PublishedDate string `json:"publishedDate,omitempty"`
|
||||
} `json:"results"`
|
||||
// UnresponsiveEngines is a list of [engine, reason] tuples returned by
|
||||
// SearXNG when one or more upstream engines fail. We surface it on empty
|
||||
// result sets to aid debugging.
|
||||
UnresponsiveEngines [][]string `json:"unresponsive_engines,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
func TestValidateSearxngBaseURL(t *testing.T) {
|
||||
// utils.ValidateURLForSSRF caches the parsed SSRF_WHITELIST via
|
||||
// sync.Once on first call. An alphabetically-earlier test in this
|
||||
// binary (TestValidateProxyURL) triggers ValidateURLForSSRF with an
|
||||
// empty whitelist and caches an empty config, so the later setenv
|
||||
// here would otherwise be ignored. Reset the singleton on both
|
||||
// entry and exit to keep the env / singleton in sync.
|
||||
utils.ResetSSRFWhitelistForTest()
|
||||
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
defer func() {
|
||||
os.Unsetenv("SSRF_WHITELIST")
|
||||
utils.ResetSSRFWhitelistForTest()
|
||||
}()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
url string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "empty", url: "", wantErr: true},
|
||||
{name: "no scheme", url: "searxng:8080", wantErr: true},
|
||||
{name: "bad scheme", url: "ftp://searxng:8080", wantErr: true},
|
||||
{name: "with query", url: "http://127.0.0.1:8080/?x=1", wantErr: true},
|
||||
{name: "with fragment", url: "http://127.0.0.1:8080/#frag", wantErr: true},
|
||||
{name: "loopback ok via whitelist", url: "http://127.0.0.1:8888", wantErr: false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := ValidateSearxngBaseURL(tc.url)
|
||||
if (err != nil) != tc.wantErr {
|
||||
t.Fatalf("ValidateSearxngBaseURL(%q) err=%v wantErr=%v", tc.url, err, tc.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSearxngDate(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
ok bool
|
||||
}{
|
||||
{"", false},
|
||||
{"2024-05-01", true},
|
||||
{"2024-05-01T12:30:45", true},
|
||||
{"2024-05-01T12:30:45Z", true},
|
||||
{"2024-05-01T12:30:45.123456789Z", true},
|
||||
{"2024-05-01 12:30:45", true},
|
||||
{"Wed, 01 May 2024 12:30:45 GMT", true},
|
||||
{"not-a-date", false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.in, func(t *testing.T) {
|
||||
_, ok := parseSearxngDate(tc.in)
|
||||
if ok != tc.ok {
|
||||
t.Fatalf("parseSearxngDate(%q) ok=%v want=%v", tc.in, ok, tc.ok)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearxngProvider_Search(t *testing.T) {
|
||||
// See TestValidateSearxngBaseURL comment — reset the SSRF whitelist
|
||||
// singleton so the setenv below is actually observed by the cached
|
||||
// ssrfWhitelistConfig in internal/utils.
|
||||
utils.ResetSSRFWhitelistForTest()
|
||||
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
defer func() {
|
||||
os.Unsetenv("SSRF_WHITELIST")
|
||||
utils.ResetSSRFWhitelistForTest()
|
||||
}()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/search" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if got := r.URL.Query().Get("format"); got != "json" {
|
||||
t.Errorf("expected format=json, got %q", got)
|
||||
}
|
||||
if got := r.URL.Query().Get("q"); got != "hello" {
|
||||
t.Errorf("expected q=hello, got %q", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"results": []map[string]any{
|
||||
{"title": "T1", "url": "https://e/1", "content": "c1", "publishedDate": "2024-05-01"},
|
||||
{"title": "", "url": "https://e/skip"},
|
||||
{"title": "T2", "url": "https://e/2", "content": "c2"},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
provider, err := NewSearxngProvider(types.WebSearchProviderParameters{BaseURL: srv.URL})
|
||||
if err != nil {
|
||||
t.Fatalf("NewSearxngProvider: %v", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
results, err := provider.Search(ctx, "hello", 5, true)
|
||||
if err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("expected 2 results (one skipped), got %d", len(results))
|
||||
}
|
||||
if results[0].PublishedAt == nil {
|
||||
t.Fatalf("expected first result PublishedAt to be set")
|
||||
}
|
||||
if got := results[0].Source; got != "searxng" {
|
||||
t.Fatalf("unexpected source: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearxngProvider_Search_EmptyWithUnresponsiveEngines(t *testing.T) {
|
||||
utils.ResetSSRFWhitelistForTest()
|
||||
os.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
defer func() {
|
||||
os.Unsetenv("SSRF_WHITELIST")
|
||||
utils.ResetSSRFWhitelistForTest()
|
||||
}()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"results": []any{},
|
||||
"unresponsive_engines": [][]string{{"google", "timeout"}},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
provider, err := NewSearxngProvider(types.WebSearchProviderParameters{BaseURL: srv.URL})
|
||||
if err != nil {
|
||||
t.Fatalf("NewSearxngProvider: %v", err)
|
||||
}
|
||||
sp := provider.(*SearxngProvider)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
results, err := sp.Search(ctx, "test", 1, false)
|
||||
if err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
if len(results) != 0 {
|
||||
t.Fatalf("expected 0 results, got %d", len(results))
|
||||
}
|
||||
if got := sp.EmptyResultDiagnostics(); !strings.Contains(got, "google (timeout)") {
|
||||
t.Fatalf("EmptyResultDiagnostics() = %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultTavilySearchURL is the hardcoded Tavily API URL.
|
||||
// Not configurable by tenants — prevents SSRF.
|
||||
defaultTavilySearchURL = "https://api.tavily.com/search"
|
||||
)
|
||||
|
||||
var (
|
||||
defaultTavilyTimeout = 15 * time.Second
|
||||
)
|
||||
|
||||
// TavilyProvider implements web search using Tavily Search API
|
||||
type TavilyProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
}
|
||||
|
||||
// NewTavilyProvider creates a new Tavily provider from parameters (no environment variables).
|
||||
func NewTavilyProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if params.APIKey == "" {
|
||||
return nil, fmt.Errorf("API key is required for Tavily provider")
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultTavilyTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &TavilyProvider{
|
||||
client: client,
|
||||
baseURL: defaultTavilySearchURL,
|
||||
apiKey: params.APIKey,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Name returns the provider name
|
||||
func (p *TavilyProvider) Name() string {
|
||||
return "tavily"
|
||||
}
|
||||
|
||||
// Search performs a web search using Tavily Search API
|
||||
func (p *TavilyProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
if len(query) == 0 {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Tavily] query=%q maxResults=%d url=%s", query, maxResults, p.baseURL)
|
||||
|
||||
reqBody := tavilySearchRequest{
|
||||
APIKey: p.apiKey,
|
||||
Query: query,
|
||||
MaxResults: maxResults,
|
||||
}
|
||||
|
||||
bodyBytes, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", p.baseURL, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
logger.Warnf(ctx, "[WebSearch][Tavily] API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("tavily API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
var respData tavilySearchResponse
|
||||
if err := json.Unmarshal(respBody, &respData); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, len(respData.Results))
|
||||
for _, item := range respData.Results {
|
||||
result := &types.WebSearchResult{
|
||||
Title: item.Title,
|
||||
URL: item.URL,
|
||||
Snippet: item.Content,
|
||||
Source: "tavily",
|
||||
}
|
||||
if includeDate && item.PublishedDate != "" {
|
||||
if t, err := time.Parse(time.RFC3339, item.PublishedDate); err == nil {
|
||||
result.PublishedAt = &t
|
||||
}
|
||||
}
|
||||
results = append(results, result)
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Tavily] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// tavilySearchRequest defines the request body for Tavily search API
|
||||
type tavilySearchRequest struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Query string `json:"query"`
|
||||
MaxResults int `json:"max_results"`
|
||||
}
|
||||
|
||||
// tavilySearchResponse defines the response structure for Tavily search API
|
||||
type tavilySearchResponse struct {
|
||||
Query string `json:"query"`
|
||||
Results []struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
Score float64 `json:"score"`
|
||||
PublishedDate string `json:"published_date,omitempty"`
|
||||
} `json:"results"`
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// emptyResultDiagnostics is implemented by providers that can explain an empty
|
||||
// search result set (used by the settings "test connection" flow).
|
||||
type emptyResultDiagnostics interface {
|
||||
EmptyResultDiagnostics() string
|
||||
}
|
||||
|
||||
// EmptyTestResultsError builds a provider-specific message when a connectivity
|
||||
// test search succeeds but returns no usable results.
|
||||
func EmptyTestResultsError(providerType string, provider any) error {
|
||||
detail := ""
|
||||
if dr, ok := provider.(emptyResultDiagnostics); ok {
|
||||
detail = dr.EmptyResultDiagnostics()
|
||||
}
|
||||
|
||||
switch types.WebSearchProviderType(providerType) {
|
||||
case types.WebSearchProviderTypeSearxng:
|
||||
if detail != "" {
|
||||
return fmt.Errorf("searxng returned 0 results: %s", detail)
|
||||
}
|
||||
return fmt.Errorf(
|
||||
"searxng returned 0 results; verify the instance URL, JSON format in settings.yml, and upstream search engine connectivity",
|
||||
)
|
||||
case types.WebSearchProviderTypeDuckDuckGo:
|
||||
return fmt.Errorf("duckduckgo returned 0 results; verify network connectivity and proxy settings")
|
||||
case types.WebSearchProviderTypeKeenable:
|
||||
return fmt.Errorf(
|
||||
"keenable returned 0 results; keyless requests are rate-limited, so verify network connectivity or set an API key to lift the cap",
|
||||
)
|
||||
case types.WebSearchProviderTypeExa:
|
||||
return fmt.Errorf(
|
||||
"exa returned 0 results; verify the API key, account quota, network connectivity, and proxy settings",
|
||||
)
|
||||
default:
|
||||
return fmt.Errorf("search returned 0 results, please verify your API key and configuration")
|
||||
}
|
||||
}
|
||||
|
||||
func formatUnresponsiveEngines(engines [][]string) string {
|
||||
if len(engines) == 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, len(engines))
|
||||
for _, e := range engines {
|
||||
if len(e) == 0 {
|
||||
continue
|
||||
}
|
||||
if len(e) == 1 {
|
||||
parts = append(parts, e[0])
|
||||
continue
|
||||
}
|
||||
parts = append(parts, e[0]+" ("+e[1]+")")
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return ""
|
||||
}
|
||||
return "unresponsive engines: " + strings.Join(parts, ", ")
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestEmptyTestResultsError_SearxngWithUnresponsiveEngines(t *testing.T) {
|
||||
provider := &SearxngProvider{
|
||||
lastUnresponsive: [][]string{
|
||||
{"google", "timeout"},
|
||||
{"bing", "blocked"},
|
||||
},
|
||||
}
|
||||
err := EmptyTestResultsError(string(types.WebSearchProviderTypeSearxng), provider)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
msg := err.Error()
|
||||
if !strings.Contains(msg, "searxng returned 0 results") {
|
||||
t.Fatalf("unexpected message: %q", msg)
|
||||
}
|
||||
if !strings.Contains(msg, "google (timeout)") || !strings.Contains(msg, "bing (blocked)") {
|
||||
t.Fatalf("expected unresponsive engine details, got: %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyTestResultsError_SearxngWithoutDiagnostics(t *testing.T) {
|
||||
provider := &SearxngProvider{}
|
||||
err := EmptyTestResultsError(string(types.WebSearchProviderTypeSearxng), provider)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
msg := err.Error()
|
||||
if !strings.Contains(msg, "JSON format is enabled in settings.yml") {
|
||||
t.Fatalf("unexpected message: %q", msg)
|
||||
}
|
||||
if strings.Contains(msg, "API key") {
|
||||
t.Fatalf("searxng message must not mention API key: %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyTestResultsError_DuckDuckGo(t *testing.T) {
|
||||
err := EmptyTestResultsError(string(types.WebSearchProviderTypeDuckDuckGo), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "duckduckgo returned 0 results") {
|
||||
t.Fatalf("unexpected message: %q", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyTestResultsError_Keenable(t *testing.T) {
|
||||
err := EmptyTestResultsError(string(types.WebSearchProviderTypeKeenable), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
msg := err.Error()
|
||||
if !strings.Contains(msg, "keenable returned 0 results") {
|
||||
t.Fatalf("unexpected message: %q", msg)
|
||||
}
|
||||
if !strings.Contains(msg, "rate-limited") {
|
||||
t.Fatalf("keenable message should hint at the keyless rate limit: %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyTestResultsError_Exa(t *testing.T) {
|
||||
err := EmptyTestResultsError(string(types.WebSearchProviderTypeExa), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
msg := err.Error()
|
||||
if !strings.Contains(msg, "exa returned 0 results") {
|
||||
t.Fatalf("unexpected message: %q", msg)
|
||||
}
|
||||
if !strings.Contains(msg, "API key") || !strings.Contains(msg, "quota") || !strings.Contains(msg, "proxy") {
|
||||
t.Fatalf("exa message should include actionable connection checks: %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyTestResultsError_APIKeyProviders(t *testing.T) {
|
||||
err := EmptyTestResultsError(string(types.WebSearchProviderTypeBing), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "API key") {
|
||||
t.Fatalf("bing message should mention API key: %q", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatUnresponsiveEngines(t *testing.T) {
|
||||
got := formatUnresponsiveEngines([][]string{
|
||||
{"google", "timeout"},
|
||||
{"wikipedia"},
|
||||
})
|
||||
want := "unresponsive engines: google (timeout), wikipedia"
|
||||
if got != want {
|
||||
t.Fatalf("formatUnresponsiveEngines() = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultZhipuSearchURL = "https://open.bigmodel.cn/api/paas/v4/web_search"
|
||||
defaultZhipuTimeout = 15 * time.Second
|
||||
defaultZhipuResults = 10
|
||||
maxZhipuResults = 50
|
||||
maxZhipuQueryRunes = 70
|
||||
maxZhipuResponseBytes = 2 << 20
|
||||
defaultZhipuSearchEngine = "search_std"
|
||||
defaultZhipuContentSize = "medium"
|
||||
)
|
||||
|
||||
var validZhipuSearchEngines = map[string]struct{}{
|
||||
"search_std": {},
|
||||
"search_pro": {},
|
||||
"search_pro_sogou": {},
|
||||
"search_pro_quark": {},
|
||||
}
|
||||
|
||||
var validZhipuContentSizes = map[string]struct{}{
|
||||
"medium": {},
|
||||
"high": {},
|
||||
}
|
||||
|
||||
// ZhipuProvider implements web search using Zhipu AI's standalone Web Search API.
|
||||
type ZhipuProvider struct {
|
||||
client *http.Client
|
||||
baseURL string
|
||||
apiKey string
|
||||
searchEngine string
|
||||
contentSize string
|
||||
}
|
||||
|
||||
// NewZhipuProvider creates a Zhipu AI web search provider from persisted parameters.
|
||||
func NewZhipuProvider(params types.WebSearchProviderParameters) (interfaces.WebSearchProvider, error) {
|
||||
if err := ValidateZhipuParameters(params); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := NewSearchHTTPClient(defaultZhipuTimeout, params.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
searchEngine, contentSize := zhipuOptions(params.ExtraConfig)
|
||||
return &ZhipuProvider{
|
||||
client: client,
|
||||
baseURL: defaultZhipuSearchURL,
|
||||
apiKey: strings.TrimSpace(params.APIKey),
|
||||
searchEngine: searchEngine,
|
||||
contentSize: contentSize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ValidateZhipuParameters validates credentials and provider-specific options.
|
||||
func ValidateZhipuParameters(params types.WebSearchProviderParameters) error {
|
||||
if strings.TrimSpace(params.APIKey) == "" {
|
||||
return fmt.Errorf("API key is required for Zhipu provider")
|
||||
}
|
||||
searchEngine, contentSize := zhipuOptions(params.ExtraConfig)
|
||||
if _, ok := validZhipuSearchEngines[searchEngine]; !ok {
|
||||
return fmt.Errorf("invalid Zhipu search engine: %s", searchEngine)
|
||||
}
|
||||
if _, ok := validZhipuContentSizes[contentSize]; !ok {
|
||||
return fmt.Errorf("invalid Zhipu content size: %s", contentSize)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func zhipuOptions(extraConfig map[string]string) (searchEngine, contentSize string) {
|
||||
searchEngine = defaultZhipuSearchEngine
|
||||
contentSize = defaultZhipuContentSize
|
||||
if value := strings.TrimSpace(extraConfig["search_engine"]); value != "" {
|
||||
searchEngine = value
|
||||
}
|
||||
if value := strings.TrimSpace(extraConfig["content_size"]); value != "" {
|
||||
contentSize = value
|
||||
}
|
||||
return searchEngine, contentSize
|
||||
}
|
||||
|
||||
// Name returns the provider name.
|
||||
func (p *ZhipuProvider) Name() string {
|
||||
return "zhipu"
|
||||
}
|
||||
|
||||
// Search performs a web search using Zhipu AI's standalone Web Search API.
|
||||
func (p *ZhipuProvider) Search(
|
||||
ctx context.Context,
|
||||
query string,
|
||||
maxResults int,
|
||||
includeDate bool,
|
||||
) ([]*types.WebSearchResult, error) {
|
||||
preparedQuery := normalizeZhipuQuery(query)
|
||||
if preparedQuery == "" {
|
||||
return nil, fmt.Errorf("query is empty")
|
||||
}
|
||||
if utf8.RuneCountInString(strings.TrimSpace(query)) > maxZhipuQueryRunes {
|
||||
logger.Infof(ctx, "[WebSearch][Zhipu] truncated query to %d characters", maxZhipuQueryRunes)
|
||||
}
|
||||
if maxResults <= 0 {
|
||||
maxResults = defaultZhipuResults
|
||||
}
|
||||
if maxResults > maxZhipuResults {
|
||||
maxResults = maxZhipuResults
|
||||
}
|
||||
|
||||
requestBody := zhipuSearchRequest{
|
||||
SearchQuery: preparedQuery,
|
||||
SearchEngine: p.searchEngine,
|
||||
SearchIntent: false,
|
||||
Count: maxResults,
|
||||
ContentSize: p.contentSize,
|
||||
}
|
||||
body, err := json.Marshal(requestBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal Zhipu request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create Zhipu request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
logger.Infof(ctx, "[WebSearch][Zhipu] query=%q maxResults=%d engine=%s", preparedQuery, maxResults, p.searchEngine)
|
||||
resp, err := p.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to execute Zhipu request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := readZhipuResponseBody(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, zhipuHTTPError(resp.StatusCode, respBody)
|
||||
}
|
||||
|
||||
var response zhipuSearchResponse
|
||||
if err := json.Unmarshal(respBody, &response); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal Zhipu response: %w", err)
|
||||
}
|
||||
if response.Error.Message != "" || response.Error.Code != "" {
|
||||
return nil, fmt.Errorf("Zhipu API error (%s): %s", response.Error.Code, response.Error.Message)
|
||||
}
|
||||
|
||||
results := make([]*types.WebSearchResult, 0, len(response.SearchResult))
|
||||
for _, item := range response.SearchResult {
|
||||
if strings.TrimSpace(item.Title) == "" && strings.TrimSpace(item.Link) == "" {
|
||||
continue
|
||||
}
|
||||
result := &types.WebSearchResult{
|
||||
Title: item.Title,
|
||||
URL: item.Link,
|
||||
Snippet: item.Content,
|
||||
Source: "zhipu",
|
||||
}
|
||||
if includeDate {
|
||||
if publishedAt, ok := parseZhipuDate(item.PublishDate); ok {
|
||||
result.PublishedAt = &publishedAt
|
||||
}
|
||||
}
|
||||
results = append(results, result)
|
||||
if len(results) >= maxResults {
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Infof(ctx, "[WebSearch][Zhipu] returned %d results", len(results))
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func normalizeZhipuQuery(query string) string {
|
||||
query = strings.TrimSpace(query)
|
||||
if utf8.RuneCountInString(query) <= maxZhipuQueryRunes {
|
||||
return query
|
||||
}
|
||||
runes := []rune(query)
|
||||
return string(runes[:maxZhipuQueryRunes])
|
||||
}
|
||||
|
||||
func parseZhipuDate(value string) (time.Time, bool) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return time.Time{}, false
|
||||
}
|
||||
for _, layout := range []string{
|
||||
time.RFC3339Nano,
|
||||
"2006-01-02 15:04:05",
|
||||
"2006-01-02 15:04",
|
||||
"2006-01-02",
|
||||
} {
|
||||
if parsed, err := time.Parse(layout, value); err == nil {
|
||||
return parsed, true
|
||||
}
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
func readZhipuResponseBody(reader io.Reader) ([]byte, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(reader, maxZhipuResponseBytes+1))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read Zhipu response: %w", err)
|
||||
}
|
||||
if len(body) > maxZhipuResponseBytes {
|
||||
return nil, fmt.Errorf("Zhipu response exceeds %d bytes", maxZhipuResponseBytes)
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func zhipuHTTPError(statusCode int, body []byte) error {
|
||||
var response zhipuSearchResponse
|
||||
if err := json.Unmarshal(body, &response); err == nil && (response.Error.Code != "" || response.Error.Message != "") {
|
||||
return fmt.Errorf("Zhipu API returned status %d (%s): %s", statusCode, response.Error.Code, response.Error.Message)
|
||||
}
|
||||
detail := strings.TrimSpace(string(body))
|
||||
if len(detail) > 4096 {
|
||||
detail = detail[:4096]
|
||||
}
|
||||
if detail == "" {
|
||||
return fmt.Errorf("Zhipu API returned status %d", statusCode)
|
||||
}
|
||||
return fmt.Errorf("Zhipu API returned status %d: %s", statusCode, detail)
|
||||
}
|
||||
|
||||
type zhipuSearchRequest struct {
|
||||
SearchQuery string `json:"search_query"`
|
||||
SearchEngine string `json:"search_engine"`
|
||||
SearchIntent bool `json:"search_intent"`
|
||||
Count int `json:"count"`
|
||||
ContentSize string `json:"content_size"`
|
||||
}
|
||||
|
||||
type zhipuSearchResponse struct {
|
||||
ID string `json:"id"`
|
||||
RequestID string `json:"request_id"`
|
||||
SearchResult []zhipuSearchResult `json:"search_result"`
|
||||
Error zhipuError `json:"error"`
|
||||
}
|
||||
|
||||
type zhipuSearchResult struct {
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
Link string `json:"link"`
|
||||
Media string `json:"media"`
|
||||
Icon string `json:"icon"`
|
||||
Refer string `json:"refer"`
|
||||
PublishDate string `json:"publish_date"`
|
||||
}
|
||||
|
||||
type zhipuError struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package web_search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestValidateZhipuParameters(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
params types.WebSearchProviderParameters
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "defaults", params: types.WebSearchProviderParameters{APIKey: "key"}},
|
||||
{
|
||||
name: "custom options",
|
||||
params: types.WebSearchProviderParameters{
|
||||
APIKey: "key",
|
||||
ExtraConfig: map[string]string{
|
||||
"search_engine": "search_pro_sogou",
|
||||
"content_size": "high",
|
||||
},
|
||||
},
|
||||
},
|
||||
{name: "missing key", params: types.WebSearchProviderParameters{}, wantErr: true},
|
||||
{
|
||||
name: "invalid engine",
|
||||
params: types.WebSearchProviderParameters{
|
||||
APIKey: "key",
|
||||
ExtraConfig: map[string]string{"search_engine": "unknown"},
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "invalid content size",
|
||||
params: types.WebSearchProviderParameters{
|
||||
APIKey: "key",
|
||||
ExtraConfig: map[string]string{"content_size": "large"},
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := ValidateZhipuParameters(tt.params)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Fatalf("ValidateZhipuParameters() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestZhipuProviderSearch(t *testing.T) {
|
||||
query := strings.Repeat("智", maxZhipuQueryRunes+1)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("method = %s, want POST", r.Method)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
|
||||
t.Errorf("Authorization = %q, want Bearer test-key", got)
|
||||
}
|
||||
var request zhipuSearchRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
||||
t.Errorf("decode request: %v", err)
|
||||
}
|
||||
if got := utf8.RuneCountInString(request.SearchQuery); got != maxZhipuQueryRunes {
|
||||
t.Errorf("query rune count = %d, want %d", got, maxZhipuQueryRunes)
|
||||
}
|
||||
if request.SearchEngine != "search_pro" {
|
||||
t.Errorf("search_engine = %q, want search_pro", request.SearchEngine)
|
||||
}
|
||||
if request.ContentSize != "high" {
|
||||
t.Errorf("content_size = %q, want high", request.ContentSize)
|
||||
}
|
||||
if request.SearchIntent {
|
||||
t.Error("search_intent = true, want false")
|
||||
}
|
||||
if request.Count != maxZhipuResults {
|
||||
t.Errorf("count = %d, want %d", request.Count, maxZhipuResults)
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "search-id",
|
||||
"request_id": "request-id",
|
||||
"search_result": []map[string]any{
|
||||
{
|
||||
"title": "Result 1",
|
||||
"link": "https://example.com/1",
|
||||
"content": "Summary 1",
|
||||
"publish_date": "2026-07-16",
|
||||
},
|
||||
{
|
||||
"title": "Result 2",
|
||||
"link": "https://example.com/2",
|
||||
"content": "Summary 2",
|
||||
},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
provider := &ZhipuProvider{
|
||||
client: srv.Client(),
|
||||
baseURL: srv.URL,
|
||||
apiKey: "test-key",
|
||||
searchEngine: "search_pro",
|
||||
contentSize: "high",
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
results, err := provider.Search(ctx, query, maxZhipuResults+1, true)
|
||||
if err != nil {
|
||||
t.Fatalf("Search() error = %v", err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("len(results) = %d, want 2", len(results))
|
||||
}
|
||||
if results[0].Snippet != "Summary 1" || results[0].Content != "" {
|
||||
t.Errorf("first result content mapping = %+v", results[0])
|
||||
}
|
||||
if results[0].Source != "zhipu" {
|
||||
t.Errorf("source = %q, want zhipu", results[0].Source)
|
||||
}
|
||||
if results[0].PublishedAt == nil || results[0].PublishedAt.Format("2006-01-02") != "2026-07-16" {
|
||||
t.Errorf("published_at = %v, want 2026-07-16", results[0].PublishedAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestZhipuProviderSearchDefaults(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var request zhipuSearchRequest
|
||||
_ = json.NewDecoder(r.Body).Decode(&request)
|
||||
if request.SearchEngine != defaultZhipuSearchEngine {
|
||||
t.Errorf("search_engine = %q, want %q", request.SearchEngine, defaultZhipuSearchEngine)
|
||||
}
|
||||
if request.ContentSize != defaultZhipuContentSize {
|
||||
t.Errorf("content_size = %q, want %q", request.ContentSize, defaultZhipuContentSize)
|
||||
}
|
||||
if request.Count != defaultZhipuResults {
|
||||
t.Errorf("count = %d, want %d", request.Count, defaultZhipuResults)
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"search_result":[]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
provider := &ZhipuProvider{
|
||||
client: srv.Client(),
|
||||
baseURL: srv.URL,
|
||||
apiKey: "test-key",
|
||||
searchEngine: defaultZhipuSearchEngine,
|
||||
contentSize: defaultZhipuContentSize,
|
||||
}
|
||||
if _, err := provider.Search(context.Background(), "test", 0, false); err != nil {
|
||||
t.Fatalf("Search() error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestZhipuProviderSearchError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = w.Write([]byte(`{"error":{"code":"1302","message":"rate limited"}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
provider := &ZhipuProvider{
|
||||
client: srv.Client(),
|
||||
baseURL: srv.URL,
|
||||
apiKey: "test-key",
|
||||
searchEngine: defaultZhipuSearchEngine,
|
||||
contentSize: defaultZhipuContentSize,
|
||||
}
|
||||
_, err := provider.Search(context.Background(), "test", 1, false)
|
||||
if err == nil || !strings.Contains(err.Error(), "rate limited") {
|
||||
t.Fatalf("Search() error = %v, want rate limited error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseZhipuDate(t *testing.T) {
|
||||
for _, value := range []string{
|
||||
"2026-07-16",
|
||||
"2026-07-16 12:30",
|
||||
"2026-07-16 12:30:45",
|
||||
"2026-07-16T12:30:45Z",
|
||||
} {
|
||||
if _, ok := parseZhipuDate(value); !ok {
|
||||
t.Errorf("parseZhipuDate(%q) failed", value)
|
||||
}
|
||||
}
|
||||
if _, ok := parseZhipuDate("not-a-date"); ok {
|
||||
t.Error("parseZhipuDate(not-a-date) unexpectedly succeeded")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user