Refactor: Extract helper functions

agent/agent.py

@@ -1,4 +1,5 @@
 import json
+import re
 from langchain_google_genai import GoogleGenerativeAI
 
 class MaintenanceAgent:
@@ -9,8 +10,9 @@ class MaintenanceAgent:
             google_api_key=api_key
         )
 
-    def run(self, phase2_output: dict) -> dict:
-        prompt = f"""
+    def _build_prompt(self, phase2_output: dict) -> str:
+        """Build the maintenance analysis prompt."""
+        return f"""
 You are a maintenance decision AI.
 You must reason ONLY from the provided JSON.
 Do NOT invent data.
@@ -29,12 +31,22 @@ OUTPUT FORMAT:
 }}
 """
 
-        response = self.llm.invoke(prompt)
+    def _parse_response(self, response: str) -> dict:
+        """Parse LLM response, handling various JSON formats."""
         try:
             return json.loads(response)
         except json.JSONDecodeError:
-            import re
-            match = re.search(r'```json\s*(.*?)\s*```', response, re.DOTALL) or re.search(r'\{.*\}', response, re.DOTALL)
+            # Try extracting JSON from markdown code blocks
+            match = re.search(r'```json\s*(.*?)\s*```', response, re.DOTALL)
+            if match:
+                return json.loads(match.group(1))
+            # Try extracting raw JSON object
+            match = re.search(r'\{.*\}', response, re.DOTALL)
             if match:
-                return json.loads(match.group(1) if '```' in response else match.group(0))
+                return json.loads(match.group(0))
             raise ValueError(f"Could not parse LLM response: {response[:200]}")
+
+    def run(self, phase2_output: dict) -> dict:
+        prompt = self._build_prompt(phase2_output)
+        response = self.llm.invoke(prompt)
+        return self._parse_response(response)

app.py

@@ -9,12 +9,55 @@ from agent.agent import MaintenanceAgent
 
 load_dotenv()
 logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = FastAPI(title="Solar PV Predictive Maintenance API", version="1.0.0")
 
 # Load ML models once on startup for production performance
 ml_engine = MLEngine()
 
+# ============ Helper Functions ============
+
+def validate_sensor_data(vdc1: list, idc1: list) -> None:
+    """Validate sensor data consistency. Raises HTTPException on error."""
+    if len(vdc1) != len(idc1):
+        raise HTTPException(status_code=400, detail="Voltage and current lists must have the same length")
+    if len(vdc1) < 3:
+        raise HTTPException(status_code=400, detail="Need at least 3 data points")
+
+def prepare_dataframe(vdc1: list, idc1: list) -> pd.DataFrame:
+    """Prepare sensor data for ML inference by padding to 100 points."""
+    return pd.DataFrame({
+        "vdc1": (vdc1 * (100 // len(vdc1) + 1))[:100],
+        "idc1": (idc1 * (100 // len(idc1) + 1))[:100]
+    })
+
+def get_agent_output(api_key: str, ml_output: dict) -> dict:
+    """Get agent analysis if API key is provided, otherwise return no-key message."""
+    if not api_key:
+        return {
+            "diagnosis": "No API key provided - LLM features disabled",
+            "urgency": "Unknown",
+            "recommended_action": "Provide Google API key in request for AI diagnosis",
+            "justification": ["Google API key required for maintenance reasoning"]
+        }
+    
+    try:
+        agent = MaintenanceAgent(
+            api_key=api_key,
+            model_name="gemini-2.5-flash-lite",
+            temperature=0.0
+        )
+        return agent.run(ml_output)
+    except Exception as e:
+        logger.warning(f"Agent initialization failed: {e}")
+        return {
+            "diagnosis": "Agent initialization failed",
+            "urgency": "Unknown",
+            "recommended_action": "Check your Google API key",
+            "justification": [f"Error: {str(e)}"]
+        }
+
 class SensorData(BaseModel):
     vdc1: list[float]
     idc1: list[float]
@@ -27,54 +70,24 @@ class AnalysisResponse(BaseModel):
 @app.post("/analyze", response_model=AnalysisResponse)
 async def analyze_sensor_data(data: SensorData):
     try:
-        logging.info(f"Processing request with {len(data.vdc1)} voltage and {len(data.idc1)} current data points")
-        
-        if len(data.vdc1) != len(data.idc1):
-            raise HTTPException(status_code=400, detail="Voltage and current lists must have the same length")
-        
-        if len(data.vdc1) < 3:
-            raise HTTPException(status_code=400, detail="Need at least 3 data points")
+        logger.info(f"Processing request with {len(data.vdc1)} voltage and {len(data.idc1)} current data points")
         
-        # Repeat to make at least 100 points if needed
-        raw_df = pd.DataFrame({
-            "vdc1": (data.vdc1 * (100 // len(data.vdc1) + 1))[:100],
-            "idc1": (data.idc1 * (100 // len(data.idc1) + 1))[:100]
-        })
+        # Validate input
+        validate_sensor_data(data.vdc1, data.idc1)
         
-        # ML Inference
-        phase2_output = ml_engine.predict_from_raw(raw_df)
+        # Prepare data and run ML inference
+        raw_df = prepare_dataframe(data.vdc1, data.idc1)
+        ml_output = ml_engine.predict_from_raw(raw_df)
         
-        # Agent Reasoning (if API key provided)
-        if data.api_key:
-            try:
-                request_agent = MaintenanceAgent(
-                    api_key=data.api_key,
-                    model_name="gemini-2.5-flash-lite",
-                    temperature=0.0
-                )
-                agent_output = request_agent.run(phase2_output)
-            except Exception as e:
-                logging.warning(f"Agent initialization failed: {e}")
-                agent_output = {
-                    "diagnosis": "Agent initialization failed",
-                    "urgency": "Unknown",
-                    "recommended_action": "Check your Google API key",
-                    "justification": [f"Error: {str(e)}"]
-                }
-        else:
-            agent_output = {
-                "diagnosis": "No API key provided - LLM features disabled",
-                "urgency": "Unknown",
-                "recommended_action": "Provide Google API key in request for AI diagnosis",
-                "justification": ["Google API key required for maintenance reasoning"]
-            }
+        # Get agent analysis
+        agent_output = get_agent_output(data.api_key, ml_output)
         
-        return AnalysisResponse(ml_output=phase2_output, agent_output=agent_output)
+        return AnalysisResponse(ml_output=ml_output, agent_output=agent_output)
     
     except HTTPException:
         raise
     except Exception as e:
-        logging.error(f"Error processing request: {e}")
+        logger.error(f"Error processing request: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")

ml/inference.py

@@ -16,6 +16,16 @@ ARTIFACTS_DIR = os.path.join(BASE_DIR, "artifacts")
 
 class MLEngine:
     def __init__(self):
+        # Load configuration
+        self._load_config()
+        # Load all models
+        self._load_scaler()
+        self._load_isolation_forest()
+        self._load_xgboost_models()
+        self._load_lstm_model()
+
+    def _load_config(self):
+        """Load ML configuration from JSON."""
         with open(os.path.join(ARTIFACTS_DIR, "ml_config.json")) as f:
             self.config = json.load(f)
 
@@ -24,33 +34,39 @@ class MLEngine:
         self.seq_len = self.config["seq_len"]
         self.design_life_days = self.config["design_life_days"]
 
-        # Load scaler from JSON
+    def _load_scaler(self):
+        """Load and reconstruct StandardScaler from JSON."""
         with open(os.path.join(ARTIFACTS_DIR, "scaler.json"), "r") as f:
             params = json.load(f)
+        
         self.scaler = StandardScaler()
         self.scaler.mean_ = np.array(params["mean"])
         self.scaler.scale_ = np.array(params["scale"])
         self.scaler.var_ = self.scaler.scale_ ** 2
         self.scaler.n_features_in_ = len(self.scaler.mean_)
 
-        # Retrain IsolationForest at startup using saved training data
+    def _load_isolation_forest(self):
+        """Load and retrain IsolationForest using saved training data."""
         self.iso = IsolationForest(
             n_estimators=200,
             contamination=0.05,
             random_state=42
         )
-        # Load training data (scaled features from Colab) and fit
         train_data = pd.read_json(os.path.join(ARTIFACTS_DIR, "training_data.json"))
         self.iso.fit(train_data[self.feature_cols])
 
-        # Load XGBoost from JSON
+    def _load_xgboost_models(self):
+        """Load XGBoost models from JSON artifacts."""
         import xgboost as xgb
+        
         self.ttf_model = xgb.XGBRegressor()
         self.ttf_model.load_model(os.path.join(ARTIFACTS_DIR, "xgb_ttf.json"))
+        
         self.fail_model = xgb.XGBClassifier()
         self.fail_model.load_model(os.path.join(ARTIFACTS_DIR, "xgb_fail.json"))
 
-        # Load LSTM from safetensors
+    def _load_lstm_model(self):
+        """Load LSTM autoencoder from safetensors."""
         self.lstm = LSTMAutoencoder(
             input_dim=len(self.feature_cols),
             hidden_dim=32
@@ -59,21 +75,12 @@ class MLEngine:
         self.lstm.load_state_dict(state_dict)
         self.lstm.eval()
 
-    def predict_from_raw(self, raw_df: pd.DataFrame):
-        # --- Feature engineering ---
-        df = build_features(raw_df, self.window)
-        df = df[self.feature_cols].dropna()
-
-        if len(df) < self.seq_len:
-            raise ValueError("Not enough data for LSTM sequence")
-
-        # --- Scaling ---
-        df_scaled = pd.DataFrame(
-            self.scaler.transform(df),
-            columns=self.feature_cols,
-            index=df.index
-        )
-
+    def _compute_anomalies(self, df_scaled: pd.DataFrame) -> tuple:
+        """Compute anomaly scores from LSTM and IsolationForest.
+        
+        Returns:
+            (anomaly_lstm, health) tuple
+        """
         # --- Isolation Forest anomaly ---
         df_scaled["anomaly_iforest"] = -self.iso.decision_function(df_scaled)
 
@@ -91,7 +98,14 @@ class MLEngine:
         anomaly_norm = min(anomaly_lstm / 1e6, 1.0)
         health = max(0.0, 1.0 - anomaly_norm)
 
-        # --- ML predictions ---
+        return anomaly_lstm, health
+
+    def _make_predictions(self, df_scaled: pd.DataFrame, anomaly_lstm: float, health: float) -> dict:
+        """Make TTF and failure probability predictions.
+        
+        Returns:
+            Dictionary with ttf, failure_prob, and rul predictions
+        """
         latest_features = df_scaled[self.feature_cols].iloc[[-1]].copy()
         latest_features["anomaly_lstm"] = anomaly_lstm
         latest_features["health_index"] = health
@@ -115,9 +129,37 @@ class MLEngine:
         )
 
         return {
-            "asset_id": "PV_INVERTER_001",
-            "failure_probability": round(failure_probability, 2),
-            "expected_ttf_days": round(expected_ttf_days, 1),
-            "expected_rul_days": round(expected_rul_days, 1),
+            "ttf_days": expected_ttf_days,
+            "failure_prob": failure_probability,
+            "rul_days": expected_rul_days,
             "confidence": confidence
         }
+
+    def predict_from_raw(self, raw_df: pd.DataFrame):
+        # --- Feature engineering ---
+        df = build_features(raw_df, self.window)
+        df = df[self.feature_cols].dropna()
+
+        if len(df) < self.seq_len:
+            raise ValueError("Not enough data for LSTM sequence")
+
+        # --- Scaling ---
+        df_scaled = pd.DataFrame(
+            self.scaler.transform(df),
+            columns=self.feature_cols,
+            index=df.index
+        )
+
+        # --- Compute anomalies ---
+        anomaly_lstm, health = self._compute_anomalies(df_scaled)
+
+        # --- Make predictions ---
+        predictions = self._make_predictions(df_scaled, anomaly_lstm, health)
+
+        return {
+            "asset_id": "PV_INVERTER_001",
+            "failure_probability": round(predictions["failure_prob"], 2),
+            "expected_ttf_days": round(predictions["ttf_days"], 1),
+            "expected_rul_days": round(predictions["rul_days"], 1),
+            "confidence": predictions["confidence"]
+        }