Maven/internal/kiwix/rewrite_test.go

package kiwix

import (
	"context"
	"testing"

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

// Bad model output must never reach Kiwix. No model needed for this.
func TestCleanQueryRejectsJunk(t *testing.T) {
	bad := []struct{ name, raw string }{
		{"empty", ""},
		{"blank", "   \n  "},
		{"russian came back", "почему небо синее"},
		{"mixed russian", "sky синее scattering"},
		{"full sentence", "The sky looks blue because of the scattering of sunlight by air molecules"},
		{"prose with quotes", `Sure! Here is a good search query: "Rayleigh scattering", which explains it.`},
	}
	for _, c := range bad {
		if got, err := CleanQuery(c.raw); err == nil {
			t.Errorf("%s: want rejection, got %q", c.name, got)
		}
	}
}

func TestCleanQueryCleans(t *testing.T) {
	ok := []struct{ raw, want string }{
		{"Rayleigh scattering sky", "Rayleigh scattering sky"},
		{"  boiled   egg  cooking \n", "boiled egg cooking"},
		{`"virtual private network"`, "virtual private network"},
		{"solid-state drive", "solid-state drive"},
		{"cat purr.", "cat purr"},
		{"hiccup\nAlso: hiccough", "hiccup"},
		// Question words are filler to a keyword ranker, so they go.
		{"why is the sky blue", "sky blue"},
		{"how much water to drink daily", "water drink daily"},
		{"SSD vs HDD comparison", "SSD HDD comparison"},
		// Nothing but filler: keep it rather than return nothing.
		{"what is it", "what is it"},
	}
	for _, c := range ok {
		got, err := CleanQuery(c.raw)
		if err != nil {
			t.Errorf("%q: %v", c.raw, err)
			continue
		}
		if got != c.want {
			t.Errorf("%q -> %q, want %q", c.raw, got, c.want)
		}
	}
}

type fakeCompleter struct {
	out string
	req llm.Req
}

func (f *fakeCompleter) Complete(_ context.Context, r llm.Req) (string, error) {
	f.req = r
	return f.out, nil
}

func TestRewriteConstrainsTheCall(t *testing.T) {
	f := &fakeCompleter{out: `{"query":"Rayleigh scattering sky"}`}
	got, err := NewRewriter(f).Rewrite(context.Background(), "почему небо синее?")
	if err != nil {
		t.Fatalf("rewrite: %v", err)
	}
	if got != "Rayleigh scattering sky" {
		t.Errorf("query = %q", got)
	}
	if f.req.Grammar == "" {
		t.Error("no grammar sent")
	}
	if f.req.MaxTokens == 0 || f.req.MaxTokens > 32 {
		t.Errorf("max_tokens = %d, want a small cap", f.req.MaxTokens)
	}
}

func TestRewriteRejectsBadModelOutput(t *testing.T) {
	bad := []string{
		`{"query":"почему небо синее"}`, // never translated
		`{"query":""}`, // empty
		`{"query":"the sky is blue because sunlight is scattered by air"}`, // an answer
		// Note: a SHORT English prose fragment ("Let me analyze this request")
		// is under the word cap and cannot be caught here. The grammar is what
		// stops that one.
	}
	for _, out := range bad {
		f := &fakeCompleter{out: out}
		if got, err := NewRewriter(f).Rewrite(context.Background(), "почему небо синее?"); err == nil {
			t.Errorf("%s: want rejection, got %q", out, got)
		}
	}
}

// A reply that is not the JSON shape but is still usable keywords should pass.
func TestRewriteFallsBackToPlainText(t *testing.T) {
	f := &fakeCompleter{out: "Rayleigh scattering sky"}
	got, err := NewRewriter(f).Rewrite(context.Background(), "почему небо синее?")
	if err != nil || got != "Rayleigh scattering sky" {
		t.Errorf("got %q, %v", got, err)
	}
}