# Created by aaronkueh on 9/19/2025 # aom/agents/maintenance.py from __future__ import annotations from datetime import datetime from typing import Dict, Any, Optional import pandas as pd from jinja2 import Template from aom.utils.logger import get_logger from aom.utils.agent_helpers import load_agent_prompts from aom.mcp_server.server import call_tool from aom.utils.contracts import MAINTRequest from aom.utils.utilities import compact import re def _parse_date_filter(date_str: Optional[str]) -> Optional[datetime]: """Parse date string to datetime object""" if not date_str: return None try: return pd.to_datetime(date_str) except: return None def _filter_by_date_range(df: pd.DataFrame, start_date: Optional[datetime], end_date: Optional[datetime], date_col: str = 'create_date') -> pd.DataFrame: """Filter dataframe by date range""" if start_date: df = df[df[date_col] >= start_date] if end_date: df = df[df[date_col] <= end_date] return df def _filter_by_description_keywords(df: pd.DataFrame, keywords: list) -> pd.DataFrame: """ Filter dataframe by description keywords. Searches for ANY keyword appearing in the description column. Args: df: DataFrame with 'description' column keywords: List of keywords to search for Returns: Filtered DataFrame containing rows where ANY keyword appears in description """ if not keywords or df.empty or 'description' not in df.columns: return df # Create a mask that matches ANY keyword mask = pd.Series([False] * len(df), index=df.index) for keyword in keywords: # Case-insensitive search for each keyword keyword_mask = df['description'].str.contains(keyword, case=False, na=False, regex=False) mask |= keyword_mask # OR operation: match if ANY keyword found return df[mask] class MAINTAgent: """ Maintenance Agent (MAINT): - Fetch maintenance work order data from asset_wo_pub table - Filter by asset_id, date range, maintenance_by, status, and description - Present results in a tabular format with natural language summary """ def __init__(self): self.log = get_logger(self.__class__.__name__) self.agent_profile = load_agent_prompts(agent='MAINT', config='agents.yaml') def run(self, req: MAINTRequest) -> Dict[str, Any]: """ Execute maintenance query based on request parameters Parameters in req: - asset_id: Filter by specific asset - start_date: Filter work orders created/scheduled after this date - end_date: Filter work orders created/scheduled before this date - maintenance_by: Filter by technician/person performing maintenance - status: Filter by work order status (COMPLETED, SCHED, INPROG, RESCHED) - description_keywords: List of keywords to search in description - db_profile: Database profile to use """ self.log.info("MAINT Agent activated") # 1) Fetch all maintenance work order data maint_params = { "profile": req.db_profile, } maint_res = call_tool("mssql_maint_query", maint_params) maint_df = pd.DataFrame(maint_res["rows"], columns=maint_res["columns"]) if maint_res.get( "rows") else pd.DataFrame() if maint_df.empty: return { "summary": "No maintenance work order data available in the system.", "table": None, } # Convert date columns to datetime date_columns = ['create_date', 'maintenance_date'] for col in date_columns: if col in maint_df.columns: maint_df[col] = pd.to_datetime(maint_df[col], errors='coerce') # 2) Apply filters based on request parameters filtered_df = maint_df.copy() filter_info = [] # Filter by asset_id if req.asset_id: filtered_df = filtered_df[filtered_df['asset_id'] == req.asset_id] filter_info.append(f"asset: {req.asset_id}") # Filter by date range (on create_date or maintenance_date) if req.start_date or req.end_date: start_dt = _parse_date_filter(req.start_date) if req.start_date else None end_dt = _parse_date_filter(req.end_date) if req.end_date else None # Use create_date by default (fallback to maintenance_date) date_col = 'create_date' if 'create_date' in filtered_df.columns else ( 'maintenance_date' if 'maintenance_date' in filtered_df.columns else None ) if date_col: filtered_df = _filter_by_date_range(filtered_df, start_dt, end_dt, date_col) if start_dt and end_dt: filter_info.append(f"date range: {start_dt.date()} to {end_dt.date()}") elif start_dt: filter_info.append(f"after: {start_dt.date()}") elif end_dt: filter_info.append(f"before: {end_dt.date()}") # Filter by work order number (WO) if req.wo_no is not None: if "wo_no" in filtered_df.columns: filtered_df = filtered_df[pd.to_numeric(filtered_df["wo_no"], errors="coerce") == int(req.wo_no)] filter_info.append(f"wo_no: {int(req.wo_no)}") # Filter by maintenance_by if req.maintenance_by: filtered_df = filtered_df[filtered_df['maintenance_by'].str.contains( req.maintenance_by, case=False, na=False)] filter_info.append(f"technician: {req.maintenance_by}") # Filter by status if req.status: filtered_df = filtered_df[filtered_df['status'].str.upper() == req.status.upper()] filter_info.append(f"status: {req.status}") print('status filter filter\n', filtered_df) if filtered_df.empty: filter_str = "; ".join(filter_info) if filter_info else "specified criteria" return { "summary": f"No maintenance work orders found matching the {filter_str}.", "table": None, } # 4) Sort by create_date or maintenance_date sort_col = 'create_date' if 'create_date' in filtered_df.columns else 'maintenance_date' filtered_df = filtered_df.sort_values(by=sort_col, ascending=False) # 5) Format the results for display display_df = filtered_df.copy() # Format datetime columns for better display for col in date_columns: if col in display_df.columns: display_df[col] = display_df[col].dt.strftime('%Y-%m-%d %H:%M') if display_df.empty: summary_text = "No maintenance work orders found matching the specified criteria." else: summary_text = 'table' # 6) Format table data for return table_data = display_df.to_dict(orient='records') return { "summary": summary_text, "table": table_data }