mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 09:53:29 +08:00
Fix(go): align ingestion pipeline with Python (parser/media dispatch + PDF coordinate chain + Chunker) (#17349)
## Summary Aligns the Go ingestion pipeline with the Python implementation, closing several behavioral gaps found during the Python→Go migration (tracked in `docs/migration_python_go_diff.md`). Covers parser/media dispatch alignment, the PDF coordinate-chain (preview images, outline→title, chunk coordinate finalization), and the Chunker Token/QA batches below. Commits are grouped as follows. ### 1. Fix parser params (c524f450e) Fixes parser/media wiring and several dispatch gaps: - **docx/pdf vision dispatch**: correct parameter handling and VLM invocation. - **markdown vision (diff 2.5)**: also enhance items whose `doc_type_kwd` is `table`, not only `image` (parser/utils.py:181). - **media audio (diff 2.11)**: when `output_format` is `json`, carry the ASR transcription as a JSON item instead of only the `Text` field (the Invoke switch had no `json` branch and dropped it). - **email (diff 2.2)**: default `output_format` is `json` (parser.py:212), not `text`. - **tokenizer**: handle empty/whitespace-only names; trim before embedding. - **extractor**: tag-matching parameter wiring. - **split**: keyword-split regex now covers CJK/English separators. - **parser.go**: parser-param plumbing. ### 2. fix parser gap (373537da1) Image dispatch now mirrors `rag/app/picture.py:chunk()`: - Always OCR the image (PaddleOCR or local ONNX). - When OCR text is short, also call VLM (`describe`) and combine `OCR + VLM` text. - Emits a **structured JSON item** carrying the image data-URI and `doc_type_kwd:"image"`, instead of a bare `Text` string. This fixes the payload being rejected downstream by OneChunker/TokenChunker (JSON=nil). ### 3. PDF coordinate-chain fixes (55367a820,727f8167c) Closes three items from the migration tracker in the chunker/tokenizer/task layer: - **(Chunker-1.3) `restore_pdf_text_previews`** — `needsCrop` now also returns true for `text` chunks that carry PDF positions (`pdfcrop_cgo.go`), so text blocks get a rendered preview image uploaded to storage via `imageUploadDecorator`/`ChunkImageUploader`, matching Python `restore_pdf_text_previews` + `image2id`. - **(Chunker-1.5) PDF outline → title levels** — `title.go` adds `outlineSimilarity` (rune-bigram Jaccard, mirroring `common.py:_outline_similarity`), `resolveOutlineLevels` (matches text lines to outline entries at similarity > 0.8, with a sparse guard `len(outline)/len(records) <= 0.03`), and `outlineFromInputs` (reads `file.outline`). Wired into `newLevelContext` in both `group.go` and `hierarchy.go`; falls back to the title-shape heuristic when no outline is present. - **(Tokenizer-(T)1) `finalize_pdf_chunk`** — the coordinate → `position_int`/`page_num_int`/`top_int` conversion is owned by the task layer (`processChunkPositions`→`AddPositions`), which runs *after* the tokenizer and consumes the tokenizer-owned fields. The tokenizer only preserves the raw `positions`/`_pdf_positions` (no duplicate conversion), pinned by `TestChunkDocsToMaps_PreservesPDFPositions`. ### 4. Integration test made environment-free (`internal/ingestion/task/pipeline_real_integration_test.go`) - Removed the `//go:build integration` tag so the contract tests run under the default `build.sh --test` (which does not pass `-tags integration`). - External dependencies replaced with in-memory substitutes so no MySQL/MinIO/ES is required: - MySQL → on-disk sqlite (`glebarez/sqlite`) with the needed tables auto-migrated. - MinIO → `storage.NewMemoryStorage()`. - Elasticsearch → chunks captured via `WithInsertFunc` instead of `engine.InsertChunks`/`Search`. - `requireTokenizerPool` still skips gracefully when the native tokenizer pool is unavailable; `WithLogCreateFunc(noop)` avoids depending on the operation-log table. - Added `taskChunkFieldEqualsStr` to tolerate `kb_id` being a `[]string`/`[]any` in the raw chunk payload (the search engine flattens it to a string on read). ### 5. TokenChunker alignment — Batch 1 (`internal/ingestion/component/chunker/token.go`) Closes four Chunker items from the migration tracker: - **(Chunker-2.1) sentence delimiter** — the boundary regex now also breaks on ASCII `!`/`?`. Extracted to a package-level `var sentenceDelimiter` and used in `mergeByTokenSize`, matching Python's full delimiter set. - **(Chunker-2.2) overlap tag leakage** — when a new chunk starts, its overlap prefix is taken from the previous chunk *after* `removeTag`, in both the text path (`mergeByTokenSize`) and the JSON path (`mergeByTokenSizeFromJSON`). Parser tags (`@@…##`) no longer leak into the overlap region (mirrors `nlp/__init__.py:1181`). - **(Chunker-2.11) empty-text merge** — merging a non-empty chunk into an empty previous chunk now assigns the text directly instead of being skipped (`mergeByTokenSizeFromJSON`), mirroring `token_chunker.py:236-239`. - **(Chunker-2.4) overlap token counting** — `takeFromEnd`/`takeFromStart` now count tokens exactly via `tokenizeStr` instead of the 4-bytes/token heuristic, fixing over-counting for CJK text. ### 6. QA Chunker alignment — Batch 2 (`internal/ingestion/component/chunker/qa.go` + `schema`) Closes three Chunker items from the migration tracker: - **(Chunker-2.13) default language** — an empty `lang` now defaults to Chinese prefixes (`问题:`/`回答:`) instead of English, matching `qa.py:299`. - **(Chunker-2.12) `rmQAPrefix` regex** — the separator is changed to `[\t:: ]+` (one-or-more), matching `qa.py:241`, so multiple separators (e.g. `Q:: answer`) are fully stripped. - **(Chunker-1.8 QA) missing chunk fields** — QA chunks now preserve: - `top_int` — the source row/record index, threaded through the tab/csv/markdown extractors (mirrors `qa.py` `beAdoc(..., row_num=i)`); - `image` + `doc_type_kwd:"image"`; - `_pdf_positions` / `positions` carried from the upstream JSON item. `schema.ChunkDoc` gains a `TopInt []int` field (serialized as `top_int`, registered in `UnmarshalJSON`). Note: the Tag/Table/Presentation/One chunker field gaps under 1.8 remain pending. ## Test plan - Added/updated unit tests: `pdfcrop_cgo_test.go` (`TestNeedsCrop`, `TestRestorePDFTextPreview`), `title_test.go` (`TestResolveOutlineLevels`, `TestResolveOutlineLevels_SparseGuard`, `TestNewLevelContext_OutlineBranch`, `TestOutlineFromInputs`), `tokenizer_unit_test.go` (`TestChunkDocsToMaps_PreservesPDFPositions`), `token_pdfpos_test.go`. - **Batch 1** — `token_batch1_test.go`: `TestSentenceDelimiterMatchesBangAndQuestion`, `TestMergeByTokenSizeFromJSON_OverlapStripsTags`, `TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk`, `TestTakeFromEndRespectsTokenCount`, `TestTakeFromStartRespectsTokenCount`. - **Batch 2** — `qa_batch2_test.go`: `TestQAChunker_DefaultLangIsChinese`, `TestRmQAPrefixStripsMultipleSeparators`, `TestQAChunker_SetsTopInt`, `TestQAChunker_CarriesImageAndPositions`. Existing `qa_test.go` expectations were updated to the corrected language default / separator behavior. - `pipeline_real_integration_test.go` (`TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline`, `TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks`, `TestRunPipeline_RealPipelineOutput_ProducesIndexFields`) now runs without any external service. - `bash build.sh --test ./internal/ingestion/...` passes. - No files deleted.
This commit is contained in:
@@ -36,6 +36,8 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/zap"
|
||||
@@ -115,7 +117,7 @@ func invokeGroup(_ context.Context, inputs map[string]any, p *titleChunkerParam)
|
||||
if len(records) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
ctx := newLevelContext(records, p)
|
||||
ctx := newLevelContext(records, outlineFromInputs(inputs), p)
|
||||
levels := ctx.Levels()
|
||||
// Count heading level distribution for debugging.
|
||||
headingCounts := make(map[int]int)
|
||||
@@ -256,7 +258,10 @@ func buildChunksFromRecordGroups(groups [][]lineRecord, p *titleChunkerParam, pl
|
||||
if len(g) == 0 {
|
||||
continue
|
||||
}
|
||||
chunk := map[string]any{"text": joinGroupText(g)}
|
||||
// Strip parser-emitted position tags (`@@...##`) from the
|
||||
// joined text (diff 1.6 / 2.8). Mirrors common.py:255
|
||||
// `RAGFlowPdfParser.remove_tag("".join(...))`.
|
||||
chunk := map[string]any{"text": removeTag(joinGroupText(g))}
|
||||
if !plain {
|
||||
first := g[0]
|
||||
if first.docType != "" {
|
||||
@@ -266,6 +271,24 @@ func buildChunksFromRecordGroups(groups [][]lineRecord, p *titleChunkerParam, pl
|
||||
chunk["img_id"] = *first.imgID
|
||||
}
|
||||
}
|
||||
// Merge PDF coordinate matrices across the merged records
|
||||
// instead of keeping only the leading record's (diff 1.6).
|
||||
// Mirrors pdf_chunk_metadata.py:127 merge_pdf_positions.
|
||||
var pdfSrc, posSrc []json.RawMessage
|
||||
for _, r := range g {
|
||||
if len(r.pdfPositions) > 0 {
|
||||
pdfSrc = append(pdfSrc, r.pdfPositions)
|
||||
}
|
||||
if len(r.positions) > 0 {
|
||||
posSrc = append(posSrc, r.positions)
|
||||
}
|
||||
}
|
||||
if m := mergePositionMatrix(pdfSrc...); m != nil {
|
||||
chunk["_pdf_positions"] = m
|
||||
}
|
||||
if m := mergePositionMatrix(posSrc...); m != nil {
|
||||
chunk["positions"] = m
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
if p.RootChunkAsHeading && len(chunks) > 1 {
|
||||
@@ -278,6 +301,61 @@ func buildChunksFromRecordGroups(groups [][]lineRecord, p *titleChunkerParam, pl
|
||||
return chunks
|
||||
}
|
||||
|
||||
// posTagRemove matches parser-emitted position tags of the form
|
||||
// `@@<page>\t<left>\t<right>\t<top>\t<bottom>##`. Mirrors Python
|
||||
// pdf_parser.py:1934 `re.sub(r"@@[\t0-9.-]+?##", "", txt)`.
|
||||
var posTagRemove = regexp.MustCompile(`@@[\t0-9.-]+?##`)
|
||||
|
||||
// removeTag strips parser-emitted position tags from a chunk's text.
|
||||
// Mirrors deepdoc RAGFlowPdfParser.remove_tag.
|
||||
func removeTag(text string) string {
|
||||
return posTagRemove.ReplaceAllString(text, "")
|
||||
}
|
||||
|
||||
// mergePositionMatrix aggregates multiple PDF coordinate matrices into a
|
||||
// single de-duplicated, sorted matrix. Mirrors Python
|
||||
// pdf_chunk_metadata.py:127 merge_pdf_positions: rows are 5-tuples
|
||||
// [page,left,right,top,bottom]; duplicates (by the first five columns)
|
||||
// are dropped and rows are sorted by (page, top, left). Returns nil when
|
||||
// no source carries any usable coordinates.
|
||||
func mergePositionMatrix(sources ...json.RawMessage) [][]float64 {
|
||||
var out [][]float64
|
||||
seen := make(map[string]bool)
|
||||
for _, src := range sources {
|
||||
if len(src) == 0 {
|
||||
continue
|
||||
}
|
||||
var mat [][]float64
|
||||
if err := json.Unmarshal(src, &mat); err != nil {
|
||||
continue
|
||||
}
|
||||
for _, row := range mat {
|
||||
if len(row) < 5 {
|
||||
continue
|
||||
}
|
||||
key := fmt.Sprintf("%v|%v|%v|%v|%v", row[0], row[1], row[2], row[3], row[4])
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
out = append(out, row)
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i][0] != out[j][0] {
|
||||
return out[i][0] < out[j][0]
|
||||
}
|
||||
if out[i][3] != out[j][3] {
|
||||
return out[i][3] < out[j][3]
|
||||
}
|
||||
return out[i][1] < out[j][1]
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// extractLineRecords reads the chunker inputs in the same order the
|
||||
// python BaseTitleChunker.extract_line_records uses:
|
||||
//
|
||||
@@ -333,12 +411,14 @@ func recordsFromStructured(items []schema.ChunkDoc) []lineRecord {
|
||||
meta[k] = json.RawMessage(v)
|
||||
}
|
||||
out = append(out, lineRecord{
|
||||
text: text,
|
||||
docType: dt,
|
||||
imgID: imgID,
|
||||
layout: it.Layout,
|
||||
ckType: it.CKType,
|
||||
parentMeta: meta,
|
||||
text: text,
|
||||
docType: dt,
|
||||
imgID: imgID,
|
||||
layout: it.Layout,
|
||||
ckType: it.CKType,
|
||||
pdfPositions: it.PDFPositions,
|
||||
positions: it.Positions,
|
||||
parentMeta: meta,
|
||||
})
|
||||
}
|
||||
return out
|
||||
|
||||
@@ -18,6 +18,7 @@ package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
@@ -230,6 +231,59 @@ func TestGroupChunker_StructuredMetadata(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestGroupChunker_MergesPDFPositionsAndRemovesTags is the TDD test for
|
||||
// migration diffs Chunker-1.6 / 2.8: when the group chunker merges
|
||||
// multiple adjacent text records into one chunk it must (a) strip the
|
||||
// parser-emitted `@@...##` position tags from the joined text, and
|
||||
// (b) MERGE (not drop) the `positions` coordinate matrices across the
|
||||
// merged records — mirroring common.py:255 remove_tag + merge.
|
||||
func TestGroupChunker_MergesPDFPositionsAndRemovesTags(t *testing.T) {
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{{`^# `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewGroupTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "# Heading", "doc_type_kwd": "text"},
|
||||
{"text": "body one @@1\t10.0\t20.0\t30.0\t40.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{1, 10, 20, 30, 40}}},
|
||||
{"text": "body two @@2\t15.0\t25.0\t35.0\t45.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{2, 15, 25, 35, 45}}},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "doc",
|
||||
"output_format": "chunks",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, _ := out["chunks"].([]map[string]any)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("no chunks emitted")
|
||||
}
|
||||
for _, ck := range chunks {
|
||||
text, _ := ck["text"].(string)
|
||||
// Only the merged body group carries both bodies.
|
||||
if !strings.Contains(text, "body one") || !strings.Contains(text, "body two") {
|
||||
continue
|
||||
}
|
||||
// (a) parser tags must be stripped from the text.
|
||||
if strings.Contains(text, "@@") {
|
||||
t.Errorf("parser position tags leaked into chunk text: %q", text)
|
||||
}
|
||||
// (b) positions must be merged across both records.
|
||||
pos, ok := ck["positions"].([][]float64)
|
||||
if !ok {
|
||||
t.Fatalf("positions missing or wrong type %T on merged group chunk", ck["positions"])
|
||||
}
|
||||
if len(pos) != 2 {
|
||||
t.Errorf("merged positions = %d groups, want 2 (both records)", len(pos))
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatal("merged body group chunk not found in output")
|
||||
}
|
||||
|
||||
func TestGroupTitleChunker_InvokeDeterministic(t *testing.T) {
|
||||
c, err := NewGroupTitleChunker(map[string]any{
|
||||
"levels": [][]string{
|
||||
|
||||
@@ -148,7 +148,7 @@ func invokeHierarchy(_ context.Context, inputs map[string]any, p *titleChunkerPa
|
||||
if len(records) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
ctx := newLevelContext(records, p)
|
||||
ctx := newLevelContext(records, outlineFromInputs(inputs), p)
|
||||
levels := ctx.Levels()
|
||||
// Count heading level distribution for debugging.
|
||||
headingCounts := make(map[int]int)
|
||||
|
||||
@@ -53,7 +53,9 @@ func newPDFEngineFromUpstream(ctx context.Context, up schema.ChunkerFromUpstream
|
||||
return deepdocpdf.NewEngine(data)
|
||||
}
|
||||
|
||||
// cropImageChunks crops image/table chunks in place. Each spanned page is
|
||||
// cropImageChunks crops image/table chunks and renders text previews (for
|
||||
// text chunks that carry PDF positions, mirroring Python
|
||||
// restore_pdf_text_previews). Each spanned page is
|
||||
// rendered at most once. Chunks arrive in document order, so we keep only a
|
||||
// sliding window of page images: once we advance past a chunk whose minimum
|
||||
// page is P, no later chunk references a page < P, and we evict those entries
|
||||
@@ -139,9 +141,14 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [
|
||||
return out
|
||||
}
|
||||
|
||||
// needsCrop reports whether a chunk should be cropped to a page-region
|
||||
// preview from its PDF positions. Image/table chunks get their media region
|
||||
// cropped; text chunks with positions get a rendered preview of the text
|
||||
// region (Python restore_pdf_text_previews). A pre-existing Image is never
|
||||
// re-cropped — cropImageChunks honors that separately.
|
||||
func needsCrop(ck schema.ChunkDoc) bool {
|
||||
switch ck.CKType {
|
||||
case "image", "table":
|
||||
case "image", "table", "text":
|
||||
return len(ck.PDFPositions) > 0 || len(ck.Positions) > 0
|
||||
default:
|
||||
return false
|
||||
|
||||
@@ -71,7 +71,7 @@ func TestNeedsCrop(t *testing.T) {
|
||||
}{
|
||||
{"image with positions", schema.ChunkDoc{CKType: "image", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, true},
|
||||
{"table with positions", schema.ChunkDoc{CKType: "table", Positions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, true},
|
||||
{"text with positions", schema.ChunkDoc{CKType: "text", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, false},
|
||||
{"text with positions", schema.ChunkDoc{CKType: "text", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, true},
|
||||
{"image without positions", schema.ChunkDoc{CKType: "image"}, false},
|
||||
{"unknown type", schema.ChunkDoc{CKType: "equation", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, false},
|
||||
}
|
||||
@@ -82,7 +82,7 @@ func TestNeedsCrop(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCropImageChunks_CropsImageAndTable(t *testing.T) {
|
||||
func TestCropImageChunks_CropsImageTableAndText(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
// 1-based JSON position (1) must be rendered as 0-based page 0.
|
||||
eng := assertZeroPageEngine{}
|
||||
@@ -91,7 +91,7 @@ func TestCropImageChunks_CropsImageAndTable(t *testing.T) {
|
||||
chunks := []schema.ChunkDoc{
|
||||
{CKType: "image", PDFPositions: pos},
|
||||
{CKType: "table", PDFPositions: pos},
|
||||
{CKType: "text", PDFPositions: pos}, // skipped (not image/table)
|
||||
{CKType: "text", PDFPositions: pos}, // restored preview (Chunker-1.3)
|
||||
{CKType: "image", Image: "data:image/png;base64,preexisting"}, // preserved
|
||||
}
|
||||
out := cropImageChunks(ctx, eng, chunks)
|
||||
@@ -101,8 +101,10 @@ func TestCropImageChunks_CropsImageAndTable(t *testing.T) {
|
||||
for i, ck := range out {
|
||||
switch ck.CKType {
|
||||
case "text":
|
||||
if ck.Image != "" {
|
||||
t.Errorf("chunk %d (text): image should stay empty, got %q", i, ck.Image)
|
||||
// Chunker-1.3: text chunks with PDF positions get a rendered
|
||||
// preview, mirroring Python restore_pdf_text_previews.
|
||||
if !strings.HasPrefix(ck.Image, "data:image/png;base64,") {
|
||||
t.Errorf("chunk %d (text): image = %q, want data:image/png;base64, prefix (preview restored)", i, ck.Image)
|
||||
}
|
||||
case "image":
|
||||
if ck.Image == "data:image/png;base64,preexisting" {
|
||||
@@ -119,6 +121,30 @@ func TestCropImageChunks_CropsImageAndTable(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestRestorePDFTextPreview covers Chunker-1.3 directly: a text chunk that
|
||||
// carries PDF positions must receive a rendered preview image, while a text
|
||||
// chunk without positions must be left untouched (no spurious preview). The
|
||||
// img_id upload is owned by imageUploadDecorator (image_upload.go) and is
|
||||
// not asserted here.
|
||||
func TestRestorePDFTextPreview(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
pos := jsonPositions(t, []float64{1, 10, 100, 10, 100})
|
||||
|
||||
withPos := schema.ChunkDoc{CKType: "text", PDFPositions: pos}
|
||||
withoutPos := schema.ChunkDoc{CKType: "text", Text: "plain text, no coordinates"}
|
||||
|
||||
out := cropImageChunks(ctx, mockCropEngine{}, []schema.ChunkDoc{withPos, withoutPos})
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("len(out) = %d, want 2", len(out))
|
||||
}
|
||||
if !strings.HasPrefix(out[0].Image, "data:image/png;base64,") {
|
||||
t.Errorf("chunk with positions: image = %q, want data:image/png;base64, prefix", out[0].Image)
|
||||
}
|
||||
if out[1].Image != "" {
|
||||
t.Errorf("chunk without positions: image = %q, want empty", out[1].Image)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCropImageChunks_NilEnginePassesThrough(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
pos := jsonPositions(t, []float64{1, 10, 100, 10, 100})
|
||||
|
||||
@@ -29,6 +29,7 @@ package chunker
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"html"
|
||||
"regexp"
|
||||
@@ -98,7 +99,9 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m
|
||||
}
|
||||
|
||||
qPrefix, aPrefix := "问题:", "回答:"
|
||||
eng := strings.EqualFold(c.param.Lang, "english") || c.param.Lang == ""
|
||||
// Python qa.py defaults to Chinese when no language is supplied; only
|
||||
// an explicit "english" switches to English prefixes (diff Chunker-2.13).
|
||||
eng := strings.EqualFold(c.param.Lang, "english")
|
||||
if eng {
|
||||
qPrefix, aPrefix = "Question: ", "Answer: "
|
||||
}
|
||||
@@ -133,6 +136,21 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m
|
||||
ContentLtks: contentLTKS,
|
||||
ContentSmLtks: contentSMLTKS,
|
||||
}
|
||||
// Restore metadata lost before diff Chunker-1.8: top_int (row
|
||||
// index), image id + coordinates carried from the source item.
|
||||
if pair.RowNum >= 0 {
|
||||
chunk.TopInt = []int{pair.RowNum}
|
||||
}
|
||||
if pair.Image != "" {
|
||||
chunk.Image = pair.Image
|
||||
chunk.DocType = "image"
|
||||
}
|
||||
if len(pair.PDFPositions) > 0 {
|
||||
chunk.PDFPositions = pair.PDFPositions
|
||||
}
|
||||
if len(pair.Positions) > 0 {
|
||||
chunk.Positions = pair.Positions
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
|
||||
@@ -148,9 +166,20 @@ func renderMarkdown(s string) string {
|
||||
type qaPair struct {
|
||||
Question string
|
||||
Answer string
|
||||
// RowNum is the 0-based source line/record index, mapped to Python's
|
||||
// top_int (qa.py beAdoc(..., row_num=i)). -1 means unset.
|
||||
RowNum int
|
||||
// Image and positions are carried from the upstream item so the QA
|
||||
// chunk preserves metadata that Python sets via beAdocPdf/beAdocDocx
|
||||
// (diff Chunker-1.8).
|
||||
Image string
|
||||
PDFPositions json.RawMessage
|
||||
Positions json.RawMessage
|
||||
}
|
||||
|
||||
var rmQAPrefixRe = regexp.MustCompile(`(?i)^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[ \t]*(?:[::]|\t)[ \t]*`)
|
||||
// rmQAPrefixRe mirrors Python qa.py:241 `[\t:: ]+` — one-or-more separator
|
||||
// chars, so "Q:: answer" is fully stripped (diff Chunker-2.12).
|
||||
var rmQAPrefixRe = regexp.MustCompile(`(?i)^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[\t:: ]+`)
|
||||
|
||||
func rmQAPrefix(txt string) string {
|
||||
return strings.TrimSpace(rmQAPrefixRe.ReplaceAllString(txt, ""))
|
||||
@@ -209,18 +238,19 @@ func extractQAMarkdown(md string) []qaPair {
|
||||
var questionStack []string
|
||||
var levelStack []int
|
||||
var answer []string
|
||||
curRow := -1
|
||||
codeBlock := false
|
||||
|
||||
flushAnswer := func() {
|
||||
joined := strings.TrimSpace(strings.Join(answer, "\n"))
|
||||
if joined != "" && len(questionStack) > 0 {
|
||||
sumQ := strings.Join(questionStack, "\n")
|
||||
pairs = append(pairs, qaPair{Question: sumQ, Answer: joined})
|
||||
pairs = append(pairs, qaPair{Question: sumQ, Answer: joined, RowNum: curRow})
|
||||
}
|
||||
answer = nil
|
||||
}
|
||||
|
||||
for _, line := range lines {
|
||||
for i, line := range lines {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "```") {
|
||||
codeBlock = !codeBlock
|
||||
@@ -239,6 +269,7 @@ func extractQAMarkdown(md string) []qaPair {
|
||||
|
||||
flushAnswer()
|
||||
question := strings.TrimSpace(line[level:])
|
||||
curRow = i
|
||||
|
||||
for len(levelStack) > 0 && level <= levelStack[len(levelStack)-1] {
|
||||
questionStack = questionStack[:len(questionStack)-1]
|
||||
@@ -273,8 +304,9 @@ func extractQAText(text string) []qaPair {
|
||||
func extractQATextTab(lines []string) []qaPair {
|
||||
var pairs []qaPair
|
||||
var question, answer string
|
||||
var row int
|
||||
|
||||
for _, line := range lines {
|
||||
for i, line := range lines {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
@@ -286,13 +318,14 @@ func extractQATextTab(lines []string) []qaPair {
|
||||
continue
|
||||
}
|
||||
if question != "" && answer != "" {
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row})
|
||||
}
|
||||
question = parts[0]
|
||||
answer = parts[1]
|
||||
row = i
|
||||
}
|
||||
if question != "" {
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row})
|
||||
}
|
||||
return pairs
|
||||
}
|
||||
@@ -321,13 +354,16 @@ func extractQATextCSV(text string, lines []string) []qaPair {
|
||||
|
||||
var pairs []qaPair
|
||||
var question, answer string
|
||||
var row int
|
||||
prevLine := 0
|
||||
recIdx := -1
|
||||
|
||||
for {
|
||||
record, err := r.Read()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
recIdx++
|
||||
|
||||
// Map InputOffset back to the physical lines consumed.
|
||||
endOff := int(r.InputOffset())
|
||||
@@ -346,13 +382,14 @@ func extractQATextCSV(text string, lines []string) []qaPair {
|
||||
continue
|
||||
}
|
||||
if question != "" && answer != "" {
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row})
|
||||
}
|
||||
question = record[0]
|
||||
answer = record[1]
|
||||
row = recIdx
|
||||
}
|
||||
if question != "" {
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row})
|
||||
}
|
||||
return pairs
|
||||
}
|
||||
@@ -385,7 +422,14 @@ func extractQAJSON(items []schema.ChunkDoc) []qaPair {
|
||||
continue
|
||||
}
|
||||
tmp := extractQAText(txt)
|
||||
pairs = append(pairs, tmp...)
|
||||
// Preserve the source item's image id and coordinates on each
|
||||
// extracted pair (diff Chunker-1.8).
|
||||
for _, p := range tmp {
|
||||
p.Image = item.Image
|
||||
p.PDFPositions = item.PDFPositions
|
||||
p.Positions = item.Positions
|
||||
pairs = append(pairs, p)
|
||||
}
|
||||
}
|
||||
return pairs
|
||||
}
|
||||
|
||||
147
internal/ingestion/component/chunker/qa_batch2_test.go
Normal file
147
internal/ingestion/component/chunker/qa_batch2_test.go
Normal file
@@ -0,0 +1,147 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// qaInvoke is a small helper that runs the QA chunker on a upstream-style
|
||||
// input map and returns the produced chunks as generic maps.
|
||||
func qaInvoke(t *testing.T, inputs map[string]any) []map[string]any {
|
||||
t.Helper()
|
||||
c, err := NewQAChunker(map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("NewQAChunker: %v", err)
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("Invoke did not return []map chunks: %#v", out["chunks"])
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
// TestQAChunker_DefaultLangIsChinese exercises migration diff Chunker-2.13:
|
||||
// when no language is supplied, Python defaults to Chinese prefixes
|
||||
// ("问题:"/"回答:"); the legacy Go code defaulted to English.
|
||||
func TestQAChunker_DefaultLangIsChinese(t *testing.T) {
|
||||
inputs := map[string]any{
|
||||
"name": "test.txt",
|
||||
"output_format": "text",
|
||||
"text": "What is RAG?\tRAG retrieves then generates.",
|
||||
}
|
||||
chunks := qaInvoke(t, inputs)
|
||||
if len(chunks) == 0 {
|
||||
t.Fatalf("no chunks produced")
|
||||
}
|
||||
cw := chunks[0]["content_with_weight"].(string)
|
||||
if !contains(cw, "问题:") || !contains(cw, "回答:") {
|
||||
t.Errorf("empty lang should default to Chinese prefixes, got %q", cw)
|
||||
}
|
||||
if contains(cw, "Question:") || contains(cw, "Answer:") {
|
||||
t.Errorf("empty lang must not use English prefixes, got %q", cw)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRmQAPrefixStripsMultipleSeparators exercises migration diff
|
||||
// Chunker-2.12: the prefix regex must allow one-or-more separator chars
|
||||
// (Python uses `[\t:: ]+`), so "Q:: answer" is fully stripped. The legacy
|
||||
// Go pattern only matched a single separator, leaving ": answer".
|
||||
func TestRmQAPrefixStripsMultipleSeparators(t *testing.T) {
|
||||
if got := rmQAPrefix("Q:: answer"); got != "answer" {
|
||||
t.Errorf("multi-separator prefix not fully stripped: got %q", got)
|
||||
}
|
||||
if got := rmQAPrefix("Question: foo"); got != "foo" {
|
||||
t.Errorf("single-separator prefix regression: got %q", got)
|
||||
}
|
||||
if got := rmQAPrefix("问:答案在此"); got != "答案在此" {
|
||||
t.Errorf("CJK prefix regression: got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQAChunker_SetsTopInt exercises migration diff Chunker-1.8 (top_int):
|
||||
// each QA chunk must carry the source row index in `top_int`, matching
|
||||
// Python beAdoc(..., row_num=i).
|
||||
func TestQAChunker_SetsTopInt(t *testing.T) {
|
||||
inputs := map[string]any{
|
||||
"name": "test.txt",
|
||||
"output_format": "text",
|
||||
"text": "Q1\tA1\nQ2\tA2",
|
||||
}
|
||||
chunks := qaInvoke(t, inputs)
|
||||
if len(chunks) != 2 {
|
||||
t.Fatalf("want 2 QA chunks, got %d", len(chunks))
|
||||
}
|
||||
for i, c := range chunks {
|
||||
raw, ok := c["top_int"]
|
||||
if !ok {
|
||||
t.Fatalf("chunk %d missing top_int field", i)
|
||||
}
|
||||
arr, ok := raw.([]any)
|
||||
if !ok || len(arr) != 1 {
|
||||
t.Fatalf("chunk %d top_int wrong shape: %#v", i, raw)
|
||||
}
|
||||
if int(arr[0].(float64)) != i {
|
||||
t.Errorf("chunk %d top_int = %v, want %d", i, arr[0], i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestQAChunker_CarriesImageAndPositions exercises migration diff
|
||||
// Chunker-1.8 (image / positions): when the upstream JSON item already
|
||||
// carries an image id and pdf positions, the QA chunk must preserve them
|
||||
// (Python beAdocPdf sets d["image"] and add_positions).
|
||||
func TestQAChunker_CarriesImageAndPositions(t *testing.T) {
|
||||
inputs := map[string]any{
|
||||
"name": "test.pdf",
|
||||
"output_format": "json",
|
||||
"json": []map[string]any{
|
||||
{
|
||||
"text": "Q\tA",
|
||||
"image": "img-42",
|
||||
"_pdf_positions": [][]int{{1, 2, 3, 4, 5}},
|
||||
},
|
||||
},
|
||||
}
|
||||
chunks := qaInvoke(t, inputs)
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("want 1 QA chunk, got %d", len(chunks))
|
||||
}
|
||||
c := chunks[0]
|
||||
if c["image"] != "img-42" {
|
||||
t.Errorf("QA chunk lost upstream image: %#v", c["image"])
|
||||
}
|
||||
if _, ok := c["_pdf_positions"]; !ok {
|
||||
t.Errorf("QA chunk lost upstream _pdf_positions")
|
||||
}
|
||||
// The prefix-stripped content must still be present.
|
||||
cw, _ := c["content_with_weight"].(string)
|
||||
if !contains(cw, "A") {
|
||||
t.Errorf("QA content missing answer: %q", cw)
|
||||
}
|
||||
}
|
||||
|
||||
// contains is a tiny helper to avoid importing strings in every test.
|
||||
func contains(s, sub string) bool {
|
||||
return strings.Contains(s, sub)
|
||||
}
|
||||
@@ -39,7 +39,7 @@ func TestQAChunker_Registered(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestQAChunker_DelimiterTab(t *testing.T) {
|
||||
comp, err := NewQAChunker(nil)
|
||||
comp, err := NewQAChunker(map[string]any{"lang": "english"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -64,7 +64,7 @@ func TestQAChunker_DelimiterTab(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestQAChunker_DelimiterComma(t *testing.T) {
|
||||
comp, err := NewQAChunker(nil)
|
||||
comp, err := NewQAChunker(map[string]any{"lang": "english"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -129,7 +129,7 @@ func TestQAChunker_HTMLTable(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestQAChunker_RmQAPrefix(t *testing.T) {
|
||||
comp, err := NewQAChunker(nil)
|
||||
comp, err := NewQAChunker(map[string]any{"lang": "english"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -170,7 +170,7 @@ func TestQAChunker_Empty(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestQAChunker_CaseInsensitivePrefix(t *testing.T) {
|
||||
comp, err := NewQAChunker(nil)
|
||||
comp, err := NewQAChunker(map[string]any{"lang": "english"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -193,8 +193,8 @@ func TestQAChunker_CaseInsensitivePrefix(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestQAChunker_PrefixRequiresColonOrTab(t *testing.T) {
|
||||
comp, err := NewQAChunker(nil)
|
||||
func TestQAChunker_PrefixSpaceSeparatorStrips(t *testing.T) {
|
||||
comp, err := NewQAChunker(map[string]any{"lang": "english"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -212,8 +212,10 @@ func TestQAChunker_PrefixRequiresColonOrTab(t *testing.T) {
|
||||
t.Fatalf("expected 1 chunk, got %d", len(chunks))
|
||||
}
|
||||
cww, _ := chunks[0]["content_with_weight"].(string)
|
||||
if cww != "Question: A language model is useful\tAnswer: Q How does it work" {
|
||||
t.Fatalf("space-only separator should not strip prefix: %q", cww)
|
||||
// Python qa.py:241 uses `[\t:: ]+`, so a space is a valid separator:
|
||||
// a leading "A"/"Q" followed by a space is stripped (diff Chunker-2.12).
|
||||
if cww != "Question: language model is useful\tAnswer: How does it work" {
|
||||
t.Fatalf("space-separator prefix not stripped: %q", cww)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -54,6 +54,7 @@ package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
@@ -261,8 +262,9 @@ func matchLayoutLevel(text, layout string, fallbackLevel int) int {
|
||||
|
||||
// resolveTitleLevels mirrors common.py:resolve_frequency_levels over the
|
||||
// full record stream. It is the "frequency" branch of
|
||||
// common.py:resolve_title_levels (the outline branch is parity-gap
|
||||
// territory per the SCOPE comment above).
|
||||
// common.py:resolve_title_levels. The outline branch
|
||||
// (common.py:resolve_outline_levels) is handled by resolveOutlineLevels and
|
||||
// tried first in newLevelContext.
|
||||
//
|
||||
// For each record:
|
||||
// - a non-text record is pinned to BODY_LEVEL directly (python skips
|
||||
@@ -339,6 +341,140 @@ func resolveTitleLevels(records []lineRecord, p *titleChunkerParam) []int {
|
||||
return out
|
||||
}
|
||||
|
||||
// outlineEntry is one PDF bookmark/heading from the parser-supplied
|
||||
// outline, mirroring Python extract_pdf_outlines' (text, level, page)
|
||||
// tuple. The page is unused by title detection.
|
||||
type outlineEntry struct {
|
||||
title string
|
||||
level int
|
||||
}
|
||||
|
||||
// outlineSimilarity mirrors common.py:_outline_similarity: the Jaccard
|
||||
// overlap of character bigrams between two strings. It is rune-based so it
|
||||
// matches Python's code-point indexing (str[i] is a Unicode character, not
|
||||
// a byte). The right-hand bigram set is capped at min(len(left), len(right)-1)
|
||||
// characters, exactly as the Python range() does.
|
||||
func outlineSimilarity(left, right string) float64 {
|
||||
lr := []rune(left)
|
||||
rr := []rune(right)
|
||||
leftPairs := make(map[string]struct{}, max(0, len(lr)-1))
|
||||
for i := 0; i+1 < len(lr); i++ {
|
||||
leftPairs[string(lr[i])+string(lr[i+1])] = struct{}{}
|
||||
}
|
||||
n := len(lr)
|
||||
if m := len(rr) - 1; m < n {
|
||||
n = m
|
||||
}
|
||||
if n < 0 {
|
||||
n = 0
|
||||
}
|
||||
rightPairs := make(map[string]struct{}, n)
|
||||
for i := 0; i < n; i++ {
|
||||
rightPairs[string(rr[i])+string(rr[i+1])] = struct{}{}
|
||||
}
|
||||
denom := len(leftPairs)
|
||||
if len(rightPairs) > denom {
|
||||
denom = len(rightPairs)
|
||||
}
|
||||
if denom == 0 {
|
||||
return 0
|
||||
}
|
||||
inter := 0
|
||||
for k := range leftPairs {
|
||||
if _, ok := rightPairs[k]; ok {
|
||||
inter++
|
||||
}
|
||||
}
|
||||
return float64(inter) / float64(denom)
|
||||
}
|
||||
|
||||
// resolveOutlineLevels mirrors common.py:resolve_outline_levels. Each text
|
||||
// record is matched against the outline by character-bigram similarity (>0.8
|
||||
// assigns level+1); unmatched records stay BODY_LEVEL. It returns ok=false
|
||||
// when there is no outline, or when the outline is too sparse relative to the
|
||||
// record count (len(outlines)/len(records) <= 0.03), in which case the
|
||||
// caller falls back to frequency-based detection. mostLevel mirrors Python's
|
||||
// max(1, max_outline_level).
|
||||
func resolveOutlineLevels(records []lineRecord, outline []outlineEntry) (levels []int, mostLevel int, ok bool) {
|
||||
if len(outline) == 0 || len(records) == 0 {
|
||||
return nil, 0, false
|
||||
}
|
||||
if float64(len(outline))/float64(len(records)) <= 0.03 {
|
||||
return nil, 0, false
|
||||
}
|
||||
maxLevel := 0
|
||||
for _, o := range outline {
|
||||
if o.level > maxLevel {
|
||||
maxLevel = o.level
|
||||
}
|
||||
}
|
||||
levels = make([]int, len(records))
|
||||
for i, rec := range records {
|
||||
if !rec.isText() {
|
||||
levels[i] = bodyLevel
|
||||
continue
|
||||
}
|
||||
matched := 0
|
||||
for _, o := range outline {
|
||||
if outlineSimilarity(o.title, rec.text) > 0.8 {
|
||||
matched = o.level + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
if matched == 0 {
|
||||
levels[i] = bodyLevel
|
||||
} else {
|
||||
levels[i] = matched
|
||||
}
|
||||
}
|
||||
return levels, max(1, maxLevel), true
|
||||
}
|
||||
|
||||
// outlineFromInputs reads the parser-supplied PDF outline from the upstream
|
||||
// file metadata (file.outline, written by the ingestion PDF parser's
|
||||
// outlinesToFileMeta) and normalizes it into the chunker's outlineEntry
|
||||
// shape. Returns nil when no outline is present, so callers fall back to
|
||||
// frequency-based title detection. Numbers are coerced from int/float64
|
||||
// because the runtime may hand the chunker a JSON-decoded payload.
|
||||
func outlineFromInputs(inputs map[string]any) []outlineEntry {
|
||||
file, _ := inputs["file"].(map[string]any)
|
||||
if file == nil {
|
||||
return nil
|
||||
}
|
||||
raw, _ := file["outline"].([]any)
|
||||
if len(raw) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]outlineEntry, 0, len(raw))
|
||||
for _, item := range raw {
|
||||
m, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
title, _ := m["title"].(string)
|
||||
if title == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, outlineEntry{title: title, level: anyToInt(m["level"])})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func anyToInt(v any) int {
|
||||
switch t := v.(type) {
|
||||
case int:
|
||||
return t
|
||||
case int64:
|
||||
return int(t)
|
||||
case float64:
|
||||
return int(t)
|
||||
case float32:
|
||||
return int(t)
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// bodyLevel is the sentinel python uses for non-heading lines. We use
|
||||
// the same large int (sys.maxsize - 1) for parity. Practically this
|
||||
// just needs to be "larger than any realistic heading level"; tests
|
||||
@@ -357,11 +493,12 @@ func lineRecordsFromText(text string) []lineRecord {
|
||||
continue
|
||||
}
|
||||
out = append(out, lineRecord{
|
||||
text: ln,
|
||||
docType: "text",
|
||||
imgID: nil,
|
||||
layout: "",
|
||||
pdfPos: nil,
|
||||
text: ln,
|
||||
docType: "text",
|
||||
imgID: nil,
|
||||
layout: "",
|
||||
pdfPositions: nil,
|
||||
positions: nil,
|
||||
})
|
||||
}
|
||||
return out
|
||||
@@ -371,13 +508,14 @@ func lineRecordsFromText(text string) []lineRecord {
|
||||
// common.py:extract_line_records yields. Used by Group/Hierarchy
|
||||
// chunk-builders.
|
||||
type lineRecord struct {
|
||||
text string
|
||||
docType string
|
||||
imgID *string
|
||||
layout string
|
||||
ckType string
|
||||
pdfPos []map[string]any
|
||||
parentMeta map[string]any
|
||||
text string
|
||||
docType string
|
||||
imgID *string
|
||||
layout string
|
||||
ckType string
|
||||
pdfPositions json.RawMessage
|
||||
positions json.RawMessage
|
||||
parentMeta map[string]any
|
||||
}
|
||||
|
||||
func (r lineRecord) textOrEmpty() string { return r.text }
|
||||
@@ -421,7 +559,13 @@ type LevelContext struct {
|
||||
mostLevel int
|
||||
}
|
||||
|
||||
func newLevelContext(records []lineRecord, p *titleChunkerParam) LevelContext {
|
||||
// newLevelContext resolves per-line heading levels, mirroring Python's
|
||||
// resolve_title_levels: try the PDF outline branch first (when an outline is
|
||||
// supplied and dense enough), otherwise fall back to frequency detection.
|
||||
func newLevelContext(records []lineRecord, outline []outlineEntry, p *titleChunkerParam) LevelContext {
|
||||
if levels, mostLevel, ok := resolveOutlineLevels(records, outline); ok {
|
||||
return LevelContext{levels: levels, mostLevel: mostLevel}
|
||||
}
|
||||
levels := resolveTitleLevels(records, p)
|
||||
// most_level is the most-frequent non-body heading level
|
||||
// (common.py:resolve_frequency_levels). Python computes this via
|
||||
|
||||
@@ -255,7 +255,7 @@ func TestNewLevelContext_MostLevelIsMode(t *testing.T) {
|
||||
{text: "## c", docType: "text"},
|
||||
{text: "### d", docType: "text"},
|
||||
}
|
||||
lc := newLevelContext(records, p)
|
||||
lc := newLevelContext(records, nil, p)
|
||||
// selectLevelGroup picks the single 3-pattern family; per-line
|
||||
// levels are [1,2,2,3], so the mode (most frequent heading level)
|
||||
// is level 2.
|
||||
@@ -371,3 +371,99 @@ func TestResolveTitleLevels_LayoutFallback(t *testing.T) {
|
||||
t.Errorf("plain body level = %d, want bodyLevel", levels[2])
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveOutlineLevels covers Chunker-1.5: a text line that matches a
|
||||
// PDF outline entry by character-bigram similarity (>0.8) is assigned the
|
||||
// outline level+1; lines that do not match stay BODY_LEVEL. The non-text
|
||||
// record is pinned to BODY_LEVEL regardless.
|
||||
func TestResolveOutlineLevels(t *testing.T) {
|
||||
outline := []outlineEntry{
|
||||
{title: "第一章 概述", level: 0},
|
||||
{title: "第二章 方法", level: 1},
|
||||
}
|
||||
records := []lineRecord{
|
||||
{text: "第一章 概述", docType: "text"},
|
||||
{text: "本章介绍背景。", docType: "text"}, // no outline match
|
||||
{text: "figure caption", docType: "image"}, // non-text
|
||||
}
|
||||
levels, mostLevel, ok := resolveOutlineLevels(records, outline)
|
||||
if !ok {
|
||||
t.Fatalf("resolveOutlineLevels ok = false, want true")
|
||||
}
|
||||
if levels[0] != 1 { // outline level 0 + 1
|
||||
t.Errorf("matched line level = %d, want 1", levels[0])
|
||||
}
|
||||
if levels[1] != bodyLevel {
|
||||
t.Errorf("unmatched body level = %d, want bodyLevel", levels[1])
|
||||
}
|
||||
if levels[2] != bodyLevel {
|
||||
t.Errorf("non-text level = %d, want bodyLevel", levels[2])
|
||||
}
|
||||
if mostLevel != 1 { // max(1, max outline level 1)
|
||||
t.Errorf("mostLevel = %d, want 1", mostLevel)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveOutlineLevels_SparseGuard ensures a too-sparse outline (ratio
|
||||
// <= 0.03) is rejected so detection falls back to frequency branch.
|
||||
func TestResolveOutlineLevels_SparseGuard(t *testing.T) {
|
||||
outline := []outlineEntry{{title: "唯一的章节标题", level: 0}}
|
||||
// 100 body records: 1/100 = 0.01 <= 0.03 -> rejected.
|
||||
records := make([]lineRecord, 100)
|
||||
for i := range records {
|
||||
records[i] = lineRecord{text: "body paragraph", docType: "text"}
|
||||
}
|
||||
if _, _, ok := resolveOutlineLevels(records, outline); ok {
|
||||
t.Errorf("sparse outline ok = true, want false")
|
||||
}
|
||||
if _, _, ok := resolveOutlineLevels(records, nil); ok {
|
||||
t.Errorf("nil outline ok = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewLevelContext_OutlineBranch pins Chunker-1.5 end-to-end: when an
|
||||
// outline is supplied, newLevelContext prefers the outline branch over the
|
||||
// regex/frequency branch.
|
||||
func TestNewLevelContext_OutlineBranch(t *testing.T) {
|
||||
p := &titleChunkerParam{TitleChunkerParam: schema.TitleChunkerParam{
|
||||
Method: "group",
|
||||
Levels: [][]string{{`^# `}},
|
||||
}}
|
||||
outline := []outlineEntry{{title: "前言", level: 0}}
|
||||
records := []lineRecord{
|
||||
{text: "前言", docType: "text"}, // matches outline -> level 1
|
||||
{text: "普通正文段落。", docType: "text"}, // no outline -> BODY_LEVEL
|
||||
}
|
||||
lc := newLevelContext(records, outline, p)
|
||||
if got := lc.Levels(); got[0] != 1 || got[1] != bodyLevel {
|
||||
t.Errorf("outline levels = %v, want [1, bodyLevel]", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOutlineFromInputs verifies the parser-supplied file.outline is parsed
|
||||
// into outlineEntry, tolerating float64-encoded levels (runtime JSON path).
|
||||
func TestOutlineFromInputs(t *testing.T) {
|
||||
inputs := map[string]any{
|
||||
"file": map[string]any{
|
||||
"outline": []any{
|
||||
map[string]any{"title": "第一章", "level": float64(0), "page_number": float64(1)},
|
||||
map[string]any{"title": "第二章", "level": float64(1), "page_number": float64(3)},
|
||||
},
|
||||
},
|
||||
}
|
||||
out := outlineFromInputs(inputs)
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("outline len = %d, want 2", len(out))
|
||||
}
|
||||
if out[0].title != "第一章" || out[0].level != 0 {
|
||||
t.Errorf("entry 0 = %+v, want {第一章 0}", out[0])
|
||||
}
|
||||
if out[1].title != "第二章" || out[1].level != 1 {
|
||||
t.Errorf("entry 1 = %+v, want {第二章 1}", out[1])
|
||||
}
|
||||
|
||||
// No file / no outline -> nil (frequency fallback).
|
||||
if got := outlineFromInputs(map[string]any{}); got != nil {
|
||||
t.Errorf("empty inputs outline = %v, want nil", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,10 +46,10 @@
|
||||
// Media-context attachment is per-item sequential; merge is
|
||||
// index-deterministic.
|
||||
//
|
||||
// - No PDF/outline awareness (Python `restore_pdf_text_previews`).
|
||||
// That depends on deepdoc/parser which is out of scope for this
|
||||
// phase; the chunker accepts the parser-style structured JSON
|
||||
// payload and runs the same logic against it.
|
||||
// - PDF text previews (Python `restore_pdf_text_previews`) are
|
||||
// generated on demand for text chunks that carry PDF positions:
|
||||
// cropImageChunks crops the text region and writes a preview image,
|
||||
// then imageUploadDecorator uploads it to img_id. See pdfcrop_cgo.go.
|
||||
package chunker
|
||||
|
||||
import (
|
||||
@@ -285,6 +285,12 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string
|
||||
return chunkOutputs(flatten(merged))
|
||||
}
|
||||
|
||||
// sentenceDelimiter is the sentence/clause-boundary regex used to split
|
||||
// oversized sections. It mirrors Python's default delimiter "\n。;!?"
|
||||
// plus the English ". " fallback, and also breaks on ASCII "!" and "?"
|
||||
// (diff Chunker-2.1).
|
||||
var sentenceDelimiter = regexp.MustCompile(`(\n|[!?。;!?]|\.\s)`)
|
||||
|
||||
// mergeByTokenSize implements exact token-based chunk merging that mirrors
|
||||
// Python's naive_merge (rag/nlp/__init__.py:1156). It uses
|
||||
// tokenizeStr (= tokenizer.NumTokensFromString, cl100k_base BPE) for
|
||||
@@ -303,8 +309,10 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r
|
||||
}
|
||||
|
||||
// Sentence/clause-boundary regex for splitting oversized sections.
|
||||
// Matches Python's default delimiter "\n。;!?" plus English ". " fallback.
|
||||
sentenceDelim := regexp.MustCompile(`(\n|[。;!?]|\.\s)`)
|
||||
// Matches Python's default delimiter "\n。;!?" plus English ". "
|
||||
// fallback. The ASCII "!" and "?" are included to match Python's
|
||||
// full delimiter set (diff Chunker-2.1).
|
||||
sentenceDelim := sentenceDelimiter
|
||||
|
||||
var cks []string // chunk texts
|
||||
var tkns []int // token counts per chunk
|
||||
@@ -319,7 +327,9 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r
|
||||
seg := segment
|
||||
segTokens := tokens
|
||||
if overlapPct > 0 && len(cks) > 0 {
|
||||
prev := cks[len(cks)-1]
|
||||
// Strip parser tags before computing the overlap suffix,
|
||||
// matching Python nlp/__init__.py:1181 (diff Chunker-2.2).
|
||||
prev := removeTag(cks[len(cks)-1])
|
||||
// Take the last overlapped_percent of the previous chunk
|
||||
// (in runes, matching Python's len(overlapped) * ratio).
|
||||
prevRunes := []rune(prev)
|
||||
@@ -656,23 +666,37 @@ func collectContext(chunks []schema.ChunkDoc, i, ctxTokens int, above bool) stri
|
||||
return strings.Join(parts, "")
|
||||
}
|
||||
|
||||
// takeFromEnd returns the last approx `tokens` worth of text (1 token
|
||||
// ≈ 4 bytes is the best-effort approximation used here; python uses
|
||||
// the actual tokenizer).
|
||||
// takeFromEnd returns the smallest tail of text whose token count is >=
|
||||
// tokens, counted exactly via tokenizeStr (diff Chunker-2.4). The previous
|
||||
// 4-bytes-per-token heuristic over-counted for CJK text.
|
||||
func takeFromEnd(text string, tokens int) string {
|
||||
bytes := tokens * 4
|
||||
if bytes >= len(text) {
|
||||
return text
|
||||
runes := []rune(text)
|
||||
// The tail runes[i:] grows as i decreases, so the first (largest i,
|
||||
// i.e. smallest tail) that meets the budget is the answer.
|
||||
for i := len(runes); i > 0; i-- {
|
||||
cand := string(runes[i:])
|
||||
if tokenizeStr(cand) >= tokens {
|
||||
return cand
|
||||
}
|
||||
}
|
||||
return text[len(text)-bytes:]
|
||||
return text
|
||||
}
|
||||
|
||||
// takeFromStart returns the smallest prefix of text whose token count is >=
|
||||
// tokens, counted exactly via tokenizeStr (diff Chunker-2.4).
|
||||
func takeFromStart(text string, tokens int) string {
|
||||
bytes := tokens * 4
|
||||
if bytes >= len(text) {
|
||||
return text
|
||||
runes := []rune(text)
|
||||
best := text
|
||||
// Prefix grows as i increases; the first (smallest) qualifying prefix
|
||||
// is the answer.
|
||||
for i := 1; i <= len(runes); i++ {
|
||||
cand := string(runes[:i])
|
||||
if tokenizeStr(cand) >= tokens {
|
||||
best = cand
|
||||
break
|
||||
}
|
||||
}
|
||||
return text[:bytes]
|
||||
return best
|
||||
}
|
||||
|
||||
// mergeByTokenSizeFromJSON mirrors `naive_merge` at
|
||||
@@ -704,7 +728,10 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over
|
||||
// t = overlapped[overlap_cut:] + t
|
||||
// tnum = num_tokens_from_string(t)
|
||||
if len(merged) > 0 && merged[len(merged)-1].CKType == "text" && overlappedPct > 0 {
|
||||
if prevText := merged[len(merged)-1].Text; prevText != "" {
|
||||
// Strip parser tags before computing the overlap
|
||||
// suffix, matching Python nlp/__init__.py:1181
|
||||
// (diff Chunker-2.2).
|
||||
if prevText := removeTag(merged[len(merged)-1].Text); prevText != "" {
|
||||
runes := []rune(prevText)
|
||||
cut := int(float64(len(runes)) * (100 - overlappedPct*100) / 100.0)
|
||||
if cut < len(runes) {
|
||||
@@ -718,10 +745,21 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over
|
||||
}
|
||||
// Merge into the accumulated text chunk.
|
||||
prev := &merged[len(merged)-1]
|
||||
if prev.Text != "" {
|
||||
// Mirror Python token_chunker.py:236-239: when the accumulated
|
||||
// chunk has empty text, assign the incoming text directly instead
|
||||
// of skipping it (diff Chunker-2.11).
|
||||
if prev.Text == "" {
|
||||
prev.Text = ck.Text
|
||||
} else {
|
||||
prev.Text = prev.Text + "\n" + ck.Text
|
||||
prev.TKNums = intPtr(intValue(prev.TKNums) + tk)
|
||||
}
|
||||
prev.TKNums = intPtr(intValue(prev.TKNums) + tk)
|
||||
// Preserve PDF coordinates across the merge: extend the
|
||||
// coordinate lists instead of dropping the incoming item's
|
||||
// positions. Mirrors Python token_chunker.py:240
|
||||
// `merged[prev][PDF_POSITIONS_KEY].extend(...)` (diffs 2.5 / 2.3).
|
||||
prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions)
|
||||
prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions)
|
||||
}
|
||||
perItem[idx] = merged
|
||||
}
|
||||
@@ -742,6 +780,14 @@ func cloneChunkDoc(in schema.ChunkDoc) schema.ChunkDoc {
|
||||
v := *in.PageNumber
|
||||
out.PageNumber = &v
|
||||
}
|
||||
// Deep-copy the coordinate byte slices so the clone does not alias
|
||||
// the source's backing array (diff 2.5 defensive fix).
|
||||
if in.PDFPositions != nil {
|
||||
out.PDFPositions = append(json.RawMessage(nil), in.PDFPositions...)
|
||||
}
|
||||
if in.Positions != nil {
|
||||
out.Positions = append(json.RawMessage(nil), in.Positions...)
|
||||
}
|
||||
if in.Extra != nil {
|
||||
out.Extra = make(map[string]json.RawMessage, len(in.Extra))
|
||||
for k, v := range in.Extra {
|
||||
@@ -751,6 +797,33 @@ func cloneChunkDoc(in schema.ChunkDoc) schema.ChunkDoc {
|
||||
return out
|
||||
}
|
||||
|
||||
// extendRawJSONArray concatenates two JSON array payloads, mirroring
|
||||
// Python's `merged[prev][KEY].extend(current[KEY])`. Either operand may be
|
||||
// empty; the result is always a valid JSON array (or an empty raw message).
|
||||
// It is used to accumulate PDF coordinate lists (`_pdf_positions`,
|
||||
// `positions`) when text chunks are merged (diffs 2.5 / 2.3).
|
||||
func extendRawJSONArray(a, b json.RawMessage) json.RawMessage {
|
||||
if len(a) == 0 {
|
||||
return b
|
||||
}
|
||||
if len(b) == 0 {
|
||||
return a
|
||||
}
|
||||
var arrA, arrB []json.RawMessage
|
||||
if err := json.Unmarshal(a, &arrA); err != nil {
|
||||
return b
|
||||
}
|
||||
if err := json.Unmarshal(b, &arrB); err != nil {
|
||||
return a
|
||||
}
|
||||
arrA = append(arrA, arrB...)
|
||||
out, err := json.Marshal(arrA)
|
||||
if err != nil {
|
||||
return a
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func flatten(perItem [][]schema.ChunkDoc) []schema.ChunkDoc {
|
||||
var out []schema.ChunkDoc
|
||||
for _, cs := range perItem {
|
||||
|
||||
128
internal/ingestion/component/chunker/token_batch1_test.go
Normal file
128
internal/ingestion/component/chunker/token_batch1_test.go
Normal file
@@ -0,0 +1,128 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
// TestSentenceDelimiterMatchesBangAndQuestion exercises migration diff
|
||||
// Chunker-2.1: the sentence/clause boundary regex used to split oversized
|
||||
// sections must also break on ASCII "!" and "?" (Python's default delimiter
|
||||
// is "\n。;!?"). The legacy Go pattern `(\n|[。;!?]|\.\s)` missed the
|
||||
// ASCII variants, so English fragments like "Hi!" / "Really?" were not
|
||||
// treated as boundaries.
|
||||
func TestSentenceDelimiterMatchesBangAndQuestion(t *testing.T) {
|
||||
// The package-level sentenceDelimiter (introduced by Fix 2.1) must
|
||||
// match ASCII bang/question.
|
||||
if !sentenceDelimiter.MatchString("Hi!") {
|
||||
t.Errorf("sentenceDelimiter should split on '!': %q", "Hi!")
|
||||
}
|
||||
if !sentenceDelimiter.MatchString("Really?") {
|
||||
t.Errorf("sentenceDelimiter should split on '?': %q", "Really?")
|
||||
}
|
||||
|
||||
// Guard: the OLD pattern must NOT match these, proving the test would
|
||||
// have failed before the fix.
|
||||
old := regexp.MustCompile(`(\n|[。;!?]|\.\s)`)
|
||||
if old.MatchString("Hi!") || old.MatchString("Really?") {
|
||||
t.Errorf("guard broken: old pattern unexpectedly matches ASCII !/?")
|
||||
}
|
||||
}
|
||||
|
||||
// TestMergeByTokenSizeFromJSON_OverlapStripsTags exercises migration diff
|
||||
// Chunker-2.2: when a new chunk is started, its overlap prefix must be taken
|
||||
// from the previous chunk AFTER remove_tag, otherwise parser tags (e.g.
|
||||
// "@@1\t2.3##") leak into the overlap region. Mirrors Python
|
||||
// nlp/__init__.py:1181 (remove_tag applied before overlap).
|
||||
func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) {
|
||||
aText := strings.Repeat("word ", 20) + "@@1\t2.3## tail"
|
||||
items := [][]schema.ChunkDoc{
|
||||
{
|
||||
{Text: aText, DocType: "text", CKType: "text", TKNums: intPtr(100)},
|
||||
{Text: "body", DocType: "text", CKType: "text", TKNums: intPtr(5)},
|
||||
},
|
||||
}
|
||||
got := mergeByTokenSizeFromJSON(items, 128, 0.3)
|
||||
merged := got[0]
|
||||
if len(merged) != 2 {
|
||||
t.Fatalf("want 2 merged chunks (overlap path), got %d", len(merged))
|
||||
}
|
||||
// The overlap prefix is prepended to the SECOND chunk. The original
|
||||
// first chunk legitimately keeps its own parser tag; only the overlap
|
||||
// region (merged[1]) must be tag-free (diff Chunker-2.2).
|
||||
if strings.Contains(merged[1].Text, "@@") || strings.Contains(merged[1].Text, "##") {
|
||||
t.Errorf("overlap prefix leaked parser tag into chunk 1: %q", merged[1].Text)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk exercises migration diff
|
||||
// Chunker-2.11: merging a non-empty chunk into an empty previous chunk must
|
||||
// assign the text directly instead of being skipped. The legacy guard
|
||||
// `if prev.Text != ""` silently dropped the incoming chunk when the previous
|
||||
// one had empty text. Mirrors Python token_chunker.py:236-239.
|
||||
func TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk(t *testing.T) {
|
||||
items := [][]schema.ChunkDoc{
|
||||
{
|
||||
{Text: "", DocType: "text", CKType: "text", TKNums: intPtr(5)},
|
||||
{Text: "keepme", DocType: "text", CKType: "text", TKNums: intPtr(5)},
|
||||
},
|
||||
}
|
||||
got := mergeByTokenSizeFromJSON(items, 128, 0)
|
||||
merged := got[0]
|
||||
if len(merged) != 1 {
|
||||
t.Fatalf("want 1 merged chunk, got %d", len(merged))
|
||||
}
|
||||
if merged[0].Text != "keepme" {
|
||||
t.Errorf("empty previous chunk dropped incoming text; got %q", merged[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTakeFromEndRespectsTokenCount and TestTakeFromStartRespectsTokenCount
|
||||
// exercise migration diff Chunker-2.4: takeFromEnd/takeFromStart used a
|
||||
// fixed 4-bytes-per-token heuristic which badly over-counts for CJK text
|
||||
// (≈3 bytes/char, 1-2 tokens/char). They must now count tokens exactly via
|
||||
// tokenizeStr so the returned slice is close to the requested token budget.
|
||||
func TestTakeFromEndRespectsTokenCount(t *testing.T) {
|
||||
const target = 20
|
||||
s := strings.Repeat("中", 60)
|
||||
got := takeFromEnd(s, target)
|
||||
if !strings.HasSuffix(s, got) {
|
||||
t.Fatalf("takeFromEnd result must be a suffix of input")
|
||||
}
|
||||
n := tokenizeStr(got)
|
||||
if n < target-3 || n > target+3 {
|
||||
t.Errorf("takeFromEnd(%d tokens) returned slice with %d tokens (want ~%d)", target, n, target)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTakeFromStartRespectsTokenCount(t *testing.T) {
|
||||
const target = 20
|
||||
s := strings.Repeat("中", 60)
|
||||
got := takeFromStart(s, target)
|
||||
if !strings.HasPrefix(s, got) {
|
||||
t.Fatalf("takeFromStart result must be a prefix of input")
|
||||
}
|
||||
n := tokenizeStr(got)
|
||||
if n < target-3 || n > target+3 {
|
||||
t.Errorf("takeFromStart(%d tokens) returned slice with %d tokens (want ~%d)", target, n, target)
|
||||
}
|
||||
}
|
||||
119
internal/ingestion/component/chunker/token_pdfpos_test.go
Normal file
119
internal/ingestion/component/chunker/token_pdfpos_test.go
Normal file
@@ -0,0 +1,119 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package chunker
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
// TestMergeByTokenSizeFromJSON_ExtendsPDFPositions is the TDD test for
|
||||
// migration diffs Chunker-2.5 / 2.3: when two JSON text items carrying
|
||||
// `_pdf_positions` / `positions` are merged into one chunk, the merged
|
||||
// chunk must extend (not drop) the coordinate lists — mirroring Python
|
||||
// token_chunker.py:240 `merged[prev][PDF_POSITIONS_KEY].extend(...)`.
|
||||
func TestMergeByTokenSizeFromJSON_ExtendsPDFPositions(t *testing.T) {
|
||||
posA := json.RawMessage(`[[1,10,20,30,40]]`)
|
||||
posB := json.RawMessage(`[[2,15,25,35,45]]`)
|
||||
items := [][]schema.ChunkDoc{
|
||||
{
|
||||
{Text: "alpha", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posA},
|
||||
{Text: "beta", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posB},
|
||||
},
|
||||
}
|
||||
got := mergeByTokenSizeFromJSON(items, 128, 0)
|
||||
merged := got[0]
|
||||
if len(merged) != 1 {
|
||||
t.Fatalf("want 1 merged chunk, got %d", len(merged))
|
||||
}
|
||||
combined := string(merged[0].PDFPositions)
|
||||
if !strings.Contains(combined, "1,10,20,30,40") {
|
||||
t.Errorf("merged chunk lost first item _pdf_positions: %s", combined)
|
||||
}
|
||||
if !strings.Contains(combined, "2,15,25,35,45") {
|
||||
t.Errorf("merged chunk dropped second item _pdf_positions (not extended): %s", combined)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMergeByTokenSizeFromJSON_ExtendsPositions covers the parallel
|
||||
// `positions` field (diff 2.3).
|
||||
func TestMergeByTokenSizeFromJSON_ExtendsPositions(t *testing.T) {
|
||||
posA := json.RawMessage(`[[1,2,3]]`)
|
||||
posB := json.RawMessage(`[[4,5,6]]`)
|
||||
items := [][]schema.ChunkDoc{
|
||||
{
|
||||
{Text: "a", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posA},
|
||||
{Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB},
|
||||
},
|
||||
}
|
||||
got := mergeByTokenSizeFromJSON(items, 128, 0)
|
||||
combined := string(got[0][0].Positions)
|
||||
if !strings.Contains(combined, "1,2,3") || !strings.Contains(combined, "4,5,6") {
|
||||
t.Errorf("merged chunk dropped/omitted `positions`: %s", combined)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloneChunkDoc_DeepCopiesPDFPositions ensures cloneChunkDoc does not
|
||||
// alias the underlying _pdf_positions / positions byte slices (diff 2.5
|
||||
// defensive fix).
|
||||
func TestCloneChunkDoc_DeepCopiesPDFPositions(t *testing.T) {
|
||||
pos := json.RawMessage(`[[1,2,3,4,5]]`)
|
||||
orig := schema.ChunkDoc{Text: "x", PDFPositions: pos, Positions: pos}
|
||||
cp := cloneChunkDoc(orig)
|
||||
// Mutate the source's backing array after the clone.
|
||||
pos[0] = '9'
|
||||
if string(cp.PDFPositions) != "[[1,2,3,4,5]]" {
|
||||
t.Errorf("clone shares _pdf_positions backing array: %s", string(cp.PDFPositions))
|
||||
}
|
||||
if string(cp.Positions) != "[[1,2,3,4,5]]" {
|
||||
t.Errorf("clone shares positions backing array: %s", string(cp.Positions))
|
||||
}
|
||||
}
|
||||
|
||||
// TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix verifies the
|
||||
// chunker-side contract for diff 1.4: preserved `positions` must decode
|
||||
// (via ChunkDoc.ToMap → decodeStructuredValue) to a [][]float64 matrix so
|
||||
// the downstream task-layer processChunkPositions → AddPositions can
|
||||
// convert it to page_num_int / top_int / position_int. The coordinate
|
||||
// conversion itself lives in internal/ingestion/task (processChunkPositions),
|
||||
// not in the chunker.
|
||||
func TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix(t *testing.T) {
|
||||
posA := json.RawMessage(`[[1,2,3,4,5]]`)
|
||||
posB := json.RawMessage(`[[6,7,8,9,10]]`)
|
||||
items := [][]schema.ChunkDoc{
|
||||
{
|
||||
{Text: "a", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posA},
|
||||
{Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB},
|
||||
},
|
||||
}
|
||||
got := mergeByTokenSizeFromJSON(items, 128, 0)
|
||||
m := got[0][0].ToMap()
|
||||
raw, ok := m["positions"]
|
||||
if !ok {
|
||||
t.Fatal("positions missing from ToMap output")
|
||||
}
|
||||
matrix, ok := raw.([][]float64)
|
||||
if !ok {
|
||||
t.Fatalf("positions decoded to %T, want [][]float64", raw)
|
||||
}
|
||||
if len(matrix) != 2 {
|
||||
t.Fatalf("positions matrix has %d groups, want 2 (both merged items)", len(matrix))
|
||||
}
|
||||
}
|
||||
@@ -16,13 +16,14 @@
|
||||
|
||||
// DOCX vision figure dispatch: enriches the parse result with
|
||||
// LLM-generated descriptions of embedded images, mirroring
|
||||
// Python's vision_figure_parser_docx_wrapper_naive in
|
||||
// deepdoc/parser/figure_parser.py.
|
||||
// Python's enhance_media_sections_with_vision in
|
||||
// rag/flow/parser/utils.py (invoked from parser.py:_doc's JSON branch).
|
||||
//
|
||||
// Unlike the PDF vision path (which replaces dispatchParse
|
||||
// entirely), DOCX vision is a post-processing step: it takes
|
||||
// the already-parsed markdown + extracted figures and augments
|
||||
// the markdown text with vision model descriptions.
|
||||
// Unlike the PDF vision path (which replaces dispatchParse entirely),
|
||||
// DOCX vision is a post-processing step. It mirrors Python exactly:
|
||||
// vision enrichment happens ONLY on the JSON output path, where each
|
||||
// item carries a doc_type_kwd and an optional image. The markdown path
|
||||
// performs no vision enrichment in Python, so it must not here either.
|
||||
|
||||
package component
|
||||
|
||||
@@ -31,7 +32,6 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
@@ -59,18 +59,20 @@ var (
|
||||
docxVisionPromptMu sync.RWMutex
|
||||
)
|
||||
|
||||
// maybeDispatchDOCXVision checks whether the dispatch result for a
|
||||
// DOCX file contains embedded image figures and, when a vision
|
||||
// model is available, enriches the markdown with AI-generated
|
||||
// figure descriptions. It mirrors the Python flow:
|
||||
// maybeDispatchDOCXVision enriches a DOCX parse result with vision-model
|
||||
// descriptions of embedded images. It mirrors Python's
|
||||
// enhance_media_sections_with_vision (rag/flow/parser/utils.py:162), which
|
||||
// runs only in the JSON output branch of parser.py:_doc.
|
||||
//
|
||||
// 1. naive_merge_docx → chunks (text + images + context)
|
||||
// 2. vision_figure_parser_docx_wrapper_naive → LLM descriptions
|
||||
// For each JSON item whose doc_type_kwd is "image" or "table" AND that
|
||||
// carries a non-empty "image" field, the vision model describes the image
|
||||
// and the description is appended to the item's text (Python:
|
||||
// item["text"] = f"{text}\n{parsed_text}" if text else parsed_text). Items
|
||||
// without an image (e.g. DOCX tables) are left untouched, exactly as Python
|
||||
// skips them via `if item.get("image") is None: continue`.
|
||||
//
|
||||
// The function is called AFTER dispatchParse so the normal parse
|
||||
// path produces figures in dispatched.File["figures"].
|
||||
// It returns (result, handled, error). handled is true when the
|
||||
// dispatched result was modified.
|
||||
// The markdown output path receives no vision enrichment — Python's DOCX
|
||||
// markdown branch only concatenates text and never calls the vision model.
|
||||
func maybeDispatchDOCXVision(
|
||||
ctx context.Context,
|
||||
fileType utility.FileType,
|
||||
@@ -81,11 +83,9 @@ func maybeDispatchDOCXVision(
|
||||
if fileType != utility.FileTypeDOCX {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
if dispatched.Err != nil || dispatched.OutputFormat != "markdown" {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
figs, hasFigures := extractDOCXFiguresFromDispatch(dispatched)
|
||||
if !hasFigures {
|
||||
// Python triggers vision enrichment only on the JSON path
|
||||
// (parser.py:_doc → enhance_media_sections_with_vision).
|
||||
if dispatched.Err != nil || dispatched.OutputFormat != "json" || len(dispatched.JSON) == 0 {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
|
||||
@@ -102,105 +102,74 @@ func maybeDispatchDOCXVision(
|
||||
return dispatched, false, nil
|
||||
}
|
||||
|
||||
descriptions := make([]string, len(figs))
|
||||
// Collect the indices of JSON items that carry an embeddable image.
|
||||
type target struct {
|
||||
idx int
|
||||
}
|
||||
var targets []target
|
||||
for i, item := range dispatched.JSON {
|
||||
kd, _ := item["doc_type_kwd"].(string)
|
||||
if kd != "image" && kd != "table" {
|
||||
continue
|
||||
}
|
||||
img, _ := item["image"].(string)
|
||||
if img == "" {
|
||||
continue
|
||||
}
|
||||
targets = append(targets, target{idx: i})
|
||||
}
|
||||
if len(targets) == 0 {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
|
||||
descriptions := make([]string, len(targets))
|
||||
var wg sync.WaitGroup
|
||||
sem := make(chan struct{}, docxVisionConcurrency)
|
||||
|
||||
for i, fig := range figs {
|
||||
for slot, tg := range targets {
|
||||
wg.Add(1)
|
||||
go func(idx int, f map[string]any) {
|
||||
go func(slot int, itemIdx int) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{}
|
||||
defer func() { <-sem }()
|
||||
|
||||
imageB64, _ := f["image"].(string)
|
||||
ctxAbove, _ := f["context_above"].(string)
|
||||
ctxBelow, _ := f["context_below"].(string)
|
||||
|
||||
if imageB64 == "" {
|
||||
img, _ := dispatched.JSON[itemIdx]["image"].(string)
|
||||
if img == "" {
|
||||
return
|
||||
}
|
||||
|
||||
prompt, err := docxVisionPromptBuilder(ctxAbove, ctxBelow)
|
||||
if err != nil {
|
||||
// DOCX JSON items have no surrounding context (unlike the
|
||||
// former markdown path), so use the bare figure prompt —
|
||||
// matching Python's VisionFigureParser(context_size=0).
|
||||
prompt, perr := docxVisionPromptBuilder("", "")
|
||||
if perr != nil {
|
||||
return
|
||||
}
|
||||
|
||||
messages := buildVisionMessages(prompt, imageB64)
|
||||
resp, err := visionChatInvoker(ctx, driver, modelName, messages, apiConfig)
|
||||
if err != nil {
|
||||
messages := buildVisionMessages(prompt, img)
|
||||
resp, ierr := visionChatInvoker(ctx, driver, modelName, messages, apiConfig)
|
||||
if ierr != nil {
|
||||
return
|
||||
}
|
||||
descriptions[idx] = extractDOCXVisionAnswer(resp)
|
||||
}(i, fig)
|
||||
descriptions[slot] = extractDOCXVisionAnswer(resp)
|
||||
}(slot, tg.idx)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Insert each description at the figure's position in the markdown,
|
||||
// matching Python's `chunks[idx]["text"] += description`.
|
||||
// Figures carry a "marker" (text immediately before the image) to
|
||||
// locate the insertion point. Process in reverse order so earlier
|
||||
// insertions don't shift later markers.
|
||||
md := dispatched.Markdown
|
||||
type indexedDesc struct {
|
||||
idx int
|
||||
desc string
|
||||
}
|
||||
var inserts []indexedDesc
|
||||
for i, d := range descriptions {
|
||||
if d = strings.TrimSpace(d); d == "" {
|
||||
modified := false
|
||||
for slot, tg := range targets {
|
||||
desc := strings.TrimSpace(descriptions[slot])
|
||||
if desc == "" {
|
||||
continue
|
||||
}
|
||||
if i >= len(figs) {
|
||||
continue
|
||||
existing, _ := dispatched.JSON[tg.idx]["text"].(string)
|
||||
if existing != "" {
|
||||
dispatched.JSON[tg.idx]["text"] = existing + "\n" + desc
|
||||
} else {
|
||||
dispatched.JSON[tg.idx]["text"] = desc
|
||||
}
|
||||
marker, _ := figs[i]["marker"].(string)
|
||||
if marker != "" {
|
||||
if pos := strings.LastIndex(md, marker); pos >= 0 {
|
||||
inserts = append(inserts, indexedDesc{idx: pos + len(marker), desc: d})
|
||||
continue
|
||||
}
|
||||
}
|
||||
// Fallback: try context_above as a search anchor.
|
||||
if ctx, _ := figs[i]["context_above"].(string); ctx != "" {
|
||||
if pos := strings.LastIndex(md, ctx); pos >= 0 {
|
||||
inserts = append(inserts, indexedDesc{idx: pos + len(ctx), desc: d})
|
||||
continue
|
||||
}
|
||||
}
|
||||
// No anchor found — append to end.
|
||||
inserts = append(inserts, indexedDesc{idx: len(md), desc: "\n\n" + d})
|
||||
modified = true
|
||||
}
|
||||
// Sort descending by position for stable insertion.
|
||||
sort.Slice(inserts, func(a, b int) bool { return inserts[a].idx > inserts[b].idx })
|
||||
for _, ins := range inserts {
|
||||
desc := ins.desc
|
||||
if !strings.HasPrefix(desc, "\n") {
|
||||
desc = "\n\n" + desc
|
||||
}
|
||||
md = md[:ins.idx] + desc + md[ins.idx:]
|
||||
}
|
||||
dispatched.Markdown = md
|
||||
|
||||
return dispatched, true, nil
|
||||
}
|
||||
|
||||
func extractDOCXFiguresFromDispatch(dispatched parserDispatchResult) ([]map[string]any, bool) {
|
||||
if dispatched.File == nil {
|
||||
return nil, false
|
||||
}
|
||||
raw, ok := dispatched.File["figures"]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
list, ok := raw.([]map[string]any)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
if len(list) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return list, true
|
||||
return dispatched, modified, nil
|
||||
}
|
||||
|
||||
// buildDOCXVisionPrompt loads the figure-describe prompt template
|
||||
|
||||
181
internal/ingestion/component/docx_vision_dispatch_test.go
Normal file
181
internal/ingestion/component/docx_vision_dispatch_test.go
Normal file
@@ -0,0 +1,181 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/entity"
|
||||
modelModule "ragflow/internal/entity/models"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
// docxVisionFakeDriver satisfies modelModule.ModelDriver but never reaches the
|
||||
// network: docxVisionCaptureInvoker intercepts the call before the driver.
|
||||
type docxVisionFakeDriver struct {
|
||||
modelModule.ModelDriver
|
||||
}
|
||||
|
||||
// docxVisionCaptureInvoker records the requested image and returns a fixed
|
||||
// description, mirroring the markdown-vision test's capture driver.
|
||||
type docxVisionCaptureInvoker struct {
|
||||
mu sync.Mutex
|
||||
images []string
|
||||
captured []modelModule.Message
|
||||
}
|
||||
|
||||
func (c *docxVisionCaptureInvoker) invoke(
|
||||
ctx context.Context,
|
||||
driver modelModule.ModelDriver,
|
||||
modelName string,
|
||||
messages []modelModule.Message,
|
||||
apiConfig *modelModule.APIConfig,
|
||||
) (*modelModule.ChatResponse, error) {
|
||||
c.mu.Lock()
|
||||
c.captured = append(c.captured, messages...)
|
||||
// Pull the data URI out of the second content part.
|
||||
if parts, ok := messages[0].Content.([]interface{}); ok && len(parts) >= 2 {
|
||||
if img, ok := parts[1].(map[string]any); ok {
|
||||
if url, ok := img["image_url"].(map[string]any); ok {
|
||||
if u, ok := url["url"].(string); ok {
|
||||
c.images = append(c.images, u)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
c.mu.Unlock()
|
||||
ans := "a diagram of a pipeline"
|
||||
return &modelModule.ChatResponse{Answer: &ans}, nil
|
||||
}
|
||||
|
||||
// TestMaybeDispatchDOCXVision_EnhancesJSONImages verifies Diff 2.4: DOCX vision
|
||||
// enhancement must trigger on the JSON output path (like Python's
|
||||
// enhance_media_sections_with_vision in parser.py:_doc) and must NOT trigger on
|
||||
// the markdown path. Image items with a non-empty `image` field get their VLM
|
||||
// description appended to `text`; table items (no image) and text items are
|
||||
// left untouched.
|
||||
func TestMaybeDispatchDOCXVision_EnhancesJSONImages(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
origInvoker := visionChatInvoker
|
||||
origPrompt := docxVisionPromptBuilder
|
||||
defer func() {
|
||||
resolveTenantModelByType = origResolver
|
||||
visionChatInvoker = origInvoker
|
||||
docxVisionPromptBuilder = origPrompt
|
||||
}()
|
||||
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return &docxVisionFakeDriver{}, "docx-vision-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
invoker := &docxVisionCaptureInvoker{}
|
||||
visionChatInvoker = invoker.invoke
|
||||
docxVisionPromptBuilder = func(string, string) (string, error) { return "describe the figure", nil }
|
||||
|
||||
dispatched := parserDispatchResult{
|
||||
OutputFormat: "json",
|
||||
DocType: "docx",
|
||||
JSON: []map[string]any{
|
||||
{"text": "Intro paragraph", "image": nil, "doc_type_kwd": "text"},
|
||||
{"text": "", "image": "aGVsbG8taW1hZ2U=", "doc_type_kwd": "image"},
|
||||
{"text": "<table></table>", "image": nil, "doc_type_kwd": "table"},
|
||||
},
|
||||
}
|
||||
|
||||
res, handled, err := maybeDispatchDOCXVision(
|
||||
context.Background(),
|
||||
utility.FileTypeDOCX,
|
||||
dispatched,
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
defaultSetups(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchDOCXVision: unexpected error: %v", err)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatal("handled = false, want true (JSON image item should be enhanced)")
|
||||
}
|
||||
if len(res.JSON) != 3 {
|
||||
t.Fatalf("JSON len = %d, want 3", len(res.JSON))
|
||||
}
|
||||
if got := res.JSON[0]["text"].(string); got != "Intro paragraph" {
|
||||
t.Errorf("text item text = %q, want unchanged", got)
|
||||
}
|
||||
// image item: VLM description appended to (empty) text.
|
||||
if got, _ := res.JSON[1]["text"].(string); got != "a diagram of a pipeline" {
|
||||
t.Errorf("image item text = %q, want appended VLM description", got)
|
||||
}
|
||||
if got, _ := res.JSON[2]["text"].(string); got != "<table></table>" {
|
||||
t.Errorf("table item text = %q, want unchanged (no image)", got)
|
||||
}
|
||||
if len(invoker.images) != 1 {
|
||||
t.Fatalf("vision invoker called %d times, want 1 (only the image item)", len(invoker.images))
|
||||
}
|
||||
if want := "data:image/png;base64,aGVsbG8taW1hZ2U="; invoker.images[0] != want {
|
||||
t.Errorf("vision image data URI = %q, want %q", invoker.images[0], want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchDOCXVision_JSONOnly verifies Diff 2.4: the markdown output
|
||||
// path must NOT be enhanced (Python's markdown/docx branch performs no vision
|
||||
// enrichment). A markdown result with embedded figures is returned untouched.
|
||||
func TestMaybeDispatchDOCXVision_JSONOnly(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
origInvoker := visionChatInvoker
|
||||
defer func() {
|
||||
resolveTenantModelByType = origResolver
|
||||
visionChatInvoker = origInvoker
|
||||
}()
|
||||
|
||||
called := false
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
called = true
|
||||
return &docxVisionFakeDriver{}, "m", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
visionChatInvoker = func(ctx context.Context, d modelModule.ModelDriver, m string, msgs []modelModule.Message, c *modelModule.APIConfig) (*modelModule.ChatResponse, error) {
|
||||
called = true
|
||||
ans := "x"
|
||||
return &modelModule.ChatResponse{Answer: &ans}, nil
|
||||
}
|
||||
|
||||
dispatched := parserDispatchResult{
|
||||
OutputFormat: "markdown",
|
||||
DocType: "docx",
|
||||
Markdown: "",
|
||||
File: map[string]any{"figures": []map[string]any{{"image": "abc", "marker": "x"}}},
|
||||
}
|
||||
|
||||
res, handled, err := maybeDispatchDOCXVision(
|
||||
context.Background(),
|
||||
utility.FileTypeDOCX,
|
||||
dispatched,
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
defaultSetups(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchDOCXVision: unexpected error: %v", err)
|
||||
}
|
||||
if handled {
|
||||
t.Error("handled = true, want false (markdown path must not be enhanced)")
|
||||
}
|
||||
if called {
|
||||
t.Error("vision model was resolved/invoked on the markdown path")
|
||||
}
|
||||
if res.Markdown != "" {
|
||||
t.Errorf("markdown mutated: %q", res.Markdown)
|
||||
}
|
||||
}
|
||||
@@ -169,6 +169,13 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) {
|
||||
}
|
||||
if v, ok := params["prompt"].(string); ok {
|
||||
p.Prompt = v
|
||||
} else if v, ok := params["prompts"].(string); ok && v != "" {
|
||||
// Python agent/component/llm.py:119-120 normalizes a bare-string
|
||||
// prompts into [{"role":"user","content":prompts}]. Mirror that
|
||||
// here so a front-end/template that emits prompts as a string
|
||||
// (the graph.nodes form / dsl testdata) is not silently dropped
|
||||
// by the .([]any) assertion on the list branch below.
|
||||
p.Prompt = v
|
||||
} else if promptsRaw, ok := params["prompts"].([]any); ok && len(promptsRaw) > 0 {
|
||||
if first, ok := promptsRaw[0].(map[string]any); ok {
|
||||
if content, ok := first["content"].(string); ok {
|
||||
|
||||
@@ -570,6 +570,48 @@ func TestNewExtractorComponent_PromptsArray_PromptWins(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_PromptsString verifies that a bare-string
|
||||
// "prompts" (the shape emitted by the front-end graph.nodes form and
|
||||
// the dsl/testdata templates) is normalized into Param.Prompt, mirroring
|
||||
// Python agent/component/llm.py:119-120 which coerces a string prompts
|
||||
// into [{"role":"user","content":prompts}]. Without this normalization
|
||||
// the string form is silently dropped (the .([]any) assertion fails).
|
||||
func TestNewExtractorComponent_PromptsString(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"prompts": "Content: {TitleChunker:FlatMiceFix@chunks}",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
want := "Content: {TitleChunker:FlatMiceFix@chunks}"
|
||||
if ec.Param.Prompt != want {
|
||||
t.Errorf("Prompt = %q, want %q (string prompts should be normalized)", ec.Param.Prompt, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_PromptsString_PromptWins verifies that
|
||||
// "prompt" (string) still takes priority over a string-form "prompts"
|
||||
// when both are present, matching the prompt>prompts precedence of
|
||||
// the list-form path (TestNewExtractorComponent_PromptsArray_PromptWins).
|
||||
func TestNewExtractorComponent_PromptsString_PromptWins(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"prompt": "Direct prompt wins.",
|
||||
"prompts": "Should be ignored.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
if ec.Param.Prompt != "Direct prompt wins." {
|
||||
t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, "Direct prompt wins.")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_SystemPromptWinsOverSysPrompt verifies
|
||||
// that "system_prompt" takes priority over "sys_prompt".
|
||||
func TestNewExtractorComponent_SystemPromptWinsOverSysPrompt(t *testing.T) {
|
||||
|
||||
@@ -78,7 +78,10 @@ func maybeDispatchMarkdownVision(
|
||||
var images []imgItem
|
||||
for i, item := range dispatched.JSON {
|
||||
kd, _ := item["doc_type_kwd"].(string)
|
||||
if kd != "image" {
|
||||
// Diff 2.5: Python enhances both image and table items
|
||||
// (parser/utils.py:181 checks {"image","table"}); only items
|
||||
// carrying an image are sent to the VLM (utils.py:183).
|
||||
if kd != "image" && kd != "table" {
|
||||
continue
|
||||
}
|
||||
img, _ := item["image"].(string)
|
||||
|
||||
@@ -168,10 +168,17 @@ func maybeDispatchImage(
|
||||
}
|
||||
}
|
||||
|
||||
outputFormat, _ := setup["output_format"].(string)
|
||||
if outputFormat == "" {
|
||||
outputFormat = "text"
|
||||
}
|
||||
// The image family always emits a structured JSON item carrying the
|
||||
// image attachment (data URI) and doc_type_kwd, mirroring Python
|
||||
// rag/app/picture.py:71-72 (doc["image"]=img, doc["doc_type_kwd"]=
|
||||
// "image"). picture.py has no "text" output mode — it always returns
|
||||
// a structured doc — so output_format is hardcoded to "json" and any
|
||||
// setup override is ignored. The former behavior returned a bare Text
|
||||
// string, which dropped the image attachment, set doc_type to "text",
|
||||
// and on the default json path produced JSON=nil so downstream
|
||||
// Chunkers rejected the payload with errRequiredField{"json"}.
|
||||
imageB64 := base64.StdEncoding.EncodeToString(binary)
|
||||
dataURI := "data:" + imageMIME(filename) + ";base64," + imageB64
|
||||
|
||||
// --- Phase 2: VLM description (when OCR text is short) ---
|
||||
// Mirrors Python's check: if (eng and len(txt.split()) > 32) or len(txt) > 32
|
||||
@@ -184,11 +191,7 @@ func maybeDispatchImage(
|
||||
charCount := len(ocrText)
|
||||
if (eng && wordCount > 32) || charCount > 32 {
|
||||
// OCR returned substantial text — skip VLM.
|
||||
return parserDispatchResult{
|
||||
OutputFormat: outputFormat,
|
||||
DocType: "image",
|
||||
Text: ocrText,
|
||||
}, true, nil
|
||||
return imageDispatchResult(ocrText, dataURI), true, nil
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,20 +200,12 @@ func maybeDispatchImage(
|
||||
if err != nil {
|
||||
// If VLM is unavailable but we have OCR text, return it.
|
||||
if ocrText != "" {
|
||||
return parserDispatchResult{
|
||||
OutputFormat: outputFormat,
|
||||
DocType: "image",
|
||||
Text: ocrText,
|
||||
}, true, nil
|
||||
return imageDispatchResult(ocrText, dataURI), true, nil
|
||||
}
|
||||
return parserDispatchResult{}, true,
|
||||
fmt.Errorf("Parser: picture image2text model: %w", err)
|
||||
}
|
||||
|
||||
imageB64 := base64.StdEncoding.EncodeToString(binary)
|
||||
mimeType := imageMIME(filename)
|
||||
dataURI := "data:" + mimeType + ";base64," + imageB64
|
||||
|
||||
prompt := "Describe this image in detail."
|
||||
// image family's contract key is system_prompt (parser.go:295),
|
||||
// mirroring Python parser.py:1119. Do NOT read setup["prompt"]
|
||||
@@ -229,11 +224,7 @@ func maybeDispatchImage(
|
||||
resp, err := driver.ChatWithMessages(ctx, modelName, messages, apiConfig, &modelModule.ChatConfig{Vision: &vision}, nil)
|
||||
if err != nil {
|
||||
if ocrText != "" {
|
||||
return parserDispatchResult{
|
||||
OutputFormat: outputFormat,
|
||||
DocType: "image",
|
||||
Text: ocrText,
|
||||
}, true, nil
|
||||
return imageDispatchResult(ocrText, dataURI), true, nil
|
||||
}
|
||||
return parserDispatchResult{}, true,
|
||||
fmt.Errorf("Parser: picture describe: %w", err)
|
||||
@@ -253,11 +244,23 @@ func maybeDispatchImage(
|
||||
combined = vlmText
|
||||
}
|
||||
}
|
||||
return imageDispatchResult(combined, dataURI), true, nil
|
||||
}
|
||||
|
||||
// imageDispatchResult builds the structured JSON payload for the image
|
||||
// family: a single item carrying the combined text, the image attachment
|
||||
// (data URI), and doc_type_kwd "image". Mirrors Python
|
||||
// rag/app/picture.py:71-72.
|
||||
func imageDispatchResult(text, dataURI string) parserDispatchResult {
|
||||
return parserDispatchResult{
|
||||
OutputFormat: outputFormat,
|
||||
OutputFormat: "json",
|
||||
DocType: "image",
|
||||
Text: combined,
|
||||
}, true, nil
|
||||
JSON: []map[string]any{{
|
||||
"text": text,
|
||||
"image": dataURI,
|
||||
"doc_type_kwd": "image",
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
// Audio dispatch: SPEECH2TEXT transcription ---
|
||||
@@ -315,6 +318,21 @@ func maybeDispatchAudio(
|
||||
if outputFormat == "" {
|
||||
outputFormat = "text"
|
||||
}
|
||||
// Diff 2.11: when output_format is "json" the transcription must be
|
||||
// carried as a JSON item. Returning it only in Text made the Invoke
|
||||
// switch silently drop it (the switch has no "json" branch and the
|
||||
// JSON slice was empty). Mirror the JSON-item shape used by the
|
||||
// other parser branches.
|
||||
if outputFormat == "json" {
|
||||
return parserDispatchResult{
|
||||
OutputFormat: "json",
|
||||
DocType: "audio",
|
||||
JSON: []map[string]any{{
|
||||
"text": transcription,
|
||||
"doc_type_kwd": "audio",
|
||||
}},
|
||||
}, true, nil
|
||||
}
|
||||
return parserDispatchResult{
|
||||
OutputFormat: outputFormat,
|
||||
DocType: "audio",
|
||||
|
||||
@@ -17,6 +17,7 @@ package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
@@ -104,8 +105,18 @@ func TestMaybeDispatchImage_UsesSystemPrompt(t *testing.T) {
|
||||
if !dispatched {
|
||||
t.Fatalf("expected dispatched=true for VISUAL file")
|
||||
}
|
||||
if res.Text == "" {
|
||||
t.Fatalf("expected non-empty combined text")
|
||||
// After the output-shape fix the image branch returns JSON items
|
||||
// (OutputFormat=="json"), not a bare Text field. The combined text
|
||||
// now lives in JSON[0]["text"]; the legacy res.Text is no longer
|
||||
// populated for the image family.
|
||||
if res.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want json (image family is always structured)", res.OutputFormat)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1 (image result must be a single JSON item)", len(res.JSON))
|
||||
}
|
||||
if txt, _ := res.JSON[0]["text"].(string); txt == "" {
|
||||
t.Fatalf("expected non-empty combined text in JSON[0][\"text\"]")
|
||||
}
|
||||
|
||||
got, ok := firstUserText(drv.captured)
|
||||
@@ -116,3 +127,238 @@ func TestMaybeDispatchImage_UsesSystemPrompt(t *testing.T) {
|
||||
t.Fatalf("VLM user text = %q, want %q (image branch must read system_prompt)", got, "自定义视觉提示")
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchImage_ReturnsJSONWithImage pins the output-shape fix:
|
||||
// the image branch must return a JSON item carrying the `image` attachment
|
||||
// (data URI) and `doc_type_kwd:"image"`, mirroring Python
|
||||
// rag/app/picture.py:71-72. Before the fix the branch returned a bare Text
|
||||
// string with JSON=nil, dropping the image attachment and (on the default
|
||||
// json path) causing OneChunker/TokenChunker to reject the payload.
|
||||
func TestMaybeDispatchImage_ReturnsJSONWithImage(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
drv := &imagePromptCaptureDriver{}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "img-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
setups := defaultSetups()
|
||||
res, dispatched, err := maybeDispatchImage(
|
||||
context.Background(),
|
||||
utility.FileTypeVISUAL,
|
||||
"test.png",
|
||||
[]byte("not-a-real-image"),
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
setups,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchImage: %v", err)
|
||||
}
|
||||
if !dispatched {
|
||||
t.Fatalf("expected dispatched=true")
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1", len(res.JSON))
|
||||
}
|
||||
item := res.JSON[0]
|
||||
if got, _ := item["doc_type_kwd"].(string); got != "image" {
|
||||
t.Errorf("doc_type_kwd = %q, want \"image\"", got)
|
||||
}
|
||||
img, _ := item["image"].(string)
|
||||
if !strings.HasPrefix(img, "data:") || !strings.Contains(img, ";base64,") {
|
||||
t.Errorf("image = %q, want a data URI (data:<mime>;base64,<b64>)", img)
|
||||
}
|
||||
if txt, _ := item["text"].(string); txt == "" {
|
||||
t.Errorf("text field empty; want non-empty combined OCR+VLM text")
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchImage_HardcodesJSONOutput verifies the image family
|
||||
// always emits json regardless of setup["output_format"]. Python
|
||||
// rag/app/picture.py:chunk() has no output_format concept — it always
|
||||
// returns a structured doc. Honoring a "text" override produced a bare
|
||||
// Text payload that lost the image attachment and set doc_type to "text".
|
||||
func TestMaybeDispatchImage_HardcodesJSONOutput(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
drv := &imagePromptCaptureDriver{}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "img-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
setups := defaultSetups()
|
||||
setups["image"]["output_format"] = "text" // legacy/override; must be ignored
|
||||
res, _, err := maybeDispatchImage(
|
||||
context.Background(),
|
||||
utility.FileTypeVISUAL,
|
||||
"test.png",
|
||||
[]byte("not-a-real-image"),
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
setups,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchImage: %v", err)
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want json (image family must ignore output_format override)", res.OutputFormat)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1 even when setup says text", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
// audioTranscribeDriver is a mock ModelDriver whose TranscribeAudio returns a
|
||||
// fixed transcription, so maybeDispatchAudio can be exercised without a real
|
||||
// ASR provider.
|
||||
type audioTranscribeDriver struct {
|
||||
modelModule.ModelDriver
|
||||
transcription string
|
||||
}
|
||||
|
||||
func (d *audioTranscribeDriver) TranscribeAudio(ctx context.Context, _ *string, _ *string, _ *modelModule.APIConfig, _ *modelModule.ASRConfig, _ *common.ModelUsage) (*modelModule.ASRResponse, error) {
|
||||
return &modelModule.ASRResponse{Text: d.transcription}, nil
|
||||
}
|
||||
|
||||
// TestMaybeDispatchAudio_JSONCarriesTranscription pins diff 2.11: when the
|
||||
// audio family's output_format is "json", the ASR transcription must be
|
||||
// carried in the JSON items (not only in the Text field). Before the fix the
|
||||
// branch returned Text only with an empty JSON slice, and the Invoke switch
|
||||
// silently dropped the transcription because it has no "json" branch.
|
||||
func TestMaybeDispatchAudio_JSONCarriesTranscription(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
const want = "hello world"
|
||||
drv := &audioTranscribeDriver{transcription: want}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "asr-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
setups := defaultSetups()
|
||||
setups["audio"]["output_format"] = "json"
|
||||
|
||||
res, dispatched, err := maybeDispatchAudio(
|
||||
context.Background(),
|
||||
utility.FileTypeAURAL,
|
||||
"test.mp3",
|
||||
[]byte("fake-audio"),
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
setups,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchAudio: %v", err)
|
||||
}
|
||||
if !dispatched {
|
||||
t.Fatalf("expected dispatched=true for AURAL file")
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1 (transcription must be carried as a JSON item)", len(res.JSON))
|
||||
}
|
||||
if got, _ := res.JSON[0]["text"].(string); got != want {
|
||||
t.Fatalf("JSON[0].text = %q, want %q", got, want)
|
||||
}
|
||||
if got, _ := res.JSON[0]["doc_type_kwd"].(string); got != "audio" {
|
||||
t.Fatalf("JSON[0].doc_type_kwd = %q, want audio", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchAudio_TextCarriesTranscription guards the text path: with
|
||||
// output_format "text" the transcription stays in the Text field and JSON is
|
||||
// empty (current default after aligning with Python parser.py:232).
|
||||
func TestMaybeDispatchAudio_TextCarriesTranscription(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
const want = "hello world"
|
||||
drv := &audioTranscribeDriver{transcription: want}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "asr-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
setups := defaultSetups()
|
||||
setups["audio"]["output_format"] = "text"
|
||||
|
||||
res, dispatched, err := maybeDispatchAudio(
|
||||
context.Background(),
|
||||
utility.FileTypeAURAL,
|
||||
"test.mp3",
|
||||
[]byte("fake-audio"),
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
setups,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchAudio: %v", err)
|
||||
}
|
||||
if !dispatched {
|
||||
t.Fatalf("expected dispatched=true for AURAL file")
|
||||
}
|
||||
if res.OutputFormat != "text" {
|
||||
t.Fatalf("OutputFormat = %q, want text", res.OutputFormat)
|
||||
}
|
||||
if res.Text != want {
|
||||
t.Fatalf("Text = %q, want %q", res.Text, want)
|
||||
}
|
||||
if len(res.JSON) != 0 {
|
||||
t.Fatalf("JSON len = %d, want 0 for text output", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchMarkdownVision_EnhancesTables pins diff 2.5: markdown
|
||||
// vision enhancement must also process items whose doc_type_kwd is "table"
|
||||
// (Python checks {"image","table"} in parser/utils.py:181), not only "image".
|
||||
// Before the fix the table item was skipped and never sent to the VLM.
|
||||
func TestMaybeDispatchMarkdownVision_EnhancesTables(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
drv := &imagePromptCaptureDriver{}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "img-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
dispatched := parserDispatchResult{
|
||||
OutputFormat: "json",
|
||||
JSON: []map[string]any{
|
||||
{"doc_type_kwd": "table", "image": "base64table", "text": ""},
|
||||
},
|
||||
}
|
||||
|
||||
res, handled, err := maybeDispatchMarkdownVision(
|
||||
context.Background(),
|
||||
utility.FileTypeMarkdown,
|
||||
dispatched,
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchMarkdownVision: %v", err)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatalf("expected handled=true for markdown with a table image")
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1", len(res.JSON))
|
||||
}
|
||||
// The table item must have been sent to the VLM and its description appended.
|
||||
if got, _ := res.JSON[0]["text"].(string); got != "captured" {
|
||||
t.Fatalf("table item text = %q, want %q (table items must be vision-enhanced)", got, "captured")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultEmailOutputFormatIsJSON pins diff 2.2: the email family default
|
||||
// output_format must be "json" (matching Python parser.py:212), not "text".
|
||||
// With "text" the structured email fields (from/to/subject/attachments/...) are
|
||||
// flattened into a blob and lost downstream.
|
||||
func TestDefaultEmailOutputFormatIsJSON(t *testing.T) {
|
||||
got, _ := defaultSetups()["email"]["output_format"].(string)
|
||||
if got != "json" {
|
||||
t.Fatalf("email default output_format = %q, want json", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -302,7 +302,7 @@ func defaultSetups() map[string]schema.ParserSetup {
|
||||
"from", "to", "cc", "bcc", "date", "subject",
|
||||
"body", "attachments", "metadata",
|
||||
},
|
||||
"output_format": "text",
|
||||
"output_format": "json",
|
||||
},
|
||||
"audio": {
|
||||
"suffix": []string{
|
||||
@@ -310,7 +310,7 @@ func defaultSetups() map[string]schema.ParserSetup {
|
||||
"aiff", "au", "midi", "wma", "realaudio", "vqf",
|
||||
"oggvorbis", "ape",
|
||||
},
|
||||
"output_format": "json",
|
||||
"output_format": "text",
|
||||
},
|
||||
"video": {
|
||||
"suffix": []string{"mp4", "avi", "mkv"},
|
||||
@@ -478,9 +478,10 @@ func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (ma
|
||||
dispatched = dispatchParse(ctx, fileTypeExt, filename, binary, c.Setups)
|
||||
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
|
||||
|
||||
// DOCX vision figure enhancement: enrich the markdown
|
||||
// with LLM-generated descriptions of embedded images.
|
||||
// Mirrors Python's vision_figure_parser_docx_wrapper_naive.
|
||||
// DOCX vision figure enhancement: on the JSON output path,
|
||||
// append vision-model descriptions to embedded image items
|
||||
// (doc_type_kwd "image"). Mirrors Python's
|
||||
// enhance_media_sections_with_vision in parser.py:_doc.
|
||||
dispatched, _, _ = maybeDispatchDOCXVision(ctx, fileTypeExt, dispatched, inputs, c.Setups)
|
||||
|
||||
// Markdown vision figure enhancement: enrich parsed
|
||||
|
||||
@@ -358,16 +358,27 @@ func resolvePDFVisionModelID(setup schema.ParserSetup) (string, bool) {
|
||||
return "", false
|
||||
}
|
||||
|
||||
// isNamedPDFParseMethod reports whether raw is a recognized named PDF
|
||||
// parse method (as opposed to a CustomVLM model name). Its membership set
|
||||
// MUST stay aligned with the PDF whitelist enforced by
|
||||
// (*ParserComponent).Check() (parser.go:200-203):
|
||||
//
|
||||
// deepdoc, plain_text, mineru, docling,
|
||||
// opendataloader, tcadp parser, paddleocr, somark
|
||||
//
|
||||
// A parse_method that Check() rejects must not be treated as a named method
|
||||
// here, otherwise it silently falls through to the CustomVLM vision path
|
||||
// instead of failing fast at construction.
|
||||
//
|
||||
// Note: "@"-suffixed spellings such as "foo@mineru" are layout_recognizer
|
||||
// selectors, not parse_method values. Check() rejects them as parse_method,
|
||||
// and the MinerU layout branch is resolved from the layout_recognizer field
|
||||
// separately (pdf_vision_dispatch.go:62-68), so they must NOT be recognized
|
||||
// here.
|
||||
func isNamedPDFParseMethod(raw string) bool {
|
||||
method := strings.ToLower(strings.TrimSpace(raw))
|
||||
switch {
|
||||
case strings.HasSuffix(method, "@paddleocr"),
|
||||
strings.HasSuffix(method, "@somark"),
|
||||
strings.HasSuffix(method, "@opendataloader"):
|
||||
return true
|
||||
}
|
||||
switch method {
|
||||
case "deepdoc", "mineru", "plain_text", "plain text", "plaintext", "paddleocr", "docling", "opendataloader", "somark", "tcadp", "tcadp parser":
|
||||
case "deepdoc", "plain_text", "mineru", "docling", "opendataloader", "tcadp parser", "paddleocr", "somark":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
|
||||
86
internal/ingestion/component/pdf_vision_dispatch_test.go
Normal file
86
internal/ingestion/component/pdf_vision_dispatch_test.go
Normal file
@@ -0,0 +1,86 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package component
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestIsNamedPDFParseMethodWhitelistAligned verifies that the runtime
|
||||
// "named parse_method" classifier agrees with (*ParserComponent).Check()'s
|
||||
// PDF whitelist (parser.go:200-203):
|
||||
//
|
||||
// deepdoc, plain_text, mineru, docling,
|
||||
// opendataloader, tcadp parser, paddleocr, somark
|
||||
//
|
||||
// Diff 2.10: a parse_method that Check() rejects must NOT be treated as a
|
||||
// recognized named method by isNamedPDFParseMethod — otherwise it silently
|
||||
// falls through to the CustomVLM vision path instead of failing fast at
|
||||
// construction (and Python would have rejected it outright).
|
||||
func TestIsNamedPDFParseMethodWhitelistAligned(t *testing.T) {
|
||||
// Values that MUST be recognized (subset of the Check() whitelist,
|
||||
// case-insensitive).
|
||||
named := []string{
|
||||
"deepdoc", "plain_text", "mineru", "docling",
|
||||
"opendataloader", "tcadp parser", "paddleocr", "somark",
|
||||
"DeepDoc", "PLAIN_TEXT", "MinerU", "DocLing",
|
||||
"OpenDataLoader", "TCADP Parser", "PaddleOCR", "SoMark",
|
||||
}
|
||||
for _, v := range named {
|
||||
if !isNamedPDFParseMethod(v) {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = false, want true (in Check() whitelist)", v)
|
||||
}
|
||||
}
|
||||
|
||||
// Values that MUST NOT be recognized. These either duplicate the
|
||||
// whitelist with non-canonical spelling ("plain text"/"plaintext")
|
||||
// or are bare-family abbreviations ("tcadp") that Check() does not
|
||||
// accept, so they should be funneled to the CustomVLM path (or fail
|
||||
// construction) rather than masquerading as a named method.
|
||||
notNamed := []string{
|
||||
"plain text", "plaintext", "tcadp",
|
||||
"CustomVLM", "some_vlm", "gpt-4o",
|
||||
"", " ",
|
||||
}
|
||||
for _, v := range notNamed {
|
||||
if isNamedPDFParseMethod(v) {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = true, want false (not in Check() whitelist)", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsNamedPDFParseMethodLayoutSuffixes verifies that "@"-suffixed
|
||||
// layout_recognizer spellings are NOT treated as named parse methods. They
|
||||
// are layout_recognizer selectors (resolved separately at
|
||||
// pdf_vision_dispatch.go:62-68), and Check() rejects them as parse_method,
|
||||
// so they must fall through to the CustomVLM/VLM path — consistent with the
|
||||
// (*ParserComponent).Check() whitelist (parser.go:200-203).
|
||||
func TestIsNamedPDFParseMethodLayoutSuffixes(t *testing.T) {
|
||||
suffixed := []string{
|
||||
"foo@mineru", "@mineru",
|
||||
"foo@paddleocr", "@paddleocr",
|
||||
"foo@somark", "@somark",
|
||||
"foo@opendataloader", "@opendataloader",
|
||||
}
|
||||
for _, v := range suffixed {
|
||||
if isNamedPDFParseMethod(v) {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = true, want false (layout_recognizer selector, not a named parse_method)", v)
|
||||
}
|
||||
}
|
||||
|
||||
// An unknown suffix is also not a named method.
|
||||
if isNamedPDFParseMethod("foo@unknown") {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = true, want false", "foo@unknown")
|
||||
}
|
||||
}
|
||||
@@ -122,6 +122,7 @@ type ChunkDoc struct {
|
||||
ContentSmLtks string `json:"content_sm_ltks,omitempty"`
|
||||
TagKwd []string `json:"tag_kwd,omitempty"`
|
||||
PageNumber *int `json:"page_number,omitempty"`
|
||||
TopInt []int `json:"top_int,omitempty"`
|
||||
PDFPositions json.RawMessage `json:"_pdf_positions,omitempty"`
|
||||
Positions json.RawMessage `json:"positions,omitempty"`
|
||||
Extra map[string]json.RawMessage `json:"-"`
|
||||
@@ -142,7 +143,7 @@ func (d *ChunkDoc) UnmarshalJSON(data []byte) error {
|
||||
"ck_type", "tk_nums", "layout", "layout_type", "layoutno", "image",
|
||||
"context_above", "context_below", "questions", "keywords", "summary",
|
||||
"chunk_order_int", "title_tks", "title_sm_tks", "content_ltks",
|
||||
"content_sm_ltks", "tag_kwd", "page_number", "_pdf_positions", "positions",
|
||||
"content_sm_ltks", "tag_kwd", "page_number", "top_int", "_pdf_positions", "positions",
|
||||
} {
|
||||
delete(raw, key)
|
||||
}
|
||||
|
||||
@@ -434,7 +434,11 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em
|
||||
if trimmedName == "" {
|
||||
log.Printf("Tokenizer: empty name provided from upstream, embedding will skip title weighting")
|
||||
} else {
|
||||
titleResults, err := encodeWithTimeout(ctx, embedder, []string{trimmedName})
|
||||
// Encode the raw name (no TrimSpace) to mirror Python
|
||||
// tokenizer.py:95 which passes name verbatim to embedding. The
|
||||
// empty-name guard above still uses TrimSpace, matching Python's
|
||||
// `.strip()==""` check at tokenizer.py:200.
|
||||
titleResults, err := encodeWithTimeout(ctx, embedder, []string{name})
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("Tokenizer: encode title: %w", err)
|
||||
}
|
||||
|
||||
@@ -591,6 +591,32 @@ func TestTokenizerComponent_Embedding_EmptyNameWarnsAndUsesContentVector(t *test
|
||||
}
|
||||
}
|
||||
|
||||
// Python tokenizer.py:95 passes the raw name to embedding without .strip();
|
||||
// Go must match — the title embedding must receive the original name, not a
|
||||
// TrimSpace'd copy. The empty-name guard still uses TrimSpace (mirroring
|
||||
// Python's `.strip()==""` check at tokenizer.py:200), but the value encoded
|
||||
// is the raw name.
|
||||
func TestTokenizerComponent_Embedding_UsesRawNameNotTrimmed(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
c, stub := withStubEmbedder(t, 2)
|
||||
|
||||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": " report.pdf ",
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "alpha"}},
|
||||
}); err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if len(stub.callInputs) < 1 {
|
||||
t.Fatalf("callInputs len = %d, want >= 1", len(stub.callInputs))
|
||||
}
|
||||
// First call is the title embedding; it must receive the raw name with
|
||||
// surrounding whitespace preserved, matching Python.
|
||||
if got := stub.callInputs[0][0]; got != " report.pdf " {
|
||||
t.Fatalf("title embedding input = %q, want %q (raw, not trimmed)", got, " report.pdf ")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerComponent_Embedding_TruncatesByMaxTokensMinus10(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
c, stub := withStubEmbedder(t, 2)
|
||||
|
||||
@@ -22,6 +22,7 @@ package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
@@ -322,3 +323,32 @@ func TestValidateTokenizerOutputs_SymbolOnlyContentLtksIsEmptyFails(t *testing.T
|
||||
t.Fatalf("err = %v, want missing full_text tokens", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChunkDocsToMaps_PreservesPDFPositions is the pool-free unit test for
|
||||
// Tokenizer-(T)1: the tokenizer emits chunks via schema.ChunkDocsToMaps
|
||||
// (ChunkDoc.ToMap), which must carry the raw `positions` / `_pdf_positions`
|
||||
// through untouched so the downstream executor stage
|
||||
// (internal/ingestion/task processChunkPositions → AddPositions) can convert
|
||||
// them into position_int / page_num_int / top_int exactly once. This does NOT
|
||||
// require the C++ analyzer pool, so it runs under plain `go test`.
|
||||
func TestChunkDocsToMaps_PreservesPDFPositions(t *testing.T) {
|
||||
pos := json.RawMessage(`[[1,10,20,30,40],[2,15,25,35,45]]`)
|
||||
chunks := []schema.ChunkDoc{
|
||||
{Text: "PDF paragraph", DocType: "text", CKType: "text",
|
||||
Positions: pos, PDFPositions: pos},
|
||||
}
|
||||
maps := schema.ChunkDocsToMaps(chunks)
|
||||
|
||||
got, ok := maps[0]["positions"].([][]float64)
|
||||
if !ok || len(got) != 2 {
|
||||
t.Fatalf("positions not preserved through tokenizer output mapping: %#v", maps[0]["positions"])
|
||||
}
|
||||
if _, ok := maps[0]["_pdf_positions"].([][]float64); !ok {
|
||||
t.Errorf("_pdf_positions not preserved through tokenizer output mapping: %#v", maps[0]["_pdf_positions"])
|
||||
}
|
||||
// Sanity: page numbers are still raw 1-indexed, i.e. not yet converted
|
||||
// to page_num_int (the executor owns that step).
|
||||
if int(got[0][0]) != 1 {
|
||||
t.Errorf("positions page already converted; want raw 1-indexed page 1, got %v", got[0][0])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,7 +50,7 @@
|
||||
}
|
||||
},
|
||||
"image": {
|
||||
"output_format": "text",
|
||||
"output_format": "json",
|
||||
"parse_method": "ocr",
|
||||
"preprocess": [
|
||||
"main_content"
|
||||
@@ -231,7 +231,7 @@
|
||||
"setups": [
|
||||
{
|
||||
"fileFormat": "image",
|
||||
"output_format": "text",
|
||||
"output_format": "json",
|
||||
"parse_method": "ocr",
|
||||
"preprocess": [
|
||||
"main_content"
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -15,19 +15,17 @@ import (
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/engine"
|
||||
enginetypes "ragflow/internal/engine/types"
|
||||
"ragflow/internal/entity"
|
||||
_ "ragflow/internal/ingestion/component"
|
||||
componentpkg "ragflow/internal/ingestion/component"
|
||||
_ "ragflow/internal/ingestion/component/chunker"
|
||||
pipelinepkg "ragflow/internal/ingestion/pipeline"
|
||||
"ragflow/internal/server"
|
||||
"ragflow/internal/storage"
|
||||
"ragflow/internal/tokenizer"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
|
||||
"go.uber.org/zap"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/gorm"
|
||||
gormlogger "gorm.io/gorm/logger"
|
||||
)
|
||||
@@ -35,30 +33,18 @@ import (
|
||||
func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
|
||||
cfg := mustLoadTaskRealIntegrationConfig(t)
|
||||
realDB := mustOpenTaskRealMySQL(t, cfg)
|
||||
if err := realDB.AutoMigrate(
|
||||
&entity.Tenant{},
|
||||
&entity.Knowledgebase{},
|
||||
&entity.Document{},
|
||||
&entity.File{},
|
||||
&entity.File2Document{},
|
||||
&entity.UserCanvas{},
|
||||
); err != nil {
|
||||
t.Fatalf("auto-migrate real mysql tables: %v", err)
|
||||
}
|
||||
|
||||
realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio)
|
||||
if err != nil {
|
||||
t.Fatalf("connect real minio: %v", err)
|
||||
}
|
||||
|
||||
mustLoadTaskTestConfig(t)
|
||||
origDB := dao.DB
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
realDB := mustOpenTaskTestDB(t)
|
||||
dao.DB = realDB
|
||||
storage.GetStorageFactory().SetStorage(realStorage)
|
||||
t.Cleanup(func() {
|
||||
dao.DB = origDB
|
||||
})
|
||||
|
||||
realStorage := storage.NewMemoryStorage()
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
storage.GetStorageFactory().SetStorage(realStorage)
|
||||
t.Cleanup(func() {
|
||||
storage.GetStorageFactory().SetStorage(origStorage)
|
||||
})
|
||||
|
||||
@@ -148,46 +134,25 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *testing.T) {
|
||||
func TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
|
||||
cfg := mustLoadTaskRealIntegrationConfig(t)
|
||||
realDB := mustOpenTaskRealMySQL(t, cfg)
|
||||
if err := realDB.AutoMigrate(
|
||||
&entity.Tenant{},
|
||||
&entity.Knowledgebase{},
|
||||
&entity.Document{},
|
||||
&entity.File{},
|
||||
&entity.File2Document{},
|
||||
&entity.UserCanvas{},
|
||||
); err != nil {
|
||||
t.Fatalf("auto-migrate real mysql tables: %v", err)
|
||||
}
|
||||
|
||||
realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio)
|
||||
if err != nil {
|
||||
t.Fatalf("connect real minio: %v", err)
|
||||
}
|
||||
if err := engine.Init(&cfg.DocEngine); err != nil {
|
||||
t.Fatalf("init real doc engine: %v", err)
|
||||
}
|
||||
if engine.Get() == nil {
|
||||
t.Fatal("doc engine is nil after init")
|
||||
}
|
||||
if engine.GetEngineType() != engine.EngineElasticsearch {
|
||||
t.Fatalf("doc engine type = %s, want %s", engine.GetEngineType(), engine.EngineElasticsearch)
|
||||
}
|
||||
|
||||
// Loads service config (server.Init side effect) without requiring any
|
||||
// external MySQL/MinIO/ES. The pipeline runs against an in-memory sqlite
|
||||
// DB and an in-memory storage backend; chunks are captured via WithInsertFunc.
|
||||
mustLoadTaskTestConfig(t)
|
||||
origDB := dao.DB
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
origDocResolver := componentpkg.ResolveDocumentStorageOverride
|
||||
realDB := mustOpenTaskTestDB(t)
|
||||
dao.DB = realDB
|
||||
storage.GetStorageFactory().SetStorage(realStorage)
|
||||
componentpkg.ResolveDocumentStorageOverride = nil
|
||||
t.Cleanup(func() {
|
||||
dao.DB = origDB
|
||||
})
|
||||
|
||||
realStorage := storage.NewMemoryStorage()
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
storage.GetStorageFactory().SetStorage(realStorage)
|
||||
t.Cleanup(func() {
|
||||
storage.GetStorageFactory().SetStorage(origStorage)
|
||||
componentpkg.ResolveDocumentStorageOverride = origDocResolver
|
||||
})
|
||||
|
||||
templatePath := filepath.Join(taskRepoRoot(t), "internal", "ingestion", "pipeline", "template", "ingestion_pipeline_general.json")
|
||||
@@ -217,7 +182,6 @@ func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *tes
|
||||
bucket := taskS3SafeBucketName(kbID)
|
||||
docName := "01_english_simple.pdf"
|
||||
objectPath := fmt.Sprintf("integration/task/%s/%s", docID, docName)
|
||||
baseName := fmt.Sprintf("ragflow_%s", tenantID)
|
||||
|
||||
mustSeedTaskRealPipelineDocumentBytes(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, ".pdf", "pdf", pdfBytes)
|
||||
if err := realDB.Model(&entity.Document{}).Where("id = ?", docID).Update("pipeline_id", canvasID).Error; err != nil {
|
||||
@@ -233,14 +197,13 @@ func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *tes
|
||||
t.Fatalf("create user canvas: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = engine.Get().DropChunkStore(context.Background(), baseName, kbID)
|
||||
_ = realDB.Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error
|
||||
cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath)
|
||||
})
|
||||
|
||||
taskCtx := &TaskContext{
|
||||
IngestionTask: &entity.IngestionTask{
|
||||
ID: "task-real-pdf-es-1",
|
||||
ID: "task-real-pdf-1",
|
||||
DocumentID: docID,
|
||||
DatasetID: kbID,
|
||||
},
|
||||
@@ -258,68 +221,66 @@ func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *tes
|
||||
Tenant: entity.Tenant{ID: tenantID},
|
||||
}
|
||||
|
||||
svc := mustNewPipelineExecutor(t, taskCtx, canvasID, 0)
|
||||
var inserted [][]map[string]any
|
||||
svc := mustNewPipelineExecutor(t, taskCtx, canvasID, 0).
|
||||
WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) {
|
||||
inserted = append(inserted, deepCopyTaskChunks(chunks))
|
||||
return nil, nil
|
||||
}).
|
||||
WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil })
|
||||
|
||||
if _, err := svc.Execute(context.Background()); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
|
||||
result, err := engine.Get().Search(context.Background(), &enginetypes.SearchRequest{
|
||||
IndexNames: []string{baseName},
|
||||
KbIDs: []string{kbID},
|
||||
Limit: 20,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("search indexed chunks: %v", err)
|
||||
if len(inserted) == 0 {
|
||||
t.Fatal("no chunks inserted")
|
||||
}
|
||||
if result == nil {
|
||||
t.Fatal("search result is nil")
|
||||
var chunks []map[string]any
|
||||
for _, batch := range inserted {
|
||||
chunks = append(chunks, batch...)
|
||||
}
|
||||
if len(result.Chunks) == 0 {
|
||||
t.Fatal("expected indexed chunks in Elasticsearch, got 0")
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("inserted 0 chunks")
|
||||
}
|
||||
for i, chunk := range result.Chunks {
|
||||
sawImage := false
|
||||
for i, chunk := range chunks {
|
||||
if got := chunk["doc_id"]; got != docID {
|
||||
t.Fatalf("result chunk[%d].doc_id = %v, want %q", i, got, docID)
|
||||
t.Fatalf("chunk[%d].doc_id = %v, want %q", i, got, docID)
|
||||
}
|
||||
if got := chunk["kb_id"]; got != kbID {
|
||||
t.Fatalf("result chunk[%d].kb_id = %v, want %q", i, got, kbID)
|
||||
if !taskChunkFieldEqualsStr(chunk["kb_id"], kbID) {
|
||||
t.Fatalf("chunk[%d].kb_id = %v, want %q", i, chunk["kb_id"], kbID)
|
||||
}
|
||||
if got := chunk["docnm_kwd"]; got != docName {
|
||||
t.Fatalf("result chunk[%d].docnm_kwd = %v, want %q", i, got, docName)
|
||||
t.Fatalf("chunk[%d].docnm_kwd = %v, want %q", i, got, docName)
|
||||
}
|
||||
if got := chunk["content_with_weight"]; got == nil || got == "" {
|
||||
t.Fatalf("result chunk[%d].content_with_weight = %v, want non-empty", i, got)
|
||||
t.Fatalf("chunk[%d].content_with_weight = %v, want non-empty", i, got)
|
||||
}
|
||||
if img, ok := chunk["img_id"]; ok && img != nil && img != "" {
|
||||
sawImage = true
|
||||
}
|
||||
}
|
||||
if !sawImage {
|
||||
t.Fatal("expected at least one chunk with img_id (pdf preview image uploaded to storage)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunPipeline_RealPipelineOutput_ProducesIndexFields(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
|
||||
cfg := mustLoadTaskRealIntegrationConfig(t)
|
||||
realDB := mustOpenTaskRealMySQL(t, cfg)
|
||||
if err := realDB.AutoMigrate(
|
||||
&entity.Tenant{},
|
||||
&entity.Knowledgebase{},
|
||||
&entity.Document{},
|
||||
&entity.File{},
|
||||
&entity.File2Document{},
|
||||
); err != nil {
|
||||
t.Fatalf("auto-migrate real mysql tables: %v", err)
|
||||
}
|
||||
|
||||
realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio)
|
||||
if err != nil {
|
||||
t.Fatalf("connect real minio: %v", err)
|
||||
}
|
||||
|
||||
mustLoadTaskTestConfig(t)
|
||||
origDB := dao.DB
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
realDB := mustOpenTaskTestDB(t)
|
||||
dao.DB = realDB
|
||||
storage.GetStorageFactory().SetStorage(realStorage)
|
||||
t.Cleanup(func() {
|
||||
dao.DB = origDB
|
||||
})
|
||||
|
||||
realStorage := storage.NewMemoryStorage()
|
||||
origStorage := storage.GetStorageFactory().GetStorage()
|
||||
storage.GetStorageFactory().SetStorage(realStorage)
|
||||
t.Cleanup(func() {
|
||||
storage.GetStorageFactory().SetStorage(origStorage)
|
||||
})
|
||||
|
||||
@@ -426,7 +387,7 @@ func taskRepoRoot(t *testing.T) string {
|
||||
return filepath.Clean(filepath.Join(wd, "..", "..", ".."))
|
||||
}
|
||||
|
||||
func mustLoadTaskRealIntegrationConfig(t *testing.T) *server.Config {
|
||||
func mustLoadTaskTestConfig(t *testing.T) *server.Config {
|
||||
t.Helper()
|
||||
if err := common.Init("info", common.FileOutput{}, ""); err != nil {
|
||||
t.Fatalf("init common logger: %v", err)
|
||||
@@ -437,27 +398,42 @@ func mustLoadTaskRealIntegrationConfig(t *testing.T) *server.Config {
|
||||
t.Fatalf("init service config from %s: %v", configPath, err)
|
||||
}
|
||||
cfg := server.GetConfig()
|
||||
if cfg == nil || cfg.Database.Host == "" || cfg.StorageEngine.Minio == nil || cfg.StorageEngine.Minio.Host == "" {
|
||||
t.Fatal("real integration config is incomplete")
|
||||
if cfg == nil {
|
||||
t.Fatal("task test config is nil after server.Init")
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
func mustOpenTaskRealMySQL(t *testing.T, cfg *server.Config) *gorm.DB {
|
||||
// mustOpenTaskTestDB opens an isolated in-memory sqlite database and migrates
|
||||
// the tables the real-pipeline contract tests seed and read. It does not touch
|
||||
// the filesystem or any external MySQL. The connection pool is pinned to a
|
||||
// single connection (MaxOpenConns(1)) so the in-memory database is shared
|
||||
// across all statements — without that, gorm's default pool would hand each
|
||||
// statement a separate connection, and ":memory:" is per-connection (so seeds
|
||||
// on one connection would be invisible to reads on another).
|
||||
func mustOpenTaskTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local",
|
||||
cfg.Database.Username,
|
||||
cfg.Database.Password,
|
||||
cfg.Database.Host,
|
||||
cfg.Database.Port,
|
||||
cfg.Database.Database,
|
||||
cfg.Database.Charset,
|
||||
)
|
||||
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
|
||||
Logger: gormlogger.Default.LogMode(gormlogger.Silent),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("connect real mysql: %v", err)
|
||||
t.Fatalf("open in-memory sqlite db: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(
|
||||
&entity.Tenant{},
|
||||
&entity.Knowledgebase{},
|
||||
&entity.Document{},
|
||||
&entity.File{},
|
||||
&entity.File2Document{},
|
||||
&entity.UserCanvas{},
|
||||
&entity.PipelineOperationLog{},
|
||||
); err != nil {
|
||||
t.Fatalf("auto-migrate sqlite tables: %v", err)
|
||||
}
|
||||
if sqlDB, err := db.DB(); err != nil {
|
||||
t.Fatalf("get sql.DB from gorm: %v", err)
|
||||
} else {
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
}
|
||||
return db
|
||||
}
|
||||
@@ -731,3 +707,18 @@ func taskS3SafeBucketName(s string) string {
|
||||
s = strings.ReplaceAll(s, "_", "-")
|
||||
return s
|
||||
}
|
||||
|
||||
// taskChunkFieldEqualsStr compares a chunk field to a plain string, tolerating
|
||||
// the slice form used internally (e.g. kb_id is []string{kbID} and survives a
|
||||
// JSON round-trip as []any{kbID}).
|
||||
func taskChunkFieldEqualsStr(v any, want string) bool {
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
return val == want
|
||||
case []string:
|
||||
return len(val) == 1 && val[0] == want
|
||||
case []any:
|
||||
return len(val) == 1 && fmt.Sprint(val[0]) == want
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -27,14 +27,15 @@ import (
|
||||
// task_executor.run_dataflow:879 re.split(r"[,,;;、\r\n]+", keywords).
|
||||
var keywordsSplitRE = regexp.MustCompile(`[,,;;、\r\n]+`)
|
||||
|
||||
// nonEmpty drops empty strings from parts and returns nil if none remain. It is
|
||||
// the shared tail of SplitKeywords and SplitQuestions: split by whatever
|
||||
// delimiter, then prune empties and collapse an all-empty result to nil so a
|
||||
// _kwd array is absent (nil) rather than [""].
|
||||
// nonEmpty drops empty and whitespace-only strings from parts and returns nil
|
||||
// if none remain. It is the shared tail of SplitKeywords and SplitQuestions:
|
||||
// split by whatever delimiter, then prune blanks and collapse an all-blank
|
||||
// result to nil so a _kwd array is absent (nil) rather than [""].
|
||||
// Mirrors Python task_executor's `if k.strip()` filter.
|
||||
func nonEmpty(parts []string) []string {
|
||||
out := make([]string, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if p != "" {
|
||||
if strings.TrimSpace(p) != "" {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,10 @@
|
||||
|
||||
package utility
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// SplitKeywords - Python task_executor.run_dataflow:879
|
||||
// re.split(r"[,,;;、\r\n]+", keywords) with empty filtering.
|
||||
@@ -49,6 +52,28 @@ func TestSplitKeywords_FiltersEmptyStrings(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Python task_executor.run_dataflow filters split parts with `if k.strip()`,
|
||||
// so whitespace-only segments (a lone space between commas, a tab, etc.) must
|
||||
// be dropped — not kept as " ". Input "a, ,b" must yield ["a","b"], matching
|
||||
// Python, not ["a"," ","b"].
|
||||
func TestSplitKeywords_FiltersWhitespaceOnly(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
want []string
|
||||
}{
|
||||
{"a, ,b", []string{"a", "b"}},
|
||||
{"a,\t,b", []string{"a", "b"}},
|
||||
{" , , ", nil},
|
||||
{"kw1, ,kw2", []string{"kw1", "kw2"}},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := SplitKeywords(c.in)
|
||||
if !reflect.DeepEqual(got, c.want) {
|
||||
t.Errorf("SplitKeywords(%q) = %v, want %v", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitKeywords_Empty(t *testing.T) {
|
||||
result := SplitKeywords("")
|
||||
if len(result) != 0 {
|
||||
|
||||
Reference in New Issue
Block a user