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..2a83efe --- /dev/null +++ b/pkg/retrieval/units.go @@ -0,0 +1,304 @@ +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 != "" { + // 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):]) + } + } + 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..d135e0e --- /dev/null +++ b/pkg/retrieval/units_test.go @@ -0,0 +1,196 @@ +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) + } +} + +// 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) + } +}