Maven/cmd/mavend/query_recall_test.go

package main

import (
	"context"
	"errors"
	"fmt"
	"math"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/kami/maven/internal/ipc"
	"github.com/kami/maven/internal/memory"
	"github.com/kami/maven/internal/phraser"
	"github.com/kami/maven/internal/router"
	"github.com/kami/maven/internal/voice"
)

// fixedEmbedder hands back a vector chosen per text, so a test can say exactly
// how close each stored memory is to the question. The real embedders make
// scores that are realistic but not controllable, and this test is about the
// gate, not about the embedder.
type fixedEmbedder struct{ vecs map[string][]float32 }

func (f *fixedEmbedder) Dim() int     { return 4 }
func (f *fixedEmbedder) Close() error { return nil }

func (f *fixedEmbedder) Embed(_ context.Context, text string) ([]float32, error) {
	v, ok := f.vecs[text]
	if !ok {
		return nil, fmt.Errorf("fixedEmbedder: no vector for %q", text)
	}
	return v, nil
}

// brokenEmbedder fails every call, which is what an ONNX session error looks
// like from the query chain's side.
type brokenEmbedder struct{}

func (brokenEmbedder) Dim() int     { return 4 }
func (brokenEmbedder) Close() error { return nil }
func (brokenEmbedder) Embed(context.Context, string) ([]float32, error) {
	return nil, errors.New("onnx: session failed")
}

// TestQueryEmbedFailureDoesNotStopTheChain — V-568. The embed source used to
// claim the turn on an embedder error, so one failing EmbedQuery answered every
// question below it with "не смогла ответить", including the ones the search
// answers without an embedder at all. A source that could not look must pass.
func TestQueryEmbedFailureDoesNotStopTheChain(t *testing.T) {
	const q = "почему небо голубое"

	h, _ := searchHandler(t, searchBody, http.StatusOK)
	h.api = ipc.NewStoreAPI(newTestStore(t))
	h.recall = recallWiring{embedder: brokenEmbedder{}, minScore: 0.55, minMargin: 0.008}
	h.now = time.Now

	// The embed source itself passes rather than claiming.
	turn := &queryTurn{dec: router.Decision{Intent: router.IntentQuery, Utterance: q}}
	if reply, ok := h.queryEmbed(context.Background(), turn); ok {
		t.Fatalf("queryEmbed claimed the turn on an embedder error: %q", reply)
	}

	// And the whole chain still reaches the search below it.
	reply := askQuery(t, h, q)
	if !strings.Contains(reply, "рэлеевского рассеяния") {
		t.Fatalf("reply = %q, want the search answer", reply)
	}
}

// The recall sources read the empty vector the failed embed left behind, and
// neither of them may turn that into an answer: no vector means they could not
// look, which is not the same as looking and finding nothing.
func TestQueryRecallPassesWithoutAVector(t *testing.T) {
	h, _ := buildRecallHandler(t, "где молоко", []recallCase{
		{text: "молоко стоит в холодильнике", score: 0.90, kind: "note"},
	})
	turn := &queryTurn{dec: router.Decision{Intent: router.IntentQuery, Utterance: "где молоко"}}
	if reply, ok := h.queryMemory(context.Background(), turn); ok {
		t.Errorf("queryMemory claimed with no vector: %q", reply)
	}
	if reply, ok := h.queryNotes(context.Background(), turn); ok {
		t.Errorf("queryNotes claimed with no vector: %q", reply)
	}
}

// scoreVec builds a unit vector whose cosine against the query vector
// (1,0,0,0) is exactly score.
func scoreVec(score float64) []float32 {
	rest := math.Sqrt(1 - score*score)
	return []float32{float32(score), float32(rest), 0, 0}
}

// recordingPhraser remembers what the query path handed it to phrase, which is
// how the test can tell which pass produced the answer.
type recordingPhraser struct {
	*phraser.Stub
	notes []string
}

func (r *recordingPhraser) PhraseQuery(ctx context.Context, utterance string, notes []string) (string, error) {
	r.notes = notes
	return r.Stub.PhraseQuery(ctx, utterance, notes)
}

// recallCase — one stored memory: its text, how close it is to the question,
// whether it is a note or a fact, and whether the notes table holds it too.
type recallCase struct {
	text  string
	score float64
	kind  string
}

// buildRecallHandler stores the given memories and returns a handler whose
// query path can be run directly. Notes go into BOTH the notes table and the
// vector index, which is what the daemon does (voice.go's IntentNote).
func buildRecallHandler(t *testing.T, question string, mems []recallCase) (*reactiveHandler, *recordingPhraser) {
	t.Helper()
	ctx := context.Background()
	st := newTestStore(t)
	emb := &fixedEmbedder{vecs: map[string][]float32{question: {1, 0, 0, 0}}}
	mem := memory.NewInMemoryStore()
	now := time.Now()

	for i, m := range mems {
		vec := scoreVec(m.score)
		emb.vecs[m.text] = vec
		id := fmt.Sprintf("%s:%d", m.kind, i)
		if m.kind == "note" {
			if _, err := st.WriteNote(ctx, now, m.text, vec, "tap:voice"); err != nil {
				t.Fatalf("WriteNote: %v", err)
			}
		}
		if err := mem.Insert(ctx, id, vec, map[string]string{"text": m.text, "type": m.kind}); err != nil {
			t.Fatalf("memory insert: %v", err)
		}
	}

	phr := &recordingPhraser{Stub: phraser.NewStub()}
	h := &reactiveHandler{
		api: ipc.NewStoreAPI(st),
		recall: recallWiring{
			embedder:  emb,
			memStore:  mem,
			minScore:  0.55,
			minMargin: 0.008,
		},
		replier:         voice.NewStubReplier(),
		phraser:         phr,
		now:             func() time.Time { return now },
		dataStore:       st,
		weatherProvider: nil,
	}
	return h, phr
}

func askQuery(t *testing.T, h *reactiveHandler, question string) string {
	t.Helper()
	return h.applyAction(context.Background(), router.Decision{
		Intent:    router.IntentQuery,
		Utterance: question,
	})
}

// TestQueryRecallNoteCanWin — the note-recall regression (Vikunja #373). Notes
// and facts share one vector index, and a note that clearly beats everything
// else must be the answer. Before the fix the memory pass only ran after the
// notes-only gate had already rejected the same note at the same score, so only
// a fact could ever come back from it.
func TestQueryRecallNoteCanWin(t *testing.T) {
	const q = "где молоко"

	t.Run("a clearly best note answers", func(t *testing.T) {
		h, phr := buildRecallHandler(t, q, []recallCase{
			{text: "молоко стоит в холодильнике", score: 0.90, kind: "note"},
			{text: "выучил пару аккордов", score: 0.50, kind: "note"},
		})
		reply := askQuery(t, h, q)
		if !phraser.IsSourcesFallback(reply, "молоко стоит в холодильнике") {
			t.Errorf("reply %q, want the note read back", reply)
		}
		// One text, the winning memory's — the answer came from the memory
		// pass, not from handing the phraser every note in the table.
		if len(phr.notes) != 1 || phr.notes[0] != "молоко стоит в холодильнике" {
			t.Errorf("phraser got %q, want just the recalled note", phr.notes)
		}
	})

	// The other half of "one gate over everything": a fact that matches better
	// than the best note now answers, instead of losing to a note that only had
	// to beat other notes.
	t.Run("the better-matching fact answers", func(t *testing.T) {
		h, _ := buildRecallHandler(t, q, []recallCase{
			{text: "молоко стоит в холодильнике", score: 0.80, kind: "note"},
			{text: "купил молоко в среду", score: 0.95, kind: "fact"},
		})
		if reply := askQuery(t, h, q); reply != "купил молоко в среду" {
			t.Errorf("reply %q, want the fact read back", reply)
		}
	})

	// The gate is untouched: two memories this close mean the embedder cannot
	// tell them apart, and silence still beats a coin flip.
	t.Run("no clear best stays silent", func(t *testing.T) {
		h, _ := buildRecallHandler(t, q, []recallCase{
			{text: "молоко стоит в холодильнике", score: 0.860, kind: "note"},
			{text: "молоко закончилось", score: 0.858, kind: "note"},
		})
		if reply := askQuery(t, h, q); !phraser.IsUnknownFallback(reply) {
			t.Errorf("reply %q, want silence", reply)
		}
	})
}