fix: ensure registration of adk demo functions; reduce warnings

[willkill07]

Description

Updates the ADK Demo example to work.

• add missing registration imports
• exactly specify framework_wrappers

Closes nvbugs-5554110

By Submitting this PR I confirm:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

Summary by CodeRabbit

• New Features

  • Enabled ADK framework support in the demo: City Time and Weather Update tools are now registered for ADK, making them discoverable and usable in ADK-based workflows and integrations.
  • Improves out-of-the-box experience for ADK users running the examples by ensuring these tools are available without additional configuration.

• Chores

  • Updated tool registrations to target the ADK framework while keeping existing function behavior and signatures unchanged.

[coderabbitai]

Walkthrough

Adds LLMFrameworkEnum imports and extends @register_function decorators in two ADK demo tools to include framework_wrappers=[LLMFrameworkEnum.ADK], scoping their registration to the ADK framework. No function bodies, signatures, or return types changed.

Changes

Cohort / File(s)	Summary of Changes
ADK framework scoping for tool registration examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py, examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py	Import LLMFrameworkEnum and update @register_function to pass framework_wrappers=[LLMFrameworkEnum.ADK]; no logic or signature changes.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant Module as ADK Demo Module (on import)
  participant Decorator as register_function(...)
  participant Registry as ToolRegistry
  participant Enum as LLMFrameworkEnum

  Note over Module: Module import
  Module->>Decorator: Apply @register_function(..., framework_wrappers=[Enum.ADK])
  Decorator->>Enum: Resolve ADK
  Decorator->>Registry: register(tool_fn, config_type, frameworks=[ADK])
  Note right of Registry: Tool is scoped to ADK framework

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested labels

bug, non-breaking

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py (1)

31-32: Consider adding a docstring for the outer function.

The weather_update function (the registration generator) lacks its own docstring. While the inner _weather_update function is well-documented, adding a docstring similar to the one in nat_time_mcp_tool.py (lines 36-42) would improve consistency and clarity.

Apply this diff to add the missing docstring:

 @register_function(config_type=WeatherToolConfig, framework_wrappers=[LLMFrameworkEnum.ADK])
 async def weather_update(_config: WeatherToolConfig, _builder: Builder) -> AsyncIterator[FunctionInfo]:
+    """
+    Register a weather_update(city: str) -> str tool for ADK.
+
+    Args:
+        _config (WeatherToolConfig): The configuration for the weather_update tool.
+        _builder (Builder): The NAT builder instance.
+    """
 
     async def _weather_update(city: str) -> str:

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b5541b5 and 0a351d5.

📒 Files selected for processing (2)

• examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py (2 hunks)
• examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py (2 hunks)

🧰 Additional context used

📓 Path-based instructions (4)

**/*.{py,yaml,yml}

📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)

**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.

Files:

• examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py
• examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py

**/*.py

📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)

**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).

**/*.py: In code comments/identifiers use NAT abbreviations as specified: nat for API namespace/CLI, nvidia-nat for package name, NAT for env var prefixes; do not use these abbreviations in documentation
Follow PEP 20 and PEP 8; run yapf with column_limit=120; use 4-space indentation; end files with a single trailing newline
Run ruff check --fix as linter (not formatter) using pyproject.toml config; fix warnings unless explicitly ignored
Respect naming: snake_case for functions/variables, PascalCase for classes, UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: use bare raise to re-raise; log with logger.error() when re-raising to avoid duplicate stack traces; use logger.exception() when catching without re-raising
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise and ending with a period; surround code entities with backticks
Validate and sanitize 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; profile CPU-heavy paths with cProfile or mprof before optimizing; cache expensive computations with functools.lru_cache or external cache; leverage NumPy vectorized operations when beneficial

Files:

• examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py
• examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.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.
  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 raise statements to maintain the original stack trace,
    and use logger.error() (not logger.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.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file 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:

• examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py
• examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py

examples/**/*

⚙️ CodeRabbit configuration file

examples/**/*: - This directory contains example code and usage scenarios for the toolkit, at a minimum an example should
contain a README.md or file README.ipynb.

• If an example contains Python code, it should be placed in a subdirectory named src/ and should
  contain a pyproject.toml file. Optionally, it might also contain scripts in a scripts/ directory.
• If an example contains YAML files, they should be placed in a subdirectory named configs/. - If an example contains sample data files, they should be placed in a subdirectory named data/, and should
  be checked into git-lfs.

Files:

• examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py
• examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py

🧬 Code graph analysis (2)

examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py (1)

src/nat/builder/framework_enum.py (1)

• LLMFrameworkEnum (19-25)

examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py (1)

src/nat/builder/framework_enum.py (1)

• LLMFrameworkEnum (19-25)

🔇 Additional comments (4)

examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py (2)

22-22: LGTM: Import added correctly.

The LLMFrameworkEnum import is properly added to support the framework-specific registration.

---

34-34: Approve ADK-specific function registrations
Verified that all functions in examples/frameworks/adk_demo/src/nat_adk_demo/ are registered with framework_wrappers=[LLMFrameworkEnum.ADK].

examples/frameworks/adk_demo/src/nat_adk_demo/weather_update_tool.py (2)

20-20: LGTM: Import added correctly.

The LLMFrameworkEnum import is properly added to support framework-specific registration.

---

30-30: LGTM: Function registration now explicitly scoped to ADK framework.

The framework_wrappers=[LLMFrameworkEnum.ADK] parameter ensures this tool is only registered for the ADK framework, which aligns with the PR objective to fix ADK demo registration.

[willkill07]

/merge