From 848d21fc3034fa1c9a80cfa2e2ac01d19ffb9733 Mon Sep 17 00:00:00 2001 From: Wang Qi Date: Fri, 23 Oct 2026 13:50:09 +0800 Subject: [PATCH] Revert "feat(ingestion): generate TOC chunks in the Go Extractor" (#21385) --- internal/ingestion/component/extractor.go | 37 +- internal/ingestion/component/extractor_toc.go | 374 ------------------ .../ingestion/component/extractor_toc_test.go | 316 --------------- .../ingestion/component/schema/extractor.go | 9 +- internal/rag/prompts/assign_toc_levels.md | 55 --- internal/rag/prompts/templates.go | 2 +- internal/rag/prompts/templates_test.go | 3 - internal/rag/prompts/toc_from_text_system.md | 119 ------ internal/rag/prompts/toc_from_text_user.md | 8 - internal/service/document/document_test.go | 33 +- internal/service/document/wiki_refresh.go | 12 +- internal/service/toc_enhancer_test.go | 14 - 12 files changed, 10 insertions(+), 972 deletions(-) delete mode 100644 internal/ingestion/component/extractor_toc.go delete mode 100644 internal/ingestion/component/extractor_toc_test.go delete mode 100644 internal/rag/prompts/assign_toc_levels.md delete mode 100644 internal/rag/prompts/toc_from_text_system.md delete mode 100644 internal/rag/prompts/toc_from_text_user.md diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 2a0a236c7f..3c13debdeb 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -59,7 +59,6 @@ import ( "ragflow/internal/entity" "ragflow/internal/entity/models" "ragflow/internal/ingestion/chunkcache" - "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/progress" "ragflow/internal/ingestion/component/schema" "ragflow/internal/tokenizer" @@ -234,7 +233,6 @@ type ExtractorComponent struct { // "tags": map[string]any, — optional auto tags extraction config // "summary": map[string]any, — optional auto summary extraction config // "metadata": map[string]any, — optional auto metadata extraction config -// "toc": map[string]any, - optional table-of-contents extraction config // } // // The keywords/questions/tags groups accept an "enabled" switch. A group @@ -323,15 +321,6 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) { p.Metadata.BuiltInMetadata = parseMetadataFieldDefs(v) } } - - // 6. TOC - if tocRaw, ok := params["toc"].(map[string]any); ok { - if v, ok := tocRaw["enabled"].(bool); ok { - p.TOC.Enabled = v - } else if v, ok := tocRaw["enabled"]; ok { - p.TOC.Enabled = mapInt(v) == 1 - } - } } if err := p.Validate(); err != nil { return nil, fmt.Errorf("extractor: param check: %w", err) @@ -395,7 +384,7 @@ func (c *ExtractorComponent) Outputs() map[string]string { return map[string]string{ "chunks": "Extraction results — input chunks, each enriched with modular extraction fields (important_kwd, question_kwd, tag_feas, summary, metadata).", "output_format": "Always \"chunks\". Parity marker for downstream consumers.", - "_ERROR": "Optional short-circuit error message.", + "_ERROR": "Optional short-circuit error message (reserved for the future TOC branch and other error paths).", } } @@ -709,18 +698,7 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map keywordsOn := c.Param.Keywords.Enabled tagsOn := c.Param.Tags.Enabled remainingOn := c.Param.Questions.Enabled || c.Param.Summary.Enabled || c.Param.Metadata.Enabled - tocOn := c.Param.TOC.Enabled - phaseWindows := progress.NewSplit(keywordsOn, tagsOn, remainingOn, tocOn) - - // The TOC chunk id derives from doc_id, and the retrieval path finds the - // TOC by doc_id, so a missing doc_id fails before any model call. - var docID string - if tocOn { - docID = globals.GlobalOrInput(ctx, inputs, "doc_id", "") - if docID == "" { - return nil, fmt.Errorf("extractor: toc: doc_id is required") - } - } + phaseWindows := progress.NewSplit(keywordsOn, tagsOn, remainingOn) // No node-level wall-clock budget (see the package doc): each LLM call // carries its own budget, and an expired call skips its chunk instead @@ -742,17 +720,6 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map if err := c.runRemainingExtractions(ctx, db, in, phaseWindows.Take(remainingOn)); err != nil { return nil, fmt.Errorf("extractor: %w", err) } - - // Phase 4: one TOC for the whole document, appended as an extra chunk. - if tocOn { - tocChunk, err := c.buildTOCChunk(ctx, db, in, docID, phaseWindows.Take(tocOn)) - if err != nil { - return nil, fmt.Errorf("extractor: toc: %w", err) - } - if tocChunk != nil { - in.chunks = append(in.chunks, tocChunk) - } - } // The tag phase swallows per-chunk failures, so a task stopped mid-run // could otherwise report success. This mirrors the old node-level // WithTimeout post-check: never report success on a dead context. diff --git a/internal/ingestion/component/extractor_toc.go b/internal/ingestion/component/extractor_toc.go deleted file mode 100644 index 494719f34e..0000000000 --- a/internal/ingestion/component/extractor_toc.go +++ /dev/null @@ -1,374 +0,0 @@ -// -// 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" - "encoding/json" - "fmt" - "math" - "regexp" - "strconv" - "strings" - "sync/atomic" - - "go.uber.org/zap" - "gorm.io/gorm" - - "ragflow/internal/common" - kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" - "ragflow/internal/ingestion/component/progress" - "ragflow/internal/rag/prompts" - "ragflow/internal/tokenizer" - "ragflow/internal/utility" -) - -// TOC extraction builds one table of contents per document. The retrieval -// path reads it in internal/service/toc_enhancer.go. -const ( - tocBatchTokenCap = 1024 - tocInputUtilization = 0.5 - tocMaxTitles = 512 - tocTitleMaxWords = 22 - tocTitleMaxWordsPrune = 12 - tocJSONAttempts = 2 - // tocPageNum sorts the TOC chunk after every content chunk. - tocPageNum = 100000000 - // tocChunkIDSeed derives the TOC chunk id from the document id alone, so - // a retried run overwrites the TOC row instead of adding a second one. The - // NUL byte keeps the seed apart from the text of any real chunk. - tocChunkIDSeed = "\x00toc" -) - -var tocNumberOnlyTitleRE = regexp.MustCompile(`^[0-9,.()/ -]+$`) - -// tocEntry is one stored TOC item. The stored level is a numeric string. -type tocEntry struct { - Level string `json:"level"` - Title string `json:"title"` - IDs []string `json:"ids,omitempty"` -} - -// tocTitle is a heading the model found, with the chunk index it reported. -type tocTitle struct { - title string - chunkID any -} - -// buildTOCChunk returns the TOC chunk for the document, or nil when the -// model finds no headings. -func (c *ExtractorComponent) buildTOCChunk(ctx context.Context, db *gorm.DB, in extractorInputs, docID string, w progress.Window) (map[string]any, error) { - texts := make([]string, len(in.chunks)) - for i, ck := range in.chunks { - texts[i] = extractorChunkText(ck) - } - entries, err := c.runTOCFromText(ctx, db, in, texts, w) - if err != nil { - return nil, err - } - if len(entries) == 0 { - return nil, nil - } - var buf bytes.Buffer - enc := json.NewEncoder(&buf) - enc.SetEscapeHTML(false) - if err := enc.Encode(entries); err != nil { - return nil, err - } - return map[string]any{ - "text": strings.TrimSuffix(buf.String(), "\n"), - "id": common.ChunkID(docID, tocChunkIDSeed), - "toc_kwd": "toc", - "available_int": 0, - "page_num_int": []int{tocPageNum}, - }, nil -} - -// runTOCFromText finds the headings in batches of chunks, then asks the model -// for the level of each heading. Each entry covers the chunks from its own -// chunk up to the chunk of the next entry. The last entry runs to the last -// chunk. -func (c *ExtractorComponent) runTOCFromText(ctx context.Context, db *gorm.DB, in extractorInputs, texts []string, w progress.Window) ([]tocEntry, error) { - loader := prompts.EmbeddedPromptLoader{} - systemPrompt, err := loader.Load("toc_from_text_system") - if err != nil { - return nil, err - } - userTemplate, err := loader.Load("toc_from_text_user") - if err != nil { - return nil, err - } - levelsPrompt, err := loader.Load("assign_toc_levels") - if err != nil { - return nil, err - } - - budget := tocBatchTokenCap - if ctxLen := extractorContextLength(ctx, db, in.llmID); ctxLen > 0 { - budget = min(int(float64(ctxLen)*tocInputUtilization)-tokenizer.NumTokensFromString(userTemplate+systemPrompt), tocBatchTokenCap) - if budget <= 0 { - return nil, fmt.Errorf("model context of %d tokens leaves no room for chunk text", ctxLen) - } - } - batches, err := tocBatchPrompts(texts, budget) - if err != nil { - return nil, err - } - - temp := 0.0 - tin := extractorInputs{llmID: in.llmID, modelID: in.modelID, temperature: &temp} - results := make([][]any, len(batches)) - var skipped atomic.Int64 - futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, len(batches)) - for i, batch := range batches { - i := i - userPrompt := strings.Replace(userTemplate, "{{text}}", batch, 1) - fn := func() error { - parsed, err := c.genTOCJSON(ctx, db, tin, systemPrompt, userPrompt) - if err != nil { - return resolveChunkCallError(ctx, i, "toc", fmt.Errorf("batch %d: %w", i, err), false, &skipped) - } - results[i], _ = parsed.([]any) - return nil - } - f, err := extractorPool.Submit(ctx, fn) - if err != nil { - return nil, err - } - futs = append(futs, f) - } - if err := awaitFutures(ctx, futs, w); err != nil { - return nil, err - } - if err := checkAllChunksSkipped("toc", len(futs), skipped.Load()); err != nil { - return nil, err - } - - var raw []any - for _, r := range results { - raw = append(raw, r...) - } - titles, err := filterTOCTitles(raw, tokenizer.New(in.lang)) - if err != nil { - return nil, err - } - if len(titles) == 0 { - return nil, nil - } - - names := make([]string, len(titles)) - for i, t := range titles { - names[i] = t.title - } - namesJSON, err := json.Marshal(names) - if err != nil { - return nil, err - } - parsed, err := c.genTOCJSON(ctx, db, tin, levelsPrompt, string(namesJSON)) - if err != nil { - return nil, err - } - levels, _ := parsed.([]any) - return mergeTOCLevels(levels, titles, in.chunks), nil -} - -// tocBatchPrompts packs chunks into batches of about budget tokens, one -// {"": text} JSON line per chunk. A chunk is never split, even when it -// alone exceeds the budget. -func tocBatchPrompts(texts []string, budget int) ([]string, error) { - var batches []string - var lines []string - tokens := 0 - for i, text := range texts { - t := tokenizer.NumTokensFromString(text) - if len(lines) > 0 && tokens+t > budget { - batches = append(batches, strings.Join(lines, "\n")) - lines, tokens = nil, 0 - } - var buf bytes.Buffer - enc := json.NewEncoder(&buf) - enc.SetEscapeHTML(false) - if err := enc.Encode(map[string]string{strconv.Itoa(i): text}); err != nil { - return nil, err - } - lines = append(lines, strings.TrimSuffix(buf.String(), "\n")) - tokens += t - } - if len(lines) > 0 { - batches = append(batches, strings.Join(lines, "\n")) - } - return batches, nil -} - -// genTOCJSON parses the model reply as JSON. When the reply is not valid JSON, -// it asks once more with the reply and the parse error appended. It returns -// nil when no attempt parses. -func (c *ExtractorComponent) genTOCJSON(ctx context.Context, db *gorm.DB, in extractorInputs, systemPrompt, userPrompt string) (any, error) { - var errs string - for attempt := 0; attempt < tocJSONAttempts; attempt++ { - ans, err := c.callText(ctx, db, in, systemPrompt, userPrompt) - if err != nil { - return nil, err - } - repaired, err := kccommon.RepairJSONText(ans) - if err == nil { - var out any - if err = json.Unmarshal([]byte(repaired), &out); err == nil { - return out, nil - } - } - errs += err.Error() - userPrompt += "\nGenerated JSON is as following:\n" + ans + "\nBut exception while loading:\n" + errs + "\nPlease reconsider and correct it." - } - common.Warn("extractor: toc: model reply is not valid JSON; ignoring it", - zap.String("component", componentNameExtractor), - zap.String("error", errs), - ) - return nil, nil -} - -// filterTOCTitles drops empty titles, long sentences and number-only titles -// from the model output. The number-only rule also drops the "-1" placeholder -// the prompt asks for on chunks with no heading. -func filterTOCTitles(raw []any, tok tokenizer.Tokenizer) ([]tocTitle, error) { - maxWords := tocTitleMaxWords - if len(raw) > tocMaxTitles { - maxWords = tocTitleMaxWordsPrune - } - var out []tocTitle - for _, item := range raw { - m, ok := item.(map[string]any) - if !ok { - continue - } - title, _ := m["title"].(string) - if title == "" { - continue - } - tks, err := tok.Tokenize(title) - if err != nil { - return nil, fmt.Errorf("tokenize title: %w", err) - } - if len(strings.Split(tks, " ")) > maxWords { - continue - } - if tocNumberOnlyTitleRE.MatchString(title) { - continue - } - out = append(out, tocTitle{title: title, chunkID: m["chunk_id"]}) - } - return out, nil -} - -// mergeTOCLevels pairs each level item with the title at the same index and -// maps the reported chunk index to chunk ids. An item with an unusable level -// is dropped, because the TOC reader rejects the whole TOC on a non-numeric -// level. -func mergeTOCLevels(levels []any, titles []tocTitle, chunks []map[string]any) []tocEntry { - type item struct { - level int - title string - src tocTitle - } - var items []item - for i := 0; i < len(levels) && i < len(titles); i++ { - var levelRaw, title any - switch v := levels[i].(type) { - case map[string]any: - levelRaw, title = v["level"], v["title"] - case []any: - if len(v) < 2 { - continue - } - levelRaw, title = v[0], v[1] - default: - continue - } - level := 0 - if levelRaw != nil { - n, ok := tocInt(levelRaw) - if !ok { - continue - } - level = n - } - titleStr, ok := title.(string) - if !ok { - titleStr = titles[i].title - } - items = append(items, item{level: level, title: titleStr, src: titles[i]}) - } - - maxLevel := math.MinInt - for _, it := range items { - maxLevel = max(maxLevel, it.level) - } - prune := len(items) > tocMaxTitles - var kept []item - for _, it := range items { - if prune && it.level >= maxLevel { - continue - } - kept = append(kept, it) - } - - entries := make([]tocEntry, len(kept)) - for i, it := range kept { - entries[i] = tocEntry{Level: strconv.Itoa(it.level), Title: it.title} - idx, ok := tocInt(it.src.chunkID) - if !ok || idx < 0 || idx >= len(chunks) { - continue - } - entries[i].IDs = appendChunkID(nil, chunks[idx]) - end := len(chunks) - 1 - if i < len(kept)-1 { - next, ok := tocInt(kept[i+1].src.chunkID) - if !ok { - continue - } - end = next - } - for j := idx + 1; j <= end && j < len(chunks); j++ { - entries[i].IDs = appendChunkID(entries[i].IDs, chunks[j]) - } - } - return entries -} - -func appendChunkID(ids []string, ck map[string]any) []string { - if id := chunkCacheID(ck); id != "" { - ids = append(ids, id) - } - return ids -} - -// tocInt reads an integer the model wrote as a JSON number or a string. -func tocInt(v any) (int, bool) { - switch n := v.(type) { - case float64: - if n != math.Trunc(n) { - return 0, false - } - return int(n), true - case string: - i, err := strconv.Atoi(strings.TrimSpace(n)) - return i, err == nil - default: - return 0, false - } -} diff --git a/internal/ingestion/component/extractor_toc_test.go b/internal/ingestion/component/extractor_toc_test.go deleted file mode 100644 index 7baedda139..0000000000 --- a/internal/ingestion/component/extractor_toc_test.go +++ /dev/null @@ -1,316 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package component - -import ( - "context" - "reflect" - "strings" - "sync/atomic" - "testing" - - "ragflow/internal/common" - "ragflow/internal/ingestion/component/schema" - "ragflow/internal/rag/prompts" - "ragflow/internal/tokenizer" -) - -func tocTestChunks(docID string, texts ...string) []map[string]any { - chunks := make([]map[string]any, len(texts)) - for i, text := range texts { - chunks[i] = map[string]any{"text": text, "id": common.ChunkID(docID, text)} - } - return chunks -} - -func useWhitespaceTokenizer(t *testing.T) { - t.Helper() - tokenizer.SetEngineType("infinity") - t.Cleanup(func() { tokenizer.SetEngineType("") }) -} - -func TestExtractorTOC_AppendsTOCChunk(t *testing.T) { - useWhitespaceTokenizer(t) - longTitle := strings.TrimSpace(strings.Repeat("word ", 23)) - stub := withStubChatInvoker(t, - stubResponse{Content: "```json\n[" + - `{"title":"Chapter 1 Intro","chunk_id":"0"},` + - `{"title":"-1","chunk_id":"1"},` + - `{"chunk_id":"1"},` + - `{"title":"` + longTitle + `","chunk_id":"1"},` + - `{"title":"Chapter 2 Methods","chunk_id":"2"},` + - `{"title":"3.1","chunk_id":"3"},` + - `{"title":"Section 2.1","chunk_id":3}` + - "]\n```"}, - stubResponse{Content: `[` + - `{"level":"1","title":"Chapter 1 Intro"},` + - `{"level":"1","title":"Chapter 2 Methods"},` + - `{"level":2,"title":"Section 2.1"}]`}, - ) - chunks := tocTestChunks("doc1", "Chapter 1 Intro", "body one", "Chapter 2 Methods", "Section 2.1 body") - ids := []string{chunks[0]["id"].(string), chunks[1]["id"].(string), chunks[2]["id"].(string), chunks[3]["id"].(string)} - - c := &ExtractorComponent{Param: schema.ExtractorParam{ - LLMID: "gpt-4o-mini", - TOC: schema.TOCExtractConfig{Enabled: true}, - }} - out, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": chunks, "doc_id": "doc1"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - got := out["chunks"].([]map[string]any) - if len(got) != 5 { - t.Fatalf("chunks len = %d, want 4 content chunks + 1 TOC chunk", len(got)) - } - if stub.Calls() != 2 { - t.Fatalf("chat calls = %d, want 1 batch + 1 level call", stub.Calls()) - } - batchPrompt := stub.requests[0].Messages[len(stub.requests[0].Messages)-1].Content - if !strings.Contains(batchPrompt, `{"0":"Chapter 1 Intro"}`+"\n"+`{"1":"body one"}`) { - t.Errorf("batch prompt does not list the chunks as index lines:\n%s", batchPrompt) - } - levelPrompt := stub.requests[1].Messages[len(stub.requests[1].Messages)-1].Content - if levelPrompt != `["Chapter 1 Intro","Chapter 2 Methods","Section 2.1"]` { - t.Errorf("level prompt = %s", levelPrompt) - } - - toc := got[4] - want := `[{"level":"1","title":"Chapter 1 Intro","ids":["` + ids[0] + `","` + ids[1] + `","` + ids[2] + `"]},` + - `{"level":"1","title":"Chapter 2 Methods","ids":["` + ids[2] + `","` + ids[3] + `"]},` + - `{"level":"2","title":"Section 2.1","ids":["` + ids[3] + `"]}]` - if toc["text"] != want { - t.Errorf("TOC text\n got: %v\nwant: %s", toc["text"], want) - } - if toc["id"] != common.ChunkID("doc1", tocChunkIDSeed) { - t.Errorf("TOC id = %v, want the per-document TOC id", toc["id"]) - } - if toc["toc_kwd"] != "toc" || toc["available_int"] != 0 || !reflect.DeepEqual(toc["page_num_int"], []int{100000000}) { - t.Errorf("TOC chunk fields = %+v", toc) - } -} - -// tocChatFunc answers each chat call from its user prompt, because batch -// calls run concurrently and arrive in no fixed order. -type tocChatFunc func(prompt string) string - -func (f tocChatFunc) Chat(_ context.Context, req extractorChatRequest) (*extractorChatResponse, error) { - return &extractorChatResponse{Content: f(req.Messages[len(req.Messages)-1].Content)}, nil -} - -func TestExtractorTOC_KeepsChunkIndexesAcrossBatches(t *testing.T) { - useWhitespaceTokenizer(t) - chunks := tocTestChunks("doc1", "Part one", "text a", "Part two", "text b") - ids := []string{chunks[0]["id"].(string), chunks[1]["id"].(string), chunks[2]["id"].(string), chunks[3]["id"].(string)} - - // Size the context so one batch holds exactly two chunks. - loader := prompts.EmbeddedPromptLoader{} - system, err := loader.Load("toc_from_text_system") - if err != nil { - t.Fatal(err) - } - user, err := loader.Load("toc_from_text_user") - if err != nil { - t.Fatal(err) - } - budget := tokenizer.NumTokensFromString("Part one") + tokenizer.NumTokensFromString("text a") - ctxLen := 2 * (tokenizer.NumTokensFromString(user+system) + budget) - SetExtractorContextLengthOverride(func(context.Context, string) int { return ctxLen }) - t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) - - var batches atomic.Int32 - prev := defaultExtractorChatInvoker - SetExtractorChatInvoker(tocChatFunc(func(prompt string) string { - switch { - case strings.Contains(prompt, `{"0":"Part one"}`): - batches.Add(1) - return `[{"title":"Part one","chunk_id":"0"},{"title":"-1","chunk_id":"1"}]` - case strings.Contains(prompt, `{"2":"Part two"}`): - batches.Add(1) - return `[{"title":"Part two","chunk_id":"2"},{"title":"-1","chunk_id":"3"}]` - case prompt == `["Part one","Part two"]`: - return `[{"level":"1","title":"Part one"},{"level":"1","title":"Part two"}]` - } - t.Errorf("unexpected prompt:\n%s", prompt) - return "[]" - })) - t.Cleanup(func() { SetExtractorChatInvoker(prev) }) - - c := &ExtractorComponent{Param: schema.ExtractorParam{ - LLMID: "gpt-4o-mini", - TOC: schema.TOCExtractConfig{Enabled: true}, - }} - out, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": chunks, "doc_id": "doc1"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if batches.Load() != 2 { - t.Fatalf("batch calls = %d, want 2", batches.Load()) - } - got := out["chunks"].([]map[string]any) - want := `[{"level":"1","title":"Part one","ids":["` + ids[0] + `","` + ids[1] + `","` + ids[2] + `"]},` + - `{"level":"1","title":"Part two","ids":["` + ids[2] + `","` + ids[3] + `"]}]` - if got[len(got)-1]["text"] != want { - t.Errorf("TOC text\n got: %v\nwant: %s", got[len(got)-1]["text"], want) - } -} - -func TestExtractorTOC_RetriesInvalidJSON(t *testing.T) { - useWhitespaceTokenizer(t) - stub := withStubChatInvoker(t, - stubResponse{Content: "no headings here"}, - stubResponse{Content: `[{"title":"Overview","chunk_id":"0"}]`}, - stubResponse{Content: `[{"level":"1","title":"Overview"}]`}, - ) - c := &ExtractorComponent{Param: schema.ExtractorParam{ - LLMID: "gpt-4o-mini", - TOC: schema.TOCExtractConfig{Enabled: true}, - }} - out, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": tocTestChunks("doc1", "Overview"), "doc_id": "doc1"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if stub.Calls() != 3 { - t.Fatalf("chat calls = %d, want 2 batch attempts + 1 level call", stub.Calls()) - } - retry := stub.requests[1].Messages[len(stub.requests[1].Messages)-1].Content - if !strings.Contains(retry, "Generated JSON is as following:\nno headings here\nBut exception while loading:") { - t.Errorf("retry prompt does not carry the invalid reply:\n%s", retry) - } - got := out["chunks"].([]map[string]any) - if len(got) != 2 { - t.Fatalf("chunks len = %d, want 2", len(got)) - } -} - -func TestExtractorTOC_NoHeadingsAppendsNothing(t *testing.T) { - useWhitespaceTokenizer(t) - stub := withStubChatInvoker(t, - stubResponse{Content: `[{"title":"-1","chunk_id":"0"},{"title":"-1","chunk_id":"1"}]`}, - ) - c := &ExtractorComponent{Param: schema.ExtractorParam{ - LLMID: "gpt-4o-mini", - TOC: schema.TOCExtractConfig{Enabled: true}, - }} - out, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": tocTestChunks("doc1", "a", "b"), "doc_id": "doc1"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["chunks"].([]map[string]any); len(got) != 2 { - t.Fatalf("chunks len = %d, want 2 (no TOC chunk)", len(got)) - } - if stub.Calls() != 1 { - t.Fatalf("chat calls = %d, want 1 (no level call without titles)", stub.Calls()) - } -} - -func TestExtractorTOC_DisabledMakesNoCalls(t *testing.T) { - stub := withStubChatInvoker(t) - c := &ExtractorComponent{Param: schema.ExtractorParam{LLMID: "gpt-4o-mini"}} - out, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": tocTestChunks("doc1", "Chapter 1"), "doc_id": "doc1"}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if got := out["chunks"].([]map[string]any); len(got) != 1 { - t.Fatalf("chunks len = %d, want 1", len(got)) - } - if stub.Calls() != 0 { - t.Fatalf("chat calls = %d, want 0", stub.Calls()) - } -} - -func TestExtractorTOC_MissingDocIDMakesNoCalls(t *testing.T) { - stub := withStubChatInvoker(t) - c := &ExtractorComponent{Param: schema.ExtractorParam{ - LLMID: "gpt-4o-mini", - Keywords: schema.KeywordExtractConfig{Enabled: true, TopN: 3}, - TOC: schema.TOCExtractConfig{Enabled: true}, - }} - _, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": tocTestChunks("doc1", "Chapter 1")}) - if err == nil || !strings.Contains(err.Error(), "doc_id is required") { - t.Fatalf("err = %v, want doc_id is required", err) - } - if stub.Calls() != 0 { - t.Fatalf("chat calls = %d, want 0", stub.Calls()) - } -} - -func TestExtractorTOC_ContextTooSmallFails(t *testing.T) { - useWhitespaceTokenizer(t) - stub := withStubChatInvoker(t) - SetExtractorContextLengthOverride(func(context.Context, string) int { return 10 }) - t.Cleanup(func() { SetExtractorContextLengthOverride(nil) }) - c := &ExtractorComponent{Param: schema.ExtractorParam{ - LLMID: "gpt-4o-mini", - TOC: schema.TOCExtractConfig{Enabled: true}, - }} - _, err := c.Invoke(t.Context(), nil, map[string]any{"chunks": tocTestChunks("doc1", "Chapter 1"), "doc_id": "doc1"}) - if err == nil || !strings.Contains(err.Error(), "model context of 10 tokens leaves no room") { - t.Fatalf("err = %v, want a context error", err) - } - if stub.Calls() != 0 { - t.Fatalf("chat calls = %d, want 0", stub.Calls()) - } -} - -func TestMergeTOCLevels_BadInput(t *testing.T) { - chunks := tocTestChunks("doc1", "a", "b") - levels := []any{ - []any{"1", "From a list"}, - map[string]any{"level": "x", "title": "Bad level"}, - map[string]any{"level": float64(2), "title": "Out of range"}, - map[string]any{"level": "1"}, - } - titles := []tocTitle{ - {title: "From a list", chunkID: "1"}, - {title: "Bad level", chunkID: "0"}, - {title: "Out of range", chunkID: "9"}, - {title: "No title in reply", chunkID: "0"}, - } - got := mergeTOCLevels(levels, titles, chunks) - want := []tocEntry{ - {Level: "1", Title: "From a list", IDs: []string{chunks[1]["id"].(string)}}, - {Level: "2", Title: "Out of range"}, - {Level: "1", Title: "No title in reply", IDs: []string{chunks[0]["id"].(string), chunks[1]["id"].(string)}}, - } - if !reflect.DeepEqual(got, want) { - t.Fatalf("mergeTOCLevels\n got: %+v\nwant: %+v", got, want) - } -} - -// The body chunks after the last heading belong to the last entry. -func TestMergeTOCLevels_LastEntryRunsToLastChunk(t *testing.T) { - chunks := tocTestChunks("doc1", "Intro", "body one", "body two") - levels := []any{map[string]any{"level": "1", "title": "Intro"}} - titles := []tocTitle{{title: "Intro", chunkID: "0"}} - got := mergeTOCLevels(levels, titles, chunks) - want := []tocEntry{ - {Level: "1", Title: "Intro", IDs: []string{chunks[0]["id"].(string), chunks[1]["id"].(string), chunks[2]["id"].(string)}}, - } - if !reflect.DeepEqual(got, want) { - t.Fatalf("mergeTOCLevels\n got: %+v\nwant: %+v", got, want) - } -} - -func TestNewExtractorComponent_TOCGroup(t *testing.T) { - c, err := NewExtractorComponent(map[string]any{"toc": map[string]any{"enabled": true}}) - if err != nil { - t.Fatalf("NewExtractorComponent: %v", err) - } - if !c.(*ExtractorComponent).Param.TOC.Enabled { - t.Fatal("toc.enabled=true did not enable TOC extraction") - } -} diff --git a/internal/ingestion/component/schema/extractor.go b/internal/ingestion/component/schema/extractor.go index 23c20db1a9..379ffb154f 100644 --- a/internal/ingestion/component/schema/extractor.go +++ b/internal/ingestion/component/schema/extractor.go @@ -116,14 +116,8 @@ type MetadataExtractConfig struct { BuiltInMetadata []common.MetadataFieldDef `json:"built_in_metadata,omitempty"` } -// TOCExtractConfig configures table-of-contents extraction. The TOC is built -// once per document from all chunks and appended as one extra chunk. -type TOCExtractConfig struct { - Enabled bool `json:"enabled"` -} - // ExtractorParam is the static configuration for the Extractor component. -// Fully modularized into base settings and 6 sub-extraction tasks. +// Fully modularized into base settings and 5 sub-extraction tasks. type ExtractorParam struct { // Base settings LLMID string `json:"llm_id,omitempty"` @@ -134,7 +128,6 @@ type ExtractorParam struct { Tags TagExtractConfig `json:"tags,omitempty"` Summary SummaryExtractConfig `json:"summary,omitempty"` Metadata MetadataExtractConfig `json:"metadata,omitempty"` - TOC TOCExtractConfig `json:"toc,omitempty"` } // Defaults returns the default ExtractorParam. diff --git a/internal/rag/prompts/assign_toc_levels.md b/internal/rag/prompts/assign_toc_levels.md deleted file mode 100644 index 56b94f2f45..0000000000 --- a/internal/rag/prompts/assign_toc_levels.md +++ /dev/null @@ -1,55 +0,0 @@ -You are given a JSON array of TOC(table of contents) items. Each item has at least {"title": string} and may include an existing title hierarchical level. - -Task -- For each item, assign a depth label using Arabic numerals only: top-level = 1, second-level = 2, third-level = 3, etc. -- Multiple items may share the same depth (e.g., many 1s, many 2s). -- Do not use dotted numbering (no 1.1/1.2). Use a single digit string per item indicating its depth only. -- Preserve the original item order exactly. Do not insert, delete, or reorder. -- Decide levels yourself to keep a coherent hierarchy. Keep peers at the same depth. - -Output -- Return a valid JSON array only (no extra text, no markdown code blocks). -- Each element MUST be a JSON object with exactly this structure: {"level": "1", "title": "some title"}. -- title must be the original title string exactly. -- DO NOT return arrays of arrays like [["1", "title"]] or other formats. -- The output must be parseable by json.loads() directly. - -Examples - -Example A (chapters with sections) -Input: -["Chapter 1 Methods", "Section 1 Definition", "Section 2 Process", "Chapter 2 Experiment"] - -Output: -[ - {"level":"1","title":"Chapter 1 Methods"}, - {"level":"2","title":"Section 1 Definition"}, - {"level":"2","title":"Section 2 Process"}, - {"level":"1","title":"Chapter 2 Experiment"} -] - -Example B (parts with chapters) -Input: -["Part I Theory", "Chapter 1 Basics", "Chapter 2 Methods", "Part II Applications", "Chapter 3 Case Studies"] - -Output: -[ - {"level":"1","title":"Part I Theory"}, - {"level":"2","title":"Chapter 1 Basics"}, - {"level":"2","title":"Chapter 2 Methods"}, - {"level":"1","title":"Part II Applications"}, - {"level":"2","title":"Chapter 3 Case Studies"} -] - -Example C (plain headings) -Input: -["Introduction", "Background and Motivation", "Related Work", "Methodology", "Evaluation"] - -Output: -[ - {"level":"1","title":"Introduction"}, - {"level":"2","title":"Background and Motivation"}, - {"level":"2","title":"Related Work"}, - {"level":"1","title":"Methodology"}, - {"level":"1","title":"Evaluation"} -] diff --git a/internal/rag/prompts/templates.go b/internal/rag/prompts/templates.go index 6c38676024..8323d6e078 100644 --- a/internal/rag/prompts/templates.go +++ b/internal/rag/prompts/templates.go @@ -25,7 +25,7 @@ import ( // Prompt templates for retrieval, chat, and ingestion are embedded here so // callers can load them independently of the working directory. -//go:embed citation_prompt.md ask_summary.md keyword_prompt.md cross_languages_sys_prompt.md cross_languages_user_prompt.md meta_filter.md related_question.md vision_llm_describe_prompt.md vision_llm_figure_describe_prompt.md toc_from_text_system.md toc_from_text_user.md assign_toc_levels.md +//go:embed citation_prompt.md ask_summary.md keyword_prompt.md cross_languages_sys_prompt.md cross_languages_user_prompt.md meta_filter.md related_question.md vision_llm_describe_prompt.md vision_llm_figure_describe_prompt.md var templatesFS embed.FS // EmbeddedPromptLoader serves the bundled .md templates by name. Load returns diff --git a/internal/rag/prompts/templates_test.go b/internal/rag/prompts/templates_test.go index cbb2de4956..195ec9edc5 100644 --- a/internal/rag/prompts/templates_test.go +++ b/internal/rag/prompts/templates_test.go @@ -35,9 +35,6 @@ func TestEmbeddedLoaderCanonicalTemplates(t *testing.T) { "related_question", "vision_llm_describe_prompt", "vision_llm_figure_describe_prompt", - "toc_from_text_system", - "toc_from_text_user", - "assign_toc_levels", } for _, name := range names { got, err := (EmbeddedPromptLoader{}).Load(name) diff --git a/internal/rag/prompts/toc_from_text_system.md b/internal/rag/prompts/toc_from_text_system.md deleted file mode 100644 index 7a12b39abc..0000000000 --- a/internal/rag/prompts/toc_from_text_system.md +++ /dev/null @@ -1,119 +0,0 @@ -You are a robust Table-of-Contents (TOC) extractor. - -GOAL -Given a dictionary of chunks {"": chunk_text}, extract TOC-like headings and return a strict JSON array of objects: -[ - {"title": "", "chunk_id": ""}, - ... -] - -FIELDS -- "title": the heading text (clean, no page numbers or leader dots). - - If any part of a chunk has no valid heading, output that part as {"title":"-1", ...}. -- "chunk_id": the chunk ID (string). - - One chunk can yield multiple JSON objects in order (unmatched text + one or more headings). - -RULES -1) Preserve input chunk order strictly. -2) If a chunk contains multiple headings, expand them in order: - - Pre-heading narrative → {"title":"-1","chunk_id":""} - - Then each heading → {"title":"...","chunk_id":""} -3) Do not merge outputs across chunks; each object refers to exactly one chunk ID. -4) "title" must be non-empty (or exactly "-1"). "chunk_id" must be a string (chunk ID). -5) When ambiguous, prefer "-1" unless the text strongly looks like a heading. - -HEADING DETECTION (cues, not hard rules) -- Appears near line start, short isolated phrase, often followed by content. -- May contain separators: — —— - : ： · • -- Numbering styles: - • 第[一二三四五六七八九十百]+(篇|章|节|条) - • [(（]?[一二三四五六七八九十]+[)）]? - • [(（]?[①②③④⑤⑥⑦⑧⑨⑩][)）]? - • ^\d+(\.\d+)*[)．.]?\s* - • ^[IVXLCDM]+[).] - • ^[A-Z][).] -- Canonical section cues (general only): - Common heading indicators include words such as: - "Overview", "Introduction", "Background", "Purpose", "Scope", "Definition", - "Method", "Procedure", "Result", "Discussion", "Summary", "Conclusion", - "Appendix", "Reference", "Annex", "Acknowledgment", "Disclaimer". - These are soft cues, not strict requirements. -- Length restriction: - • Chinese heading: ≤25 characters - • English heading: ≤80 characters -- Exclude long narrative sentences, continuous prose, or bullet-style lists → output as "-1". - -OUTPUT FORMAT -- Return ONLY a valid JSON array of {"title","chunk_id"} objects. -- No reasoning or commentary. - -EXAMPLES - -Example 1 — No heading -Input: -[{"0": "Copyright page · Publication info (ISBN 123-456). All rights reserved."}, ...] -Output: -[ - {"title":"-1","chunk_id":"0"}, - ... -] - -Example 2 — One heading -Input: -[{"1": "Chapter 1: General Provisions This chapter defines the overall rules…"}, ...] -Output: -[ - {"title":"Chapter 1: General Provisions","chunk_id":"1"}, - ... -] - -Example 3 — Narrative + heading -Input: -[{"2": "This paragraph introduces the background and goals. Section 2: Definitions Key terms are explained…"}, ...] -Output: -[ - {"title":"Section 2: Definitions","chunk_id":"2"}, - ... -] - -Example 4 — Multiple headings in one chunk -Input: -[{"3": "Declarations and Commitments (I) Party B commits… (II) Party C commits… Appendix A Data Specification"}, ...] -Output: -[ - {"title":"Declarations and Commitments","chunk_id":"3"}, - {"title":"(I) Party B commits","chunk_id":"3"}, - {"title":"(II) Party C commits","chunk_id":"3"}, - {"title":"Appendix A Data Specification","chunk_id":"3"}, - ... -] - -Example 5 — Numbering styles -Input: -[{"4": "1. Scope: Defines boundaries. 2) Definitions: Terms used. III) Methods Overview."}, ...] -Output: -[ - {"title":"1. Scope","chunk_id":"4"}, - {"title":"2) Definitions","chunk_id":"4"}, - {"title":"III) Methods Overview","chunk_id":"4"}, - ... -] - -Example 6 — Long list (NOT headings) -Input: -{"5": "Item list: apples, bananas, strawberries, blueberries, mangos, peaches"}, ...] -Output: -[ - {"title":"-1","chunk_id":"5"}, - ... -] - -Example 7 — Mixed Chinese/English -Input: -{"6": "（出版信息略）This standard follows industry practices. Chapter 1: Overview 摘要… 第2节：术语与缩略语"}, ...] -Output: -[ - {"title":"Chapter 1: Overview","chunk_id":"6"}, - {"title":"第2节：术语与缩略语","chunk_id":"6"}, - ... -] diff --git a/internal/rag/prompts/toc_from_text_user.md b/internal/rag/prompts/toc_from_text_user.md deleted file mode 100644 index 952d8eff25..0000000000 --- a/internal/rag/prompts/toc_from_text_user.md +++ /dev/null @@ -1,8 +0,0 @@ -OUTPUT FORMAT -- Return ONLY the JSON array. -- Use double quotes. -- No extra commentary. -- Keep language of "title" the same as the input. - -INPUT -{{text}} diff --git a/internal/service/document/document_test.go b/internal/service/document/document_test.go index 8be4c21abd..6b9c283f36 100644 --- a/internal/service/document/document_test.go +++ b/internal/service/document/document_test.go @@ -27,7 +27,6 @@ import ( "net/http/httptest" "path/filepath" "reflect" - "slices" "strings" "sync" "testing" @@ -334,8 +333,7 @@ func (e *sourceAvailabilityDocEngine) Search(_ context.Context, req *types.Searc {"id": "wiki-1", "compile_kwd": "wiki_page"}, {"id": "map-1", "compile_kwd": []interface{}{"wiki_map_active"}}, {"id": "nav-1", "compile_kwd": "dataset_nav"}, - {"id": "toc-1", "toc_kwd": []interface{}{"toc"}}, - }, Total: 8}, nil + }, Total: 7}, nil } func (e *sourceAvailabilityDocEngine) UpdateChunks(_ context.Context, condition map[string]interface{}, newValue map[string]interface{}, _, _ string) error { @@ -2779,35 +2777,6 @@ func TestUpdateDocumentChunkAvailabilityTogglesFinalProducts(t *testing.T) { } } -// The Extractor writes the TOC row with available_int=0, and only the TOC -// enhancer reads it. Re-enabling a document must not expose it to retrieval. -func TestUpdateDocumentChunkAvailabilityKeepsTOCHidden(t *testing.T) { - docEngine := &sourceAvailabilityDocEngine{} - svc := testDocumentService(t) - svc.docEngine = docEngine - - for _, status := range []int{0, 1} { - if err := svc.updateDocumentChunkAvailability(t.Context(), "tenant-1", "kb-1", "doc-1", status); err != nil { - t.Fatalf("updateDocumentChunkAvailability(%d) failed: %v", status, err) - } - } - if !slices.Contains(docEngine.search.SelectFields, "toc_kwd") { - t.Fatalf("SelectFields = %v, want toc_kwd", docEngine.search.SelectFields) - } - if len(docEngine.updateConditions) != 2 { - t.Fatalf("UpdateChunks calls = %d, want 2", len(docEngine.updateConditions)) - } - for i, want := range []int{0, 1} { - ids, _ := docEngine.updateConditions[i]["id"].([]string) - if slices.Contains(ids, "toc-1") { - t.Fatalf("call %d updated ids = %v, want the TOC row excluded", i, ids) - } - if got := docEngine.updateValues[i]["available_int"]; got != want { - t.Fatalf("call %d available_int = %#v, want %d", i, got, want) - } - } -} - func TestDocumentKnowledgeCompileTypesIncludesUnavailableProducts(t *testing.T) { docEngine := &sourceAvailabilityDocEngine{} svc := testDocumentService(t) diff --git a/internal/service/document/wiki_refresh.go b/internal/service/document/wiki_refresh.go index ac7ee508b0..91ec114c8a 100644 --- a/internal/service/document/wiki_refresh.go +++ b/internal/service/document/wiki_refresh.go @@ -35,8 +35,7 @@ const sourceChunkAvailabilityBatchSize = 1000 // disabled/enabled document is allowed to expose: its ordinary source chunks and // its per-document final compiled products (tree / structure / mindmap — the wiki // variant is hidden at write time, see markCompiledProductsHidden). Wiki staging -// rows, unknown-kwd rows and the TOC row are skipped; see -// loadAvailabilityToggleChunkIDs. +// rows and unknown-kwd rows are skipped; see loadAvailabilityToggleChunkIDs. func (s *DocumentService) updateDocumentChunkAvailability(ctx context.Context, tenantID, datasetID, documentID string, available int) error { if s.docEngine == nil { return fmt.Errorf("document engine not initialized") @@ -64,9 +63,8 @@ func (s *DocumentService) updateDocumentChunkAvailability(ctx context.Context, t // availability follows the document status: source chunks (no compile_kwd) and // final compiled products (a compile_kwd mapping to a non-wiki variant). Skipped: // the wiki variant (wiki_page / wiki_section), whose per-document rows are -// staging for the dataset-level merge, unknown-kwd rows such as the -// wiki_map_active state row or the legacy, KB-scoped wiki_page_graph blob, and -// the TOC row (toc_kwd=toc), which only the TOC enhancer reads and stays hidden. +// staging for the dataset-level merge, and unknown-kwd rows such as the +// wiki_map_active state row or the legacy, KB-scoped wiki_page_graph blob. func (s *DocumentService) loadAvailabilityToggleChunkIDs(ctx context.Context, indexName, datasetID, documentID string) ([]string, error) { ids := make([]string, 0) for offset := 0; ; offset += sourceChunkAvailabilityBatchSize { @@ -76,7 +74,7 @@ func (s *DocumentService) loadAvailabilityToggleChunkIDs(ctx context.Context, in KbIDs: []string{datasetID}, Offset: offset, Limit: sourceChunkAvailabilityBatchSize, - SelectFields: []string{"id", "compile_kwd", "type_kwd", "toc_kwd"}, + SelectFields: []string{"id", "compile_kwd", "type_kwd"}, Filter: map[string]any{"doc_id": []string{documentID}}, }) cancel() @@ -87,7 +85,7 @@ func (s *DocumentService) loadAvailabilityToggleChunkIDs(ctx context.Context, in break } for _, row := range result.Chunks { - if enginetypes.IsNavigationRow(row) || strings.TrimSpace(documentStoreString(row["toc_kwd"])) == "toc" { + if enginetypes.IsNavigationRow(row) { continue } if kwd := strings.TrimSpace(documentStoreString(row["compile_kwd"])); kwd != "" { diff --git a/internal/service/toc_enhancer_test.go b/internal/service/toc_enhancer_test.go index 9bf564feb1..faf65b77fe 100644 --- a/internal/service/toc_enhancer_test.go +++ b/internal/service/toc_enhancer_test.go @@ -251,20 +251,6 @@ func TestParseTOCEntries_MultiEntry(t *testing.T) { } } -// The Extractor writes the level as a numeric string. -func TestParseTOCEntries_StringLevel(t *testing.T) { - chunks := []map[string]interface{}{ - {"content_with_weight": `[{"level":"1","title":"Intro","ids":["c1","c2"]},{"level":"2","title":"Details"}]`}, - } - got := parseTOCEntries(chunks) - if len(got) != 2 { - t.Fatalf("expected 2 entries, got %d", len(got)) - } - if got[0].Level != 1 || len(got[0].IDs) != 2 || got[1].Level != 2 || got[1].IDs != nil { - t.Errorf("unexpected entries: %+v", got) - } -} - func TestParseTOCEntries_InvalidJSON(t *testing.T) { chunks := []map[string]interface{}{ {"content_with_weight": `not json`},