From 92a5c288e1905914f6a221fa0066d83494738cef Mon Sep 17 00:00:00 2001 From: hallelx2 Date: Sat, 17 Oct 2026 10:53:27 +0100 Subject: [PATCH] feat(retrieval): return a short document whole, and the top evidence page's neighbour Three of Vectorless's nine FinanceBench declines in the 2.5 Flash-Lite pilot were retrieval misses. Two were short filings: Foot Locker's 8-K came back as 2 sections, and MGM's earnings release as 5 small sections without the financial table. The third was Boeing, where the walk returned page 7 and the answer is on page 8. - A document whose text fits in judgewalk.whole_document_tokens (default 12,000; negative turns it off) comes back entire, whatever max_sections is. Navigation's ranked units lead, the rest follow in document order, each marked whole_document. - Slots left under max_sections fill first with the pages beside the top evidence page (next, then previous) and then with runner-ups. Each is marked neighbour_of. unit_stats counts neighbours and whole documents. --- cmd/server/main.go | 9 ++- internal/handler/answer.go | 4 +- internal/handler/answer_pages.go | 5 +- internal/handler/query.go | 6 ++ pkg/config/config.go | 9 +++ pkg/retrieval/judgewalk.go | 122 ++++++++++++++++++++++++++----- pkg/retrieval/strategy.go | 18 +++++ pkg/retrieval/units.go | 62 ++++++++++++++-- pkg/retrieval/units_test.go | 73 ++++++++++++++++++ 9 files changed, 278 insertions(+), 30 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index 4046bc8..da4ee1f 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -244,7 +244,7 @@ func run() error { // Corpus mode (HAL-3075): one question over a whole collection, routed // and read on the Judge. Its library caches each document's leaves, // pages and routing card across questions. - corpusLibrary := retrieval.NewCorpusLibrary(buildJudgeWalkStrategy(judge, store, pool), pool.LoadTree) + corpusLibrary := retrieval.NewCorpusLibrary(buildJudgeWalkStrategy(cfg.Engine.Retrieval.JudgeWalk, judge, store, pool), pool.LoadTree) corpusNavigator := retrieval.CorpusNavigator{Judge: judge} // Pre-built set of selectable strategies, keyed by config name. @@ -773,7 +773,7 @@ func buildStrategy(c enginecfg.RetrievalConfig, client llmgate.Client, judge llm // The Judge's provider refusing service (no credit, a rejected key, // an outage) or no Judge being configured right now falls back to // treewalk instead of failing every query. - return retrieval.NewFallbackStrategy(buildJudgeWalkStrategy(judge, store, pool), + return retrieval.NewFallbackStrategy(buildJudgeWalkStrategy(c.JudgeWalk, judge, store, pool), buildTreeWalkStrategy(c, client, store, pool), nil) case "single-pass": return retrieval.NewSinglePass(client) @@ -821,7 +821,7 @@ func buildStrategySet(c enginecfg.RetrievalConfig, client llmgate.Client, judge "auto": retrieval.NewAuto(retrieval.NewSinglePass(client), buildTreeWalkStrategy(c, client, store, pool)), } if judge != nil { - set["judgewalk"] = retrieval.NewFallbackStrategy(buildJudgeWalkStrategy(judge, store, pool), set["treewalk"], nil) + set["judgewalk"] = retrieval.NewFallbackStrategy(buildJudgeWalkStrategy(c.JudgeWalk, judge, store, pool), set["treewalk"], nil) } else { // The same fallback the default builder applies: a request that // names judgewalk on a server with no Judge gets treewalk, not @@ -835,8 +835,9 @@ func buildStrategySet(c enginecfg.RetrievalConfig, client llmgate.Client, judge // leaves ranked in one request, the best sections' bodies ranked in one // or two more, no generative call (HAL-1371). Selectable per request as // strategy=judgewalk whenever llm.judge is configured. -func buildJudgeWalkStrategy(judge llmgate.Judge, store storage.Storage, pool *db.Pool) *retrieval.JudgeWalkStrategy { +func buildJudgeWalkStrategy(c enginecfg.JudgeWalkBlock, judge llmgate.Judge, store storage.Storage, pool *db.Pool) *retrieval.JudgeWalkStrategy { s := retrieval.NewJudgeWalkStrategy(judge) + s.WholeDocumentTokens = c.WholeDocumentTokens s.PageLoader = storagePageLoader{s: store} if pool != nil { s.TOC = dbTOCProvider{db: pool} diff --git a/internal/handler/answer.go b/internal/handler/answer.go index 0d0c869..63ecf23 100644 --- a/internal/handler/answer.go +++ b/internal/handler/answer.go @@ -344,7 +344,9 @@ func (h *AnswerHandler) answerFromSections(ctx context.Context, t *tree.Tree, bo // Each section the answer is written from carries body text; a // heading-only one, or a repeat, gives its slot to the next-ranked // leaf (HAL-3160). - units, _ := retrieval.AssembleUnits(ctx, t, &retrieval.Result{SelectedIDs: ids, RunnerUpIDs: res.RunnerUpIDs}, maxSections, + units, _ := retrieval.AssembleUnits(ctx, t, &retrieval.Result{ + SelectedIDs: ids, RunnerUpIDs: res.RunnerUpIDs, WholeDocument: res.WholeDocument && len(res.DocumentIDs) > 0, DocumentIDs: res.DocumentIDs, + }, maxSections, func(ctx context.Context, ref string) (string, error) { return readStored(ctx, h.storage, ref) }) enriched := make([]answerSection, 0, len(units)) for _, u := range units { diff --git a/internal/handler/answer_pages.go b/internal/handler/answer_pages.go index d6e8520..cadd047 100644 --- a/internal/handler/answer_pages.go +++ b/internal/handler/answer_pages.go @@ -59,7 +59,10 @@ func (h *AnswerHandler) answerFromPages( // fewer than k above its threshold, the pages it ranked next, so a // question that needs two statements is not answered from one // (HAL-3160). - units, _ := retrieval.AssembleUnits(ctx, t, &retrieval.Result{EvidencePages: res.EvidencePages, RunnerUpPages: res.RunnerUpPages}, k, nil) + units, _ := retrieval.AssembleUnits(ctx, t, &retrieval.Result{ + EvidencePages: res.EvidencePages, RunnerUpPages: res.RunnerUpPages, NeighbourPages: res.NeighbourPages, + WholeDocument: res.WholeDocument, DocumentPages: res.DocumentPages, + }, k, nil) evidence := make([]retrieval.EvidencePage, 0, len(units)) for _, u := range units { evidence = append(evidence, u.Page) diff --git a/internal/handler/query.go b/internal/handler/query.go index 8ba0357..07cbb81 100644 --- a/internal/handler/query.go +++ b/internal/handler/query.go @@ -278,6 +278,12 @@ func unitMaps(units []retrieval.Unit) []map[string]any { if u.RunnerUp { m["runner_up"] = true } + if u.Kind == retrieval.UnitPage && u.Page.NeighbourOf > 0 { + m["neighbour_of"] = u.Page.NeighbourOf + } + if u.Whole { + m["whole_document"] = true + } out = append(out, m) } return out diff --git a/pkg/config/config.go b/pkg/config/config.go index 26b2fe7..5b45901 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -491,6 +491,15 @@ type RetrievalConfig struct { Replay ReplayBlock `yaml:"replay"` Abstain AbstainBlock `yaml:"abstain"` TreeWalk TreeWalkBlock `yaml:"treewalk"` + JudgeWalk JudgeWalkBlock `yaml:"judgewalk"` +} + +// JudgeWalkBlock tunes the judgewalk strategy. +type JudgeWalkBlock struct { + // WholeDocumentTokens is the size up to which a document is returned + // entire instead of as a selection (an 8-K, an earnings release). + // Zero takes the strategy's default (12,000); negative turns it off. + WholeDocumentTokens int `yaml:"whole_document_tokens"` } // TreeWalkBlock configures the TreeWalk page-based agentic diff --git a/pkg/retrieval/judgewalk.go b/pkg/retrieval/judgewalk.go index 09c18e3..801e42b 100644 --- a/pkg/retrieval/judgewalk.go +++ b/pkg/retrieval/judgewalk.go @@ -910,6 +910,24 @@ type JudgeWalkStrategy struct { // be unreliable — and scored 0.65 where the pages scored 0.90. TOC TOCProvider Pages PageStore + + // WholeDocumentTokens is the size up to which a document is returned + // entire rather than as a selection: an 8-K or an earnings release fits + // in any answer budget, and picking two of its sections lost the answer + // twice in the pilot (HAL-3246). Zero selects + // DefaultWholeDocumentTokens; negative turns it off. + WholeDocumentTokens int +} + +// DefaultWholeDocumentTokens is the document size returned entire when the +// strategy names none: about twelve 10-K pages. +const DefaultWholeDocumentTokens = 12000 + +func (s *JudgeWalkStrategy) wholeDocumentTokens() int { + if s.WholeDocumentTokens == 0 { + return DefaultWholeDocumentTokens + } + return s.WholeDocumentTokens } // PageStore serves a document's per-page text, as ingest persisted it. @@ -943,6 +961,41 @@ func (s *JudgeWalkStrategy) Select(ctx context.Context, t *tree.Tree, query stri return res.SelectedIDs, nil } +// neighbourPages are the pages beside the top evidence page, the next then +// the previous, that are not evidence themselves. A page navigation also +// scored keeps its score. +func neighbourPages(nav *NavResult, pages []NavPage, leafTitle map[string]string) []EvidencePage { + if nav == nil || len(nav.Evidence) == 0 { + return nil + } + byNum := map[int]NavPage{} + for _, p := range pages { + byNum[p.Number] = p + } + evidence := map[int]bool{} + for _, ev := range nav.Evidence { + evidence[ev.Page.Number] = true + } + score := map[int]float64{} + for _, ps := range nav.Pages { + score[ps.Page.Number] = ps.P + } + top := nav.Evidence[0].Page.Number + var out []EvidencePage + for _, n := range []int{top + 1, top - 1} { + p, ok := byNum[n] + if !ok || evidence[n] { + continue + } + title := leafTitle[p.LeafID] + if title == "" { + title = leafTitle[nav.Evidence[0].Page.LeafID] + } + out = append(out, EvidencePage{Page: n, Title: title, Text: p.Text, Confidence: score[n], NeighbourOf: top}) + } + return out +} + // runnerUps are the pages navigation read in full that did not make the // evidence, best first, at most maxRunnerUps of them. func runnerUps(nav *NavResult) []PageScore { @@ -1045,16 +1098,34 @@ func (s *JudgeWalkStrategy) selectOnPersistedPages(ctx context.Context, t *tree. for _, ps := range runnerUps(nav) { runnerUp = append(runnerUp, EvidencePage{Page: ps.Page.Number, Title: leafTitle[ps.Page.LeafID], Text: ps.Page.Text, Confidence: ps.P}) } + neighbours := neighbourPages(nav, pages, leafTitle) + var whole []EvidencePage + if limit := s.wholeDocumentTokens(); limit > 0 { + chars := 0 + for _, p := range pages { + chars += len(p.Text) + } + if chars/4 <= limit { + sorted := append([]NavPage(nil), pages...) + sort.Slice(sorted, func(i, j int) bool { return sorted[i].Number < sorted[j].Number }) + for _, p := range sorted { + whole = append(whole, EvidencePage{Page: p.Number, Title: leafTitle[p.LeafID], Text: p.Text}) + } + } + } return &Result{ - SelectedIDs: ids, - Confidences: conf, - Confidence: best, - CitedPages: rangesToPairs(ranges), - EvidencePages: evidence, - RunnerUpPages: runnerUp, - ModelUsed: "judge", - Usage: nav.Usage, - HopsTaken: nav.Requests, + SelectedIDs: ids, + Confidences: conf, + Confidence: best, + CitedPages: rangesToPairs(ranges), + EvidencePages: evidence, + RunnerUpPages: runnerUp, + NeighbourPages: neighbours, + WholeDocument: len(whole) > 0, + DocumentPages: whole, + ModelUsed: "judge", + Usage: nav.Usage, + HopsTaken: nav.Requests, }, true } @@ -1266,15 +1337,32 @@ func (s *JudgeWalkStrategy) selectOnSectionTree(ctx context.Context, t *tree.Tre if ordinal { ranges = nil // the units were numbered by position; there are no pages to cite } + var whole []tree.SectionID + if limit := s.wholeDocumentTokens(); limit > 0 { + tokens := 0 + var order []tree.SectionID + t.Root.Walk(func(sec *tree.Section) bool { + if sec.ContentRef != "" { + tokens += sec.TokenCount + order = append(order, sec.ID) + } + return true + }) + if tokens > 0 && tokens <= limit { + whole = order + } + } return &Result{ - SelectedIDs: ids, - RunnerUpIDs: runnerUp, - Confidences: conf, - Confidence: best, - CitedPages: rangesToPairs(ranges), - ModelUsed: "judge", - Usage: nav.Usage, - HopsTaken: nav.Requests, + SelectedIDs: ids, + RunnerUpIDs: runnerUp, + WholeDocument: len(whole) > 0, + DocumentIDs: whole, + Confidences: conf, + Confidence: best, + CitedPages: rangesToPairs(ranges), + ModelUsed: "judge", + Usage: nav.Usage, + HopsTaken: nav.Requests, }, nil } diff --git a/pkg/retrieval/strategy.go b/pkg/retrieval/strategy.go index daa3c85..b581ef6 100644 --- a/pkg/retrieval/strategy.go +++ b/pkg/retrieval/strategy.go @@ -121,6 +121,21 @@ type Result struct { // best first, for strategies whose units are sections. RunnerUpIDs []tree.SectionID `json:"runner_up_ids,omitempty"` + // NeighbourPages are the pages beside the top evidence page that are not + // evidence themselves: the next page, then the previous one. A statement + // often runs across a page break, and a heading often sits a page from + // its answer (HAL-3246: Boeing's cyclicality is on page 8; the walk + // returned page 7). + NeighbourPages []EvidencePage `json:"neighbour_pages,omitempty"` + + // WholeDocument is true when the document is short enough to return + // entire (JudgeWalkStrategy.WholeDocumentTokens). DocumentPages (a paged + // document) or DocumentIDs (a section tree) then hold all of it, in + // document order; the navigation's evidence still leads. + WholeDocument bool `json:"whole_document,omitempty"` + DocumentPages []EvidencePage `json:"document_pages,omitempty"` + DocumentIDs []tree.SectionID `json:"document_ids,omitempty"` + // CitedPages is the FINAL set of page ranges the answer commits // to — the model's cited_pages after dedup and the confidence cap, // NOT every page it read (that is PagesRead). Page-based strategies @@ -146,6 +161,9 @@ type EvidencePage struct { Title string `json:"title,omitempty"` // the section the page belongs to Text string `json:"text"` Confidence float64 `json:"confidence"` + // NeighbourOf is the evidence page this one sits beside, when it was + // added as a neighbour rather than judged. + NeighbourOf int `json:"neighbour_of,omitempty"` } // PageReadEntry is one get_pages tool call that materialised during a diff --git a/pkg/retrieval/units.go b/pkg/retrieval/units.go index 2a83efe..dd5d22c 100644 --- a/pkg/retrieval/units.go +++ b/pkg/retrieval/units.go @@ -58,14 +58,19 @@ type Unit struct { // RunnerUp is true when the unit was ranked below the strategy's // evidence and fills a slot the evidence left open. RunnerUp bool + // Whole is true when the unit is part of a short document returned + // entire, past the strategy's own ranking. + Whole bool } // UnitStats counts what assembly left out and filled in. type UnitStats struct { - Empty int `json:"empty"` // sections dropped for having no body text - Duplicate int `json:"duplicate"` // sections dropped as repeats of units already returned - Expanded int `json:"expanded"` // heading-only sections filled from the text beneath them - RunnerUps int `json:"runner_ups"` // slots filled from below the evidence + Empty int `json:"empty"` // sections dropped for having no body text + Duplicate int `json:"duplicate"` // sections dropped as repeats of units already returned + Expanded int `json:"expanded"` // heading-only sections filled from the text beneath them + RunnerUps int `json:"runner_ups"` // slots filled from below the evidence + Neighbours int `json:"neighbours"` // slots filled by a page beside the top evidence page + Whole bool `json:"whole_document"` } // ContentLoader reads a section's stored text by its content reference. @@ -78,14 +83,21 @@ type ContentLoader func(ctx context.Context, ref string) (string, error) // takes the text beneath it when that is compact, and is dropped when // it is not; a section whose text the units before it already carry is // dropped as a repeat; -// 3. when k is set and slots remain, the pages and then the leaves the -// strategy ranked just below its evidence, under the same rules. +// 3. when k is set and slots remain, the pages beside the top evidence +// page, then the pages and the leaves the strategy ranked just below +// its evidence, under the same rules. // // k <= 0 returns every unit of steps 1 and 2 and fills nothing. +// +// A short document (res.WholeDocument) comes back entire whatever k is: its +// ranked units first, then the rest of it in document order. func AssembleUnits(ctx context.Context, t *tree.Tree, res *Result, k int, load ContentLoader) ([]Unit, UnitStats) { if res == nil { return nil, UnitStats{} } + if res.WholeDocument { + return assembleWhole(ctx, t, res, load) + } a := NewUnitAssembler(t, k, load) for _, ep := range res.EvidencePages { a.AddPage(ep, false) @@ -94,6 +106,9 @@ func AssembleUnits(ctx context.Context, t *tree.Tree, res *Result, k int, load C a.AddSection(ctx, id, false) } if k > 0 { + for _, ep := range res.NeighbourPages { + a.AddPage(ep, true) + } for _, ep := range res.RunnerUpPages { a.AddPage(ep, true) } @@ -104,6 +119,36 @@ func AssembleUnits(ctx context.Context, t *tree.Tree, res *Result, k int, load C return a.Units(), a.Stats() } +// assembleWhole returns a short document entire: what navigation ranked, +// best first, then the rest in document order, with no cap. +func assembleWhole(ctx context.Context, t *tree.Tree, res *Result, load ContentLoader) ([]Unit, UnitStats) { + a := NewUnitAssembler(t, 0, load) + for _, ep := range res.EvidencePages { + a.AddPage(ep, false) + } + for _, id := range res.SelectedIDs { + a.AddSection(ctx, id, false) + } + for _, ep := range res.RunnerUpPages { + a.AddPage(ep, true) + } + for _, id := range res.RunnerUpIDs { + a.AddSection(ctx, id, true) + } + ranked := len(a.units) + for _, ep := range res.DocumentPages { + a.AddPage(ep, false) + } + for _, id := range res.DocumentIDs { + a.AddSection(ctx, id, false) + } + for i := ranked; i < len(a.units); i++ { + a.units[i].Whole = true + } + a.stats.Whole = true + return a.Units(), a.Stats() +} + // UnitAssembler applies AssembleUnits's rules one unit at a time, for a // caller that receives units as they are found (a stream). type UnitAssembler struct { @@ -142,7 +187,10 @@ func (a *UnitAssembler) AddPage(ep EvidencePage, runnerUp bool) (Unit, bool) { u := Unit{Kind: UnitPage, Page: ep, Content: ep.Text, RunnerUp: runnerUp} a.units = append(a.units, u) addShingles(a.seen, ep.Text) - if runnerUp { + switch { + case ep.NeighbourOf > 0: + a.stats.Neighbours++ + case runnerUp: a.stats.RunnerUps++ } return u, true diff --git a/pkg/retrieval/units_test.go b/pkg/retrieval/units_test.go index d135e0e..aeab388 100644 --- a/pkg/retrieval/units_test.go +++ b/pkg/retrieval/units_test.go @@ -3,6 +3,7 @@ package retrieval import ( "context" "errors" + "fmt" "strings" "testing" @@ -194,3 +195,75 @@ func TestSubstantiveSurvivesCaseFoldingThatChangesLength(t *testing.T) { t.Errorf("substantive = %q", got) } } + +// A short document comes back entire whatever the cap: what navigation +// ranked first, then the rest in document order (HAL-3246). +func TestShortDocumentComesBackWhole(t *testing.T) { + res := &Result{ + WholeDocument: true, + EvidencePages: []EvidencePage{{Page: 63, Text: runnerUpPage, Confidence: 0.8}}, + DocumentPages: []EvidencePage{{Page: 62, Text: incomePage}, {Page: 63, Text: runnerUpPage}, {Page: 64, Text: "Page 64. " + cashText}}, + } + units, stats := AssembleUnits(context.Background(), filingTree(), res, 2, filingStore()) + got := []int{} + for _, u := range units { + got = append(got, u.Page.Page) + } + if fmt.Sprint(got) != "[63 62 64]" { + t.Fatalf("pages = %v, want the evidence first then the rest in order, past the cap of 2", got) + } + if units[0].Whole || !units[1].Whole || !units[2].Whole || !stats.Whole { + t.Errorf("whole marks: %+v %+v", units, stats) + } +} + +func TestShortSectionDocumentComesBackWhole(t *testing.T) { + res := &Result{WholeDocument: true, SelectedIDs: []tree.SectionID{"sec_cash"}, + DocumentIDs: []tree.SectionID{"sec_risk", "sec_income", "sec_cash", "sec_note_body"}} + units, _ := AssembleUnits(context.Background(), filingTree(), res, 1, filingStore()) + if got := strings.Join(ids(units), ","); got != "sec_cash,sec_risk,sec_income,sec_note_body" { + t.Fatalf("units = %s", got) + } +} + +// The page beside the top evidence page fills before pages ranked from +// elsewhere in the document. +func TestNeighbourFillsBeforeRunnerUps(t *testing.T) { + res := &Result{ + EvidencePages: []EvidencePage{{Page: 62, Text: incomePage, Confidence: 0.9}}, + NeighbourPages: []EvidencePage{{Page: 63, Text: runnerUpPage, NeighbourOf: 62}}, + RunnerUpPages: []EvidencePage{{Page: 9, Text: "Page 9. " + riskText, Confidence: 0.2}}, + } + units, stats := AssembleUnits(context.Background(), filingTree(), res, 2, filingStore()) + if len(units) != 2 || units[1].Page.Page != 63 || units[1].Page.NeighbourOf != 62 { + t.Fatalf("units = %+v", units) + } + if stats.Neighbours != 1 || stats.RunnerUps != 0 { + t.Errorf("stats = %+v", stats) + } +} + +func TestNeighbourPagesAreNextThenPreviousAndNeverEvidence(t *testing.T) { + p := func(n int) NavPage { return NavPage{Number: n, Text: fmt.Sprintf("page %d text", n), LeafID: "l"} } + pages := []NavPage{p(6), p(7), p(8), p(9)} + nav := &NavResult{ + Evidence: []PageScore{{Page: p(7), P: 0.9}, {Page: p(6), P: 0.7}}, + Pages: []PageScore{{Page: p(7), P: 0.9}, {Page: p(6), P: 0.7}, {Page: p(8), P: 0.3}}, + } + got := neighbourPages(nav, pages, map[string]string{"l": "Business"}) + if len(got) != 1 || got[0].Page != 8 || got[0].NeighbourOf != 7 || got[0].Confidence != 0.3 || got[0].Title != "Business" { + t.Fatalf("neighbours = %+v, want page 8 only (6 is evidence)", got) + } + if neighbourPages(&NavResult{}, pages, nil) != nil { + t.Error("no evidence, no neighbours") + } +} + +func TestWholeDocumentThreshold(t *testing.T) { + if (&JudgeWalkStrategy{}).wholeDocumentTokens() != DefaultWholeDocumentTokens { + t.Error("zero takes the default") + } + if (&JudgeWalkStrategy{WholeDocumentTokens: -1}).wholeDocumentTokens() > 0 { + t.Error("negative turns it off") + } +}