Implement Function Policy component and DynamicFunctionMiddleware #1216
Conversation
…tion Signed-off-by: Eric Evans <[email protected]>
ericevans-nv
added
feature request
WalkthroughAdds a Function Policies subsystem (interfaces, configs, logging policy, registration and CLI), builder and type-registry support, DynamicFunctionMiddleware for runtime discovery/wrapping, generalizes middleware signatures to variadic *args/**kwargs, refactors cache middleware into its own package, updates docs, and adds extensive tests. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Builder
participant DynamicMiddleware
participant Registry
participant PrePolicy as Pre-Invoke Policy
participant TargetFn as Target Function
participant PostPolicy as Post-Invoke Policy
Client->>Builder: retrieve or call function
Builder->>DynamicMiddleware: intercept / wrap function (patch)
DynamicMiddleware->>Registry: discover & load policies
Registry-->>DynamicMiddleware: list of configured policies
Client->>TargetFn: call wrapped_func(*args, **kwargs)
wrappedFunc->>DynamicMiddleware: run middleware chain
DynamicMiddleware->>PrePolicy: on_pre_invoke(context)
PrePolicy-->>DynamicMiddleware: (maybe) modified args
DynamicMiddleware->>TargetFn: call_next(*modified_args, **kwargs)
TargetFn-->>DynamicMiddleware: output
DynamicMiddleware->>PostPolicy: on_post_invoke(context_with_output)
PostPolicy-->>DynamicMiddleware: (maybe) modified output
DynamicMiddleware-->>Client: final_output
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Areas requiring special attention:
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…eature/dynamic-middleware-function-policies
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (20)
src/nat/function_policy/__init__.py (1)
21-26: Sort__all__alphabetically.The
__all__list should be sorted alphabetically per the static analysis hint (RUF022).Apply this diff to sort the exports:
__all__ = [ "FunctionPolicyBase", "PolicyConfigT", - "PreInvokeContext", "PostInvokeContext", + "PreInvokeContext", ]src/nat/middleware/register.py (1)
27-38: Silence Ruff ARG001 for unusedbuilderwithout changing behavior.
builderis intentionally unused here but required by theregister_middlewarecomponent pattern, so Ruff’s ARG001 is expected. Easiest fix is to rename the parameter to start with an underscore so tools treat it as intentionally unused:-@register_middleware(config_type=CacheMiddlewareConfig) -async def cache_middleware(config: CacheMiddlewareConfig, builder: Builder): +@register_middleware(config_type=CacheMiddlewareConfig) +async def cache_middleware(config: CacheMiddlewareConfig, _builder: Builder):This keeps the calling convention (positional args) intact and aligns with linting guidelines.
src/nat/cli/commands/info/list_function_policies.py (1)
19-160: CLI for function policies is coherent and matches registry semantics.The
function-policiesandfunction-policycommands correctly:
- force plugin discovery before querying the registry,
- handle the “no policies” case with a clear message,
- and surface useful metadata and config schema details when available.
If you find yourself adding more policy-related commands later, consider extracting the common “discover plugins + fetch policy list” logic into a small helper to avoid duplication, but it’s not necessary for this PR.
docs/source/reference/function-policies.md (4)
22-23: Use full toolkit name on first reference.Per coding guidelines, the first reference should use "NVIDIA NeMo Agent toolkit" before subsequent abbreviated uses.
-Function policies provide a standardized way to intercept and modify function calls in the NeMo Agent toolkit. They work by wrapping functions with hooks that execute before (`on_pre_invoke`) and after (`on_post_invoke`) the function runs. +Function policies provide a standardized way to intercept and modify function calls in the NVIDIA NeMo Agent toolkit. They work by wrapping functions with hooks that execute before (`on_pre_invoke`) and after (`on_post_invoke`) the function runs.Based on coding guidelines.
164-172: Add language specifier to fenced code block.Static analysis flagged this as missing a language specifier. Since this is a text-based flow diagram, use
textas the language identifier.-``` +```text Original Input (10) ↓ Policy 1: Add 5 → Returns 15 ↓ Policy 2: Multiply by 2 → Returns 30 ↓ Function receives: 30--- `180-188`: **Add language specifier to fenced code block.** Same issue - use `text` for this ASCII flow diagram. ```diff -``` +```text Function returns: 100 ↓ Policy 1: Log result → Returns 100 ↓ Policy 2: Round to nearest 10 → Returns 100 ↓ User receives: 100--- `463-480`: **Add language specifier to ASCII diagram.** Use `text` for consistency with other ASCII diagrams. ```diff -``` +```text ┌─────────────────────────────────────────┐ │ Dynamic Middleware │src/nat/function_policy/interface.py (3)
88-103: Remove extra blank lines in docstring.There are extra blank lines before the closing triple quotes that should be removed for cleaner formatting.
Returns: Transformed input value, or None to preserve current value - - """ pass
105-121: Remove extra blank lines in docstring.Same issue with extra blank lines before the closing quotes.
Returns: Transformed output value, or None to preserve current value - - """ pass
123-128: Consider sorting__all__for consistency.Static analysis suggests sorting
__all__alphabetically, which is a common Python convention.__all__ = [ "FunctionPolicyBase", + "PolicyConfigT", + "PostInvokeContext", "PreInvokeContext", - "PostInvokeContext", - "PolicyConfigT", ]src/nat/middleware/dynamic_function_middleware.py (4)
100-104: Duplicate "already registered" checks in discovery methods.Each
_discover_and_register_*method checks if a component is already in the inventory, but_should_intercept_*methods also perform this same check at lines 393-394, 413-414, etc. This creates redundant checks when both code paths execute sequentially.While not a bug, this adds minor overhead on every discovery call. Consider removing the duplicate check from either location.
Also applies to: 139-140, 178-179, 216-217, 255-256, 295-296
116-119: Intentional error resilience with broad exception handling.The bare
Exceptioncatches flagged by static analysis are intentional design choices for error resilience. Per the PR description, failed policies should be logged and skipped rather than crashing the workflow. Thelogger.debugcalls provide visibility without noise.However, consider using
logger.warninginstead oflogger.debugfor registration failures, as these may indicate configuration issues worth surfacing.except Exception as e: - logger.debug("Failed to register component function '%s' on LLM '%s': %s", function_name, llm_name, e) + logger.warning("Failed to register component function '%s' on LLM '%s': %s", function_name, llm_name, e)Also applies to: 152-159, 191-198, 229-236, 269-276, 309-316
801-833: Remove unusedcomponent_typeparameter or document its future use.Static analysis correctly identifies that
component_typeis never used in_get_callable_functions. Either remove it or add a docstring note explaining if it's reserved for future use.- def _get_callable_functions(self, instance: Any, component_type: str | None = None) -> set[str]: + def _get_callable_functions(self, instance: Any) -> set[str]: """Get all callable functions from component instance that can be safely wrapped. This discovers ALL potentially wrappable component functions without allowlist filtering. Safety checks ensure only valid, callable, bound functions are included. Args: instance: The component instance to introspect - component_type: Type of component (for logging/metadata, not filtering) Returns: Set of all valid component function names that could be wrapped """If removing the parameter, also update the call sites at lines 106, 142, 181, 219, 258, 298.
889-998: Builder patching infrastructure is comprehensive but tightly coupled.The patching mechanism works correctly, but the tight coupling to specific builder methods means:
- Adding new component types requires adding new patch methods
- If builder method signatures change, this code must be updated
Consider documenting this coupling in the class docstring so future maintainers understand the relationship.
src/nat/middleware/utils/workflow_inventory.py (1)
29-88: Allowlist constant looks good; consider documenting evolution/extension strategy.The
COMPONENT_FUNCTION_ALLOWLISTSmapping cleanly captures the default surfaces perComponentGroupand aligns with the test expectations. As a minor improvement, you might add a brief module-level comment or docstring noting that these names must stay in sync with the public client APIs, and that new client methods should be added here (rather than ad hoc in middleware) so discovery and tests remain consistent.src/nat/data_models/middleware.py (1)
52-155: Dynamic middleware config and allowlist merging are consistent and safe.
AllowedComponentFunctions.merge_with_defaults()correctly unions user-provided sets withCOMPONENT_FUNCTION_ALLOWLISTS, and theDynamicMiddlewareConfig.allowed_component_functionsdocstring matches this “extend, don’t override” behavior. If you want to make this even clearer to readers, you could (optionally) note in the docstring that shrinking the default allowlist isn’t supported yet and add a-> AllowedComponentFunctionsreturn annotation on the validator, but the current implementation is sound.tests/nat/middleware/test_dynamic_middleware.py (3)
278-281: Silence unusedvalueparameters in localmock_call_nexthelpers.Ruff flags these
valueparameters as unused (ARG001). Since the tests don’t depend on them, you can mark them as intentionally unused by renaming to_:- async def mock_call_next(value): + async def mock_call_next(_): return expected_outputApply the same pattern to the other
mock_call_next(value)definitions in this file (invoke/stream, pre/post failure paths) to satisfy linting without changing behavior.Also applies to: 378-380, 442-444, 471-474, 550-552, 680-683, 783-785
811-817: Silence unused lambda parameters in disabled streaming policy helpers.The lambdas passed as
pre_invoke_modifierandpost_invoke_modifierdon’t use their arguments, triggering Ruff ARG005. You can make that explicit by naming the parameter_:- DisabledPrePolicy = create_tracking_policy("disabled_pre", - execution_order, - pre_invoke_modifier=lambda inp: {"should_not_see": "this"}) - DisabledPostPolicy = create_tracking_policy("disabled_post", - execution_order, - post_invoke_modifier=lambda out: "MODIFIED") + DisabledPrePolicy = create_tracking_policy("disabled_pre", + execution_order, + pre_invoke_modifier=lambda _: {"should_not_see": "this"}) + DisabledPostPolicy = create_tracking_policy("disabled_post", + execution_order, + post_invoke_modifier=lambda _: "MODIFIED")
1667-1707: Usellm_allowlistin the test or drop it to avoid dead code.
llm_allowlist = COMPONENT_FUNCTION_ALLOWLISTS[ComponentGroup.LLMS]is never used, triggering Ruff F841. You can turn this into a useful assertion that ties the test to the configured allowlist:- # Get the allowlist for LLMs - llm_allowlist = COMPONENT_FUNCTION_ALLOWLISTS[ComponentGroup.LLMS] - - # Verify only allowlisted functions are registered + # Get the allowlist for LLMs and ensure expected names stay covered + llm_allowlist = COMPONENT_FUNCTION_ALLOWLISTS[ComponentGroup.LLMS] + assert {"generate", "agenerate", "predict"}.issubset(llm_allowlist) + + # Verify only allowlisted functions are registeredAlternatively, if you don’t want this coupling, you can simply remove the unused variable and its import.
src/nat/cli/type_registry.py (1)
967-1041: Consider exposing function policies via the genericget_infos_by_typehelpers.
ComponentEnum.FUNCTION_POLICYis used when registering policies, butget_infos_by_type/get_registered_types_by_component_typedon’t currently handle that enum value. If any CLI or tooling later calls these helpers withComponentEnum.FUNCTION_POLICY, it will hit the genericValueError. You can wire it in similarly to other components:def get_infos_by_type(self, component_type: ComponentEnum) -> dict: @@ - if component_type == ComponentEnum.MIDDLEWARE: - return self._registered_middleware + if component_type == ComponentEnum.MIDDLEWARE: + return self._registered_middleware + + if component_type == ComponentEnum.FUNCTION_POLICY: + return self._registered_function_policies @@ def get_registered_types_by_component_type(self, component_type: ComponentEnum) -> list[str]: @@ - if component_type == ComponentEnum.TTC_STRATEGY: - return [i.static_type() for i in self._registered_ttc_strategies] + if component_type == ComponentEnum.TTC_STRATEGY: + return [i.static_type() for i in self._registered_ttc_strategies] + + if component_type == ComponentEnum.FUNCTION_POLICY: + return [i.static_type() for i in self._registered_function_policies]If your new list/show CLI for function policies already calls the more specific
get_registered_function_policies(), this is not strictly required for correctness but will keep the generic discovery API complete.Also applies to: 1045-1095
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (27)
docs/source/reference/cli.md(2 hunks)docs/source/reference/function-policies.md(1 hunks)docs/source/reference/middleware.md(1 hunks)src/nat/builder/builder.py(4 hunks)src/nat/builder/component_utils.py(5 hunks)src/nat/builder/workflow_builder.py(9 hunks)src/nat/cli/commands/info/info.py(2 hunks)src/nat/cli/commands/info/list_function_policies.py(1 hunks)src/nat/cli/register_workflow.py(3 hunks)src/nat/cli/type_registry.py(8 hunks)src/nat/data_models/component.py(2 hunks)src/nat/data_models/component_ref.py(1 hunks)src/nat/data_models/config.py(6 hunks)src/nat/data_models/function_policy.py(1 hunks)src/nat/data_models/middleware.py(2 hunks)src/nat/function_policy/__init__.py(1 hunks)src/nat/function_policy/interface.py(1 hunks)src/nat/middleware/__init__.py(1 hunks)src/nat/middleware/dynamic_function_middleware.py(1 hunks)src/nat/middleware/register.py(2 hunks)src/nat/middleware/utils/__init__.py(1 hunks)src/nat/middleware/utils/workflow_inventory.py(1 hunks)tests/nat/builder/test_builder.py(10 hunks)tests/nat/builder/test_component_utils.py(3 hunks)tests/nat/cli/commands/info/test_list_function_policies.py(1 hunks)tests/nat/function_policy/test_policy_interface.py(1 hunks)tests/nat/middleware/test_dynamic_middleware.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (16)
**/*.{md,rst,py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NVIDIA NeMo Agent toolkit' on first use, then 'NeMo Agent toolkit' for subsequent references
Files:
docs/source/reference/middleware.mdtests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pydocs/source/reference/function-policies.mdsrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pydocs/source/reference/cli.mdsrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NeMo Agent Toolkit' (capitalize 'T') when the name appears in headings
Files:
docs/source/reference/middleware.mddocs/source/reference/function-policies.mddocs/source/reference/cli.md
docs/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never use deprecated names: Agent Intelligence toolkit, aiqtoolkit, AgentIQ, AIQ/aiq in documentation (unless intentionally referring to deprecated versions)
Files:
docs/source/reference/middleware.mddocs/source/reference/function-policies.mddocs/source/reference/cli.md
**/*.{py,js,ts,yaml,yml,json,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces, never tabs, and ensure every file ends with a single newline
Files:
docs/source/reference/middleware.mdtests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pydocs/source/reference/function-policies.mdsrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pydocs/source/reference/cli.mdsrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
docs/source/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
docs/source/**/*.{md,rst}: Documentation must be clear and comprehensive, without TODOs, FIXMEs, or placeholder text
Ensure documentation is free of offensive or outdated terms
Ensure documentation is free of spelling mistakes and do not use words listed in ci/vale/styles/config/vocabularies/nat/reject.txt
Files:
docs/source/reference/middleware.mddocs/source/reference/function-policies.mddocs/source/reference/cli.md
**/*.{py,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived automatically by setuptools-scm; never hard-code them in code or docs
Files:
docs/source/reference/middleware.mdtests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pydocs/source/reference/function-policies.mdsrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pydocs/source/reference/cli.mdsrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
**/*
⚙️ CodeRabbit configuration file
**/*: # Code Review Instructions
- Ensure the code follows best practices and coding standards. - For Python code, follow
PEP 20 and
PEP 8 for style guidelines.- Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values (except for return values of
None,
in that situation no return type hint is needed).
Example:def my_function(param1: int, param2: str) -> bool: pass- For Python exception handling, ensure proper stack trace preservation:
- When re-raising exceptions: use bare
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.exception()) to avoid duplicate stack trace output.- When catching and logging exceptions without re-raising: always use
logger.exception()
to capture the full stack trace information.Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any
words listed in the
ci/vale/styles/config/vocabularies/nat/reject.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile are OK.Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,
and should contain an Apache License 2.0 header comment at the top of each file.
- Confirm that copyright years are up-to date whenever a file is changed.
Files:
docs/source/reference/middleware.mdtests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pydocs/source/reference/function-policies.mdsrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pydocs/source/reference/cli.mdsrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
docs/source/**/*
⚙️ CodeRabbit configuration file
This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the
docs/source/_staticdirectory.
Files:
docs/source/reference/middleware.mddocs/source/reference/function-policies.mddocs/source/reference/cli.md
**/*.{py,toml,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' for environment variable prefixes and informal comments
Files:
tests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pysrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Follow PEP 20 and PEP 8 for Python style guidelines
Run yapf second (PEP 8 base, column_limit = 120) for Python formatting
Use ruff check --fix (via pre-commit) for linting using configuration embedded in pyproject.toml, fix warnings unless explicitly ignored
Use snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
All public APIs require Python 3.11+ type hints on parameters and return values
Prefer collections.abc / typing abstractions (Sequence over list) for type hints
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Preserve stack traces and prevent duplicate logging when handling exceptions; use bare raise statements and logger.error() when re-raising
When catching and logging exceptions without re-raising, always use logger.exception() to capture the full stack trace information
Provide Google-style docstrings for every public module, class, function and CLI command
The first line of docstrings must be a concise description ending with a period (Vale checks this)
Surround code entities with backticks in docstrings to avoid Vale false-positives
Validate and sanitise all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work (HTTP, DB, file reads)
Cache expensive computations with functools.lru_cache or an external cache when appropriate
Leverage NumPy vectorised operations whenever beneficial and feasible
Files:
tests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pysrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
**/test_*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/test_*.py: Use pytest with pytest-asyncio for asynchronous code testing
Test functions should be named using the test_ prefix, using snake_case
Extract frequently repeated code into pytest fixtures, which should be named using the fixture_ prefix and define the name argument in the decorator
Mock external services with pytest_httpserver or unittest.mock instead of hitting live endpoints
Mark slow tests with @pytest.mark.slow so they can be skipped in the default test suite
Mark integration tests requiring external services with @pytest.mark.integration so they can be skipped in the default test suite
Files:
tests/nat/cli/commands/info/test_list_function_policies.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.py
**/*.{py,env,toml,yaml,yml,json}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never commit API keys, credentials or personal data; use environment variables or .env files excluded from Git
Files:
tests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pysrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
**/*.{py,js,ts,java,cpp,c,go,rb,php}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Every file must start with the standard SPDX Apache-2.0 header
Files:
tests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pysrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
**/*.{py,js,ts,java,cpp,c,go,rb,php,sh}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All source files must include the SPDX Apache-2.0 header template
Files:
tests/nat/cli/commands/info/test_list_function_policies.pysrc/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pysrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pysrc/nat/cli/commands/info/info.pytests/nat/builder/test_builder.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
tests/**/*.py
⚙️ CodeRabbit configuration file
tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using thetest_prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using thefixture_prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass
Files:
tests/nat/cli/commands/info/test_list_function_policies.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pytests/nat/function_policy/test_policy_interface.pytests/nat/middleware/test_dynamic_middleware.py
src/nat/**/*
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/data_models/function_policy.pysrc/nat/middleware/register.pysrc/nat/middleware/dynamic_function_middleware.pysrc/nat/middleware/utils/__init__.pysrc/nat/function_policy/interface.pysrc/nat/builder/component_utils.pysrc/nat/data_models/component.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/data_models/component_ref.pysrc/nat/data_models/middleware.pysrc/nat/cli/type_registry.pysrc/nat/cli/commands/info/info.pysrc/nat/function_policy/__init__.pysrc/nat/middleware/__init__.pysrc/nat/cli/commands/info/list_function_policies.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
🧠 Learnings (4)
📚 Learning: 2025-08-28T14:28:34.779Z
Learnt from: willkill07
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 684
File: src/nat/builder/component_utils.py:110-114
Timestamp: 2025-08-28T14:28:34.779Z
Learning: FunctionGroupBaseConfig and FunctionBaseConfig are sibling classes that both inherit from (TypedBaseModel, BaseModelRegistryTag). FunctionGroupBaseConfig does not inherit from FunctionBaseConfig, so isinstance checks between them won't interfere with each other and the order doesn't matter for classification purposes.
Applied to files:
src/nat/data_models/function_policy.pysrc/nat/builder/component_utils.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
📚 Learning: 2025-08-28T14:28:34.779Z
Learnt from: willkill07
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 684
File: src/nat/builder/component_utils.py:110-114
Timestamp: 2025-08-28T14:28:34.779Z
Learning: FunctionGroupBaseConfig and FunctionBaseConfig are separate class hierarchies that both inherit from (TypedBaseModel, BaseModelRegistryTag). FunctionGroupBaseConfig does not inherit from FunctionBaseConfig, so isinstance checks between them won't interfere with each other.
Applied to files:
src/nat/data_models/function_policy.pysrc/nat/builder/component_utils.pysrc/nat/cli/type_registry.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/builder.pysrc/nat/builder/workflow_builder.pysrc/nat/cli/register_workflow.pysrc/nat/data_models/config.py
📚 Learning: 2025-11-24T18:56:53.109Z
Learnt from: CR
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 0
File: .cursor/rules/general.mdc:0-0
Timestamp: 2025-11-24T18:56:53.109Z
Learning: Applies to **/*.{py,js,ts,java,cpp,c,go,rb,php} : Every file must start with the standard SPDX Apache-2.0 header
Applied to files:
src/nat/middleware/utils/__init__.py
📚 Learning: 2025-11-24T18:56:53.109Z
Learnt from: CR
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 0
File: .cursor/rules/general.mdc:0-0
Timestamp: 2025-11-24T18:56:53.109Z
Learning: Applies to **/*.py : Use typing.Annotated for units or extra metadata when useful
Applied to files:
src/nat/data_models/config.py
🧬 Code graph analysis (15)
tests/nat/cli/commands/info/test_list_function_policies.py (2)
src/nat/cli/entrypoint.py (1)
cli(84-102)src/nat/builder/context.py (1)
output(60-61)
src/nat/middleware/register.py (3)
src/nat/builder/builder.py (1)
Builder(74-383)src/nat/data_models/middleware.py (1)
DynamicMiddlewareConfig(96-156)src/nat/middleware/dynamic_function_middleware.py (1)
DynamicFunctionMiddleware(45-998)
src/nat/builder/component_utils.py (1)
src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)
src/nat/middleware/utils/workflow_inventory.py (2)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/data_models/function.py (1)
FunctionBaseConfig(26-36)
src/nat/data_models/component_ref.py (2)
src/nat/utils/type_utils.py (1)
override(56-57)src/nat/data_models/component.py (1)
ComponentGroup(48-60)
src/nat/cli/type_registry.py (4)
src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/function_policy/interface.py (1)
FunctionPolicyBase(63-120)src/nat/cli/register_workflow.py (1)
register_function_policy(267-303)src/nat/builder/workflow_builder.py (2)
get_function_policy(1129-1144)get_function_policy(1590-1592)
tests/nat/builder/test_component_utils.py (5)
src/nat/data_models/authentication.py (1)
AuthProviderBaseConfig(31-37)src/nat/data_models/function.py (1)
FunctionGroupBaseConfig(39-68)src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/data_models/middleware.py (1)
MiddlewareBaseConfig(36-42)src/nat/data_models/component.py (1)
ComponentGroup(48-60)
src/nat/cli/commands/info/info.py (1)
src/nat/cli/commands/info/list_function_policies.py (2)
list_function_policies(25-74)show_function_policy(79-160)
src/nat/function_policy/__init__.py (1)
src/nat/function_policy/interface.py (3)
FunctionPolicyBase(63-120)PostInvokeContext(46-60)PreInvokeContext(29-42)
src/nat/builder/builder.py (4)
src/nat/data_models/component_ref.py (1)
FunctionPolicyRef(193-201)src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/function_policy/interface.py (1)
FunctionPolicyBase(63-120)src/nat/builder/workflow_builder.py (6)
add_function_policy(1096-1126)add_function_policy(1584-1587)get_function_policy(1129-1144)get_function_policy(1590-1592)get_function_policy_config(1147-1162)get_function_policy_config(1595-1597)
src/nat/builder/workflow_builder.py (5)
src/nat/data_models/component_ref.py (13)
FunctionPolicyRef(193-201)component_group(69-76)component_group(90-91)component_group(101-102)component_group(112-113)component_group(123-124)component_group(134-135)component_group(145-146)component_group(156-157)component_group(167-168)component_group(178-179)component_group(189-190)component_group(200-201)src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/function_policy/interface.py (1)
FunctionPolicyBase(63-120)src/nat/builder/builder.py (3)
add_function_policy(348-359)get_function_policy(362-371)get_function_policy_config(374-383)src/nat/cli/type_registry.py (1)
get_function_policy(597-613)
tests/nat/function_policy/test_policy_interface.py (3)
src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/function_policy/interface.py (5)
FunctionPolicyBase(63-120)PostInvokeContext(46-60)PreInvokeContext(29-42)on_pre_invoke(89-103)on_post_invoke(106-120)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)
tests/nat/middleware/test_dynamic_middleware.py (4)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/data_models/middleware.py (1)
DynamicMiddlewareConfig(96-156)src/nat/function_policy/interface.py (3)
FunctionPolicyBase(63-120)PostInvokeContext(46-60)PreInvokeContext(29-42)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)
src/nat/cli/register_workflow.py (2)
src/nat/cli/type_registry.py (4)
register_function_policy(580-595)GlobalTypeRegistry(1176-1197)RegisteredFunctionPolicyInfo(209-215)get(1181-1182)src/nat/data_models/discovery_metadata.py (2)
DiscoveryMetadata(53-305)from_config_type(152-189)
src/nat/data_models/config.py (2)
src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/cli/type_registry.py (3)
get(1181-1182)get_registered_function_policies(615-621)compute_annotation(1123-1173)
🪛 markdownlint-cli2 (0.18.1)
docs/source/reference/function-policies.md
164-164: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
180-180: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
463-463: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🪛 Ruff (0.14.7)
src/nat/middleware/register.py
28-28: Unused function argument: builder
(ARG001)
src/nat/middleware/dynamic_function_middleware.py
118-118: Do not catch blind exception: Exception
(BLE001)
155-155: Do not catch blind exception: Exception
(BLE001)
194-194: Do not catch blind exception: Exception
(BLE001)
232-232: Do not catch blind exception: Exception
(BLE001)
272-272: Do not catch blind exception: Exception
(BLE001)
312-312: Do not catch blind exception: Exception
(BLE001)
535-535: Avoid specifying long messages outside the exception class
(TRY003)
539-541: Avoid specifying long messages outside the exception class
(TRY003)
583-584: Avoid specifying long messages outside the exception class
(TRY003)
660-660: Consider moving this statement to an else block
(TRY300)
661-661: Do not catch blind exception: Exception
(BLE001)
801-801: Unused method argument: component_type
(ARG002)
829-831: try-except-continue detected, consider logging the exception
(S112)
829-829: Do not catch blind exception: Exception
(BLE001)
884-884: Consider moving this statement to an else block
(TRY300)
886-886: Do not catch blind exception: Exception
(BLE001)
902-902: Avoid specifying long messages outside the exception class
(TRY003)
910-910: Avoid specifying long messages outside the exception class
(TRY003)
918-918: Avoid specifying long messages outside the exception class
(TRY003)
926-926: Avoid specifying long messages outside the exception class
(TRY003)
934-935: Avoid specifying long messages outside the exception class
(TRY003)
943-943: Avoid specifying long messages outside the exception class
(TRY003)
951-951: Avoid specifying long messages outside the exception class
(TRY003)
961-961: Avoid specifying long messages outside the exception class
(TRY003)
967-967: Avoid specifying long messages outside the exception class
(TRY003)
973-973: Avoid specifying long messages outside the exception class
(TRY003)
979-979: Avoid specifying long messages outside the exception class
(TRY003)
985-985: Avoid specifying long messages outside the exception class
(TRY003)
991-991: Avoid specifying long messages outside the exception class
(TRY003)
997-997: Avoid specifying long messages outside the exception class
(TRY003)
src/nat/function_policy/interface.py
123-128: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
src/nat/cli/type_registry.py
590-591: Avoid specifying long messages outside the exception class
(TRY003)
612-613: Avoid specifying long messages outside the exception class
(TRY003)
tests/nat/builder/test_builder.py
299-299: Unused function argument: b
(ARG001)
src/nat/function_policy/__init__.py
21-26: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
src/nat/builder/workflow_builder.py
1111-1111: Avoid specifying long messages outside the exception class
(TRY003)
1123-1123: Consider moving this statement to an else block
(TRY300)
1125-1125: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
1142-1142: Avoid specifying long messages outside the exception class
(TRY003)
1160-1160: Avoid specifying long messages outside the exception class
(TRY003)
tests/nat/middleware/test_dynamic_middleware.py
278-278: Unused function argument: value
(ARG001)
378-378: Unused function argument: value
(ARG001)
442-442: Unused function argument: value
(ARG001)
471-471: Unused function argument: value
(ARG001)
550-550: Unused function argument: value
(ARG001)
680-680: Unused function argument: value
(ARG001)
783-783: Unused function argument: value
(ARG001)
813-813: Unused lambda argument: inp
(ARG005)
816-816: Unused lambda argument: out
(ARG005)
1698-1698: Local variable llm_allowlist is assigned to but never used
Remove assignment to unused variable llm_allowlist
(F841)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI Pipeline / Check
|
What is the difference between function policy and function middleware? Seems to me like the do the same thing:
Is there a specific functionality policies provide that middleware cannot? If so, would it be worth extending a function middleware instead so we don't have duplicated functionality? |
…port Signed-off-by: Eric Evans <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (17)
src/nat/function_policy/register.py (1)
15-18: Missing module-level docstring.Per coding guidelines, all public modules require Google-style docstrings. This file currently lacks a module docstring explaining its purpose as a re-export module.
# flake8: noqa +"""Public registration exports for function policies.""" from nat.function_policy.logging import register as register_loggingsrc/nat/function_policy/logging/register.py (1)
23-34: Consider adding return type hint for the async generator.The function signature lacks a return type annotation. For async generators yielding policy instances, adding the type hint improves IDE support and type checking.
+from collections.abc import AsyncGenerator + + @register_function_policy(config_type=LoggingPolicyConfig) -async def logging_policy(config: LoggingPolicyConfig, builder: Builder): +async def logging_policy(config: LoggingPolicyConfig, builder: Builder) -> AsyncGenerator[LoggingPolicy, None]: """Build a logging policy from configuration.Note: The unused
builderparameter is intentional per the component pattern interface and is properly documented.src/nat/function_policy/logging/logging_policy_config.py (1)
22-26: Consider usingLiteralfor stricter log level validation.The
log_levelfield is typed asstrbut only specific values are valid per the description. UsingLiteralwould catch invalid configurations early during parsing rather than at runtime inLoggingPolicy.on_pre_invoke.-from pydantic import Field +from typing import Literal + +from pydantic import Field from nat.data_models.function_policy import FunctionPolicyBaseConfig class LoggingPolicyConfig(FunctionPolicyBaseConfig, name="logging"): """Configuration for logging policy.""" - log_level: str = Field(default="INFO", description="Logging level (DEBUG, INFO, WARNING, ERROR)") + log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR"] = Field( + default="INFO", description="Logging level") max_value_length: int = Field(default=200, description="Maximum length for logged values (truncated if exceeded)")src/nat/middleware/cache/register.py (1)
23-34: Consider adding return type hint for the async generator.Similar to the logging policy registration, adding a return type annotation improves type safety and IDE support.
+from collections.abc import AsyncGenerator + + @register_middleware(config_type=CacheMiddlewareConfig) -async def cache_middleware(config: CacheMiddlewareConfig, builder: Builder): +async def cache_middleware(config: CacheMiddlewareConfig, builder: Builder) -> AsyncGenerator[CacheMiddleware, None]: """Build a cache middleware from configuration.Note: The unused
builderparameter follows the established component pattern and is properly documented.docs/source/reference/function-policies.md (1)
187-195: Optional: Add language identifiers to code blocks.The markdown linter flags three code blocks without language identifiers. These blocks contain ASCII diagrams/flow representations. Adding
textas the language identifier would silence the linting warnings while maintaining readability.For example:
-``` +```text Original Args ((10,)) ↓Also applies to: 203-211, 495-512
src/nat/function_policy/interface.py (1)
126-131: Optional: Sort__all__exports alphabetically.The
__all__list is not alphabetically sorted. While this doesn't affect functionality, sorting it would align with Python conventions and silence the linting warning.Apply this diff:
__all__ = [ "FunctionPolicyBase", + "PolicyConfigT", "PreInvokeContext", "PostInvokeContext", - "PolicyConfigT", ]src/nat/middleware/dynamic/dynamic_function_middleware.py (6)
45-49: Module docstring missing for new file.Per coding guidelines, public modules require Google-style docstrings. Add a module-level docstring describing the purpose of this middleware.
+"""Dynamic function middleware for runtime component discovery and policy orchestration. + +This module provides `DynamicFunctionMiddleware`, which intercepts workflow builder +methods to discover and wrap component functions with policy-based pre/post invoke hooks. +""" + class DynamicFunctionMiddleware(FunctionMiddleware):
115-120: Silent failure during component function registration.Registration failures are logged at DEBUG level and silently skipped. While resilient, this could mask configuration errors. Consider logging at WARNING level when in development/debug mode.
The same pattern repeats across all
_discover_and_register_*methods (lines 152-159, 191-198, 229-236, 269-276, 309-316). The blind exception catches are intentional for resilience but could be more informative.except Exception as e: - logger.debug("Failed to register component function '%s' on LLM '%s': %s", function_name, llm_name, e) + logger.warning("Failed to register component function '%s' on LLM '%s': %s", function_name, llm_name, e)
240-243: Docstring formatting: extra blank line.Minor formatting issue - there's an extra blank line between the docstring summary and Args section.
async def _discover_and_register_object_store(self, object_store_name: str) -> Any: """Intercept object store creation and register allowlisted component functions with middleware. - Args:
280-283: Docstring formatting: extra blank line.Same formatting issue as above.
async def _discover_and_register_auth_provider(self, auth_provider_name: str) -> Any: """Intercept auth provider creation and register allowlisted component functions with middleware. - Args:
817-817: Unused parametercomponent_typeshould be removed or prefixed.Static analysis correctly flags this as unused. Either use it for logging/filtering or prefix with underscore.
- def _get_callable_functions(self, instance: Any, component_type: str | None = None) -> set[str]: + def _get_callable_functions(self, instance: Any, _component_type: str | None = None) -> set[str]:Or remove the parameter entirely if not needed for future use.
845-847: Silent exception during introspection could hide issues.The
try-except-continuepattern silently swallows exceptions during function introspection. Consider logging at DEBUG level for diagnosability.except Exception: # Skip functions that raise errors during introspection + logger.debug("Skipping function '%s' due to introspection error", function_name, exc_info=True) continuesrc/nat/middleware/utils/workflow_inventory.py (1)
123-124: Missing__all__export and module docstring.Per coding guidelines, public modules should have docstrings and explicit exports.
Add at the top of the file after imports:
"""Workflow inventory models for dynamic component discovery. This module provides data models for tracking discovered workflow components and their callable functions, along with default allowlists for secure function wrapping. """ __all__ = [ "COMPONENT_FUNCTION_ALLOWLISTS", "DiscoveredComponent", "DiscoveredFunction", "WorkflowInventory", ]tests/nat/builder/test_builder.py (1)
298-313: Unused parameterbin registration function.Static analysis correctly flags the unused
bparameter. Per coding guidelines, prefix with underscore.@register_function_policy(config_type=TFunctionPolicyConfig) - async def register_function_policy_fn(config: TFunctionPolicyConfig, b: Builder): + async def register_function_policy_fn(config: TFunctionPolicyConfig, _builder: Builder):src/nat/middleware/dynamic/dynamic_middleware_config.py (2)
15-15: Module docstring is brief but present.The docstring at line 15 is minimal. Consider expanding for clarity.
-"""Configuration for dynamic middleware.""" +"""Configuration models for dynamic middleware. + +This module provides configuration classes for the dynamic middleware system, +including component function allowlists and the main middleware configuration. +"""
136-138: Missing__all__export for public API.Per coding guidelines, consider adding explicit exports.
Add after imports:
__all__ = ["AllowedComponentFunctions", "DynamicMiddlewareConfig"]src/nat/middleware/middleware.py (1)
125-125: **Consider documenting expected kwargs more specifically.The docstrings describe
**kwargsas "Additional function arguments" which is accurate but generic. Middleware implementers might benefit from knowing what contextual arguments they can expect (e.g., policy contexts, session metadata). However, keeping the documentation generic maintains flexibility as the framework evolves.Also applies to: 161-161
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (26)
docs/source/reference/function-policies.md(1 hunks)src/nat/data_models/middleware.py(1 hunks)src/nat/function_policy/interface.py(1 hunks)src/nat/function_policy/logging/__init__.py(1 hunks)src/nat/function_policy/logging/logging_policy.py(1 hunks)src/nat/function_policy/logging/logging_policy_config.py(1 hunks)src/nat/function_policy/logging/register.py(1 hunks)src/nat/function_policy/register.py(1 hunks)src/nat/middleware/__init__.py(1 hunks)src/nat/middleware/cache/__init__.py(1 hunks)src/nat/middleware/cache/cache_middleware.py(6 hunks)src/nat/middleware/cache/cache_middleware_config.py(1 hunks)src/nat/middleware/cache/register.py(1 hunks)src/nat/middleware/dynamic/__init__.py(1 hunks)src/nat/middleware/dynamic/dynamic_function_middleware.py(1 hunks)src/nat/middleware/dynamic/dynamic_middleware_config.py(1 hunks)src/nat/middleware/dynamic/register.py(1 hunks)src/nat/middleware/function_middleware.py(3 hunks)src/nat/middleware/middleware.py(5 hunks)src/nat/middleware/register.py(1 hunks)src/nat/middleware/utils/workflow_inventory.py(1 hunks)tests/nat/builder/test_builder.py(10 hunks)tests/nat/function_policy/test_policy_interface.py(1 hunks)tests/nat/middleware/test_cache_middleware.py(11 hunks)tests/nat/middleware/test_dynamic_middleware.py(1 hunks)tests/nat/middleware/test_middleware_components.py(3 hunks)
✅ Files skipped from review due to trivial changes (3)
- src/nat/middleware/dynamic/init.py
- src/nat/middleware/cache/init.py
- src/nat/function_policy/logging/init.py
🚧 Files skipped from review as they are similar to previous changes (3)
- src/nat/middleware/register.py
- tests/nat/function_policy/test_policy_interface.py
- src/nat/data_models/middleware.py
🧰 Additional context used
📓 Path-based instructions (16)
**/*.{md,rst,py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NVIDIA NeMo Agent toolkit' on first use, then 'NeMo Agent toolkit' for subsequent references
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pydocs/source/reference/function-policies.mdtests/nat/builder/test_builder.py
**/*.{py,toml,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' for environment variable prefixes and informal comments
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pytests/nat/builder/test_builder.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Follow PEP 20 and PEP 8 for Python style guidelines
Run yapf second (PEP 8 base, column_limit = 120) for Python formatting
Use ruff check --fix (via pre-commit) for linting using configuration embedded in pyproject.toml, fix warnings unless explicitly ignored
Use snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
All public APIs require Python 3.11+ type hints on parameters and return values
Prefer collections.abc / typing abstractions (Sequence over list) for type hints
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Preserve stack traces and prevent duplicate logging when handling exceptions; use bare raise statements and logger.error() when re-raising
When catching and logging exceptions without re-raising, always use logger.exception() to capture the full stack trace information
Provide Google-style docstrings for every public module, class, function and CLI command
The first line of docstrings must be a concise description ending with a period (Vale checks this)
Surround code entities with backticks in docstrings to avoid Vale false-positives
Validate and sanitise all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work (HTTP, DB, file reads)
Cache expensive computations with functools.lru_cache or an external cache when appropriate
Leverage NumPy vectorised operations whenever beneficial and feasible
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pytests/nat/builder/test_builder.py
**/*.{py,js,ts,yaml,yml,json,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces, never tabs, and ensure every file ends with a single newline
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pydocs/source/reference/function-policies.mdtests/nat/builder/test_builder.py
**/*.{py,env,toml,yaml,yml,json}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never commit API keys, credentials or personal data; use environment variables or .env files excluded from Git
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pytests/nat/builder/test_builder.py
**/*.{py,js,ts,java,cpp,c,go,rb,php}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Every file must start with the standard SPDX Apache-2.0 header
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pytests/nat/builder/test_builder.py
**/*.{py,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived automatically by setuptools-scm; never hard-code them in code or docs
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pydocs/source/reference/function-policies.mdtests/nat/builder/test_builder.py
**/*.{py,js,ts,java,cpp,c,go,rb,php,sh}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All source files must include the SPDX Apache-2.0 header template
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pytests/nat/builder/test_builder.py
**/*
⚙️ CodeRabbit configuration file
**/*: # Code Review Instructions
- Ensure the code follows best practices and coding standards. - For Python code, follow
PEP 20 and
PEP 8 for style guidelines.- Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values (except for return values of
None,
in that situation no return type hint is needed).
Example:def my_function(param1: int, param2: str) -> bool: pass- For Python exception handling, ensure proper stack trace preservation:
- When re-raising exceptions: use bare
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.exception()) to avoid duplicate stack trace output.- When catching and logging exceptions without re-raising: always use
logger.exception()
to capture the full stack trace information.Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any
words listed in the
ci/vale/styles/config/vocabularies/nat/reject.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile are OK.Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,
and should contain an Apache License 2.0 header comment at the top of each file.
- Confirm that copyright years are up-to date whenever a file is changed.
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pytests/nat/middleware/test_middleware_components.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pytests/nat/middleware/test_cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.pydocs/source/reference/function-policies.mdtests/nat/builder/test_builder.py
src/nat/**/*
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/function_policy/logging/logging_policy_config.pysrc/nat/function_policy/logging/logging_policy.pysrc/nat/middleware/cache/register.pysrc/nat/function_policy/register.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pysrc/nat/middleware/cache/cache_middleware.pysrc/nat/middleware/utils/workflow_inventory.pysrc/nat/middleware/__init__.pysrc/nat/middleware/dynamic/dynamic_middleware_config.pysrc/nat/middleware/dynamic/register.pysrc/nat/function_policy/logging/register.pysrc/nat/middleware/middleware.pysrc/nat/middleware/cache/cache_middleware_config.pysrc/nat/middleware/function_middleware.py
**/test_*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/test_*.py: Use pytest with pytest-asyncio for asynchronous code testing
Test functions should be named using the test_ prefix, using snake_case
Extract frequently repeated code into pytest fixtures, which should be named using the fixture_ prefix and define the name argument in the decorator
Mock external services with pytest_httpserver or unittest.mock instead of hitting live endpoints
Mark slow tests with @pytest.mark.slow so they can be skipped in the default test suite
Mark integration tests requiring external services with @pytest.mark.integration so they can be skipped in the default test suite
Files:
tests/nat/middleware/test_middleware_components.pytests/nat/middleware/test_dynamic_middleware.pytests/nat/middleware/test_cache_middleware.pytests/nat/builder/test_builder.py
tests/**/*.py
⚙️ CodeRabbit configuration file
tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using thetest_prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using thefixture_prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass
Files:
tests/nat/middleware/test_middleware_components.pytests/nat/middleware/test_dynamic_middleware.pytests/nat/middleware/test_cache_middleware.pytests/nat/builder/test_builder.py
**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NeMo Agent Toolkit' (capitalize 'T') when the name appears in headings
Files:
docs/source/reference/function-policies.md
docs/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never use deprecated names: Agent Intelligence toolkit, aiqtoolkit, AgentIQ, AIQ/aiq in documentation (unless intentionally referring to deprecated versions)
Files:
docs/source/reference/function-policies.md
docs/source/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
docs/source/**/*.{md,rst}: Documentation must be clear and comprehensive, without TODOs, FIXMEs, or placeholder text
Ensure documentation is free of offensive or outdated terms
Ensure documentation is free of spelling mistakes and do not use words listed in ci/vale/styles/config/vocabularies/nat/reject.txt
Files:
docs/source/reference/function-policies.md
docs/source/**/*
⚙️ CodeRabbit configuration file
This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the
docs/source/_staticdirectory.
Files:
docs/source/reference/function-policies.md
🧠 Learnings (2)
📚 Learning: 2025-08-28T14:28:34.779Z
Learnt from: willkill07
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 684
File: src/nat/builder/component_utils.py:110-114
Timestamp: 2025-08-28T14:28:34.779Z
Learning: FunctionGroupBaseConfig and FunctionBaseConfig are sibling classes that both inherit from (TypedBaseModel, BaseModelRegistryTag). FunctionGroupBaseConfig does not inherit from FunctionBaseConfig, so isinstance checks between them won't interfere with each other and the order doesn't matter for classification purposes.
Applied to files:
tests/nat/builder/test_builder.py
📚 Learning: 2025-08-28T14:28:34.779Z
Learnt from: willkill07
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 684
File: src/nat/builder/component_utils.py:110-114
Timestamp: 2025-08-28T14:28:34.779Z
Learning: FunctionGroupBaseConfig and FunctionBaseConfig are separate class hierarchies that both inherit from (TypedBaseModel, BaseModelRegistryTag). FunctionGroupBaseConfig does not inherit from FunctionBaseConfig, so isinstance checks between them won't interfere with each other.
Applied to files:
tests/nat/builder/test_builder.py
🧬 Code graph analysis (14)
src/nat/function_policy/logging/logging_policy_config.py (1)
src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)
src/nat/middleware/cache/register.py (2)
src/nat/middleware/cache/cache_middleware.py (1)
CacheMiddleware(46-235)src/nat/middleware/cache/cache_middleware_config.py (1)
CacheMiddlewareConfig(24-44)
tests/nat/middleware/test_middleware_components.py (3)
src/nat/middleware/cache/cache_middleware_config.py (1)
CacheMiddlewareConfig(24-44)src/nat/middleware/cache/register.py (1)
cache_middleware(24-34)src/nat/middleware/cache/cache_middleware.py (1)
CacheMiddleware(46-235)
src/nat/middleware/dynamic/dynamic_function_middleware.py (5)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/middleware/dynamic/dynamic_middleware_config.py (1)
DynamicMiddlewareConfig(78-138)src/nat/middleware/function_middleware.py (3)
FunctionMiddleware(43-110)function_middleware_invoke(69-88)function_middleware_stream(90-110)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)src/nat/middleware/utils/workflow_inventory.py (3)
DiscoveredComponent(55-74)DiscoveredFunction(77-90)WorkflowInventory(93-124)
tests/nat/middleware/test_dynamic_middleware.py (5)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/middleware/dynamic/dynamic_function_middleware.py (10)
function_middleware_invoke(666-739)function_middleware_stream(741-813)_get_callable_functions(817-849)_discover_and_register_llm(87-121)_discover_and_register_embedder(123-161)_discover_and_register_memory(202-238)_discover_and_register_auth_provider(280-318)_discover_functions(358-379)_register_function(503-521)_register_component_function(523-565)src/nat/middleware/dynamic/dynamic_middleware_config.py (1)
DynamicMiddlewareConfig(78-138)src/nat/middleware/utils/workflow_inventory.py (2)
DiscoveredComponent(55-74)DiscoveredFunction(77-90)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)
src/nat/middleware/cache/cache_middleware.py (4)
src/nat/middleware/dynamic/dynamic_function_middleware.py (1)
function_middleware_invoke(666-739)src/nat/middleware/function_middleware.py (1)
function_middleware_invoke(69-88)tests/nat/middleware/test_middleware_components.py (2)
function_middleware_invoke(47-51)function_middleware_invoke(650-659)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)
tests/nat/middleware/test_cache_middleware.py (2)
src/nat/middleware/cache/cache_middleware.py (2)
CacheMiddleware(46-235)function_middleware_invoke(144-200)tests/nat/middleware/test_middleware_components.py (2)
function_middleware_invoke(47-51)function_middleware_invoke(650-659)
src/nat/middleware/utils/workflow_inventory.py (2)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/data_models/function.py (1)
FunctionBaseConfig(26-36)
src/nat/middleware/dynamic/dynamic_middleware_config.py (2)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/data_models/component_ref.py (6)
AuthenticationRef(160-168)EmbedderRef(83-91)LLMRef(116-124)MemoryRef(127-135)ObjectStoreRef(138-146)RetrieverRef(149-157)
src/nat/middleware/dynamic/register.py (3)
src/nat/builder/builder.py (1)
Builder(74-383)src/nat/middleware/dynamic/dynamic_function_middleware.py (1)
DynamicFunctionMiddleware(45-1014)src/nat/middleware/dynamic/dynamic_middleware_config.py (1)
DynamicMiddlewareConfig(78-138)
src/nat/function_policy/logging/register.py (2)
src/nat/function_policy/logging/logging_policy.py (1)
LoggingPolicy(28-74)src/nat/function_policy/logging/logging_policy_config.py (1)
LoggingPolicyConfig(22-26)
src/nat/middleware/middleware.py (3)
src/nat/middleware/function_middleware.py (2)
middleware_invoke(52-58)middleware_stream(60-67)src/nat/runtime/session.py (1)
context(96-97)src/nat/runtime/runner.py (1)
context(97-98)
src/nat/middleware/function_middleware.py (2)
src/nat/runtime/session.py (1)
context(96-97)src/nat/runtime/runner.py (1)
context(97-98)
tests/nat/builder/test_builder.py (5)
src/nat/cli/register_workflow.py (1)
register_function_policy(267-303)src/nat/cli/type_registry.py (2)
register_function_policy(580-595)get_function_policy(597-613)src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)src/nat/function_policy/interface.py (5)
FunctionPolicyBase(67-123)PostInvokeContext(48-64)PreInvokeContext(29-44)on_pre_invoke(93-106)on_post_invoke(109-123)src/nat/builder/workflow_builder.py (4)
add_function_policy(1096-1126)add_function_policy(1584-1587)get_function_policy(1129-1144)get_function_policy(1590-1592)
🪛 markdownlint-cli2 (0.18.1)
docs/source/reference/function-policies.md
187-187: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
203-203: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
495-495: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🪛 Ruff (0.14.7)
src/nat/middleware/cache/register.py
24-24: Unused function argument: builder
(ARG001)
src/nat/middleware/dynamic/dynamic_function_middleware.py
118-118: Do not catch blind exception: Exception
(BLE001)
155-155: Do not catch blind exception: Exception
(BLE001)
194-194: Do not catch blind exception: Exception
(BLE001)
232-232: Do not catch blind exception: Exception
(BLE001)
272-272: Do not catch blind exception: Exception
(BLE001)
312-312: Do not catch blind exception: Exception
(BLE001)
535-535: Avoid specifying long messages outside the exception class
(TRY003)
539-541: Avoid specifying long messages outside the exception class
(TRY003)
583-584: Avoid specifying long messages outside the exception class
(TRY003)
660-660: Consider moving this statement to an else block
(TRY300)
661-661: Do not catch blind exception: Exception
(BLE001)
700-701: Abstract raise to an inner function
(TRY301)
700-701: Avoid specifying long messages outside the exception class
(TRY003)
775-776: Abstract raise to an inner function
(TRY301)
775-776: Avoid specifying long messages outside the exception class
(TRY003)
817-817: Unused method argument: component_type
(ARG002)
845-847: try-except-continue detected, consider logging the exception
(S112)
845-845: Do not catch blind exception: Exception
(BLE001)
900-900: Consider moving this statement to an else block
(TRY300)
902-902: Do not catch blind exception: Exception
(BLE001)
918-918: Avoid specifying long messages outside the exception class
(TRY003)
926-926: Avoid specifying long messages outside the exception class
(TRY003)
934-934: Avoid specifying long messages outside the exception class
(TRY003)
942-942: Avoid specifying long messages outside the exception class
(TRY003)
950-951: Avoid specifying long messages outside the exception class
(TRY003)
959-959: Avoid specifying long messages outside the exception class
(TRY003)
967-967: Avoid specifying long messages outside the exception class
(TRY003)
977-977: Avoid specifying long messages outside the exception class
(TRY003)
983-983: Avoid specifying long messages outside the exception class
(TRY003)
989-989: Avoid specifying long messages outside the exception class
(TRY003)
995-995: Avoid specifying long messages outside the exception class
(TRY003)
1001-1001: Avoid specifying long messages outside the exception class
(TRY003)
1007-1007: Avoid specifying long messages outside the exception class
(TRY003)
1013-1013: Avoid specifying long messages outside the exception class
(TRY003)
tests/nat/middleware/test_dynamic_middleware.py
76-76: Unused method argument: messages
(ARG002)
76-76: Unused method argument: kwargs
(ARG002)
79-79: Unused method argument: messages
(ARG002)
79-79: Unused method argument: kwargs
(ARG002)
82-82: Unused method argument: messages
(ARG002)
82-82: Unused method argument: kwargs
(ARG002)
85-85: Unused method argument: messages
(ARG002)
85-85: Unused method argument: kwargs
(ARG002)
97-97: Unused method argument: text
(ARG002)
100-100: Unused method argument: text
(ARG002)
112-112: Unused method argument: query
(ARG002)
112-112: Unused method argument: kwargs
(ARG002)
124-124: Unused method argument: query
(ARG002)
124-124: Unused method argument: top_k
(ARG002)
124-124: Unused method argument: kwargs
(ARG002)
169-169: Unused method argument: user_id
(ARG002)
169-169: Unused method argument: kwargs
(ARG002)
224-224: Consider (modified_first, *context.function_args[1:]) instead of concatenation
Replace with (modified_first, *context.function_args[1:])
(RUF005)
251-251: Unused function argument: kwargs
(ARG001)
278-278: Unused function argument: kwargs
(ARG001)
305-305: Unused function argument: args
(ARG001)
305-305: Unused function argument: kwargs
(ARG001)
326-326: Unused function argument: args
(ARG001)
326-326: Unused function argument: kwargs
(ARG001)
341-341: Unused function argument: args
(ARG001)
341-341: Unused function argument: kwargs
(ARG001)
370-370: Unused function argument: kwargs
(ARG001)
390-390: Unused function argument: args
(ARG001)
390-390: Unused function argument: kwargs
(ARG001)
409-409: Unused function argument: args
(ARG001)
409-409: Unused function argument: kwargs
(ARG001)
434-434: Unused function argument: args
(ARG001)
434-434: Unused function argument: kwargs
(ARG001)
src/nat/function_policy/interface.py
126-131: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
tests/nat/middleware/test_cache_middleware.py
81-81: Unused function argument: kwargs
(ARG001)
116-116: Unused function argument: args
(ARG001)
116-116: Unused function argument: kwargs
(ARG001)
151-151: Unused function argument: args
(ARG001)
151-151: Unused function argument: kwargs
(ARG001)
190-190: Unused function argument: args
(ARG001)
190-190: Unused function argument: kwargs
(ARG001)
221-221: Unused function argument: args
(ARG001)
221-221: Unused function argument: kwargs
(ARG001)
256-256: Unused function argument: args
(ARG001)
256-256: Unused function argument: kwargs
(ARG001)
302-302: Unused function argument: args
(ARG001)
302-302: Unused function argument: kwargs
(ARG001)
src/nat/function_policy/logging/register.py
24-24: Unused function argument: builder
(ARG001)
src/nat/middleware/function_middleware.py
72-72: Unused method argument: context
(ARG002)
93-93: Unused method argument: context
(ARG002)
tests/nat/builder/test_builder.py
299-299: Unused function argument: b
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI Pipeline / Check
🔇 Additional comments (38)
src/nat/middleware/__init__.py (1)
25-33: LGTM!The
__all__exports are properly alphabetized, and the removal ofCacheMiddlewarefrom this module aligns with its migration to the dedicatednat.middleware.cachesubpackage.src/nat/middleware/cache/cache_middleware_config.py (1)
24-44: Well-structured configuration class.The configuration follows Pydantic best practices with proper type hints (
Literalfor the mode enum,floatwithge/leconstraints for the threshold). The docstring properly documents the behavior of each mode.tests/nat/middleware/test_middleware_components.py (1)
271-302: LGTM! Import paths correctly updated.The import paths have been properly updated to reflect the new module structure for
CacheMiddlewareConfigandCacheMiddleware, aligning with the cache middleware reorganization.src/nat/middleware/dynamic/register.py (1)
23-34: LGTM! Registration function follows the standard pattern.The
dynamic_middlewareregistration function correctly follows the component registration pattern with proper decorator usage, async generator pattern, and comprehensive docstring.src/nat/middleware/function_middleware.py (2)
52-110: LGTM! Signature updates enable flexible argument passing.The updated signatures to accept
*argsand**kwargscorrectly enable the middleware to support multi-argument function targets and policy-driven invocations. The changes maintain backward compatibility while providing the flexibility needed for the Function Policies framework.Note: The static analysis warnings about unused
contextparameters are false positives—contextis required by the middleware interface and is passed through to the wrapper functions.
139-171: LGTM! Chain building correctly forwards expanded arguments.The
build_singleandbuild_streammethods have been properly updated to forward*argsand**kwargsthrough the middleware chain, ensuring compatibility with the new flexible signatures.src/nat/function_policy/logging/logging_policy.py (1)
28-74: LGTM! Clean logging policy implementation.The
LoggingPolicyclass is well-implemented with proper inheritance, clear separation of concerns, and appropriate use of the policy interface. The policy correctly returnsNoneto indicate no modifications, and the_truncatehelper ensures log output is kept manageable.docs/source/reference/function-policies.md (1)
18-609: Excellent comprehensive documentation for Function Policies.This documentation provides thorough coverage of the Function Policies framework, including:
- Clear conceptual overview with pipeline visualization
- Step-by-step implementation guide
- Comprehensive examples for common use cases
- Testing guidance and best practices
- CLI command reference
- Troubleshooting section
The structure and content are well-organized and will help users understand and effectively use the new Function Policies feature.
tests/nat/middleware/test_cache_middleware.py (2)
26-26: LGTM! Import path correctly updated.The import path has been properly updated to reflect the new module structure for
CacheMiddleware.
81-308: LGTM! Test signatures correctly adapted to new middleware interface.The test functions have been properly updated to use keyword arguments (
call_nextandcontext) instead of positional arguments, aligning with the new middleware signatures. The mock functions correctly accept*argsand**kwargsto support the flexible argument passing pattern.Note: Static analysis warnings about unused
argsandkwargsin mock functions are expected and can be safely ignored—these parameters are required for signature compatibility.tests/nat/middleware/test_dynamic_middleware.py (4)
41-237: LGTM! Well-designed test fixtures.The fixtures provide comprehensive mocking for builders and various component types (LLM, embedder, retriever, memory, object store, auth provider). The
tracking_policyfixture is particularly well-designed, offering flexible configuration for testing different policy behaviors.
243-444: LGTM! Comprehensive policy execution tests.These tests thoroughly validate the middleware's policy orchestration capabilities:
- Invoke/stream delegation with no policies
- Pre/post-invoke policy execution and modification
- Policy execution order (pre → function → post)
- Error handling (skipping failing policies)
- Disabled policy behavior
The test coverage ensures the policy framework works correctly across various scenarios.
449-559: LGTM! Thorough component discovery tests.These tests validate discovery and registration for all supported component types (LLMs, embedders, retrievers, memory, object stores, auth providers). The tests also correctly verify that components are skipped when not configured and that duplicates are handled properly.
564-630: LGTM! Registration mechanics properly validated.These tests ensure workflow function discovery and registration work correctly, including:
- Discovery from builder
- Conditional discovery based on configuration
- Duplicate prevention for both functions and component functions
The tests provide good coverage of the registration lifecycle.
src/nat/function_policy/interface.py (2)
28-64: LGTM! Well-designed context objects.The context objects are thoughtfully designed:
PreInvokeContextis mutable to allow middleware to updatefunction_argsas policies executePostInvokeContextis frozen to provide an immutable audit trail of the complete execution- Both include comprehensive attributes for function metadata, original args, current args, and kwargs
This design supports the policy pipeline model effectively.
67-123: LGTM! Clear abstract base class contract.The
FunctionPolicyBaseabstract base class provides a clear contract for policy implementations:
- Generic typing with
PolicyConfigTfor type-safe configuration- Abstract methods for
on_pre_invokeandon_post_invokewith comprehensive docstrings- Clear guidance on return value semantics (None preserves current value, tuple/value transforms)
This interface will guide policy authors effectively.
src/nat/middleware/dynamic/dynamic_function_middleware.py (2)
696-710: Pre-invoke policy error handling logs but continues - verify this is intentional.When a pre-invoke policy fails, the error is logged and the policy is skipped. This provides resilience but may hide critical policy failures. The same pattern applies to post-invoke (lines 731-737) and streaming (lines 778-784, 806-812).
The error-resilient approach is documented in the PR objectives ("failed policies logged and skipped"). The
exc_info=Trueensures full stack traces are captured.
905-914: LGTM - Builder method patching is well-structured.The patching logic properly stores original methods before replacing them, enabling proper delegation. Error handling for missing methods is appropriate.
src/nat/middleware/cache/cache_middleware.py (3)
144-148: Signature updated to match new FunctionMiddleware interface.The change from a single
valueparameter to*args, **kwargsaligns with the broader middleware interface changes shown in the relevant code snippets.
170-172: Cache key extraction assumes at least one positional argument.When
argsis empty,args[0]will raiseIndexError. The current code handles this by settingvalue = None, but_serialize_input(None)may produce unexpected cache behavior.Verify that functions wrapped by this cache middleware always receive at least one positional argument, or consider adding explicit handling:
# Use first arg as cache key (primary input) - value = args[0] if args else None + if not args: + logger.debug("No positional arguments for function %s, bypassing cache", context.name) + return await call_next(*args, **kwargs) + value = args[0]
202-206: Streaming method signature correctly updated.The streaming method signature mirrors the invoke method changes, maintaining consistency.
src/nat/middleware/utils/workflow_inventory.py (4)
1-14: License header and imports are correct.The SPDX Apache-2.0 header is present and properly formatted.
28-52: Component function allowlists are well-defined.The allowlists follow a secure-by-default approach, only exposing commonly-used component methods. This aligns with the PR objectives mentioning "allowlists for secure-by-default wrapping."
55-74: DiscoveredComponent model is well-structured.The model correctly uses
arbitrary_types_allowed=Trueto store component instances and includes proper field descriptions.
93-124: WorkflowInventory provides a clean container for discovered components.The inventory structure mirrors the component groups and provides a clear API for tracking discovered items.
tests/nat/builder/test_builder.py (6)
118-119: Test configuration class follows existing patterns.The
TFunctionPolicyConfigmirrors other test config classes in the file (e.g.,TLLMProviderConfig).
303-313: Test policy implementation is correctly minimal.The
TestPolicyclass properly implements the abstract methods with passthrough behavior, which is appropriate for testing the registration and lifecycle.
920-932: Function policy tests follow established patterns.
test_add_function_policymirrors the structure oftest_add_llm,test_add_embedder, etc., testing:
- Successful addition
- Error during creation (raise_error=True)
- Duplicate name rejection
934-948: test_get_function_policy validates retrieval and type.Good coverage of retrieval, config consistency, and
isinstancecheck.
951-962: test_get_function_policy_config follows existing patterns.Consistent with other
test_get_*_configtests.
976-1011: test_built_config correctly extended for function policies.The test properly adds function policies and verifies they appear in the workflow config.
src/nat/middleware/dynamic/dynamic_middleware_config.py (4)
34-76: AllowedComponentFunctions merge logic is clean and correct.The
model_validatorproperly merges user-provided values with defaults, using set union (|) for combining allowlists.
87-101: Usingdefault_factory=listis correct for mutable defaults.This avoids the mutable default argument pitfall in Pydantic.
105-124: Boolean flags with| Nonetype anddefault=False.The
bool | Nonetype allows for explicit None values in config, but the default is False. This is intentional for YAML configs where omission means use default.
128-132: Policy configuration fields are well-defined.Using
list[FunctionPolicyRef]ensures type safety when referencing policies.src/nat/middleware/middleware.py (3)
35-38: Type signature relaxation enables flexible argument passing.The change from
Callable[[Any], ...]toCallable[..., ...]intentionally reduces type specificity to support passing multiple arguments (value + kwargs) through the middleware chain. This is necessary for the Function Policies framework to thread contextual arguments but means the type checker can no longer validate call arity or argument types at these boundaries.
150-184: Streaming path correctly mirrors the non-streaming changes.The
middleware_streamupdates maintain symmetry withmiddleware_invoke, ensuring**kwargsare propagated through both execution paths. The example code demonstrates the correct pattern for transforming stream chunks while preserving context.
114-148: Verify all Middleware subclasses accept **kwargs.The addition of
**kwargsto the signature enables threading contextual arguments (e.g., function policy contexts) through the middleware chain. However, any existingMiddlewaresubclass that overridesmiddleware_invokewithout accepting**kwargswill raise aTypeErrorat runtime if the caller passes extra arguments. This represents a potential breaking change for the core toolkit.Ensure that:
- All existing
Middlewaresubclasses in the codebase accept**kwargsin theirmiddleware_invokeoverride- If subclasses exist without
**kwargs, either update them or implement a compatibility layer- Document this requirement for users extending the middleware system
Signed-off-by: Eric Evans <[email protected]>
…eature/dynamic-middleware-function-policies
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
src/nat/middleware/dynamic/dynamic_function_middleware.py (2)
608-616: Double registration tracking in_configure_component_function_middleware.The registration key is added to
_registered_callablesat line 614, but the caller_register_component_functionalso adds it at line 564. This creates redundant tracking.if inspect.isasyncgenfunction(original_function): wrapped_function = chain.build_stream(original_function) else: wrapped_function = chain.build_single(original_function) - # Track registration - self._registered_callables.add(registration_key) - return wrapped_function
62-67: Potential KeyError if policy reference is missing from builder.Accessing
builder._function_policies[str(ref)]will raiseKeyErrorif the policy isn't registered. Consider adding validation or a more descriptive error message as previously suggested.
🧹 Nitpick comments (4)
src/nat/middleware/dynamic/dynamic_function_middleware.py (4)
116-120: Consider usinglogger.exceptionfor better debugging of registration failures.The pattern of catching all exceptions and logging at debug level is repeated throughout the discovery methods. While the resilient approach is appropriate, using
logger.exceptioninstead oflogger.debugwould capture the full stack trace when needed.try: self._register_component_function(discovered_component, function_name) except Exception as e: - logger.debug("Failed to register component function '%s' on LLM '%s': %s", function_name, llm_name, e) + logger.debug("Failed to register component function '%s' on LLM '%s': %s", + function_name, llm_name, e, exc_info=True)This applies similarly to lines 155-159, 194-198, 232-236, 272-276, and 312-316.
819-851: Unused parametercomponent_typeand silent exception swallowing.Two issues flagged by static analysis:
The
component_typeparameter (line 819) is documented as "for logging/metadata, not filtering" but is never used. Either use it in the debug logging or remove it.Lines 847-849 silently swallow exceptions during introspection without logging. Consider adding minimal logging for debugging purposes.
def _get_callable_functions(self, instance: Any, component_type: str | None = None) -> set[str]: """Get all callable functions from component instance that can be safely wrapped. This discovers ALL potentially wrappable component functions without allowlist filtering. Safety checks ensure only valid, callable, bound functions are included. Args: instance: The component instance to introspect - component_type: Type of component (for logging/metadata, not filtering) + component_type: Type of component (unused, reserved for future use) Returns: Set of all valid component function names that could be wrapped """ functions = set() for function_name in dir(instance): # Skip private/dunder functions if function_name.startswith('_'): continue try: # Must pass basic validity checks (no errors) if not self._is_valid_wrappable_function(instance, function_name): continue # Passed all safety checks - this component function CAN be wrapped functions.add(function_name) - except Exception: - # Skip functions that raise errors during introspection + except Exception as e: + # Skip functions that raise errors during introspection + logger.debug("Skipping function '%s' due to introspection error: %s", function_name, e) continue return functions
45-49: Consider expanding the class docstring.Per coding guidelines, public classes require comprehensive Google-style docstrings. The current docstring briefly describes functionality but could benefit from additional details about usage and attributes.
class DynamicFunctionMiddleware(FunctionMiddleware): - """Middleware extends FunctionMiddleware to provide dynamic discovery and - interception of all workflow components, including functions and components, without requiring explicit - per-component configuration. + """Dynamic middleware for automatic component discovery and function interception. + + This middleware extends FunctionMiddleware to provide dynamic discovery and + interception of workflow components (LLMs, embedders, retrievers, memory, + object stores, auth providers) and functions without requiring explicit + per-component configuration. It patches builder methods to intercept + component creation and automatically applies function policies. + + Attributes: + _config: The dynamic middleware configuration. + _builder: Reference to the workflow builder. + _pre_invoke_policies: List of policies to run before function execution. + _post_invoke_policies: List of policies to run after function execution. + _workflow_inventory: Inventory of discovered components and functions. """
163-172: Missing return type hint.The method signature is missing the return type annotation. Per coding guidelines, all public APIs require type hints.
- async def _discover_and_register_retriever(self, retriever_name: str, wrapper_type: Any = None): + async def _discover_and_register_retriever(self, retriever_name: str, wrapper_type: Any = None) -> Any:
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
docs/source/reference/cli.md(2 hunks)docs/source/reference/middleware.md(6 hunks)src/nat/middleware/dynamic/dynamic_function_middleware.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/source/reference/cli.md
🧰 Additional context used
📓 Path-based instructions (14)
**/*.{md,rst,py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NVIDIA NeMo Agent toolkit' on first use, then 'NeMo Agent toolkit' for subsequent references
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pydocs/source/reference/middleware.md
**/*.{py,toml,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' for environment variable prefixes and informal comments
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Follow PEP 20 and PEP 8 for Python style guidelines
Run yapf second (PEP 8 base, column_limit = 120) for Python formatting
Use ruff check --fix (via pre-commit) for linting using configuration embedded in pyproject.toml, fix warnings unless explicitly ignored
Use snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
All public APIs require Python 3.11+ type hints on parameters and return values
Prefer collections.abc / typing abstractions (Sequence over list) for type hints
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Preserve stack traces and prevent duplicate logging when handling exceptions; use bare raise statements and logger.error() when re-raising
When catching and logging exceptions without re-raising, always use logger.exception() to capture the full stack trace information
Provide Google-style docstrings for every public module, class, function and CLI command
The first line of docstrings must be a concise description ending with a period (Vale checks this)
Surround code entities with backticks in docstrings to avoid Vale false-positives
Validate and sanitise all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work (HTTP, DB, file reads)
Cache expensive computations with functools.lru_cache or an external cache when appropriate
Leverage NumPy vectorised operations whenever beneficial and feasible
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.py
**/*.{py,js,ts,yaml,yml,json,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces, never tabs, and ensure every file ends with a single newline
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pydocs/source/reference/middleware.md
**/*.{py,env,toml,yaml,yml,json}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never commit API keys, credentials or personal data; use environment variables or .env files excluded from Git
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.py
**/*.{py,js,ts,java,cpp,c,go,rb,php}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Every file must start with the standard SPDX Apache-2.0 header
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.py
**/*.{py,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived automatically by setuptools-scm; never hard-code them in code or docs
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pydocs/source/reference/middleware.md
**/*.{py,js,ts,java,cpp,c,go,rb,php,sh}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All source files must include the SPDX Apache-2.0 header template
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.py
**/*
⚙️ CodeRabbit configuration file
**/*: # Code Review Instructions
- Ensure the code follows best practices and coding standards. - For Python code, follow
PEP 20 and
PEP 8 for style guidelines.- Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values (except for return values of
None,
in that situation no return type hint is needed).
Example:def my_function(param1: int, param2: str) -> bool: pass- For Python exception handling, ensure proper stack trace preservation:
- When re-raising exceptions: use bare
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.exception()) to avoid duplicate stack trace output.- When catching and logging exceptions without re-raising: always use
logger.exception()
to capture the full stack trace information.Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any
words listed in the
ci/vale/styles/config/vocabularies/nat/reject.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile are OK.Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,
and should contain an Apache License 2.0 header comment at the top of each file.
- Confirm that copyright years are up-to date whenever a file is changed.
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pydocs/source/reference/middleware.md
src/nat/**/*
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.py
**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NeMo Agent Toolkit' (capitalize 'T') when the name appears in headings
Files:
docs/source/reference/middleware.md
docs/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never use deprecated names: Agent Intelligence toolkit, aiqtoolkit, AgentIQ, AIQ/aiq in documentation (unless intentionally referring to deprecated versions)
Files:
docs/source/reference/middleware.md
docs/source/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
docs/source/**/*.{md,rst}: Documentation must be clear and comprehensive, without TODOs, FIXMEs, or placeholder text
Ensure documentation is free of offensive or outdated terms
Ensure documentation is free of spelling mistakes and do not use words listed in ci/vale/styles/config/vocabularies/nat/reject.txt
Files:
docs/source/reference/middleware.md
docs/source/**/*
⚙️ CodeRabbit configuration file
This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the
docs/source/_staticdirectory.
Files:
docs/source/reference/middleware.md
🧬 Code graph analysis (1)
src/nat/middleware/dynamic/dynamic_function_middleware.py (4)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/middleware/dynamic/dynamic_middleware_config.py (1)
DynamicMiddlewareConfig(78-138)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)src/nat/middleware/utils/workflow_inventory.py (2)
DiscoveredComponent(55-74)DiscoveredFunction(77-90)
🪛 Ruff (0.14.7)
src/nat/middleware/dynamic/dynamic_function_middleware.py
118-118: Do not catch blind exception: Exception
(BLE001)
155-155: Do not catch blind exception: Exception
(BLE001)
194-194: Do not catch blind exception: Exception
(BLE001)
232-232: Do not catch blind exception: Exception
(BLE001)
272-272: Do not catch blind exception: Exception
(BLE001)
312-312: Do not catch blind exception: Exception
(BLE001)
535-535: Avoid specifying long messages outside the exception class
(TRY003)
539-541: Avoid specifying long messages outside the exception class
(TRY003)
583-584: Avoid specifying long messages outside the exception class
(TRY003)
660-660: Consider moving this statement to an else block
(TRY300)
661-661: Do not catch blind exception: Exception
(BLE001)
700-701: Abstract raise to an inner function
(TRY301)
700-701: Avoid specifying long messages outside the exception class
(TRY003)
775-776: Abstract raise to an inner function
(TRY301)
775-776: Avoid specifying long messages outside the exception class
(TRY003)
819-819: Unused method argument: component_type
(ARG002)
847-849: try-except-continue detected, consider logging the exception
(S112)
847-847: Do not catch blind exception: Exception
(BLE001)
902-902: Consider moving this statement to an else block
(TRY300)
904-904: Do not catch blind exception: Exception
(BLE001)
920-920: Avoid specifying long messages outside the exception class
(TRY003)
928-928: Avoid specifying long messages outside the exception class
(TRY003)
936-936: Avoid specifying long messages outside the exception class
(TRY003)
944-944: Avoid specifying long messages outside the exception class
(TRY003)
952-953: Avoid specifying long messages outside the exception class
(TRY003)
961-961: Avoid specifying long messages outside the exception class
(TRY003)
969-969: Avoid specifying long messages outside the exception class
(TRY003)
979-979: Avoid specifying long messages outside the exception class
(TRY003)
985-985: Avoid specifying long messages outside the exception class
(TRY003)
991-991: Avoid specifying long messages outside the exception class
(TRY003)
997-997: Avoid specifying long messages outside the exception class
(TRY003)
1003-1003: Avoid specifying long messages outside the exception class
(TRY003)
1009-1009: Avoid specifying long messages outside the exception class
(TRY003)
1015-1015: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (7)
src/nat/middleware/dynamic/dynamic_function_middleware.py (3)
1-14: LGTM: License header is correct.The Apache 2.0 SPDX header is properly formatted with correct copyright year (2025).
698-710: Validation errors are silently logged and skipped instead of propagating.When a policy returns an incorrect number of args (lines 699-702), a
ValueErroris raised but immediately caught by the outer exception handler and logged. This means malformed policy responses are silently ignored rather than failing fast.If validation errors should propagate, consider restructuring:
try: modified_args = await policy.on_pre_invoke(pre_context) if modified_args is not None: if len(modified_args) != len(pre_context.function_args): - raise ValueError(f"Policy '{policy.name}' returned {len(modified_args)} args, " - f"expected {len(pre_context.function_args)}") + logger.error("Policy '%s' returned %d args, expected %d - skipping policy", + policy.name, len(modified_args), len(pre_context.function_args)) + continue pre_context.function_args = modified_args except Exception as e:Alternatively, if validation errors should fail the call, move the validation outside the try block.
907-973: LGTM: Patch methods are well-structured with proper guard checks.The patch methods correctly:
- Verify the builder has the expected methods before patching
- Store references to original methods
- Replace with interception wrappers
The long exception messages (flagged by TRY003) are acceptable here as they provide helpful debugging context for misconfiguration errors.
docs/source/reference/middleware.md (4)
129-172: LGTM: Middleware examples correctly demonstrate the variadic signature pattern.The logging middleware example properly shows:
*args: Anyand**kwargs: Anyin method signatures- Forwarding via
call_next(*args, **kwargs)- Accessing args for logging purposes
This aligns with the new function policy framework's design.
345-359: LGTM: Validation middleware example correctly demonstrates args handling.The example properly shows:
- Extracting the first arg for validation:
value = args[0] if args else None- Forwarding with modified first arg:
return await call_next(validated, *args[1:], **kwargs)This is a good pattern for middleware that needs to transform specific positional arguments.
529-538: LGTM: Test example correctly demonstrates the new invocation pattern.The mock function and test invocation properly show:
async def mock_next(*args, **kwargs)pattern- Accessing
args[0]within the mock- Calling
function_middleware_invoke(5, call_next=mock_next, context=context)
673-674: Verify the updated cache middleware import paths.The API references have been updated to:
nat.middleware.cache.cache_middleware_config.CacheMiddlewareConfignat.middleware.cache.cache_middleware.CacheMiddlewareEnsure these paths match the actual module structure after the cache middleware refactor.
Signed-off-by: Eric Evans <[email protected]>
|
Function policies are a higher-level abstraction that runs within DynamicFunctionMiddleware, which handles auto-discovery, allowlist security, and policy orchestration for both sync and streaming functions. Policies provide a simple on_pre_invoke/on_post_invoke interface for inspecting or transforming inputs and outputs, users just receive a context object and return a value. Unlike middleware, policies don't require knowing which functions to intercept at compile time; the dynamic middleware discovers and wraps components at runtime, then applies policies to them. This lets users build and share reusable policy logic without writing custom middleware or manually registering functions. @dnandakumar-nv |
…ration of discovered functions Signed-off-by: Eric Evans <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/source/reference/middleware.md (1)
20-30: Align product naming with branding guidelines.In Line 22, the first mention should follow the guideline: use "
NVIDIA NeMo Agent toolkit" on first use, and "NeMo Agent toolkit" for subsequent mentions. Also, reserve "NeMo Agent Toolkit" (capitalT) for headings only. Please adjust the body text accordingly.
🧹 Nitpick comments (2)
src/nat/middleware/dynamic/dynamic_function_middleware.py (1)
690-713: Helper introspection utilities are defensive; minor cleanups possible.
_get_callable_functions’scomponent_typeparameter is currently unused. Either:
- Use it for debug logging (e.g., to annotate which component type is being scanned), or
- Drop the parameter to avoid
ARG002from ruff.
_get_callable_functionsand_is_valid_wrappable_functionintentionally catchExceptionto be robust against arbitrary component implementations; that’s reasonable for a plugin-style system. If you want better observability, you could log atDEBUGwhen introspection of a particular attribute fails.
_extract_component_attributesalso swallows all exceptions and returnsNone. Given this is purely for metadata enrichment, this fail-closed behavior is acceptable, but consider at least alogger.debugwhen an unexpected exception occurs to help debugging odd components.Also applies to: 869-901, 903-956
tests/nat/middleware/test_dynamic_middleware.py (1)
243-417: Confirm async test execution strategy (pytest-asyncio vs asyncio_mode).Most tests in this file are declared as
async def(e.g., Lines 243–353, 336–353, 358–417, 422–444, 449–559) but aren’t decorated with@pytest.mark.asyncio, even though the project guideline calls out pytest-asyncio for async tests.If the repo is using
pytest-asynciowithasyncio_mode = "auto"(or similar) inpytest.ini, this is fine. Otherwise, these tests will fail to run. Either:
- Add
pytestmark = pytest.mark.asyncioat the module level, or- Decorate each async test with
@pytest.mark.asyncio.Also applies to: 422-444, 449-559
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
docs/source/reference/middleware.md(7 hunks)src/nat/middleware/dynamic/dynamic_function_middleware.py(1 hunks)src/nat/middleware/utils/workflow_inventory.py(1 hunks)tests/nat/middleware/test_dynamic_middleware.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (16)
**/*.{md,rst,py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NVIDIA NeMo Agent toolkit' on first use, then 'NeMo Agent toolkit' for subsequent references
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pydocs/source/reference/middleware.mdsrc/nat/middleware/utils/workflow_inventory.py
**/*.{py,toml,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' for environment variable prefixes and informal comments
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/middleware/utils/workflow_inventory.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Follow PEP 20 and PEP 8 for Python style guidelines
Run yapf second (PEP 8 base, column_limit = 120) for Python formatting
Use ruff check --fix (via pre-commit) for linting using configuration embedded in pyproject.toml, fix warnings unless explicitly ignored
Use snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
All public APIs require Python 3.11+ type hints on parameters and return values
Prefer collections.abc / typing abstractions (Sequence over list) for type hints
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Preserve stack traces and prevent duplicate logging when handling exceptions; use bare raise statements and logger.error() when re-raising
When catching and logging exceptions without re-raising, always use logger.exception() to capture the full stack trace information
Provide Google-style docstrings for every public module, class, function and CLI command
The first line of docstrings must be a concise description ending with a period (Vale checks this)
Surround code entities with backticks in docstrings to avoid Vale false-positives
Validate and sanitise all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work (HTTP, DB, file reads)
Cache expensive computations with functools.lru_cache or an external cache when appropriate
Leverage NumPy vectorised operations whenever beneficial and feasible
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/middleware/utils/workflow_inventory.py
**/*.{py,js,ts,yaml,yml,json,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces, never tabs, and ensure every file ends with a single newline
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pydocs/source/reference/middleware.mdsrc/nat/middleware/utils/workflow_inventory.py
**/*.{py,env,toml,yaml,yml,json}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never commit API keys, credentials or personal data; use environment variables or .env files excluded from Git
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/middleware/utils/workflow_inventory.py
**/*.{py,js,ts,java,cpp,c,go,rb,php}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Every file must start with the standard SPDX Apache-2.0 header
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/middleware/utils/workflow_inventory.py
**/*.{py,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived automatically by setuptools-scm; never hard-code them in code or docs
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pydocs/source/reference/middleware.mdsrc/nat/middleware/utils/workflow_inventory.py
**/*.{py,js,ts,java,cpp,c,go,rb,php,sh}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All source files must include the SPDX Apache-2.0 header template
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pysrc/nat/middleware/utils/workflow_inventory.py
**/*
⚙️ CodeRabbit configuration file
**/*: # Code Review Instructions
- Ensure the code follows best practices and coding standards. - For Python code, follow
PEP 20 and
PEP 8 for style guidelines.- Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values (except for return values of
None,
in that situation no return type hint is needed).
Example:def my_function(param1: int, param2: str) -> bool: pass- For Python exception handling, ensure proper stack trace preservation:
- When re-raising exceptions: use bare
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.exception()) to avoid duplicate stack trace output.- When catching and logging exceptions without re-raising: always use
logger.exception()
to capture the full stack trace information.Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any
words listed in the
ci/vale/styles/config/vocabularies/nat/reject.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile are OK.Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,
and should contain an Apache License 2.0 header comment at the top of each file.
- Confirm that copyright years are up-to date whenever a file is changed.
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pytests/nat/middleware/test_dynamic_middleware.pydocs/source/reference/middleware.mdsrc/nat/middleware/utils/workflow_inventory.py
src/nat/**/*
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/middleware/utils/workflow_inventory.py
**/test_*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/test_*.py: Use pytest with pytest-asyncio for asynchronous code testing
Test functions should be named using the test_ prefix, using snake_case
Extract frequently repeated code into pytest fixtures, which should be named using the fixture_ prefix and define the name argument in the decorator
Mock external services with pytest_httpserver or unittest.mock instead of hitting live endpoints
Mark slow tests with @pytest.mark.slow so they can be skipped in the default test suite
Mark integration tests requiring external services with @pytest.mark.integration so they can be skipped in the default test suite
Files:
tests/nat/middleware/test_dynamic_middleware.py
tests/**/*.py
⚙️ CodeRabbit configuration file
tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using thetest_prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using thefixture_prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass
Files:
tests/nat/middleware/test_dynamic_middleware.py
**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NeMo Agent Toolkit' (capitalize 'T') when the name appears in headings
Files:
docs/source/reference/middleware.md
docs/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never use deprecated names: Agent Intelligence toolkit, aiqtoolkit, AgentIQ, AIQ/aiq in documentation (unless intentionally referring to deprecated versions)
Files:
docs/source/reference/middleware.md
docs/source/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
docs/source/**/*.{md,rst}: Documentation must be clear and comprehensive, without TODOs, FIXMEs, or placeholder text
Ensure documentation is free of offensive or outdated terms
Ensure documentation is free of spelling mistakes and do not use words listed in ci/vale/styles/config/vocabularies/nat/reject.txt
Files:
docs/source/reference/middleware.md
docs/source/**/*
⚙️ CodeRabbit configuration file
This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the
docs/source/_staticdirectory.
Files:
docs/source/reference/middleware.md
🧬 Code graph analysis (2)
tests/nat/middleware/test_dynamic_middleware.py (3)
src/nat/middleware/dynamic/dynamic_function_middleware.py (4)
DynamicFunctionMiddleware(45-1066)function_middleware_invoke(716-789)function_middleware_stream(791-865)_discover_and_register_llm(107-141)src/nat/middleware/utils/workflow_inventory.py (3)
DiscoveredFunction(83-96)RegisteredFunction(110-113)RegisteredComponentMethod(116-121)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)
src/nat/middleware/utils/workflow_inventory.py (2)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/data_models/function.py (1)
FunctionBaseConfig(26-36)
🪛 Ruff (0.14.7)
src/nat/middleware/dynamic/dynamic_function_middleware.py
101-102: Avoid specifying long messages outside the exception class
(TRY003)
138-138: Do not catch blind exception: Exception
(BLE001)
175-175: Do not catch blind exception: Exception
(BLE001)
214-214: Do not catch blind exception: Exception
(BLE001)
252-252: Do not catch blind exception: Exception
(BLE001)
292-292: Do not catch blind exception: Exception
(BLE001)
332-332: Do not catch blind exception: Exception
(BLE001)
340-340: Undefined name FunctionRef
(F821)
555-555: Avoid specifying long messages outside the exception class
(TRY003)
559-561: Avoid specifying long messages outside the exception class
(TRY003)
603-603: Avoid specifying long messages outside the exception class
(TRY003)
636-637: Avoid specifying long messages outside the exception class
(TRY003)
652-652: Undefined name FunctionMiddlewareContext
(F821)
710-710: Consider moving this statement to an else block
(TRY300)
711-711: Do not catch blind exception: Exception
(BLE001)
719-719: Undefined name FunctionMiddlewareContext
(F821)
750-751: Abstract raise to an inner function
(TRY301)
750-751: Avoid specifying long messages outside the exception class
(TRY003)
794-794: Undefined name FunctionMiddlewareContext
(F821)
825-826: Abstract raise to an inner function
(TRY301)
825-826: Avoid specifying long messages outside the exception class
(TRY003)
869-869: Unused method argument: component_type
(ARG002)
897-899: try-except-continue detected, consider logging the exception
(S112)
897-897: Do not catch blind exception: Exception
(BLE001)
952-952: Consider moving this statement to an else block
(TRY300)
954-954: Do not catch blind exception: Exception
(BLE001)
970-970: Avoid specifying long messages outside the exception class
(TRY003)
978-978: Avoid specifying long messages outside the exception class
(TRY003)
986-986: Avoid specifying long messages outside the exception class
(TRY003)
994-994: Avoid specifying long messages outside the exception class
(TRY003)
1002-1003: Avoid specifying long messages outside the exception class
(TRY003)
1011-1011: Avoid specifying long messages outside the exception class
(TRY003)
1019-1019: Avoid specifying long messages outside the exception class
(TRY003)
1029-1029: Avoid specifying long messages outside the exception class
(TRY003)
1035-1035: Avoid specifying long messages outside the exception class
(TRY003)
1041-1041: Avoid specifying long messages outside the exception class
(TRY003)
1047-1047: Avoid specifying long messages outside the exception class
(TRY003)
1053-1053: Avoid specifying long messages outside the exception class
(TRY003)
1059-1059: Avoid specifying long messages outside the exception class
(TRY003)
1065-1065: Avoid specifying long messages outside the exception class
(TRY003)
tests/nat/middleware/test_dynamic_middleware.py
76-76: Unused method argument: messages
(ARG002)
76-76: Unused method argument: kwargs
(ARG002)
79-79: Unused method argument: messages
(ARG002)
79-79: Unused method argument: kwargs
(ARG002)
82-82: Unused method argument: messages
(ARG002)
82-82: Unused method argument: kwargs
(ARG002)
85-85: Unused method argument: messages
(ARG002)
85-85: Unused method argument: kwargs
(ARG002)
97-97: Unused method argument: text
(ARG002)
100-100: Unused method argument: text
(ARG002)
112-112: Unused method argument: query
(ARG002)
112-112: Unused method argument: kwargs
(ARG002)
124-124: Unused method argument: query
(ARG002)
124-124: Unused method argument: top_k
(ARG002)
124-124: Unused method argument: kwargs
(ARG002)
169-169: Unused method argument: user_id
(ARG002)
169-169: Unused method argument: kwargs
(ARG002)
224-224: Consider (modified_first, *context.function_args[1:]) instead of concatenation
Replace with (modified_first, *context.function_args[1:])
(RUF005)
251-251: Unused function argument: kwargs
(ARG001)
278-278: Unused function argument: kwargs
(ARG001)
305-305: Unused function argument: args
(ARG001)
305-305: Unused function argument: kwargs
(ARG001)
326-326: Unused function argument: args
(ARG001)
326-326: Unused function argument: kwargs
(ARG001)
341-341: Unused function argument: args
(ARG001)
341-341: Unused function argument: kwargs
(ARG001)
370-370: Unused function argument: kwargs
(ARG001)
390-390: Unused function argument: args
(ARG001)
390-390: Unused function argument: kwargs
(ARG001)
409-409: Unused function argument: args
(ARG001)
409-409: Unused function argument: kwargs
(ARG001)
434-434: Unused function argument: args
(ARG001)
434-434: Unused function argument: kwargs
(ARG001)
716-716: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI Pipeline / Check
🔇 Additional comments (5)
docs/source/reference/middleware.md (2)
129-172: Variadic middleware examples and tests look consistent.The updated
LoggingMiddlewareandValidationMiddlewareexamples correctly use*args/**kwargsand forward all arguments tocall_next, and the unit-test snippet reflects the new call signature. This matches the new middleware API and preserves backward compatibility for call sites.Also applies to: 345-359, 529-537
703-711: API reference additions for cache middleware and plugin link look good.The new API entries for
CacheMiddlewareConfig/CacheMiddlewareand the See Also link to the plugin system are accurate and helpful cross-references to the reorganized middleware API.src/nat/middleware/dynamic/dynamic_function_middleware.py (2)
51-67: DynamicFunctionMiddleware initialization and policy wiring look solid.The constructor cleanly wires pre/post policy instances via
_get_policy_instance, initializes the workflow inventory and component allowlists, and kicks off discovery/patching. The helper_get_policy_instancenow raises a clearValueErrorwith available policy names, which should make misconfiguration easier to diagnose.Also applies to: 79-82, 85-104
957-1067: Builder patching and original-method getters look correct and self-consistent.The
_patch_get_*methods preserve the original builder methods and replace them with dynamic wrappers, and the_get_builder_get_*accessors guard against use before patching. The error messages for missing builder methods are clear, and this pattern should behave well with both production builders and the mocked builders in tests.src/nat/middleware/utils/workflow_inventory.py (1)
29-53: Workflow inventory and allowlists are well-structured and match dynamic middleware usage.The
COMPONENT_FUNCTION_ALLOWLISTSconstant and theDiscovered*,Registered*, andWorkflowInventorymodels provide a clear, Pydantic-validated structure for discovery and registration. The field descriptions andarbitrary_types_allowedconfig fit the dynamic component surface, and the sets of allowed method names align with the documented component method allowlists.Also applies to: 56-155
Signed-off-by: Eric Evans <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
docs/source/reference/middleware.md (1)
18-31: Dynamic middleware unregister section now aligns with the public API.The new “Dynamic Middleware: Unregistering Callables” section:
- Uses
get_registered()/get_registered_keys()instead of reaching into_registered_callables.- Correctly describes behavior for workflow functions vs component methods.
- Documents the
RegisteredFunctionandRegisteredComponentMethodmodels in a concise way.This is much clearer and better aligned with the actual surface of
DynamicFunctionMiddleware.Also applies to: 639-672
🧹 Nitpick comments (7)
src/nat/middleware/dynamic/dynamic_function_middleware.py (4)
53-71: Policy resolution helper looks good; consider tightening typing and encapsulation.The
_get_policy_instancehelper cleanly addresses missing-policy diagnostics and makes the pre/post policy lists safer to build. Two small follow‑ups you might consider:
- Type
builderasBuilderinstead ofAnyso pyright can catch misuse.- If
_function_policiesis meant to remain a private detail ofBuilder, you could optionally front it with a small public accessor onBuilderand call that here to reduce direct reliance on a private attribute.These are style/maintainability points only; behavior as written is sound.
Also applies to: 87-105
545-595: Registration failure handling is intentionally tolerant; consider richer logging when debugging.Catching
Exceptionaround_register_component_functionand logging at debug level matches the “best effort” discovery story, but it can make diagnosing misconfigurations or unexpected errors harder in production.If you hit debuggability issues later, you might:
- Narrow the exception type for expected discovery errors (e.g.,
ValueError) and- Use
logger.exception(perhaps at debug level) for genuinely unexpected failures so the stack trace is still available when needed.Given the design goals, the current behavior is acceptable.
691-709: Allowlist building and callable discovery are robust; tidy up a couple of nits.The allowlist construction and
_get_callable_functions/_is_valid_wrappable_functionlogic gives a clear two‑stage filter: “can be wrapped” vs “allowed to be auto‑wrapped”. That’s a good separation.Two minor cleanups you might consider:
_get_callable_functions’scomponent_typeargument is currently unused; either remove it or leverage it in debug logs for easier troubleshooting.- In
_get_callable_functionsand_is_valid_wrappable_function, you silently swallow all exceptions; if you ever need to debug odd components, adding alogger.debug/logger.exceptionwith thecomponent_typeandfunction_namewould help without changing behavior.These are small maintainability tweaks, not correctness issues.
Also applies to: 880-913, 914-967
968-1078: Builder patching and original‑getter helpers are straightforward; consider adding return type hints.The
_patch_*methods and their_get_builder_get_*counterparts clearly separate the patched entrypoints from the saved originals and fail fast if a getter hasn’t been patched.To align more closely with the typing guideline, you could add explicit return types such as
Callable[..., Any]to the_get_builder_get_*helpers. This would help pyright and readers without changing behavior.src/nat/function_policy/interface.py (1)
69-125: FunctionPolicyBase contract is clear; small typing/style polish possible.The abstract
on_pre_invoke/on_post_invokesignatures and docstrings clearly define:
- Return‑
Noneas “no change”.- Tuple length invariants for args.
- Generic
PolicyConfigTfor config typing.One minor nit: ruff flags
__all__ordering; you may want to sort it (e.g.,["FunctionPolicyBase", "PolicyConfigT", "PostInvokeContext", "PreInvokeContext"]) to keep style checks happy.tests/nat/middleware/test_dynamic_middleware.py (1)
449-560: Component discovery tests hit the main component types; consider expanding to cover allowlists if they change often.The discovery tests for LLMs, embedders, retrievers, memory, object stores, and auth providers validate:
- That discovery only occurs when the corresponding
register_*flag is enabled.- That duplicates are not added twice.
If you later evolve
COMPONENT_FUNCTION_ALLOWLISTSor per‑type allowlist config, it may be worth adding a focused test that asserts only allowlisted methods end up registered (not just discovered), but the current coverage is already good for this PR.docs/source/reference/middleware.md (1)
22-31: Minor product‑name wording tweak to match style guide.The first mention says “functions in the NeMo Agent Toolkit”. Per the guidelines, prose should use “NVIDIA NeMo Agent toolkit” on first mention and “NeMo Agent toolkit” thereafter (docs headings already follow the capital‑T variant).
If you touch this section again, consider updating that phrase to match the project’s preferred wording.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
ci/scripts/path_checks.py(1 hunks)ci/vale/styles/config/vocabularies/nat/accept.txt(1 hunks)docs/source/reference/middleware.md(7 hunks)src/nat/function_policy/interface.py(1 hunks)src/nat/middleware/dynamic/dynamic_function_middleware.py(1 hunks)tests/nat/middleware/test_dynamic_middleware.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (16)
**/*.{md,rst,py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NVIDIA NeMo Agent toolkit' on first use, then 'NeMo Agent toolkit' for subsequent references
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pydocs/source/reference/middleware.mdtests/nat/middleware/test_dynamic_middleware.py
**/*.{py,toml,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' for environment variable prefixes and informal comments
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pytests/nat/middleware/test_dynamic_middleware.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Follow PEP 20 and PEP 8 for Python style guidelines
Run yapf second (PEP 8 base, column_limit = 120) for Python formatting
Use ruff check --fix (via pre-commit) for linting using configuration embedded in pyproject.toml, fix warnings unless explicitly ignored
Use snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
All public APIs require Python 3.11+ type hints on parameters and return values
Prefer collections.abc / typing abstractions (Sequence over list) for type hints
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Preserve stack traces and prevent duplicate logging when handling exceptions; use bare raise statements and logger.error() when re-raising
When catching and logging exceptions without re-raising, always use logger.exception() to capture the full stack trace information
Provide Google-style docstrings for every public module, class, function and CLI command
The first line of docstrings must be a concise description ending with a period (Vale checks this)
Surround code entities with backticks in docstrings to avoid Vale false-positives
Validate and sanitise all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work (HTTP, DB, file reads)
Cache expensive computations with functools.lru_cache or an external cache when appropriate
Leverage NumPy vectorised operations whenever beneficial and feasible
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pytests/nat/middleware/test_dynamic_middleware.py
**/*.{py,js,ts,yaml,yml,json,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces, never tabs, and ensure every file ends with a single newline
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pydocs/source/reference/middleware.mdtests/nat/middleware/test_dynamic_middleware.py
**/*.{py,env,toml,yaml,yml,json}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never commit API keys, credentials or personal data; use environment variables or .env files excluded from Git
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pytests/nat/middleware/test_dynamic_middleware.py
**/*.{py,js,ts,java,cpp,c,go,rb,php}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Every file must start with the standard SPDX Apache-2.0 header
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pytests/nat/middleware/test_dynamic_middleware.py
**/*.{py,md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived automatically by setuptools-scm; never hard-code them in code or docs
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pydocs/source/reference/middleware.mdtests/nat/middleware/test_dynamic_middleware.py
**/*.{py,js,ts,java,cpp,c,go,rb,php,sh}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All source files must include the SPDX Apache-2.0 header template
Files:
ci/scripts/path_checks.pysrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pytests/nat/middleware/test_dynamic_middleware.py
**/*
⚙️ CodeRabbit configuration file
**/*: # Code Review Instructions
- Ensure the code follows best practices and coding standards. - For Python code, follow
PEP 20 and
PEP 8 for style guidelines.- Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values (except for return values of
None,
in that situation no return type hint is needed).
Example:def my_function(param1: int, param2: str) -> bool: pass- For Python exception handling, ensure proper stack trace preservation:
- When re-raising exceptions: use bare
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.exception()) to avoid duplicate stack trace output.- When catching and logging exceptions without re-raising: always use
logger.exception()
to capture the full stack trace information.Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any
words listed in the
ci/vale/styles/config/vocabularies/nat/reject.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile are OK.Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,
and should contain an Apache License 2.0 header comment at the top of each file.
- Confirm that copyright years are up-to date whenever a file is changed.
Files:
ci/scripts/path_checks.pyci/vale/styles/config/vocabularies/nat/accept.txtsrc/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.pydocs/source/reference/middleware.mdtests/nat/middleware/test_dynamic_middleware.py
src/nat/**/*
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/middleware/dynamic/dynamic_function_middleware.pysrc/nat/function_policy/interface.py
**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use 'NeMo Agent Toolkit' (capitalize 'T') when the name appears in headings
Files:
docs/source/reference/middleware.md
docs/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never use deprecated names: Agent Intelligence toolkit, aiqtoolkit, AgentIQ, AIQ/aiq in documentation (unless intentionally referring to deprecated versions)
Files:
docs/source/reference/middleware.md
docs/source/**/*.{md,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
docs/source/**/*.{md,rst}: Documentation must be clear and comprehensive, without TODOs, FIXMEs, or placeholder text
Ensure documentation is free of offensive or outdated terms
Ensure documentation is free of spelling mistakes and do not use words listed in ci/vale/styles/config/vocabularies/nat/reject.txt
Files:
docs/source/reference/middleware.md
docs/source/**/*
⚙️ CodeRabbit configuration file
This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the
docs/source/_staticdirectory.
Files:
docs/source/reference/middleware.md
**/test_*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/test_*.py: Use pytest with pytest-asyncio for asynchronous code testing
Test functions should be named using the test_ prefix, using snake_case
Extract frequently repeated code into pytest fixtures, which should be named using the fixture_ prefix and define the name argument in the decorator
Mock external services with pytest_httpserver or unittest.mock instead of hitting live endpoints
Mark slow tests with @pytest.mark.slow so they can be skipped in the default test suite
Mark integration tests requiring external services with @pytest.mark.integration so they can be skipped in the default test suite
Files:
tests/nat/middleware/test_dynamic_middleware.py
tests/**/*.py
⚙️ CodeRabbit configuration file
tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using thetest_prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using thefixture_prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass
Files:
tests/nat/middleware/test_dynamic_middleware.py
🧠 Learnings (1)
📚 Learning: 2025-11-24T18:56:53.109Z
Learnt from: CR
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 0
File: .cursor/rules/general.mdc:0-0
Timestamp: 2025-11-24T18:56:53.109Z
Learning: Words listed in ci/vale/styles/config/vocabularies/nat/accept.txt are acceptable even if they appear to be spelling mistakes
Applied to files:
ci/vale/styles/config/vocabularies/nat/accept.txt
🧬 Code graph analysis (2)
src/nat/middleware/dynamic/dynamic_function_middleware.py (2)
src/nat/data_models/component.py (1)
ComponentGroup(48-60)src/nat/middleware/function_middleware.py (1)
FunctionMiddleware(43-110)
tests/nat/middleware/test_dynamic_middleware.py (7)
src/nat/data_models/authentication.py (2)
AuthProviderBaseConfig(31-37)AuthResult(194-245)src/nat/middleware/dynamic/dynamic_function_middleware.py (3)
DynamicFunctionMiddleware(47-1077)function_middleware_invoke(737-803)function_middleware_stream(805-876)src/nat/middleware/dynamic/dynamic_middleware_config.py (1)
DynamicMiddlewareConfig(78-138)src/nat/middleware/utils/workflow_inventory.py (4)
DiscoveredComponent(66-80)DiscoveredFunction(83-96)RegisteredFunction(110-113)RegisteredComponentMethod(116-121)src/nat/builder/workflow_builder.py (14)
get_llm(709-727)get_llm(1454-1459)get_embedder(825-843)get_embedder(1470-1475)get_retriever(932-953)get_retriever(1547-1550)get_memory_client(869-876)get_memory_client(1486-1494)get_object_store_client(901-905)get_object_store_client(1505-1513)get_auth_provider(781-806)get_auth_provider(1450-1451)get_function(567-573)get_function(1382-1388)src/nat/middleware/middleware.py (1)
FunctionMiddlewareContext(42-66)src/nat/data_models/function_policy.py (1)
FunctionPolicyBaseConfig(24-31)
🪛 Ruff (0.14.7)
src/nat/middleware/dynamic/dynamic_function_middleware.py
103-104: Avoid specifying long messages outside the exception class
(TRY003)
140-140: Do not catch blind exception: Exception
(BLE001)
177-177: Do not catch blind exception: Exception
(BLE001)
216-216: Do not catch blind exception: Exception
(BLE001)
254-254: Do not catch blind exception: Exception
(BLE001)
294-294: Do not catch blind exception: Exception
(BLE001)
334-334: Do not catch blind exception: Exception
(BLE001)
557-557: Avoid specifying long messages outside the exception class
(TRY003)
561-563: Avoid specifying long messages outside the exception class
(TRY003)
624-624: Avoid specifying long messages outside the exception class
(TRY003)
657-658: Avoid specifying long messages outside the exception class
(TRY003)
731-731: Consider moving this statement to an else block
(TRY300)
732-732: Do not catch blind exception: Exception
(BLE001)
771-772: Abstract raise to an inner function
(TRY301)
771-772: Avoid specifying long messages outside the exception class
(TRY003)
839-840: Abstract raise to an inner function
(TRY301)
839-840: Avoid specifying long messages outside the exception class
(TRY003)
880-880: Unused method argument: component_type
(ARG002)
908-910: try-except-continue detected, consider logging the exception
(S112)
908-908: Do not catch blind exception: Exception
(BLE001)
963-963: Consider moving this statement to an else block
(TRY300)
965-965: Do not catch blind exception: Exception
(BLE001)
981-981: Avoid specifying long messages outside the exception class
(TRY003)
989-989: Avoid specifying long messages outside the exception class
(TRY003)
997-997: Avoid specifying long messages outside the exception class
(TRY003)
1005-1005: Avoid specifying long messages outside the exception class
(TRY003)
1013-1014: Avoid specifying long messages outside the exception class
(TRY003)
1022-1022: Avoid specifying long messages outside the exception class
(TRY003)
1030-1030: Avoid specifying long messages outside the exception class
(TRY003)
1040-1040: Avoid specifying long messages outside the exception class
(TRY003)
1046-1046: Avoid specifying long messages outside the exception class
(TRY003)
1052-1052: Avoid specifying long messages outside the exception class
(TRY003)
1058-1058: Avoid specifying long messages outside the exception class
(TRY003)
1064-1064: Avoid specifying long messages outside the exception class
(TRY003)
1070-1070: Avoid specifying long messages outside the exception class
(TRY003)
1076-1076: Avoid specifying long messages outside the exception class
(TRY003)
src/nat/function_policy/interface.py
128-133: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
tests/nat/middleware/test_dynamic_middleware.py
76-76: Unused method argument: messages
(ARG002)
76-76: Unused method argument: kwargs
(ARG002)
79-79: Unused method argument: messages
(ARG002)
79-79: Unused method argument: kwargs
(ARG002)
82-82: Unused method argument: messages
(ARG002)
82-82: Unused method argument: kwargs
(ARG002)
85-85: Unused method argument: messages
(ARG002)
85-85: Unused method argument: kwargs
(ARG002)
97-97: Unused method argument: text
(ARG002)
100-100: Unused method argument: text
(ARG002)
112-112: Unused method argument: query
(ARG002)
112-112: Unused method argument: kwargs
(ARG002)
124-124: Unused method argument: query
(ARG002)
124-124: Unused method argument: top_k
(ARG002)
124-124: Unused method argument: kwargs
(ARG002)
169-169: Unused method argument: user_id
(ARG002)
169-169: Unused method argument: kwargs
(ARG002)
224-224: Consider (modified_first, *context.function_args[1:]) instead of concatenation
Replace with (modified_first, *context.function_args[1:])
(RUF005)
251-251: Unused function argument: kwargs
(ARG001)
278-278: Unused function argument: kwargs
(ARG001)
305-305: Unused function argument: args
(ARG001)
305-305: Unused function argument: kwargs
(ARG001)
326-326: Unused function argument: args
(ARG001)
326-326: Unused function argument: kwargs
(ARG001)
341-341: Unused function argument: args
(ARG001)
341-341: Unused function argument: kwargs
(ARG001)
370-370: Unused function argument: kwargs
(ARG001)
390-390: Unused function argument: args
(ARG001)
390-390: Unused function argument: kwargs
(ARG001)
409-409: Unused function argument: args
(ARG001)
409-409: Unused function argument: kwargs
(ARG001)
434-434: Unused function argument: args
(ARG001)
434-434: Unused function argument: kwargs
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI Pipeline / Check
🔇 Additional comments (10)
ci/vale/styles/config/vocabularies/nat/accept.txt (1)
173-176: Vocabulary additions are correct and relevant to the PR.The four new entries are properly formatted Go regex patterns and directly support the PR's focus on function policies with dynamic middleware and Python kwargs/args forwarding:
[Ss]anitizationfor data sanitization features[Uu]nregister(ing|ed|s)?for policy/function unregistration[Aa]rgsand[Kk]wargsfor variable argument handlingci/scripts/path_checks.py (1)
114-114: LGTM!The addition of "Enable/disable" to the allowlist is appropriate given the Function Policies feature being introduced. The capitalization pattern and alphabetical placement are consistent with existing entries.
src/nat/middleware/dynamic/dynamic_function_middleware.py (3)
525-595: Registration/unregistration API is consistent and matches docs.
_register_function,_register_component_function,get_registered*, andunregistercooperate cleanly:
- Duplicate registration is idempotent and logged at debug level.
- Component methods correctly store the original callable and restore it on unregister.
- Workflow functions have this middleware removed from their chain without disturbing others.
This aligns well with the documented unregister example and keeps the surface area small and explicit.
Also applies to: 614-639
641-688: Component function wrapping correctly avoids double‑wrapping and preserves metadata.
_configure_component_function_middlewarecarefully:
- Skips re‑wrapping functions already present in
_registered_callables.- Extracts
config,description, and schema attributes defensively.- Chooses between
build_singleandbuild_streambased oninspect.isasyncgenfunction.This matches the intended FunctionMiddlewareChain usage and should keep wrapped methods predictable and debuggable.
737-804: Policy orchestration semantics for invoke/stream are consistent with the interface.
function_middleware_invokeandfunction_middleware_stream:
- Correctly thread a mutable
PreInvokeContextandPostInvokeContextthrough the chain.- Enforce that pre‑invoke policies preserve arg tuple length.
- Treat
Nonefrom pre/post hooks as “no change”.- Use
logger.exceptionwhen swallowing policy errors, matching the logging guideline.The output‑chaining semantics for post‑invoke policies now reflect the documented contract (each policy sees the latest
function_output), and the streaming path mirrors the single‑invoke path per chunk. This looks solid.Also applies to: 805-877
src/nat/function_policy/interface.py (1)
28-67: Context dataclasses align well with middleware usage.
PreInvokeContextandPostInvokeContextmatch the way DynamicFunctionMiddleware uses them:
original_argsstays immutable whilefunction_argsandfunction_outputare updated along the chain.- Types are explicit and compatible with the variadic
*args, **kwargssignatures.This provides a clear, policy‑friendly surface without leaking implementation details from the middleware.
tests/nat/middleware/test_dynamic_middleware.py (3)
41-55: Fixtures for builder and tracking policies provide good coverage of the dynamic path.
mock_builderand thetracking_policyfactory fixture do a nice job of:
- Exercising both pre‑ and post‑invoke hooks without going through the config‑based policy resolution.
- Verifying order, enable/disable flags, and error‑tolerant behavior.
This gives strong confidence in the DynamicFunctionMiddleware orchestration logic.
Also applies to: 178-237
243-354: Invoke/stream and error‑handling tests are thorough and match the design goals.The
function_middleware_invoke/function_middleware_streamtests:
- Validate the no‑policy passthrough behavior.
- Confirm that pre policies can modify inputs and post policies can modify outputs.
- Assert ordering (
pre→ function →post) and that failing policies are skipped while others still run.This directly tests the error‑resilient semantics described in the PR and looks solid.
Also applies to: 358-418
564-595: Workflow function and unregister tests validate the new registration API well.The tests around:
- Discovering existing workflow functions from
builder._functions.- Preventing duplicate function/component registration.
- Unregistering both workflow functions and component methods, including error cases.
all line up with the new
get_registered/unregisterbehavior and ensure state is correctly restored (e.g., original callables coming back on components).Also applies to: 599-635, 640-717
docs/source/reference/middleware.md (1)
111-172: Middleware signature examples are up to date with the variadic API.The LoggingMiddleware and ValidationMiddleware examples, along with the unit‑testing snippet, now:
- Use
*args, **kwargsand forward them throughcall_next.- Demonstrate how to pull the primary payload from
args[0]while preserving the rest.- Match the actual
FunctionMiddlewareprotocol used by DynamicFunctionMiddleware.These examples will help users avoid subtle signature mismatches when writing custom middleware.
Also applies to: 345-359, 529-538
| @pytest.fixture | ||
| def mock_builder(): | ||
| """Create a mock builder with all required methods.""" | ||
| builder = Mock() | ||
| builder._functions = {} | ||
| builder._function_policies = {} | ||
| builder.get_llm = AsyncMock() | ||
| builder.get_embedder = AsyncMock() | ||
| builder.get_retriever = AsyncMock() | ||
| builder.get_memory_client = AsyncMock() | ||
| builder.get_object_store_client = AsyncMock() | ||
| builder.get_auth_provider = AsyncMock() | ||
| builder.get_function = AsyncMock() | ||
| builder.get_function_config = Mock() | ||
| return builder |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, search for pytest fixtures to understand naming patterns used
rg -n "@pytest\.fixture" --type py | head -50Repository: NVIDIA/NeMo-Agent-Toolkit
Length of output: 194
🏁 Script executed:
#!/bin/bash
# Check the pyproject.toml for any pre-commit or linting configs that might enforce naming
find . -name "pyproject.toml" -o -name ".pre-commit-config.yaml" -o -name "ruff.toml" | head -5Repository: NVIDIA/NeMo-Agent-Toolkit
Length of output: 194
🏁 Script executed:
#!/bin/bash
# Look at the coding guidelines file mentioned in learnings
cat .cursor/rules/general.mdc 2>/dev/null | grep -A 5 "fixture_"Repository: NVIDIA/NeMo-Agent-Toolkit
Length of output: 194
Rename fixtures to follow the project's documented naming convention with fixture_ prefix and explicit name= argument.
The project guidelines require fixtures to be named with a fixture_ prefix and to define the name argument in the @pytest.fixture decorator. Update mock_builder (and the fixtures at lines 58-68, 160-173) to follow this convention:
@pytest.fixture(name="mock_builder")
def fixture_mock_builder():
"""Create a mock builder with all required methods."""
# ...🤖 Prompt for AI Agents
In tests/nat/middleware/test_dynamic_middleware.py around lines 41-55 (and also
update the fixtures at lines 58-68 and 160-173), rename the pytest fixtures to
use the documented naming convention by adding the fixture_ prefix to the
function name and specifying the fixture name via the decorator; e.g., change
the decorator to @pytest.fixture(name="mock_builder") and rename the function to
fixture_mock_builder, keeping the existing docstring and body unchanged, and
apply the same pattern to the other two fixtures mentioned.
3 participants
Description
This PR introduces Function Policies, a component-based framework for intercepting and modifying function execution in the NeMo Agent Toolkit. Function policies are first-class components that wrap function calls with pre-invoke and post-invoke hooks, enabling features like input validation, output transformation, audit logging, data sanitization, and monitoring.
Function policies remove code duplication and tightly coupled logic by making these concerns configurable, reusable, and chainable.
Key Features
Policy Framework:
enabledconfig flagDynamic Middleware:
Example Configuration
Step 1: Define Configuration
Step 2: Implement Policy
Step 3: Register
CLI Commands
Component Discovery
Dynamic middleware automatically discovers and wraps component methods based on allowlists.
Allowlisted Functions by Component
LLMs:
invoke, ainvoke, stream, astream
Embedders:
embed_query, aembed_query
Retrievers:
search
Memory:
search, add_items, remove_items
Object Stores:
put_object, get_object, delete_object, upsert_object
Auth Providers:
authenticate
Users may extend allowlists in configuration to support additional component methods.
Summary by CodeRabbit
New Features
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.