mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 09:53:29 +08:00
## 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.
809 lines
26 KiB
Go
809 lines
26 KiB
Go
//go:build integration
|
||
// +build integration
|
||
|
||
//
|
||
// 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 (
|
||
"bytes"
|
||
"context"
|
||
"errors"
|
||
"log"
|
||
"math"
|
||
"reflect"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"ragflow/internal/ingestion/component/schema"
|
||
"ragflow/internal/tokenizer"
|
||
)
|
||
|
||
func requireTokenizerPool(t *testing.T) {
|
||
t.Helper()
|
||
if err := tokenizer.Init(&tokenizer.PoolConfig{
|
||
DictPath: "/usr/share/infinity/resource",
|
||
MinSize: 1,
|
||
MaxSize: 2,
|
||
IdleTimeout: 30 * time.Second,
|
||
AcquireTimeout: 5 * time.Second,
|
||
}); err != nil {
|
||
t.Skipf("tokenizer pool unavailable: %v", err)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_HappyPath drives three chunks
|
||
// through both full_text tokenization and embedding. Verifies that
|
||
// every chunk gains `content_ltks`, `content_sm_ltks`, and a
|
||
// `q_<n>_vec` vector keyed by the embedder's vector dimension.
|
||
func TestTokenizerComponent_Invoke_HappyPath(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
const dim = 4
|
||
c, stub := withStubEmbedder(t, dim)
|
||
|
||
_ = stub
|
||
var err error
|
||
if err != nil {
|
||
t.Fatalf("NewTokenizerComponent: %v", err)
|
||
}
|
||
|
||
chunks := []map[string]any{
|
||
{"text": "alpha chunk text"},
|
||
{"text": "bravo chunk text"},
|
||
{"text": "charlie chunk text"},
|
||
}
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"tenant_id": "t1",
|
||
"model_id": "embd-1",
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": chunks,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
|
||
gotChunks, ok := out["chunks"].([]map[string]any)
|
||
if !ok {
|
||
t.Fatalf("chunks type = %T, want []map[string]any", out["chunks"])
|
||
}
|
||
if len(gotChunks) != 3 {
|
||
t.Fatalf("len(chunks) = %d, want 3", len(gotChunks))
|
||
}
|
||
for i, ck := range gotChunks {
|
||
if ck["content_ltks"] == nil {
|
||
t.Errorf("chunk[%d].content_ltks missing", i)
|
||
}
|
||
if ck["content_sm_ltks"] == nil {
|
||
t.Errorf("chunk[%d].content_sm_ltks missing", i)
|
||
}
|
||
if ck["title_tks"] == nil {
|
||
t.Errorf("chunk[%d].title_tks missing", i)
|
||
}
|
||
key := "q_4_vec"
|
||
v, ok := ck[key].([]float64)
|
||
if !ok {
|
||
t.Errorf("chunk[%d].%s missing or wrong type: %T", i, key, ck[key])
|
||
continue
|
||
}
|
||
if len(v) != dim {
|
||
t.Errorf("chunk[%d].%s len = %d, want %d", i, key, len(v), dim)
|
||
}
|
||
if v[0] == 0 {
|
||
t.Errorf("chunk[%d].%s[0] = %v, want non-zero", i, key, v[0])
|
||
}
|
||
}
|
||
if out["output_format"] != "chunks" {
|
||
t.Errorf("output_format = %v, want chunks", out["output_format"])
|
||
}
|
||
if out["embedding_token_consumption"] == nil {
|
||
t.Error("embedding_token_consumption missing")
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_Unicode asserts CJK input
|
||
// produces finite, non-negative token counts (plan §8 Q2: Go
|
||
// `NumTokensFromString` over-counts CJK on tiktoken-init failure;
|
||
// Python returns 0 — both are valid as long as the count is
|
||
// finite).
|
||
func TestTokenizerComponent_Invoke_Unicode(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 4)
|
||
_ = stub
|
||
|
||
inputs := []string{
|
||
"中文测试文本",
|
||
"こんにちは世界",
|
||
"한국어 텍스트",
|
||
"English mixed 中文 français 日本語",
|
||
}
|
||
chunks := make([]map[string]any, len(inputs))
|
||
for i, txt := range inputs {
|
||
chunks[i] = map[string]any{"text": txt}
|
||
}
|
||
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": chunks,
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
gotChunks, _ := out["chunks"].([]map[string]any)
|
||
if len(gotChunks) != len(inputs) {
|
||
t.Fatalf("chunks len = %d, want %d", len(gotChunks), len(inputs))
|
||
}
|
||
for i, ck := range gotChunks {
|
||
// Direct call to verify the count contract.
|
||
tokens := tokenizer.NumTokensFromString(inputs[i])
|
||
if tokens < 0 {
|
||
t.Errorf("chunk[%d] token count negative: %d", i, tokens)
|
||
}
|
||
if ck["content_ltks"] == nil {
|
||
t.Errorf("chunk[%d].content_ltks missing", i)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Invoke_TextPayload(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
_, _ = withStubEmbedder(t, 4)
|
||
c, _ := NewTokenizerComponent(map[string]any{
|
||
"search_method": []any{"full_text"},
|
||
})
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "note.txt",
|
||
"output_format": "text",
|
||
"text": "plain payload",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
if len(got) != 1 {
|
||
t.Fatalf("chunks len = %d, want 1", len(got))
|
||
}
|
||
if got[0]["text"] != "plain payload" {
|
||
t.Errorf("text = %v, want plain payload", got[0]["text"])
|
||
}
|
||
if got[0]["content_ltks"] == nil {
|
||
t.Errorf("content_ltks missing: %v", got[0])
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Invoke_JSONPayload(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
_, _ = withStubEmbedder(t, 4)
|
||
c, _ := NewTokenizerComponent(map[string]any{
|
||
"search_method": []any{"full_text"},
|
||
})
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "note.pdf",
|
||
"output_format": "json",
|
||
"json": []map[string]any{{"text": "row one"}, {"text": "row two"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
if len(got) != 2 {
|
||
t.Fatalf("chunks len = %d, want 2", len(got))
|
||
}
|
||
if got[0]["content_ltks"] == nil || got[1]["content_ltks"] == nil {
|
||
t.Errorf("content_ltks missing: %v", got)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_BatchedEmbedding asserts the
|
||
// embedding client is called once for the title plus once for the
|
||
// chunk batch when all chunks fit into a single batch.
|
||
func TestTokenizerComponent_Invoke_BatchedEmbedding(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 8)
|
||
chunks := []map[string]any{
|
||
{"text": "one"},
|
||
{"text": "two"},
|
||
{"text": "three"},
|
||
}
|
||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.txt",
|
||
"output_format": "chunks",
|
||
"chunks": chunks,
|
||
}); err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if got := stub.calls.Load(); got != 2 {
|
||
t.Errorf("embedder calls = %d, want 2 (title + single content batch)", got)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_FullTextOnly covers the
|
||
// search_method=["full_text"] branch: no embedding, no encoder
|
||
// call, but tokenized fields present.
|
||
func TestTokenizerComponent_Invoke_FullTextOnly(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
_, stub := withStubEmbedder(t, 4)
|
||
c, _ := NewTokenizerComponent(map[string]any{
|
||
"search_method": []any{"full_text"},
|
||
})
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha bravo"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if stub.calls.Load() != 0 {
|
||
t.Errorf("embedder should not be called, got %d", stub.calls.Load())
|
||
}
|
||
if out["embedding_token_consumption"] != nil {
|
||
t.Errorf("embedding_token_consumption should be absent, got %v", out["embedding_token_consumption"])
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
if len(got) == 0 || got[0]["content_ltks"] == nil {
|
||
t.Errorf("content_ltks missing: %v", got)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_KeywordSplitCJK verifies important_kwd is
|
||
// split by the full ASCII+CJK delimiter set, not just ASCII comma. A Chinese
|
||
// LLM commonly emits CJK commas/semicolons even when asked for
|
||
// "comma-separated"; ASCII-only splitting would leave keywords glued together.
|
||
func TestTokenizerComponent_Invoke_KeywordSplitCJK(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
_, stub := withStubEmbedder(t, 4)
|
||
c, _ := NewTokenizerComponent(map[string]any{
|
||
"search_method": []any{"full_text"},
|
||
})
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha", "keywords": "kw1,kw2;kw3"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if stub.calls.Load() != 0 {
|
||
t.Errorf("embedder should not be called in full_text-only mode, got %d", stub.calls.Load())
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
if len(got) != 1 {
|
||
t.Fatalf("chunks len = %d, want 1", len(got))
|
||
}
|
||
kwd, ok := got[0]["important_kwd"].([]string)
|
||
if !ok {
|
||
t.Fatalf("important_kwd should be []string, got %T", got[0]["important_kwd"])
|
||
}
|
||
if len(kwd) != 3 {
|
||
t.Errorf("important_kwd must split CJK delimiters into 3 elements, got %d: %v", len(kwd), kwd)
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Invoke_FullTextAndEmbedding(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, _ := withStubEmbedder(t, 4)
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha bravo"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
if len(got) != 1 {
|
||
t.Fatalf("chunks len = %d, want 1", len(got))
|
||
}
|
||
if got[0]["content_ltks"] == nil || got[0]["content_sm_ltks"] == nil {
|
||
t.Fatalf("full-text tokens missing: %v", got[0])
|
||
}
|
||
if got[0]["q_4_vec"] == nil {
|
||
t.Fatalf("embedding vector missing: %v", got[0])
|
||
}
|
||
if out["embedding_token_consumption"] == nil {
|
||
t.Fatal("embedding_token_consumption missing")
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_EmbedNoResolver covers the
|
||
// "embedding requested but no embedder resolver configured" branch
|
||
// (explicit resolver nil and DefaultEmbedderResolver unset) — must
|
||
// return a clear error, not panic.
|
||
func TestTokenizerComponent_Invoke_EmbedNoResolver(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, _ := NewTokenizerComponent(nil)
|
||
_, err := c.Invoke(context.Background(), map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}},
|
||
})
|
||
if err == nil {
|
||
t.Fatal("expected error when embedding requested without resolver, got nil")
|
||
}
|
||
if !strings.Contains(err.Error(), "resolver") {
|
||
t.Errorf("error should mention resolver: %v", err)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_EmbedderError covers the
|
||
// propagation of an error from the embedding driver.
|
||
func TestTokenizerComponent_Invoke_EmbedderError(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 4)
|
||
stub.err = errors.New("simulated upstream error")
|
||
|
||
_, err := c.Invoke(context.Background(), map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}},
|
||
})
|
||
if err == nil {
|
||
t.Fatal("expected error from embedder, got nil")
|
||
}
|
||
if !strings.Contains(err.Error(), "simulated upstream error") {
|
||
t.Errorf("error should chain embedder error: %v", err)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_EncoderCountMismatch covers the
|
||
// "embedder returned wrong number of vectors" defensive branch.
|
||
func TestTokenizerComponent_Invoke_EncoderCountMismatch(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
_, stub := withStubEmbedder(t, 4)
|
||
// Inject an embedder that returns the wrong number of vectors
|
||
// regardless of input.
|
||
wrong := &countMismatchedEmbedder{want: 1}
|
||
cIntf, err := NewTokenizerComponentWithResolver(nil, func(_, _, _ string) (Embedder, error) { return wrong, nil })
|
||
if err != nil {
|
||
t.Fatalf("NewTokenizerComponentWithResolver: %v", err)
|
||
}
|
||
c := cIntf.(*TokenizerComponent)
|
||
_ = stub
|
||
_, err = c.Invoke(context.Background(), map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "a"}, {"text": "b"}, {"text": "c"}},
|
||
})
|
||
if err == nil {
|
||
t.Fatal("expected error from count mismatch, got nil")
|
||
}
|
||
if !strings.Contains(err.Error(), "vectors") {
|
||
t.Errorf("error should mention vectors: %v", err)
|
||
}
|
||
}
|
||
|
||
type countMismatchedEmbedder struct{ want int }
|
||
|
||
func (c *countMismatchedEmbedder) MaxTokens() int { return 0 }
|
||
|
||
func (c *countMismatchedEmbedder) Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error) {
|
||
out := make([]EmbeddingResult, c.want)
|
||
for i := range out {
|
||
out[i] = EmbeddingResult{Vector: make([]float64, 4), TokenCount: 1}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// TestTokenizerComponent_Invoke_HonorsTimeout installs an
|
||
// embedder that blocks past a (test-shrunk) tokenizerTimeout and
|
||
// asserts the component returns context.DeadlineExceeded.
|
||
func TestTokenizerComponent_Invoke_HonorsTimeout(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
t.Setenv("COMPONENT_EXEC_TIMEOUT_TOKENIZER", "1")
|
||
|
||
c, stub := withStubEmbedder(t, 4)
|
||
stub.delay = 2 * time.Second
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
defer cancel()
|
||
|
||
_, err := c.Invoke(ctx, map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}},
|
||
})
|
||
if err == nil {
|
||
t.Fatal("expected timeout error, got nil")
|
||
}
|
||
if !errors.Is(err, context.DeadlineExceeded) {
|
||
t.Errorf("expected context.DeadlineExceeded, got %v", err)
|
||
}
|
||
}
|
||
|
||
// TestTokenizerComponent_Smoke_EndToEnd is the BLOCKER smoke test
|
||
// (plan §8 R3). Drives 1 chunk of ~1000 tokens through the real
|
||
// tokenizer and a stub embedder with no artificial latency, then
|
||
// asserts:
|
||
//
|
||
// - non-zero vector returned
|
||
// - latency well under 5s (real stub returns in <1ms; we assert
|
||
// < 5s as the §R3 ceiling)
|
||
// - no panic
|
||
//
|
||
// Documented result: this stub embedding completes in well under
|
||
// 5s (typical observed latency < 5ms on the test host). The
|
||
// production path against a real embedding API was not exercised
|
||
// in this CI sandbox; the helper `withStubEmbedder` deliberately
|
||
// avoids the network round-trip while still exercising the full
|
||
// wiring (TrackElapsed, WithTimeout, batched Encode, vector
|
||
// stamping).
|
||
|
||
func TestTokenizerComponent_Smoke_EndToEnd(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
const dim = 1024
|
||
c, _ := withStubEmbedder(t, dim)
|
||
|
||
words := make([]string, 0, 1000)
|
||
for i := 0; i < 1000; i++ {
|
||
words = append(words, "ragflow")
|
||
}
|
||
chunkText := strings.Join(words, " ")
|
||
preflightTokens := tokenizer.NumTokensFromString(chunkText)
|
||
if preflightTokens < 100 || preflightTokens > 5000 {
|
||
t.Logf("preflight token count = %d (acceptable range 100-5000)", preflightTokens)
|
||
}
|
||
|
||
start := time.Now()
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"tenant_id": "tenant-smoke",
|
||
"model_id": "embd-smoke",
|
||
"name": "smoke.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": chunkText}},
|
||
})
|
||
elapsed := time.Since(start)
|
||
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if elapsed >= 5*time.Second {
|
||
t.Errorf("elapsed %v exceeds 5s ceiling", elapsed)
|
||
}
|
||
|
||
got, ok := out["chunks"].([]map[string]any)
|
||
if !ok || len(got) != 1 {
|
||
t.Fatalf("chunks output malformed: %v", out["chunks"])
|
||
}
|
||
vec, ok := got[0]["q_1024_vec"].([]float64)
|
||
if !ok {
|
||
t.Fatalf("q_1024_vec missing or wrong type: %T", got[0]["q_1024_vec"])
|
||
}
|
||
if len(vec) != dim {
|
||
t.Errorf("vector len = %d, want %d", len(vec), dim)
|
||
}
|
||
nonZero := 0
|
||
for _, x := range vec {
|
||
if x != 0 {
|
||
nonZero++
|
||
}
|
||
}
|
||
if nonZero == 0 {
|
||
t.Error("vector is all zeros")
|
||
}
|
||
|
||
t.Logf("smoke complete: chunks=%d elapsed=%v tokens=%d vec_dim=%d",
|
||
len(got), elapsed, preflightTokens, len(vec))
|
||
}
|
||
|
||
func TestTokenizerComponent_Embedding_MergesTitleAndContentVectors(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 2)
|
||
stub.resultsByCall = []embeddingCallResult{
|
||
{vectors: [][]float64{{10, 20}}, tokenCount: 7},
|
||
{vectors: [][]float64{{1, 2}, {3, 4}}, tokenCount: 11},
|
||
}
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}, {"text": "beta"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
want0 := []float64{1.9, 3.8}
|
||
want1 := []float64{3.7, 5.6}
|
||
if !floatSliceClose(got[0]["q_2_vec"].([]float64), want0) {
|
||
t.Fatalf("chunk[0] q_2_vec = %v, want %v", got[0]["q_2_vec"], want0)
|
||
}
|
||
if !floatSliceClose(got[1]["q_2_vec"].([]float64), want1) {
|
||
t.Fatalf("chunk[1] q_2_vec = %v, want %v", got[1]["q_2_vec"], want1)
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Embedding_UsesFilenameWeight(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
cIntf, err := NewTokenizerComponentWithResolver(map[string]any{
|
||
"filename_embd_weight": 0.25,
|
||
}, func(_, _, _ string) (Embedder, error) {
|
||
stub := newStubEmbedder(2)
|
||
stub.resultsByCall = []embeddingCallResult{
|
||
{vectors: [][]float64{{8, 8}}, tokenCount: 3},
|
||
{vectors: [][]float64{{2, 2}}, tokenCount: 5},
|
||
}
|
||
return stub, nil
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewTokenizerComponentWithResolver: %v", err)
|
||
}
|
||
c := cIntf.(*TokenizerComponent)
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
want := []float64{3.5, 3.5}
|
||
if !floatSliceClose(got[0]["q_2_vec"].([]float64), want) {
|
||
t.Fatalf("q_2_vec = %v, want %v", got[0]["q_2_vec"], want)
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Embedding_EmptyNameWarnsAndUsesContentVector(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 2)
|
||
stub.resultsByCall = []embeddingCallResult{{vectors: [][]float64{{2, 4}}, tokenCount: 5}}
|
||
|
||
var buf bytes.Buffer
|
||
prevWriter := log.Writer()
|
||
prevFlags := log.Flags()
|
||
log.SetOutput(&buf)
|
||
log.SetFlags(0)
|
||
t.Cleanup(func() {
|
||
log.SetOutput(prevWriter)
|
||
log.SetFlags(prevFlags)
|
||
})
|
||
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": " ",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if got := stub.calls.Load(); got != 1 {
|
||
t.Fatalf("embedder calls = %d, want 1 (content only)", got)
|
||
}
|
||
if !strings.Contains(buf.String(), "empty name provided from upstream") {
|
||
t.Fatalf("log output = %q, want empty-name warning", buf.String())
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
want := []float64{2, 4}
|
||
if !floatSliceClose(got[0]["q_2_vec"].([]float64), want) {
|
||
t.Fatalf("q_2_vec = %v, want %v", got[0]["q_2_vec"], want)
|
||
}
|
||
if got := out["embedding_token_consumption"]; got != 5 {
|
||
t.Fatalf("embedding_token_consumption = %v, want 5", got)
|
||
}
|
||
}
|
||
|
||
// 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)
|
||
stub.maxTokens = 12
|
||
longText := strings.Repeat("hello world ", 20)
|
||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": longText}},
|
||
}); err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if len(stub.callInputs) != 2 {
|
||
t.Fatalf("callInputs len = %d, want 2", len(stub.callInputs))
|
||
}
|
||
if len(stub.callInputs[1]) != 1 {
|
||
t.Fatalf("content batch size = %d, want 1", len(stub.callInputs[1]))
|
||
}
|
||
if got := stub.callInputs[1][0]; len(got) >= len(longText) {
|
||
t.Fatalf("content text was not truncated: original=%d got=%d", len(longText), len(got))
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Embedding_SkipsEmptyCleanedTextsButReturnsZeroWhenAllSkipped(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 2)
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{
|
||
{"text": "<table><tr><td></td></tr></table>"},
|
||
{"text": " "},
|
||
},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if got := stub.calls.Load(); got != 0 {
|
||
t.Fatalf("embedder calls = %d, want 0", got)
|
||
}
|
||
if got := out["embedding_token_consumption"]; got != 0 {
|
||
t.Fatalf("embedding_token_consumption = %v, want 0", got)
|
||
}
|
||
got, _ := out["chunks"].([]map[string]any)
|
||
for i, ck := range got {
|
||
if _, ok := ck["q_2_vec"]; ok {
|
||
t.Fatalf("chunk[%d] should not have vector: %v", i, ck)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Embedding_SetsTokenConsumptionIncludingTitleCall(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 2)
|
||
stub.callTokens = []int{3, 5, 7}
|
||
prevBatchSize := tokenizerEmbeddingBatchSize
|
||
tokenizerEmbeddingBatchSize = 1
|
||
t.Cleanup(func() { tokenizerEmbeddingBatchSize = prevBatchSize })
|
||
|
||
out, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{{"text": "alpha"}, {"text": "beta"}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if got := out["embedding_token_consumption"]; got != 15 {
|
||
t.Fatalf("embedding_token_consumption = %v, want 15", got)
|
||
}
|
||
}
|
||
|
||
func TestTokenizerComponent_Embedding_BatchesByConfiguredBatchSize(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
c, stub := withStubEmbedder(t, 2)
|
||
prevBatchSize := tokenizerEmbeddingBatchSize
|
||
tokenizerEmbeddingBatchSize = 2
|
||
t.Cleanup(func() { tokenizerEmbeddingBatchSize = prevBatchSize })
|
||
|
||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||
"name": "doc.pdf",
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{
|
||
{"text": "one"},
|
||
{"text": "two"},
|
||
{"text": "three"},
|
||
{"text": "four"},
|
||
{"text": "five"},
|
||
},
|
||
}); err != nil {
|
||
t.Fatalf("Invoke: %v", err)
|
||
}
|
||
if got := stub.calls.Load(); got != 4 {
|
||
t.Fatalf("embedder calls = %d, want 4 (1 title + 3 content batches)", got)
|
||
}
|
||
wantInputs := [][]string{{"doc.pdf"}, {"one", "two"}, {"three", "four"}, {"five"}}
|
||
if !reflect.DeepEqual(stub.callInputs, wantInputs) {
|
||
t.Fatalf("call inputs = %#v, want %#v", stub.callInputs, wantInputs)
|
||
}
|
||
}
|
||
|
||
func floatSliceClose(got, want []float64) bool {
|
||
if len(got) != len(want) {
|
||
return false
|
||
}
|
||
for i := range got {
|
||
if math.Abs(got[i]-want[i]) > 1e-9 {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
func TestTokenizerComponent_InstanceResolversDoNotLeakAcrossComponents(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
compAIntf, err := NewTokenizerComponentWithResolver(nil, func(_, _, _ string) (Embedder, error) {
|
||
stub := newStubEmbedder(2)
|
||
stub.resultsByCall = []embeddingCallResult{{vectors: [][]float64{{10, 10}}, tokenCount: 1}, {vectors: [][]float64{{1, 1}}, tokenCount: 1}}
|
||
return stub, nil
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewTokenizerComponentWithResolver(A): %v", err)
|
||
}
|
||
compBIntf, err := NewTokenizerComponentWithResolver(nil, func(_, _, _ string) (Embedder, error) {
|
||
stub := newStubEmbedder(2)
|
||
stub.resultsByCall = []embeddingCallResult{{vectors: [][]float64{{20, 20}}, tokenCount: 1}, {vectors: [][]float64{{2, 2}}, tokenCount: 1}}
|
||
return stub, nil
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("NewTokenizerComponentWithResolver(B): %v", err)
|
||
}
|
||
compA := compAIntf.(*TokenizerComponent)
|
||
compB := compBIntf.(*TokenizerComponent)
|
||
|
||
outA, err := compA.Invoke(context.Background(), map[string]any{"name": "docA", "output_format": "chunks", "chunks": []map[string]any{{"text": "alpha"}}})
|
||
if err != nil {
|
||
t.Fatalf("Invoke A: %v", err)
|
||
}
|
||
outB, err := compB.Invoke(context.Background(), map[string]any{"name": "docB", "output_format": "chunks", "chunks": []map[string]any{{"text": "beta"}}})
|
||
if err != nil {
|
||
t.Fatalf("Invoke B: %v", err)
|
||
}
|
||
vecA := outA["chunks"].([]map[string]any)[0]["q_2_vec"].([]float64)
|
||
vecB := outB["chunks"].([]map[string]any)[0]["q_2_vec"].([]float64)
|
||
if reflect.DeepEqual(vecA, vecB) {
|
||
t.Fatalf("instance resolvers leaked: vecA=%v vecB=%v", vecA, vecB)
|
||
}
|
||
}
|
||
|
||
func TestTokenizeChunks_SymbolOnlyTextFallsBackToRawText(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
chunks := []schema.ChunkDoc{
|
||
{Text: "·"}, // middle dot · — seen in production chunk[15]
|
||
{Text: ")"},
|
||
{Text: "("},
|
||
{Text: "*"},
|
||
}
|
||
err := tokenizeChunks(chunks, "test", "English")
|
||
if err != nil {
|
||
t.Fatalf("tokenizeChunks: %v", err)
|
||
}
|
||
for i, ck := range chunks {
|
||
t.Logf("chunk[%d]: text=%q content_ltks=%q content_sm_ltks=%q",
|
||
i, ck.Text, ck.ContentLtks, ck.ContentSmLtks)
|
||
// After fix: Tokenize returns empty for symbol-only text,
|
||
// but the fallback sets ContentLtks = raw text.
|
||
if strings.TrimSpace(ck.ContentLtks) == "" {
|
||
t.Errorf("chunk[%d]: expected non-empty ContentLtks (raw text fallback) for %q, got empty",
|
||
i, ck.Text)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestTokenizeChunks_WhitespaceSummaryShadowsTextBug(t *testing.T) {
|
||
requireTokenizerPool(t)
|
||
chunks := []schema.ChunkDoc{
|
||
{Summary: " ", Text: "real content here"},
|
||
}
|
||
err := tokenizeChunks(chunks, "test", "English")
|
||
if err != nil {
|
||
t.Fatalf("tokenizeChunks: %v", err)
|
||
}
|
||
// After fix: TrimSpace(" ") is empty, so the Summary branch is skipped.
|
||
// The Text branch is entered and "real content here" is tokenized normally.
|
||
if strings.TrimSpace(chunks[0].ContentLtks) == "" {
|
||
t.Errorf("whitespace Summary should be skipped, Text %q should be tokenized, but ContentLtks is empty",
|
||
chunks[0].Text)
|
||
}
|
||
}
|