Add trace-based behavioral tests with Monocle Test Tools

.gitignore

@@ -40,6 +40,9 @@ env/
 .env
 .env.local
 .env.*.local
+
+# Monocle traces
+.monocle/
 # Logs
 *.log
 logs/

configs/config_openai_cli.yml

@@ -0,0 +1,129 @@
+# This is the default configuration for the CLI mode.
+# It has the following features:
+# - Human-in-the-loop clarification and plan approval before deep research starts
+# - Web search and Paper search tools by default
+# - There is no knowledge retrieval
+
+general:
+  telemetry:
+    logging:
+      console:
+        _type: console
+        level: INFO
+    # tracing:
+    #   langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=<your-langsmith-api-key>`
+    #     _type: langsmith
+    #     project: nvidia-aiq
+
+llms:
+  nemotron_llm_intent:
+    _type: openai
+    model_name: gpt-4o
+    api_key: ${OPENAI_API_KEY}
+    temperature: 0.5
+    top_p: 0.9
+    max_tokens: 4096
+    num_retries: 5
+
+  nemotron_super_llm:
+    _type: openai
+    model_name: gpt-4o
+    api_key: ${OPENAI_API_KEY}
+    temperature: 0.7
+    top_p: 0.7
+    max_tokens: 16384
+    num_retries: 5
+
+functions:
+  # =========================================================================
+  # Data Source Registry
+  # =========================================================================
+  # Central registry that controls:
+  #   1. UI toggles — each source appears as an on/off switch in the frontend
+  #   2. Per-message filtering — users can select active sources per request
+  #   3. Tool auto-inheritance — agents with no explicit `tools` list receive
+  #      every tool listed here (use `exclude_tools` on agents to specialize)
+  #
+  # Source entry fields:
+  #   id, name, description, tools, requires_auth (default: false),
+  #   default_enabled (default: true)
+  #
+  # See docs/source/customization/tools-and-sources.md for full details.
+  # =========================================================================
+  data_sources:
+    _type: data_source_registry
+    sources:
+      - id: web_search
+        name: "Web Search"
+        description: "Search the web for real-time information."
+        tools:
+          - web_search_tool
+          - advanced_web_search_tool
+      # - id: paper_search
+      #   name: "Academic Papers"
+      #   description: "Search academic papers and scientific publications."
+      #   tools:
+      #     - paper_search_tool
+
+  web_search_tool:
+    _type: exa_web_search
+    max_results: 5
+    max_content_length: 1000
+
+  advanced_web_search_tool:
+    _type: exa_web_search
+    max_results: 2
+
+  # Paper Search (optional - requires SERPER_API_KEY)
+  # Uncomment the block below and set SERPER_API_KEY to enable academic paper search.
+  # paper_search_tool:
+  #   _type: paper_search
+  #   max_results: 5
+  #   serper_api_key: ${SERPER_API_KEY}
+
+  # =========================================================================
+  # Agents — inherit all registry tools; use exclude_tools to specialize
+  # =========================================================================
+  intent_classifier:
+    _type: intent_classifier
+    llm: nemotron_llm_intent
+    # tools: omitted -> inherits all from data_source_registry
+    # exclude_tools: []
+    # llm_timeout: 90  # optional; seconds for intent LLM call (default 90)
+
+  clarifier_agent:
+    _type: clarifier_agent
+    llm: nemotron_super_llm
+    planner_llm: nemotron_super_llm
+    # tools: omitted -> inherits all from data_source_registry
+    # exclude_tools: []
+    max_turns: 3
+    enable_plan_approval: true
+    log_response_max_chars: 2000
+    verbose: true
+
+  shallow_research_agent:
+    _type: shallow_research_agent
+    llm: nemotron_super_llm
+    # tools: omitted -> inherits all from data_source_registry
+    exclude_tools:
+      - advanced_web_search_tool
+    max_llm_turns: 10
+    max_tool_iterations: 5
+
+  deep_research_agent:
+    _type: deep_research_agent
+    enable_citation_verification: true
+    orchestrator_llm: nemotron_super_llm
+    source_router_llm: nemotron_super_llm
+    researcher_llm: nemotron_super_llm
+    planner_llm: nemotron_super_llm
+    writer_llm: nemotron_super_llm
+    exclude_tools:
+      - web_search_tool
+
+workflow:
+  _type: chat_deepresearcher_agent
+  enable_escalation: true
+  enable_clarifier: true
+  checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db}

tests/monocle/README.md

@@ -0,0 +1,73 @@
+# AI-Q behavioural tests (Monocle Test Tools)
+
+Trace-based tests that lock in the AI-Q research agent's behaviour. AI-Q runs on
+the NeMo Agent Toolkit (NAT); Monocle records each run as a structured trace --
+the agent invocation, every tool call, and timings -- and each test asserts
+against that trace: which agent ran, which tools it called, what it was asked,
+what it produced, and its duration cost. A later prompt, model, or tool change
+that regresses the behaviour fails here.
+
+## Layout
+
+- `test_nvidiaaiq.py` — the suite: two offline tests + one live test
+- `conftest.py` — Monocle setup, `.env` loading, and `run_nvidiaaiq()`
+- `traces/` — recorded good-trace fixtures the offline tests replay
+- `requirements.txt` — dependencies
+
+## Tests
+
+| Test | Scenario | What it shows |
+|---|---|---|
+| `test_capabilities_intro` | "Hi, what can you do?" | direct answer, a `does_not_call_tool` negative, budget |
+| `test_nvda_stock_lookup` | Current NVIDIA stock price | `web_search_tool` call, input/output, budget |
+| `test_nvda_stock_lookup_live` | The stock question, run live | live run, web-search path, structure + budget |
+
+The offline tests replay recorded traces with duration budgets measured from
+those runs (rounded up with headroom). The live test drives the agent
+end-to-end and asserts structure and budget only, since the output legitimately
+varies run to run.
+
+Note on budgets: NAT traces do not carry token metadata -- the `inference.*`
+spans record only finish reasons, no token counts -- so `under_token_limit(...)`
+would sum 0 and always pass. This suite omits it and budgets on
+`under_duration(..., span_type="workflow")` instead. The agent name asserted
+(`LangGraph`) is the real `entity.1.name` on the `agentic.invocation` spans (NAT
+runs its research agents on a LangGraph runtime); the web tool is
+`web_search_tool`.
+
+## Run
+
+```bash
+pip install -r requirements.txt
+pytest tests/monocle/ -k "not live"   # offline, no network, no keys
+```
+
+The live test is opt-in (`RUN_LIVE_NVIDIAAIQ=1`) and skipped by default.
+NAT binds its async singletons to the first event loop, so only one in-process
+live run works per process, and NAT leaves non-daemon threads that keep the
+interpreter from exiting cleanly. So run the live test in its own process, with
+keys in `deploy/.env` (`OPENAI_API_KEY` plus a search key, `EXA_API_KEY`/`SERPER_API_KEY`):
+
+```bash
+RUN_LIVE_NVIDIAAIQ=1 pytest tests/monocle/ -k nvda_stock_lookup_live -s
+```
+
+They drive the workflow via `configs/config_openai_cli.yml` with an
+auto-approving `user_input_callback`, so a human-in-the-loop clarification or
+plan-approval interrupt is answered automatically instead of blocking the run.
+
+## Add your own test
+
+1. Run AI-Q under Monocle and capture a trace of a run you're happy with
+   (Monocle writes trace JSON to `.monocle/` by default).
+2. Move it into `traces/` and load it with
+   `monocle_trace_asserter.with_trace_source("file", trace_path=path)`.
+3. Assert with the fluent API — `called_agent(...)`, `called_tool(...)`,
+   `contains_input/output(...)`, `under_duration(..., span_type="workflow")` —
+   then add it alongside the others.
+
+## Evaluations (optional)
+
+Each test carries a commented-out `check_eval("hallucination", ...)` chain.
+Monocle can run evaluation checks against a trace; set `OKAHU_API_KEY` and
+uncomment to enable.

tests/monocle/conftest.py

@@ -0,0 +1,57 @@
+"""Pytest scaffold for the AI-Q (NeMo Agent Toolkit) Monocle test suite.
+
+Enables Monocle tracing, loads the repo `deploy/.env`, and exposes
+``run_nvidiaaiq`` -- the single entry the live tests use to drive the agent
+under instrumentation and return its final answer text.
+"""
+import asyncio
+import os
+import uuid
+from pathlib import Path
+
+from dotenv import load_dotenv
+from monocle_apptrace import setup_monocle_telemetry
+
+HERE = Path(__file__).resolve().parent
+TRACES = HERE / "traces"
+REPO_ROOT = HERE.parent.parent
+
+# This repo runs pytest with `--import-mode=importlib`, so a test's directory is
+# not auto-added to sys.path and `import conftest` would load a *second* copy of
+# this module -- double-instrumenting the tracer. Guard the one-time setup.
+if not os.environ.get("_NVIDIAAIQ_MONOCLE_READY"):
+    setup_monocle_telemetry(workflow_name="nvidia-aiq")
+    load_dotenv(REPO_ROOT / "deploy" / ".env")
+    os.environ["_NVIDIAAIQ_MONOCLE_READY"] = "1"
+
+# The CLI-mode workflow used for live runs. It exposes the web/paper search
+# tools and a human-in-the-loop clarification + plan-approval step.
+LIVE_CONFIG = REPO_ROOT / "configs" / "config_openai_cli.yml"
+
+
+def run_nvidiaaiq(message: str) -> str:
+    """Run the AI-Q (NAT) workflow once and return its final response text.
+
+    NAT's CLI workflow can raise a human-in-the-loop interrupt (clarification or
+    plan approval). We pass an auto-approving ``user_input_callback`` so the run
+    never blocks waiting on a human -- a short affirmative satisfies both the
+    clarifier and plan approval.
+    """
+    from nat.builder.context import ContextState
+    from nat.data_models.interactive import HumanResponseText
+    from nat.runtime.loader import load_workflow
+
+    async def _auto_approve(prompt):
+        return HumanResponseText(text="Yes, proceed with the research plan.")
+
+    async def _run() -> str:
+        async with load_workflow(str(LIVE_CONFIG)) as session_manager:
+            try:
+                ContextState.get().conversation_id.set(str(uuid.uuid4()))
+            except Exception:
+                pass
+            async with session_manager.session(user_input_callback=_auto_approve) as session:
+                async with session.run(message) as runner:
+                    return await runner.result(to_type=str)
+
+    return asyncio.run(_run())

tests/monocle/requirements.txt

@@ -0,0 +1,5 @@
+# Installing monocle_test_tools pulls in everything this suite needs
+# (pytest, pytest-asyncio, and monocle_apptrace come transitively).
+monocle_test_tools
+# Auto-loads the repo deploy/.env for the live tests.
+python-dotenv

tests/monocle/test_nvidiaaiq.py

@@ -0,0 +1,109 @@
+"""Trace-based behavioural tests for AI-Q (NeMo Agent Toolkit), using Monocle
+Test Tools.
+
+Each test asserts against the Monocle trace a run emits -- which agent ran,
+which tools it called, what it was asked, what it produced, and its duration
+cost. Two offline tests replay recorded good traces (fast, no keys); one live
+test runs the agent end-to-end.
+
+    pytest tests/monocle/ -k "not live"   # offline, no keys
+    RUN_LIVE_NVIDIAAIQ=1 pytest tests/monocle/ -k live -s   # live (needs OPENAI + a search key)
+
+NOTE ON TOKENS: no `under_token_limit` here. NAT calls the OpenAI models in
+streaming mode without `stream_options.include_usage`, so the response carries
+no usage and the `inference.*` spans record only {finish_reason, finish_type}.
+This is upstream (NAT), not a Monocle gap -- Monocle captures token counts
+whenever the response includes usage. So the suite budgets on
+`under_duration(..., span_type="workflow")` instead. The agent name is
+"LangGraph" (NAT runs on a LangGraph runtime); the web tool is `web_search_tool`.
+"""
+import os
+import sys
+from pathlib import Path
+
+import pytest
+from monocle_test_tools import TraceAssertion
+
+# This repo uses `--import-mode=importlib`, which does not put a test's own
+# directory on sys.path; make the sibling conftest importable by name.
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from conftest import TRACES, run_nvidiaaiq  # noqa: E402
+
+# Recorded good traces (captured from this repo under monocle_apptrace 0.8.8).
+# Q: "Hi, what can you do?" -- meta/capabilities, answered directly (no tools).
+TRACE_INTRO = str(TRACES / "monocle_trace_nvidia-aiq_4a94876861e3ade549dc74fb785e94d9_2026-07-09_13.17.46.json")
+# Q: "What is the current stock price of NVIDIA today?" -- live-fact lookup that
+# calls web_search_tool once.
+TRACE_STOCK = str(TRACES / "monocle_trace_nvidia-aiq_e2961458302a780a27f91a0d86b290dd_2026-07-09_13.25.14.json")
+
+
+# --- Offline: replay recorded good traces ---------------------------------
+
+def test_capabilities_intro(monocle_trace_asserter: TraceAssertion):
+    """Meta/capabilities question. The orchestrator answers directly (routes it
+    as a meta_response) -- no tool is called. Real trace: 5 spans, ~2.0s."""
+    monocle_trace_asserter.with_trace_source("file", trace_path=TRACE_INTRO)
+
+    monocle_trace_asserter.called_agent("LangGraph").contains_input("what can you do")
+    monocle_trace_asserter.contains_output("AI Research Assistant")
+    monocle_trace_asserter.does_not_call_tool("web_search_tool", "LangGraph")
+    monocle_trace_asserter.under_duration(10, span_type="workflow")
+
+    # Eval layer (deferred -- set OKAHU_API_KEY and uncomment to enable):
+    # monocle_trace_asserter.with_evaluation("okahu").check_eval("hallucination", "no_hallucination") \
+    #     .check_eval("contextual_precision", "high_precision") \
+    #     .check_eval("sentiment", "positive") \
+    #     .check_eval("bias", "unbiased")
+
+
+def test_nvda_stock_lookup(monocle_trace_asserter: TraceAssertion):
+    """Live-fact lookup. The agent calls web_search_tool once and reports the
+    price it found ($202.78, the captured value). Real trace: 11 spans, ~7.7s."""
+    monocle_trace_asserter.with_trace_source("file", trace_path=TRACE_STOCK)
+
+    monocle_trace_asserter.called_agent("LangGraph").contains_input("stock price of NVIDIA")
+    monocle_trace_asserter.contains_output("202.78")
+    monocle_trace_asserter.called_tool("web_search_tool", "LangGraph")
+    monocle_trace_asserter.under_duration(15, span_type="workflow")
+
+    # monocle_trace_asserter.with_evaluation("okahu").check_eval("hallucination", "no_hallucination") \
+    #     .check_eval("contextual_precision", "high_precision") \
+    #     .check_eval("sentiment", "positive") \
+    #     .check_eval("bias", "unbiased")
+
+
+# --- Live: run the agent end-to-end ---------------------------------------
+# The web-search fact lookup (web_search_tool over the configured provider).
+# Output text varies run to run, so it asserts structure + budget with
+# contains_any_output kept phrasing-robust.
+#
+# Opt-in (RUN_LIVE_NVIDIAAIQ=1), one run per process: NAT binds its async
+# singletons (module-level locks) to the first event loop, so a second
+# asyncio.run() in the same process raises "bound to a different event loop",
+# and NAT leaves non-daemon worker threads that keep the interpreter from
+# exiting cleanly. So a default `pytest tests/monocle/` skips this (offline
+# stays green and exits clean); run it explicitly, e.g.:
+#     RUN_LIVE_NVIDIAAIQ=1 pytest tests/monocle/ -k live -s
+_LIVE = pytest.mark.skipif(
+    os.environ.get("RUN_LIVE_NVIDIAAIQ") != "1",
+    reason="opt-in live run (set RUN_LIVE_NVIDIAAIQ=1; run one live test per process -- see note above)",
+)
+
+
+@_LIVE
+def test_nvda_stock_lookup_live(monocle_trace_asserter: TraceAssertion):
+    """Web-search path: the NVIDIA stock question, run live (web_search_tool)."""
+    monocle_trace_asserter.validator.test_workflow(
+        run_nvidiaaiq,
+        {"test_input": ("What is the current stock price of NVIDIA today?",)},
+    )
+
+    monocle_trace_asserter.called_agent("LangGraph")
+    monocle_trace_asserter.contains_any_output("NVIDIA", "NVDA", "stock", "price")
+    monocle_trace_asserter.called_tool("web_search_tool", "LangGraph")
+    monocle_trace_asserter.under_duration(120, span_type="workflow")
+
+    # monocle_trace_asserter.with_evaluation("okahu").check_eval("hallucination", "no_hallucination") \
+    #     .check_eval("contextual_precision", "high_precision") \
+    #     .check_eval("sentiment", "positive") \
+    #     .check_eval("bias", "unbiased")