Added routing layer in agentic RAG

Author: smasurekar

deploy/compose/.env

@@ -101,3 +101,39 @@ export OCR_MS_GPU_ID=0
 # ==========================
 
 export PROMPT_CONFIG_FILE=${PWD}/src/nvidia_rag/rag_server/prompt.yaml
+
+# ==========================
+# Agentic RAG
+# ==========================
+
+# Post-synthesis verification pass (re-checks the answer and back-fills with
+# follow-up retrieval). Enabled to improve groundedness on queries where the
+# context was retrieved but the synthesis refused/under-answered.
+export AGENTIC_VERIFICATION_ENABLED=false
+
+# --- Per-role reasoning config (active = DEFAULT all-on for latency probe) ---
+# Default = thinking=true, budget=0 (unbounded), low_effort=false on ALL roles = max reasoning.
+export AGENTIC_PLANNER_LLM_ENABLE_THINKING=true
+export AGENTIC_PLANNER_LLM_REASONING_BUDGET=0
+export AGENTIC_PLANNER_LLM_LOW_EFFORT=false
+export AGENTIC_SYNTHESIS_LLM_ENABLE_THINKING=true
+export AGENTIC_SYNTHESIS_LLM_REASONING_BUDGET=0
+export AGENTIC_SYNTHESIS_LLM_LOW_EFFORT=false
+export AGENTIC_SEED_GEN_LLM_ENABLE_THINKING=true
+export AGENTIC_SEED_GEN_LLM_REASONING_BUDGET=0
+export AGENTIC_SEED_GEN_LLM_LOW_EFFORT=false
+export AGENTIC_TASK_LLM_ENABLE_THINKING=true
+export AGENTIC_TASK_LLM_REASONING_BUDGET=0
+export AGENTIC_TASK_LLM_LOW_EFFORT=false
+
+# --- Per-query reasoning routing (classifier picks thinking on/off per query) ---
+# When enabled, a lightweight classifier labels each query simple/complex and
+# overrides the per-role ENABLE_THINKING above for that request only (uncertain
+# -> thinking-ON for safety). ROLES limits which roles the router controls.
+export AGENTIC_QUERY_ROUTING_ENABLED=false
+export AGENTIC_QUERY_ROUTING_ROLES=planner,task,seed_gen,synthesis
+
+# --- Cycle 3: non-reasoning latency lever (fewer scope-discovery rounds) ---
+# Default 2; reduced to 1 to cut sequential planner round-trips. Recall is saturated
+# (recall@10=0.94) so a 2nd scope round adds little retrieval coverage.
+export AGENTIC_PLANNER_MAX_SCOPE_ROUNDS=1

deploy/compose/docker-compose-rag-server.yaml

@@ -97,6 +97,12 @@ services:
       AGENTIC_SYNTHESIS_LLM_REASONING_BUDGET: ${AGENTIC_SYNTHESIS_LLM_REASONING_BUDGET:-0}
       AGENTIC_SYNTHESIS_LLM_LOW_EFFORT: ${AGENTIC_SYNTHESIS_LLM_LOW_EFFORT:-false}
 
+      # Per-query reasoning routing (off by default). When enabled, a lightweight
+      # classifier labels each query simple/complex and overrides the per-role
+      # ENABLE_THINKING for that request only (uncertain -> thinking-ON for safety).
+      AGENTIC_QUERY_ROUTING_ENABLED: ${AGENTIC_QUERY_ROUTING_ENABLED:-false}
+      AGENTIC_QUERY_ROUTING_ROLES: ${AGENTIC_QUERY_ROUTING_ROLES:-planner,task,seed_gen,synthesis}
+
       ##===Query Rewriter Model specific configurations===
       APP_QUERYREWRITER_MODELNAME: ${APP_QUERYREWRITER_MODELNAME:-"nvidia/nemotron-3-super-120b-a12b"}
       # url on which query rewriter model is hosted. If "", Nvidia hosted API is used
@@ -261,10 +267,15 @@ services:
 
       # Verification pass (disabled by default)
       AGENTIC_VERIFICATION_ENABLED: ${AGENTIC_VERIFICATION_ENABLED:-false}
+      AGENTIC_VERIFICATION_MAX_TASKS: ${AGENTIC_VERIFICATION_MAX_TASKS:-3}
 
       # Context window budget for retrieved chunks
       AGENTIC_CONTEXT_MAX_TOKENS: ${AGENTIC_CONTEXT_MAX_TOKENS:-100000}
 
+      # Planner decomposition budget (surfaced for tuning; defaults match code)
+      AGENTIC_PLANNER_MAX_TASKS: ${AGENTIC_PLANNER_MAX_TASKS:-5}
+      AGENTIC_PLANNER_MAX_SCOPE_ROUNDS: ${AGENTIC_PLANNER_MAX_SCOPE_ROUNDS:-2}
+
     ports:
       - "8081:8081"
     expose:

deploy/helm/nvidia-blueprint-rag/values.yaml

@@ -292,6 +292,12 @@ envVars:
   AGENTIC_SYNTHESIS_LLM_REASONING_BUDGET: "0"
   AGENTIC_SYNTHESIS_LLM_LOW_EFFORT: "false"
 
+  # Per-query reasoning routing (off by default). When enabled, a lightweight
+  # classifier labels each query simple/complex and overrides the per-role
+  # ENABLE_THINKING for that request only (uncertain -> thinking-ON for safety).
+  AGENTIC_QUERY_ROUTING_ENABLED: "false"
+  AGENTIC_QUERY_ROUTING_ROLES: "planner,task,seed_gen,synthesis"
+
   NEMO_GUARDRAILS_URL: "nemo-guardrails:7331"
 
   # enable iterative query decomposition

frontend/src/hooks/__tests__/useChatStream.test.ts

@@ -448,4 +448,46 @@ describe('useChatStream — agentic streaming path (PR #512)', () => {
     ]);
     expect(last.reasoning_steps?.[0].status).toBe('done');
   });
+
+  it('renders the query_routing stage (start+end) as one completed step', async () => {
+    // Backend emits a query_routing stage_start/stage_end pair as the first
+    // two chunks when per-query reasoning routing is enabled. No frontend
+    // code changes are needed — it flows through the generic stage handling.
+    const { last } = await drain([
+      {
+        choices: [
+          {
+            delta: {
+              event_type: 'stage_start',
+              stage: 'query_routing',
+              reasoning_content:
+                'Analyzing the question to choose a reasoning strategy…',
+            },
+          },
+        ],
+      },
+      {
+        choices: [
+          {
+            delta: {
+              event_type: 'stage_end',
+              stage: 'query_routing',
+              reasoning_content:
+                'Classified as a multi-step question — enabling step-by-step reasoning.',
+            },
+          },
+        ],
+      },
+      { choices: [{ delta: {}, finish_reason: 'stop' }] },
+    ]);
+    const step = last.reasoning_steps?.[0];
+    expect(step?.stage).toBe('query_routing');
+    expect(step?.status).toBe('done');
+    expect(step?.label).toBe(
+      'Analyzing the question to choose a reasoning strategy…'
+    );
+    expect(step?.summary).toBe(
+      'Classified as a multi-step question — enabling step-by-step reasoning.'
+    );
+  });
 });

frontend/src/hooks/__tests__/useCitationUtils.test.ts

@@ -34,6 +34,7 @@ describe('useCitationUtils', () => {
       const { result } = renderHook(() => useCitationUtils());
       expect(result.current.formatStage('initial_retrieval')).toBe('Initial retrieval');
       expect(result.current.formatStage('verify_execute')).toBe('Verify execute');
+      expect(result.current.formatStage('query_routing')).toBe('Query routing');
     });
 
     it('humanises kebab-case stage identifiers', () => {

src/nvidia_rag/rag_server/agentic_rag/agentic_rag.py

@@ -259,7 +259,16 @@ def _resolve_role_llm(
         model_or_endpoint_overridden = (
             overrides.model is not None or overrides.llm_endpoint is not None
         )
-        cache_key = "__shared__" if model_or_endpoint_overridden else role_name
+        # When the query router supplies per-role thinking overrides, every role
+        # may resolve to a different enable_thinking, so they cannot share one
+        # client — force a per-role cache key even if model/endpoint is also
+        # overridden (each role still uses the overridden model/endpoint, just
+        # with its own thinking flag).
+        role_thinking_overridden = bool(overrides.enable_thinking_by_role)
+        if model_or_endpoint_overridden and not role_thinking_overridden:
+            cache_key = "__shared__"
+        else:
+            cache_key = role_name
 
         cache = overrides._built_llms
         if cache_key in cache:
@@ -356,10 +365,17 @@ def _resolve_gen(field_name: str, override_val: Any) -> Any:
         temperature = _resolve_gen("temperature", overrides.temperature)
         top_p = _resolve_gen("top_p", overrides.top_p)
         max_tokens = _resolve_gen("max_tokens", overrides.max_tokens)
-        # Thinking params are never overridden by the per-request API; always
-        # resolved from the per-role agentic config so that AGENTIC_*_LLM_ENABLE_THINKING
-        # is respected even when temperature/top_p/max_tokens cause this override path.
-        enable_thinking = _resolve_gen("enable_thinking", None)
+        # enable_thinking may be overridden per-role by the query router; when
+        # the router didn't decide for this role (or routing is off) the value
+        # is None and we fall back to the per-role AGENTIC_*_LLM_ENABLE_THINKING
+        # config. reasoning_budget / low_effort are never overridden per-request
+        # and always resolve from config.
+        role_thinking_override = (
+            overrides.enable_thinking_by_role.get(role_name)
+            if overrides.enable_thinking_by_role
+            else None
+        )
+        enable_thinking = _resolve_gen("enable_thinking", role_thinking_override)
         reasoning_budget = _resolve_gen("reasoning_budget", None)
         low_effort = _resolve_gen("low_effort", None)
 

src/nvidia_rag/rag_server/agentic_rag/builder.py

@@ -104,18 +104,24 @@ class AgenticLLMOverrides:
     temperature: float | None = None
     top_p: float | None = None
     max_tokens: int | None = None
+    # Per-role enable_thinking override set by the query router (see
+    # query_router.build_thinking_overrides). Maps role name (planner / task /
+    # seed_gen / synthesis) → desired thinking bool for this request only.
+    # A role absent from the dict keeps its server-wide AGENTIC_*_LLM_ENABLE_THINKING
+    # default. None means "router inactive — no thinking override".
+    enable_thinking_by_role: dict[str, bool] | None = None
     # Per-request cache: maps cache key → built LLM client. Key is "__shared__"
     # when model/endpoint is overridden (all 4 roles collapse to one client);
     # otherwise it's the role name (planner / task / seed_gen / synthesis) so
     # each role gets its own client with that role's configured model/endpoint
     # but the request-provided generation params applied.
     # Mutable; not part of equality / hashing.
-    _built_llms: dict[str, Any] = field(
-        default_factory=dict, repr=False, compare=False
-    )
+    _built_llms: dict[str, Any] = field(default_factory=dict, repr=False, compare=False)
 
     def has_any_override(self) -> bool:
         """True if at least one non-None override is set."""
+        if self.enable_thinking_by_role:
+            return True
         return any(
             v is not None
             for v in (
@@ -319,9 +325,7 @@ def _make_role_llm(
     temperature = _resolve_role_generation_param(
         role_cfg, fallback_cfg, rag_config, "temperature"
     )
-    top_p = _resolve_role_generation_param(
-        role_cfg, fallback_cfg, rag_config, "top_p"
-    )
+    top_p = _resolve_role_generation_param(role_cfg, fallback_cfg, rag_config, "top_p")
     max_tokens = _resolve_role_generation_param(
         role_cfg, fallback_cfg, rag_config, "max_tokens"
     )