feat(retrieval): return a short document whole, and the top evidence page's neighbour

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}

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 {

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)

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

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

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
 }
 

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

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