Skip to content

Commit a5c4d8b

Browse files
committed
feat(config): retrieval.rerank block + VLE_RETRIEVAL_RERANK_* env overrides
Adds ReRankBlock under RetrievalConfig with Enabled / Model / MaxContentChars (default 2000) / TopK (default 0 = keep all) fields, plus matching VLE_RETRIEVAL_RERANK_ENABLED, _MODEL, _MAX_CONTENT_CHARS, and _TOP_K env overrides following the same shape as the existing planning + agentic + answer_span env handling. Re-rank is opt-in: defaults are Enabled=false / TopK=0 so existing deployments behave identically. Per-request `enable_rerank` body field (wired in the next commit) overrides this block when set. Validate rejects negative MaxContentChars and TopK so a stray env or YAML typo can't silently flip the model into "send no content" mode. Zero is valid for both: TopK=0 means keep all candidates, MaxContentChars=0 falls through to the ReRanker's compiled default. Four new tests: default values, env override happy path, env disable of a YAML-set true, env override rejects bad numeric values, and validation of negative inputs.
1 parent 123be8d commit a5c4d8b

2 files changed

Lines changed: 183 additions & 0 deletions

File tree

pkg/config/config.go

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,44 @@ type RetrievalConfig struct {
204204
AnswerSpan AnswerSpanBlock `yaml:"answer_span"`
205205
Answer AnswerBlock `yaml:"answer"`
206206
Planning PlanningBlock `yaml:"planning"`
207+
ReRank ReRankBlock `yaml:"rerank"`
208+
}
209+
210+
// ReRankBlock configures the Phase 2.3 content-aware re-rank pass.
211+
//
212+
// When enabled, every /v1/query and /v1/answer request that returns
213+
// candidate sections runs one extra LLM call: the candidates' first
214+
// MaxContentChars chars of content are fed to the model with the
215+
// query, and the model returns a per-section relevance score
216+
// (0-100). Sections are reordered by score; if TopK > 0 the response
217+
// is truncated to the top K.
218+
//
219+
// The pass is opt-in. Per-request `enable_rerank` body field
220+
// overrides this block.
221+
//
222+
// Re-rank failures never drop sections — at worst the original
223+
// strategy order is preserved and the request returns unchanged from
224+
// the no-rerank path. See pkg/retrieval/rerank.go for the exact
225+
// degradation contract.
226+
type ReRankBlock struct {
227+
// Enabled toggles re-rank at the server level. Default: false.
228+
Enabled bool `yaml:"enabled"`
229+
230+
// Model overrides the re-rank LLM model. Empty means use the
231+
// request's model (which itself falls back to the engine default).
232+
// Point this at a small/fast model — the re-rank prompt is short
233+
// and shouldn't burn the flagship model's budget.
234+
Model string `yaml:"model"`
235+
236+
// MaxContentChars caps how many characters of each candidate's
237+
// content are sent to the model. Default: 2000.
238+
MaxContentChars int `yaml:"max_content_chars"`
239+
240+
// TopK caps the number of sections kept after re-ranking. 0 means
241+
// keep all candidates (re-rank only reorders). Useful when the
242+
// strategy is configured to return a wide candidate list and the
243+
// re-rank pass picks the focused top-k for synthesis.
244+
TopK int `yaml:"top_k"`
207245
}
208246

209247
// PlanningBlock configures Phase 2.1 query planning + Phase 2.2 multi-hop
@@ -368,6 +406,11 @@ func Default() Config {
368406
CacheSize: 128,
369407
Decompose: true,
370408
},
409+
ReRank: ReRankBlock{
410+
Enabled: false,
411+
MaxContentChars: 2000,
412+
TopK: 0,
413+
},
371414
},
372415
Ingest: IngestConfig{
373416
GlobalLLMConcurrency: 12,
@@ -556,6 +599,27 @@ func applyEnvOverrides(c *Config) {
556599
c.Retrieval.Planning.Decompose = false
557600
}
558601
}
602+
if v := os.Getenv("VLE_RETRIEVAL_RERANK_ENABLED"); v != "" {
603+
switch strings.ToLower(strings.TrimSpace(v)) {
604+
case "1", "true", "yes", "on":
605+
c.Retrieval.ReRank.Enabled = true
606+
case "0", "false", "no", "off":
607+
c.Retrieval.ReRank.Enabled = false
608+
}
609+
}
610+
if v := os.Getenv("VLE_RETRIEVAL_RERANK_MODEL"); v != "" {
611+
c.Retrieval.ReRank.Model = v
612+
}
613+
if v := os.Getenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS"); v != "" {
614+
if n, err := strconv.Atoi(v); err == nil && n > 0 {
615+
c.Retrieval.ReRank.MaxContentChars = n
616+
}
617+
}
618+
if v := os.Getenv("VLE_RETRIEVAL_RERANK_TOP_K"); v != "" {
619+
if n, err := strconv.Atoi(v); err == nil && n >= 0 {
620+
c.Retrieval.ReRank.TopK = n
621+
}
622+
}
559623
}
560624

561625
// Validate checks that required fields for the selected drivers are set.
@@ -642,5 +706,12 @@ func (c Config) Validate() error {
642706
return fmt.Errorf("retrieval.planning.cache_size must be >= 0, got %d", c.Retrieval.Planning.CacheSize)
643707
}
644708

709+
if c.Retrieval.ReRank.MaxContentChars < 0 {
710+
return fmt.Errorf("retrieval.rerank.max_content_chars must be >= 0, got %d", c.Retrieval.ReRank.MaxContentChars)
711+
}
712+
if c.Retrieval.ReRank.TopK < 0 {
713+
return fmt.Errorf("retrieval.rerank.top_k must be >= 0, got %d", c.Retrieval.ReRank.TopK)
714+
}
715+
645716
return nil
646717
}

pkg/config/config_test.go

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,123 @@ func TestDefaultValues(t *testing.T) {
4646
if !cfg.Retrieval.Planning.Decompose {
4747
t.Error("retrieval.planning.decompose should default to true (when planning is enabled)")
4848
}
49+
if cfg.Retrieval.ReRank.Enabled {
50+
t.Error("retrieval.rerank.enabled should default to false (opt-in)")
51+
}
52+
if cfg.Retrieval.ReRank.MaxContentChars != 2000 {
53+
t.Errorf("retrieval.rerank.max_content_chars = %d, want 2000", cfg.Retrieval.ReRank.MaxContentChars)
54+
}
55+
if cfg.Retrieval.ReRank.TopK != 0 {
56+
t.Errorf("retrieval.rerank.top_k = %d, want 0 (keep all)", cfg.Retrieval.ReRank.TopK)
57+
}
4958
if cfg.Log.Level != "info" {
5059
t.Errorf("log.level = %q, want info", cfg.Log.Level)
5160
}
5261
}
5362

63+
func TestReRankEnvOverride(t *testing.T) {
64+
// Not parallel — mutates env. Restore on exit.
65+
prevEnabled := os.Getenv("VLE_RETRIEVAL_RERANK_ENABLED")
66+
prevModel := os.Getenv("VLE_RETRIEVAL_RERANK_MODEL")
67+
prevMax := os.Getenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS")
68+
prevTopK := os.Getenv("VLE_RETRIEVAL_RERANK_TOP_K")
69+
defer func() {
70+
os.Setenv("VLE_RETRIEVAL_RERANK_ENABLED", prevEnabled)
71+
os.Setenv("VLE_RETRIEVAL_RERANK_MODEL", prevModel)
72+
os.Setenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS", prevMax)
73+
os.Setenv("VLE_RETRIEVAL_RERANK_TOP_K", prevTopK)
74+
}()
75+
76+
os.Setenv("VLE_RETRIEVAL_RERANK_ENABLED", "true")
77+
os.Setenv("VLE_RETRIEVAL_RERANK_MODEL", "gemini-2.0-flash")
78+
os.Setenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS", "1500")
79+
os.Setenv("VLE_RETRIEVAL_RERANK_TOP_K", "3")
80+
81+
cfg := Default()
82+
applyEnvOverrides(&cfg)
83+
84+
if !cfg.Retrieval.ReRank.Enabled {
85+
t.Error("VLE_RETRIEVAL_RERANK_ENABLED=true should enable rerank")
86+
}
87+
if cfg.Retrieval.ReRank.Model != "gemini-2.0-flash" {
88+
t.Errorf("rerank model = %q, want gemini-2.0-flash", cfg.Retrieval.ReRank.Model)
89+
}
90+
if cfg.Retrieval.ReRank.MaxContentChars != 1500 {
91+
t.Errorf("rerank max_content_chars = %d, want 1500", cfg.Retrieval.ReRank.MaxContentChars)
92+
}
93+
if cfg.Retrieval.ReRank.TopK != 3 {
94+
t.Errorf("rerank top_k = %d, want 3", cfg.Retrieval.ReRank.TopK)
95+
}
96+
}
97+
98+
func TestReRankEnvOverrideDisable(t *testing.T) {
99+
// Toggle off via env: start from a config that defaults to false,
100+
// then set =false explicitly; verify the path executes (not just
101+
// that the default value is preserved).
102+
prev := os.Getenv("VLE_RETRIEVAL_RERANK_ENABLED")
103+
defer os.Setenv("VLE_RETRIEVAL_RERANK_ENABLED", prev)
104+
105+
cfg := Default()
106+
cfg.Retrieval.ReRank.Enabled = true // simulate a YAML-set true
107+
os.Setenv("VLE_RETRIEVAL_RERANK_ENABLED", "false")
108+
applyEnvOverrides(&cfg)
109+
if cfg.Retrieval.ReRank.Enabled {
110+
t.Error("VLE_RETRIEVAL_RERANK_ENABLED=false should disable rerank even when YAML set it true")
111+
}
112+
}
113+
114+
func TestReRankEnvOverrideRejectsBad(t *testing.T) {
115+
// Garbage env values should be rejected, not silently zero the field.
116+
prevMax := os.Getenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS")
117+
prevTopK := os.Getenv("VLE_RETRIEVAL_RERANK_TOP_K")
118+
defer func() {
119+
os.Setenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS", prevMax)
120+
os.Setenv("VLE_RETRIEVAL_RERANK_TOP_K", prevTopK)
121+
}()
122+
123+
os.Setenv("VLE_RETRIEVAL_RERANK_MAX_CONTENT_CHARS", "not-a-number")
124+
os.Setenv("VLE_RETRIEVAL_RERANK_TOP_K", "abc")
125+
126+
cfg := Default()
127+
applyEnvOverrides(&cfg)
128+
if cfg.Retrieval.ReRank.MaxContentChars != 2000 {
129+
t.Errorf("bad max_content_chars env should preserve default, got %d", cfg.Retrieval.ReRank.MaxContentChars)
130+
}
131+
if cfg.Retrieval.ReRank.TopK != 0 {
132+
t.Errorf("bad top_k env should preserve default, got %d", cfg.Retrieval.ReRank.TopK)
133+
}
134+
}
135+
136+
func TestValidateReRankNegatives(t *testing.T) {
137+
t.Parallel()
138+
139+
// Negative max_content_chars rejected.
140+
cfg := Default()
141+
cfg.Database.URL = "postgres://localhost/test"
142+
cfg.Retrieval.ReRank.MaxContentChars = -1
143+
if err := cfg.Validate(); err == nil {
144+
t.Error("negative max_content_chars should fail validation")
145+
}
146+
147+
// Negative top_k rejected.
148+
cfg2 := Default()
149+
cfg2.Database.URL = "postgres://localhost/test"
150+
cfg2.Retrieval.ReRank.TopK = -1
151+
if err := cfg2.Validate(); err == nil {
152+
t.Error("negative top_k should fail validation")
153+
}
154+
155+
// Zero on both is valid (TopK=0 means "keep all", MaxContentChars=0
156+
// means "use default").
157+
cfg3 := Default()
158+
cfg3.Database.URL = "postgres://localhost/test"
159+
cfg3.Retrieval.ReRank.MaxContentChars = 0
160+
cfg3.Retrieval.ReRank.TopK = 0
161+
if err := cfg3.Validate(); err != nil {
162+
t.Errorf("zero rerank values should pass validation: %v", err)
163+
}
164+
}
165+
54166
func TestPlanningEnvOverride(t *testing.T) {
55167
// Not parallel — mutates env. Restore on exit.
56168
prevEnabled := os.Getenv("VLE_RETRIEVAL_PLANNING_ENABLED")

0 commit comments

Comments
 (0)