docs/source/improve-workflows/profiler.md (1)

74-74: Pre-existing: capitalize Toolkit in body text.

"the toolkit will automatically apply" → "the Toolkit will automatically apply". This is a pre-existing occurrence not introduced by this PR, but since the file is being modified it is worth cleaning up.

✏️ Proposed fix

-Simply specify which NeMo Agent Toolkit supported frameworks you will be using anywhere in your workflow (including tools) upon registration and the toolkit will automatically apply the appropriate profiling decorators at build time.
+Simply specify which NeMo Agent Toolkit supported frameworks you will be using anywhere in your workflow (including tools) upon registration and the Toolkit will automatically apply the appropriate profiling decorators at build time.

Based on learnings: "Starting from commit 37d9f33 (Jan 29, 2026), the lowercase 'toolkit' should no longer be used in body text, titles, or headings."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@docs/source/improve-workflows/profiler.md` at line 74, The sentence
containing "the toolkit will automatically apply" should use the capitalized
form "Toolkit" per the new style rule; update that exact phrase to "the Toolkit
will automatically apply" (and search the surrounding body text in this document
for other lowercase "toolkit" occurrences and capitalize them) so all body text,
titles, and headings follow the new convention introduced in commit 37d9f337.

packages/nvidia_nat_core/src/nat/builder/context.py (1)

392-399: Add a Returns: section to the docstring to comply with Google-style guidelines.

The sibling latency_sensitivity property (lines 379–390) includes a Returns: block; this new property should match for consistency.

📝 Proposed docstring fix

     `@property`
     def has_manual_latency_sensitivity(self) -> bool:
-        """True if any `@latency_sensitive` decorator is active in the current scope.
-
-        The default stack is [2] (length 1). Any push_latency_sensitivity call
-        adds to the stack, making length > 1.
-        """
+        """Returns `True` if any `@latency_sensitive` decorator is active in the current scope.
+
+        The default stack is `[2]` (length 1). Any `push_latency_sensitivity` call
+        adds to the stack, making length > 1.
+
+        Returns:
+            bool: `True` when at least one manual latency sensitivity push is active,
+                `False` otherwise.
+        """
         return len(self._context_state.latency_sensitivity_stack.get()) > 1

As per coding guidelines: "Provide Google-style docstrings for every public module, class, function and CLI command."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_core/src/nat/builder/context.py` around lines 392 - 399,
The docstring for the property has_manual_latency_sensitivity should include a
Google-style "Returns:" section like its sibling latency_sensitivity; update the
docstring in the has_manual_latency_sensitivity property to add a "Returns:"
block that states it returns a bool indicating whether any `@latency_sensitive`
decorator is active in the current scope, referencing the property name
has_manual_latency_sensitivity and keeping the existing brief description
intact.

packages/nvidia_nat_core/tests/nat/builder/test_context.py (1)

20-62: Extract the repeated ContextState reset into a pytest fixture.

The 3-line setup block (state = ContextState.get(); state._latency_sensitivity_stack.set(None); ctx = Context.get()) is copied into all four tests. Per coding guidelines this should live in a fixture, which also provides proper teardown isolation.

♻️ Proposed fixture refactor

+import pytest
+
 from nat.builder.context import Context
 from nat.builder.context import ContextState
 
 
+@pytest.fixture(name="fresh_ctx")
+def fixture_fresh_ctx():
+    """Yield a Context with a clean latency sensitivity stack."""
+    state = ContextState.get()
+    token = state._latency_sensitivity_stack.set(None)
+    try:
+        yield Context.get()
+    finally:
+        state._latency_sensitivity_stack.reset(token)
+
+
-def test_has_manual_latency_sensitivity_false_by_default():
+def test_has_manual_latency_sensitivity_false_by_default(fresh_ctx):
     """Default stack [2] means no manual decorator is active."""
-    state = ContextState.get()
-    # Reset to ensure fresh state
-    state._latency_sensitivity_stack.set(None)
-
-    ctx = Context.get()
-    assert ctx.has_manual_latency_sensitivity is False
+    assert fresh_ctx.has_manual_latency_sensitivity is False
 
 
-def test_has_manual_latency_sensitivity_true_when_pushed():
+def test_has_manual_latency_sensitivity_true_when_pushed(fresh_ctx):
     """After push_latency_sensitivity, a manual decorator is active."""
-    state = ContextState.get()
-    state._latency_sensitivity_stack.set(None)
-
-    ctx = Context.get()
-    with ctx.push_latency_sensitivity(5):
-        assert ctx.has_manual_latency_sensitivity is True
+    with fresh_ctx.push_latency_sensitivity(5):
+        assert fresh_ctx.has_manual_latency_sensitivity is True

Apply the same pattern to the remaining two tests.

As per coding guidelines: "Any frequently repeated code should be extracted into pytest fixtures" and "Pytest fixtures should define the name argument … fixture_ prefix."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_core/tests/nat/builder/test_context.py` around lines 20 -
62, Extract the repeated setup into a pytest fixture named fixture_reset_context
that gets the ContextState singleton, resets state._latency_sensitivity_stack to
None, yields a fresh Context via Context.get(), and then resets
state._latency_sensitivity_stack to None again for teardown; update each test
(test_has_manual_latency_sensitivity_false_by_default,
test_has_manual_latency_sensitivity_true_when_pushed,
test_has_manual_latency_sensitivity_false_after_pop,
test_has_manual_latency_sensitivity_nested) to accept the fixture_reset_context
argument and use the yielded Context instead of creating ContextState/Context
inline.

packages/nvidia_nat_core/src/nat/profiler/prediction_trie/trie_builder.py (1)

29-36: SensitivityConfig docstring is missing a Google-style Attributes: section.

The class docstring is present but the three fields (scale, w_critical, w_fanout, w_position) have no documentation, violating the requirement for Google-style docstrings on every public API.

♻️ Suggested docstring update

 `@dataclass`
 class SensitivityConfig:
-    """Configuration for auto-sensitivity scoring."""
+    """Configuration for auto-sensitivity scoring.
+
+    Attributes:
+        scale: Integer upper bound for sensitivity scores; must be >= 1.
+        w_critical: Weight applied to the critical-path signal (call_duration / workflow_duration).
+        w_fanout: Weight applied to the downstream fan-out signal (remaining_calls / max_remaining).
+        w_position: Weight applied to the U-shaped call-position signal.
+    """

As per coding guidelines, "Provide Google-style docstrings for every public module, class, function and CLI command."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_core/src/nat/profiler/prediction_trie/trie_builder.py`
around lines 29 - 36, The SensitivityConfig dataclass docstring lacks a
Google-style Attributes: section documenting its public fields; update the
docstring for SensitivityConfig to include an "Attributes:" block that briefly
describes each field (scale, w_critical, w_fanout, w_position), their types and
their default/meaning so it conforms to the project Google-style docstring
requirement and documents the public API.

packages/nvidia_nat_core/src/nat/data_models/profiler.py (1)

43-50: New PredictionTrieConfig fields lack docstring Attributes: section.

None of the five new public fields (auto_sensitivity, sensitivity_scale, w_critical, w_fanout, w_position) are documented in the class docstring or via Field(description=…). The same weight defaults are also silently duplicated from SensitivityConfig, creating a DRY risk if they diverge.

♻️ Suggested improvement

 class PredictionTrieConfig(BaseModel):
+    """Configuration for the prediction trie profiler component.
+
+    Attributes:
+        enable: Whether to enable prediction trie building.
+        output_filename: Filename for the serialized trie JSON.
+        auto_sensitivity: Automatically compute latency sensitivity scores during profiling.
+        sensitivity_scale: Integer scale for sensitivity scores; must be >= 1.
+        w_critical: Weight for the critical-path signal in [0, 1].
+        w_fanout: Weight for the downstream fan-out signal in [0, 1].
+        w_position: Weight for the U-shaped call-position signal in [0, 1].
+    """
     enable: bool = False
     output_filename: str = "prediction_trie.json"
-    auto_sensitivity: bool = True
-    sensitivity_scale: int = 5
-    w_critical: float = 0.5
-    w_fanout: float = 0.3
-    w_position: float = 0.2
+    auto_sensitivity: bool = True
+    sensitivity_scale: int = Field(default=5, ge=1, description="Integer scale for sensitivity scores.")
+    w_critical: float = Field(default=0.5, ge=0.0, description="Weight for critical-path signal.")
+    w_fanout: float = Field(default=0.3, ge=0.0, description="Weight for fan-out signal.")
+    w_position: float = Field(default=0.2, ge=0.0, description="Weight for call-position signal.")

As per coding guidelines, "Provide Google-style docstrings for every public module, class, function and CLI command" and "All public APIs require Python 3.11+ type hints on parameters and return values."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_core/src/nat/data_models/profiler.py` around lines 43 -
50, The new PredictionTrieConfig class has five undocumented public fields
(auto_sensitivity, sensitivity_scale, w_critical, w_fanout, w_position) and
duplicates default weight values from SensitivityConfig; update the
PredictionTrieConfig docstring to include a Google-style "Attributes:" section
describing each field (purpose, types, and default), and/or add
Field(description=...) to each Pydantic field; to avoid DRY, reference or import
the canonical defaults from SensitivityConfig (or a shared constant) instead of
hardcoding the same numeric defaults in PredictionTrieConfig so the defaults
remain single-source-of-truth.

packages/nvidia_nat_core/tests/nat/profiler/prediction_trie/test_trie_builder.py (1)

144-162: Tests exercise private method _extract_llm_contexts directly.

Both test_extract_contexts_include_call_duration and test_extract_contexts_include_workflow_duration call builder._extract_llm_contexts(simple_trace). Testing private methods couples the tests to internal implementation details; any rename or signature change breaks the tests without breaking public behavior.

The same assertions can be reached through the public add_trace → build() path by checking the computed call_duration_s effect on latency_sensitivity values, or by exposing the needed invariants via a thin public helper if they're genuinely important to contract-test.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@packages/nvidia_nat_core/tests/nat/profiler/prediction_trie/test_trie_builder.py`
around lines 144 - 162, Tests currently call the private method
_extract_llm_contexts on PredictionTrieBuilder; replace this with the public API
by feeding the trace through add_trace(simple_trace) and then calling build()
(or otherwise using the builder's public methods) to obtain the resulting
contexts/results and assert on call_duration_s/workflow_duration_s effects
(e.g., via the produced LLMCallContext-like objects or latency_sensitivity
values); alternatively, if the duration values must be directly asserted, add a
small public helper method (e.g., extract_llm_contexts or get_llm_contexts) on
PredictionTrieBuilder that wraps the private implementation, then update tests
to call that public method instead of _extract_llm_contexts to avoid coupling to
the private implementation.

docs/source/improve-workflows/profiler.md (1)

219-222: Inconsistent formatting: `Interarrival` uses backticks while sibling labels do not.

Interarrival is a plain English term in this context, not a code identifier. Wrapping it in backticks is inconsistent with the surrounding items.

✏️ Proposed fix

-- **`Interarrival` time**: Expected time in milliseconds until the next LLM call.
+- **Interarrival time**: Expected time in milliseconds until the next LLM call.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@docs/source/improve-workflows/profiler.md` around lines 219 - 222, The list
item uses backticks around `Interarrival` which is inconsistent with the other
plain English labels; update the bullet label from "`Interarrival` time" to
"Interarrival time" (leave the surrounding items "Remaining calls", "Output
tokens", and "Latency sensitivity" unchanged) so all sibling labels share the
same plain-text formatting.

packages/nvidia_nat_core/src/nat/profiler/prediction_trie/trie_builder.py (1)

29-36: Consider adding __post_init__ validation for SensitivityConfig.

There are no guards against invalid inputs (e.g., sensitivity_scale <= 0 or negative weights). A sensitivity_scale of 0 would make _score_to_sensitivity always return 1, and negative weights would invert signal contributions. A lightweight __post_init__ would catch misconfigurations early.

Proposed validation

 `@dataclass`
 class SensitivityConfig:
     """Configuration for auto-sensitivity scoring."""
 
     sensitivity_scale: int = 5
     w_critical: float = 0.5
     w_fanout: float = 0.3
     w_position: float = 0.2
+
+    def __post_init__(self) -> None:
+        if self.sensitivity_scale < 1:
+            raise ValueError(f"sensitivity_scale must be >= 1, got {self.sensitivity_scale}")
+        if any(w < 0 for w in (self.w_critical, self.w_fanout, self.w_position)):
+            raise ValueError("Sensitivity weights must be non-negative")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_core/src/nat/profiler/prediction_trie/trie_builder.py`
around lines 29 - 36, Add a __post_init__ to SensitivityConfig that validates
inputs: ensure sensitivity_scale is an int > 0, each weight (w_critical,
w_fanout, w_position) is >= 0, and that the sum of the weights is > 0 (or
normalize the weights to sum to 1 if you prefer consistent scoring). If any
check fails, raise a ValueError with a clear message. Implement this in the
SensitivityConfig dataclass so _score_to_sensitivity and other consumers get
guaranteed valid parameters.

packages/nvidia_nat_eval/src/nat/plugins/eval/profiler/profile_runner.py (1)

282-294: Inconsistent import path for SensitivityConfig.

PredictionTrieBuilder and save_prediction_trie are imported from the public package (nat.profiler.prediction_trie), but SensitivityConfig is imported directly from the internal module (nat.profiler.prediction_trie.trie_builder). For consistency and to reduce coupling to internal module structure, consider re-exporting SensitivityConfig from packages/nvidia_nat_core/src/nat/profiler/prediction_trie/__init__.py and importing it from the public package instead.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_eval/src/nat/plugins/eval/profiler/profile_runner.py`
around lines 282 - 294, The imports are inconsistent: PredictionTrieBuilder and
save_prediction_trie come from nat.profiler.prediction_trie but
SensitivityConfig is pulled from the internal trie_builder module; re-export
SensitivityConfig from the public package (add SensitivityConfig to
nat.profiler.prediction_trie.__init__.py exports) and then update the import in
profile_runner.py to import SensitivityConfig from nat.profiler.prediction_trie
alongside PredictionTrieBuilder and save_prediction_trie so all three symbols
come from the public API.