mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-23 08:56:42 +08:00
fix(ingestion): align laws DSL with Python — heading fallback, colon-title, short-line filter, remove_toc, and image extension mapping (#17200)
## Summary
This PR aligns the Go ingestion pipeline's **Laws** DSL template with
the Python implementation by fixing heading-detection gaps, adds
image-extension support, refactors the **Extractor** component's LLM
resolution, hardens heading detection for CJK text, and makes the
Extractor accept the Python DSL prompt key names
(`sys_prompt`/`prompts`) alongside the Go names.
## Changes
### 1. Picture file-type detection (`internal/utility/file.go`)
Adds explicit mapping for common image extensions (png, jpg, jpeg, gif,
bmp, tiff, tif, webp, svg, ico, avif, heic, apng) → `FileTypeVISUAL`,
with regression tests.
### 2. Laws DSL heading-detection alignment
(`internal/ingestion/component/chunker/`)
Four fixes to `resolveTitleLevels`:
| Fix | What changed | Why |
|-----|-------------|-----|
| **DOCX `ck_type` fallback** | `ckType` field on `lineRecord`,
propagated from `ChunkDoc.CKType` in `recordsFromStructured`. When
`ck_type=="heading"`, assign `fallbackLevel`. | office_oxide extracts
DOCX heading metadata, but the info was lost before reaching the heading
detector. Word headings whose text doesn't match any regex (e.g.
"Introduction") were treated as body. |
| **`make_colon_as_title` promotion** | `isColonTitle()`: promotes lines
ending with `:`/`:` that have sentence-ending punctuation before the
colon and ≥32 runes between them. | Mirrors Python's
`make_colon_as_title` in `rag/nlp/__init__.py`. Triple guard prevents
false positives. |
| **Short/numeric line filter** | Lines with ≤1 rune or purely numeric
are pinned to body level. | Mirrors Python `tree_merge`'s filter of
`sections` where `len(...) <= 1` or `re.match(r"[0-9]+$", ...)`. |
| **PDF `remove_toc`** | `"remove_toc": true` added to the PDF parser
setup in `ingestion_pipeline_laws.json`. | The Go PDF parser already
supports TOC removal; the Book template already enables it. |
### 3. Extractor llm_id resolution
(`internal/ingestion/component/extractor.go`)
Refactored to handle both **bare tenant_model UUIDs** and **composite
model@provider** strings via the shared `resolveModelConfig`
(`dispatch_model.go`):
- **`resolveExtractorChatConfig`** — UUID path calls
`resolveModelConfigByID` directly (one DB hit); composite path goes
through `resolveModelConfig`. Added `isBareTenantModelID` pre-check for
clear errors when a UUID doesn't exist.
- **`resolveExtractorChatTarget`** — propagates resolution errors
instead of silently returning empty driver.
- **`Chat()`** — removed `driver = "dummy"` fallback. Missing driver is
now an explicit error.
- **Removed dead code**: `splitExtractorLLID`,
`findExtractorSoleActiveInstance`.
### 4. `InjectExtractorLLMID` — fallback when no user config
(`internal/common/parser_config.go`)
Injects the tenant's global default LLM into extractor components **only
when their `llm_id` is empty**. Preserves user-selected UUID or
model@provider values.
Priority: user-configured llm_id > tenant global default > error (no
silent dummy fallback).
### 5. `ResponseHeaderTimeout` increase
(`internal/entity/models/base_model.go`)
`ResponseHeaderTimeout` 60s → 120s in `NewDriverHTTPClient`. Reasoning
models with large extraction prompts can take longer than 60s to produce
the first response token.
### 6. CJK rune-aware heading detection
(`internal/ingestion/component/chunker/title.go`)
Two byte-vs-rune bugs that only manifest on CJK text:
| Fix | What changed | Why |
|-----|-------------|-----|
| **`isColonTitle` byte offset** | `body[lastPunct+1:]` →
`body[lastPunct+runeLen:]` via `utf8.DecodeRuneInString` |
`strings.LastIndexAny` returns a byte index; `+1` skips only 1 byte,
corrupting multi-byte CJK punctuation (e.g. `。` = 3 bytes) and inflating
the rune count past the 32-rune threshold → false-positive heading
promotion. |
| **Short-line filter byte count** | `len(text) <= 1` →
`utf8.RuneCountInString(text) <= 1` | Go `len` is UTF-8 bytes; a single
CJK char (3 bytes) passed the filter, but Python's `len` returns 1 →
mismatch. |
### 7. `extractor_tag.go` — log error when llm fails
When `resolveExtractorChatTarget` returned an error, `runAutoTags` will
log error.
### 8. Python DSL prompt-key compatibility
(`internal/ingestion/component/extractor.go`)
The Resume DSL template uses Python-side key names (`sys_prompt`,
`prompts`). `NewExtractorComponent` now accepts them as fallbacks
alongside the Go names:
- `system_prompt` (Go) ← `sys_prompt` (Python) as fallback
- `prompt` (Go string) ← `prompts` (Python array `[{"role","content"}]`,
takes `[0].content`) as fallback
Mirrors the alias pattern already in `internal/agent/component/llm.go`.
`resolveInputs` accepts per-call `sys_prompt` override too.
## Remaining gaps vs Python
| Gap | Scope | Impact |
|-----|-------|--------|
| **TOC removal for TXT/MD/HTML** | Python's `remove_contents_table`
works on all text formats; Go's `remove_toc` is PDF-only. | Low —
plain-text documents rarely contain structured TOCs. |
| **Regex pattern details** | Minor differences in quantifiers, missing
H5/H6 markdown patterns, missing 4-level numbering pattern. | Low — Go's
variants are stricter; DOCX headings are covered by `ck_type` fallback.
|
## Testing
- `TestHierarchyTitleChunker_CKTypeHeadingFallback` — DOCX `ck_type`
heading promotion
- `TestHierarchyTitleChunker_ColonTitlePromotion` /
`_ColonTitleShortLine_Negative` — colon-title promotion + guard
- `TestIsColonTitle_CJKEdgeCase` / `TestIsColonTitle_ASCII_NoRegression`
— CJK byte-offset fix + ASCII regression
- `TestHierarchyTitleChunker_ColonTitlePromotion_CJK_EdgeCase` — CJK
colon edge case through full pipeline
- `TestHierarchyTitleChunker_ShortSingleCJKLineFilter` — single CJK char
filtered to body
- `TestHierarchyTitleChunker_ShortNumericLineFilter` — purely numeric
lines filtered
- `TestGetFileType_ImageExtensions` / `_ExistingFormats_NoRegression` —
image extension mapping
- `TestInjectExtractorLLMID_SkipWhenUUID` / `_SkipWhenComposite` /
`_InjectWhenEmpty` — llm_id injection guard
- `TestIsBareTenantModelID` — UUID detection
- `TestResolveExtractorChatTarget_AtSplitFallback` / `_NoDriver` — @
split fallback without DB
- `TestNewExtractorComponent_SysPromptAlias` / `_PromptsArray` /
`_PromptsArray_PromptWins` / `_SystemPromptWinsOverSysPrompt` — Python
key compatibility
- `TestBuildDOCXJSONSections_List` / `_TextBox` / `_MixedWithList` —
DOCX list/text_box parsing
- Full ingestion test suite passes (chunker, pipeline, task, service,
component packages)
This commit is contained in:
@@ -17,8 +17,10 @@ func InjectExtractorLLMID(parserConfig map[string]interface{}, llmID string) boo
|
||||
}
|
||||
cidLower := strings.ToLower(cid)
|
||||
if strings.HasPrefix(cidLower, "extractor:") || strings.HasPrefix(cidLower, "extractor_") {
|
||||
compMap["llm_id"] = llmID
|
||||
updated = true
|
||||
if current, ok := compMap["llm_id"].(string); !ok || current == "" {
|
||||
compMap["llm_id"] = llmID
|
||||
updated = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return updated
|
||||
|
||||
@@ -1,215 +1,53 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
import "testing"
|
||||
|
||||
func TestDeepMergeMaps_Nil(t *testing.T) {
|
||||
result := DeepMergeMaps(nil, nil)
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil for nil inputs")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepMergeMaps_Override(t *testing.T) {
|
||||
base := map[string]interface{}{
|
||||
"a": 1,
|
||||
"b": "hello",
|
||||
}
|
||||
override := map[string]interface{}{
|
||||
"a": 2,
|
||||
"c": "world",
|
||||
}
|
||||
result := DeepMergeMaps(base, override)
|
||||
if result["a"] != 2 {
|
||||
t.Fatalf("expected a=2, got %v", result["a"])
|
||||
}
|
||||
if result["b"] != "hello" {
|
||||
t.Fatalf("expected b=hello, got %v", result["b"])
|
||||
}
|
||||
if result["c"] != "world" {
|
||||
t.Fatalf("expected c=world, got %v", result["c"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepMergeMaps_Nested(t *testing.T) {
|
||||
base := map[string]interface{}{
|
||||
"parser": map[string]interface{}{
|
||||
"lang": "en",
|
||||
"model": "gpt4",
|
||||
func TestInjectExtractorLLMID_SkipWhenUUID(t *testing.T) {
|
||||
uuid := "9e819c2442b14f9dab46062916e29195"
|
||||
pc := map[string]interface{}{
|
||||
"Extractor:A": map[string]interface{}{
|
||||
"llm_id": uuid,
|
||||
},
|
||||
}
|
||||
override := map[string]interface{}{
|
||||
"parser": map[string]interface{}{
|
||||
"lang": "zh",
|
||||
InjectExtractorLLMID(pc, "Qwen/Qwen3-8B@siliconflow")
|
||||
id := pc["Extractor:A"].(map[string]interface{})["llm_id"].(string)
|
||||
if id != uuid {
|
||||
t.Fatalf("expected UUID preserved, got %q", id)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInjectExtractorLLMID_SkipWhenComposite(t *testing.T) {
|
||||
composite := "Qwen/Qwen3-8B@siliconflow"
|
||||
pc := map[string]interface{}{
|
||||
"Extractor:B": map[string]interface{}{
|
||||
"llm_id": composite,
|
||||
},
|
||||
}
|
||||
result := DeepMergeMaps(base, override)
|
||||
parser, _ := result["parser"].(map[string]interface{})
|
||||
if parser["lang"] != "zh" {
|
||||
t.Fatalf("expected lang=zh, got %v", parser["lang"])
|
||||
}
|
||||
if parser["model"] != "gpt4" {
|
||||
t.Fatalf("expected model=gpt4, got %v", parser["model"])
|
||||
InjectExtractorLLMID(pc, "DeepSeek@siliconflow")
|
||||
id := pc["Extractor:B"].(map[string]interface{})["llm_id"].(string)
|
||||
if id != composite {
|
||||
t.Fatalf("expected composite preserved, got %q", id)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetParserConfig_Naive(t *testing.T) {
|
||||
result := GetParserConfig("naive", nil)
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil")
|
||||
func TestInjectExtractorLLMID_InjectWhenEmpty(t *testing.T) {
|
||||
defaultLLM := "Qwen/Qwen3-8B@siliconflow"
|
||||
pc := map[string]interface{}{
|
||||
"Extractor:C": map[string]interface{}{},
|
||||
}
|
||||
if result["chunk_token_num"] != 512 {
|
||||
t.Fatalf("expected chunk_token_num=512, got %v", result["chunk_token_num"])
|
||||
}
|
||||
if result["layout_recognize"] != "DeepDOC" {
|
||||
t.Fatalf("expected layout_recognize=DeepDOC, got %v", result["layout_recognize"])
|
||||
InjectExtractorLLMID(pc, defaultLLM)
|
||||
id := pc["Extractor:C"].(map[string]interface{})["llm_id"].(string)
|
||||
if id != defaultLLM {
|
||||
t.Fatalf("expected %q injected, got %q", defaultLLM, id)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetParserConfig_QA(t *testing.T) {
|
||||
result := GetParserConfig("qa", nil)
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetParserConfig_Override(t *testing.T) {
|
||||
override := map[string]interface{}{
|
||||
"chunk_token_num": 256,
|
||||
}
|
||||
result := GetParserConfig("naive", override)
|
||||
if result["chunk_token_num"] != 256 {
|
||||
t.Fatalf("expected chunk_token_num=256, got %v", result["chunk_token_num"])
|
||||
}
|
||||
if result["layout_recognize"] != "DeepDOC" {
|
||||
t.Fatalf("expected layout_recognize preserved, got %v", result["layout_recognize"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractPipelineDefaults_Nil(t *testing.T) {
|
||||
if result := ExtractPipelineDefaults(nil); result != nil {
|
||||
t.Fatalf("expected nil for nil input, got %v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractPipelineDefaults_Basic(t *testing.T) {
|
||||
dsl := map[string]interface{}{
|
||||
"components": map[string]interface{}{
|
||||
"Parser:abc": map[string]interface{}{
|
||||
"obj": map[string]interface{}{
|
||||
"component_name": "Parser",
|
||||
"params": map[string]interface{}{
|
||||
"pdf": map[string]interface{}{
|
||||
"parse_method": "DeepDOC",
|
||||
"lang": "en",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := ExtractPipelineDefaults(dsl)
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil")
|
||||
}
|
||||
params, ok := result["Parser:abc"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatal("expected Parser:abc in result")
|
||||
}
|
||||
pdf, _ := params["pdf"].(map[string]interface{})
|
||||
if pdf["parse_method"] != "DeepDOC" {
|
||||
t.Fatalf("expected parse_method=DeepDOC, got %v", pdf["parse_method"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractPipelineDefaults_SkipsFile(t *testing.T) {
|
||||
dsl := map[string]interface{}{
|
||||
"components": map[string]interface{}{
|
||||
"File:xyz": map[string]interface{}{
|
||||
"obj": map[string]interface{}{
|
||||
"component_name": "File",
|
||||
"params": map[string]interface{}{
|
||||
"path": "/tmp/test",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := ExtractPipelineDefaults(dsl)
|
||||
if result != nil {
|
||||
t.Fatal("expected nil when only File component exists")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractPipelineDefaults_StripsOutputs(t *testing.T) {
|
||||
dsl := map[string]interface{}{
|
||||
"components": map[string]interface{}{
|
||||
"Tokenizer:abc": map[string]interface{}{
|
||||
"obj": map[string]interface{}{
|
||||
"component_name": "Tokenizer",
|
||||
"params": map[string]interface{}{
|
||||
"fields": "text",
|
||||
"search_method": []interface{}{"embedding"},
|
||||
"outputs": map[string]interface{}{"chunks": "data"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := ExtractPipelineDefaults(dsl)
|
||||
params := result["Tokenizer:abc"].(map[string]interface{})
|
||||
if _, ok := params["outputs"]; ok {
|
||||
t.Fatal("expected outputs to be stripped")
|
||||
}
|
||||
if params["fields"] != "text" {
|
||||
t.Fatalf("expected fields=text, got %v", params["fields"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractPipelineDefaults_WithDSLWrapper(t *testing.T) {
|
||||
dsl := map[string]interface{}{
|
||||
"dsl": map[string]interface{}{
|
||||
"components": map[string]interface{}{
|
||||
"Extractor:xyz": map[string]interface{}{
|
||||
"obj": map[string]interface{}{
|
||||
"component_name": "Extractor",
|
||||
"params": map[string]interface{}{
|
||||
"auto_keywords": float64(5),
|
||||
"auto_questions": float64(3),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := ExtractPipelineDefaults(dsl)
|
||||
raw := result["Extractor:xyz"].(map[string]interface{})
|
||||
if raw["auto_keywords"] != float64(5) {
|
||||
t.Fatalf("expected auto_keywords=5, got %v", raw["auto_keywords"])
|
||||
}
|
||||
if raw["auto_questions"] != float64(3) {
|
||||
t.Fatalf("expected auto_questions=3, got %v", raw["auto_questions"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractPipelineDefaults_Delimiters(t *testing.T) {
|
||||
dsl := map[string]interface{}{
|
||||
"components": map[string]interface{}{
|
||||
"TokenChunker:abc": map[string]interface{}{
|
||||
"obj": map[string]interface{}{
|
||||
"component_name": "TokenChunker",
|
||||
"params": map[string]interface{}{
|
||||
"delimiters": []interface{}{"\n", ".", " "},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := ExtractPipelineDefaults(dsl)
|
||||
raw := result["TokenChunker:abc"].(map[string]interface{})
|
||||
delims, _ := raw["delimiters"].([]interface{})
|
||||
if len(delims) != 3 {
|
||||
t.Fatalf("expected 3 delimiters, got %v", delims)
|
||||
func TestInjectExtractorLLMID_NoExtractor(t *testing.T) {
|
||||
pc := map[string]interface{}{
|
||||
"Parser:X": map[string]interface{}{"llm_id": ""},
|
||||
}
|
||||
InjectExtractorLLMID(pc, "default@provider")
|
||||
if _, ok := pc["Parser:X"]; !ok {
|
||||
t.Fatal("expected Parser:X still present")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -193,7 +193,7 @@ func NewDriverHTTPClient() *http.Client {
|
||||
t.MaxIdleConnsPerHost = 10
|
||||
t.IdleConnTimeout = 90 * time.Second
|
||||
t.DisableCompression = false
|
||||
t.ResponseHeaderTimeout = 60 * time.Second
|
||||
t.ResponseHeaderTimeout = 2 * 60 * time.Second
|
||||
t.TLSHandshakeTimeout = 30 * time.Second
|
||||
return &http.Client{Transport: t}
|
||||
}
|
||||
|
||||
@@ -337,6 +337,7 @@ func recordsFromStructured(items []schema.ChunkDoc) []lineRecord {
|
||||
docType: dt,
|
||||
imgID: imgID,
|
||||
layout: it.Layout,
|
||||
ckType: it.CKType,
|
||||
parentMeta: meta,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -227,6 +227,270 @@ func TestHierarchyTitleChunker_InvokeDeterministic(t *testing.T) {
|
||||
// preceding text run on the first non-text and is a no-op for the next
|
||||
// non-text, yielding [..., N1, N2, run, ...]; the old loop flushed the
|
||||
// trailing run early, yielding [..., N1, run, N2].
|
||||
// TestHierarchyTitleChunker_ColonTitlePromotion verifies that a line
|
||||
// ending with colon, having sentence-ending punctuation before it, and
|
||||
// at least 32 chars between them, is promoted to heading level (mirroring
|
||||
// Python make_colon_as_title intent). Two colon headings produce 2 chunks
|
||||
// while without the fix all text would merge into 1 chunk.
|
||||
func TestHierarchyTitleChunker_ColonTitlePromotion(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}}, // won't match our plain text
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "Introductory section providing background. The scope and purpose of this document are defined as follows:", "doc_type_kwd": "text"},
|
||||
{"text": "Body one.", "doc_type_kwd": "text"},
|
||||
{"text": "Another section continuing the discussion. The key provisions are outlined below:", "doc_type_kwd": "text"},
|
||||
{"text": "Body two.", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "test.txt",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// Without fix: 1 chunk (all body). With fix: 2 chunks (colon heading + body each).
|
||||
if len(chunks) != 2 {
|
||||
t.Fatalf("len(chunks) = %d, want 2 (colon titles should each produce a chunk)", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestHierarchyTitleChunker_ColonTitleShortLine_Negative verifies that
|
||||
// short colon-ended lines (e.g. "Note:") are NOT promoted, avoiding
|
||||
// false positives.
|
||||
func TestHierarchyTitleChunker_ColonTitleShortLine_Negative(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "Note:", "doc_type_kwd": "text"},
|
||||
{"text": "Body text.", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "test.txt",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// Short colon line must NOT be promoted: all body → 1 chunk.
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("len(chunks) = %d, want 1 (short colon line must not be promoted)", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestHierarchyTitleChunker_ShortNumericLineFilter verifies that
|
||||
// purely numeric short lines are filtered to body level even when
|
||||
// they match a heading regex (mirroring Python tree_merge's
|
||||
// line filter: len(t.split("@")[0].strip()) > 1 + not purely numeric).
|
||||
func TestHierarchyTitleChunker_ShortNumericLineFilter(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^[0-9]+$`}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "1", "doc_type_kwd": "text"},
|
||||
{"text": "Introduction paragraph.", "doc_type_kwd": "text"},
|
||||
{"text": "2", "doc_type_kwd": "text"},
|
||||
{"text": "Methodology paragraph.", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "test.txt",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// Without fix: "1" and "2" match ^[0-9]+$ → 2 chunks
|
||||
// With fix: "1" and "2" are purely numeric → filtered to body → 1 chunk
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("len(chunks) = %d, want 1 (purely numeric lines filtered to body)", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsColonTitle_CJKEdgeCase verifies the isColonTitle function with
|
||||
// CJK sentence-ending punctuation (。). When there are exactly 31 ASCII
|
||||
// characters between the 。 and the :, the function must return false
|
||||
// (31 < 32 rune threshold). The Go byte-offset bug (body[lastPunct+1:])
|
||||
// would corrupt the CJK punctuation bytes and artifactually inflate the
|
||||
// rune count past 32, causing a false positive.
|
||||
func TestIsColonTitle_CJKEdgeCase(t *testing.T) {
|
||||
// 31 ASCII chars between 。 and : → < 32 runes → must not promote
|
||||
line := "abc。1234567890123456789012345678901:"
|
||||
if isColonTitle(line) {
|
||||
t.Error("isColonTitle should be false: 31 ASCII chars between CJK punct and colon is < 32")
|
||||
}
|
||||
|
||||
// 32 ASCII chars → exactly 32 → must promote
|
||||
line32 := "abc。12345678901234567890123456789012:"
|
||||
if !isColonTitle(line32) {
|
||||
t.Error("isColonTitle should be true: 32 ASCII chars between CJK punct and colon is >= 32")
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsColonTitle_ASCII_NoRegression verifies the existing ASCII-only
|
||||
// isColonTitle path still works correctly (regression guard).
|
||||
func TestIsColonTitle_ASCII_NoRegression(t *testing.T) {
|
||||
// 32 ASCII chars between . and : → must promote
|
||||
line := "abc.12345678901234567890123456789012:"
|
||||
if !isColonTitle(line) {
|
||||
t.Error("isColonTitle should be true: 32 ASCII chars between . and :")
|
||||
}
|
||||
// 31 ASCII chars → must NOT promote
|
||||
line31 := "abc.1234567890123456789012345678901:"
|
||||
if isColonTitle(line31) {
|
||||
t.Error("isColonTitle should be false: 31 ASCII chars between . and :")
|
||||
}
|
||||
// Short colon line → must NOT promote
|
||||
if isColonTitle("Note:") {
|
||||
t.Error("isColonTitle should be false for short colon line")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHierarchyTitleChunker_ColonTitlePromotion_CJK_EdgeCase verifies
|
||||
// the CJK colon-title edge case through the full hierarchy chunker
|
||||
// pipeline. With 31 ASCII chars between 。 and : the line must NOT be
|
||||
// promoted (1 chunk), but the byte-offset bug causes a false promotion
|
||||
// (2 chunks).
|
||||
func TestHierarchyTitleChunker_ColonTitlePromotion_CJK_EdgeCase(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "abc。1234567890123456789012345678901:", "doc_type_kwd": "text"},
|
||||
{"text": "Body one.", "doc_type_kwd": "text"},
|
||||
{"text": "def。1234567890123456789012345678901:", "doc_type_kwd": "text"},
|
||||
{"text": "Body two.", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "test.txt",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// With fix: both have 31 < 32 → NOT promoted → 1 chunk (all body).
|
||||
// Without fix: byte-offset corruption inflates rune count past 32 → both promoted → 2 chunks.
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("len(chunks) = %d, want 1 (31 ASCII chars between CJK punct and colon is < 32, must NOT promote)", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestHierarchyTitleChunker_ShortSingleCJKLineFilter verifies that a
|
||||
// single CJK character (3 bytes UTF-8) is filtered to body level.
|
||||
// The Go byte-count bug (len(text) <= 1) would let it through because
|
||||
// len("案") = 3 > 1, but Python's Unicode-aware len returns 1, so the
|
||||
// line should be filtered to body. Without the fix the single CJK char
|
||||
// would be promoted when it matches a heading regex.
|
||||
func TestHierarchyTitleChunker_ShortSingleCJKLineFilter(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^..?$`}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "案", "doc_type_kwd": "text"},
|
||||
{"text": "Body one.", "doc_type_kwd": "text"},
|
||||
{"text": "例", "doc_type_kwd": "text"},
|
||||
{"text": "Body two.", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "test.txt",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// With fix: "案" is 1 rune → len <= 1 → filtered to body → 1 chunk.
|
||||
// Without fix: len("案") = 3 bytes > 1 → passes filter → ^..?$ matches → 2 chunks.
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("len(chunks) = %d, want 1 (single CJK char must be filtered to body)", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
// TestHierarchyTitleChunker_CKTypeHeadingFallback verifies that
|
||||
// records with ck_type "heading" (from office_oxide DOCX parsing)
|
||||
// are treated as heading nodes in the tree, even when their text
|
||||
// doesn't match any regex pattern. Without the fix, such records
|
||||
// are treated as body text — all content merges into one chunk.
|
||||
func TestHierarchyTitleChunker_CKTypeHeadingFallback(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 2,
|
||||
"levels": [][]string{{`^# `}, {`^## `}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewHierarchyTitleChunker: %v", err)
|
||||
}
|
||||
items := []map[string]any{
|
||||
{"text": "Introduction", "doc_type_kwd": "text", "ck_type": "heading"},
|
||||
{"text": "Intro body.", "doc_type_kwd": "text"},
|
||||
{"text": "Background", "doc_type_kwd": "text", "ck_type": "heading"},
|
||||
{"text": "Background body.", "doc_type_kwd": "text"},
|
||||
{"text": "Conclusion", "doc_type_kwd": "text", "ck_type": "heading"},
|
||||
{"text": "Final words.", "doc_type_kwd": "text"},
|
||||
}
|
||||
out, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": "test.docx",
|
||||
"chunks": items,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
chunks, ok := out["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) == 0 {
|
||||
t.Fatal("chunks: want >=1, got 0")
|
||||
}
|
||||
// Without the fix: 1 chunk (all text merged as body under root).
|
||||
// With the fix: 3 chunks (one per heading + its body).
|
||||
if len(chunks) != 3 {
|
||||
t.Fatalf("len(chunks) = %d, want 3 (3 ck_type=headings should produce 3 chunks)", len(chunks))
|
||||
}
|
||||
for _, ck := range chunks {
|
||||
text, _ := ck["text"].(string)
|
||||
if text == "" {
|
||||
t.Error("chunk text is empty")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestHierarchyTitleChunker_ConsecutiveNonTextOrder pins Gap G: two
|
||||
func TestHierarchyTitleChunker_ConsecutiveNonTextOrder(t *testing.T) {
|
||||
c, err := NewHierarchyTitleChunker(map[string]any{
|
||||
"hierarchy": 1,
|
||||
|
||||
@@ -221,6 +221,7 @@ var (
|
||||
notTitleException = regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+条`)
|
||||
notTitlePunct = regexp.MustCompile(`[,;,。;!!]`)
|
||||
layoutHeadingRe = regexp.MustCompile(`(?i)(section|title|head)`)
|
||||
numericOnly = regexp.MustCompile(`^[0-9]+$`)
|
||||
)
|
||||
|
||||
// notTitle mirrors rag/nlp.not_title. Returns true when the line looks
|
||||
@@ -271,6 +272,33 @@ func matchLayoutLevel(text, layout string, fallbackLevel int) int {
|
||||
// BODY_LEVEL.
|
||||
//
|
||||
// fallback_level is len(selected_group) + 1, exactly as in python.
|
||||
// isColonTitle mirrors Python make_colon_as_title intent: returns true
|
||||
// when a line ends with colon, has sentence-ending punctuation before
|
||||
// the colon, and the text between them is at least 32 runes long.
|
||||
func isColonTitle(s string) bool {
|
||||
if s == "" {
|
||||
return false
|
||||
}
|
||||
if !strings.HasSuffix(s, ":") && !strings.HasSuffix(s, ":") {
|
||||
return false
|
||||
}
|
||||
body := strings.TrimSuffix(s, ":")
|
||||
body = strings.TrimSuffix(body, ":")
|
||||
if body == s {
|
||||
return false
|
||||
}
|
||||
lastPunct := strings.LastIndexAny(body, "。?!!?;;.")
|
||||
if lastPunct < 0 {
|
||||
return false
|
||||
}
|
||||
// Use rune-aware slicing: body[lastPunct+1] would be wrong for CJK
|
||||
// punctuation (e.g. "。" is 3 bytes in UTF-8). Decode the rune at
|
||||
// lastPunct to get the correct byte width and skip the full rune.
|
||||
_, runeLen := utf8.DecodeRuneInString(body[lastPunct:])
|
||||
between := strings.TrimSpace(body[lastPunct+runeLen:])
|
||||
return utf8.RuneCountInString(between) >= 32
|
||||
}
|
||||
|
||||
func resolveTitleLevels(records []lineRecord, p *titleChunkerParam) []int {
|
||||
lines := make([]string, len(records))
|
||||
for i, r := range records {
|
||||
@@ -284,10 +312,28 @@ func resolveTitleLevels(records []lineRecord, p *titleChunkerParam) []int {
|
||||
out[i] = bodyLevel
|
||||
continue
|
||||
}
|
||||
// Python tree_merge short/numeric line filter:
|
||||
// sections = [s for s in sections if len(s.split("@")[0].strip()) > 1
|
||||
// and not re.match(r"[0-9]+$", s.split("@")[0].strip())]
|
||||
if text := beforeAt(rec.text); utf8.RuneCountInString(text) <= 1 || numericOnly.MatchString(text) {
|
||||
out[i] = bodyLevel
|
||||
continue
|
||||
}
|
||||
if lvl := matchRegexLevel(rec.text, group); lvl != 0 {
|
||||
out[i] = lvl
|
||||
continue
|
||||
}
|
||||
if rec.ckType == "heading" {
|
||||
out[i] = fallbackLevel
|
||||
continue
|
||||
}
|
||||
// Python make_colon_as_title: promote lines ending with colon
|
||||
// that have sentence-ending punctuation before it and at least
|
||||
// 32 chars between the punctuation and the colon.
|
||||
if text := beforeAt(rec.text); isColonTitle(text) {
|
||||
out[i] = fallbackLevel
|
||||
continue
|
||||
}
|
||||
out[i] = matchLayoutLevel(rec.text, rec.layout, fallbackLevel)
|
||||
}
|
||||
return out
|
||||
@@ -329,6 +375,7 @@ type lineRecord struct {
|
||||
docType string
|
||||
imgID *string
|
||||
layout string
|
||||
ckType string
|
||||
pdfPos []map[string]any
|
||||
parentMeta map[string]any
|
||||
}
|
||||
|
||||
@@ -75,7 +75,6 @@ import (
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/entity/models"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
@@ -88,7 +87,7 @@ const componentNameExtractor = "Extractor"
|
||||
// `@timeout(60)` default at rag/flow/base.py:60. The pipeline
|
||||
// orchestrator (Phase 3) overrides this if a stage-level ceiling
|
||||
// is configured.
|
||||
const extractorTimeout = 60 * time.Second
|
||||
const extractorTimeout = 600 * time.Second
|
||||
|
||||
const (
|
||||
autoKeywordPrompt = `## Role
|
||||
@@ -165,9 +164,17 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) {
|
||||
}
|
||||
if v, ok := params["system_prompt"].(string); ok {
|
||||
p.SystemPrompt = v
|
||||
} else if v, ok := params["sys_prompt"].(string); ok {
|
||||
p.SystemPrompt = v
|
||||
}
|
||||
if v, ok := params["prompt"].(string); ok {
|
||||
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 {
|
||||
p.Prompt = content
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := params["auto_keywords"]; ok {
|
||||
p.AutoKeywords = mapInt(v)
|
||||
@@ -330,8 +337,9 @@ func (e *einoExtractorChatInvoker) Chat(ctx context.Context, req extractorChatRe
|
||||
}
|
||||
}
|
||||
if driver == "" {
|
||||
driver = "dummy"
|
||||
return nil, fmt.Errorf("extractor: chat: no driver resolved for model %q", modelName)
|
||||
}
|
||||
common.Info(fmt.Sprintf("extractor: chat: driver=%s modelName=%s baseUrl=%s", driver, modelName, req.BaseURL))
|
||||
d, err := models.GetPreconfiguredDriver(driver, req.BaseURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("extractor: resolve driver %q: %w", driver, err)
|
||||
@@ -347,11 +355,13 @@ func (e *einoExtractorChatInvoker) Chat(ctx context.Context, req extractorChatRe
|
||||
wrapper := models.NewEinoChatModel(cm, chatCfg)
|
||||
// Honour ctx cancel up front so the caller's WithTimeout(...)
|
||||
// is observed even when the driver layer doesn't take a ctx.
|
||||
common.Info(fmt.Sprintf("try to chat with message: %v", req.Messages))
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, err := wrapper.Generate(ctx, toExtractorEinoMessages(req.Messages))
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("error when chat with message: %v", req.Messages), err)
|
||||
return nil, err
|
||||
}
|
||||
return &extractorChatResponse{Content: out.Content}, nil
|
||||
@@ -434,6 +444,8 @@ func (c *ExtractorComponent) resolveInputs(inputs map[string]any) extractorInput
|
||||
}
|
||||
if v, ok := inputs["system_prompt"].(string); ok && v != "" {
|
||||
out.systemPrompt = v
|
||||
} else if v, ok := inputs["sys_prompt"].(string); ok && v != "" {
|
||||
out.systemPrompt = v
|
||||
}
|
||||
if v, ok := inputs["lang"].(string); ok && v != "" {
|
||||
out.lang = v
|
||||
@@ -675,7 +687,10 @@ func splitKeywords(s string) []string {
|
||||
// raw string from the model — JSON parsing happens here so
|
||||
// callers can rely on a structured value downstream.
|
||||
func (c *ExtractorComponent) call(ctx context.Context, in extractorInputs, chunkText string) (any, error) {
|
||||
driver, modelName, apiKey, baseURL := resolveExtractorChatTarget(ctx, in.llmID)
|
||||
driver, modelName, apiKey, baseURL, err := resolveExtractorChatTarget(ctx, in.llmID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
msgs := buildExtractorMessages(in.systemPrompt, in.prompt, chunkText, in.chunks)
|
||||
inv := getExtractorChatInvoker()
|
||||
resp, err := inv.Chat(ctx, extractorChatRequest{
|
||||
@@ -704,43 +719,34 @@ func (c *ExtractorComponent) call(ctx context.Context, in extractorInputs, chunk
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
// resolveExtractorChatTarget splits a composite llm_id
|
||||
// "model@provider" into driver / model / api_key / base_url
|
||||
// using the tenant-scoped provider → instance → model lookup.
|
||||
func resolveExtractorChatTarget(ctx context.Context, llmID string) (driver, modelName, apiKey, baseURL string) {
|
||||
// resolveExtractorChatTarget resolves the llm_id into driver / model /
|
||||
// api_key / base_url. The llm_id may be a bare tenant_model UUID or
|
||||
// a composite "model@provider" string. Errors from DAO resolution are
|
||||
// propagated so the caller sees the real failure reason.
|
||||
func resolveExtractorChatTarget(ctx context.Context, llmID string) (driver, modelName, apiKey, baseURL string, err error) {
|
||||
if override := getExtractorChatTargetResolverOverride(); override != nil {
|
||||
if driver, modelName, apiKey, baseURL, ok := override(llmID); ok {
|
||||
return driver, modelName, apiKey, baseURL
|
||||
return driver, modelName, apiKey, baseURL, nil
|
||||
}
|
||||
}
|
||||
if llmID == "" {
|
||||
return "", "", "", ""
|
||||
}
|
||||
|
||||
// Derive driver and model name from the composite llm_id.
|
||||
modelName = llmID
|
||||
parsedModel, _, parsedProvider := splitExtractorLLID(llmID)
|
||||
if parsedProvider != "" {
|
||||
modelName = parsedModel
|
||||
driver = strings.ToLower(parsedProvider)
|
||||
cfg, cfgErr := resolveExtractorChatConfig(ctx, llmID)
|
||||
if cfgErr != nil {
|
||||
return "", "", "", "", cfgErr
|
||||
}
|
||||
|
||||
// Override with tenant-scoped credentials when available.
|
||||
cfg := resolveExtractorChatConfig(ctx, llmID)
|
||||
if cfg.driver != "" {
|
||||
driver = cfg.driver
|
||||
}
|
||||
if cfg.modelName != "" {
|
||||
modelName = cfg.modelName
|
||||
}
|
||||
if cfg.apiKey != "" {
|
||||
apiKey = cfg.apiKey
|
||||
}
|
||||
if cfg.baseURL != "" {
|
||||
baseURL = cfg.baseURL
|
||||
return cfg.driver, cfg.modelName, cfg.apiKey, cfg.baseURL, nil
|
||||
}
|
||||
|
||||
return driver, modelName, apiKey, baseURL
|
||||
// Fallback: when tenant credentials are not available
|
||||
// (no canvas state / no DB), fall back to the basic @ split
|
||||
// so callers can still use model@provider format in tests.
|
||||
if bare, provider, ok := splitExtractorLLIDPair(llmID); ok {
|
||||
return strings.ToLower(provider), bare, "", "", nil
|
||||
}
|
||||
// Nothing left to try — let Chat() surface a clear error when
|
||||
// the driver ends up empty.
|
||||
return "", llmID, "", "", nil
|
||||
}
|
||||
|
||||
// extractorChatConfig holds the resolved chat model configuration.
|
||||
@@ -752,145 +758,75 @@ type extractorChatConfig struct {
|
||||
}
|
||||
|
||||
// resolveExtractorChatConfig resolves tenant-scoped credentials for
|
||||
// the given composite llm_id using tenant_model_provider → instance → model.
|
||||
func resolveExtractorChatConfig(ctx context.Context, compositeLLMID string) extractorChatConfig {
|
||||
// the given llm_id via the shared resolveModelConfig helper.
|
||||
//
|
||||
// - Bare UUID → DAO lookup via resolveModelConfigByID.
|
||||
// - "model@provider" → parsed via resolveModelConfigFromProviderInstance.
|
||||
//
|
||||
// Returns nil error when there is no canvas state (unit tests) —
|
||||
// the caller's @ split fallback handles that case.
|
||||
func resolveExtractorChatConfig(ctx context.Context, compositeLLMID string) (extractorChatConfig, error) {
|
||||
state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
|
||||
if err != nil || state == nil {
|
||||
return extractorChatConfig{}
|
||||
return extractorChatConfig{}, nil
|
||||
}
|
||||
tidVal, _ := state.GetGlobal("tenant_id")
|
||||
tid, _ := tidVal.(string)
|
||||
if tid == "" {
|
||||
return extractorChatConfig{}
|
||||
return extractorChatConfig{}, nil
|
||||
}
|
||||
|
||||
// Split the composite id with rsplit from the right
|
||||
// (preserves embedded '@' in model names).
|
||||
pureModelName, instanceName, providerName := splitExtractorLLID(compositeLLMID)
|
||||
if providerName == "" {
|
||||
return extractorChatConfig{}
|
||||
}
|
||||
|
||||
// 1. Lookup provider (case-insensitive).
|
||||
provider, err := dao.NewTenantModelProviderDAO().GetByTenantIDAndProviderName(tid, providerName)
|
||||
if err != nil || provider == nil {
|
||||
common.Debug("extractor credentials: provider not found",
|
||||
zap.String("provider", providerName),
|
||||
zap.Error(err))
|
||||
return extractorChatConfig{}
|
||||
}
|
||||
|
||||
// 2. Lookup instance (with "default" → sole active fallback).
|
||||
instance, err := dao.NewTenantModelInstanceDAO().GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
||||
if err != nil || instance == nil {
|
||||
if instanceName == "default" {
|
||||
if fallback := findExtractorSoleActiveInstance(provider.ID); fallback != nil {
|
||||
common.Debug("extractor credentials: remapped default instance to sole active",
|
||||
zap.String("instance", fallback.InstanceName),
|
||||
zap.String("provider", providerName))
|
||||
instance = fallback
|
||||
err = nil
|
||||
}
|
||||
var driver models.ModelDriver
|
||||
var modelName string
|
||||
var apiConfig *models.APIConfig
|
||||
if isBareTenantModelID(compositeLLMID) {
|
||||
// UUID path: resolveModelConfigByID does a single GetByID and
|
||||
// returns a clear error if the record doesn't exist. No need
|
||||
// for a separate pre-check — resolveModelConfig's redundant
|
||||
// GetByID dispatch check is also bypassed.
|
||||
driver, modelName, apiConfig, _, err = resolveModelConfigByID(tid, entity.ModelTypeChat, compositeLLMID)
|
||||
if err != nil {
|
||||
return extractorChatConfig{}, fmt.Errorf("extractor: tenant model %q not found or not usable: %w", compositeLLMID, err)
|
||||
}
|
||||
} else {
|
||||
// Composite "model@provider" path: delegate to the shared dispatcher.
|
||||
driver, modelName, apiConfig, _, err = resolveModelConfig(tid, entity.ModelTypeChat, compositeLLMID)
|
||||
if err != nil {
|
||||
return extractorChatConfig{}, fmt.Errorf("extractor: resolve model %q: %w", compositeLLMID, err)
|
||||
}
|
||||
}
|
||||
if err != nil || instance == nil {
|
||||
common.Debug("extractor credentials: instance not found",
|
||||
zap.String("instance", instanceName),
|
||||
zap.String("provider", providerName))
|
||||
return extractorChatConfig{}
|
||||
}
|
||||
|
||||
// 3. Lookup tenant_model to validate the model is registered.
|
||||
model, modelErr := dao.NewTenantModelDAO().GetByProviderIDAndInstanceIDAndModelTypeAndModelName(
|
||||
provider.ID, instance.ID, int(entity.ModelTypeChat), pureModelName)
|
||||
if modelErr != nil || model == nil {
|
||||
common.Debug("extractor credentials: tenant_model not found",
|
||||
zap.String("model", pureModelName),
|
||||
zap.Error(modelErr))
|
||||
return extractorChatConfig{}
|
||||
}
|
||||
canonicalModelName := pureModelName
|
||||
if model.ModelName != "" {
|
||||
canonicalModelName = model.ModelName
|
||||
}
|
||||
common.Debug("extractor credentials: tenant_model found",
|
||||
zap.String("model", canonicalModelName))
|
||||
|
||||
// 4. Assemble config from instance.
|
||||
if instance.APIKey == "" {
|
||||
common.Debug("extractor credentials: instance has no api_key",
|
||||
zap.String("instance", instance.InstanceName),
|
||||
zap.String("provider", providerName))
|
||||
return extractorChatConfig{}
|
||||
}
|
||||
apiKey := ""
|
||||
baseURL := ""
|
||||
if instance.Extra != "" {
|
||||
var extra map[string]string
|
||||
if err := json.Unmarshal([]byte(instance.Extra), &extra); err == nil {
|
||||
baseURL = extra["base_url"]
|
||||
if apiConfig != nil {
|
||||
if apiConfig.ApiKey != nil {
|
||||
apiKey = *apiConfig.ApiKey
|
||||
}
|
||||
if apiConfig.BaseURL != nil {
|
||||
baseURL = *apiConfig.BaseURL
|
||||
}
|
||||
}
|
||||
|
||||
common.Debug("extractor credentials: resolved",
|
||||
zap.String("llm_factory", provider.ProviderName),
|
||||
zap.String("instance", instance.InstanceName),
|
||||
zap.String("model", canonicalModelName),
|
||||
zap.Bool("api_key_present", true),
|
||||
zap.Bool("base_url_present", baseURL != ""))
|
||||
return extractorChatConfig{
|
||||
driver: provider.ProviderName,
|
||||
modelName: canonicalModelName,
|
||||
apiKey: instance.APIKey,
|
||||
driver: strings.ToLower(driver.Name()),
|
||||
modelName: modelName,
|
||||
apiKey: apiKey,
|
||||
baseURL: baseURL,
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
|
||||
// splitExtractorLLID splits a composite llm_id using rsplit("@", 2)
|
||||
// from the right to preserve embedded '@' in model names.
|
||||
//
|
||||
// "model@provider" → ("model", "default", "provider")
|
||||
// "model@instance@provider" → ("model", "instance", "provider")
|
||||
// "a@b@c@d" → ("a@b", "c", "d")
|
||||
// "plain" → ("plain", "", "")
|
||||
func splitExtractorLLID(s string) (modelName, instanceName, providerName string) {
|
||||
parts := strings.SplitN(strings.TrimSpace(s), "@", -1)
|
||||
// rsplit("@", 2) from the right
|
||||
n := len(parts)
|
||||
if n >= 3 {
|
||||
// Last segment = provider, second-to-last = instance,
|
||||
// everything to the left = model name (rejoined with @).
|
||||
providerName = parts[n-1]
|
||||
instanceName = parts[n-2]
|
||||
modelName = strings.Join(parts[:n-2], "@")
|
||||
return
|
||||
// isBareTenantModelID reports whether s is a 32-character hex string
|
||||
// (a tenant_model primary key), as opposed to a composite "model@provider".
|
||||
func isBareTenantModelID(s string) bool {
|
||||
s = strings.TrimSpace(s)
|
||||
if len(s) != 32 {
|
||||
return false
|
||||
}
|
||||
if n == 2 {
|
||||
return parts[0], "default", parts[1]
|
||||
}
|
||||
return s, "", ""
|
||||
}
|
||||
|
||||
// findExtractorSoleActiveInstance returns the sole active instance
|
||||
// for a provider when the "default" instance is not found.
|
||||
func findExtractorSoleActiveInstance(providerID string) *entity.TenantModelInstance {
|
||||
instances, err := dao.NewTenantModelInstanceDAO().GetAllInstancesByProviderID(providerID)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var active []*entity.TenantModelInstance
|
||||
for _, inst := range instances {
|
||||
if inst == nil {
|
||||
continue
|
||||
for _, r := range s {
|
||||
if !((r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') || (r >= 'A' && r <= 'F')) {
|
||||
return false
|
||||
}
|
||||
if strings.EqualFold(strings.TrimSpace(inst.Status), "inactive") {
|
||||
continue
|
||||
}
|
||||
active = append(active, inst)
|
||||
}
|
||||
if len(active) != 1 {
|
||||
return nil
|
||||
}
|
||||
return active[0]
|
||||
return true
|
||||
}
|
||||
|
||||
// buildExtractorMessages assembles system + user messages for
|
||||
|
||||
@@ -171,7 +171,10 @@ func (c *ExtractorComponent) runAutoTags(ctx context.Context, in extractorInputs
|
||||
}
|
||||
|
||||
if len(docsToTag) > 0 && in.llmID != "" {
|
||||
driver, model, apiKey, baseURL := resolveExtractorChatTarget(ctx, in.llmID)
|
||||
driver, model, apiKey, baseURL, err := resolveExtractorChatTarget(ctx, in.llmID)
|
||||
if err != nil {
|
||||
common.Warn("extractor tag: resolve model failed, skipping LLM tagging", zap.Error(err))
|
||||
}
|
||||
if driver != "" && model != "" {
|
||||
inv := getExtractorChatInvoker()
|
||||
sem := make(chan struct{}, taggerLLMConcurrency)
|
||||
|
||||
@@ -506,6 +506,88 @@ func TestExtractorComponent_NewExtractorComponent_Happy(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_SysPromptAlias verifies that "sys_prompt"
|
||||
// (the Python DSL name) is accepted as a fallback for SystemPrompt.
|
||||
func TestNewExtractorComponent_SysPromptAlias(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"sys_prompt": "You are a Python DSL prompt.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
if ec.Param.SystemPrompt != "You are a Python DSL prompt." {
|
||||
t.Errorf("SystemPrompt = %q, want %q", ec.Param.SystemPrompt, "You are a Python DSL prompt.")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_PromptsArray verifies that the Python DSL
|
||||
// "prompts" array format is parsed into Param.Prompt.
|
||||
func TestNewExtractorComponent_PromptsArray(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"prompts": []any{
|
||||
map[string]any{
|
||||
"content": "Analyze: {TitleChunker:FlatMiceFix@chunks}",
|
||||
"role": "user",
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
want := "Analyze: {TitleChunker:FlatMiceFix@chunks}"
|
||||
if ec.Param.Prompt != want {
|
||||
t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_PromptsArray_PromptWins verifies that
|
||||
// "prompt" (string) takes priority over "prompts" (array) when both
|
||||
// are present in the DSL params.
|
||||
func TestNewExtractorComponent_PromptsArray_PromptWins(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"prompt": "Direct prompt wins.",
|
||||
"prompts": []any{
|
||||
map[string]any{
|
||||
"content": "Should be ignored.",
|
||||
"role": "user",
|
||||
},
|
||||
},
|
||||
})
|
||||
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) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"system_prompt": "system_prompt wins.",
|
||||
"sys_prompt": "sys_prompt ignored.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
if ec.Param.SystemPrompt != "system_prompt wins." {
|
||||
t.Errorf("SystemPrompt = %q, want %q", ec.Param.SystemPrompt, "system_prompt wins.")
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractorComponent_InputsOutputs_NonEmpty is the shape
|
||||
// assertion Phase 4's API endpoint relies on.
|
||||
func TestExtractorComponent_InputsOutputs_NonEmpty(t *testing.T) {
|
||||
@@ -637,3 +719,60 @@ func TestExtractorComponent_ConcurrentInvoke(t *testing.T) {
|
||||
// (it currently isn't, but pinning the import keeps test-side
|
||||
// imports honest if helpers move around in future revisions).
|
||||
var _ = eschema.Message{}
|
||||
|
||||
// TestIsBareTenantModelID verifies UUID detection.
|
||||
func TestIsBareTenantModelID(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want bool
|
||||
}{
|
||||
{"9e819c2442b14f9dab46062916e29195", true},
|
||||
{"ABCDEFabcdef01234567890123456789", true},
|
||||
{"9e819c2442b14f9dab46062916e2919", false}, // 31 chars
|
||||
{"9e819c2442b14f9dab46062916e29195X", false}, // 33 chars
|
||||
{"gpt-4o-mini@openai", false},
|
||||
{"", false},
|
||||
{"not-a-uuid", false},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
got := isBareTenantModelID(tc.input)
|
||||
if got != tc.want {
|
||||
t.Errorf("isBareTenantModelID(%q) = %v, want %v", tc.input, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveExtractorChatTarget_AtSplitFallback verifies the @ split
|
||||
// fallback path works without canvas state (unit test compatibility).
|
||||
func TestResolveExtractorChatTarget_AtSplitFallback(t *testing.T) {
|
||||
driver, modelName, apiKey, baseURL, err := resolveExtractorChatTarget(
|
||||
context.Background(), "gpt-4o-mini@openai")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if driver != "openai" {
|
||||
t.Errorf("driver = %q, want openai", driver)
|
||||
}
|
||||
if modelName != "gpt-4o-mini" {
|
||||
t.Errorf("modelName = %q, want gpt-4o-mini", modelName)
|
||||
}
|
||||
if apiKey != "" || baseURL != "" {
|
||||
t.Errorf("apiKey/baseURL should be empty in fallback path")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveExtractorChatTarget_NoDriver verifies a non-@ plain string
|
||||
// without canvas state returns no driver (passes through to Chat()).
|
||||
func TestResolveExtractorChatTarget_NoDriver(t *testing.T) {
|
||||
driver, modelName, _, _, err := resolveExtractorChatTarget(
|
||||
context.Background(), "plain-name")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if driver != "" {
|
||||
t.Errorf("driver should be empty for plain name, got %q", driver)
|
||||
}
|
||||
if modelName != "plain-name" {
|
||||
t.Errorf("modelName = %q, want plain-name", modelName)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -99,6 +99,7 @@
|
||||
"preprocess": [
|
||||
"main_content"
|
||||
],
|
||||
"remove_toc": true,
|
||||
"suffix": [
|
||||
"pdf"
|
||||
],
|
||||
|
||||
@@ -146,8 +146,7 @@
|
||||
"^\\s*(?:\\d+[\\.、\\)]\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\s*[::]?\\s*$"
|
||||
]
|
||||
],
|
||||
"method": "hierarchy",
|
||||
"root_chunk_as_heading": true
|
||||
"method": "hierarchy"
|
||||
}
|
||||
},
|
||||
"upstream": [
|
||||
@@ -352,7 +351,6 @@
|
||||
}
|
||||
},
|
||||
"promote_first_heading_to_root": false,
|
||||
"root_chunk_as_heading": true,
|
||||
"rules": [
|
||||
{
|
||||
"levels": [
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -139,7 +139,7 @@ func extractDOCXFiguresFromIR(irJSON string) []DOCXFigure {
|
||||
flat = append(flat, flatBlock{image: b64})
|
||||
continue
|
||||
}
|
||||
text := joinDOCXIRRuns(el.Content)
|
||||
text := joinDOCXIRRuns(el.contentRuns())
|
||||
flat = append(flat, flatBlock{text: text})
|
||||
}
|
||||
}
|
||||
@@ -237,6 +237,44 @@ func joinDOCXIRRuns(runs []docxIRRun) string {
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// extractTextFromListItem extracts the plain text content from a list item.
|
||||
// Each list item contains block-level elements (typically a Paragraph),
|
||||
// whose text runs are concatenated.
|
||||
func extractTextFromListItem(item docxIRListItem) string {
|
||||
var parts []string
|
||||
for _, el := range item.Content {
|
||||
if el.Type == "paragraph" || el.Type == "heading" {
|
||||
t := joinDOCXIRRuns(el.contentRuns())
|
||||
if t != "" {
|
||||
parts = append(parts, t)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(strings.Join(parts, "\n"))
|
||||
}
|
||||
|
||||
// extractTextFromBlockElements extracts text from a slice of block-level
|
||||
// elements (paragraphs/headings), used by text_box and other compound
|
||||
// element types.
|
||||
func extractTextFromBlockElements(blocks []docxIRElement) string {
|
||||
var parts []string
|
||||
for _, el := range blocks {
|
||||
if el.Type == "paragraph" || el.Type == "heading" {
|
||||
t := joinDOCXIRRuns(el.contentRuns())
|
||||
if t != "" {
|
||||
parts = append(parts, t)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(strings.Join(parts, "\n"))
|
||||
}
|
||||
|
||||
// buildFiguresMap converts the internal DOCXFigure slice to the
|
||||
// map form attached to fileMeta["figures"].
|
||||
func buildFiguresMap(figures []DOCXFigure) []map[string]any {
|
||||
@@ -257,7 +295,7 @@ func buildFiguresMap(figures []DOCXFigure) []map[string]any {
|
||||
func joinCellText(cell docxIRCell) string {
|
||||
var parts []string
|
||||
for _, el := range cell.Content {
|
||||
if text := joinDOCXIRRuns(el.Content); text != "" {
|
||||
if text := joinDOCXIRRuns(el.contentRuns()); text != "" {
|
||||
parts = append(parts, text)
|
||||
}
|
||||
}
|
||||
@@ -294,7 +332,7 @@ func buildDOCXJSONSections(irJSON string) []map[string]any {
|
||||
for _, el := range sec.Elements {
|
||||
switch el.Type {
|
||||
case "paragraph", "heading":
|
||||
text := joinDOCXIRRuns(el.Content)
|
||||
text := joinDOCXIRRuns(el.contentRuns())
|
||||
if strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
@@ -326,6 +364,30 @@ func buildDOCXJSONSections(irJSON string) []map[string]any {
|
||||
"image": nil,
|
||||
"doc_type_kwd": "table",
|
||||
})
|
||||
|
||||
case "list":
|
||||
for _, item := range el.Items {
|
||||
text := extractTextFromListItem(item)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
sections = append(sections, map[string]any{
|
||||
"text": text,
|
||||
"image": nil,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
|
||||
case "text_box":
|
||||
text := extractTextFromBlockElements(el.contentBlocks())
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
sections = append(sections, map[string]any{
|
||||
"text": text,
|
||||
"image": nil,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -344,12 +406,38 @@ type docxIRSection struct {
|
||||
}
|
||||
|
||||
type docxIRElement struct {
|
||||
Type string `json:"type"` // "paragraph", "heading", "table", "image"
|
||||
Level int `json:"level"` // heading level (1-6)
|
||||
Style string `json:"style"` // Word style name (e.g. "Normal", "Heading 1")
|
||||
Content []docxIRRun `json:"content"` // rich text runs
|
||||
Data []byte `json:"data"` // raw image bytes (for "image" type)
|
||||
Rows []docxIRRow `json:"rows"` // table rows
|
||||
Type string `json:"type"` // "paragraph", "heading", "table", "image", "list", "text_box", ...
|
||||
Level int `json:"level"` // heading level (1-6) or list nesting level
|
||||
Style string `json:"style"` // Word style name (e.g. "Normal", "Heading 1")
|
||||
Content json.RawMessage `json:"content"` // rich text runs or block-level content; decoded per type
|
||||
Data []byte `json:"data"` // raw image bytes (for "image" type)
|
||||
Rows []docxIRRow `json:"rows"` // table rows
|
||||
Ordered bool `json:"ordered"` // true=numbered list, false=bullet list (for "list" type)
|
||||
Items []docxIRListItem `json:"items"` // list items (for "list" type)
|
||||
}
|
||||
|
||||
// contentRuns decodes Content as flat text runs (paragraph/heading type).
|
||||
func (e docxIRElement) contentRuns() []docxIRRun {
|
||||
var runs []docxIRRun
|
||||
if len(e.Content) > 0 {
|
||||
_ = json.Unmarshal(e.Content, &runs)
|
||||
}
|
||||
return runs
|
||||
}
|
||||
|
||||
// contentBlocks decodes Content as block-level elements (text_box type).
|
||||
func (e docxIRElement) contentBlocks() []docxIRElement {
|
||||
var blocks []docxIRElement
|
||||
if len(e.Content) > 0 {
|
||||
_ = json.Unmarshal(e.Content, &blocks)
|
||||
}
|
||||
return blocks
|
||||
}
|
||||
|
||||
// docxIRListItem represents one item in an ordered/unordered list.
|
||||
type docxIRListItem struct {
|
||||
Content []docxIRElement `json:"content"` // block-level content (typically a single Paragraph)
|
||||
Nested json.RawMessage `json:"nested,omitempty"` // nested sub-list (stored as raw JSON for now)
|
||||
}
|
||||
|
||||
type docxIRRun struct {
|
||||
|
||||
@@ -20,10 +20,16 @@ package parser
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func rawJSON(v any) json.RawMessage {
|
||||
data, _ := json.Marshal(v)
|
||||
return json.RawMessage(data)
|
||||
}
|
||||
|
||||
func TestBuildDOCXJSONSections_Paragraphs(t *testing.T) {
|
||||
ir := `{"sections":[{"title":"","elements":[
|
||||
{"type":"paragraph","content":[{"type":"text","text":"Hello world"}],"style":"Normal"}
|
||||
@@ -144,6 +150,65 @@ func TestBuildDOCXJSONSections_MixedContent(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDOCXJSONSections_List(t *testing.T) {
|
||||
ir := `{"sections":[{"title":"","elements":[
|
||||
{"type":"list","ordered":false,"items":[
|
||||
{"content":[{"type":"paragraph","content":[{"type":"text","text":"Item 1"}]}]},
|
||||
{"content":[{"type":"paragraph","content":[{"type":"text","text":"Item 2"}]}]},
|
||||
{"content":[{"type":"paragraph","content":[{"type":"text","text":"Item 3"}]}]}
|
||||
],"level":0}
|
||||
]}]}`
|
||||
sections := buildDOCXJSONSections(ir)
|
||||
if len(sections) != 3 {
|
||||
t.Fatalf("got %d sections, want 3 (each list item should be a section)", len(sections))
|
||||
}
|
||||
for i, want := range []string{"Item 1", "Item 2", "Item 3"} {
|
||||
if got, _ := sections[i]["text"].(string); got != want {
|
||||
t.Errorf("section[%d].text = %q, want %q", i, got, want)
|
||||
}
|
||||
if got, ok := sections[i]["doc_type_kwd"].(string); !ok || got != "text" {
|
||||
t.Errorf("section[%d].doc_type_kwd = %q, want %q", i, got, "text")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDOCXJSONSections_TextBox(t *testing.T) {
|
||||
ir := `{"sections":[{"title":"","elements":[
|
||||
{"type":"text_box","content":[{"type":"paragraph","content":[{"type":"text","text":"Boxed paragraph"}]}],"width_emu":null}
|
||||
]}]}`
|
||||
sections := buildDOCXJSONSections(ir)
|
||||
if len(sections) != 1 {
|
||||
t.Fatalf("got %d sections, want 1 (text_box content should become a section)", len(sections))
|
||||
}
|
||||
if got, _ := sections[0]["text"].(string); got != "Boxed paragraph" {
|
||||
t.Errorf("text = %q, want %q", got, "Boxed paragraph")
|
||||
}
|
||||
if got, ok := sections[0]["doc_type_kwd"].(string); !ok || got != "text" {
|
||||
t.Errorf("doc_type_kwd = %q, want %q", got, "text")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDOCXJSONSections_MixedWithList(t *testing.T) {
|
||||
ir := `{"sections":[{"title":"","elements":[
|
||||
{"type":"paragraph","content":[{"type":"text","text":"Preamble"}],"style":"Normal"},
|
||||
{"type":"list","ordered":false,"items":[
|
||||
{"content":[{"type":"paragraph","content":[{"type":"text","text":"Bullet A"}]}]},
|
||||
{"content":[{"type":"paragraph","content":[{"type":"text","text":"Bullet B"}]}]}
|
||||
],"level":0},
|
||||
{"type":"paragraph","content":[{"type":"text","text":"Trailer"}],"style":"Normal"}
|
||||
]}]}`
|
||||
sections := buildDOCXJSONSections(ir)
|
||||
if len(sections) != 4 {
|
||||
t.Fatalf("got %d sections, want 4 (para + 2 list items + para)", len(sections))
|
||||
}
|
||||
wants := []string{"Preamble", "Bullet A", "Bullet B", "Trailer"}
|
||||
for i, want := range wants {
|
||||
if got, _ := sections[i]["text"].(string); got != want {
|
||||
t.Errorf("section[%d].text = %q, want %q", i, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildDOCXJSONSections_EmptySkipped(t *testing.T) {
|
||||
ir := `{"sections":[{"title":"","elements":[
|
||||
{"type":"paragraph","content":[{"type":"text","text":""}]},
|
||||
@@ -174,7 +239,7 @@ func TestDocxIRTableToHTML_Single(t *testing.T) {
|
||||
Cells: []docxIRCell{{
|
||||
Content: []docxIRElement{{
|
||||
Type: "paragraph",
|
||||
Content: []docxIRRun{{Type: "text", Text: "hello"}},
|
||||
Content: json.RawMessage(`[{"type":"text","text":"hello"}]`),
|
||||
}},
|
||||
}},
|
||||
}},
|
||||
@@ -191,7 +256,7 @@ func TestDocxIRTableToHTML_MultiRowCol(t *testing.T) {
|
||||
return docxIRCell{
|
||||
Content: []docxIRElement{{
|
||||
Type: "paragraph",
|
||||
Content: []docxIRRun{{Type: "text", Text: text}},
|
||||
Content: rawJSON([]docxIRRun{{Type: "text", Text: text}}),
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -114,6 +114,12 @@ func GetFileType(filename string) FileType {
|
||||
case "mp4", "avi", "mkv", "mov", "webm", "flv", "mpeg",
|
||||
"mpg", "wmv", "3gp", "3gpp":
|
||||
return FileTypeVIDEO
|
||||
case "png", "jpg", "jpeg", "gif", "bmp":
|
||||
return FileTypeVISUAL
|
||||
case "tiff", "tif", "webp", "svg", "ico":
|
||||
return FileTypeVISUAL
|
||||
case "avif", "heic", "apng":
|
||||
return FileTypeVISUAL
|
||||
default:
|
||||
return FileTypeOTHER
|
||||
}
|
||||
|
||||
128
internal/utility/file_test.go
Normal file
128
internal/utility/file_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 utility
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetFileType_ImageExtensions(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
filename string
|
||||
want FileType
|
||||
}{
|
||||
// Image formats that should resolve to VISUAL.
|
||||
{"png", "image.png", FileTypeVISUAL},
|
||||
{"jpg", "photo.jpg", FileTypeVISUAL},
|
||||
{"jpeg", "photo.jpeg", FileTypeVISUAL},
|
||||
{"gif", "animation.gif", FileTypeVISUAL},
|
||||
{"bmp", "bitmap.bmp", FileTypeVISUAL},
|
||||
{"tiff", "image.tiff", FileTypeVISUAL},
|
||||
{"tif", "image.tif", FileTypeVISUAL},
|
||||
{"webp", "image.webp", FileTypeVISUAL},
|
||||
{"svg", "vector.svg", FileTypeVISUAL},
|
||||
{"ico", "icon.ico", FileTypeVISUAL},
|
||||
{"avif", "image.avif", FileTypeVISUAL},
|
||||
{"heic", "photo.heic", FileTypeVISUAL},
|
||||
{"apng", "animation.apng", FileTypeVISUAL},
|
||||
|
||||
// Capitalised extension should still match.
|
||||
{"png uppercase", "image.PNG", FileTypeVISUAL},
|
||||
{"jpg mixed case", "photo.JpG", FileTypeVISUAL},
|
||||
|
||||
// Path with directory should still resolve.
|
||||
{"png with path", "/path/to/image.png", FileTypeVISUAL},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := GetFileType(tc.filename)
|
||||
if got != tc.want {
|
||||
t.Errorf("GetFileType(%q) = %q, want %q", tc.filename, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetFileType_ExistingFormats_NoRegression(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
filename string
|
||||
want FileType
|
||||
}{
|
||||
{"pdf", "doc.pdf", FileTypePDF},
|
||||
{"doc", "old.doc", FileTypeDOC},
|
||||
{"docx", "report.docx", FileTypeDOCX},
|
||||
{"xls", "spreadsheet.xls", FileTypeXLS},
|
||||
{"xlsx", "spreadsheet.xlsx", FileTypeXLSX},
|
||||
{"csv", "data.csv", FileTypeCSV},
|
||||
{"ppt", "slides.ppt", FileTypePPT},
|
||||
{"pptx", "slides.pptx", FileTypePPTX},
|
||||
{"html", "page.html", FileTypeHTML},
|
||||
{"htm", "page.htm", FileTypeHTML},
|
||||
{"md", "readme.md", FileTypeMarkdown},
|
||||
{"markdown", "readme.markdown", FileTypeMarkdown},
|
||||
{"txt", "notes.txt", FileTypeTXT},
|
||||
{"py", "script.py", FileTypeTXT},
|
||||
{"js", "script.js", FileTypeTXT},
|
||||
{"go", "main.go", FileTypeTXT},
|
||||
{"java", "Main.java", FileTypeTXT},
|
||||
{"epub", "book.epub", FileTypeEPUB},
|
||||
{"json", "data.json", FileTypeJSON},
|
||||
{"jsonl", "data.jsonl", FileTypeJSON},
|
||||
{"eml", "email.eml", FileTypeEMAIL},
|
||||
{"msg", "email.msg", FileTypeEMAIL},
|
||||
{"mp3", "audio.mp3", FileTypeAURAL},
|
||||
{"wav", "audio.wav", FileTypeAURAL},
|
||||
{"flac", "audio.flac", FileTypeAURAL},
|
||||
{"mp4", "video.mp4", FileTypeVIDEO},
|
||||
{"avi", "video.avi", FileTypeVIDEO},
|
||||
{"mkv", "video.mkv", FileTypeVIDEO},
|
||||
{"mov", "video.mov", FileTypeVIDEO},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := GetFileType(tc.filename)
|
||||
if got != tc.want {
|
||||
t.Errorf("GetFileType(%q) = %q, want %q (regression)", tc.filename, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetFileType_UnknownExtension_ReturnsOTHER(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
filename string
|
||||
}{
|
||||
{"no extension", "Makefile"},
|
||||
{"unknown extension", "data.xyz"},
|
||||
{"dotfile", ".hidden"},
|
||||
{"empty string", ""},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := GetFileType(tc.filename)
|
||||
if got != FileTypeOTHER {
|
||||
t.Errorf("GetFileType(%q) = %q, want %q (FileTypeOTHER)", tc.filename, got, FileTypeOTHER)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user