Maven/internal/kiwix/rewrite.go

package kiwix

// Turning a Russian question into an English Kiwix search.
//
// Kiwix ranks by keyword, not by meaning. "why is the sky blue" returns a TV
// episode; "Rayleigh scattering sky blue" returns the right article. So the
// model's job here is NOT translation — it is naming the English article the
// answer lives in.
//
// The output space is a handful of words, so it is worth locking down hard: a
// GBNF grammar for the shape, a tiny token cap, and a cleanup pass that throws
// away anything odd rather than handing junk to Kiwix.

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"unicode"

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

// Completer — the LLM seam, so tests can fake it. *llm.Client satisfies it.
type Completer interface {
	Complete(ctx context.Context, r llm.Req) (string, error)
}

// queryGrammar — one JSON object holding 1..6 keyword words. Latin letters,
// digits and hyphens only, so the model physically cannot answer the question
// or reply in Russian.
//
// Why the JSON wrapper: this model always thinks out loud and this llama-server
// build ignores the thinking switch (see docs/evals/2026-07-31-routing.md). A bare
// word-list grammar just captured the reasoning — every case came back as
// "Let me analyze this request carefully". Demanding JSON, like routeGrammar and
// responseGrammar already do, gives the reasoning nowhere to go.
const queryGrammar = `
root ::= "{" ws "\"query\"" ws ":" ws "\"" word (" " word){0,5} "\"" ws "}"
word ::= [A-Za-z0-9] [A-Za-z0-9-]{0,23}
ws   ::= [ \t\n]*
`

// rewriteSystem — asks for search keywords, not an answer and not a translation.
const rewriteSystem = `You turn a question into a search query for English Wikipedia.

Rules:
- Output ONLY English search keywords. Never an answer, never an explanation.
- Do NOT translate the sentence. Name the thing the answer is about.
- The output must be a noun phrase, like a Wikipedia article title.
- Never use question words: no why, how, what, when, which, "how much",
  "how long", "how to", "vs", "reason", "difference".
- 2 to 4 words.

Reply with JSON: {"query":"<keywords>"}

Good:
"почему листья желтеют осенью?" -> {"query":"leaf senescence autumn"}
"как работает микроволновка?" -> {"query":"microwave oven"}
"не могли бы вы объяснить, что такое блокчейн?" -> {"query":"blockchain"}
"сколько живут собаки?" -> {"query":"dog lifespan"}
"как избавиться от комаров в квартире?" -> {"query":"mosquito control"}
"чем чай отличается от кофе?" -> {"query":"tea"}

Only JSON, no explanation.`

// maxQueryTokens — the output is a few words plus the JSON wrapper. A tight cap
// is the cheapest guard against the model rambling into an answer.
const maxQueryTokens = 32

// Rewriter asks the resident model for English search keywords.
type Rewriter struct{ c Completer }

func NewRewriter(c Completer) *Rewriter { return &Rewriter{c: c} }

// Rewrite returns English keywords for a question in any language.
// It errors rather than returning something Kiwix should not see.
func (r *Rewriter) Rewrite(ctx context.Context, question string) (string, error) {
	raw, err := r.c.Complete(ctx, llm.Req{
		System:        rewriteSystem,
		User:          strings.TrimSpace(question),
		Grammar:       queryGrammar,
		MaxTokens:     maxQueryTokens,
		RepeatPenalty: 1.15,
	})
	if err != nil {
		return "", err
	}
	return CleanQuery(unwrapJSON(raw))
}

// unwrapJSON pulls the query out of {"query":"..."}. If the reply is not that
// shape it is returned as-is, and CleanQuery decides whether it is usable.
func unwrapJSON(raw string) string {
	s := strings.TrimSpace(raw)
	if !strings.HasPrefix(s, "{") {
		return s
	}
	var got struct{ Query string }
	if err := json.Unmarshal([]byte(s), &got); err != nil {
		return s
	}
	return got.Query
}

// maxQueryWords matches the grammar's bound. Anything longer is prose.
const maxQueryWords = 6

// CleanQuery checks and tidies whatever the model produced. The grammar makes
// bad output unlikely, not impossible (a server without grammar support, a
// different model), so this is the real gate in front of Kiwix.
//
// Exported so it can be tested without a model.
func CleanQuery(raw string) (string, error) {
	s := strings.TrimSpace(raw)
	// Models like to wrap answers in quotes. Drop surrounding ones.
	s = strings.Trim(s, "\"'`")
	// Keep the first line only: everything after it is prose.
	if i := strings.IndexAny(s, "\r\n"); i >= 0 {
		s = s[:i]
	}
	// Keep letters, digits, spaces and hyphens; anything else becomes a space.
	var b strings.Builder
	for _, ru := range s {
		switch {
		case unicode.IsLetter(ru) || unicode.IsDigit(ru) || ru == '-':
			b.WriteRune(ru)
		default:
			b.WriteRune(' ')
		}
	}
	words := strings.Fields(b.String())
	if len(words) == 0 {
		return "", fmt.Errorf("kiwix rewrite: empty query")
	}
	if len(words) > maxQueryWords {
		return "", fmt.Errorf("kiwix rewrite: %d words, want at most %d (looks like prose)", len(words), maxQueryWords)
	}
	words = dropStopWords(words)
	out := strings.Join(words, " ")
	// The ZIMs are English. Non-Latin letters mean the model ignored the ask.
	for _, ru := range out {
		if unicode.IsLetter(ru) && !isLatin(ru) {
			return "", fmt.Errorf("kiwix rewrite: query is not English: %q", out)
		}
	}
	return out, nil
}

// stopWords — question words and filler. The model keeps writing question-shaped
// queries ("why is the sky blue", "how much water to drink daily") no matter how
// the prompt is worded, and Kiwix ranks on every word, so those words drag in
// song and episode titles. Dropping them in code is not a style preference: a
// keyword ranker gets nothing from them.
var stopWords = map[string]bool{
	"a": true, "an": true, "the": true, "is": true, "are": true, "was": true,
	"do": true, "does": true, "did": true, "to": true, "of": true, "in": true,
	"on": true, "for": true, "and": true, "or": true, "my": true, "me": true,
	"i": true, "it": true, "its": true, "be": true, "been": true, "get": true,
	"how": true, "why": true, "what": true, "when": true, "which": true,
	"who": true, "where": true, "much": true, "many": true, "long": true,
	"vs": true, "than": true, "rid": true, "from": true, "about": true,
}

// dropStopWords removes filler, but never everything: if the query was nothing
// but stop words there is nothing better to search, so the original is kept and
// the caller sees whatever Kiwix makes of it.
func dropStopWords(words []string) []string {
	kept := make([]string, 0, len(words))
	for _, w := range words {
		if !stopWords[strings.ToLower(w)] {
			kept = append(kept, w)
		}
	}
	if len(kept) == 0 {
		return words
	}
	return kept
}

func isLatin(ru rune) bool {
	return (ru >= 'a' && ru <= 'z') || (ru >= 'A' && ru <= 'Z')
}