Refactor for better prompt and tool description organization

examples/alert_triage_agent/src/aiq_alert_triage_agent/categorizer.py

@@ -28,13 +28,13 @@
 from aiq.data_models.function import FunctionBaseConfig
 
 from . import utils
-from .prompts import PipelineNodePrompts
+from .prompts import CategorizerPrompts
 
 
 class CategorizerToolConfig(FunctionBaseConfig, name="categorizer"):
-    description: str = Field(default="This is a categorization tool used at the end of the pipeline.",
-                             description="Description of the tool.")
+    description: str = Field(default=CategorizerPrompts.TOOL_DESCRIPTION, description="Description of the tool.")
     llm_name: LLMRef
+    prompt: str = Field(default=CategorizerPrompts.PROMPT, description="Main prompt for the categorization task.")
 
 
 def _extract_markdown_heading_level(report: str) -> str:
@@ -48,8 +48,7 @@ def _extract_markdown_heading_level(report: str) -> str:
 async def categorizer_tool(config: CategorizerToolConfig, builder: Builder):
     # Set up LLM and chain
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-    prompt_template = ChatPromptTemplate([("system", PipelineNodePrompts.CATEGORIZER_PROMPT),
-                                          MessagesPlaceholder("msgs")])
+    prompt_template = ChatPromptTemplate([("system", config.prompt), MessagesPlaceholder("msgs")])
     categorization_chain = prompt_template | llm
 
     async def _arun(report: str) -> str:

examples/alert_triage_agent/src/aiq_alert_triage_agent/hardware_check_tool.py

@@ -24,15 +24,13 @@
 from aiq.data_models.function import FunctionBaseConfig
 
 from . import utils
-from .prompts import ToolReasoningLayerPrompts
+from .prompts import HardwareCheckPrompts
 
 
 class HardwareCheckToolConfig(FunctionBaseConfig, name="hardware_check"):
-    description: str = Field(
-        default=("This tool checks hardware health status using IPMI monitoring to detect power state, "
-                 "hardware degradation, and anomalies that could explain alerts. Args: host_id: str"),
-        description="Description of the tool for the agent.")
+    description: str = Field(default=HardwareCheckPrompts.TOOL_DESCRIPTION, description="Description of the tool.")
     llm_name: LLMRef
+    prompt: str = Field(default=HardwareCheckPrompts.PROMPT, description="Main prompt for the hardware check task.")
     offline_mode: bool = Field(default=True, description="Whether to run in offline model")
 
 
@@ -94,7 +92,7 @@ async def _arun(host_id: str) -> str:
                 # Additional LLM reasoning layer on playbook output to provide a summary of the results
                 utils.log_header("LLM Reasoning", dash_length=50)
 
-                prompt = ToolReasoningLayerPrompts.HARDWARE_CHECK.format(input_data=monitoring_data)
+                prompt = config.prompt.format(input_data=monitoring_data)
 
                 # Get analysis from LLM
                 conclusion = await utils.llm_ainvoke(config, builder, prompt)

examples/alert_triage_agent/src/aiq_alert_triage_agent/host_performance_check_tool.py

@@ -23,15 +23,17 @@
 
 from . import utils
 from .playbooks import HOST_PERFORMANCE_CHECK_PLAYBOOK
-from .prompts import ToolReasoningLayerPrompts
+from .prompts import HostPerformanceCheckPrompts
 
 
 class HostPerformanceCheckToolConfig(FunctionBaseConfig, name="host_performance_check"):
-    description: str = Field(
-        default=("This is the Host Performance Check Tool. This tool retrieves CPU usage, memory usage, "
-                 "and hardware I/O usage details for a given host. Args: host_id: str"),
-        description="Description of the tool for the agent.")
+    description: str = Field(default=HostPerformanceCheckPrompts.TOOL_DESCRIPTION,
+                             description="Description of the tool.")
     llm_name: LLMRef
+    parsing_prompt: str = Field(default=HostPerformanceCheckPrompts.PARSING_PROMPT,
+                                description="Prompt for parsing the raw host performance data.")
+    analysis_prompt: str = Field(default=HostPerformanceCheckPrompts.ANALYSIS_PROMPT,
+                                 description="Prompt for analyzing the parsed host performance data.")
     offline_mode: bool = Field(default=True, description="Whether to run in offline model")
 
 
@@ -97,7 +99,7 @@ async def _parse_stdout_lines(config, builder, stdout_lines):
         # Join the list of lines into a single text block
         input_data = "\n".join(stdout_lines) if stdout_lines else ""
 
-        prompt = ToolReasoningLayerPrompts.HOST_PERFORMANCE_CHECK_PARSING.format(input_data=input_data)
+        prompt = config.parsing_prompt.format(input_data=input_data)
 
         response = await utils.llm_ainvoke(config=config, builder=builder, user_prompt=prompt)
     except Exception as e:
@@ -146,7 +148,7 @@ async def _arun(host_id: str) -> str:
             # Additional LLM reasoning layer on playbook output to provide a summary of the results
             utils.log_header("LLM Reasoning", dash_length=50)
 
-            prompt_template = ToolReasoningLayerPrompts.HOST_PERFORMANCE_CHECK_ANALYSIS.format(input_data=output)
+            prompt_template = config.analysis_prompt.format(input_data=output)
 
             conclusion = await utils.llm_ainvoke(config, builder, user_prompt=prompt_template)
 

examples/alert_triage_agent/src/aiq_alert_triage_agent/maintenance_check.py

@@ -31,17 +31,16 @@
 from aiq.data_models.function import FunctionBaseConfig
 
 from . import utils
-from .prompts import PipelineNodePrompts
+from .prompts import MaintenanceCheckPrompts
 
 NO_ONGOING_MAINTENANCE_STR = "No ongoing maintenance found for the host."
 
 
 class MaintenanceCheckToolConfig(FunctionBaseConfig, name="maintenance_check"):
-    description: str = Field(
-        default=("Check if a host is under maintenance during the time of an alert to help determine "
-                 "if the alert can be deprioritized."),
-        description="Description of the tool for the agent.")
+    description: str = Field(default=MaintenanceCheckPrompts.TOOL_DESCRIPTION, description="Description of the tool.")
     llm_name: LLMRef
+    prompt: str = Field(default=MaintenanceCheckPrompts.PROMPT,
+                        description="Main prompt for the maintenance check task.")
     static_data_path: str | None = Field(
         default="examples/alert_triage_agent/data/maintenance_static_dataset.csv",
         description=(
@@ -167,12 +166,13 @@ def _get_active_maintenance(df: pd.DataFrame, host_id: str, alert_time: datetime
     return start_time_str, end_time_str
 
 
-def _summarize_alert(llm, alert, maintenance_start_str, maintenance_end_str):
+def _summarize_alert(llm, prompt_template, alert, maintenance_start_str, maintenance_end_str):
     """
     Generate a summary report for an alert when the affected host is under maintenance.
 
     Args:
         llm: The language model to use for generating the summary
+        prompt_template: The prompt template to use for generating the summary
         alert (dict): Dictionary containing the alert details
         maintenance_start_str (str): Start time of maintenance window in "YYYY-MM-DD HH:MM:SS" format
         maintenance_end_str (str): End time of maintenance window in "YYYY-MM-DD HH:MM:SS" format,
@@ -181,8 +181,8 @@ def _summarize_alert(llm, alert, maintenance_start_str, maintenance_end_str):
     Returns:
         str: A markdown-formatted report summarizing the alert and maintenance status
     """
-    sys_prompt = PipelineNodePrompts.MAINTENANCE_CHECK_PROMPT.format(maintenance_start_str=maintenance_start_str,
-                                                                     maintenance_end_str=maintenance_end_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)
@@ -249,7 +249,11 @@ async def _arun(input_message: str) -> str:
         # maintenance info found, summarize alert and return a report (agent execution will be skipped)
         utils.logger.info("Host: [%s] is under maintenance according to the maintenance database", host)
 
-        report = _summarize_alert(llm, alert, maintenance_start_str, maintenance_end_str)
+        report = _summarize_alert(llm=llm,
+                                  prompt_template=config.prompt,
+                                  alert=alert,
+                                  maintenance_start_str=maintenance_start_str,
+                                  maintenance_end_str=maintenance_end_str)
 
         utils.log_footer()
         return report

examples/alert_triage_agent/src/aiq_alert_triage_agent/monitoring_process_check_tool.py

@@ -23,14 +23,15 @@
 
 from . import utils
 from .playbooks import MONITOR_PROCESS_CHECK_PLAYBOOK
-from .prompts import ToolReasoningLayerPrompts
+from .prompts import MonitoringProcessCheckPrompts
 
 
 class MonitoringProcessCheckToolConfig(FunctionBaseConfig, name="monitoring_process_check"):
-    description: str = Field(default=("This tool checks the status of critical monitoring processes and services "
-                                      "on a target host by executing system commands. Args: host_id: str"),
-                             description="Description of the tool for the agent.")
+    description: str = Field(default=MonitoringProcessCheckPrompts.TOOL_DESCRIPTION,
+                             description="Description of the tool.")
     llm_name: LLMRef
+    prompt: str = Field(default=MonitoringProcessCheckPrompts.PROMPT,
+                        description="Main prompt for the monitoring process check task.")
     offline_mode: bool = Field(default=True, description="Whether to run in offline model")
 
 
@@ -104,7 +105,7 @@ async def _arun(host_id: str) -> str:
             # Additional LLM reasoning layer on playbook output to provide a summary of the results
             utils.log_header("LLM Reasoning", dash_length=50)
 
-            prompt = ToolReasoningLayerPrompts.MONITORING_PROCESS_CHECK.format(input_data=output_for_prompt)
+            prompt = config.prompt.format(input_data=output_for_prompt)
 
             conclusion = await utils.llm_ainvoke(config, builder, prompt)
 

examples/alert_triage_agent/src/aiq_alert_triage_agent/network_connectivity_check_tool.py

@@ -25,15 +25,15 @@
 from aiq.data_models.function import FunctionBaseConfig
 
 from . import utils
-from .prompts import ToolReasoningLayerPrompts
+from .prompts import NetworkConnectivityCheckPrompts
 
 
 class NetworkConnectivityCheckToolConfig(FunctionBaseConfig, name="network_connectivity_check"):
-    description: str = Field(
-        default=("This tool checks network connectivity of a host by running ping and socket connection tests. "
-                 "Args: host_id: str"),
-        description="Description of the tool for the agent.")
+    description: str = Field(default=NetworkConnectivityCheckPrompts.TOOL_DESCRIPTION,
+                             description="Description of the tool.")
     llm_name: LLMRef
+    prompt: str = Field(default=NetworkConnectivityCheckPrompts.PROMPT,
+                        description="Main prompt for the network connectivity check task.")
     offline_mode: bool = Field(default=True, description="Whether to run in offline model")
 
 
@@ -106,8 +106,7 @@ async def _arun(host_id: str) -> str:
             # Additional LLM reasoning layer on playbook output to provide a summary of the results
             utils.log_header("LLM Reasoning", dash_length=50)
 
-            prompt = ToolReasoningLayerPrompts.NETWORK_CONNECTIVITY_CHECK.format(ping_data=ping_data,
-                                                                                 telnet_data=telnet_data)
+            prompt = config.prompt.format(ping_data=ping_data, telnet_data=telnet_data)
             conclusion = await utils.llm_ainvoke(config, builder, prompt)
 
             utils.logger.debug(conclusion)