mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-10 21:32:43 +08:00
Revert "feat(ingestion): generate TOC chunks in the Go Extractor" (#21385)
This commit is contained in:
Wang Qi authored and 
GitHub committed
1 parent
7d188e1abe
commit
848d21fc30
12 files changed
+10
-972
No files matched your search
@@ -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.
|
||||
|
||||
@@ -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
|
||||
// {"<index>": 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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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"}
|
||||
]
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -1,119 +0,0 @@
|
||||
You are a robust Table-of-Contents (TOC) extractor.
|
||||
|
||||
GOAL
|
||||
Given a dictionary of chunks {"<chunk_ID>": 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":"<chunk_ID>"}
|
||||
- Then each heading → {"title":"...","chunk_id":"<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"},
|
||||
...
|
||||
]
|
||||
@@ -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}}
|
||||
@@ -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)
|
||||
|
||||
@@ -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 != "" {
|
||||
|
||||
@@ -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`},
|
||||
|
||||
Reference in new issue
Block a user