Maven/internal/router/llmrouter_test.go

package router

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/kami/maven/internal/llm"
)

type mockLLM struct {
	out string
	err error
	got *llm.Req // last request, when the test wants to inspect it
}

func (m mockLLM) Complete(_ context.Context, r llm.Req) (string, error) {
	if m.got != nil {
		*m.got = r
	}
	return m.out, m.err
}

// Without a repeat penalty the model loops inside the text field until MaxTokens
// and the truncated JSON fails to parse.
func TestLLMRouterSetsRepeatPenalty(t *testing.T) {
	var got llm.Req
	lr := NewLLMRouter(mockLLM{out: `{"intent":"chat","text":"привет"}`, got: &got})
	if _, _, err := lr.Route(context.Background(), "привет", time.Now()); err != nil {
		t.Fatalf("route: %v", err)
	}
	if got.RepeatPenalty <= 1 {
		t.Fatalf("want repeat penalty above 1, got %v", got.RepeatPenalty)
	}
}

// An unbounded string rule lets one field eat the whole token budget.
func TestRouteGrammarBoundsStrings(t *testing.T) {
	if !strings.Contains(routeGrammar, `string ::= "\"" ([^"\\] | "\\" .){0,120} "\""`) {
		t.Fatal("grammar string rule lost its length bound")
	}
}

// A question naming a fact key used to be stored as a fact because the fact rule
// was tested first. Keep the query rule above it.
func TestRoutePromptTestsQueryBeforeFact(t *testing.T) {
	query := strings.Index(routeSystem, "→ query")
	fact := strings.Index(routeSystem, "состояние/событие → fact")
	if query < 0 || fact < 0 {
		t.Fatalf("prompt lost a rule: query=%d fact=%d", query, fact)
	}
	if query > fact {
		t.Fatal("query rule must come before the fact rule")
	}
	if !strings.Contains(routeSystem, "Задаёт вопрос") {
		t.Fatal("prompt lost the explicit question test")
	}
}

func TestLLMRouterFactMapping(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: `{"intent":"fact","key":"water","value":"выпил"}`})
	d, ok, err := lr.Route(context.Background(), "я выпил воду", time.Now())
	if err != nil || !ok {
		t.Fatalf("ok=%v err=%v", ok, err)
	}
	if d.Intent != IntentFact || d.Slots.Key != "water" || !d.Slots.HasKey {
		t.Fatalf("bad decision %+v", d)
	}
}

func TestLLMRouterNoteMapping(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: `{"intent":"note","text":"кофе закончился"}`})
	d, ok, err := lr.Route(context.Background(), "запомни что кофе закончился", time.Now())
	if err != nil || !ok {
		t.Fatalf("ok=%v err=%v", ok, err)
	}
	if d.Intent != IntentNote || d.Slots.Text != "кофе закончился" {
		t.Fatalf("bad decision %+v", d)
	}
}

func TestLLMRouterBadJSONFallsBack(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: `garbage`})
	_, ok, err := lr.Route(context.Background(), "x", time.Now())
	if ok || err == nil {
		t.Fatal("want ok=false, err!=nil on bad json")
	}
}

func TestLLMRouterReminderMapping(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: `{"intent":"reminder","text":"позвонить маме"}`})
	d, ok, err := lr.Route(context.Background(), "напомни позвонить маме", time.Now())
	if err != nil || !ok {
		t.Fatalf("ok=%v err=%v", ok, err)
	}
	if d.Intent != IntentReminder || d.Slots.Text != "позвонить маме" {
		t.Fatalf("bad decision %+v", d)
	}
}

// An intent name that is not in the contract at all (as opposed to "unknown",
// which is a real refusal) still defaults to chat.
func TestLLMRouterChatFallback(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: `{"intent":"banana"}`})
	d, ok, err := lr.Route(context.Background(), "как дела?", time.Now())
	if err != nil || !ok {
		t.Fatalf("ok=%v err=%v", ok, err)
	}
	if d.Intent != IntentChat {
		t.Fatalf("unknown intent should default to chat, got %s", d.Intent)
	}
}

// The model must be able to say "I could not route this".
func TestRouteGrammarAllowsUnknown(t *testing.T) {
	if !strings.Contains(routeGrammar, `"\"unknown\""`) {
		t.Fatal("grammar cannot express a refusal")
	}
}

// If the prompt does not tell the model when to refuse, it never will.
func TestRoutePromptExplainsUnknown(t *testing.T) {
	if !strings.Contains(routeSystem, "unknown") {
		t.Fatal("prompt never mentions the unknown intent")
	}
	if !strings.Contains(routeSystem, `"сделай это" → {"intent":"unknown"}`) {
		t.Fatal("prompt lost its worked refusal example")
	}
	// A refusal-only router is useless, so the prompt must also show cases that
	// look ambiguous but are not.
	if !strings.Contains(routeSystem, "здесь unknown не нужен") {
		t.Fatal("prompt lost its counter-examples")
	}
}

// A refusal is not an error. It reports "no decision" so the cascade moves on.
func TestLLMRouterUnknownRefuses(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: `{"intent":"unknown"}`})
	_, ok, err := lr.Route(context.Background(), "сделай это", time.Now())
	if ok {
		t.Fatal("a refusal must not produce a usable decision")
	}
	if err != nil {
		t.Fatalf("a refusal is not an error, got %v", err)
	}
}

// The whole point of the refusal: the turn keeps going on the classifier, the
// same way it does when the model returns garbage.
func TestRouterFallsBackWhenLLMRefuses(t *testing.T) {
	c := NewClassifier(NewHashEmbedder(1024))
	seedClassifier(t, c)
	r := New(Config{
		Classifier: c,
		Extractor:  Extractor{Time: StubDateTimeParser{}, Facts: DefaultFactParser{}},
		Threshold:  0.4,
		LLM:        NewLLMRouter(mockLLM{out: `{"intent":"unknown"}`}),
	})
	d, err := r.Route(context.Background(), "напомни позвонить маме", refNow())
	if err != nil {
		t.Fatalf("route: %v", err)
	}
	// Stage 1 is the LLM's own answer; the classifier lands on stage 2 or 3.
	if d.Stage < 2 {
		t.Fatalf("want the classifier to decide, got stage %d (%+v)", d.Stage, d)
	}
}

func TestLLMRouterLLMError(t *testing.T) {
	lr := NewLLMRouter(mockLLM{out: "", err: fmt.Errorf("llm down")})
	_, ok, err := lr.Route(context.Background(), "x", time.Now())
	if ok || err == nil {
		t.Fatal("want ok=false, err!=nil on llm error")
	}
}

// --- slot extraction on top of an LLM decision --------------------------------

// newLLMTestRouter — a router whose route always comes from the mock model.
func newLLMTestRouter(t *testing.T, out string) *Router {
	t.Helper()
	c := NewClassifier(NewHashEmbedder(1024))
	seedClassifier(t, c)
	acts := DefaultActMatcher{Fns: []string{"restart", "stop", "run", "backup"}}
	return New(Config{
		Classifier: c,
		Extractor:  Extractor{Time: StubDateTimeParser{}, Acts: acts, Facts: DefaultFactParser{}},
		Threshold:  0.4,
		LLM:        NewLLMRouter(mockLLM{out: out}),
	})
}

// The model cannot produce a fire time, so without extraction every LLM-routed
// reminder was dropped as "no time".
func TestLLMDecisionGetsReminderTime(t *testing.T) {
	r := newLLMTestRouter(t, `{"intent":"reminder","text":"позвонить маме"}`)
	d, err := r.Route(context.Background(), "напомни позвонить маме через 2 часа", refNow())
	if err != nil {
		t.Fatalf("route: %v", err)
	}
	if d.Intent != IntentReminder {
		t.Fatalf("want reminder, got %v", d.Intent)
	}
	if !d.Slots.HasTime || !d.Slots.Time.Equal(refNow().Add(2*time.Hour)) {
		t.Fatalf("want time now+2h, got %+v", d.Slots)
	}
	if d.Slots.Text != "позвонить маме" {
		t.Fatalf("extraction overwrote the model's text: %q", d.Slots.Text)
	}
}

// No time in the utterance ⇒ no time in the slots. Do not invent one; the
// daemon says it could not read the time.
func TestLLMReminderWithoutTimeStaysEmpty(t *testing.T) {
	r := newLLMTestRouter(t, `{"intent":"reminder","text":"позвонить маме"}`)
	d, err := r.Route(context.Background(), "напомни позвонить маме", refNow())
	if err != nil {
		t.Fatalf("route: %v", err)
	}
	if d.Slots.HasTime {
		t.Fatalf("invented a time: %v", d.Slots.Time)
	}
}

// An act decision arrived with no Fn, so the tool never ran.
func TestLLMDecisionGetsActFn(t *testing.T) {
	r := newLLMTestRouter(t, `{"intent":"act","verb":"restart nginx"}`)
	d, err := r.Route(context.Background(), "слушай, restart nginx пожалуйста", refNow())
	if err != nil {
		t.Fatalf("route: %v", err)
	}
	if !d.Slots.HasFn || d.Slots.Fn != "restart" || len(d.Slots.Args) != 1 || d.Slots.Args[0] != "nginx" {
		t.Fatalf("want fn=restart args=[nginx], got %+v", d.Slots)
	}
}

// The model's own slots win; extraction only fills gaps.
func TestLLMSlotsWinOverExtraction(t *testing.T) {
	r := newLLMTestRouter(t, `{"intent":"fact","key":"hydration","value":"выпил"}`)
	d, err := r.Route(context.Background(), "я выпил воду", refNow())
	if err != nil {
		t.Fatalf("route: %v", err)
	}
	if d.Slots.Key != "hydration" {
		t.Fatalf("extraction overwrote the model's key: %q", d.Slots.Key)
	}
}

// A fact the model left keyless still gets one from the parser.
func TestLLMFactGetsKeyFromParser(t *testing.T) {
	r := newLLMTestRouter(t, `{"intent":"fact","text":"я выпил воду"}`)
	d, err := r.Route(context.Background(), "я выпил воду", refNow())
	if err != nil {
		t.Fatalf("route: %v", err)
	}
	if !d.Slots.HasKey || d.Slots.Key != "water" {
		t.Fatalf("want key=water, got %+v", d.Slots)
	}
}