diff --git a/cmd/engine/main.go b/cmd/engine/main.go index dc126d6..3b26287 100644 --- a/cmd/engine/main.go +++ b/cmd/engine/main.go @@ -249,6 +249,8 @@ func run() error { }) if cfg.Ingest.Mode == ingest.ModeMinimal { logger.Info("ingest: MINIMAL mode — parse→persist→ready; skipping summarize/HyDE/multi-axis/TOC + table extraction") + } else if cfg.Ingest.Mode == ingest.ModeTOC { + logger.Info("ingest: TOC mode — parse→table of contents→persist→ready; skipping summarize/HyDE/multi-axis + table extraction") } else if cfg.Ingest.Tables.Enabled { logger.Info("ingest: pdf table extraction enabled", "vertical_strategy", cfg.Ingest.Tables.VerticalStrategy, diff --git a/cmd/server/main.go b/cmd/server/main.go index 48fe422..905c93c 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -241,6 +241,8 @@ func run() error { }) if cfg.Engine.Ingest.Mode == ingest.ModeMinimal { logger.Info("ingest: MINIMAL mode — parse→persist→ready; skipping summarize/HyDE/multi-axis/TOC + table extraction") + } else if cfg.Engine.Ingest.Mode == ingest.ModeTOC { + logger.Info("ingest: TOC mode — parse→table of contents→persist→ready; skipping summarize/HyDE/multi-axis + table extraction") } else if cfg.Engine.Ingest.Tables.Enabled { logger.Info("ingest: pdf table extraction enabled", "vertical_strategy", cfg.Engine.Ingest.Tables.VerticalStrategy, @@ -536,6 +538,11 @@ func buildStrategySet(c enginecfg.RetrievalConfig, client llmgate.Client, judge } if judge != nil { set["judgewalk"] = buildJudgeWalkStrategy(judge, store) + } else { + // The same fallback the default builder applies: a request that + // names judgewalk on a server with no Judge gets treewalk, not + // "unknown strategy". + set["judgewalk"] = set["treewalk"] } return set } diff --git a/config.example.yaml b/config.example.yaml index be9e3f3..5336fd6 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -362,6 +362,13 @@ ingest: # summary-dependent strategies (chunked-tree, agentic) # degrade to titles + raw content with no summaries. # + # toc parse -> build tree -> table of contents -> persist -> + # ready. The page-based pipeline and nothing else: on a + # Judge the TOC stage is three requests and seconds, and + # page-based retrieval (treewalk, judgewalk) needs none of + # the per-section enrichment. Table extraction skipped as + # in minimal. The mode the FinanceBench evaluations use. + # # Override per-process with VLE_INGEST_MODE; on the deployed # vectorless-server use VLS_INGEST_MODE=minimal (no secret edit needed). mode: "full" diff --git a/internal/api/abstention_test.go b/internal/api/abstention_test.go index 2018cdf..48ed4d6 100644 --- a/internal/api/abstention_test.go +++ b/internal/api/abstention_test.go @@ -198,7 +198,7 @@ func TestRespondAbstained(t *testing.T) { confidences := map[tree.SectionID]float64{"sec_a": 0.12, "sec_b": 0.30} rec := httptest.NewRecorder() - d.respondAbstained(rec, tree.DocumentID("doc_x"), "what is x?", confidences, nil) + d.respondAbstained(rec, tree.DocumentID("doc_x"), "what is x?", "", confidences, nil, retrieval.Usage{}) if rec.Code != http.StatusOK { t.Fatalf("status = %d, want 200", rec.Code) @@ -279,7 +279,7 @@ func TestRespondAbstainedTraceTokenAbsent(t *testing.T) { Abstain: config.AbstainBlock{Enabled: true, Below: 0.4}, } rec := httptest.NewRecorder() - d.respondAbstained(rec, tree.DocumentID("doc_x"), "q", map[tree.SectionID]float64{"a": 0.1}, nil) + d.respondAbstained(rec, tree.DocumentID("doc_x"), "q", "", map[tree.SectionID]float64{"a": 0.1}, nil, retrieval.Usage{}) var body map[string]any _ = json.Unmarshal(rec.Body.Bytes(), &body) diff --git a/internal/api/server.go b/internal/api/server.go index c619488..2660d75 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -600,13 +600,26 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { started := time.Now() - plan, _ := d.runPlanner(r.Context(), body.Query, body.EnablePlanning) - ids, confidences, err := d.runSelection(r.Context(), t, plan, body.Query, budget) + // Usage accumulates the way /v1/answer's does: planner, selection, + // re-rank. Span extraction is per section and not counted on either + // endpoint today. + totalUsage := retrieval.Usage{} + plan, planUsage := d.runPlanner(r.Context(), body.Query, body.EnablePlanning) + totalUsage.Add(planUsage) + ids, confidences, selUsage, err := d.runSelectionWithUsage(r.Context(), t, plan, body.Query, budget) if err != nil { d.Logger.Error("query: strategy failed", "err", err, "document_id", body.DocumentID) writeErr(w, http.StatusInternalServerError, "retrieval failed: "+err.Error()) return } + totalUsage.Add(selUsage) + // The model the caller named, or — when it named none, as a + // Judge-navigated query need not — the strategy, so the field is + // never empty and a client can always tell what answered. + modelUsed := body.Model + if modelUsed == "" { + modelUsed = d.Strategy.Name() + } // Phase 2.4 abstention: if every confident pick is below the // configured threshold, refuse to ground an answer in evidence @@ -615,7 +628,7 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { // responses (no confidences) always fall through to the normal // path so older models keep working. if d.abstentionEnabled(body.EnableAbstain) && shouldAbstain(confidences, d.Abstain.Below) { - d.respondAbstained(w, body.DocumentID, body.Query, confidences, plan) + d.respondAbstained(w, body.DocumentID, body.Query, modelUsed, confidences, plan, totalUsage) return } @@ -647,7 +660,9 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { // never drop sections — at worst the strategy's order is // preserved (see retrieval.ReRanker.ReRank). if d.reRankEnabled(body.EnableReRank) { - enriched, _ = d.runReRank(r.Context(), enriched, body.Query, body.Model) + var reRankUsage retrieval.Usage + enriched, reRankUsage = d.runReRank(r.Context(), enriched, body.Query, body.Model) + totalUsage.Add(reRankUsage) } // Optional: per-section answer-span extraction. Opt-in via config — @@ -678,10 +693,13 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { "document_id": body.DocumentID, "query": body.Query, "strategy": d.Strategy.Name(), - "model": body.Model, + "model": modelUsed, "sections": sections, "elapsed_ms": time.Since(started).Milliseconds(), "trace_token": traceToken, + // What retrieval cost. /v1/answer always reported this; /v1/query + // dropped it, so a caller benchmarking retrieval alone saw $0. + "usage": usageJSON(totalUsage), } if plan != nil { resp["plan"] = plan @@ -1650,11 +1668,13 @@ const abstentionAnswerText = "I cannot answer this question from the supplied do // the response in the replay log because there's no meaningful // retrieval result to reproduce. Callers replaying an abstention // will simply re-run /v1/query. -func (d Deps) respondAbstained(w http.ResponseWriter, docID tree.DocumentID, query string, confidences map[tree.SectionID]float64, plan *retrieval.Plan) { +func (d Deps) respondAbstained(w http.ResponseWriter, docID tree.DocumentID, query string, model string, confidences map[tree.SectionID]float64, plan *retrieval.Plan, usage retrieval.Usage) { resp := map[string]any{ "document_id": docID, "query": query, "strategy": d.Strategy.Name(), + "model": model, + "usage": usageJSON(usage), "sections": []any{}, "abstained": true, "abstention_reason": abstentionReason, @@ -1667,6 +1687,18 @@ func (d Deps) respondAbstained(w http.ResponseWriter, docID tree.DocumentID, que writeJSON(w, http.StatusOK, resp) } +// usageJSON is the wire shape of retrieval.Usage on /v1/query, the same +// keys /v1/answer uses for total_usage. +func usageJSON(u retrieval.Usage) map[string]any { + return map[string]any{ + "input_tokens": u.InputTokens, + "output_tokens": u.OutputTokens, + "total_tokens": u.TotalTokens, + "cost_usd": u.CostUSD, + "llm_calls": u.LLMCalls, + } +} + // respondAbstainedAnswer writes the abstention shape for /v1/answer. // The answer text is the canonical refusal; citations is empty; // usage carries the LLM tokens spent up to the abstention point diff --git a/pkg/config/config.go b/pkg/config/config.go index b334469..d9b5cdd 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1319,7 +1319,7 @@ func (c Config) Validate() error { } switch c.Retrieval.Strategy { - case "auto", "single-pass", "chunked-tree", "agentic", "treewalk": + case "auto", "single-pass", "chunked-tree", "agentic", "treewalk", "judgewalk": default: return fmt.Errorf("unknown retrieval.strategy: %q", c.Retrieval.Strategy) } @@ -1339,9 +1339,9 @@ func (c Config) Validate() error { } switch c.Ingest.Mode { - case "", "full", "minimal": + case "", "full", "minimal", "toc": default: - return fmt.Errorf("ingest.mode must be one of full|minimal, got %q", c.Ingest.Mode) + return fmt.Errorf("ingest.mode must be one of full|minimal|toc, got %q", c.Ingest.Mode) } if c.Ingest.HyDE.NumQuestions < 0 { diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 9a3bb34..7761240 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -652,7 +652,7 @@ func TestValidateLLMDrivers(t *testing.T) { func TestValidateRetrievalStrategy(t *testing.T) { t.Parallel() - for _, s := range []string{"auto", "single-pass", "chunked-tree", "agentic", "treewalk"} { + for _, s := range []string{"auto", "single-pass", "chunked-tree", "agentic", "treewalk", "judgewalk"} { cfg := Default() cfg.Database.URL = "postgres://localhost/test" cfg.Retrieval.Strategy = s diff --git a/pkg/ingest/ingest.go b/pkg/ingest/ingest.go index 4d2b6e0..8462fa7 100644 --- a/pkg/ingest/ingest.go +++ b/pkg/ingest/ingest.go @@ -50,6 +50,14 @@ import ( // and table extraction. Any other value runs the full pipeline. const ModeMinimal = "minimal" +// ModeTOC is parse → build tree → table of contents → persist → ready: +// the page-based pipeline and nothing else. On a Judge the TOC stage is +// three requests and seconds; the per-section generative enrichment +// (summarize, HyDE, multi-axis) that full mode adds is minutes and is +// not used by page-based retrieval. Table extraction is skipped as in +// minimal mode: the page text still carries the tables' text. +const ModeTOC = "toc" + // docPersister is the narrow slice of *db.Pool the parse → persist → // ready path depends on. Declaring it here (rather than threading the // concrete *db.Pool) lets the minimal-mode runner be exercised with a @@ -377,7 +385,7 @@ func (p *Pipeline) Handler() queue.Handler { // tree → persist → ready, with no LLM enrichment and no table // extraction. Otherwise it runs the full enrichment pipeline below. func (p *Pipeline) Run(ctx context.Context, pl Payload) error { - if p.Mode == ModeMinimal { + if p.Mode == ModeMinimal || p.Mode == ModeTOC { return p.runMinimal(ctx, p.DB, pl) } @@ -730,13 +738,20 @@ func (p *Pipeline) runMinimal(ctx context.Context, store docPersister, pl Payloa return err } - // Skip summarize / HyDE / multi-axis / TOC entirely — flip straight - // to ready. The document is now queryable via the page-based - // strategy (synthesised TOC + raw page reads). + // Minimal mode skips summarize / HyDE / multi-axis / TOC entirely + // and flips straight to ready; the document is queryable via the + // page-based strategy on a TOC synthesised from the section tree. + // TOC mode builds the real table of contents first — same builder + // and persistence as full mode, non-fatal for the same reason. + if p.Mode == ModeTOC && pl.ContentType == "application/pdf" { + if err := p.runTOCBuilder(ctx, pl.DocumentID, parsed, log); err != nil { + log.Warn("ingest: toc-builder failed; falling back to NULL toc_tree", "err", err) + } + } if err := store.SetDocumentStatus(ctx, pl.DocumentID, db.StatusReady, ""); err != nil { return err } - log.Info("ingest: ready (minimal mode)") + log.Info("ingest: ready (" + p.Mode + " mode)") return nil } diff --git a/pkg/ingest/minimal_mode_test.go b/pkg/ingest/minimal_mode_test.go index e753b1b..c5940fb 100644 --- a/pkg/ingest/minimal_mode_test.go +++ b/pkg/ingest/minimal_mode_test.go @@ -315,3 +315,38 @@ func reconstructTree(_ tree.DocumentID, title string, rows []db.Section) *tree.S return &tree.Section{Title: title, Children: topLevel} } } + +// TOC mode on a non-PDF is minimal mode: nothing for the TOC builder to +// do, no model call, ready. +func TestTOCModeOnMarkdownMakesNoLLMCall(t *testing.T) { + ctx := context.Background() + st, err := storage.NewLocal(t.TempDir()) + if err != nil { + t.Fatal(err) + } + body := []byte("# Title\n\n## Section A\n\nAlpha.\n\n## Section B\n\nBeta.\n") + docID := NewDocumentID() + srcKey := SourceKey(docID, "doc.md") + if err := st.Put(ctx, srcKey, bytes.NewReader(body), storage.Metadata{ContentType: "text/markdown", Size: int64(len(body))}); err != nil { + t.Fatal(err) + } + p := NewPipeline(Pipeline{ + Storage: st, + LLM: &failIfCalledLLM{t: t}, + Parsers: DefaultRegistry(), + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + Mode: ModeTOC, + TOCEnabled: true, + }) + fake := &fakeDocStore{} + if err := p.runMinimal(ctx, fake, Payload{DocumentID: docID, ContentType: "text/markdown", Filename: "doc.md", SourceRef: srcKey}); err != nil { + t.Fatalf("run: %v", err) + } + status, errMsg, sections := fake.snapshot() + if status != db.StatusReady { + t.Fatalf("status %q (%q) want ready", status, errMsg) + } + if len(sections) == 0 { + t.Fatal("no sections persisted") + } +}