From 4d186e55f856b1a091d3d77db19bf7a2f8915210 Mon Sep 17 00:00:00 2001 From: hallelx2 Date: Sat, 17 Oct 2026 00:33:14 +0100 Subject: [PATCH 1/2] fix(retrieval): return only units with body text, filling open slots from the next-ranked evidence On paged documents the walk returned its evidence pages plus every section overlapping them. That set includes cover and part headings that span the whole filing ("THE BOEING COMPANY", "PART II") and carry no body text, and sections the pages already carry. On the pilot's 50 FinanceBench questions, 39 returned units were under 150 characters and 29 questions had at least one, so the answer step often saw two pages and three headings. Units are now assembled in one place (retrieval.AssembleUnits) for /v1/query, its stream, /v1/query/multi, /v1/answer and the standalone server. A heading-only section takes the text beneath it when it is compact (two pages, 6,000 characters) and is dropped otherwise. A section the earlier units already carry is dropped as a repeat. Slots left under max_sections fill from the pages and leaves the walk ranked just below its evidence, which it had read and scored and then discarded. Responses carry unit_stats, and fill-ins are marked runner_up. --- internal/api/server.go | 123 +++++----- internal/handler/answer.go | 29 +-- internal/handler/answer_pages.go | 11 +- internal/handler/query.go | 102 +++++--- internal/handler/query_multi.go | 33 +-- internal/handler/query_strategy_test.go | 10 +- internal/handler/query_stream.go | 28 +-- pkg/retrieval/judgewalk.go | 39 +++ pkg/retrieval/strategy.go | 10 + pkg/retrieval/units.go | 301 ++++++++++++++++++++++++ pkg/retrieval/units_test.go | 185 +++++++++++++++ 11 files changed, 698 insertions(+), 173 deletions(-) create mode 100644 pkg/retrieval/units.go create mode 100644 pkg/retrieval/units_test.go diff --git a/internal/api/server.go b/internal/api/server.go index 34eeb3c..b083cbc 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -636,27 +636,14 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { return } - if body.MaxSections > 0 && len(ids) > body.MaxSections { - ids = ids[:body.MaxSections] - } - - enriched := make([]sectionWithContent, 0, len(ids)) - for _, id := range ids { - sec := t.FindByID(id) - if sec == nil { - continue - } - var content string - if sec.ContentRef != "" { - rc, _, err := d.Storage.Get(r.Context(), sec.ContentRef) - if err == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } - enriched = append(enriched, sectionWithContent{sec: sec, content: content}) + // Every unit carries body text, and max_sections caps pages and + // sections together: a heading-only section, or one the evidence pages + // already carry, gives its slot to the next-ranked evidence (HAL-3160). + res := selResult + if res == nil { + res = &retrieval.Result{SelectedIDs: ids} } + pages, enriched, runnerUpPages, unitStats := d.assembleUnits(r.Context(), t, res, body.MaxSections) // Optional: content-aware re-rank pass. One LLM call that scores // each loaded section against the query and re-orders the slice @@ -677,18 +664,21 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { runSpansConcurrent(r.Context(), extractor, body.Query, enriched, d.AnswerSpan.MaxConcurrency, d.Logger) } - sections := make([]map[string]any, 0, len(enriched)) + sections := make([]map[string]any, 0, len(pages)+len(enriched)+len(runnerUpPages)) finalIDs := make([]tree.SectionID, 0, len(enriched)) // A page-based strategy's evidence pages lead the list, as-is: the // page the Judge found holds the answer; a section mapped to it by - // the parser's page attribution may not (HAL-1390). - if selResult != nil { - sections = append(sections, evidencePageSections(selResult.EvidencePages)...) - } + // the parser's page attribution may not (HAL-1390). Pages ranked just + // below the evidence close the list when there was room for them. + sections = append(sections, evidencePageSections(pages)...) for _, e := range enriched { sections = append(sections, sectionWithContentToMap(e)) finalIDs = append(finalIDs, e.sec.ID) } + for _, m := range evidencePageSections(runnerUpPages) { + m["runner_up"] = true + sections = append(sections, m) + } // Trace token is computed over the FINAL IDs that ship in the // response (after max_sections + ReRank truncation). Recomputing @@ -709,7 +699,8 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { "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), + "usage": usageJSON(totalUsage), + "unit_stats": unitStats, } if plan != nil { resp["plan"] = plan @@ -740,6 +731,41 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) { }) } +// assembleUnits builds the units a response returns (retrieval.AssembleUnits) +// and splits them the way the handlers use them: evidence pages, sections +// with their text (re-rank and span extraction work on these), and pages +// ranked below the evidence that fill the slots left over. +func (d Deps) assembleUnits(ctx context.Context, t *tree.Tree, res *retrieval.Result, k int) ([]retrieval.EvidencePage, []sectionWithContent, []retrieval.EvidencePage, retrieval.UnitStats) { + units, stats := retrieval.AssembleUnits(ctx, t, res, k, d.loadContent) + var pages, runnerUps []retrieval.EvidencePage + var secs []sectionWithContent + for _, u := range units { + switch { + case u.Kind == retrieval.UnitPage && u.RunnerUp: + runnerUps = append(runnerUps, u.Page) + case u.Kind == retrieval.UnitPage: + pages = append(pages, u.Page) + default: + secs = append(secs, sectionWithContent{sec: u.Section, content: u.Content}) + } + } + return pages, secs, runnerUps, stats +} + +// loadContent reads a section's stored text. +func (d Deps) loadContent(ctx context.Context, ref string) (string, error) { + rc, _, err := d.Storage.Get(ctx, ref) + if err != nil { + return "", err + } + defer rc.Close() // best-effort close + raw, err := io.ReadAll(rc) + if err != nil { + return "", err + } + return string(raw), nil +} + // sectionWithContent bundles a tree section with its loaded content // and optional re-rank score / answer-span. Used by /v1/query and // /v1/answer. @@ -946,28 +972,9 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) { if maxSections <= 0 { maxSections = 5 } - if len(ids) > maxSections { - ids = ids[:maxSections] - } - - // Load each section's content. - enriched := make([]sectionWithContent, 0, len(ids)) - for _, id := range ids { - sec := t.FindByID(id) - if sec == nil { - continue - } - var content string - if sec.ContentRef != "" { - rc, _, err := d.Storage.Get(r.Context(), sec.ContentRef) - if err == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } - enriched = append(enriched, sectionWithContent{sec: sec, content: content}) - } + // Each section the answer is written from carries body text + // (HAL-3160). + _, enriched, _, _ := d.assembleUnits(r.Context(), t, &retrieval.Result{SelectedIDs: ids}, maxSections) // Optional: content-aware re-rank before synthesis sees the // evidence. When TopK is set the synthesis prompt only ever sees @@ -1194,20 +1201,9 @@ func (d Deps) handleQueryMulti(w http.ResponseWriter, r *http.Request) { docs := make([]map[string]any, 0, len(result.Documents)) for docID, dr := range result.Documents { sections := make([]map[string]any, 0, len(dr.SelectedIDs)) - for _, sid := range dr.SelectedIDs { - sec := dr.Tree.FindByID(sid) - if sec == nil { - continue - } - var content string - if sec.ContentRef != "" { - rc, _, err := d.Storage.Get(r.Context(), sec.ContentRef) - if err == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } + _, secs, _, _ := d.assembleUnits(r.Context(), dr.Tree, &retrieval.Result{SelectedIDs: dr.SelectedIDs}, body.MaxSections) + for _, e := range secs { + sec, content := e.sec, e.content s := map[string]any{ "id": sec.ID, "parent_id": sec.ParentID, @@ -1226,9 +1222,6 @@ func (d Deps) handleQueryMulti(w http.ResponseWriter, r *http.Request) { s["candidate_questions"] = sec.CandidateQuestions } sections = append(sections, s) - if body.MaxSections > 0 && len(sections) >= body.MaxSections { - break - } } docs = append(docs, map[string]any{ "document_id": docID, diff --git a/internal/handler/answer.go b/internal/handler/answer.go index 48d88c7..4b3968b 100644 --- a/internal/handler/answer.go +++ b/internal/handler/answer.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "io" "log/slog" "net/http" "strings" @@ -341,26 +340,14 @@ func (h *AnswerHandler) answerFromSections(ctx context.Context, t *tree.Tree, bo if maxSections <= 0 { maxSections = 5 } - if len(ids) > maxSections { - ids = ids[:maxSections] - } - - enriched := make([]answerSection, 0, len(ids)) - for _, id := range ids { - sec := t.FindByID(id) - if sec == nil { - continue - } - var content string - if sec.ContentRef != "" { - rc, _, getErr := h.storage.Get(ctx, sec.ContentRef) - if getErr == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } - enriched = append(enriched, answerSection{sec: sec, content: content}) + // 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, + func(ctx context.Context, ref string) (string, error) { return readStored(ctx, h.storage, ref) }) + enriched := make([]answerSection, 0, len(units)) + for _, u := range units { + enriched = append(enriched, answerSection{sec: u.Section, content: u.Content}) } spanExtractor := h.spanExtractor(body.Model) diff --git a/internal/handler/answer_pages.go b/internal/handler/answer_pages.go index 1bb3b63..350f26b 100644 --- a/internal/handler/answer_pages.go +++ b/internal/handler/answer_pages.go @@ -54,9 +54,14 @@ func (h *AnswerHandler) answerFromPages( if k <= 0 { k = defaultEvidencePages } - evidence := res.EvidencePages - if len(evidence) > k { - evidence = evidence[:k] + // The answer reads the evidence pages and, when the navigator found + // 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) + evidence := make([]retrieval.EvidencePage, 0, len(units)) + for _, u := range units { + evidence = append(evidence, u.Page) } retrievalMS := time.Since(started).Milliseconds() diff --git a/internal/handler/query.go b/internal/handler/query.go index 53d3763..8ba0357 100644 --- a/internal/handler/query.go +++ b/internal/handler/query.go @@ -161,6 +161,7 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) { ids []tree.SectionID usage *retrieval.Usage evidence []retrieval.EvidencePage + full *retrieval.Result ) if cs, ok := strategy.(retrieval.CostStrategy); ok { result, err := cs.SelectWithCost(r.Context(), t, body.Query, budget) @@ -175,6 +176,7 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) { ids = result.SelectedIDs usage = &result.Usage evidence = result.EvidencePages + full = result } else { var err error ids, err = strategy.Select(r.Context(), t, body.Query, budget) @@ -192,41 +194,15 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) { // page the Judge chose holds the answer, while the sections overlapping // it come back in document order (HAL-1390). max_sections caps the // combined list, so a small cap keeps the evidence, not the first - // sections of the document (HAL-3071). - sections := evidencePageSections(evidence) - if body.MaxSections > 0 { - if len(sections) > body.MaxSections { - sections = sections[:body.MaxSections] - } - if room := body.MaxSections - len(sections); len(ids) > room { - ids = ids[:room] - } - } - - for _, id := range ids { - sec := t.FindByID(id) - if sec == nil { - continue - } - var content string - if sec.ContentRef != "" { - rc, _, getErr := h.storage.Get(r.Context(), sec.ContentRef) - if getErr == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } - sections = append(sections, map[string]any{ - "id": sec.ID, - "parent_id": sec.ParentID, - "title": sec.Title, - "summary": sec.Summary, - "token_count": sec.TokenCount, - "content": content, - }) + // sections of the document (HAL-3071). Every unit carries body text: a + // heading-only section, or one the pages already carry, gives its slot + // to the next-ranked evidence (HAL-3160). + res := &retrieval.Result{SelectedIDs: ids, EvidencePages: evidence} + if full != nil { + res = full } - + units, unitStats := retrieval.AssembleUnits(r.Context(), t, res, body.MaxSections, h.loadContent) + sections := unitMaps(units) resp := map[string]any{ "document_id": body.DocumentID, "query": body.Query, @@ -234,6 +210,7 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) { "model": body.Model, "sections": sections, "elapsed_ms": time.Since(started).Milliseconds(), + "unit_stats": unitStats, } if usage != nil { resp["usage"] = map[string]any{ @@ -249,6 +226,63 @@ func (h *QueryHandler) HandleQuery(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, resp) } +// loadContent reads a section's stored text. +func (h *QueryHandler) loadContent(ctx context.Context, ref string) (string, error) { + return readStored(ctx, h.storage, ref) +} + +// readStored reads one stored object as text. +func readStored(ctx context.Context, store storage.Storage, ref string) (string, error) { + rc, _, err := store.Get(ctx, ref) + if err != nil { + return "", err + } + defer rc.Close() // best-effort close + raw, err := io.ReadAll(rc) + if err != nil { + return "", err + } + return string(raw), nil +} + +// unitMaps renders assembled units in the response shape: a page as +// evidencePageSections renders it, a section with its tree fields. A unit +// that fills a slot below the evidence says so (runner_up), and a heading +// whose text was gathered from the sections beneath it says that +// (expanded). +func unitMaps(units []retrieval.Unit) []map[string]any { + out := make([]map[string]any, 0, len(units)) + for _, u := range units { + var m map[string]any + if u.Kind == retrieval.UnitPage { + m = evidencePageSections([]retrieval.EvidencePage{u.Page})[0] + } else { + m = map[string]any{ + "id": u.Section.ID, + "parent_id": u.Section.ParentID, + "title": u.Section.Title, + "summary": u.Section.Summary, + "token_count": u.Section.TokenCount, + "content": u.Content, + } + if u.Section.PageStart > 0 { + m["page_start"] = u.Section.PageStart + } + if u.Section.PageEnd > 0 { + m["page_end"] = u.Section.PageEnd + } + if u.Expanded { + m["expanded"] = true + } + } + if u.RunnerUp { + m["runner_up"] = true + } + out = append(out, m) + } + return out +} + // evidencePageSections turns a page-based strategy's evidence pages into // response units: id "page_", the owning section's title with the page // ("
· p."), the page text as content. The same shape the diff --git a/internal/handler/query_multi.go b/internal/handler/query_multi.go index 61ccdc6..76f0a44 100644 --- a/internal/handler/query_multi.go +++ b/internal/handler/query_multi.go @@ -1,8 +1,8 @@ package handler import ( + "context" "encoding/json" - "io" "log/slog" "net/http" "time" @@ -95,33 +95,10 @@ func (h *QueryMultiHandler) HandleQueryMulti(w http.ResponseWriter, r *http.Requ // Build per-document response. docs := make([]map[string]any, 0, len(result.Documents)) for docID, dr := range result.Documents { - sections := make([]map[string]any, 0, len(dr.SelectedIDs)) - for _, sid := range dr.SelectedIDs { - sec := dr.Tree.FindByID(sid) - if sec == nil { - continue - } - var content string - if sec.ContentRef != "" { - rc, _, getErr := h.storage.Get(r.Context(), sec.ContentRef) - if getErr == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } - sections = append(sections, map[string]any{ - "id": sec.ID, - "parent_id": sec.ParentID, - "title": sec.Title, - "summary": sec.Summary, - "token_count": sec.TokenCount, - "content": content, - }) - if body.MaxSections > 0 && len(sections) >= body.MaxSections { - break - } - } + // Every section carries body text (HAL-3160). + units, _ := retrieval.AssembleUnits(r.Context(), dr.Tree, &retrieval.Result{SelectedIDs: dr.SelectedIDs}, body.MaxSections, + func(ctx context.Context, ref string) (string, error) { return readStored(ctx, h.storage, ref) }) + sections := unitMaps(units) docs = append(docs, map[string]any{ "document_id": docID, "sections": sections, diff --git a/internal/handler/query_strategy_test.go b/internal/handler/query_strategy_test.go index 8ca5c84..1387da9 100644 --- a/internal/handler/query_strategy_test.go +++ b/internal/handler/query_strategy_test.go @@ -81,6 +81,12 @@ func (m *memStorage) SignedURL(ctx context.Context, key string, expiry time.Dura return "", nil } +// Section bodies long enough to be returned as units (retrieval.MinUnitChars). +const ( + sectionAText = "section A content: revenue rose to $4.2 billion on higher volumes in the Americas, while gross margin narrowed by forty basis points as freight and labour costs climbed through the year." + sectionBText = "section B content: capital expenditure was $1,577 million, mostly new distribution centres, and the company expects spending to stay near that level while it finishes the network build-out." +) + func queryStrategyTestTree() *tree.Tree { a := &tree.Section{ID: "sec_a", Title: "A", ContentRef: "a_ref", PageStart: 1, PageEnd: 2} b := &tree.Section{ID: "sec_b", Title: "B", ContentRef: "b_ref", PageStart: 3, PageEnd: 4} @@ -93,8 +99,8 @@ func queryStrategyTestTree() *tree.Tree { // so HandleQuery runs end-to-end via httptest without a real DB. func newQueryStrategyHandler(def retrieval.Strategy, set map[string]retrieval.Strategy) *QueryHandler { store := &memStorage{data: map[string][]byte{ - "a_ref": []byte("section A content"), - "b_ref": []byte("section B content"), + "a_ref": []byte(sectionAText), + "b_ref": []byte(sectionBText), }} h := NewQueryHandler(slog.Default(), nil, store, def, set) h.treeLoader = func(ctx context.Context, orgID, storeID string, docID tree.DocumentID) (*tree.Tree, error) { diff --git a/internal/handler/query_stream.go b/internal/handler/query_stream.go index d7a7e3d..b65824e 100644 --- a/internal/handler/query_stream.go +++ b/internal/handler/query_stream.go @@ -1,10 +1,10 @@ package handler import ( + "context" "encoding/json" "errors" "fmt" - "io" "log/slog" "net/http" "time" @@ -120,6 +120,7 @@ func (h *QueryStreamHandler) HandleQueryStream(w http.ResponseWriter, r *http.Re started := time.Now() events := ss.SelectStream(ctx, t, body.Query, budget) + units := retrieval.NewUnitAssembler(t, 0, func(ctx context.Context, ref string) (string, error) { return readStored(ctx, h.storage, ref) }) for evt := range events { sse := sseEvent{ @@ -136,26 +137,13 @@ func (h *QueryStreamHandler) HandleQueryStream(w http.ResponseWriter, r *http.Re sse.Query = body.Query case retrieval.EventSectionSelected: - sec := t.FindByID(evt.SectionID) - if sec != nil { - var content string - if sec.ContentRef != "" { - rc, _, getErr := h.storage.Get(ctx, sec.ContentRef) - if getErr == nil { - raw, _ := io.ReadAll(rc) - _ = rc.Close() // best-effort close - content = string(raw) - } - } - sse.Section = map[string]any{ - "id": sec.ID, - "parent_id": sec.ParentID, - "title": sec.Title, - "summary": sec.Summary, - "token_count": sec.TokenCount, - "content": content, - } + // A section with no body text, or one an earlier event already + // carried, is not sent (HAL-3160). + u, ok := units.AddSection(ctx, evt.SectionID, false) + if !ok { + continue } + sse.Section = unitMaps([]retrieval.Unit{u})[0] case retrieval.EventCompleted: sse.ElapsedMS = time.Since(started).Milliseconds() diff --git a/pkg/retrieval/judgewalk.go b/pkg/retrieval/judgewalk.go index 6912287..09c18e3 100644 --- a/pkg/retrieval/judgewalk.go +++ b/pkg/retrieval/judgewalk.go @@ -943,6 +943,32 @@ func (s *JudgeWalkStrategy) Select(ctx context.Context, t *tree.Tree, query stri return res.SelectedIDs, nil } +// 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 { + type key struct { + leaf string + page int + } + in := map[key]bool{} + for _, ev := range nav.Evidence { + in[key{ev.Page.LeafID, ev.Page.Number}] = true + } + var out []PageScore + for _, ps := range nav.Pages { + k := key{ps.Page.LeafID, ps.Page.Number} + if in[k] { + continue + } + in[k] = true + out = append(out, ps) + if len(out) == maxRunnerUps { + break + } + } + return out +} + // SelectWithCost runs navigation and reports usage. func (s *JudgeWalkStrategy) SelectWithCost(ctx context.Context, t *tree.Tree, query string, budget ContextBudget) (*Result, error) { if t == nil || t.Root == nil { @@ -1015,12 +1041,17 @@ func (s *JudgeWalkStrategy) selectOnPersistedPages(ctx context.Context, t *tree. for _, ev := range nav.Evidence { evidence = append(evidence, EvidencePage{Page: ev.Page.Number, Title: leafTitle[ev.Page.LeafID], Text: ev.Page.Text, Confidence: ev.P}) } + var runnerUp []EvidencePage + 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}) + } return &Result{ SelectedIDs: ids, Confidences: conf, Confidence: best, CitedPages: rangesToPairs(ranges), EvidencePages: evidence, + RunnerUpPages: runnerUp, ModelUsed: "judge", Usage: nav.Usage, HopsTaken: nav.Requests, @@ -1225,11 +1256,19 @@ func (s *JudgeWalkStrategy) selectOnSectionTree(ctx context.Context, t *tree.Tre if len(nav.Evidence) > 0 { best = nav.Evidence[0].P } + var runnerUp []tree.SectionID + for _, ps := range runnerUps(nav) { + if id := tree.SectionID(ps.Page.LeafID); !seen[ps.Page.LeafID] { + seen[ps.Page.LeafID] = true + runnerUp = append(runnerUp, id) + } + } if ordinal { ranges = nil // the units were numbered by position; there are no pages to cite } return &Result{ SelectedIDs: ids, + RunnerUpIDs: runnerUp, Confidences: conf, Confidence: best, CitedPages: rangesToPairs(ranges), diff --git a/pkg/retrieval/strategy.go b/pkg/retrieval/strategy.go index 76d5761..daa3c85 100644 --- a/pkg/retrieval/strategy.go +++ b/pkg/retrieval/strategy.go @@ -111,6 +111,16 @@ type Result struct { // strategies. EvidencePages []EvidencePage `json:"evidence_pages,omitempty"` + // RunnerUpPages are pages a page-based strategy read in full and + // ranked just below its evidence, best first. A response that has + // room after the evidence fills it from here rather than from + // sections the evidence already covers (HAL-3160). + RunnerUpPages []EvidencePage `json:"runner_up_pages,omitempty"` + + // RunnerUpIDs are the leaves ranked just below the selected ones, + // best first, for strategies whose units are sections. + RunnerUpIDs []tree.SectionID `json:"runner_up_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 diff --git a/pkg/retrieval/units.go b/pkg/retrieval/units.go new file mode 100644 index 0000000..a79f6e5 --- /dev/null +++ b/pkg/retrieval/units.go @@ -0,0 +1,301 @@ +package retrieval + +import ( + "context" + "strings" + "unicode" + + "github.com/hallelx2/vectorless-engine/pkg/tree" +) + +// MinUnitChars is the least body text a returned unit carries, after its +// own heading is taken off. A shorter unit is a heading ("THE BOEING +// COMPANY", "PART II") or a stub, and it would take a top-k slot from real +// evidence (HAL-3160: 17% of FinanceBench units were under this, 18 of 239 +// empty). +const MinUnitChars = 150 + +// unitExpandChars bounds how much text a heading-only section may gather +// from its descendants. A heading over a compact block becomes that block; +// a cover heading over the whole document is dropped instead of becoming +// the document's first pages. +const unitExpandChars = 6000 + +// unitExpandPages is the most pages a heading may span and still be stood +// in for by the text beneath it. +const unitExpandPages = 2 + +// maxRunnerUps caps the runner-up pages and leaves a strategy reports. +const maxRunnerUps = 8 + +// duplicateShare is the share of a section's word shingles that, when +// already present in the units returned before it, makes it a repeat of +// them. +const duplicateShare = 0.8 + +const shingleWords = 6 + +// UnitKind says what a returned unit is. +type UnitKind string + +const ( + // UnitPage is one page of the document, as a page-based strategy read it. + UnitPage UnitKind = "page" + // UnitSection is one section of the document tree. + UnitSection UnitKind = "section" +) + +// Unit is one piece of evidence a query returns: a page or a section, +// with the text the caller reads. +type Unit struct { + Kind UnitKind + Page EvidencePage // Kind == UnitPage + Section *tree.Section // Kind == UnitSection + Content string + // Expanded is true when a heading-only section's text was gathered + // from the sections beneath it. + Expanded bool + // RunnerUp is true when the unit was ranked below the strategy's + // evidence and fills a slot the evidence left open. + RunnerUp 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 +} + +// ContentLoader reads a section's stored text by its content reference. +type ContentLoader func(ctx context.Context, ref string) (string, error) + +// AssembleUnits builds the evidence a query returns, best first: +// +// 1. the strategy's evidence pages, as found; +// 2. its selected sections, each with body text: a heading-only section +// 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. +// +// k <= 0 returns every unit of steps 1 and 2 and fills nothing. +func AssembleUnits(ctx context.Context, t *tree.Tree, res *Result, k int, load ContentLoader) ([]Unit, UnitStats) { + if res == nil { + return nil, UnitStats{} + } + a := NewUnitAssembler(t, k, load) + for _, ep := range res.EvidencePages { + a.AddPage(ep, false) + } + for _, id := range res.SelectedIDs { + a.AddSection(ctx, id, false) + } + if k > 0 { + for _, ep := range res.RunnerUpPages { + a.AddPage(ep, true) + } + for _, id := range res.RunnerUpIDs { + a.AddSection(ctx, id, 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 { + t *tree.Tree + k int + load ContentLoader + units []Unit + stats UnitStats + seen map[string]bool + pages map[int]bool + ids map[tree.SectionID]bool +} + +// NewUnitAssembler starts an empty assembly capped at k units (k <= 0: +// no cap). +func NewUnitAssembler(t *tree.Tree, k int, load ContentLoader) *UnitAssembler { + return &UnitAssembler{t: t, k: k, load: load, seen: map[string]bool{}, pages: map[int]bool{}, ids: map[tree.SectionID]bool{}} +} + +// Full reports whether the cap is reached. +func (a *UnitAssembler) Full() bool { return a.k > 0 && len(a.units) >= a.k } + +// Units are the units accepted so far, in order. +func (a *UnitAssembler) Units() []Unit { return a.units } + +// Stats counts what was left out and filled in. +func (a *UnitAssembler) Stats() UnitStats { return a.stats } + +// AddPage accepts a page unless the cap is reached, the page is already in, +// or it has no text. +func (a *UnitAssembler) AddPage(ep EvidencePage, runnerUp bool) (Unit, bool) { + if a.Full() || a.pages[ep.Page] || len(substantive(ep.Text, "")) == 0 { + return Unit{}, false + } + a.pages[ep.Page] = true + u := Unit{Kind: UnitPage, Page: ep, Content: ep.Text, RunnerUp: runnerUp} + a.units = append(a.units, u) + addShingles(a.seen, ep.Text) + if runnerUp { + a.stats.RunnerUps++ + } + return u, true +} + +// AddSection accepts a section with its body text, gathering a compact +// heading's text from beneath it; it refuses one with no body text, a +// repeat of the units already in, or any once the cap is reached. +func (a *UnitAssembler) AddSection(ctx context.Context, id tree.SectionID, runnerUp bool) (Unit, bool) { + if a.Full() || a.ids[id] || a.t == nil { + return Unit{}, false + } + sec := a.t.FindByID(id) + if sec == nil { + return Unit{}, false + } + a.ids[id] = true + content := loadContent(ctx, sec, a.load) + expanded := false + if len(substantive(content, sec.Title)) < MinUnitChars { + gathered, ok := gatherBeneath(ctx, sec, a.load) + if !ok || len(substantive(gathered, sec.Title)) < MinUnitChars { + a.stats.Empty++ + return Unit{}, false + } + content, expanded = gathered, true + } + if repeats(a.seen, content) { + a.stats.Duplicate++ + return Unit{}, false + } + u := Unit{Kind: UnitSection, Section: sec, Content: content, Expanded: expanded, RunnerUp: runnerUp} + a.units = append(a.units, u) + addShingles(a.seen, content) + if expanded { + a.stats.Expanded++ + } + if runnerUp { + a.stats.RunnerUps++ + } + return u, true +} + +func loadContent(ctx context.Context, sec *tree.Section, load ContentLoader) string { + if sec.ContentRef == "" || load == nil { + return "" + } + s, err := load(ctx, sec.ContentRef) + if err != nil { + return "" + } + return s +} + +// gatherBeneath is a heading's text: the stored text of every section +// beneath it, in document order, each under its own heading. ok is false +// when the heading spans more than unitExpandPages pages or its text would +// exceed unitExpandChars, so a heading over a large part of the document +// is not stood in for by that part's opening pages. +func gatherBeneath(ctx context.Context, sec *tree.Section, load ContentLoader) (string, bool) { + if len(sec.Children) == 0 { + return "", false + } + // A heading over more than unitExpandPages pages is a part or a cover, + // not a heading over one block. + if sec.PageEnd > 0 && sec.PageEnd-sec.PageStart+1 > unitExpandPages { + return "", false + } + estimate := 0 + sec.Walk(func(s *tree.Section) bool { + if s != sec && s.ContentRef != "" { + estimate += s.TokenCount * 4 + } + return estimate <= unitExpandChars + }) + if estimate > unitExpandChars { + return "", false + } + var b strings.Builder + over := false + sec.Walk(func(s *tree.Section) bool { + if s == sec || over { + return !over + } + text := strings.TrimSpace(loadContent(ctx, s, load)) + if text == "" { + return true + } + if b.Len() > 0 { + b.WriteString("\n\n") + } + if title := strings.TrimSpace(s.Title); title != "" && !strings.HasPrefix(text, title) { + b.WriteString(title) + b.WriteString("\n") + } + b.WriteString(text) + if b.Len() > unitExpandChars { + over = true + } + return !over + }) + if over { + return "", false + } + return b.String(), true +} + +// substantive is text with its leading heading and surrounding space taken +// off: what a reader gets beyond the title. +func substantive(text, title string) string { + s := strings.TrimSpace(text) + if title = strings.TrimSpace(title); title != "" { + if strings.HasPrefix(strings.ToLower(s), strings.ToLower(title)) { + s = strings.TrimSpace(s[len(title):]) + } + } + return strings.TrimFunc(s, func(r rune) bool { return unicode.IsSpace(r) || unicode.IsPunct(r) }) +} + +func shingles(text string) []string { + words := strings.FieldsFunc(strings.ToLower(text), func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsDigit(r) + }) + if len(words) < shingleWords { + if len(words) == 0 { + return nil + } + return []string{strings.Join(words, " ")} + } + out := make([]string, 0, len(words)-shingleWords+1) + for i := 0; i+shingleWords <= len(words); i++ { + out = append(out, strings.Join(words[i:i+shingleWords], " ")) + } + return out +} + +func addShingles(seen map[string]bool, text string) { + for _, s := range shingles(text) { + seen[s] = true + } +} + +// repeats reports whether the units already returned carry most of text. +func repeats(seen map[string]bool, text string) bool { + sh := shingles(text) + if len(sh) == 0 || len(seen) == 0 { + return false + } + hit := 0 + for _, s := range sh { + if seen[s] { + hit++ + } + } + return float64(hit) >= duplicateShare*float64(len(sh)) +} diff --git a/pkg/retrieval/units_test.go b/pkg/retrieval/units_test.go new file mode 100644 index 0000000..5f35e08 --- /dev/null +++ b/pkg/retrieval/units_test.go @@ -0,0 +1,185 @@ +package retrieval + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/hallelx2/vectorless-engine/pkg/tree" +) + +// A filing shaped like the FinanceBench ones HAL-3160 measured: a cover +// heading over the whole document, a part heading, a compact heading over +// one short note, and real statement sections. +const ( + incomeText = "Consolidated Statement of Income. Net revenue was $86,392 million in 2022 against $79,474 million in 2021; cost of sales was $40,576 million and operating profit $11,512 million, after restructuring charges of $411 million." + cashText = "Consolidated Statement of Cash Flows. Depreciation and amortization was $2,763 million in 2022; net cash provided by operating activities was $10,811 million and capital spending was $5,207 million for the year." + noteText = "Restructuring. In 2022 the company recorded charges of $411 million under its productivity plan, mainly severance and asset impairments in Europe and Latin America, with most cash payments due within twelve months." + riskText = "Risk factors. Demand for our products depends on consumer spending, which falls in a downturn; commodity and freight costs can rise faster than we can raise prices, and currency moves change reported results." + incomePage = "Page 62 header. " + incomeText + " Earnings per share were $6.42." + runnerUpPage = "Page 63. Consolidated Balance Sheet. Total assets were $92,187 million and total liabilities $74,914 million at year end, with inventories of $5,222 million and accounts payable of $9,201 million." +) + +func filingTree() *tree.Tree { + income := &tree.Section{ID: "sec_income", Title: "Consolidated Statement of Income", ContentRef: "income", TokenCount: len(incomeText) / 4, PageStart: 62, PageEnd: 62} + cash := &tree.Section{ID: "sec_cash", Title: "Consolidated Statement of Cash Flows", ContentRef: "cash", TokenCount: len(cashText) / 4, PageStart: 64, PageEnd: 64} + noteBody := &tree.Section{ID: "sec_note_body", Title: "", ContentRef: "note", TokenCount: len(noteText) / 4, PageStart: 78, PageEnd: 78} + note := &tree.Section{ID: "sec_note", Title: "Note 3. Restructuring", ContentRef: "note_heading", PageStart: 78, PageEnd: 78, Children: []*tree.Section{noteBody}} + risk := &tree.Section{ID: "sec_risk", Title: "Risk Factors", ContentRef: "risk", TokenCount: len(riskText) / 4, PageStart: 9, PageEnd: 12} + part := &tree.Section{ID: "sec_part2", Title: "PART II", PageStart: 30, PageEnd: 90, Children: []*tree.Section{income, cash, note}} + cover := &tree.Section{ID: "sec_cover", Title: "PEPSICO, INC.", ContentRef: "cover", PageStart: 1, PageEnd: 120, Children: []*tree.Section{risk, part}} + return &tree.Tree{DocumentID: "doc_pep", Title: "PepsiCo 10-K", Root: cover} +} + +func filingStore() ContentLoader { + data := map[string]string{ + "income": incomeText, + "cash": cashText, + "note": noteText, + "note_heading": "Note 3. Restructuring", + "risk": riskText, + "cover": "PEPSICO, INC.", + } + return func(ctx context.Context, ref string) (string, error) { + s, ok := data[ref] + if !ok { + return "", errors.New("not found") + } + return s, nil + } +} + +func ids(units []Unit) []string { + var out []string + for _, u := range units { + if u.Kind == UnitPage { + out = append(out, "page") + } else { + out = append(out, string(u.Section.ID)) + } + } + return out +} + +func TestHeadingOnlyTopHitGivesItsSlotToRealEvidence(t *testing.T) { + res := &Result{SelectedIDs: []tree.SectionID{"sec_cover", "sec_part2", "sec_cash", "sec_risk"}} + units, stats := AssembleUnits(context.Background(), filingTree(), res, 2, filingStore()) + if got := strings.Join(ids(units), ","); got != "sec_cash,sec_risk" { + t.Fatalf("units = %s, want the two sections with body text", got) + } + if stats.Empty != 2 { + t.Errorf("empty = %d, want the cover and part headings counted", stats.Empty) + } +} + +// A cover heading spans the whole filing: gathering its text would return +// the document's opening pages, so it is dropped rather than expanded. +func TestCoverTitleIsDroppedNotExpanded(t *testing.T) { + units, stats := AssembleUnits(context.Background(), filingTree(), &Result{SelectedIDs: []tree.SectionID{"sec_cover"}}, 0, filingStore()) + if len(units) != 0 { + t.Fatalf("units = %v, want none for a cover heading", ids(units)) + } + if stats.Expanded != 0 || stats.Empty != 1 { + t.Errorf("stats = %+v", stats) + } +} + +func TestCompactHeadingTakesTheTextBeneathIt(t *testing.T) { + units, stats := AssembleUnits(context.Background(), filingTree(), &Result{SelectedIDs: []tree.SectionID{"sec_note"}}, 5, filingStore()) + if len(units) != 1 || units[0].Section.ID != "sec_note" || !units[0].Expanded { + t.Fatalf("units = %+v, want the note heading expanded", units) + } + if !strings.Contains(units[0].Content, "$411 million") { + t.Errorf("content = %q, want the note's body", units[0].Content) + } + if stats.Expanded != 1 { + t.Errorf("expanded = %d", stats.Expanded) + } +} + +func TestSectionRepeatingAnEvidencePageIsDropped(t *testing.T) { + res := &Result{ + EvidencePages: []EvidencePage{{Page: 62, Text: incomePage, Confidence: 0.9}}, + SelectedIDs: []tree.SectionID{"sec_income", "sec_cash"}, + } + units, stats := AssembleUnits(context.Background(), filingTree(), res, 5, filingStore()) + if got := strings.Join(ids(units), ","); got != "page,sec_cash" { + t.Fatalf("units = %s, want the page and the section it does not carry", got) + } + if stats.Duplicate != 1 { + t.Errorf("duplicate = %d", stats.Duplicate) + } +} + +func TestOpenSlotsFillFromRunnerUpsUnderTheCap(t *testing.T) { + res := &Result{ + EvidencePages: []EvidencePage{{Page: 62, Text: incomePage, Confidence: 0.9}}, + SelectedIDs: []tree.SectionID{"sec_part2", "sec_income"}, + RunnerUpPages: []EvidencePage{{Page: 62, Text: incomePage}, {Page: 63, Text: runnerUpPage, Confidence: 0.3}}, + RunnerUpIDs: []tree.SectionID{"sec_cash", "sec_risk"}, + } + units, stats := AssembleUnits(context.Background(), filingTree(), res, 3, filingStore()) + if got := strings.Join(ids(units), ","); got != "page,page,sec_cash" { + t.Fatalf("units = %s", got) + } + if !units[1].RunnerUp || units[1].Page.Page != 63 || !units[2].RunnerUp { + t.Errorf("runner-ups not marked: %+v", units) + } + if stats.RunnerUps != 2 { + t.Errorf("runner_ups = %d", stats.RunnerUps) + } + // Without a cap nothing is filled in. + all, _ := AssembleUnits(context.Background(), filingTree(), res, 0, filingStore()) + if got := strings.Join(ids(all), ","); got != "page" { + t.Errorf("uncapped units = %s, want only the evidence", got) + } +} + +func TestCapCountsPagesAndSectionsTogether(t *testing.T) { + res := &Result{ + EvidencePages: []EvidencePage{{Page: 62, Text: incomePage}, {Page: 63, Text: runnerUpPage}}, + SelectedIDs: []tree.SectionID{"sec_cash", "sec_risk"}, + } + units, _ := AssembleUnits(context.Background(), filingTree(), res, 3, filingStore()) + if got := strings.Join(ids(units), ","); got != "page,page,sec_cash" { + t.Fatalf("units = %s", got) + } +} + +func TestUnitAssemblerAcceptsOneUnitAtATime(t *testing.T) { + a := NewUnitAssembler(filingTree(), 0, filingStore()) + if _, ok := a.AddSection(context.Background(), "sec_part2", false); ok { + t.Error("part heading accepted") + } + if _, ok := a.AddSection(context.Background(), "sec_income", false); !ok { + t.Error("income statement refused") + } + if _, ok := a.AddSection(context.Background(), "sec_income", false); ok { + t.Error("the same section accepted twice") + } + if _, ok := a.AddPage(EvidencePage{Page: 62, Text: " "}, false); ok { + t.Error("blank page accepted") + } + if n := len(a.Units()); n != 1 { + t.Errorf("units = %d", n) + } +} + +func TestAssembleUnitsNilResult(t *testing.T) { + if units, _ := AssembleUnits(context.Background(), filingTree(), nil, 5, filingStore()); units != nil { + t.Errorf("units = %v", units) + } +} + +func TestRunnerUpsSkipTheEvidenceAndKeepRankOrder(t *testing.T) { + p := func(n int, prob float64) PageScore { return PageScore{Page: NavPage{Number: n, LeafID: "l"}, P: prob} } + nav := &NavResult{ + Pages: []PageScore{p(5, 0.9), p(7, 0.6), p(2, 0.4), p(9, 0.1)}, + Evidence: []PageScore{p(5, 0.9)}, + } + got := runnerUps(nav) + if len(got) != 3 || got[0].Page.Number != 7 || got[1].Page.Number != 2 || got[2].Page.Number != 9 { + t.Fatalf("runner-ups = %+v", got) + } +} From 5ac7ef394cdaf13bf230edce157e6705be408b01 Mon Sep 17 00:00:00 2001 From: hallelx2 Date: Sat, 17 Oct 2026 00:49:01 +0100 Subject: [PATCH 2/2] fix(retrieval): compare a unit's heading without changing its byte length --- pkg/retrieval/units.go | 5 ++++- pkg/retrieval/units_test.go | 11 +++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/pkg/retrieval/units.go b/pkg/retrieval/units.go index a79f6e5..2a83efe 100644 --- a/pkg/retrieval/units.go +++ b/pkg/retrieval/units.go @@ -255,7 +255,10 @@ func gatherBeneath(ctx context.Context, sec *tree.Section, load ContentLoader) ( func substantive(text, title string) string { s := strings.TrimSpace(text) if title = strings.TrimSpace(title); title != "" { - if strings.HasPrefix(strings.ToLower(s), strings.ToLower(title)) { + // EqualFold on the same byte span: lower-casing can change a + // rune's byte length (U+212A KELVIN SIGN is 3 bytes, "k" is 1), so + // slicing s by a lower-cased length could panic. + if len(s) >= len(title) && strings.EqualFold(s[:len(title)], title) { s = strings.TrimSpace(s[len(title):]) } } diff --git a/pkg/retrieval/units_test.go b/pkg/retrieval/units_test.go index 5f35e08..d135e0e 100644 --- a/pkg/retrieval/units_test.go +++ b/pkg/retrieval/units_test.go @@ -183,3 +183,14 @@ func TestRunnerUpsSkipTheEvidenceAndKeepRankOrder(t *testing.T) { t.Fatalf("runner-ups = %+v", got) } } + +// A title whose case folding changes its byte length must not panic the +// request (U+212A KELVIN SIGN lower-cases to the 1-byte "k"). +func TestSubstantiveSurvivesCaseFoldingThatChangesLength(t *testing.T) { + if got := substantive("k", "\u212a"); got != "k" { + t.Errorf("substantive = %q", got) + } + if got := substantive("\u212aelvin scale readings follow", "\u212aelvin"); got != "scale readings follow" { + t.Errorf("substantive = %q", got) + } +}