Skip to content

Commit 4f1f49d

Browse files
author
Halleluyah Oladipo
committed
feat(ingest,config): wire TOC builder into the pipeline + config
- Adds IngestConfig.TOC (Enabled / Model / Concurrency / TOCCheckPages) with defaults (Enabled=true, Concurrency=4, TOCCheckPages=20), env overrides (VLE_INGEST_TOC_{ENABLED,MODEL,CONCURRENCY,TOC_CHECK_PAGES}), validation, and example-config documentation. - Pipeline.Run calls the new builder after summarize+HyDE for PDF inputs, persists the result via UpdateDocumentTOCTree, and logs the LLM-call accounting. Failures are non-fatal — they leave documents.toc_tree NULL and the document remains fully retrievable via the existing sections tree. - assemblePagesFromSections groups parsed sections by PageStart to reconstruct per-page text the builder can reason over. PageStart==0 sections are skipped so the builder never sees ambiguous page numbers. - cmd/server wires the new config block into the pipeline literal.
1 parent 766c981 commit 4f1f49d

6 files changed

Lines changed: 394 additions & 0 deletions

File tree

cmd/server/main.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,10 @@ func run() error {
169169
SummaryAxesMaxTopics: cfg.Engine.Ingest.SummaryAxes.MaxTopics,
170170
SummaryAxesMaxEntities: cfg.Engine.Ingest.SummaryAxes.MaxEntities,
171171
SummaryAxesMaxNumbers: cfg.Engine.Ingest.SummaryAxes.MaxNumbers,
172+
TOCEnabled: cfg.Engine.Ingest.TOC.Enabled,
173+
TOCModel: cfg.Engine.Ingest.TOC.Model,
174+
TOCConcurrency: cfg.Engine.Ingest.TOC.Concurrency,
175+
TOCCheckPages: cfg.Engine.Ingest.TOC.TOCCheckPages,
172176
GlobalLLMConcurrency: cfg.Engine.Ingest.GlobalLLMConcurrency,
173177
})
174178
if cfg.Engine.Ingest.Tables.Enabled {

config.example.yaml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,32 @@ ingest:
297297
max_entities: 8
298298
max_numbers: 6
299299

300+
# LLM-built table-of-contents tree (PageIndex-style). Runs after
301+
# summarize+HyDE on PDF inputs and persists a hierarchical TOC on
302+
# documents.toc_tree (JSONB). The tree is small (a few KB even
303+
# for 300-page filings) and is intended as a higher-level map
304+
# retrieval strategies can reason over before drilling into the
305+
# parser-derived sections tree.
306+
#
307+
# ENABLED BY DEFAULT for PDFs. Non-PDF documents skip the stage
308+
# unconditionally. Builder failures are non-fatal — the document
309+
# remains fully retrievable via the existing sections tree.
310+
toc:
311+
enabled: true
312+
# Override the LLM model used by the builder; empty inherits
313+
# the summary model. Point this at a reasoning-capable model —
314+
# the no-TOC generator has to find hierarchy in raw body text,
315+
# which a small/fast model often botches.
316+
model: ""
317+
# Cap on parallel LLM calls during the verification phase
318+
# (one call per leaf node).
319+
concurrency: 4
320+
# The detector scans the first N pages for a table of
321+
# contents. PageIndex defaults this to 20 — financial filings
322+
# put their TOC inside the first dozen pages and a document
323+
# without one by page 20 almost never has one further in.
324+
toc_check_pages: 20
325+
300326
log:
301327
level: "info" # debug | info | warn | error
302328
format: "json" # json | console

pkg/config/config.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,13 @@ type IngestConfig struct {
4848
// populate it).
4949
SummaryAxes SummaryAxesBlock `yaml:"summary_axes"`
5050

51+
// TOC configures the PageIndex-style LLM-built table-of-contents
52+
// tree stage. Enabled by default for PDF inputs; the resulting
53+
// tree is persisted on documents.toc_tree (JSONB). Failures are
54+
// non-fatal — they leave the column NULL and the document fully
55+
// retrievable via the existing sections tree.
56+
TOC TOCBlock `yaml:"toc"`
57+
5158
// GlobalLLMConcurrency caps the total number of LLM calls in flight
5259
// across the summarize and HyDE stages combined, which now run
5360
// concurrently. Each stage still respects its own per-stage cap
@@ -60,6 +67,39 @@ type IngestConfig struct {
6067
GlobalLLMConcurrency int `yaml:"global_llm_concurrency"`
6168
}
6269

70+
// TOCBlock configures the LLM-driven table-of-contents tree
71+
// builder. The builder reads page-by-page text from a freshly-
72+
// ingested PDF and emits a hierarchical TOC (PageIndex-style),
73+
// persisted on documents.toc_tree (JSONB).
74+
//
75+
// Enabled by default for PDF inputs; non-PDF documents skip the
76+
// stage unconditionally. Builder failures never break ingest —
77+
// the document remains fully retrievable via the existing
78+
// sections tree.
79+
type TOCBlock struct {
80+
// Enabled toggles the stage. Default: true. Flip to false to
81+
// skip the extra LLM round-trip when ingest budget matters
82+
// more than having a TOC tree for retrieval to reason over.
83+
Enabled bool `yaml:"enabled"`
84+
85+
// Model overrides the LLM model used by the builder. Empty
86+
// inherits the engine's configured default. Point this at a
87+
// reasoning-capable model — the no-TOC generator has to find
88+
// hierarchy in raw body text, which a small/fast model often
89+
// botches.
90+
Model string `yaml:"model"`
91+
92+
// Concurrency caps parallel LLM calls during the verification
93+
// phase (one call per leaf node). Default: 4.
94+
Concurrency int `yaml:"concurrency"`
95+
96+
// TOCCheckPages bounds the leading prefix the detector scans
97+
// for a table of contents. Default: 20 — financial filings
98+
// put their TOC inside the first dozen pages and a document
99+
// without one by page 20 almost never has one further in.
100+
TOCCheckPages int `yaml:"toc_check_pages"`
101+
}
102+
63103
// SummaryAxesBlock configures the Phase 2.5 structured summarizer.
64104
//
65105
// When enabled, the summarize stage runs in JSON mode and produces
@@ -584,6 +624,11 @@ func Default() Config {
584624
MaxEntities: 8,
585625
MaxNumbers: 6,
586626
},
627+
TOC: TOCBlock{
628+
Enabled: true,
629+
Concurrency: 4,
630+
TOCCheckPages: 20,
631+
},
587632
},
588633
Log: LogConfig{Level: "info", Format: "json"},
589634
}
@@ -767,6 +812,30 @@ func applyEnvOverrides(c *Config) {
767812
c.Ingest.SummaryAxes.MaxNumbers = n
768813
}
769814
}
815+
// LLM-built TOC tree (PageIndex-style). Same truthy-string set
816+
// as the other ingest toggles; numeric overrides require a
817+
// positive int so a typo doesn't silently flip the default.
818+
if v := os.Getenv("VLE_INGEST_TOC_ENABLED"); v != "" {
819+
switch strings.ToLower(strings.TrimSpace(v)) {
820+
case "1", "true", "yes", "on":
821+
c.Ingest.TOC.Enabled = true
822+
case "0", "false", "no", "off":
823+
c.Ingest.TOC.Enabled = false
824+
}
825+
}
826+
if v := os.Getenv("VLE_INGEST_TOC_MODEL"); v != "" {
827+
c.Ingest.TOC.Model = v
828+
}
829+
if v := os.Getenv("VLE_INGEST_TOC_CONCURRENCY"); v != "" {
830+
if n, err := strconv.Atoi(v); err == nil && n > 0 {
831+
c.Ingest.TOC.Concurrency = n
832+
}
833+
}
834+
if v := os.Getenv("VLE_INGEST_TOC_TOC_CHECK_PAGES"); v != "" {
835+
if n, err := strconv.Atoi(v); err == nil && n > 0 {
836+
c.Ingest.TOC.TOCCheckPages = n
837+
}
838+
}
770839
if v := os.Getenv("VLE_RETRIEVAL_ANSWER_SPAN_ENABLED"); v != "" {
771840
switch strings.ToLower(strings.TrimSpace(v)) {
772841
case "1", "true", "yes", "on":
@@ -978,6 +1047,13 @@ func (c Config) Validate() error {
9781047
return fmt.Errorf("ingest.summary_axes.max_numbers must be >= 0, got %d", c.Ingest.SummaryAxes.MaxNumbers)
9791048
}
9801049

1050+
if c.Ingest.TOC.Concurrency < 0 {
1051+
return fmt.Errorf("ingest.toc.concurrency must be >= 0, got %d", c.Ingest.TOC.Concurrency)
1052+
}
1053+
if c.Ingest.TOC.TOCCheckPages < 0 {
1054+
return fmt.Errorf("ingest.toc.toc_check_pages must be >= 0, got %d", c.Ingest.TOC.TOCCheckPages)
1055+
}
1056+
9811057
if c.Retrieval.Planning.CacheSize < 0 {
9821058
return fmt.Errorf("retrieval.planning.cache_size must be >= 0, got %d", c.Retrieval.Planning.CacheSize)
9831059
}

pkg/config/config_test.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,55 @@ func TestDefaultValues(t *testing.T) {
7373
if cfg.Log.Level != "info" {
7474
t.Errorf("log.level = %q, want info", cfg.Log.Level)
7575
}
76+
if !cfg.Ingest.TOC.Enabled {
77+
t.Error("ingest.toc.enabled should default to true (opt-out)")
78+
}
79+
if cfg.Ingest.TOC.Concurrency != 4 {
80+
t.Errorf("ingest.toc.concurrency = %d, want 4", cfg.Ingest.TOC.Concurrency)
81+
}
82+
if cfg.Ingest.TOC.TOCCheckPages != 20 {
83+
t.Errorf("ingest.toc.toc_check_pages = %d, want 20", cfg.Ingest.TOC.TOCCheckPages)
84+
}
85+
}
86+
87+
func TestTOCEnvOverride(t *testing.T) {
88+
// Mutates env — restore on exit. Not parallel.
89+
keys := []string{
90+
"VLE_INGEST_TOC_ENABLED",
91+
"VLE_INGEST_TOC_MODEL",
92+
"VLE_INGEST_TOC_CONCURRENCY",
93+
"VLE_INGEST_TOC_TOC_CHECK_PAGES",
94+
}
95+
prev := make(map[string]string, len(keys))
96+
for _, k := range keys {
97+
prev[k] = os.Getenv(k)
98+
}
99+
defer func() {
100+
for k, v := range prev {
101+
os.Setenv(k, v)
102+
}
103+
}()
104+
105+
os.Setenv("VLE_INGEST_TOC_ENABLED", "false")
106+
os.Setenv("VLE_INGEST_TOC_MODEL", "gemini-2.5-pro")
107+
os.Setenv("VLE_INGEST_TOC_CONCURRENCY", "12")
108+
os.Setenv("VLE_INGEST_TOC_TOC_CHECK_PAGES", "30")
109+
110+
cfg := Default()
111+
applyEnvOverrides(&cfg)
112+
113+
if cfg.Ingest.TOC.Enabled {
114+
t.Error("VLE_INGEST_TOC_ENABLED=false should disable the stage")
115+
}
116+
if cfg.Ingest.TOC.Model != "gemini-2.5-pro" {
117+
t.Errorf("VLE_INGEST_TOC_MODEL not applied, got %q", cfg.Ingest.TOC.Model)
118+
}
119+
if cfg.Ingest.TOC.Concurrency != 12 {
120+
t.Errorf("VLE_INGEST_TOC_CONCURRENCY=12 not applied, got %d", cfg.Ingest.TOC.Concurrency)
121+
}
122+
if cfg.Ingest.TOC.TOCCheckPages != 30 {
123+
t.Errorf("VLE_INGEST_TOC_TOC_CHECK_PAGES=30 not applied, got %d", cfg.Ingest.TOC.TOCCheckPages)
124+
}
76125
}
77126

78127
func TestAbstainEnvOverride(t *testing.T) {

pkg/ingest/ingest.go

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,29 @@ type Pipeline struct {
134134
// per-stage semaphore). Default applied by NewPipeline: 12.
135135
GlobalLLMConcurrency int
136136

137+
// TOCEnabled toggles the LLM-built table-of-contents stage. The
138+
// stage runs after summarize+HyDE on PDF inputs and persists the
139+
// resulting tree on documents.toc_tree (JSONB). Failures are
140+
// non-fatal — they leave the column NULL.
141+
//
142+
// Defaulted to true by config wiring; left as the Go zero value
143+
// (false) when Pipeline is constructed directly, so unit tests
144+
// with no LLM can opt out by simply not setting it.
145+
TOCEnabled bool
146+
147+
// TOCModel overrides the LLM model used by the TOC builder.
148+
// Empty inherits SummaryModel (which itself falls back to the
149+
// client default).
150+
TOCModel string
151+
152+
// TOCConcurrency caps parallel LLM calls during the TOC
153+
// verification phase. Default: 4.
154+
TOCConcurrency int
155+
156+
// TOCCheckPages bounds the leading prefix the detector scans
157+
// for a table of contents. Default: 20.
158+
TOCCheckPages int
159+
137160
// globalLLMSem is the lazily-initialized shared semaphore enforcing
138161
// GlobalLLMConcurrency. nil means "no global cap" — callers fall back
139162
// to per-stage limits only.
@@ -168,6 +191,12 @@ func NewPipeline(p Pipeline) *Pipeline {
168191
if p.HyDEConcurrency <= 0 {
169192
p.HyDEConcurrency = 4
170193
}
194+
if p.TOCConcurrency <= 0 {
195+
p.TOCConcurrency = 4
196+
}
197+
if p.TOCCheckPages <= 0 {
198+
p.TOCCheckPages = 20
199+
}
171200
// Default the global cap to a value that comfortably exceeds the
172201
// sum of the two default per-stage caps (4 + 4 = 8) while leaving
173202
// some headroom — but stays well below typical provider per-tenant
@@ -265,13 +294,131 @@ func (p *Pipeline) Run(ctx context.Context, pl Payload) error {
265294
}
266295
log.Info("ingest: summarize+hyde complete", "elapsed", time.Since(stageStart))
267296

297+
// LLM-built TOC tree (PageIndex-style). PDF-only because it
298+
// relies on the parser's PageStart/PageEnd attribution to
299+
// reconstruct per-page text. Non-fatal: a builder failure
300+
// leaves documents.toc_tree NULL and the document remains
301+
// fully retrievable via the sections tree above.
302+
if p.TOCEnabled && pl.ContentType == "application/pdf" {
303+
if err := p.runTOCBuilder(ctx, pl.DocumentID, parsed, log); err != nil {
304+
log.Warn("ingest: toc-builder failed; falling back to NULL toc_tree", "err", err)
305+
}
306+
}
307+
268308
if err := p.DB.SetDocumentStatus(ctx, pl.DocumentID, db.StatusReady, ""); err != nil {
269309
return err
270310
}
271311
log.Info("ingest: ready")
272312
return nil
273313
}
274314

315+
// runTOCBuilder assembles per-page text from the parsed PDF, runs
316+
// the LLM-driven TOC builder over it, and persists the result.
317+
// Returns an error only on a transport-level builder failure or a
318+
// JSON-marshal blip; the caller logs and continues either way.
319+
//
320+
// A nil-result (no usable nodes) is treated as success and writes
321+
// SQL NULL to documents.toc_tree (which is the column's default,
322+
// so this is also the no-op).
323+
func (p *Pipeline) runTOCBuilder(ctx context.Context, docID tree.DocumentID, parsed *parser.ParsedDoc, log *slog.Logger) error {
324+
pages := assemblePagesFromSections(parsed.Sections)
325+
if len(pages) == 0 {
326+
log.Info("ingest: toc-builder skipped; no per-page text available")
327+
return nil
328+
}
329+
model := p.TOCModel
330+
if model == "" {
331+
model = p.SummaryModel
332+
}
333+
builder := &TOCBuilder{
334+
LLM: p.LLM,
335+
Model: model,
336+
Concurrency: p.TOCConcurrency,
337+
TOCCheckPages: p.TOCCheckPages,
338+
}
339+
nodes, usage, err := builder.Build(ctx, pages)
340+
if err != nil {
341+
return err
342+
}
343+
log.Info("ingest: toc-builder done",
344+
"top_level_nodes", len(nodes),
345+
"llm_calls", usage.LLMCalls,
346+
"input_tokens", usage.InputTokens,
347+
"output_tokens", usage.OutputTokens,
348+
)
349+
if len(nodes) == 0 {
350+
return nil
351+
}
352+
treeJSON, err := json.Marshal(nodes)
353+
if err != nil {
354+
return fmt.Errorf("marshal toc tree: %w", err)
355+
}
356+
if err := p.DB.UpdateDocumentTOCTree(ctx, docID, treeJSON); err != nil {
357+
return fmt.Errorf("persist toc tree: %w", err)
358+
}
359+
return nil
360+
}
361+
362+
// assemblePagesFromSections groups the parsed sections' text by
363+
// their PageStart, producing PageText entries the TOC builder can
364+
// reason over. Sections that span multiple pages collapse onto
365+
// their starting page — perfect page reconstruction would need
366+
// raw glyph-level coordinates the parser doesn't currently
367+
// surface, but the title-on-claimed-page heuristic still works
368+
// because section starts (where the LLM looks for titles) live
369+
// on PageStart.
370+
//
371+
// Sections with PageStart == 0 are skipped (the parser couldn't
372+
// place them) so the builder never sees ambiguous page numbers.
373+
func assemblePagesFromSections(secs []parser.Section) []PageText {
374+
pageText := map[int]*strings.Builder{}
375+
pages := []int{}
376+
var walk func([]parser.Section)
377+
walk = func(ss []parser.Section) {
378+
for _, s := range ss {
379+
if s.PageStart > 0 {
380+
b, ok := pageText[s.PageStart]
381+
if !ok {
382+
b = &strings.Builder{}
383+
pageText[s.PageStart] = b
384+
pages = append(pages, s.PageStart)
385+
}
386+
if title := strings.TrimSpace(s.Title); title != "" {
387+
if b.Len() > 0 {
388+
b.WriteByte('\n')
389+
}
390+
b.WriteString(title)
391+
b.WriteByte('\n')
392+
}
393+
if body := strings.TrimSpace(s.Content); body != "" {
394+
b.WriteString(body)
395+
b.WriteByte('\n')
396+
}
397+
}
398+
walk(s.Children)
399+
}
400+
}
401+
walk(secs)
402+
// Sort the page-number index in place.
403+
sortIntsAscending(pages)
404+
out := make([]PageText, 0, len(pages))
405+
for _, p := range pages {
406+
out = append(out, PageText{PageNumber: p, Text: pageText[p].String()})
407+
}
408+
return out
409+
}
410+
411+
// sortIntsAscending sorts a slice of ints in place. Insertion sort
412+
// is fine here — pages slice is typically a few hundred items
413+
// at most.
414+
func sortIntsAscending(xs []int) {
415+
for i := 1; i < len(xs); i++ {
416+
for j := i; j > 0 && xs[j-1] > xs[j]; j-- {
417+
xs[j-1], xs[j] = xs[j], xs[j-1]
418+
}
419+
}
420+
}
421+
275422
// runParallelStages runs summarize and HyDE concurrently, returning each
276423
// stage's error independently so callers can log them separately. A nil
277424
// hydeFn skips the HyDE stage (returns nil for hydeErr).

0 commit comments

Comments
 (0)