Skip to content

Commit bb32d4e

Browse files
committed
feat(api,config): wire planner + decomposer into /v1/query + /v1/answer
Server-side opt-in for Phase 2.1 + 2.2. New PlanningBlock under retrieval (enabled, model, cache_size, decompose; env: VLE_RETRIEVAL_PLANNING_*). Default disabled at both config and per-request levels, so existing callers see no behaviour change. Wiring: - api.Deps gains Planner + Planning fields. The body-level enable_planning field (pointer-bool to disambiguate absent from explicit-false) overrides the config block. - handleQuery / handleAnswer route through a small set of helpers (runPlanner, runSelection, runSelectionWithUsage) that fold the Planner output into selection. Multi-hop plans go through a Decomposer wrapping the active Strategy when planning.decompose is true. - /v1/answer's synthesis prompt grows a short "Planner notes" block (intent, entities, expected doc areas, sub-questions) when a plan is present, so the model reasons with the same understanding the retrieval pipeline used. - Both endpoints surface the sanitised Plan in the response under "plan" (omitempty) when planning ran. - cmd/engine instantiates a Planner whenever LLM is configured, so per-request opt-in still works even with planning.enabled=false. - Planner transport errors are LOGGED but not propagated — a planner blip cannot 500 an otherwise-working retrieval request. OpenAPI: Plan schema added; QueryRequest/AnswerRequest get enable_planning; QueryResponse/AnswerResponse get plan ($ref Plan, omitempty). config.example.yaml gets a documented retrieval.planning block. Tests: planning defaults + env-override coverage added to pkg/config/config_test.go. All existing tests still pass.
1 parent 71c4aac commit bb32d4e

6 files changed

Lines changed: 392 additions & 27 deletions

File tree

cmd/engine/main.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,26 @@ func run() error {
109109
// Multi-document query dispatcher.
110110
multiDoc := retrieval.NewMultiDoc(strategy, pool.LoadTree)
111111

112+
// Planner: opt-in Phase 2.1. When disabled at boot we still
113+
// instantiate it lazily — the per-request `enable_planning` body
114+
// field overrides the config, so a server with planning.enabled=false
115+
// but a Planner configured can still serve opt-in callers.
116+
var planner *retrieval.Planner
117+
if llmClient != nil {
118+
plannerModel := cfg.Retrieval.Planning.Model
119+
if plannerModel == "" {
120+
plannerModel = modelFor(cfg.LLM)
121+
}
122+
planner = retrieval.NewPlannerWithCacheSize(llmClient, plannerModel, cfg.Retrieval.Planning.CacheSize)
123+
if cfg.Retrieval.Planning.Enabled {
124+
logger.Info("retrieval: planner enabled",
125+
"model", plannerModel,
126+
"cache_size", cfg.Retrieval.Planning.CacheSize,
127+
"decompose", cfg.Retrieval.Planning.Decompose,
128+
)
129+
}
130+
}
131+
112132
pipeline := ingest.NewPipeline(ingest.Pipeline{
113133
DB: pool,
114134
Storage: store,
@@ -135,6 +155,8 @@ func run() error {
135155
LLMModel: modelFor(cfg.LLM),
136156
AnswerSpan: cfg.Retrieval.AnswerSpan,
137157
Answer: cfg.Retrieval.Answer,
158+
Planner: planner,
159+
Planning: cfg.Retrieval.Planning,
138160
}
139161

140162
srv := &http.Server{

config.example.yaml

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,29 @@ retrieval:
129129
max_sections: 5
130130
max_answer_tokens: 1024
131131

132+
# planning: Phase 2.1 query planning + Phase 2.2 multi-hop decomposition.
133+
# When enabled, every /v1/query and /v1/answer request first issues a
134+
# short LLM call that returns a structured Plan (intent, entities,
135+
# expected document areas, multi-hop flag, sub-questions). Multi-hop
136+
# plans fan retrieval out one selection call per sub-question and
137+
# union the results.
138+
#
139+
# OPT-IN. Default disabled. Per-request `enable_planning` body field
140+
# overrides this block, so callers can experiment without a restart.
141+
# Plans are cached in a per-process LRU keyed on (query, model);
142+
# repeated questions don't burn extra LLM budget.
143+
planning:
144+
enabled: false
145+
# Override the planner's model; empty inherits the engine's
146+
# configured default. Point this at a small/fast model — planning
147+
# is a short prompt that shouldn't run on the flagship model.
148+
model: ""
149+
cache_size: 128
150+
# decompose: when planning runs, multi-hop plans fan retrieval
151+
# out per sub-question. Set false to validate the planner in
152+
# isolation (plan returned, but retrieval uses the original query).
153+
decompose: true
154+
132155
ingest:
133156
# The summarize and HyDE stages run concurrently. This caps the total
134157
# number of LLM calls in flight across both stages combined, so the

internal/api/server.go

Lines changed: 165 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,16 @@ type Deps struct {
6363
// values (AnswerSpan disabled, Answer.MaxSections=5) are safe.
6464
AnswerSpan config.AnswerSpanBlock
6565
Answer config.AnswerBlock
66+
67+
// Planner runs one LLM call before retrieval to build a structured
68+
// Plan (intent + entities + multi-hop sub-questions). Nil disables
69+
// planning even when a request opts in via `enable_planning`.
70+
Planner *retrieval.Planner
71+
72+
// Planning carries the server-side planning config. The body-level
73+
// `enable_planning` field on /v1/query and /v1/answer overrides
74+
// Planning.Enabled.
75+
Planning config.PlanningBlock
6676
}
6777

6878
// Router builds and returns the chi router wired with v1 routes.
@@ -368,8 +378,15 @@ func (d Deps) handleGetSection(w http.ResponseWriter, r *http.Request) {
368378
// --- query ---
369379

370380
// handleQuery accepts { document_id, query, model?, max_tokens?,
371-
// reserved_for_prompt?, max_parallel_calls?, max_sections? } and runs the
372-
// configured retrieval.Strategy against the document's tree.
381+
// reserved_for_prompt?, max_parallel_calls?, max_sections?,
382+
// enable_planning? } and runs the configured retrieval.Strategy against
383+
// the document's tree.
384+
//
385+
// When `enable_planning` is true (or `retrieval.planning.enabled` is on
386+
// at config level) the request first issues a planning LLM call. The
387+
// resulting Plan is surfaced in the response under "plan". If the plan
388+
// is multi-hop and decomposition is enabled, retrieval fans out one
389+
// strategy call per sub-question and unions the results.
373390
func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
374391
var body struct {
375392
DocumentID tree.DocumentID `json:"document_id"`
@@ -379,6 +396,10 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
379396
ReservedForPrompt int `json:"reserved_for_prompt"`
380397
MaxParallelCalls int `json:"max_parallel_calls"`
381398
MaxSections int `json:"max_sections"`
399+
// EnablePlanning opts this request into the Phase 2.1 query
400+
// planner. A pointer so we can distinguish "absent" from
401+
// "explicit false" — absent falls back to the server config.
402+
EnablePlanning *bool `json:"enable_planning"`
382403
}
383404
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
384405
writeErr(w, http.StatusBadRequest, "invalid json: "+err.Error())
@@ -420,7 +441,9 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
420441
}
421442

422443
started := time.Now()
423-
ids, err := d.Strategy.Select(r.Context(), t, body.Query, budget)
444+
445+
plan, _ := d.runPlanner(r.Context(), body.Query, body.EnablePlanning)
446+
ids, err := d.runSelection(r.Context(), t, plan, body.Query, budget)
424447
if err != nil {
425448
d.Logger.Error("query: strategy failed", "err", err, "document_id", body.DocumentID)
426449
writeErr(w, http.StatusInternalServerError, "retrieval failed: "+err.Error())
@@ -461,14 +484,18 @@ func (d Deps) handleQuery(w http.ResponseWriter, r *http.Request) {
461484
sections = append(sections, sectionWithContentToMap(e))
462485
}
463486

464-
writeJSON(w, http.StatusOK, map[string]any{
487+
resp := map[string]any{
465488
"document_id": body.DocumentID,
466489
"query": body.Query,
467490
"strategy": d.Strategy.Name(),
468491
"model": body.Model,
469492
"sections": sections,
470493
"elapsed_ms": time.Since(started).Milliseconds(),
471-
})
494+
}
495+
if plan != nil {
496+
resp["plan"] = plan
497+
}
498+
writeJSON(w, http.StatusOK, resp)
472499
}
473500

474501
// sectionWithContent bundles a tree section with its loaded content
@@ -590,6 +617,9 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
590617
MaxParallelCalls int `json:"max_parallel_calls"`
591618
MaxSections int `json:"max_sections"`
592619
MaxAnswerTokens int `json:"max_answer_tokens"`
620+
// EnablePlanning opts this request into the Phase 2.1 query
621+
// planner. See handleQuery for the same field's semantics.
622+
EnablePlanning *bool `json:"enable_planning"`
593623
}
594624
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
595625
writeErr(w, http.StatusBadRequest, "invalid json: "+err.Error())
@@ -629,22 +659,13 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
629659
started := time.Now()
630660
totalUsage := retrieval.Usage{}
631661

632-
var ids []tree.SectionID
633-
var retrievalUsage retrieval.Usage
634-
if cs, ok := d.Strategy.(retrieval.CostStrategy); ok {
635-
res, err := cs.SelectWithCost(r.Context(), t, body.Query, budget)
636-
if err != nil {
637-
writeErr(w, http.StatusInternalServerError, "retrieval failed: "+err.Error())
638-
return
639-
}
640-
ids, retrievalUsage = res.SelectedIDs, res.Usage
641-
} else {
642-
picks, err := d.Strategy.Select(r.Context(), t, body.Query, budget)
643-
if err != nil {
644-
writeErr(w, http.StatusInternalServerError, "retrieval failed: "+err.Error())
645-
return
646-
}
647-
ids = picks
662+
plan, planUsage := d.runPlanner(r.Context(), body.Query, body.EnablePlanning)
663+
totalUsage.Add(planUsage)
664+
665+
ids, retrievalUsage, err := d.runSelectionWithUsage(r.Context(), t, plan, body.Query, budget)
666+
if err != nil {
667+
writeErr(w, http.StatusInternalServerError, "retrieval failed: "+err.Error())
668+
return
648669
}
649670
totalUsage.Add(retrievalUsage)
650671

@@ -700,7 +721,7 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
700721
maxAnswerTokens = 1024
701722
}
702723

703-
answerText, synthUsage, err := synthesiseAnswer(r.Context(), d.LLM, synthModel, body.Query, enriched, maxAnswerTokens)
724+
answerText, synthUsage, err := synthesiseAnswer(r.Context(), d.LLM, synthModel, body.Query, plan, enriched, maxAnswerTokens)
704725
if err != nil {
705726
writeErr(w, http.StatusInternalServerError, "synthesis failed: "+err.Error())
706727
return
@@ -729,7 +750,7 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
729750
citations = append(citations, c)
730751
}
731752

732-
writeJSON(w, http.StatusOK, map[string]any{
753+
resp := map[string]any{
733754
"document_id": body.DocumentID,
734755
"query": body.Query,
735756
"answer": answerText,
@@ -744,17 +765,27 @@ func (d Deps) handleAnswer(w http.ResponseWriter, r *http.Request) {
744765
"llm_calls": totalUsage.LLMCalls,
745766
},
746767
"elapsed_ms": time.Since(started).Milliseconds(),
747-
})
768+
}
769+
if plan != nil {
770+
resp["plan"] = plan
771+
}
772+
writeJSON(w, http.StatusOK, resp)
748773
}
749774

750775
// synthesiseAnswer runs one LLM call producing the final answer from
751776
// retrieved sections + their extracted spans. The model is told to
752-
// cite by section ID.
753-
func synthesiseAnswer(ctx context.Context, client llmgate.Client, model, query string, secs []sectionWithContent, maxAnswerTokens int) (string, retrieval.Usage, error) {
777+
// cite by section ID. When plan is non-nil its structured hints
778+
// (intent, entities, expected_doc_areas, sub_questions) are folded
779+
// into the prompt as a short "Planner notes" block so the model can
780+
// reason with the same understanding the retrieval pipeline used.
781+
func synthesiseAnswer(ctx context.Context, client llmgate.Client, model, query string, plan *retrieval.Plan, secs []sectionWithContent, maxAnswerTokens int) (string, retrieval.Usage, error) {
754782
var b strings.Builder
755783
b.WriteString("You are answering a user's question using ONLY the evidence below.\n\n")
756784
b.WriteString("User query:\n")
757785
b.WriteString(query)
786+
if plan != nil {
787+
writePlanHints(&b, plan)
788+
}
758789
b.WriteString("\n\nRetrieved evidence (each block is a section of the document):\n")
759790
for i, e := range secs {
760791
fmt.Fprintf(&b, "\n[%d] section_id=%s, title=%q", i+1, e.sec.ID, e.sec.Title)
@@ -1018,6 +1049,114 @@ func (d Deps) handleQueueWebhook(w http.ResponseWriter, r *http.Request) {
10181049
w.WriteHeader(http.StatusNoContent)
10191050
}
10201051

1052+
// --- planning helpers ---
1053+
1054+
// planningEnabled reports whether the request should go through the
1055+
// planner. The per-request body field (when present) wins over the
1056+
// server-side config; a nil body field falls back to the config.
1057+
func (d Deps) planningEnabled(bodyOverride *bool) bool {
1058+
if d.Planner == nil {
1059+
return false
1060+
}
1061+
if bodyOverride != nil {
1062+
return *bodyOverride
1063+
}
1064+
return d.Planning.Enabled
1065+
}
1066+
1067+
// runPlanner issues the planning LLM call when planning is enabled.
1068+
// Returns (nil, zero usage) when planning is off, the query is empty,
1069+
// the planner is missing, or the planner gracefully degraded to a
1070+
// no-plan result. Transport errors from the planner are LOGGED but not
1071+
// propagated — the engine continues with the original query rather
1072+
// than 500ing a working retrieval request because of a planner blip.
1073+
func (d Deps) runPlanner(ctx context.Context, query string, bodyOverride *bool) (*retrieval.Plan, retrieval.Usage) {
1074+
if !d.planningEnabled(bodyOverride) {
1075+
return nil, retrieval.Usage{}
1076+
}
1077+
plan, usage, err := d.Planner.Plan(ctx, query)
1078+
if err != nil {
1079+
if d.Logger != nil {
1080+
d.Logger.Warn("planner: failed; continuing without plan", "err", err)
1081+
}
1082+
return nil, usage
1083+
}
1084+
return plan, usage
1085+
}
1086+
1087+
// runSelection picks section IDs for the query, optionally going
1088+
// through the Decomposer when the plan is multi-hop AND planning-level
1089+
// decomposition is enabled. Returns the same []SectionID Strategy.Select
1090+
// would.
1091+
func (d Deps) runSelection(ctx context.Context, t *tree.Tree, plan *retrieval.Plan, query string, budget retrieval.ContextBudget) ([]tree.SectionID, error) {
1092+
if d.shouldDecompose(plan) {
1093+
ids, _, err := retrieval.NewDecomposer(d.Strategy).DecomposedSelect(ctx, t, plan, query, budget)
1094+
return ids, err
1095+
}
1096+
return d.Strategy.Select(ctx, t, query, budget)
1097+
}
1098+
1099+
// runSelectionWithUsage is the cost-tracking variant used by /v1/answer.
1100+
// Returns the selected IDs plus the Usage accumulated during selection
1101+
// (across all sub-questions for multi-hop plans).
1102+
func (d Deps) runSelectionWithUsage(ctx context.Context, t *tree.Tree, plan *retrieval.Plan, query string, budget retrieval.ContextBudget) ([]tree.SectionID, retrieval.Usage, error) {
1103+
if d.shouldDecompose(plan) {
1104+
return retrieval.NewDecomposer(d.Strategy).DecomposedSelect(ctx, t, plan, query, budget)
1105+
}
1106+
if cs, ok := d.Strategy.(retrieval.CostStrategy); ok {
1107+
res, err := cs.SelectWithCost(ctx, t, query, budget)
1108+
if err != nil {
1109+
return nil, retrieval.Usage{}, err
1110+
}
1111+
if res == nil {
1112+
return nil, retrieval.Usage{}, nil
1113+
}
1114+
return res.SelectedIDs, res.Usage, nil
1115+
}
1116+
ids, err := d.Strategy.Select(ctx, t, query, budget)
1117+
if err != nil {
1118+
return nil, retrieval.Usage{}, err
1119+
}
1120+
return ids, retrieval.Usage{}, nil
1121+
}
1122+
1123+
// shouldDecompose returns true when the plan is multi-hop AND
1124+
// decomposition is enabled at the config level. The Decomposer
1125+
// itself short-circuits to Strategy.Select when the plan is missing
1126+
// or non-multi-hop, but we duplicate that check here so we avoid
1127+
// allocating a Decomposer when it would be a no-op.
1128+
func (d Deps) shouldDecompose(plan *retrieval.Plan) bool {
1129+
if plan == nil || !plan.IsMultiHop || len(plan.SubQuestions) == 0 {
1130+
return false
1131+
}
1132+
return d.Planning.Decompose
1133+
}
1134+
1135+
// writePlanHints appends a short, model-readable "Planner notes" block
1136+
// describing the structured plan. Synthesis uses this to orient itself
1137+
// before reading the evidence.
1138+
func writePlanHints(b *strings.Builder, plan *retrieval.Plan) {
1139+
if plan == nil {
1140+
return
1141+
}
1142+
b.WriteString("\n\nPlanner notes (structured understanding of the query):")
1143+
if plan.Intent != "" {
1144+
fmt.Fprintf(b, "\n- intent: %s", plan.Intent)
1145+
}
1146+
if len(plan.Entities) > 0 {
1147+
fmt.Fprintf(b, "\n- entities: %s", strings.Join(plan.Entities, ", "))
1148+
}
1149+
if len(plan.ExpectedDocAreas) > 0 {
1150+
fmt.Fprintf(b, "\n- expected document areas: %s", strings.Join(plan.ExpectedDocAreas, ", "))
1151+
}
1152+
if plan.IsMultiHop && len(plan.SubQuestions) > 0 {
1153+
b.WriteString("\n- sub-questions:")
1154+
for _, q := range plan.SubQuestions {
1155+
fmt.Fprintf(b, "\n - %s", q)
1156+
}
1157+
}
1158+
}
1159+
10211160
// --- helpers ---
10221161

10231162
func writeJSON(w http.ResponseWriter, status int, v any) {

0 commit comments

Comments
 (0)