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

This commit is contained in:
yawan
2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -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))
}
}
+241
View File
@@ -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)
}
}
}
+914
View File
@@ -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
+363
View File
@@ -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)
}
}
+159
View File
@@ -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
// `![alt](images/image-N.ext)` 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, &section, 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 := "![Shipping chart](images/image-1.png)"
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: "![page](images/page_1.jpg)",
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("![%s](%s)", 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("![%s](%s)", 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
}
+21
View File
@@ -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{
"![](images/foo.png)",
`<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 ![](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("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 := "![](DATA:image/png;BASE64," + b64 + ")"
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 := `![](data:image/x-emf;base64,` + b64 + `)`
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, "![image](local://test/") {
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`,
`![](data:image/png;base64,` + b64 + `)`,
`![C:\path\file]name.png](data:image/png;base64,` + b64 + `)`,
`![](data:image/x-emf;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: `![图片](images/test.png "图片")`,
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, `![图片](local://test/`) || !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: `![图片](images/test.png '图片')`,
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, `![图片](local://test/`) || !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: `![](images/第 1 页.png)`,
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, `![](local://test/`) {
t.Fatalf("spaced filename path was not replaced: %s", out)
}
}
func TestResolveAndStoreMarkdownImageTitleContainingRightParen(t *testing.T) {
png := createTestPNG(200, 150)
result := &types.ReadResult{
MarkdownContent: `![图片](images/test.png "阶段 1) 图片")`,
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, `![图片](local://test/`) || !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: "![图片](images/test.png\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(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, "![图片](local://test/") || !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: `![图片](<images/第 1 页 (测试).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, `![图片](<local://test/`) || !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{
`![图片](images/test.png "第一处")`,
`![图片](images/test.png "第二处")`,
}, "\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: `![图片](images/missing.png "图片")`,
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, `![图片](local://test/`) {
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![图片](local://test/") {
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 ![a](images/a.png) after`,
targets: []string{`images/a.png`},
},
{
name: "double quoted title",
input: `![a](images/a.png "图")`,
targets: []string{`images/a.png "图"`},
},
{
name: "single quoted title",
input: `![a](images/a.png '图')`,
targets: []string{`images/a.png '图'`},
},
{
name: "parenthesized title",
input: `![a](images/a.png (图))`,
targets: []string{`images/a.png (图)`},
},
{
name: "title containing right paren",
input: `![a](images/a.png "阶段 1) 结果")`,
targets: []string{`images/a.png "阶段 1) 结果"`},
},
{
name: "title containing both parens",
input: `![a](images/a.png '阶段 (1) 结果')`,
targets: []string{`images/a.png '阶段 (1) 结果'`},
},
{
name: "escaped quote in title",
input: `![a](images/a.png "阶段 \"1\"")`,
targets: []string{`images/a.png "阶段 \"1\""`},
},
{
name: "multiline title",
input: "![a](images/a.png\n \"多行 title\")",
targets: []string{"images/a.png\n \"多行 title\""},
},
{
name: "spaced path",
input: `![a](images/第 1 页.png)`,
targets: []string{`images/第 1 页.png`},
},
{
name: "path containing balanced parens",
input: `![a](images/a_(1).png "title")`,
targets: []string{`images/a_(1).png "title"`},
},
{
name: "angle destination",
input: `![a](<images/a b.png> "title")`,
targets: []string{`<images/a b.png> "title"`},
},
{
name: "malformed image is skipped",
input: `![a](images/a.png "unterminated title) after ![b](images/b.png)`,
targets: []string{`images/b.png`},
},
{
name: "escaped image marker is skipped",
input: `\![a](images/a.png) ![b](images/b.png)`,
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 `![alt](url)` 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, "![$1]($2)")
}
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",
"",
"![](images/cover.jpg)",
"",
`<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, "![](images/cover.jpg)") {
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{
"![](images/cover.jpg)",
`<img src="./images/profile.jpg"/>`,
`![](plain.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 := "![](images/第 1 页.jpg)"
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, "![") {
t.Errorf("tag was converted to markdown syntax: %s", updated)
}
}
func TestResolveRemoteImages_MarkdownAndHTMLBothResolved(t *testing.T) {
allowLocalhost(t)
ts, _ := remoteImageServer(t)
markdown := fmt.Sprintf("![md](%s/md.png)\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, "![md%d](%s/md%d.png)\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("![x](%s/md-up.png)", 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&amp;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&amp;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, "&amp;") {
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![photo](%s/image.png)\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 := "![evil](http://127.0.0.1:8080/secret.png)\n\n![also-evil](http://169.254.169.254/metadata)"
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("![bad](%s/page.html)", 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 := "![already](local://images/abc.png)\n![also](minio://bucket/key.jpg)"
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("![img1](%s/a.png)\n\ntext\n\n![img2](%s/b.png)\n\n![img3](%s/c.png)",
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("![missing](%s/nope.png)", 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 },
}
}
+265
View File
@@ -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"`
}
+199
View File
@@ -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)
}
+181
View File
@@ -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)
}
}
+271
View File
@@ -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")
}
}