Track agent system prompt in config and add config to skip maintenance check #724
Conversation
…heck; add unit test; fix unit tests Signed-off-by: Hsin Chen <[email protected]>
Signed-off-by: Hsin Chen <[email protected]>
WalkthroughAdds a boolean skip_maintenance_check flag to the maintenance tool with an early-return in the async path, exposes agent_prompt in the workflow config for the system message, consolidates prompt/message imports, and updates tests and test data path resolutions to cover the new flag and paths. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Caller
participant Workflow
participant MaintenanceCheckTool
participant DataStore as MaintenanceData
Caller->>Workflow: Run alert triage
Workflow->>MaintenanceCheckTool: arun(config, inputs)
alt config.skip_maintenance_check == True
Note over MaintenanceCheckTool: Early return (skip)
MaintenanceCheckTool-->>Workflow: NO_ONGOING_MAINTENANCE_STR
else config.skip_maintenance_check == False
MaintenanceCheckTool->>DataStore: Load maintenance data
DataStore-->>MaintenanceCheckTool: Data
MaintenanceCheckTool->>MaintenanceCheckTool: Parse & evaluate
MaintenanceCheckTool-->>Workflow: Result string
end
Workflow-->>Caller: Triage outcome
sequenceDiagram
autonumber
actor Caller
participant Workflow
participant LLM as Assistant
Caller->>Workflow: Start workflow
Note over Workflow: Build system message from config.agent_prompt
Workflow->>LLM: System + user messages
LLM-->>Workflow: Response
Workflow-->>Caller: Output
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested labels
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
hsin-c
added
improvement
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (7)
examples/advanced_agents/alert_triage_agent/tests/test_utils.py (1)
103-107: Make data path resolution robust (avoid hard-coded ../../../ levels).Using "../../../../.." is brittle across packaging/working-directory contexts. Prefer resolving relative paths from CWD or anchoring to repo root deterministically.
Apply this minimal change if tests are always run from the repo root:
- offline_data_path_abs = importlib.resources.files(package_name).joinpath("../../../../../", - offline_data_path).absolute() - benign_fallback_data_path_abs = importlib.resources.files(package_name).joinpath( - "../../../../../", benign_fallback_data_path).absolute() + offline_data_path_abs = Path(offline_data_path).resolve() + benign_fallback_data_path_abs = Path(benign_fallback_data_path).resolve()If not guaranteed, consider deriving from the config file location (e.g., Path(config_file).parents[n] / offline_data_path) or a helper that locates the repo root (presence of pyproject.toml).
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py (3)
21-23: Top-level imports: OK; prefer ChatPromptTemplate.from_messages for clarity.Moving these imports up is fine. When composing prompts, the canonical API is ChatPromptTemplate.from_messages([...]) for forward compatibility.
210-212: Early return when skipping: LGTM; consider slightly richer log.Behavior matches the new flag. Optional: include a short reason in the log (e.g., “Skipping maintenance check (config flag)”) to ease trace filtering.
174-195: Avoid variable shadowing and use ChatPromptTemplate.from_messages.Current code reuses the name prompt_template both as a str and then as a ChatPromptTemplate, which hurts readability. Also, prefer the from_messages constructor.
Apply:
-def _summarize_alert(llm, prompt_template, alert, maintenance_start_str, maintenance_end_str): +def _summarize_alert(llm, prompt: str, alert: dict, maintenance_start_str: str, maintenance_end_str: str) -> str: @@ - sys_prompt = prompt_template.format(maintenance_start_str=maintenance_start_str, - maintenance_end_str=maintenance_end_str) - prompt_template = ChatPromptTemplate([("system", sys_prompt), MessagesPlaceholder("msgs")]) - summarization_chain = prompt_template | llm - alert_json_str = json.dumps(alert) - result = summarization_chain.invoke({"msgs": [HumanMessage(content=alert_json_str)]}).content - return result + sys_prompt = prompt.format(maintenance_start_str=maintenance_start_str, + maintenance_end_str=maintenance_end_str) + prompt_tmpl = ChatPromptTemplate.from_messages([("system", sys_prompt), MessagesPlaceholder("msgs")]) + summarization_chain = prompt_tmpl | llm + alert_json_str = json.dumps(alert, ensure_ascii=False) + return summarization_chain.invoke({"msgs": [HumanMessage(content=alert_json_str)]}).contentAnd update the call site:
- report = _summarize_alert(llm=llm, - prompt_template=config.prompt, + report = _summarize_alert(llm=llm, + prompt=config.prompt, alert=alert, maintenance_start_str=maintenance_start_str, maintenance_end_str=maintenance_end_str)Optional (outside this hunk): add from typing import Any and annotate llm if you want stricter typing.
examples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.py (2)
50-51: Harden maintenance data path resolution.Same concern as other tests: "../../../../.." is fragile.
- maintenance_data_path_abs = importlib.resources.files(package_name).joinpath("../../../../../", - maintenance_data_path).absolute() + maintenance_data_path_abs = Path(maintenance_data_path).resolve()If the test runner might not start at repo root, anchor relative to config_file or implement a small helper to discover project root.
254-255: Be careful mutating config after tool registration.Depending on WorkflowBuilder/get_tool internals, updating config.skip_maintenance_check post-registration may not propagate to the running tool (copy vs reference). Safer to instantiate a fresh config (or update the tool via a supported API) per case.
Example approach (conceptual):
- Move MaintenanceCheckToolConfig creation into the loop and re-add/replace the function, or
- If supported, call a builder.update_tool_config("maintenance_check", new_config) before invoking.
examples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.py (1)
42-42: Stabilize input file path resolution.The extra “../” level is brittle. Prefer resolving from CWD (if guaranteed) or anchoring to the config file’s directory.
- input_filepath_abs = importlib.resources.files(package_name).joinpath("../../../../../", input_filepath).absolute() + input_filepath_abs = Path(input_filepath).resolve()If CWD isn’t stable, compute relative to config_file (e.g., Path(config_file).parents[n] / input_filepath) or add a helper.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py(3 hunks)examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py(2 hunks)examples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.py(1 hunks)examples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.py(3 hunks)examples/advanced_agents/alert_triage_agent/tests/test_utils.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (7)
{src/**/*.py,packages/*/src/**/*.py,examples/*/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' only for env var prefixes and informal code comments (never in documentation)
Files:
examples/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Format Python code with yapf (PEP8 base, column_limit=120) run second in the toolchain
Run ruff (ruff check --fix) using configuration in pyproject.toml; fix warnings unless explicitly ignored; ruff is a linter only (not for formatting)
Respect naming: snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
Prefer typing/collections.abc abstractions (e.g., Sequence over list)
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Prefer httpx with SSL verification enabled by default for HTTP requests; follow OWASP Top-10
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 vectorized operations when beneficial
Files:
examples/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py
**/*.{py,md,sh,yml,yaml,toml,json,ini,cfg,txt,rst,pyi}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces (no tabs) and ensure every file ends with a single trailing newline
Files:
examples/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py
**/*.{py,sh,md,yml,yaml,toml,json,ini,txt,cfg,rst,pyi}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.{py,sh,md,yml,yaml,toml,json,ini,txt,cfg,rst,pyi}: Every file must start with the standard SPDX Apache-2.0 header
All source files must include the SPDX Apache-2.0 header template
Files:
examples/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py
**/*.{py,md}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived by setuptools-scm; never hard-code versions in code or documentation
Files:
examples/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.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: passDocumentation 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 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/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.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 apyproject.tomlfile. Optionally, it might also contain scripts in ascripts/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 nameddata/, and should
be checked into git-lfs.
Files:
examples/advanced_agents/alert_triage_agent/tests/test_utils.pyexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.pyexamples/advanced_agents/alert_triage_agent/tests/test_alert_triage_agent_workflow.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py
🔇 Additional comments (4)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py (1)
48-52: New config flag skip_maintenance_check: LGTM.Good default and clear description. This enables faster offline tests without maintenance data.
examples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.py (1)
218-223: Skip flag test case: LGTM.Covers the new branch and asserts _summarize_alert isn’t called. Good addition.
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (2)
63-64: Configurable agent system prompt: LGTM.Good defaulting to ALERT_TRIAGE_AGENT_PROMPT and exposing via config improves reproducibility.
96-96: Using config.agent_prompt in SystemMessage: LGTM.Correctly replaces the previous hard-coded prompt.
Signed-off-by: Hsin Chen <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py (2)
21-23: Simplify prompt wiring; drop HumanMessage/MessagesPlaceholder to reduce deps and boilerplateSwitch to a simple human-slot in the prompt and pass the JSON via a variable. This removes two imports and keeps LCEL usage straightforward.
Apply:
-import pandas as pd -from langchain_core.messages import HumanMessage -from langchain_core.prompts import ChatPromptTemplate -from langchain_core.prompts import MessagesPlaceholder +import pandas as pd +from typing import Any +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.runnables import Runnable- prompt = ChatPromptTemplate.from_messages([("system", sys_prompt), MessagesPlaceholder("msgs")]) - summarization_chain = prompt | llm - alert_json_str = json.dumps(alert) - result = summarization_chain.invoke({"msgs": [HumanMessage(content=alert_json_str)]}).content + prompt = ChatPromptTemplate.from_messages([("system", sys_prompt), ("human", "{alert_json_str}")]) + summarization_chain = prompt | llm + alert_json_str = json.dumps(alert) + result = summarization_chain.invoke({"alert_json_str": alert_json_str}).contentAlso applies to: 191-195
174-175: Add type hints to satisfy repo policy and pyrightAnnotate inputs/return, using Runnable for the LCEL-compatible llm and Any for the alert payload.
Apply:
-def _summarize_alert(llm, prompt_template, alert, maintenance_start_str, maintenance_end_str): +def _summarize_alert( + llm: Runnable, + prompt_template: str, + alert: dict[str, Any], + maintenance_start_str: str, + maintenance_end_str: str, +) -> str:
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (7)
{src/**/*.py,packages/*/src/**/*.py,examples/*/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Use abbreviations: 'nat' for API namespace and CLI tool, 'nvidia-nat' for package name, 'NAT' only for env var prefixes and informal code comments (never in documentation)
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Format Python code with yapf (PEP8 base, column_limit=120) run second in the toolchain
Run ruff (ruff check --fix) using configuration in pyproject.toml; fix warnings unless explicitly ignored; ruff is a linter only (not for formatting)
Respect naming: snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
Prefer typing/collections.abc abstractions (e.g., Sequence over list)
Use typing.Annotated for units or extra metadata when useful
Treat pyright warnings (configured in pyproject.toml) as errors during development
Prefer httpx with SSL verification enabled by default for HTTP requests; follow OWASP Top-10
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 vectorized operations when beneficial
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py
**/*.{py,md,sh,yml,yaml,toml,json,ini,cfg,txt,rst,pyi}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Indent with 4 spaces (no tabs) and ensure every file ends with a single trailing newline
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py
**/*.{py,sh,md,yml,yaml,toml,json,ini,txt,cfg,rst,pyi}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.{py,sh,md,yml,yaml,toml,json,ini,txt,cfg,rst,pyi}: Every file must start with the standard SPDX Apache-2.0 header
All source files must include the SPDX Apache-2.0 header template
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py
**/*.{py,md}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Version numbers are derived by setuptools-scm; never hard-code versions in code or documentation
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.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: passDocumentation 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 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/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.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 apyproject.tomlfile. Optionally, it might also contain scripts in ascripts/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 nameddata/, and should
be checked into git-lfs.
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py
🔇 Additional comments (2)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/maintenance_check.py (2)
48-52: LGTM: Config flag is clear and documentedDefaults, description, and behavior are sensible and align with the PR goals.
200-203: Defer LLM construction until it’s actually neededTests in
test_maintenance_check.pyare correctly asserting that, when there’s no active maintenance,_arunreturnsNO_ONGOING_MAINTENANCE_STRand never calls the summarizer (see lines 268–272). However, because you currently call# Set up LLM llm = await builder.get_llm(…, wrapper_type=…)outside of the “summarize-only-if-needed” branch, you still pay the cost of initializing the LLM even on every early‐exit/skip path.
Move the LLM instantiation from the top of the class/method into the
_arunfunction immediately before you call_summarize_alert. That way:
- builder.get_llm is only invoked when summarization is actually needed
- the skip/fast path returns immediately without any network/setup overhead
Suggested diff for the first occurrence (apply similarly at lines 210–212 and 260–265):
- # Set up LLM - llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN) - maintenance_data_path = config.static_data_path # …later in _arun, just before summarizing… - report = _summarize_alert(llm=llm, + llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + report = _summarize_alert(llm=llm, prompt_template=config.prompt, alert=alert, maintenance_start_str=maintenance_start_str, maintenance_end_str=maintenance_end_str)Please verify that no other code paths depend on
self.llmhaving been set before_arunis invoked, and confirm all test cases still pass after making this change.
|
/merge |
2 participants
Changes:
All unit tests passed.
Summary by CodeRabbit
New Features
Tests