feat(shallow-research): replace escalation heuristics with validated assessment - #329
Conversation
Signed-off-by: Kyle Zheng <kyzheng@nvidia.com>
Signed-off-by: Kyle Zheng <kyzheng@nvidia.com>
Signed-off-by: Kyle Zheng <kyzheng@nvidia.com>
WalkthroughThe PR adds a bounded post-shallow escalation assessment with structured statuses, fail-closed parsing, source tracking, assessment-based chat routing, evaluation tooling, packaged prompts, and comprehensive tests and documentation. ChangesShallow escalation workflow
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 1 warning)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
KyleZheng1284
changed the title
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
|
@coderabbitai full review |
✅ Action performedFull review finished. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/source/resources/faq.md`:
- Around line 43-50: Update the escalation behavior description in the FAQ to
state that a material source conflict must be supported by at least two unique
retrieved sources; conflicts with fewer than two unique sources must fail
closed. Preserve the existing explanation of escalation triggers and routing.
In `@src/aiq_agent/agents/chat_researcher/agent.py`:
- Around line 296-326: Update the shallow research result handling after
computing new_messages so an empty new_messages case returns the same error
response as the existing shallow failure path, including an AIMessage, high
confidence, and no escalation. Preserve assessment handling for successful
responses, and add coverage for a shallow function that returns the unchanged
state.
In `@src/aiq_agent/agents/shallow_researcher/agent.py`:
- Line 24: Canonicalize retrieved URL identities before source counting in the
shallow researcher flow around the uniqueness check at lines 64–70: normalize
the URL scheme and authority, remove any fragment, and use the canonical value
when adding source identities. Extend the uniqueness test to cover host-case and
fragment-only URL variants so repeated retrievals cannot satisfy the
two-unique-source guard.
In `@src/aiq_agent/agents/shallow_researcher/escalation.py`:
- Around line 225-231: Update the shallow escalation assessment flow and its
assessor input to carry bounded, invocation-local source IDs or evidence instead
of relying on the aggregate source_count. For material_conflict, require and
validate at least two distinct cited source IDs that support the conflict before
escalating; otherwise retain the sufficient fallback. Add focused tests and
documentation covering source attribution and rejecting unrelated retrieved
sources.
- Around line 218-224: Broaden _UNAVAILABLE_OUTREACH_STRATEGY, used in the
deep_research_strategy check, to recognize additional human-contact phrasing
such as requesting confirmation from a vendor spokesperson while preserving
existing unavailable-outreach matches. Add a regression test covering a distinct
outreach wording and verify it returns ShallowEscalationAssessment.sufficient().
In `@src/aiq_agent/agents/shallow_researcher/prompts/escalation_assessment.j2`:
- Around line 9-15: Update the required output shape in the escalation
assessment prompt to show a syntactically valid JSON example, replacing schema
notation such as union bars and bare string types with representative quoted
values or null. Preserve the exact field names and permitted status values while
ensuring models can copy the example as valid JSON.
In `@src/aiq_agent/evaluation/shallow_escalation.py`:
- Around line 108-119: Update acceptance_passed to avoid hardcoded total,
shallow, and escalation case counts; derive the expected counts from the loaded
case-list partitions or pass them explicitly as parameters. Ensure the
acceptance checks compare summary values against those derived expectations
while preserving the existing routing, parsing, call-count, and minimum
escalation correctness thresholds.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 82c8d281-9ba4-48bc-8e02-b468b6bc8412
📒 Files selected for processing (22)
docs/source/architecture/agents/shallow-researcher.mddocs/source/architecture/overview.mddocs/source/customization/configuration-reference.mddocs/source/customization/prompts.mddocs/source/resources/faq.mdpyproject.tomlscripts/evaluate_shallow_escalation.pysrc/aiq_agent/agents/chat_researcher/agent.pysrc/aiq_agent/agents/chat_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/agent.pysrc/aiq_agent/agents/shallow_researcher/escalation.pysrc/aiq_agent/agents/shallow_researcher/models/__init__.pysrc/aiq_agent/agents/shallow_researcher/models/escalation.pysrc/aiq_agent/agents/shallow_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/prompts/escalation_assessment.j2src/aiq_agent/evaluation/__init__.pysrc/aiq_agent/evaluation/shallow_escalation.pytests/aiq_agent/agents/chat_researcher/test_escalation_routing.pytests/aiq_agent/agents/shallow_researcher/fixtures/escalation_cases.jsontests/aiq_agent/agents/shallow_researcher/test_agent.pytests/aiq_agent/agents/shallow_researcher/test_escalation.pytests/aiq_agent/test_package_assets.py
📜 Review details
🧰 Additional context used
📓 Path-based instructions (8)
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Run ruff check and ruff format validation for Python code changes
**/*.py: Python code must be linted and formatted with Ruff using line length 120, target Python 3.11, rule sets E,F,W,I,PL,UP, and isort force-single-line configuration
Never commit secrets, tokens, or environment-specific hostnames; use environment variables and SecretStr instead, resolving API keys at runtime
Never print or log secret values, including in tool output or error messages
Missing-secret paths must degrade gracefully (stub/skip), not crash or leak
Do not hand-reformat unrelated code when making changes; match the existing import and formatting style
Files:
src/aiq_agent/agents/shallow_researcher/models/__init__.pysrc/aiq_agent/evaluation/__init__.pytests/aiq_agent/test_package_assets.pysrc/aiq_agent/agents/chat_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/models/escalation.pysrc/aiq_agent/agents/shallow_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/escalation.pysrc/aiq_agent/evaluation/shallow_escalation.pysrc/aiq_agent/agents/shallow_researcher/agent.pyscripts/evaluate_shallow_escalation.pysrc/aiq_agent/agents/chat_researcher/agent.pytests/aiq_agent/agents/chat_researcher/test_escalation_routing.pytests/aiq_agent/agents/shallow_researcher/test_escalation.pytests/aiq_agent/agents/shallow_researcher/test_agent.py
src/aiq_agent/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/aiq_agent/**/*.py: Respect authenticated data sources by honoring requires_auth, per-user token pass-through, and backend token validators; apply owner guardrails before loading protected report or artifact context into an agent
Do not weaken or bypass AuthMiddleware, validators, or auth gating without a prior design discussion
Files:
src/aiq_agent/agents/shallow_researcher/models/__init__.pysrc/aiq_agent/evaluation/__init__.pysrc/aiq_agent/agents/chat_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/models/escalation.pysrc/aiq_agent/agents/shallow_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/escalation.pysrc/aiq_agent/evaluation/shallow_escalation.pysrc/aiq_agent/agents/shallow_researcher/agent.pysrc/aiq_agent/agents/chat_researcher/agent.py
**
⚙️ CodeRabbit configuration file
**:AI-Q Agent Guidance
Repository-global instructions for coding agents and for humans reviewing
agent-authored changes. These rules apply to every task in this repository.
Task-specific runbooks live in.agents/skills/— load the
relevant skill before starting a workflow it covers.Project overview
AI-Q is an NVIDIA AI Blueprint: an enterprise research agent built on the
NeMo Agent Toolkit (NAT). The deployed product is a research blueprint, not
a general skill runtime. New retrieval sources and tools are NAT functions;
agent behavior is driven by workflow YAML, Jinja2 prompts, and a data-source
registry — not by hard-coded logic.Primary boundaries:
- Backend Python package:
src/aiq_agent/.- Data-source and tool packages:
sources/(each is its own package).- Frontends and tooling:
frontends/(web UI infrontends/ui/, eval harnesses
infrontends/benchmarks/).- Configs, deployment, docs:
configs/,deploy/,docs/.Stay inside this repository. If your workspace also contains adjacent repos
(for example a sibling NeMo-Relay checkout), do not edit them as part of an AI-Q
change. Treatsources/*as independent packages: prefer the smallest change
scoped to the package you are touching.Repository structure
Path Purpose src/aiq_agent/Backend agent, FastAPI extensions, auth, observability, knowledge sources/Data-source / tool packages (e.g. tavily_web_search,google_scholar_paper_search)configs/Workflow YAML configs (e.g. config_cli_default.yml)frontends/ui/Next.js / React / TypeScript / Tailwind / KUI web UI frontends/benchmarks/Eval harnesses: freshqa,deepsearch_qa,deepresearch_benchdeploy/Docker Compose and Helm/Kubernetes assets; deploy/.envfor secretsdocs/source/...
Files:
src/aiq_agent/agents/shallow_researcher/models/__init__.pysrc/aiq_agent/evaluation/__init__.pypyproject.tomltests/aiq_agent/test_package_assets.pydocs/source/customization/prompts.mddocs/source/resources/faq.mddocs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/shallow_researcher/prompts/escalation_assessment.j2tests/aiq_agent/agents/shallow_researcher/fixtures/escalation_cases.jsonsrc/aiq_agent/agents/chat_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/models/escalation.pysrc/aiq_agent/agents/shallow_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/escalation.pydocs/source/architecture/overview.mdsrc/aiq_agent/evaluation/shallow_escalation.pysrc/aiq_agent/agents/shallow_researcher/agent.pydocs/source/architecture/agents/shallow-researcher.mdscripts/evaluate_shallow_escalation.pysrc/aiq_agent/agents/chat_researcher/agent.pytests/aiq_agent/agents/chat_researcher/test_escalation_routing.pytests/aiq_agent/agents/shallow_researcher/test_escalation.pytests/aiq_agent/agents/shallow_researcher/test_agent.py
src/aiq_agent/agents/**/*
⚙️ CodeRabbit configuration file
src/aiq_agent/agents/**/*: Review agent changes for research workflow correctness, graph state transitions, prompt/tool contracts,
HITL behavior, and failure handling. Flag changes that weaken source attribution, report generation,
async cancellation, checkpointing, or data-source selection without focused tests and docs.
Files:
src/aiq_agent/agents/shallow_researcher/models/__init__.pysrc/aiq_agent/agents/shallow_researcher/prompts/escalation_assessment.j2src/aiq_agent/agents/chat_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/models/escalation.pysrc/aiq_agent/agents/shallow_researcher/models/state.pysrc/aiq_agent/agents/shallow_researcher/escalation.pysrc/aiq_agent/agents/shallow_researcher/agent.pysrc/aiq_agent/agents/chat_researcher/agent.py
{.github/**,ci/**,.pre-commit-config.yaml,pyproject.toml,uv.lock}
⚙️ CodeRabbit configuration file
{.github/**,ci/**,.pre-commit-config.yaml,pyproject.toml,uv.lock}: Review automation and packaging changes for least-privilege permissions, pinned versions where appropriate,
copy-pr-bot pull-request/ branch behavior, reproducible uv/npm setup, secret handling, and consistency with
the documented validation matrix.
Files:
pyproject.toml
**/*test*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Run pytest for all behavior changes in Python code
Files:
tests/aiq_agent/test_package_assets.pytests/aiq_agent/agents/chat_researcher/test_escalation_routing.pytests/aiq_agent/agents/shallow_researcher/test_escalation.pytests/aiq_agent/agents/shallow_researcher/test_agent.py
docs/source/**/*
📄 CodeRabbit inference engine (AGENTS.md)
Update the docs under docs/source/ when behavior, configuration, or workflows change
Files:
docs/source/customization/prompts.mddocs/source/resources/faq.mddocs/source/customization/configuration-reference.mddocs/source/architecture/overview.mddocs/source/architecture/agents/shallow-researcher.md
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}
⚙️ CodeRabbit configuration file
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}: Review documentation for command accuracy, branch-name consistency, current CI and copy-pr-bot behavior, public
vs internal boundary clarity, stale examples, and links that no longer match the repository layout.
Files:
docs/source/customization/prompts.mddocs/source/resources/faq.mddocs/source/customization/configuration-reference.mddocs/source/architecture/overview.mddocs/source/architecture/agents/shallow-researcher.md
🧠 Learnings (1)
📚 Learning: 2026-07-06T23:55:42.908Z
Learnt from: cdgamarose-nv
Repo: NVIDIA-AI-Blueprints/aiq PR: 311
File: src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2:74-81
Timestamp: 2026-07-06T23:55:42.908Z
Learning: In agent test files (e.g., tests/aiq_agent/agents/*/test_agent.py), avoid brittle assertions that match exact substrings from prompt template files (such as *.j2 prompt wording). Prompt wording can change frequently, so instead assert structural/behavioral properties (e.g., that the prompt builder is called, that required sections/fields are present via stable markers, that the model output/agent behavior conforms to an expected schema, or that key actions are taken) rather than matching literal prompt text.
Applied to files:
tests/aiq_agent/agents/shallow_researcher/test_agent.py
🪛 ast-grep (0.44.1)
src/aiq_agent/agents/shallow_researcher/escalation.py
[info] 45-45: use jsonify instead of json.dumps for JSON output
Context: json.dumps(value, ensure_ascii=False)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 78-86: use jsonify instead of json.dumps for JSON output
Context: json.dumps(
{
"user_query": _bound_text(query, max_json_chars=_MAX_QUERY_JSON_CHARS),
"shallow_answer": _bound_text(answer, max_json_chars=_MAX_ANSWER_JSON_CHARS),
"retrieved_source_count": source_count,
"tool_budget_exhausted": tool_budget_exhausted,
},
ensure_ascii=False,
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
src/aiq_agent/evaluation/shallow_escalation.py
[info] 34-41: use jsonify instead of json.dumps for JSON output
Context: json.dumps(
{
"status": self.expected_status,
"unresolved_requirement": self.expected_detail if self.expected_status == "material_gap" else None,
"material_conflict": self.expected_detail if self.expected_status == "material_conflict" else None,
"deep_research_strategy": self.expected_strategy,
}
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
scripts/evaluate_shallow_escalation.py
[info] 37-37: Do not hardcode temporary file or directory names
Context: "/tmp/aiq-shallow-escalation-eval.json"
Note: [CWE-377] Insecure Temporary File.
(hardcoded-tmp-file)
[info] 243-243: use jsonify instead of json.dumps for JSON output
Context: json.dumps(report, indent=2, ensure_ascii=False)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 244-244: use jsonify instead of json.dumps for JSON output
Context: json.dumps({"output": str(args.output), "acceptance": report["acceptance"], "summary": summary}, indent=2)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
tests/aiq_agent/agents/shallow_researcher/test_agent.py
[info] 370-377: use jsonify instead of json.dumps for JSON output
Context: json.dumps(
{
"status": "sufficient",
"unresolved_requirement": None,
"material_conflict": None,
"deep_research_strategy": None,
}
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 411-418: use jsonify instead of json.dumps for JSON output
Context: json.dumps(
{
"status": "material_conflict",
"unresolved_requirement": None,
"material_conflict": "The specifications disagree.",
"deep_research_strategy": "Compare both primary specifications.",
}
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🪛 LanguageTool
docs/source/resources/faq.md
[style] ~45-~45: Consider a different adjective to strengthen your wording.
Context: ...onclusion-critical source conflict that deep research has a concrete way to address;...
(DEEP_PROFOUND)
[style] ~47-~47: Consider a different adjective to strengthen your wording.
Context: ... through the clarifier node and then to deep research. The clarifier asks only for m...
(DEEP_PROFOUND)
docs/source/customization/configuration-reference.md
[style] ~459-~459: Consider a different adjective to strengthen your wording.
Context: ...w a material shallow result to route to deep research. This setting does not affect ...
(DEEP_PROFOUND)
[style] ~459-~459: Consider a different adjective to strengthen your wording.
Context: ...e intent classifier initially routes to deep research. | | enable_clarifier | `boo...
(DEEP_PROFOUND)
docs/source/architecture/overview.md
[style] ~65-~65: Consider a different adjective to strengthen your wording.
Context: ...her the response warrants escalation to deep research. When escalation is enabled...
(DEEP_PROFOUND)
[style] ~85-~85: Consider a different adjective to strengthen your wording.
Context: ...herwise, the node continues directly to deep research. Research planning then occ...
(DEEP_PROFOUND)
[style] ~145-~145: Consider a different adjective to strengthen your wording.
Context: ...ved for material gaps or conflicts that deep research has a concrete way to addres...
(DEEP_PROFOUND)
docs/source/architecture/agents/shallow-researcher.md
[style] ~82-~82: Consider a different adjective to strengthen your wording.
Context: ...swer can be evaluated for escalation to deep research. This assessment is not part o...
(DEEP_PROFOUND)
[style] ~99-~99: Consider a different adjective to strengthen your wording.
Context: ...est, or any limitation does not justify deep research. | | material_gap | An expli...
(DEEP_PROFOUND)
[style] ~109-~109: Consider a different adjective to strengthen your wording.
Context: ...tructure failures, and information that deep research has no credible way to obtain ...
(DEEP_PROFOUND)
[style] ~122-~122: Consider a different adjective to strengthen your wording.
Context: ...endation routes to the clarifier before deep research; it never invokes the deep res...
(DEEP_PROFOUND)
[style] ~124-~124: Consider a different adjective to strengthen your wording.
Context: ...mplex queries should still be routed to deep research by the intent classifier rathe...
(DEEP_PROFOUND)
🔇 Additional comments (22)
src/aiq_agent/evaluation/__init__.py (1)
1-5: LGTM!src/aiq_agent/evaluation/shallow_escalation.py (1)
1-106: LGTM!tests/aiq_agent/agents/shallow_researcher/fixtures/escalation_cases.json (1)
1-202: LGTM!scripts/evaluate_shallow_escalation.py (2)
1-46: LGTM!Also applies to: 59-207, 219-251
209-218: 🗄️ Data Integrity & IntegrationDrop this dependency concern
langchain_nvidia_ai_endpointsis already pulled in transitively vianvidia-nat[langchain,...]inpyproject.tomland resolved inuv.lock, so this import is covered by the project environment.> Likely an incorrect or invalid review comment.pyproject.toml (1)
25-25: LGTM!tests/aiq_agent/test_package_assets.py (1)
1-25: LGTM!src/aiq_agent/agents/shallow_researcher/models/escalation.py (1)
1-60: LGTM!src/aiq_agent/agents/shallow_researcher/models/__init__.py (1)
20-26: LGTM!src/aiq_agent/agents/shallow_researcher/models/state.py (1)
27-28: LGTM!Also applies to: 42-44, 54-56
src/aiq_agent/agents/shallow_researcher/prompts/escalation_assessment.j2 (1)
1-8: LGTM!Also applies to: 16-49
docs/source/architecture/agents/shallow-researcher.md (1)
79-126: LGTM!Also applies to: 140-142
docs/source/customization/prompts.md (1)
16-16: LGTM!Also applies to: 33-34
docs/source/architecture/overview.md (1)
65-89: LGTM!Also applies to: 104-104, 142-146
src/aiq_agent/agents/shallow_researcher/escalation.py (2)
44-217: LGTM!
233-239: LGTM!tests/aiq_agent/agents/shallow_researcher/test_agent.py (1)
18-18: LGTM!Also applies to: 30-31, 337-548
src/aiq_agent/agents/chat_researcher/agent.py (2)
46-46: LGTM!Also applies to: 76-89, 211-286
487-492: LGTM!src/aiq_agent/agents/chat_researcher/models/state.py (1)
25-25: LGTM!Also applies to: 56-56, 72-72
docs/source/customization/configuration-reference.md (1)
459-459: LGTM!tests/aiq_agent/agents/chat_researcher/test_escalation_routing.py (1)
1-233: LGTM!
| If `enable_escalation: true` in the workflow config, a successful shallow | ||
| answer receives one bounded, tool-free assessment call. Escalation is reserved | ||
| for a material unmet requirement or a conclusion-critical source conflict that | ||
| deep research has a concrete way to address; response length and keywords do | ||
| not trigger it. A material result routes through the clarifier node and then to | ||
| deep research. The clarifier asks only for missing context or output-shape | ||
| preferences when it is enabled and not skipped; planning happens inside the | ||
| deep-research workflow. Refer to [Architecture Overview](../architecture/overview.md). |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the conflict’s two-source threshold.
A material_conflict fails closed when fewer than two unique sources were retrieved, but this routing condition is absent here. State that a conflict must be supported by at least two unique retrieved sources.
🧰 Tools
🪛 LanguageTool
[style] ~45-~45: Consider a different adjective to strengthen your wording.
Context: ...onclusion-critical source conflict that deep research has a concrete way to address;...
(DEEP_PROFOUND)
[style] ~47-~47: Consider a different adjective to strengthen your wording.
Context: ... through the clarifier node and then to deep research. The clarifier asks only for m...
(DEEP_PROFOUND)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/source/resources/faq.md` around lines 43 - 50, Update the escalation
behavior description in the FAQ to state that a material source conflict must be
supported by at least two unique retrieved sources; conflicts with fewer than
two unique sources must fail closed. Preserve the existing explanation of
escalation triggers and routing.
| error_message = "An error occurred during shallow research." | ||
| return { | ||
| "messages": [AIMessage(content=error_message)], | ||
| "shallow_assessment": None, | ||
| "shallow_result": ShallowResult( | ||
| answer="An error occurred during shallow research.", | ||
| confidence="low", | ||
| escalate_to_deep=True, | ||
| escalation_reason="Shallow research encountered an error", | ||
| ) | ||
| answer=error_message, | ||
| confidence="high", | ||
| escalate_to_deep=False, | ||
| ), | ||
| } | ||
| assessment = getattr(result, "escalation_assessment", None) | ||
| if not isinstance(assessment, ShallowEscalationAssessment): | ||
| assessment = None | ||
| new_messages = result.messages[len(trimmed_messages) :] | ||
| final_ai_message = next( | ||
| (m for m in reversed(new_messages) if isinstance(m, AIMessage) and not m.tool_calls), | ||
| None, | ||
| ) | ||
| if final_ai_message: | ||
| return {"messages": [final_ai_message], "shallow_result": None} | ||
| return { | ||
| "messages": [final_ai_message], | ||
| "shallow_assessment": assessment, | ||
| "shallow_result": None, | ||
| } | ||
| if new_messages: | ||
| return {"messages": [new_messages[-1]], "shallow_result": None} | ||
| return {"messages": [], "shallow_result": None} | ||
| return { | ||
| "messages": [new_messages[-1]], | ||
| "shallow_assessment": assessment, | ||
| "shallow_result": None, | ||
| } | ||
| return {"messages": [], "shallow_assessment": None, "shallow_result": None} |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Treat “no new message” as a shallow failure.
Line 294 only catches an entirely empty result.messages. A normal state that preserves trimmed_messages but appends no response reaches Line 326 and returns an empty update; the message reducer retains the user’s prior message and the turn ends without an assistant response or error. Handle not new_messages with the same error result, and test a shallow function that returns the unchanged state.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiq_agent/agents/chat_researcher/agent.py` around lines 296 - 326, Update
the shallow research result handling after computing new_messages so an empty
new_messages case returns the same error response as the existing shallow
failure path, including an AIMessage, high confidence, and no escalation.
Preserve assessment handling for successful responses, and add coverage for a
shallow function that returns the unchanged state.
Source: Path instructions
| import os | ||
| import re | ||
| from collections.abc import Sequence | ||
| from contextvars import ContextVar |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Canonicalize URL identities before counting sources.
Line 67 treats URLs differing only by host casing or a #fragment as separate sources. Repeated retrieval of one document can therefore satisfy the two-unique-source conflict guard and permit an unsupported material_conflict escalation. Normalize scheme/authority and drop fragments before adding identities; add a variant-URL case to the uniqueness test.
Proposed fix
from contextvars import ContextVar
+from urllib.parse import urlsplit
+from urllib.parse import urlunsplit
+def _canonical_source_url(url: str) -> str:
+ parsed = urlsplit(url.strip())
+ if not parsed.scheme or not parsed.netloc:
+ return url.strip()
+ return urlunsplit(
+ (parsed.scheme.lower(), parsed.netloc.lower(), parsed.path, parsed.query, "")
+ )
+
def _source_identity(source: SourceEntry) -> tuple[str, str] | None:
"""Return a stable invocation-local identity for a retrieved source."""
if source.url:
- return ("url", source.url.strip())
+ return ("url", _canonical_source_url(source.url))Also applies to: 64-70
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiq_agent/agents/shallow_researcher/agent.py` at line 24, Canonicalize
retrieved URL identities before source counting in the shallow researcher flow
around the uniqueness check at lines 64–70: normalize the URL scheme and
authority, remove any fragment, and use the canonical value when adding source
identities. Extend the uniqueness test to cover host-case and fragment-only URL
variants so repeated retrievals cannot satisfy the two-unique-source guard.
| if assessment.deep_research_strategy and _UNAVAILABLE_OUTREACH_STRATEGY.search( | ||
| assessment.deep_research_strategy | ||
| ): | ||
| logger.warning( | ||
| "Shallow escalation assessment fallback status=sufficient reason=unavailable_outreach_strategy" | ||
| ) | ||
| return ShallowEscalationAssessment.sufficient() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
ast-grep outline src/aiq_agent/agents/shallow_researcher/models/escalation.py --items all
rg -n -C3 'deep_research_strategy|validator|material_gap|material_conflict' \
src/aiq_agent/agents/shallow_researcher/models/escalation.py \
tests/aiq_agent/agents/shallow_researcher/test_escalation.pyRepository: NVIDIA-AI-Blueprints/aiq
Length of output: 12670
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n### src/aiq_agent/agents/shallow_researcher/escalation.py (180-250)\n'
sed -n '180,250p' src/aiq_agent/agents/shallow_researcher/escalation.py
printf '\n### tests/aiq_agent/agents/shallow_researcher/test_escalation.py (200-260)\n'
sed -n '200,260p' tests/aiq_agent/agents/shallow_researcher/test_escalation.pyRepository: NVIDIA-AI-Blueprints/aiq
Length of output: 5064
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C4 '_UNAVAILABLE_OUTREACH_STRATEGY|deep_research_strategy' src/aiq_agent/agents/shallow_researcher/escalation.pyRepository: NVIDIA-AI-Blueprints/aiq
Length of output: 1117
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C3 'deep_research_strategy|research strategy|outreach|vendor spokesperson|contact the vendor' \
src/aiq_agent/agents/shallow_researcher \
tests/aiq_agent/agents/shallow_researcherRepository: NVIDIA-AI-Blueprints/aiq
Length of output: 12287
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C3 'deep_research_strategy' src/aiq_agent tests/aiq_agent | sed -n '1,240p'Repository: NVIDIA-AI-Blueprints/aiq
Length of output: 12725
Expand the outreach rejection to cover more human-contact phrasings (src/aiq_agent/agents/shallow_researcher/escalation.py:218-224). deep_research_strategy is still free text, so the current regex can miss variants like “Request confirmation from a vendor spokesperson.” Broaden the rejection rule or constrain the field to supported deep-research actions, and add a regression test for a second outreach wording.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiq_agent/agents/shallow_researcher/escalation.py` around lines 218 -
224, Broaden _UNAVAILABLE_OUTREACH_STRATEGY, used in the deep_research_strategy
check, to recognize additional human-contact phrasing such as requesting
confirmation from a vendor spokesperson while preserving existing
unavailable-outreach matches. Add a regression test covering a distinct outreach
wording and verify it returns ShallowEscalationAssessment.sufficient().
Source: Path instructions
| if assessment.status == "material_conflict" and source_count < 2: | ||
| logger.warning( | ||
| "Shallow escalation assessment fallback status=sufficient " | ||
| "reason=insufficient_sources_for_conflict retrieved_source_count=%d", | ||
| source_count, | ||
| ) | ||
| return ShallowEscalationAssessment.sufficient() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Validate conflict evidence, not the invocation-wide count.
Line 225 accepts any conflict when two sources were retrieved overall. The assessor receives only source_count, so it cannot establish that two distinct sources support the asserted conflict; unrelated sources satisfy this guard. Pass bounded invocation-local source IDs/evidence, require two distinct cited IDs in material_conflict, and validate them before escalating.
As per path instructions, “flag changes that weaken source attribution ... without focused tests and docs.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiq_agent/agents/shallow_researcher/escalation.py` around lines 225 -
231, Update the shallow escalation assessment flow and its assessor input to
carry bounded, invocation-local source IDs or evidence instead of relying on the
aggregate source_count. For material_conflict, require and validate at least two
distinct cited source IDs that support the conflict before escalating; otherwise
retain the sufficient fallback. Add focused tests and documentation covering
source attribution and rejecting unrelated retrieved sources.
Source: Path instructions
| Return exactly one JSON object with this shape: | ||
| { | ||
| "status": "sufficient" | "material_gap" | "material_conflict", | ||
| "unresolved_requirement": string | null, | ||
| "material_conflict": string | null, | ||
| "deep_research_strategy": string | null | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use a valid JSON example for the required output.
This block is not valid JSON: | unions and bare string values are schema notation. Because malformed responses fail closed to sufficient, a model may copy this format and silently suppress valid escalation.
Proposed fix
-Return exactly one JSON object with this shape:
+Return exactly one valid JSON object. The "status" field must be one of
+"sufficient", "material_gap", or "material_conflict".
{
- "status": "sufficient" | "material_gap" | "material_conflict",
- "unresolved_requirement": string | null,
- "material_conflict": string | null,
- "deep_research_strategy": string | null
+ "status": "sufficient",
+ "unresolved_requirement": null,
+ "material_conflict": null,
+ "deep_research_strategy": null
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Return exactly one JSON object with this shape: | |
| { | |
| "status": "sufficient" | "material_gap" | "material_conflict", | |
| "unresolved_requirement": string | null, | |
| "material_conflict": string | null, | |
| "deep_research_strategy": string | null | |
| } | |
| Return exactly one valid JSON object. The "status" field must be one of | |
| "sufficient", "material_gap", or "material_conflict". | |
| { | |
| "status": "sufficient", | |
| "unresolved_requirement": null, | |
| "material_conflict": null, | |
| "deep_research_strategy": null | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiq_agent/agents/shallow_researcher/prompts/escalation_assessment.j2`
around lines 9 - 15, Update the required output shape in the escalation
assessment prompt to show a syntactically valid JSON example, replacing schema
notation such as union bars and bare string types with representative quoted
values or null. Preserve the exact field names and permitted status values while
ensuring models can copy the example as valid JSON.
| def acceptance_passed(summary: dict[str, Any], *, min_escalation_correct: int = 8) -> bool: | ||
| """Apply the documented conservative acceptance thresholds.""" | ||
| return bool( | ||
| summary["case_count"] == 20 | ||
| and summary["shallow_case_count"] == 10 | ||
| and summary["escalation_case_count"] == 10 | ||
| and summary["shallow_route_correct"] == 10 | ||
| and summary["escalation_route_correct"] >= min_escalation_correct | ||
| and summary["parse_success_count"] == 20 | ||
| and summary["cases_with_exactly_one_call"] == 20 | ||
| and summary["logical_model_call_count"] == 20 | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Hardcoded case counts make acceptance_passed brittle for fixture growth.
case_count == 20, shallow_case_count == 10, and escalation_case_count == 10 are magic numbers tied to the current fixture. If escalation_cases.json gains a case later, this always returns False regardless of actual routing quality, since the counts no longer match. Consider deriving expected splits from the loaded case list itself (e.g., pass expected shallow/escalation counts as parameters, or compute len(shallow)/len(escalation) symmetrically instead of literal 10s).
♻️ Example refactor
-def acceptance_passed(summary: dict[str, Any], *, min_escalation_correct: int = 8) -> bool:
+def acceptance_passed(
+ summary: dict[str, Any],
+ *,
+ min_escalation_correct: int = 8,
+ expected_shallow_count: int = 10,
+ expected_escalation_count: int = 10,
+) -> bool:
"""Apply the documented conservative acceptance thresholds."""
return bool(
- summary["case_count"] == 20
- and summary["shallow_case_count"] == 10
- and summary["escalation_case_count"] == 10
- and summary["shallow_route_correct"] == 10
+ summary["case_count"] == expected_shallow_count + expected_escalation_count
+ and summary["shallow_case_count"] == expected_shallow_count
+ and summary["escalation_case_count"] == expected_escalation_count
+ and summary["shallow_route_correct"] == expected_shallow_count
and summary["escalation_route_correct"] >= min_escalation_correct
- and summary["parse_success_count"] == 20
- and summary["cases_with_exactly_one_call"] == 20
- and summary["logical_model_call_count"] == 20
+ and summary["parse_success_count"] == summary["case_count"]
+ and summary["cases_with_exactly_one_call"] == summary["case_count"]
+ and summary["logical_model_call_count"] == summary["case_count"]
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def acceptance_passed(summary: dict[str, Any], *, min_escalation_correct: int = 8) -> bool: | |
| """Apply the documented conservative acceptance thresholds.""" | |
| return bool( | |
| summary["case_count"] == 20 | |
| and summary["shallow_case_count"] == 10 | |
| and summary["escalation_case_count"] == 10 | |
| and summary["shallow_route_correct"] == 10 | |
| and summary["escalation_route_correct"] >= min_escalation_correct | |
| and summary["parse_success_count"] == 20 | |
| and summary["cases_with_exactly_one_call"] == 20 | |
| and summary["logical_model_call_count"] == 20 | |
| ) | |
| def acceptance_passed( | |
| summary: dict[str, Any], | |
| *, | |
| min_escalation_correct: int = 8, | |
| expected_shallow_count: int = 10, | |
| expected_escalation_count: int = 10, | |
| ) -> bool: | |
| """Apply the documented conservative acceptance thresholds.""" | |
| return bool( | |
| summary["case_count"] == expected_shallow_count + expected_escalation_count | |
| and summary["shallow_case_count"] == expected_shallow_count | |
| and summary["escalation_case_count"] == expected_escalation_count | |
| and summary["shallow_route_correct"] == expected_shallow_count | |
| and summary["escalation_route_correct"] >= min_escalation_correct | |
| and summary["parse_success_count"] == summary["case_count"] | |
| and summary["cases_with_exactly_one_call"] == summary["case_count"] | |
| and summary["logical_model_call_count"] == summary["case_count"] | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiq_agent/evaluation/shallow_escalation.py` around lines 108 - 119,
Update acceptance_passed to avoid hardcoded total, shallow, and escalation case
counts; derive the expected counts from the loaded case-list partitions or pass
them explicitly as parameters. Ensure the acceptance checks compare summary
values against those derived expectations while preserving the existing routing,
parsing, call-count, and minimum escalation correctness thresholds.
1 participant
Overview
Replace shallow research's response-keyword escalation heuristic with one conservative, bounded assessment after a successful shallow answer when the existing
enable_escalationsetting is enabled.The assessment is intentionally a routing decision, not another research pass:
sufficient,material_gap, ormaterial_conflictresult. Invalid, empty, timed-out, or malformed responses fail closed and preserve the shallow answer.This keeps
enable_escalationas the single feature control. The trade-off is one additional model call for each successful shallow answer while that setting is enabled. In return, escalation no longer depends on accidental wording in the answer, and the decision remains cheap relative to starting a deep workflow. Classification is still probabilistic, so deterministic capability/source guards and fail-closed behavior favor precision over recall.The PR also adds a reusable 20-case fixture set, deterministic parser/routing/integration coverage, an opt-in assessment-only evaluator, and architecture/configuration documentation.
Validation
Deterministic repository checks:
The 20 table-driven cases contain 10 answers that must remain shallow and 10 shallow attempts that should escalate. Tests isolate the post-shallow decision and verify one assessment call on success, zero calls when disabled or when shallow execution fails, no bound/executed tools, clarifier-first deep routing, legacy-keyword removal, budget behavior, and fail-closed parsing.
Opt-in live validation (requires
NVIDIA_API_KEYindeploy/.env; this runs assessments only—no search tools or deep workflows):Final baseline using one model,
nvidia/nemotron-3-super-120b-a12b:20 assessment calls; 0 tool calls; 0 deep workflows.
10/10 shallow cases remained shallow; 10/10 deep-worthy cases escalated.
20/20 routing decisions correct; 20/20 responses parsed successfully.
19/20 exact status labels (one conflict was conservatively labeled a material gap, with the correct escalation route).
15,228 input tokens; 1,269 output tokens; p50 1.314 seconds; p95 4.0072 seconds.
I ran the relevant local checks or explained why they are not applicable.
I added or updated tests for behavior changes.
I updated documentation for user-facing or contributor-facing changes.
I confirmed this PR does not include secrets, credentials, or internal-only data.
I certify this contribution under the Developer Certificate of Origin (DCO) and signed my commits with
git commit -sor an equivalent sign-off.Where should reviewers start?
src/aiq_agent/agents/shallow_researcher/escalation.pyfor assessment validation, deterministic guards, bounded invocation, and fail-closed behavior.src/aiq_agent/agents/chat_researcher/agent.pyfor the parent-graph routing policy.tests/aiq_agent/agents/shallow_researcher/fixtures/escalation_cases.jsonand the escalation tests for the acceptance contract.Related Issues
Summary by CodeRabbit
New Features
Bug Fixes