feat(answer): do an answer's arithmetic in code and trace each operand to its page

internal/api/server.go

@@ -21,6 +21,7 @@ import (
 	"github.com/go-chi/chi/v5/middleware"
 	"github.com/hallelx2/llmgate"
 
+	"github.com/hallelx2/vectorless-engine/pkg/calc"
 	"github.com/hallelx2/vectorless-engine/pkg/config"
 	"github.com/hallelx2/vectorless-engine/pkg/db"
 	"github.com/hallelx2/vectorless-engine/pkg/ingest"
@@ -1014,6 +1015,16 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	totalUsage.Add(synthUsage)
+	// The model set its calculations up; code works them out (HAL-3161).
+	answerText, calcs := calc.Resolve(answerText)
+	sources := make([]calc.Source, len(enriched))
+	for i, e := range enriched {
+		sources[i] = calc.Source{Page: e.sec.PageStart, Text: e.content}
+	}
+	calc.Trace(calcs, sources)
+	if calcs == nil {
+		calcs = []calc.Calculation{}
+	}
 
 	citations := make([]map[string]any, 0, len(enriched))
 	finalIDs := make([]tree.SectionID, 0, len(enriched))
@@ -1049,12 +1060,13 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
 	traceToken := retrieval.ComputeTraceToken(body.DocumentID, "1", synthModel, finalIDs)
 
 	resp := map[string]any{
-		"document_id": body.DocumentID,
-		"query":       body.Query,
-		"answer":      answerText,
-		"citations":   citations,
-		"strategy":    retrieval.PublicName(d.Strategy.Name()),
-		"model":       synthModel,
+		"document_id":  body.DocumentID,
+		"query":        body.Query,
+		"answer":       answerText,
+		"calculations": calcs,
+		"citations":    citations,
+		"strategy":     retrieval.PublicName(d.Strategy.Name()),
+		"model":        synthModel,
 		"usage": map[string]any{
 			"input_tokens":  totalUsage.InputTokens,
 			"output_tokens": totalUsage.OutputTokens,
@@ -1120,7 +1132,8 @@ func synthesiseAnswer(ctx context.Context, client llmgate.Client, model, query s
 		}
 	}
 	b.WriteString("\nWrite a concise answer to the user's query. ")
-	b.WriteString("If the evidence does not contain an answer, say so. ")
+	b.WriteString(calc.Instructions)
+	b.WriteString(" If the evidence does not contain an answer, say so. ")
 	b.WriteString("Inline citations should reference the section_id values shown above. ")
 	b.WriteString("Output plain prose; no JSON.")
 

internal/handler/answer.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/hallelx2/llmgate"
 
+	"github.com/hallelx2/vectorless-engine/pkg/calc"
 	enginecfg "github.com/hallelx2/vectorless-engine/pkg/config"
 	"github.com/hallelx2/vectorless-engine/pkg/db"
 	"github.com/hallelx2/vectorless-engine/pkg/pincite"
@@ -359,6 +360,12 @@ func (h *AnswerHandler) answerFromSections(ctx context.Context, t *tree.Tree, bo
 		return nil, nil, "", "", fmt.Errorf("synthesis failed: %w", err)
 	}
 	totalUsage.Add(synthUsage)
+	answerText, calcs := calc.Resolve(answerText)
+	sources := make([]calc.Source, len(enriched))
+	for i, e := range enriched {
+		sources[i] = calc.Source{Page: e.sec.PageStart, Text: e.content}
+	}
+	calc.Trace(calcs, sources)
 
 	citations := make([]map[string]any, 0, len(enriched))
 	finalIDs := make([]tree.SectionID, 0, len(enriched))
@@ -390,15 +397,16 @@ func (h *AnswerHandler) answerFromSections(ctx context.Context, t *tree.Tree, bo
 	traceToken := retrieval.ComputeTraceToken(body.DocumentID, "1", synthModel, finalIDs)
 
 	resp := map[string]any{
-		"document_id": body.DocumentID,
-		"query":       body.Query,
-		"answer":      answerText,
-		"citations":   citations,
-		"strategy":    retrieval.PublicName(h.strategy.Name()),
-		"model":       synthModel,
-		"usage":       usageMap(totalUsage),
-		"elapsed_ms":  time.Since(started).Milliseconds(),
-		"trace_token": traceToken,
+		"document_id":  body.DocumentID,
+		"query":        body.Query,
+		"answer":       answerText,
+		"citations":    citations,
+		"calculations": nonNilCalcs(calcs),
+		"strategy":     retrieval.PublicName(h.strategy.Name()),
+		"model":        synthModel,
+		"usage":        usageMap(totalUsage),
+		"elapsed_ms":   time.Since(started).Milliseconds(),
+		"trace_token":  traceToken,
 	}
 	return resp, finalIDs, synthModel, traceToken, nil
 }
@@ -537,7 +545,8 @@ func synthesiseAnswer(ctx context.Context, client llmgate.Client, model, query s
 		}
 	}
 	b.WriteString("\nWrite a concise answer to the user's query. ")
-	b.WriteString("If the evidence does not contain an answer, say so. ")
+	b.WriteString(calc.Instructions)
+	b.WriteString(" If the evidence does not contain an answer, say so. ")
 	b.WriteString("Inline citations should reference the section_id values shown above. ")
 	b.WriteString("Output plain prose; no JSON.")
 

internal/handler/answer_calc_test.go

@@ -0,0 +1,72 @@
+package handler
+
+import (
+	"context"
+	"log/slog"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/hallelx2/llmgate"
+
+	"github.com/hallelx2/vectorless-engine/pkg/calc"
+	enginecfg "github.com/hallelx2/vectorless-engine/pkg/config"
+	"github.com/hallelx2/vectorless-engine/pkg/pincite"
+	"github.com/hallelx2/vectorless-engine/pkg/retrieval"
+)
+
+// The answer model sets a calculation up; the engine works it out and
+// finds each input on the page that states it (HAL-3161). FinanceBench's
+// PepsiCo FY2022 EBITDA margin, from the income statement and the cash
+// flow statement.
+func TestPageAnswerComputesItsCalculations(t *testing.T) {
+	var prompt string
+	llm := &llmgate.Mock{Respond: func(_ context.Context, req llmgate.Request) (*llmgate.Response, error) {
+		prompt = req.Messages[len(req.Messages)-1].Content
+		return &llmgate.Response{Content: "Unadjusted EBITDA margin was <<(11,512 + 2,763) / 86,392 * 100>>% [1][2]."}, nil
+	}}
+	h := NewAnswerHandler(slog.Default(), nil, nil, evidenceStrategy{}, llm, "m",
+		enginecfg.AnswerSpanBlock{}, enginecfg.AnswerBlock{}, nil)
+	res := &retrieval.Result{EvidencePages: []retrieval.EvidencePage{
+		{Page: 62, Title: "Income statement", Text: "Net Revenue $ 86,392 $ 79,474. Operating Profit 11,512 11,162.", Confidence: 0.9},
+		{Page: 64, Title: "Cash flows", Text: "Depreciation and amortization 2,763 2,710.", Confidence: 0.8},
+	}}
+	resp, _, _, _, err := h.answerFromPages(context.Background(), queryStrategyTestTree(), pincite.Source{},
+		answerRequest{DocumentID: "doc_x", Query: "PepsiCo FY2022 unadjusted EBITDA % margin?"}, res, time.Now(), func(string, any) {})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(prompt, "<<expression>>") {
+		t.Error("the answer prompt does not tell the model how to set a calculation up")
+	}
+	answer := resp["answer"].(string)
+	if !strings.Contains(answer, "16.5235 (= (11,512 + 2,763) / 86,392 * 100)%") {
+		t.Errorf("answer = %q, want the computed value beside its set-up", answer)
+	}
+	calcs := resp["calculations"].([]calc.Calculation)
+	if len(calcs) != 1 || calcs[0].Value == nil {
+		t.Fatalf("calculations = %+v", calcs)
+	}
+	ops := calcs[0].Operands
+	if !ops[0].Found || ops[0].Page != 62 || !ops[1].Found || ops[1].Page != 64 || !ops[2].Found || ops[2].Page != 62 {
+		t.Errorf("operands not traced to their pages: %+v", ops)
+	}
+	if ops[3].Found {
+		t.Errorf("100 is the formula's own constant, on no page: %+v", ops[3])
+	}
+}
+
+func TestPageAnswerWithoutCalculationsSaysSo(t *testing.T) {
+	llm := &llmgate.Mock{Reply: "3M spent $1,577 million [1]."}
+	h := NewAnswerHandler(slog.Default(), nil, nil, evidenceStrategy{}, llm, "m",
+		enginecfg.AnswerSpanBlock{}, enginecfg.AnswerBlock{}, nil)
+	res := &retrieval.Result{EvidencePages: []retrieval.EvidencePage{{Page: 39, Text: "Total Company 1,577 1,373"}}}
+	resp, _, _, _, err := h.answerFromPages(context.Background(), queryStrategyTestTree(), pincite.Source{},
+		answerRequest{DocumentID: "doc_x", Query: "capex?"}, res, time.Now(), func(string, any) {})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c := resp["calculations"].([]calc.Calculation); c == nil || len(c) != 0 {
+		t.Errorf("calculations = %#v, want an empty array", c)
+	}
+}

internal/handler/answer_pages.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/hallelx2/llmgate"
 
+	"github.com/hallelx2/vectorless-engine/pkg/calc"
 	"github.com/hallelx2/vectorless-engine/pkg/pincite"
 	"github.com/hallelx2/vectorless-engine/pkg/retrieval"
 	"github.com/hallelx2/vectorless-engine/pkg/tree"
@@ -82,6 +83,14 @@ func (h *AnswerHandler) answerFromPages(
 		return nil, nil, "", "", fmt.Errorf("synthesis failed: %w", err)
 	}
 	total.Add(synthUsage)
+	// The model set its calculations up; code works them out, and each
+	// operand is found on the evidence page that states it (HAL-3161).
+	answer, calcs := calc.Resolve(answer)
+	sources := make([]calc.Source, len(evidence))
+	for i, ev := range evidence {
+		sources[i] = calc.Source{Page: ev.Page, Text: ev.Text}
+	}
+	calc.Trace(calcs, sources)
 	answerMS := time.Since(started).Milliseconds() - retrievalMS
 
 	var layout *pincite.Layout
@@ -215,6 +224,7 @@ func (h *AnswerHandler) answerFromPages(
 		"citations":        citations,
 		"pincites":         pincites,
 		"evidence":         evList,
+		"calculations":     nonNilCalcs(calcs),
 		"markers_inline":   inline,
 		"markers_stripped": stripped,
 		"strategy":         retrieval.PublicName(h.strategy.Name()),
@@ -233,6 +243,15 @@ func (h *AnswerHandler) answerFromPages(
 	return resp, finalIDs, model, traceToken, nil
 }
 
+// nonNilCalcs keeps "calculations" an array in the response when the answer
+// set none up.
+func nonNilCalcs(c []calc.Calculation) []calc.Calculation {
+	if c == nil {
+		return []calc.Calculation{}
+	}
+	return c
+}
+
 // publicSelector names how a pincite's sentence was chosen, without
 // exposing which model makes the choice: "model" or "overlap" (the
 // lexical fallback).
@@ -283,7 +302,8 @@ func synthesiseFromPages(ctx context.Context, client llmgate.Client, model, quer
 	b.WriteString("End every sentence that states a fact from the evidence with the marker of the evidence block it came from, like [2]; ")
 	b.WriteString("use only the numbers shown above, and use two markers like [1][3] only when a sentence draws on two blocks. ")
 	b.WriteString("Give figures exactly as the page states them, with their units and periods. ")
-	b.WriteString("If the evidence does not answer the question, say so plainly and do not guess.")
+	b.WriteString(calc.Instructions)
+	b.WriteString(" If the evidence does not answer the question, say so plainly and do not guess.")
 
 	resp, err := client.Complete(ctx, llmgate.Request{
 		Model: model,