diff --git a/cmd/server/main.go b/cmd/server/main.go index a51c889..a1e7e7e 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -241,6 +241,12 @@ func run() error { // Multi-document query dispatcher. multiDoc := retrieval.NewMultiDoc(strategy, pool.LoadTree) + // 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) + corpusNavigator := retrieval.CorpusNavigator{Judge: judge} + // Pre-built set of selectable strategies, keyed by config name. // Backs the per-request "strategy" override on /v1/query so the // benchmark can A/B chunked-tree vs treewalk against this same @@ -384,6 +390,8 @@ func run() error { Queue: q, Strategy: strategy, MultiDoc: multiDoc, + CorpusLibrary: corpusLibrary, + CorpusNavigator: corpusNavigator, Version: buildVersion(), Config: cfg, Strategies: strategies, diff --git a/internal/handler/corpus.go b/internal/handler/corpus.go new file mode 100644 index 0000000..1b68765 --- /dev/null +++ b/internal/handler/corpus.go @@ -0,0 +1,300 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "strings" + "sync" + "time" + + "golang.org/x/sync/errgroup" + + "github.com/hallelx2/vectorless-engine/pkg/db" + "github.com/hallelx2/vectorless-engine/pkg/retrieval" + "github.com/hallelx2/vectorless-engine/pkg/tree" +) + +// corpusCatalog is the slice of *db.Pool corpus mode needs: list a +// scope's documents, or fetch named ones. +type corpusCatalog interface { + ListDocuments(ctx context.Context, o db.ListDocumentsOpts) ([]db.Document, time.Time, error) + GetDocument(ctx context.Context, id tree.DocumentID, orgID, storeID string) (*db.Document, error) +} + +// CorpusHandler serves corpus mode (HAL-3075): one question over a whole +// collection, routed, read and merged on the Judge, with bridge hops for +// multi-hop questions. /v1/query/corpus returns the merged evidence; +// /v1/answer/corpus also writes one answer over it. +type CorpusHandler struct { + logger *slog.Logger + catalog corpusCatalog + library *retrieval.CorpusLibrary + navigator retrieval.CorpusNavigator + // answers writes the answer and places its citations; nil makes + // /v1/answer/corpus answer 501. + answers *AnswerStoreHandler +} + +// NewCorpusHandler builds the handler. A nil library (no Judge +// configured) makes both routes answer 501. +func NewCorpusHandler(logger *slog.Logger, catalog corpusCatalog, library *retrieval.CorpusLibrary, navigator retrieval.CorpusNavigator, answers *AnswerStoreHandler) *CorpusHandler { + return &CorpusHandler{logger: logger, catalog: catalog, library: library, navigator: navigator, answers: answers} +} + +type corpusRequest struct { + Query string `json:"query"` + // DocumentIDs limits the collection; empty means every ready document + // in the caller's org and store. + DocumentIDs []tree.DocumentID `json:"document_ids"` + MaxDocuments int `json:"max_documents"` + MaxPagesPerDocument int `json:"max_pages_per_document"` + MaxEvidence int `json:"max_evidence"` + MaxHops int `json:"max_hops"` + // Answer-only fields. + Model string `json:"model"` + MaxTokens int `json:"max_tokens"` +} + +const ( + corpusMaxDocumentsCap = 32 + corpusMaxHopsCap = 5 + // The merged evidence goes into one answer prompt: bound it server-side. + corpusMaxPagesPerDocCap = 16 + corpusMaxEvidenceCap = 64 + corpusListPage = 200 + // corpusDeclineBelow is the coverage under which /v1/answer/corpus + // tells the generator the collection may not answer the question. + corpusDeclineBelow = 0.25 +) + +func (h *CorpusHandler) decode(w http.ResponseWriter, r *http.Request) (string, corpusRequest, bool) { + orgID, ok := requireOrgID(w, r) + if !ok { + return "", corpusRequest{}, false + } + var body corpusRequest + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeErr(w, http.StatusBadRequest, "invalid json: "+err.Error()) + return "", body, false + } + body.Query = strings.TrimSpace(body.Query) + if body.Query == "" { + writeErr(w, http.StatusBadRequest, "query is required") + return "", body, false + } + if h.library == nil || h.navigator.Judge == nil { + writeErr(w, http.StatusNotImplemented, "corpus mode needs a Judge (llm.judge) configured") + return "", body, false + } + return orgID, body, true +} + +// run resolves the collection, builds its documents and navigates. +func (h *CorpusHandler) run(ctx context.Context, orgID, store string, body corpusRequest) (*retrieval.CorpusResult, error) { + refs, err := h.refs(ctx, orgID, store, body.DocumentIDs) + if err != nil { + return nil, err + } + docs, buildErrs := h.library.Docs(ctx, orgID, store, refs) + nav := h.navigator + if body.MaxDocuments > 0 { + nav.MaxDocuments = min(body.MaxDocuments, corpusMaxDocumentsCap) + } + if body.MaxPagesPerDocument > 0 { + nav.MaxPagesPerDocument = min(body.MaxPagesPerDocument, corpusMaxPagesPerDocCap) + } + if body.MaxEvidence > 0 { + nav.MaxEvidence = min(body.MaxEvidence, corpusMaxEvidenceCap) + } + if body.MaxHops > 0 { + nav.MaxHops = min(body.MaxHops, corpusMaxHopsCap) + } + res, err := nav.Navigate(ctx, body.Query, docs) + if err != nil { + return nil, err + } + res.Considered = len(refs) + for id, e := range buildErrs { + if res.Errors == nil { + res.Errors = map[tree.DocumentID]string{} + } + res.Errors[id] = e.Error() + } + return res, nil +} + +// refs lists the collection: the named documents, or every ready one in +// scope. Only ready documents are read. +func (h *CorpusHandler) refs(ctx context.Context, orgID, store string, ids []tree.DocumentID) ([]retrieval.CorpusRef, error) { + toRef := func(d db.Document) retrieval.CorpusRef { + return retrieval.CorpusRef{ID: d.ID, Title: d.Title, UpdatedAt: d.UpdatedAt, TOC: d.TOCTree} + } + if len(ids) > 0 { + out := make([]retrieval.CorpusRef, len(ids)) + keep := make([]bool, len(ids)) + var mu sync.Mutex + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(16) + for i, id := range ids { + g.Go(func() error { + d, err := h.catalog.GetDocument(gctx, id, orgID, store) + if errors.Is(err, db.ErrNotFound) { + return nil // not in the caller's scope: not part of the collection + } + if err != nil { + // An outage or a cancelled request is not "this document is + // absent": answering from the rest would look like a real, + // empty answer. + return fmt.Errorf("corpus: look up %s: %w", id, err) + } + if d == nil || d.Status != db.StatusReady { + return nil // still ingesting, or failed: not readable yet + } + mu.Lock() + out[i], keep[i] = toRef(*d), true + mu.Unlock() + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, err + } + refs := make([]retrieval.CorpusRef, 0, len(ids)) + for i := range out { + if keep[i] { + refs = append(refs, out[i]) + } + } + return refs, nil + } + var ( + refs []retrieval.CorpusRef + cursor time.Time + ) + for { + page, next, err := h.catalog.ListDocuments(ctx, db.ListDocumentsOpts{ + OrgID: orgID, StoreID: store, Limit: corpusListPage, Cursor: cursor, Status: db.StatusReady, + }) + if err != nil { + return nil, err + } + for _, d := range page { + refs = append(refs, toRef(d)) + } + if next.IsZero() || len(page) == 0 { + break + } + cursor = next + } + return refs, nil +} + +func corpusResponse(query string, res *retrieval.CorpusResult, started time.Time) map[string]any { + docs := res.Documents + if len(docs) > 25 { + docs = docs[:25] // the full ranking of a large collection is noise past the top + } + return map[string]any{ + "query": query, + "strategy": "corpus", + "evidence": res.Evidence, + "documents": docs, + "documents_considered": res.Considered, + "hops": res.Hops, + "coverage": res.Coverage, + "errors": res.Errors, + "judge_requests": res.Requests, + "usage": usageMap(res.Usage), + "elapsed_ms": time.Since(started).Milliseconds(), + } +} + +// HandleQueryCorpus is POST /v1/query/corpus. +func (h *CorpusHandler) HandleQueryCorpus(w http.ResponseWriter, r *http.Request) { + orgID, body, ok := h.decode(w, r) + if !ok { + return + } + started := time.Now() + res, err := h.run(r.Context(), orgID, storeID(r), body) + if err != nil { + h.logger.Error("query/corpus: failed", "err", err) + writeErr(w, http.StatusBadGateway, "corpus retrieval failed: "+err.Error()) + return + } + resp := corpusResponse(body.Query, res, started) + withTiming(r, resp) + writeJSON(w, http.StatusOK, resp) +} + +// HandleAnswerCorpus is POST /v1/answer/corpus: the evidence of +// /v1/query/corpus, then one generated answer citing document and page, +// each citation placed on its page. +func (h *CorpusHandler) HandleAnswerCorpus(w http.ResponseWriter, r *http.Request) { + orgID, body, ok := h.decode(w, r) + if !ok { + return + } + if h.answers == nil || h.answers.llm == nil { + writeErr(w, http.StatusNotImplemented, "corpus answers require an LLM-backed configuration") + return + } + started := time.Now() + store := storeID(r) + res, err := h.run(r.Context(), orgID, store, body) + if err != nil { + h.logger.Error("answer/corpus: retrieval failed", "err", err) + writeErr(w, http.StatusBadGateway, "corpus retrieval failed: "+err.Error()) + return + } + resp := corpusResponse(body.Query, res, started) + if len(res.Evidence) == 0 { + resp["answer"] = "No document in this collection appears to hold what the question needs." + resp["citations"] = []any{} + resp["elapsed_ms"] = time.Since(started).Milliseconds() + withTiming(r, resp) + writeJSON(w, http.StatusOK, resp) + return + } + sections := make([]relevantSection, 0, len(res.Evidence)) + for i, e := range res.Evidence { + sections = append(sections, relevantSection{ + index: i + 1, docID: e.DocumentID, docTitle: e.DocumentTitle, sectionTitle: e.SectionTitle, + startPage: e.Page, endPage: e.Page, content: e.Text, + }) + } + query := body.Query + if res.Coverage < corpusDeclineBelow { + // The Judge found the evidence incomplete: tell the generator, so a + // question the collection does not answer is declined rather than + // completed from the model's own knowledge. + query += "\n\n(Retrieval judged these sections unlikely to contain everything the question needs. " + + "If they do not answer it, say so and state what is missing.)" + } + model := body.Model + if model == "" { + model = h.answers.defaultModel() + } + answer, cited, synthUsage, err := h.answers.generate(r.Context(), model, query, sections, body.MaxTokens) + if err != nil { + h.logger.Error("answer/corpus: generate failed", "err", err) + writeErr(w, http.StatusInternalServerError, "answer generation failed: "+err.Error()) + return + } + total := res.Usage + total.Add(synthUsage) + citations := buildStoreCitations(sections, cited) + h.answers.annotate(r.Context(), orgID, store, citations) + resp["answer"] = answer + resp["citations"] = citations + resp["usage"] = usageMap(total) + resp["retrieval_usage"] = usageMap(res.Usage) + resp["elapsed_ms"] = time.Since(started).Milliseconds() + setTokenHeaders(w.Header(), resp) + withTiming(r, resp) + writeJSON(w, http.StatusOK, resp) +} diff --git a/internal/handler/router.go b/internal/handler/router.go index 03e5b1c..49d2860 100644 --- a/internal/handler/router.go +++ b/internal/handler/router.go @@ -33,7 +33,13 @@ type Deps struct { Strategy retrieval.Strategy MultiDoc *retrieval.MultiDoc Version string - Config config.Config + + // CorpusLibrary and CorpusNavigator serve corpus mode (HAL-3075): + // /v1/query/corpus and /v1/answer/corpus. A nil library makes both + // answer 501. + CorpusLibrary *retrieval.CorpusLibrary + CorpusNavigator retrieval.CorpusNavigator + Config config.Config // Strategies is the pre-built set of selectable retrieval // strategies keyed by config name. It backs the per-request @@ -187,6 +193,11 @@ func Router(d Deps) http.Handler { answerTreeWalk.pincites = d.Pincites answerStore := NewAnswerStoreHandler(d.Logger, d.DB, d.Storage, d.LLM, d.LLMModel).WithModelFunc(d.LLMModelFunc) answerStore.pincites = d.Pincites + var catalog corpusCatalog + if d.DB != nil { + catalog = d.DB + } + corpus := NewCorpusHandler(d.Logger, catalog, d.CorpusLibrary, d.CorpusNavigator, answerStore) pages := NewPincitesHandler(d.Logger, d.DB, d.Pincites) webhook := NewWebhookHandler(d.Logger, d.Queue) @@ -251,6 +262,8 @@ func Router(d Deps) http.Handler { r.Post("/stream", queryStream.HandleQueryStream) r.Post("/multi", queryMulti.HandleQueryMulti) r.Post("/multi/stream", queryStreamMulti.HandleQueryStreamMulti) + // One question over a whole collection (HAL-3075). + r.Post("/corpus", corpus.HandleQueryCorpus) }) // Answer: retrieval + synthesis in one round-trip. /answer @@ -261,6 +274,9 @@ func Router(d Deps) http.Handler { // Store answer: map-reduce a question across many documents into // ONE synthesised answer with per-document citations. r.With(needsModel).Post("/answer/store", answerStore.HandleAnswerStore) + // Corpus answer: route, read and merge evidence across the + // collection on the Judge, with bridge hops, then one answer. + r.With(needsModel).Post("/answer/corpus", corpus.HandleAnswerCorpus) // Which model providers are in effect, for any authenticated // caller: the dashboard's setup banner reads it. No keys. diff --git a/internal/handler/router_parity_test.go b/internal/handler/router_parity_test.go index 7e9140e..cdb34fb 100644 --- a/internal/handler/router_parity_test.go +++ b/internal/handler/router_parity_test.go @@ -47,6 +47,8 @@ func TestRouterParity(t *testing.T) { "POST /v1/query/", "POST /v1/answer", "POST /v1/answer/treewalk", + "POST /v1/query/corpus", + "POST /v1/answer/corpus", } for _, route := range want { if !got[route] { diff --git a/pkg/retrieval/corpus.go b/pkg/retrieval/corpus.go new file mode 100644 index 0000000..696f69f --- /dev/null +++ b/pkg/retrieval/corpus.go @@ -0,0 +1,681 @@ +package retrieval + +import ( + "context" + "fmt" + "sort" + "strings" + "sync" + + "github.com/hallelx2/llmgate" + + "github.com/hallelx2/vectorless-engine/pkg/tree" +) + +// Corpus mode: one question over a whole collection (HAL-3075). +// +// JudgeNavigator answers a question about one document. A collection +// question ("did the two outlets agree on the figure?", "which company +// named in both reports...") needs facts from several documents, and +// nobody says which. Corpus mode keeps every decision a closed-set +// judgement on the Judge, and needs no index and no embedding: +// +// 1. Route. Every document is shown to the Judge as a short card — +// its title, its top-level contents, how it opens — and the Judge +// gives the probability that it holds any part of what the question +// needs. +// 2. Read. The best documents are navigated in parallel with the +// ordinary JudgeNavigator, in multipart wording: a page counts when +// it states any of the needed facts. +// 3. Merge. The Judge's probabilities are calibrated against the same +// question and the same criteria, so pages from different documents +// rank on one scale. They are merged into one list, a few pages per +// document at most, every contributing document represented. +// 4. Hop. The Judge is asked whether what was found contains every fact +// the question needs. When it does not, the unread documents are +// routed again with what was found in the state — would this one +// supply what is still missing? — and the best are read. This is the +// bridge step of a multi-hop question, taken by selection: no +// sub-question is written. +// 5. Coverage. The final completeness probability is returned, so the +// caller can decline a question the collection does not answer. + +// CorpusDoc is one document a corpus question may read. +type CorpusDoc struct { + ID tree.DocumentID + Title string + // Card is what routing sees: title, top-level contents, opening text. + Card string + // Inputs are the document's leaves and page loader. + Inputs *NavInputs +} + +// CorpusEvidence is one page in the merged evidence list. +type CorpusEvidence struct { + DocumentID tree.DocumentID `json:"document_id"` + DocumentTitle string `json:"document_title,omitempty"` + Page int `json:"page"` + SectionTitle string `json:"section_title,omitempty"` + Text string `json:"text"` + Confidence float64 `json:"confidence"` + Hop int `json:"hop"` + // Ordinal marks a document without pages: Page is then the section's + // position in the document, and SectionTitle names it. + Ordinal bool `json:"ordinal,omitempty"` +} + +// CorpusDocScore is a document's routing score: the probability that it +// holds part of what the question needs, from the hop that scored it last. +type CorpusDocScore struct { + DocumentID tree.DocumentID `json:"document_id"` + Title string `json:"title,omitempty"` + Score float64 `json:"score"` + Hop int `json:"hop"` + Read bool `json:"read"` +} + +// CorpusHop is one round of route, read, merge and coverage. +type CorpusHop struct { + Hop int `json:"hop"` + // Routed is how many documents the Judge scored this hop. + Routed int `json:"routed"` + // Read are the documents navigated this hop, best route score first. + Read []tree.DocumentID `json:"read"` + // NewEvidence is how many pages this hop added to the merged list. + NewEvidence int `json:"new_evidence"` + // Coverage is the Judge's probability, after this hop, that the merged + // evidence holds every fact the question needs. + Coverage float64 `json:"coverage"` +} + +// CorpusResult is what a corpus question found. +type CorpusResult struct { + Evidence []CorpusEvidence `json:"evidence"` + Documents []CorpusDocScore `json:"documents"` + Hops []CorpusHop `json:"hops"` + // Coverage is the final completeness probability. + Coverage float64 `json:"coverage"` + // Considered is how many documents the question was asked over. + Considered int `json:"documents_considered"` + // Errors are documents that could not be read, by ID; the question is + // still answered from the rest. + Errors map[tree.DocumentID]string `json:"errors,omitempty"` + Usage Usage `json:"usage"` + Requests int `json:"judge_requests"` +} + +// CorpusStep reports one stage as it completes, for display. +type CorpusStep struct { + // Kind is "route", "read", "merge" or "coverage". + Kind string `json:"kind"` + Hop int `json:"hop"` + Documents []CorpusDocScore `json:"documents,omitempty"` + Evidence int `json:"evidence,omitempty"` + Coverage float64 `json:"coverage,omitempty"` + Read []tree.DocumentID `json:"read,omitempty"` +} + +// CorpusNavigator answers one question over many documents. +type CorpusNavigator struct { + Judge llmgate.Judge + + // Navigator is the template each document is read with. Multipart is + // forced on; SkimAll is set per document when its pages are in memory. + Navigator JudgeNavigator + + // OnStep, when set, is told about each stage as it completes. + OnStep func(CorpusStep) + + // RouteThreshold is the routing probability at or above which a + // document is read. Zero selects 0.5. + RouteThreshold float64 + // MinDocuments are read on the first hop even below the threshold: + // a low-confidence best guess beats reading nothing. Zero selects 3. + MinDocuments int + // MaxDocuments bounds the documents read per hop. Zero selects 8. + MaxDocuments int + // MaxPagesPerDocument bounds one document's pages in the merged list. + // Zero selects 4. + MaxPagesPerDocument int + // MaxEvidence bounds the merged list. Zero selects 16. + MaxEvidence int + // EvidenceThreshold is the page probability at or above which a page + // joins the merged list. Zero selects the navigator's threshold (0.5). + EvidenceThreshold float64 + // MaxHops bounds the rounds. Zero selects 3; 1 disables the bridge. + MaxHops int + // CoverageThreshold is the completeness probability at or above which + // no further hop is taken. Zero selects 0.5. + CoverageThreshold float64 + // CardChars bounds one card. Zero selects 900. + CardChars int + // FoundChars bounds the found-so-far state of the coverage and bridge + // judgements. Zero selects 6000. + FoundChars int + // ReadConcurrency bounds documents navigated at once. Zero selects 8. + ReadConcurrency int +} + +const ( + defaultCorpusRouteThreshold = 0.5 + defaultCorpusMinDocs = 3 + defaultCorpusMaxDocs = 8 + defaultCorpusPagesPerDoc = 4 + defaultCorpusMaxEvidence = 16 + defaultCorpusMaxHops = 3 + defaultCorpusCoverage = 0.5 + defaultCorpusCardChars = 900 + defaultCorpusFoundChars = 6000 + defaultCorpusReadPar = 8 + corpusFoundPageChars = 900 +) + +func corpusOr[T int | float64](v, def T) T { + if v > 0 { + return v + } + return def +} + +func (c *CorpusNavigator) step(s CorpusStep) { + if c.OnStep != nil { + c.OnStep(s) + } +} + +// Navigate runs the corpus question. +func (c *CorpusNavigator) Navigate(ctx context.Context, query string, docs []CorpusDoc) (*CorpusResult, error) { + if c.Judge == nil { + return nil, fmt.Errorf("corpus: no Judge configured") + } + res := &CorpusResult{ + Considered: len(docs), Errors: map[tree.DocumentID]string{}, + Evidence: []CorpusEvidence{}, Documents: []CorpusDocScore{}, Hops: []CorpusHop{}, + } + if len(docs) == 0 { + return res, nil + } + byID := make(map[tree.DocumentID]*CorpusDoc, len(docs)) + for i := range docs { + byID[docs[i].ID] = &docs[i] + } + scores := map[tree.DocumentID]*CorpusDocScore{} + read := map[tree.DocumentID]bool{} + var pool []CorpusEvidence // every page any read document judged, before the merge + + maxHops := corpusOr(c.MaxHops, defaultCorpusMaxHops) + for hop := 1; hop <= maxHops; hop++ { + var unread []*CorpusDoc + for i := range docs { + if !read[docs[i].ID] { + unread = append(unread, &docs[i]) + } + } + if len(unread) == 0 { + break + } + found := "" + if hop > 1 { + found = c.foundState(res.Evidence) + } + routed, u, n, err := c.route(ctx, query, found, unread) + if err != nil { + return nil, fmt.Errorf("corpus: route (hop %d): %w", hop, err) + } + res.Usage.Add(u) + res.Requests += n + for _, r := range routed { + r.Hop = hop + cp := r + scores[r.DocumentID] = &cp + } + chosen := c.choose(routed, hop) + c.step(CorpusStep{Kind: "route", Hop: hop, Documents: topScores(routed, 8)}) + if len(chosen) == 0 { + // A bridge hop where nothing unread looks useful: what was found + // is all the collection has. + res.Hops = append(res.Hops, CorpusHop{Hop: hop, Routed: len(routed), Read: []tree.DocumentID{}, Coverage: res.Coverage}) + break + } + for _, id := range chosen { + read[id] = true + scores[id].Read = true + } + c.step(CorpusStep{Kind: "read", Hop: hop, Read: chosen}) + pages, ru, rn, errs := c.read(ctx, query, chosen, byID, hop) + res.Usage.Add(ru) + res.Requests += rn + for id, e := range errs { + res.Errors[id] = e.Error() + } + pool = append(pool, pages...) + before := evidenceKeys(res.Evidence) + res.Evidence = c.merge(pool) + added := 0 + for _, e := range res.Evidence { + if !before[evidenceKey(e)] { + added++ + } + } + c.step(CorpusStep{Kind: "merge", Hop: hop, Evidence: len(res.Evidence)}) + + cov, cu, cn, err := c.coverage(ctx, query, res.Evidence) + if err != nil { + return nil, fmt.Errorf("corpus: coverage (hop %d): %w", hop, err) + } + res.Usage.Add(cu) + res.Requests += cn + res.Coverage = cov + res.Hops = append(res.Hops, CorpusHop{Hop: hop, Routed: len(routed), Read: chosen, NewEvidence: added, Coverage: cov}) + c.step(CorpusStep{Kind: "coverage", Hop: hop, Coverage: cov}) + if cov >= corpusOr(c.CoverageThreshold, defaultCorpusCoverage) || (added == 0 && hop > 1) { + break + } + } + + for _, s := range scores { + res.Documents = append(res.Documents, *s) + } + sort.SliceStable(res.Documents, func(i, j int) bool { + if res.Documents[i].Score != res.Documents[j].Score { + return res.Documents[i].Score > res.Documents[j].Score + } + return res.Documents[i].DocumentID < res.Documents[j].DocumentID + }) + if len(res.Errors) == 0 { + res.Errors = nil + } + return res, nil +} + +// choose picks the documents to read this hop from the routing scores, +// best first: those at or above the threshold, at most MaxDocuments, and +// on the first hop at least MinDocuments. +func (c *CorpusNavigator) choose(routed []CorpusDocScore, hop int) []tree.DocumentID { + th := corpusOr(c.RouteThreshold, defaultCorpusRouteThreshold) + maxDocs := corpusOr(c.MaxDocuments, defaultCorpusMaxDocs) + minDocs := 0 + if hop == 1 { + minDocs = min(corpusOr(c.MinDocuments, defaultCorpusMinDocs), maxDocs) + } + var out []tree.DocumentID + for i, r := range routed { + if len(out) >= maxDocs { + break + } + if r.Score >= th || i < minDocs { + out = append(out, r.DocumentID) + } + } + return out +} + +// route scores every card against the question, and on a bridge hop +// against what was found so far. Cards are packed under the request +// budget and the batches sent together. +func (c *CorpusNavigator) route(ctx context.Context, query, found string, docs []*CorpusDoc) ([]CorpusDocScore, Usage, int, error) { + var usage Usage + out := make([]CorpusDocScore, len(docs)) + for i, d := range docs { + out[i] = CorpusDocScore{DocumentID: d.ID, Title: d.Title} + } + cardChars := corpusOr(c.CardChars, defaultCorpusCardChars) + budget := c.Navigator.reqTokens() + type batch struct { + state map[string]any + questions map[string]llmgate.Question + } + var batches []batch + for start := 0; start < len(docs); { + state := map[string]any{"question": query} + if found != "" { + state["found"] = found + } + questions := map[string]llmgate.Question{} + used := countTokens(query) + countTokens(found) + 100 + end := start + for end < len(docs) && end-start < navLeafBatch { + card := docs[end].Card + if len(card) > cardChars { + card = card[:cardChars] + } + cost := countTokens(card) + 30 + if end > start && used+cost > budget { + break + } + qk := fmt.Sprintf("d_%d", end) + state[qk] = map[string]any{"title": docs[end].Title, "card": card} + questions[qk] = routeQuestion(qk, found != "") + used += cost + end++ + } + batches = append(batches, batch{state, questions}) + start = end + } + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + var ( + wg sync.WaitGroup + mu sync.Mutex + firstErr error + requests int + ) + for _, b := range batches { + wg.Add(1) + go func(b batch) { + defer wg.Done() + jr, err := c.Judge.Judge(ctx, llmgate.JudgeRequest{State: b.state, Questions: b.questions}) + mu.Lock() + defer mu.Unlock() + if err != nil { + if firstErr == nil { + firstErr = err + cancel() + } + return + } + requests++ + usage.Add(judgeUsage(jr)) + for qk := range b.questions { + p, err := jr.Noul(qk) + if err != nil { + continue + } + var i int + _, _ = fmt.Sscanf(qk, "d_%d", &i) + out[i].Score = p + } + }(b) + } + wg.Wait() + if firstErr != nil { + return nil, usage, requests, firstErr + } + sort.SliceStable(out, func(i, j int) bool { + if out[i].Score != out[j].Score { + return out[i].Score > out[j].Score + } + return out[i].DocumentID < out[j].DocumentID + }) + return out, usage, requests, nil +} + +func routeQuestion(qk string, bridge bool) llmgate.Noul { + if bridge { + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is to be answered from a collection of documents. `found` is what has "+ + "already been retrieved from some of them. `%s` is another document in the "+ + "collection: its title and a card showing its contents and how it opens. Would "+ + "this document supply information the question still needs that `found` does not "+ + "already contain — the other entity, event, figure or source the question asks "+ + "about?", qk), + Criteria: &llmgate.NoulCriteria{ + True: "This document would add something the question needs and found lacks", + False: "This document adds nothing the question still needs", + }, + } + } + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is to be answered from a collection of documents; its answer may draw on "+ + "several of them. `%s` is one document in the collection: its title and a card "+ + "showing its contents and how it opens. Would this document hold any part of the "+ + "information the question needs — one of the entities, events, figures, sources or "+ + "claims it asks about?", qk), + Criteria: &llmgate.NoulCriteria{ + True: "This document holds at least part of what the question needs", + False: "Nothing the question needs is in this document", + }, + } +} + +// read navigates the chosen documents in parallel and returns every page +// each judged, tagged with its document and hop. A document that fails is +// reported and skipped. +func (c *CorpusNavigator) read(ctx context.Context, query string, ids []tree.DocumentID, byID map[tree.DocumentID]*CorpusDoc, hop int) ([]CorpusEvidence, Usage, int, map[tree.DocumentID]error) { + var ( + mu sync.Mutex + wg sync.WaitGroup + usage Usage + requests int + out []CorpusEvidence + errs = map[tree.DocumentID]error{} + ) + sem := make(chan struct{}, corpusOr(c.ReadConcurrency, defaultCorpusReadPar)) + results := make([][]CorpusEvidence, len(ids)) + for i, id := range ids { + d := byID[id] + if d == nil || d.Inputs == nil || len(d.Inputs.Leaves) == 0 { + continue + } + wg.Add(1) + go func(i int, d *CorpusDoc) { + defer wg.Done() + select { + case sem <- struct{}{}: + defer func() { <-sem }() + case <-ctx.Done(): + return + } + nav := c.Navigator + nav.Judge = c.Judge + nav.Multipart = true + nav.OnStep = nil + nav.SkimAll = d.Inputs.Pages != nil + r, err := nav.Navigate(ctx, query, d.Inputs.Leaves, d.Inputs.Load) + mu.Lock() + defer mu.Unlock() + if err != nil { + errs[d.ID] = err + return + } + usage.Add(r.Usage) + requests += r.Requests + title := map[string]string{} + for _, l := range d.Inputs.Leaves { + title[l.ID] = l.Title + } + var ev []CorpusEvidence + for _, p := range r.Evidence { + ev = append(ev, CorpusEvidence{ + DocumentID: d.ID, DocumentTitle: d.Title, Page: p.Page.Number, + SectionTitle: title[p.Page.LeafID], Text: p.Page.Text, Confidence: p.P, Hop: hop, + Ordinal: d.Inputs.Ordinal, + }) + } + results[i] = ev + }(i, d) + } + wg.Wait() + for _, ev := range results { // in route order, so equal scores stay deterministic + out = append(out, ev...) + } + return out, usage, requests, errs +} + +// merge builds the evidence list from every judged page: pages at or +// above the evidence threshold, best first, at most MaxPagesPerDocument +// from one document and MaxEvidence in all. Every document with a page +// over the threshold contributes its best page before any document +// contributes a second, so a strong document cannot crowd out the other +// half of a comparison. When no page clears the threshold the best +// pages are kept, flagged by their probability, as one document's +// navigation does. +func (c *CorpusNavigator) merge(pool []CorpusEvidence) []CorpusEvidence { + th := c.EvidenceThreshold + if th <= 0 { + th = c.Navigator.threshold() + } + perDoc := corpusOr(c.MaxPagesPerDocument, defaultCorpusPagesPerDoc) + limit := corpusOr(c.MaxEvidence, defaultCorpusMaxEvidence) + + // The same page judged twice (read on two hops) keeps its best score. + best := map[string]CorpusEvidence{} + var keys []string + for _, e := range pool { + k := evidenceKey(e) + if old, ok := best[k]; !ok || e.Confidence > old.Confidence { + if !ok { + keys = append(keys, k) + } + best[k] = e + } + } + all := make([]CorpusEvidence, 0, len(keys)) + for _, k := range keys { + all = append(all, best[k]) + } + sort.SliceStable(all, func(i, j int) bool { return all[i].Confidence > all[j].Confidence }) + + var above []CorpusEvidence + for _, e := range all { + if e.Confidence >= th { + above = append(above, e) + } + } + if len(above) == 0 { + return append([]CorpusEvidence{}, all[:min(navMinEvidencePages, len(all))]...) + } + + count := map[tree.DocumentID]int{} + taken := map[string]bool{} + out := []CorpusEvidence{} + // First pass: each document's best page. + for _, e := range above { + if len(out) >= limit { + break + } + if count[e.DocumentID] == 0 { + count[e.DocumentID]++ + taken[evidenceKey(e)] = true + out = append(out, e) + } + } + // Second pass: fill by probability under the per-document cap. + for _, e := range above { + if len(out) >= limit { + break + } + if taken[evidenceKey(e)] || count[e.DocumentID] >= perDoc { + continue + } + count[e.DocumentID]++ + taken[evidenceKey(e)] = true + out = append(out, e) + } + sort.SliceStable(out, func(i, j int) bool { return out[i].Confidence > out[j].Confidence }) + return out +} + +// coverage asks whether the merged evidence holds every fact the question +// needs. +func (c *CorpusNavigator) coverage(ctx context.Context, query string, ev []CorpusEvidence) (float64, Usage, int, error) { + if len(ev) == 0 { + return 0, Usage{}, 0, nil + } + state := map[string]any{"question": query, "found": c.foundState(ev)} + q := map[string]llmgate.Question{ + "complete": llmgate.Noul{ + Instructions: "`question` is to be answered from a collection of documents. `found` lists " + + "the passages retrieved for it, each with its document and page. Do these passages, " + + "taken together, state every fact the question needs — each entity, event, figure, " + + "date or source it asks about? Judge only what the passages state, not what could be " + + "guessed. A question about something the passages never mention is not covered.", + Criteria: &llmgate.NoulCriteria{ + True: "Every fact the question needs is stated in the passages", + False: "At least one fact the question needs is missing from the passages", + }, + }, + } + jr, err := c.Judge.Judge(ctx, llmgate.JudgeRequest{State: state, Questions: q}) + if err != nil { + return 0, Usage{}, 0, err + } + p, err := jr.Noul("complete") + if err != nil { + return 0, judgeUsage(jr), 1, nil + } + return p, judgeUsage(jr), 1, nil +} + +// foundState renders the merged evidence for a judgement's state, best +// first, each page trimmed, the whole bounded by FoundChars. +func (c *CorpusNavigator) foundState(ev []CorpusEvidence) string { + limit := corpusOr(c.FoundChars, defaultCorpusFoundChars) + var b strings.Builder + for _, e := range ev { + text := strings.Join(strings.Fields(e.Text), " ") + if len(text) > corpusFoundPageChars { + text = text[:corpusFoundPageChars] + } + title := e.DocumentTitle + if title == "" { + title = string(e.DocumentID) + } + entry := fmt.Sprintf("[%s, p.%d] %s\n", title, e.Page, text) + if b.Len()+len(entry) > limit { + break + } + b.WriteString(entry) + } + return b.String() +} + +func evidenceKey(e CorpusEvidence) string { return fmt.Sprintf("%s#%d", e.DocumentID, e.Page) } + +func evidenceKeys(ev []CorpusEvidence) map[string]bool { + out := make(map[string]bool, len(ev)) + for _, e := range ev { + out[evidenceKey(e)] = true + } + return out +} + +func topScores(s []CorpusDocScore, k int) []CorpusDocScore { + return append([]CorpusDocScore(nil), s[:min(k, len(s))]...) +} + +// CorpusCard renders a document's routing card: its title, its top-level +// contents, and how it opens. +func CorpusCard(title string, in *NavInputs, opening string, chars int) string { + if chars <= 0 { + chars = defaultCorpusCardChars + } + var b strings.Builder + b.WriteString(strings.TrimSpace(title)) + if in != nil && len(in.Leaves) > 1 { + var tops []string + seen := map[string]bool{} + for _, l := range in.Leaves { + top := l.Path + if i := strings.Index(top, " > "); i >= 0 { + top = top[:i] + } + top = strings.TrimSpace(top) + if top == "" || seen[top] { + continue + } + seen[top] = true + tops = append(tops, top) + } + if len(tops) > 0 { + contents := strings.Join(tops, "; ") + if len(contents) > chars/2 { + contents = contents[:chars/2] + } + b.WriteString("\nContents: ") + b.WriteString(contents) + } + } + if o := strings.Join(strings.Fields(opening), " "); o != "" { + b.WriteString("\nOpens: ") + b.WriteString(o) + } + s := b.String() + if len(s) > chars { + s = s[:chars] + } + return s +} diff --git a/pkg/retrieval/corpus_library.go b/pkg/retrieval/corpus_library.go new file mode 100644 index 0000000..f104c40 --- /dev/null +++ b/pkg/retrieval/corpus_library.go @@ -0,0 +1,233 @@ +package retrieval + +import ( + "container/list" + "context" + "errors" + "fmt" + "sync" + "time" + + "golang.org/x/sync/errgroup" + + "github.com/hallelx2/vectorless-engine/pkg/storage" + "github.com/hallelx2/vectorless-engine/pkg/tree" +) + +// ErrNoPages is what a PageStore may return for a document that has no +// persisted pages, as an alternative to storage.ErrNotFound. +var ErrNoPages = errors.New("retrieval: document has no persisted pages") + +// CorpusRef is a document a corpus question may read, as the caller's +// org- and store-scoped listing returned it. The library only builds +// documents it is handed, so the caller's scoping is the access control. +type CorpusRef struct { + ID tree.DocumentID + Title string + UpdatedAt time.Time + // TOC is the persisted contents tree, when ingest built one. + TOC []byte +} + +// CorpusLibrary turns document references into what corpus navigation +// reads — leaves, a page loader and a routing card — and keeps them in a +// byte-bounded cache keyed by document and its last update, so a +// collection's second question does not reload six hundred documents. +type CorpusLibrary struct { + // Walk supplies the page sources: Pages for persisted per-page text, + // PageLoader for section bodies. + Walk *JudgeWalkStrategy + // LoadTree reads a document's section tree, for documents without + // persisted pages. + LoadTree TreeLoader + // CardChars bounds one routing card. Zero selects 900. + CardChars int + // BuildConcurrency bounds documents built at once. Zero selects 16. + BuildConcurrency int + // MaxBytes bounds the cache, counted as page text held. Zero selects + // 512 MiB; a negative value disables caching. + MaxBytes int64 + + mu sync.Mutex + lru *list.List + items map[tree.DocumentID]*list.Element + bytes int64 +} + +type corpusEntry struct { + id tree.DocumentID + updated time.Time + doc CorpusDoc + size int64 +} + +const defaultCorpusCacheBytes = 512 << 20 + +// NewCorpusLibrary builds a library over a judgewalk strategy's page +// sources and a tree loader. +func NewCorpusLibrary(walk *JudgeWalkStrategy, loadTree TreeLoader) *CorpusLibrary { + return &CorpusLibrary{Walk: walk, LoadTree: loadTree} +} + +// Docs builds every referenced document, in the order given. A document +// that cannot be built is reported by ID and left out. +func (l *CorpusLibrary) Docs(ctx context.Context, orgID, storeID string, refs []CorpusRef) ([]CorpusDoc, map[tree.DocumentID]error) { + out := make([]*CorpusDoc, len(refs)) + errs := map[tree.DocumentID]error{} + var mu sync.Mutex + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(corpusOr(l.BuildConcurrency, 16)) + for i, ref := range refs { + g.Go(func() error { + if d, ok := l.cached(ref); ok { + out[i] = &d + return nil + } + d, size, err := l.build(gctx, orgID, storeID, ref) + if err != nil { + mu.Lock() + errs[ref.ID] = err + mu.Unlock() + return nil + } + l.store(ref, d, size) + out[i] = &d + return nil + }) + } + _ = g.Wait() + docs := make([]CorpusDoc, 0, len(refs)) + for _, d := range out { + if d != nil { + docs = append(docs, *d) + } + } + return docs, errs +} + +// build reads one document's inputs: the persisted contents over the +// persisted pages when both exist, the pages alone as one leaf when only +// they exist, and the section tree otherwise. +func (l *CorpusLibrary) build(ctx context.Context, orgID, storeID string, ref CorpusRef) (CorpusDoc, int64, error) { + var ( + in *NavInputs + pages []NavPage + ) + if l.Walk != nil && l.Walk.Pages != nil { + ps, err := l.Walk.Pages.LoadPages(ctx, ref.ID) + switch { + case err == nil: + pages = ps + case errors.Is(err, storage.ErrNotFound) || errors.Is(err, ErrNoPages): + // No persisted pages: an unpaged document, read from its tree. + default: + // A storage failure is not "no pages": falling back would read + // a paged document through its weaker section tree, silently. + return CorpusDoc{}, 0, fmt.Errorf("corpus: load pages %s: %w", ref.ID, err) + } + } + if len(ref.TOC) > 0 && len(pages) > 0 { + in = PersistedNavInputs(ref.TOC, pages) + } + if in == nil && len(pages) > 0 { + in = PagedNavInputs(ref.Title, pages) + } + if in == nil { + if l.LoadTree == nil || l.Walk == nil { + return CorpusDoc{}, 0, fmt.Errorf("corpus: %s has no pages and no tree loader", ref.ID) + } + t, err := l.LoadTree(ctx, ref.ID, orgID, storeID) + if err != nil { + return CorpusDoc{}, 0, fmt.Errorf("corpus: load tree %s: %w", ref.ID, err) + } + in = l.Walk.SectionTreeNavInputs(t) + if in == nil { + return CorpusDoc{}, 0, fmt.Errorf("corpus: %s has no navigable sections", ref.ID) + } + } + + var size int64 + opening := "" + if in.Pages != nil { + for _, p := range in.Pages { + size += int64(len(p.Text)) + } + opening = in.Pages[0].Text + } else if ps, err := in.Load(ctx, in.Leaves[0]); err == nil && len(ps) > 0 { + opening = ps[0].Text + } + cardChars := corpusOr(l.CardChars, defaultCorpusCardChars) + if len(opening) > cardChars { + opening = opening[:cardChars] + } + card := CorpusCard(ref.Title, in, opening, cardChars) + size += int64(len(card)) + return CorpusDoc{ID: ref.ID, Title: ref.Title, Card: card, Inputs: in}, size, nil +} + +func (l *CorpusLibrary) limit() int64 { + if l.MaxBytes == 0 { + return defaultCorpusCacheBytes + } + return l.MaxBytes +} + +func (l *CorpusLibrary) cached(ref CorpusRef) (CorpusDoc, bool) { + if l.limit() < 0 { + return CorpusDoc{}, false + } + l.mu.Lock() + defer l.mu.Unlock() + if l.items == nil { + return CorpusDoc{}, false + } + el, ok := l.items[ref.ID] + if !ok { + return CorpusDoc{}, false + } + e := el.Value.(*corpusEntry) + if !e.updated.Equal(ref.UpdatedAt) { // re-ingested since: rebuild + l.lru.Remove(el) + delete(l.items, ref.ID) + l.bytes -= e.size + return CorpusDoc{}, false + } + l.lru.MoveToFront(el) + return e.doc, true +} + +func (l *CorpusLibrary) store(ref CorpusRef, d CorpusDoc, size int64) { + limit := l.limit() + if limit < 0 || size > limit { + return + } + l.mu.Lock() + defer l.mu.Unlock() + if l.items == nil { + l.items = map[tree.DocumentID]*list.Element{} + l.lru = list.New() + } + if el, ok := l.items[ref.ID]; ok { + l.bytes -= el.Value.(*corpusEntry).size + l.lru.Remove(el) + } + l.items[ref.ID] = l.lru.PushFront(&corpusEntry{id: ref.ID, updated: ref.UpdatedAt, doc: d, size: size}) + l.bytes += size + for l.bytes > limit && l.lru.Len() > 0 { + back := l.lru.Back() + e := back.Value.(*corpusEntry) + l.lru.Remove(back) + delete(l.items, e.id) + l.bytes -= e.size + } +} + +// CachedDocuments reports how many documents and bytes the cache holds. +func (l *CorpusLibrary) CachedDocuments() (int, int64) { + l.mu.Lock() + defer l.mu.Unlock() + if l.lru == nil { + return 0, 0 + } + return l.lru.Len(), l.bytes +} diff --git a/pkg/retrieval/corpus_test.go b/pkg/retrieval/corpus_test.go new file mode 100644 index 0000000..212f628 --- /dev/null +++ b/pkg/retrieval/corpus_test.go @@ -0,0 +1,454 @@ +package retrieval + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/hallelx2/llmgate" + + "github.com/hallelx2/vectorless-engine/pkg/storage" + "github.com/hallelx2/vectorless-engine/pkg/tree" +) + +// corpusJudge scripts a collection where one question needs two hops: +// two articles report Acme's figure, one of them names Beta Corp, and only +// a third article, about Beta Corp, holds the rest. Cards and pages that +// mention Acme score high; a card about Beta scores high only once what +// was found mentions Beta; the evidence is complete once it holds the +// Beta fact. +func corpusJudge(t *testing.T) (*llmgate.MockJudge, *corpusLog) { + log := &corpusLog{} + j := &llmgate.MockJudge{Respond: func(_ context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + st := req.State.(map[string]any) + found, _ := st["found"].(string) + ans := map[string]llmgate.Answer{} + for id, q := range req.Questions { + log.add(id, fmt.Sprint(q.(llmgate.Noul).Instructions)) + p := 0.1 + switch { + case id == "complete": + if strings.Contains(found, "Beta-fact") { + p = 0.9 + } else { + p = 0.2 + } + case strings.HasPrefix(id, "d_"): + card := st[id].(map[string]any)["card"].(string) + if found == "" && strings.Contains(card, "Acme") { + p = 0.9 + } + if found != "" && strings.Contains(card, "Beta") && strings.Contains(found, "Beta Corp") { + p = 0.9 + } + case strings.HasPrefix(id, "l_"): + p = 0.9 + case strings.HasPrefix(id, "p_"): + text := st[id].(map[string]any)["text"].(string) + if strings.Contains(text, "Acme") || strings.Contains(text, "Beta-fact") { + p = 0.95 + } + } + ans[id] = llmgate.NoulAnswer{Noul: p} + } + return &llmgate.Judgment{Model: "mock", Answers: ans, Usage: llmgate.Usage{InputTokens: 100, TotalTokens: 100, TokensReported: true}}, nil + }} + return j, log +} + +type corpusLog struct { + mu sync.Mutex + inst map[string][]string // question id prefix -> instructions seen +} + +func (l *corpusLog) add(id, inst string) { + l.mu.Lock() + defer l.mu.Unlock() + if l.inst == nil { + l.inst = map[string][]string{} + } + k := id + if i := strings.IndexByte(id, '_'); i > 0 { + k = id[:i] + } + l.inst[k] = append(l.inst[k], inst) +} + +func article(id, title, text string) CorpusDoc { + in := PagedNavInputs(title, []NavPage{{Number: 1, Text: text}}) + return CorpusDoc{ID: tree.DocumentID(id), Title: title, Card: CorpusCard(title, in, text, 0), Inputs: in} +} + +func newsCorpus() []CorpusDoc { + return []CorpusDoc{ + article("a", "TechNews: Acme results", "Acme reported revenue of 10 billion. Acme partnered with Beta Corp on chips."), + article("b", "BizWire: Acme quarter", "Acme revenue came in at 10 billion, BizWire reports."), + article("c", "Beta Corp profile", "Beta Corp is led by Dana Ruiz. Beta-fact: founded 1999."), + article("d", "Sports roundup", "The home team won on Sunday."), + article("e", "Weather", "Rain is expected over the weekend."), + article("f", "Film review", "The sequel is longer than the original."), + } +} + +func TestCorpusBridgesToTheMissingDocument(t *testing.T) { + j, log := corpusJudge(t) + nav := CorpusNavigator{Judge: j, MinDocuments: 2} + var steps []string + nav.OnStep = func(s CorpusStep) { steps = append(steps, s.Kind) } + res, err := nav.Navigate(context.Background(), "Which company did Acme partner with, and when was it founded?", newsCorpus()) + if err != nil { + t.Fatal(err) + } + if len(res.Hops) != 2 { + t.Fatalf("want two hops (read Acme articles, then bridge to Beta), got %+v", res.Hops) + } + if got := res.Hops[0].Read; len(got) != 2 || got[0] != "a" || got[1] != "b" { + t.Fatalf("hop 1 should read the two Acme articles in route order, got %v", got) + } + if got := res.Hops[1].Read; len(got) != 1 || got[0] != "c" { + t.Fatalf("hop 2 should read only the Beta article, got %v", got) + } + if res.Coverage < 0.5 { + t.Fatalf("coverage after the bridge should be high, got %v", res.Coverage) + } + var bridged bool + for _, e := range res.Evidence { + if e.DocumentID == "c" && e.Hop == 2 { + bridged = true + } + } + if !bridged { + t.Fatalf("evidence should carry the Beta page from hop 2: %+v", res.Evidence) + } + if res.Considered != 6 || res.Usage.LLMCalls == 0 || res.Requests == 0 { + t.Fatalf("accounting missing: considered=%d usage=%+v requests=%d", res.Considered, res.Usage, res.Requests) + } + for _, d := range res.Documents { + if (d.DocumentID == "a" || d.DocumentID == "c") && !d.Read { + t.Fatalf("document %s was read but not marked: %+v", d.DocumentID, d) + } + } + if strings.Join(steps, ",") != "route,read,merge,coverage,route,read,merge,coverage" { + t.Fatalf("steps reported out of order: %v", steps) + } + // Pages and leaves inside a corpus read are asked in multipart wording. + for _, inst := range log.inst["p"] { + if !strings.Contains(inst, "several documents") { + t.Fatalf("page judged with single-document wording: %q", inst) + } + } + if len(log.inst["l"]) == 0 || !strings.Contains(log.inst["l"][0], "several documents") { + t.Fatalf("leaf judged with single-document wording: %v", log.inst["l"]) + } +} + +func TestCorpusStopsWhenCoverageIsReached(t *testing.T) { + j, _ := corpusJudge(t) + nav := CorpusNavigator{Judge: j, MinDocuments: 2} + docs := newsCorpus() + docs[0] = article("a", "TechNews: Acme results", "Acme reported revenue of 10 billion. Beta-fact: none needed.") + res, err := nav.Navigate(context.Background(), "What revenue did Acme report?", docs) + if err != nil { + t.Fatal(err) + } + if len(res.Hops) != 1 { + t.Fatalf("complete evidence after hop 1 should end the walk, got %d hops", len(res.Hops)) + } +} + +func TestCorpusWithNothingRelevantReportsLowCoverage(t *testing.T) { + j, _ := corpusJudge(t) + nav := CorpusNavigator{Judge: j} + docs := newsCorpus()[3:] // sports, weather, film + res, err := nav.Navigate(context.Background(), "What revenue did Acme report?", docs) + if err != nil { + t.Fatal(err) + } + if res.Coverage >= 0.5 { + t.Fatalf("a collection with no answer must not look covered: %v", res.Coverage) + } + if len(res.Hops) > 2 { + t.Fatalf("with nothing useful left to route, the bridge must stop: %+v", res.Hops) + } + for _, e := range res.Evidence { + if e.Confidence >= 0.5 { + t.Fatalf("no page should clear the threshold: %+v", e) + } + } +} + +func TestCorpusMergeRepresentsEveryDocumentBeforeASecondPage(t *testing.T) { + nav := CorpusNavigator{MaxPagesPerDocument: 4, MaxEvidence: 3} + pool := []CorpusEvidence{ + {DocumentID: "x", Page: 1, Confidence: 0.99}, + {DocumentID: "x", Page: 2, Confidence: 0.98}, + {DocumentID: "x", Page: 3, Confidence: 0.97}, + {DocumentID: "x", Page: 4, Confidence: 0.96}, + {DocumentID: "x", Page: 5, Confidence: 0.95}, + {DocumentID: "y", Page: 9, Confidence: 0.6}, + {DocumentID: "z", Page: 1, Confidence: 0.2}, // under threshold + {DocumentID: "x", Page: 1, Confidence: 0.5}, // the same page judged again, lower + } + got := nav.merge(pool) + if len(got) != 3 { + t.Fatalf("want 3 pages, got %+v", got) + } + if got[0].DocumentID != "x" || got[0].Confidence != 0.99 { + t.Fatalf("best page first, keeping its best score: %+v", got[0]) + } + var hasY bool + for _, e := range got { + hasY = hasY || e.DocumentID == "y" + if e.DocumentID == "z" { + t.Fatalf("a page under the threshold joined the list: %+v", got) + } + } + if !hasY { + t.Fatalf("the weaker document must still be represented: %+v", got) + } + + capNav := CorpusNavigator{MaxPagesPerDocument: 2, MaxEvidence: 10} + capped := capNav.merge(pool) + n := 0 + for _, e := range capped { + if e.DocumentID == "x" { + n++ + } + } + if n != 2 { + t.Fatalf("per-document cap ignored: %d pages from x", n) + } +} + +func TestCorpusMergeKeepsBestPagesWhenNoneClear(t *testing.T) { + var nav CorpusNavigator + got := nav.merge([]CorpusEvidence{ + {DocumentID: "a", Page: 1, Confidence: 0.3}, + {DocumentID: "b", Page: 1, Confidence: 0.4}, + {DocumentID: "c", Page: 1, Confidence: 0.1}, + }) + if len(got) != 2 || got[0].DocumentID != "b" { + t.Fatalf("want the two best low-confidence pages, got %+v", got) + } +} + +func TestCorpusRouteBatchesUnderTheBudget(t *testing.T) { + var requests atomic.Int32 + j := &llmgate.MockJudge{Respond: func(_ context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + requests.Add(1) + ans := map[string]llmgate.Answer{} + for id := range req.Questions { + ans[id] = llmgate.NoulAnswer{Noul: 0.1} + } + return &llmgate.Judgment{Answers: ans}, nil + }} + var docs []*CorpusDoc + for i := 0; i < 300; i++ { + d := article(string(rune('a'+i%26))+strings.Repeat("x", i), "doc", strings.Repeat("word ", 180)) + docs = append(docs, &d) + } + nav := CorpusNavigator{Judge: j, Navigator: JudgeNavigator{RequestBudgetTokens: 24000}} + scores, _, n, err := nav.route(context.Background(), "q", "", docs) + if err != nil { + t.Fatal(err) + } + if len(scores) != 300 { + t.Fatalf("every document scored once, got %d", len(scores)) + } + if n < 3 || int(requests.Load()) != n { + t.Fatalf("300 cards of ~225 tokens need several requests under 24k, got %d", n) + } +} + +func TestCorpusCardShowsContentsAndOpening(t *testing.T) { + in := &NavInputs{Leaves: []NavLeaf{ + {ID: "1", Title: "Business", Path: "PART I > Business"}, + {ID: "2", Title: "Risk", Path: "PART I > Risk"}, + {ID: "3", Title: "Statements", Path: "PART II > Statements"}, + }} + card := CorpusCard("Acme 10-K", in, "Acme Corp annual report for fiscal 2023", 0) + for _, want := range []string{"Acme 10-K", "Contents: PART I; PART II", "Opens: Acme Corp annual report"} { + if !strings.Contains(card, want) { + t.Fatalf("card missing %q:\n%s", want, card) + } + } + if len(CorpusCard("t", in, strings.Repeat("a", 5000), 300)) > 300 { + t.Fatal("card not bounded") + } +} + +type countingPages struct { + calls atomic.Int32 + pages map[tree.DocumentID][]NavPage +} + +func (c *countingPages) LoadPages(_ context.Context, id tree.DocumentID) ([]NavPage, error) { + c.calls.Add(1) + return c.pages[id], nil +} + +func TestCorpusLibraryCachesUntilTheDocumentChanges(t *testing.T) { + ps := &countingPages{pages: map[tree.DocumentID][]NavPage{ + "a": {{Number: 2, Text: "second"}, {Number: 1, Text: "Acme opens here"}}, + }} + lib := NewCorpusLibrary(&JudgeWalkStrategy{Pages: ps}, nil) + t0 := time.Unix(100, 0) + ref := CorpusRef{ID: "a", Title: "Acme", UpdatedAt: t0} + docs, errs := lib.Docs(context.Background(), "org", "", []CorpusRef{ref}) + if len(errs) != 0 || len(docs) != 1 { + t.Fatalf("build failed: %v", errs) + } + if !strings.Contains(docs[0].Card, "Opens: Acme opens here") { + t.Fatalf("card should open on the first page, got %q", docs[0].Card) + } + if docs[0].Inputs.Pages[0].Number != 1 { + t.Fatal("pages should be in order") + } + _, _ = lib.Docs(context.Background(), "org", "", []CorpusRef{ref}) + if ps.calls.Load() != 1 { + t.Fatalf("second question should hit the cache, loaded %d times", ps.calls.Load()) + } + ref.UpdatedAt = t0.Add(time.Hour) + _, _ = lib.Docs(context.Background(), "org", "", []CorpusRef{ref}) + if ps.calls.Load() != 2 { + t.Fatal("a re-ingested document must be rebuilt") + } + if n, b := lib.CachedDocuments(); n != 1 || b <= 0 { + t.Fatalf("cache should hold one document, holds %d (%d bytes)", n, b) + } + + lib.MaxBytes = -1 + _, _ = lib.Docs(context.Background(), "org", "", []CorpusRef{{ID: "a", Title: "Acme", UpdatedAt: t0.Add(2 * time.Hour)}}) + _, _ = lib.Docs(context.Background(), "org", "", []CorpusRef{{ID: "a", Title: "Acme", UpdatedAt: t0.Add(2 * time.Hour)}}) + if ps.calls.Load() != 4 { + t.Fatalf("a disabled cache must rebuild every time, loaded %d", ps.calls.Load()) + } +} + +func TestCorpusLibraryReportsADocumentItCannotRead(t *testing.T) { + lib := NewCorpusLibrary(&JudgeWalkStrategy{Pages: &countingPages{}}, func(context.Context, tree.DocumentID, string, string) (*tree.Tree, error) { + return &tree.Tree{Root: &tree.Section{}}, nil + }) + docs, errs := lib.Docs(context.Background(), "org", "", []CorpusRef{{ID: "empty"}}) + if len(docs) != 0 || errs["empty"] == nil { + t.Fatalf("an unreadable document must be reported, not returned: docs=%v errs=%v", docs, errs) + } +} + +// A Markdown or text document has no pages. It must still be navigable: +// its content sections, in order, are the units (HAL-3075's MultiHop smoke +// returned nothing for 609 articles until this held). +func TestSectionTreeWithoutPagesIsNavigable(t *testing.T) { + tr := &tree.Tree{DocumentID: "md", Root: &tree.Section{ID: "root", Children: []*tree.Section{ + {ID: "intro", Title: "Intro", ContentRef: "ref/intro"}, + {ID: "body", Title: "Acme deal", ContentRef: "ref/body"}, + {ID: "empty", Title: "Heading only"}, + }}} + loader := pageLoaderFunc(func(_ context.Context, ref string) ([]byte, error) { + if ref == "ref/body" { + return []byte("Acme partnered with Beta Corp."), nil + } + return []byte("Welcome."), nil + }) + s := &JudgeWalkStrategy{PageLoader: loader} + in := s.SectionTreeNavInputs(tr) + if in == nil || len(in.Leaves) != 2 || !in.Ordinal { + t.Fatalf("want two ordinal leaves, got %+v", in) + } + if in.Leaves[1].Start != 2 || in.Leaves[1].Title != "Acme deal" { + t.Fatalf("leaves numbered in document order: %+v", in.Leaves) + } + ps, err := in.Load(context.Background(), in.Leaves[1]) + if err != nil || len(ps) != 1 || !strings.Contains(ps[0].Text, "Beta Corp") { + t.Fatalf("load: %v %+v", err, ps) + } + + j, _ := corpusJudge(t) + nav := CorpusNavigator{Judge: j, MinDocuments: 1} + res, err := nav.Navigate(context.Background(), "Who did Acme partner with?", + []CorpusDoc{{ID: "md", Title: "Acme news", Card: "Acme news", Inputs: in}}) + if err != nil { + t.Fatal(err) + } + if len(res.Evidence) == 0 || !res.Evidence[0].Ordinal || res.Evidence[0].SectionTitle != "Acme deal" { + t.Fatalf("an unpaged document should yield ordinal evidence naming its section: %+v", res.Evidence) + } + + // The single-document strategy runs over the same units, and cites no + // invented pages. + single := &JudgeWalkStrategy{Navigator: JudgeNavigator{Judge: j}, PageLoader: loader} + r, err := single.SelectWithCost(context.Background(), tr, "Who did Acme partner with?", ContextBudget{}) + if err != nil { + t.Fatal(err) + } + if len(r.SelectedIDs) == 0 || len(r.CitedPages) != 0 { + t.Fatalf("single-document navigation of an unpaged tree: ids=%v cited=%v", r.SelectedIDs, r.CitedPages) + } +} + +func TestCorpusEmptyListsAreNotNull(t *testing.T) { + j, _ := corpusJudge(t) + res, err := (&CorpusNavigator{Judge: j}).Navigate(context.Background(), "q", nil) + if err != nil { + t.Fatal(err) + } + if res.Evidence == nil || res.Documents == nil || res.Hops == nil { + t.Fatalf("empty lists must encode as [], got %+v", res) + } +} + +type pageLoaderFunc func(ctx context.Context, ref string) ([]byte, error) + +func (f pageLoaderFunc) Load(ctx context.Context, ref string) ([]byte, error) { return f(ctx, ref) } + +// Every list in a corpus response encodes as a list: a bridge hop that +// reads nothing is "read": [], not null (the SDK rejected null). +func TestCorpusJSONHasNoNullLists(t *testing.T) { + j, _ := corpusJudge(t) + nav := CorpusNavigator{Judge: j, MinDocuments: 2} + all := newsCorpus() + docs := append([]CorpusDoc{all[0], all[1]}, all[3:]...) // no Beta article: the bridge finds nothing + res, err := nav.Navigate(context.Background(), "Which company did Acme partner with, and when was it founded?", docs) + if err != nil { + t.Fatal(err) + } + if len(res.Hops) < 2 { + t.Fatalf("expected a bridge hop that reads nothing, got %+v", res.Hops) + } + raw, err := json.Marshal(res) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(raw), "null") { + t.Fatalf("null in corpus JSON: %s", raw) + } +} + +type failingPages struct{ err error } + +func (f failingPages) LoadPages(context.Context, tree.DocumentID) ([]NavPage, error) { + return nil, f.err +} + +func TestCorpusLibraryReportsAStorageFailureInsteadOfFallingBack(t *testing.T) { + treeCalls := 0 + loadTree := func(context.Context, tree.DocumentID, string, string) (*tree.Tree, error) { + treeCalls++ + return &tree.Tree{Root: &tree.Section{Children: []*tree.Section{{ID: "s", Title: "S", ContentRef: "r"}}}}, nil + } + lib := NewCorpusLibrary(&JudgeWalkStrategy{Pages: failingPages{fmt.Errorf("s3: timeout")}}, loadTree) + docs, errs := lib.Docs(context.Background(), "org", "", []CorpusRef{{ID: "pdf"}}) + if len(docs) != 0 || errs["pdf"] == nil || treeCalls != 0 { + t.Fatalf("a storage failure must be reported, not read through the tree: docs=%d errs=%v tree=%d", len(docs), errs, treeCalls) + } + missing := NewCorpusLibrary(&JudgeWalkStrategy{Pages: failingPages{fmt.Errorf("x: %w", storage.ErrNotFound)}}, loadTree) + docs, errs = missing.Docs(context.Background(), "org", "", []CorpusRef{{ID: "md"}}) + if len(docs) != 1 || len(errs) != 0 { + t.Fatalf("no persisted pages means an unpaged document, read from its tree: docs=%d errs=%v", len(docs), errs) + } +} diff --git a/pkg/retrieval/judgewalk.go b/pkg/retrieval/judgewalk.go index 2529875..6912287 100644 --- a/pkg/retrieval/judgewalk.go +++ b/pkg/retrieval/judgewalk.go @@ -201,6 +201,17 @@ type JudgeNavigator struct { // longer one falls back to the sequential skim, so the extra heads // stay bounded. Zero selects defaultNavSkimAllMax. SkimAllMaxPages int + + // Multipart words every judgement for a question whose answer is + // spread over several documents (corpus mode, HAL-3075): a section + // or page counts when it holds ANY of the facts the question needs, + // not only when it answers the whole question. Asked whether a page + // "contains the specific facts needed to answer" a comparison of two + // companies, a page carrying one company's figure is rightly judged + // False, and the comparison can never be assembled. The wording is + // also domain-neutral: the default leaf question carries 10-K hints + // (statements, MD&A) that mean nothing to a news article. + Multipart bool } const ( @@ -339,20 +350,7 @@ func (n *JudgeNavigator) RankLeaves(ctx context.Context, query string, leaves [] item["summary"] = s } state[qk] = item - questions[qk] = llmgate.Noul{ - Instructions: fmt.Sprintf( - "`question` is a question about a document. `%s` is one section of that "+ - "document's table of contents: its title, where it sits, and its page range. "+ - "Would the information needed to answer the question be found in this "+ - "section? Judge from what such a section of such a document contains: "+ - "figures live in the statements and their notes, discussion of results in "+ - "the MD&A, customers, competition and outlook in the business and risk "+ - "sections.", qk), - Criteria: &llmgate.NoulCriteria{ - True: "This section is where a reader would look for the answer", - False: "The answer would not be in this section", - }, - } + questions[qk] = n.leafQuestion(qk) } batches = append(batches, batch{state, questions}) } @@ -402,6 +400,91 @@ func (n *JudgeNavigator) RankLeaves(ctx context.Context, query string, leaves [] return scores, usage, requests, nil } +// leafQuestion is the judgement asked of one leaf. +func (n *JudgeNavigator) leafQuestion(qk string) llmgate.Noul { + if n.Multipart { + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is a question whose answer may draw on several documents. `%s` is one "+ + "section of one document: its title, where it sits, and its page range. Would this "+ + "section hold any of the information the question needs — one of the entities, "+ + "events, figures or claims it asks about — even if not all of it?", qk), + Criteria: &llmgate.NoulCriteria{ + True: "A reader would look here for at least part of what the question needs", + False: "Nothing the question needs would be in this section", + }, + } + } + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is a question about a document. `%s` is one section of that "+ + "document's table of contents: its title, where it sits, and its page range. "+ + "Would the information needed to answer the question be found in this "+ + "section? Judge from what such a section of such a document contains: "+ + "figures live in the statements and their notes, discussion of results in "+ + "the MD&A, customers, competition and outlook in the business and risk "+ + "sections.", qk), + Criteria: &llmgate.NoulCriteria{ + True: "This section is where a reader would look for the answer", + False: "The answer would not be in this section", + }, + } +} + +// pageQuestion is the judgement asked of one page: of its head for the +// coarse pass, of its whole text otherwise. +func (n *JudgeNavigator) pageQuestion(qk string, coarse bool) llmgate.Noul { + switch { + case coarse && n.Multipart: + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is a question whose answer may draw on several documents. `%s.text` is the "+ + "START of page `%s.page` of one of them. Judging from what it opens with, could this "+ + "page hold any of the facts the question needs?", qk, qk), + Criteria: &llmgate.NoulCriteria{ + True: "This page's opening says it would hold at least part of what the question needs", + False: "This page opens on something unrelated", + }, + } + case coarse: + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is a question about a document. `%s.text` is the START of page "+ + "`%s.page` — its heading and first lines. Could the facts or figures needed "+ + "to answer the question be on this page, judging from what it opens with?", qk, qk), + Criteria: &llmgate.NoulCriteria{ + True: "This page's opening says it is the kind of page that would hold the answer", + False: "This page opens on something unrelated", + }, + } + case n.Multipart: + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is a question whose answer may draw on several documents. `%s.text` is the "+ + "text of page `%s.page` of one of them. Does this page state at least one of the "+ + "specific facts the question needs — a figure, a name, a date, an event, a claim "+ + "attributed to someone — even if the rest is elsewhere? A page that only mentions "+ + "the topic does not.", qk, qk), + Criteria: &llmgate.NoulCriteria{ + True: "At least one fact the question needs is stated on this page", + False: "The page states none of the facts the question needs", + }, + } + default: + return llmgate.Noul{ + Instructions: fmt.Sprintf( + "`question` is a question about a document. `%s.text` is the text of page "+ + "`%s.page`. Does this page contain the specific facts or figures needed to "+ + "answer the question — the number, the statement, the table row? A page that "+ + "only mentions the topic, or refers the reader elsewhere, does not.", qk, qk), + Criteria: &llmgate.NoulCriteria{ + True: "The answer, or a figure it is computed from, is on this page", + False: "The page is about something else, or only mentions the topic", + }, + } + } +} + // RankPages asks, for every page, whether it contains the facts or // figures the question asks for. Pages are batched so the request's // whole state — query plus every page in the batch — stays under the @@ -446,30 +529,7 @@ func (n *JudgeNavigator) rankPages(ctx context.Context, query string, pages []Na } qk := fmt.Sprintf("p_%d", end) state[qk] = map[string]any{"page": pages[end].Number, "text": text} - if coarse { - questions[qk] = llmgate.Noul{ - Instructions: fmt.Sprintf( - "`question` is a question about a document. `%s.text` is the START of page "+ - "`%s.page` — its heading and first lines. Could the facts or figures needed "+ - "to answer the question be on this page, judging from what it opens with?", qk, qk), - Criteria: &llmgate.NoulCriteria{ - True: "This page's opening says it is the kind of page that would hold the answer", - False: "This page opens on something unrelated", - }, - } - } else { - questions[qk] = llmgate.Noul{ - Instructions: fmt.Sprintf( - "`question` is a question about a document. `%s.text` is the text of page "+ - "`%s.page`. Does this page contain the specific facts or figures needed to "+ - "answer the question — the number, the statement, the table row? A page that "+ - "only mentions the topic, or refers the reader elsewhere, does not.", qk, qk), - Criteria: &llmgate.NoulCriteria{ - True: "The answer, or a figure it is computed from, is on this page", - False: "The page is about something else, or only mentions the topic", - }, - } - } + questions[qk] = n.pageQuestion(qk, coarse) used += cost end++ } @@ -906,48 +966,15 @@ func (s *JudgeWalkStrategy) selectOnPersistedPages(ctx context.Context, t *tree. if err != nil || len(raw) == 0 { return nil, false } - var nodes []tree.TOCNode - if err := json.Unmarshal(raw, &nodes); err != nil || len(nodes) == 0 { - return nil, false - } pages, err := s.Pages.LoadPages(ctx, t.DocumentID) if err != nil || len(pages) == 0 { return nil, false } - byNum := make(map[int]NavPage, len(pages)) - for _, p := range pages { - byNum[p.Number] = p - } - var leaves []NavLeaf - var walk func(ns []tree.TOCNode, path string) - walk = func(ns []tree.TOCNode, path string) { - for _, n := range ns { - p := n.Title - if path != "" { - p = path + " > " + n.Title - } - if len(n.Nodes) > 0 { - walk(n.Nodes, p) - continue - } - if n.StartPage > 0 && n.EndPage >= n.StartPage { - leaves = append(leaves, NavLeaf{ID: n.NodeID, Title: n.Title, Path: p, Start: n.StartPage, End: n.EndPage, Summary: n.Summary}) - } - } - } - walk(nodes, "") - if len(leaves) == 0 { + in := PersistedNavInputs(raw, pages) + if in == nil { return nil, false } - load := func(_ context.Context, leaf NavLeaf) ([]NavPage, error) { - var ps []NavPage - for n := leaf.Start; n <= leaf.End; n++ { - if p, ok := byNum[n]; ok { - ps = append(ps, p) - } - } - return ps, nil - } + leaves, load := in.Leaves, in.Load // Every page is already in memory here, so skimming them all costs // tokens but no loads — the case SkimAll was built for. navigator := s.Navigator @@ -1000,10 +1027,159 @@ func (s *JudgeWalkStrategy) selectOnPersistedPages(ctx context.Context, t *tree. }, true } +// NavInputs is what navigation runs over for one document: its leaves +// and how to read a leaf's pages. Pages is every page in order when they +// are already in memory (the persisted-pages path), nil when each leaf +// is read from storage on demand. +type NavInputs struct { + Leaves []NavLeaf + Load func(ctx context.Context, leaf NavLeaf) ([]NavPage, error) + Pages []NavPage + // Ordinal is set when the document has no pages and each unit's number + // is its section's position in the document, not a page (navSections). + Ordinal bool +} + +// PersistedNavInputs builds navigation inputs from a persisted table of +// contents and the persisted per-page text. Nil when the contents carry +// no leaf with a page range. +func PersistedNavInputs(tocRaw []byte, pages []NavPage) *NavInputs { + var nodes []tree.TOCNode + if err := json.Unmarshal(tocRaw, &nodes); err != nil || len(nodes) == 0 || len(pages) == 0 { + return nil + } + byNum := make(map[int]NavPage, len(pages)) + for _, p := range pages { + byNum[p.Number] = p + } + var leaves []NavLeaf + var walk func(ns []tree.TOCNode, path string) + walk = func(ns []tree.TOCNode, path string) { + for _, n := range ns { + p := n.Title + if path != "" { + p = path + " > " + n.Title + } + if len(n.Nodes) > 0 { + walk(n.Nodes, p) + continue + } + if n.StartPage > 0 && n.EndPage >= n.StartPage { + leaves = append(leaves, NavLeaf{ID: n.NodeID, Title: n.Title, Path: p, Start: n.StartPage, End: n.EndPage, Summary: n.Summary}) + } + } + } + walk(nodes, "") + if len(leaves) == 0 { + return nil + } + return &NavInputs{Leaves: leaves, Load: pagesInRange(byNum), Pages: sortedPages(pages)} +} + +// PagedNavInputs covers a document with per-page text but no contents +// tree: one leaf spanning every page, titled with the document's title. +func PagedNavInputs(title string, pages []NavPage) *NavInputs { + if len(pages) == 0 { + return nil + } + ps := sortedPages(pages) + byNum := make(map[int]NavPage, len(ps)) + for _, p := range ps { + byNum[p.Number] = p + } + leaf := NavLeaf{ID: "doc", Title: title, Path: title, Start: ps[0].Number, End: ps[len(ps)-1].Number} + return &NavInputs{Leaves: []NavLeaf{leaf}, Load: pagesInRange(byNum), Pages: ps} +} + +// SectionTreeNavInputs builds navigation inputs from the parser's section +// tree: each section with content is a leaf, its body read through the +// strategy's PageLoader and chunked into page-sized units. +func (s *JudgeWalkStrategy) SectionTreeNavInputs(t *tree.Tree) *NavInputs { + if t == nil || t.Root == nil { + return nil + } + sections, ordinal := navSections(t) + byID := map[string]sectionPageEntry{} + paths := sectionPaths(t) + var leaves []NavLeaf + for _, sec := range sections { + byID[string(sec.id)] = sec + leaves = append(leaves, NavLeaf{ + ID: string(sec.id), Title: sec.title, Path: paths[sec.id], + Start: sec.start, End: sec.end, Summary: sec.summary, + }) + } + if len(leaves) == 0 { + return nil + } + chars := s.Navigator.pageChars() + load := func(ctx context.Context, leaf NavLeaf) ([]NavPage, error) { + sec := byID[leaf.ID] + if s.PageLoader == nil || sec.contentRef == "" { + return nil, nil + } + b, err := s.PageLoader.Load(ctx, sec.contentRef) + if err != nil { + return nil, err + } + return chunkBody(string(b), sec.start, sec.end, chars), nil + } + return &NavInputs{Leaves: leaves, Load: load, Ordinal: ordinal} +} + +func pagesInRange(byNum map[int]NavPage) func(context.Context, NavLeaf) ([]NavPage, error) { + return func(_ context.Context, leaf NavLeaf) ([]NavPage, error) { + var ps []NavPage + for n := leaf.Start; n <= leaf.End; n++ { + if p, ok := byNum[n]; ok { + ps = append(ps, p) + } + } + return ps, nil + } +} + +func sortedPages(pages []NavPage) []NavPage { + out := append([]NavPage(nil), pages...) + sort.SliceStable(out, func(i, j int) bool { return out[i].Number < out[j].Number }) + return out +} + +// navSections lists the sections navigation runs over: every section with +// a page range when the document has pages, and otherwise — Markdown, plain +// text, HTML, a document with no page numbers at all — every section that +// carries content, in document order, numbered 1..n as one unit each. +// ordinal reports the second case, where a unit's number is its position, +// not a page: citations then name sections, never invented pages. +// +// Before this, a document without pages gave navigation nothing at all, and +// /v1/query and corpus mode returned no evidence for every Markdown or text +// document (HAL-3075's MultiHop smoke: 609 articles, "no navigable sections"). +func navSections(t *tree.Tree) (sections []sectionPageEntry, ordinal bool) { + if s := flattenSectionsByPage(t); len(s) > 0 { + return s, false + } + if t == nil || t.Root == nil { + return nil, false + } + n := 0 + t.Root.Walk(func(s *tree.Section) bool { + if strings.TrimSpace(s.ContentRef) == "" { + return true + } + n++ + sections = append(sections, sectionPageEntry{ + id: s.ID, title: s.Title, summary: s.Summary, contentRef: s.ContentRef, start: n, end: n, + }) + return true + }) + return sections, len(sections) > 0 +} + // selectOnSectionTree is the fallback: the parser's section tree, each // section's body chunked into page-sized units. func (s *JudgeWalkStrategy) selectOnSectionTree(ctx context.Context, t *tree.Tree, query string, _ ContextBudget) (*Result, error) { - sections := flattenSectionsByPage(t) + sections, ordinal := navSections(t) byID := map[string]sectionPageEntry{} paths := sectionPaths(t) var leaves []NavLeaf @@ -1049,6 +1225,9 @@ func (s *JudgeWalkStrategy) selectOnSectionTree(ctx context.Context, t *tree.Tre if len(nav.Evidence) > 0 { best = nav.Evidence[0].P } + if ordinal { + ranges = nil // the units were numbered by position; there are no pages to cite + } return &Result{ SelectedIDs: ids, Confidences: conf,