feat(engine): add LLM query-rewrite-before-retrieval strategy

Add a KeygenLLM interface (mirroring the ConsolidationLLM idiom that
keeps yaad free of any direct LLM dependency) and RewriteQuery /
Engine.RecallWithKeygen, which let a caller turn a natural-language
query into tight search keywords with an LLM before recall.

Three strategies: KeygenSynonyms (default, reuses the existing ExpandQuery
synonym map, no LLM), KeygenLLMThenSynonyms (union+dedup), and
KeygenLLMOnly. The step is best-effort: a nil LLM or any LLM error
degrades gracefully to local synonym expansion rather than failing the
recall.

Author: Patel230

engine/query_keygen.go

@@ -0,0 +1,102 @@
+package engine
+
+import (
+	"context"
+	"strings"
+)
+
+// KeygenLLM is the interface a caller implements to let yaad rewrite a natural-
+// language query into retrieval keywords with an LLM before searching. Like
+// ConsolidationLLM, it keeps yaad free of any direct LLM dependency — the engine
+// only ever sees this small interface.
+//
+// GenerateKeywords receives the user's raw query and returns a space-separated
+// (or newline-separated) set of keywords/phrases to search for. Returning an
+// empty string or an error signals "no rewrite"; the caller falls back to the
+// query unchanged (after synonym expansion).
+type KeygenLLM interface {
+	GenerateKeywords(ctx context.Context, query string) (string, error)
+}
+
+// KeygenStrategy controls how RewriteQuery turns a raw query into search terms.
+type KeygenStrategy int
+
+const (
+	// KeygenSynonyms uses only the built-in dependency-free synonym expansion
+	// (ExpandQuery). This is the default and requires no LLM.
+	KeygenSynonyms KeygenStrategy = iota
+	// KeygenLLMThenSynonyms asks the LLM to extract keywords first, then unions
+	// the result with synonym expansion of the original query. This mirrors the
+	// "generate keywords before retrieval" pattern: separate the LLM-driven
+	// keyword step from the retrieval step so a verbose question becomes tight
+	// search terms without the retriever ever seeing the prose.
+	KeygenLLMThenSynonyms
+	// KeygenLLMOnly uses only the LLM keyword output, falling back to synonym
+	// expansion if the LLM is absent or returns nothing.
+	KeygenLLMOnly
+)
+
+// RewriteQuery turns a raw user query into the search string used for recall,
+// applying the given strategy. llm may be nil; when it is, the function behaves
+// exactly as KeygenSynonyms regardless of strategy, so callers can wire an LLM
+// optionally without branching. The returned string is always non-empty as long
+// as query is non-empty.
+//
+// Defaults here are yaad's own — deliberately not borrowed from any external
+// framework's RAG config — and the whole step is best-effort: any LLM failure
+// degrades gracefully to local synonym expansion rather than erroring out of a
+// recall.
+func RewriteQuery(ctx context.Context, llm KeygenLLM, strategy KeygenStrategy, query string) string {
+	q := strings.TrimSpace(query)
+	if q == "" {
+		return ""
+	}
+
+	synonyms := ExpandQuery(q)
+
+	if llm == nil || strategy == KeygenSynonyms {
+		return synonyms
+	}
+
+	kw, err := llm.GenerateKeywords(ctx, q)
+	kw = strings.TrimSpace(kw)
+	if err != nil || kw == "" {
+		// Best-effort: fall back to local expansion.
+		return synonyms
+	}
+
+	if strategy == KeygenLLMOnly {
+		return kw
+	}
+
+	// KeygenLLMThenSynonyms: union the LLM keywords with synonym expansion,
+	// de-duplicated, preserving first-seen order (LLM terms first).
+	return unionTerms(kw, synonyms)
+}
+
+// RecallWithKeygen rewrites opts.Query via RewriteQuery and then runs Recall on
+// the rewritten query, leaving every other RecallOpts field untouched. It is a
+// thin convenience over Recall for callers that have an LLM available; callers
+// without one can keep using Recall directly.
+func (e *Engine) RecallWithKeygen(ctx context.Context, llm KeygenLLM, strategy KeygenStrategy, opts RecallOpts) (*RecallResult, error) {
+	if opts.Query != "" {
+		opts.Query = RewriteQuery(ctx, llm, strategy, opts.Query)
+	}
+	return e.Recall(ctx, opts)
+}
+
+// unionTerms merges two whitespace-separated term lists, lower-casing for the
+// dedup key but emitting the first-seen surface form, preserving order.
+func unionTerms(a, b string) string {
+	seen := make(map[string]bool)
+	var out []string
+	for _, field := range strings.Fields(a + " " + b) {
+		key := strings.ToLower(field)
+		if seen[key] {
+			continue
+		}
+		seen[key] = true
+		out = append(out, field)
+	}
+	return strings.Join(out, " ")
+}

engine/query_keygen_test.go

@@ -0,0 +1,75 @@
+package engine
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"testing"
+)
+
+type fakeKeygen struct {
+	out string
+	err error
+}
+
+func (f fakeKeygen) GenerateKeywords(_ context.Context, _ string) (string, error) {
+	return f.out, f.err
+}
+
+func TestRewriteQuery_NilLLMUsesSynonyms(t *testing.T) {
+	got := RewriteQuery(context.Background(), nil, KeygenLLMThenSynonyms, "auth bug")
+	// ExpandQuery expands "auth" → adds authentication/authorize/login.
+	if !strings.Contains(got, "auth") || !strings.Contains(got, "authentication") {
+		t.Errorf("expected synonym expansion, got %q", got)
+	}
+}
+
+func TestRewriteQuery_EmptyQuery(t *testing.T) {
+	if got := RewriteQuery(context.Background(), fakeKeygen{out: "x"}, KeygenLLMOnly, "   "); got != "" {
+		t.Errorf("empty query should rewrite to empty, got %q", got)
+	}
+}
+
+func TestRewriteQuery_LLMOnly(t *testing.T) {
+	got := RewriteQuery(context.Background(), fakeKeygen{out: "login session cookie"}, KeygenLLMOnly, "why am I logged out?")
+	if got != "login session cookie" {
+		t.Errorf("LLMOnly should return raw keywords, got %q", got)
+	}
+}
+
+func TestRewriteQuery_LLMErrorFallsBack(t *testing.T) {
+	got := RewriteQuery(context.Background(), fakeKeygen{err: errors.New("boom")}, KeygenLLMOnly, "db config")
+	// Falls back to synonym expansion of the original query.
+	if !strings.Contains(got, "database") || !strings.Contains(got, "configuration") {
+		t.Errorf("expected fallback to synonyms, got %q", got)
+	}
+}
+
+func TestRewriteQuery_LLMThenSynonymsUnionDedup(t *testing.T) {
+	// LLM returns "auth token"; synonym expansion of "auth" also yields auth/...
+	got := RewriteQuery(context.Background(), fakeKeygen{out: "auth token"}, KeygenLLMThenSynonyms, "auth")
+	fields := strings.Fields(got)
+	seen := map[string]int{}
+	for _, f := range fields {
+		seen[strings.ToLower(f)]++
+	}
+	for term, n := range seen {
+		if n > 1 {
+			t.Errorf("term %q appears %d times, want deduped: %q", term, n, got)
+		}
+	}
+	// LLM terms come first.
+	if fields[0] != "auth" {
+		t.Errorf("expected LLM terms first, got %q", got)
+	}
+	if !strings.Contains(got, "token") {
+		t.Errorf("expected LLM keyword 'token' retained, got %q", got)
+	}
+}
+
+func TestRewriteQuery_SynonymsStrategyIgnoresLLM(t *testing.T) {
+	got := RewriteQuery(context.Background(), fakeKeygen{out: "SHOULD_NOT_APPEAR"}, KeygenSynonyms, "config")
+	if strings.Contains(got, "SHOULD_NOT_APPEAR") {
+		t.Errorf("KeygenSynonyms must not call the LLM, got %q", got)
+	}
+}