src/rag/llm_explainer.py

"""
LLM Explanation Module for Predictive Maintenance
Author: Antigravity AI
Date: August 2026

This module loads a lightweight instruction LLM (SmolLM2-135M, ~270MB) to run locally
on CPU. It formats current sensor telemetry, classification probabilities, RUL forecasts,
and retrieved FAISS historical logs into structured prompts, generating plain-language
reports for production managers.
"""

import os
import re
import time
import logging
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
try:
    from utils.project_paths import LOGS_DIR
except ImportError:
    from project_paths import LOGS_DIR

# Setup logging
os.makedirs(LOGS_DIR, exist_ok=True)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOGS_DIR, "llm_explainer.log"))
    ]
)
logger = logging.getLogger(__name__)

# Constants
DEFAULT_LOCAL_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"


def clean_report(text: str) -> str:
    """Post-processing to remove prompt echo, brackets, duplicate lines, and raw timestamps."""
    # Remove raw timestamp patterns
    text = re.sub(r'\b\d{4}-\d{2}-\d{2}(?:\s+\d{2}:\d{2}:\d{2})?\b', '', text)
    text = re.sub(r'Timestamp:\s*', '', text)
    
    # Remove instruction brackets like [1-2 sentences...], [XX%], [Healthy/At Risk/Critical]
    text = re.sub(r'\[(?:1-2|2-3|3-4|\d+-\d+|XX%|Healthy|At Risk|Critical|One clear|Explain|Immediate|Monitoring|Preventive)[^\]]*\]', '', text, flags=re.IGNORECASE)
    
    lines = text.strip().split("\n")
    seen = set()
    cleaned = []
    
    for line in lines:
        stripped = line.strip()
        # Remove lines that are just leftover bracket fragments or prompt echos
        if stripped in ["[", "]", "[]"] or stripped.startswith("[Machine Status:") or stripped.startswith("[Past Reference"):
            continue
            
        normalized = stripped.lower()
        if normalized:
            if normalized not in seen:
                seen.add(normalized)
                cleaned.append(stripped)
        else:
            if cleaned and cleaned[-1] != "":
                cleaned.append("")
                
    result = "\n".join(cleaned).strip()
    return result


class LLMExplainer:
    """Explainer client utilizing a local CPU instruction model with structured prompt formatting."""
    def __init__(self, model_name: str = DEFAULT_LOCAL_MODEL, force_fallback: bool = False):
        self.force_fallback = force_fallback
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.generator = None
        self.is_local_loaded = False
        
        if not force_fallback:
            try:
                logger.info(f"Initializing local LLM explainer ({model_name}) on CPU...")
                start_time = time.time()
                
                # Load tokenizer and model
                self.tokenizer = AutoTokenizer.from_pretrained(model_name)
                self.model = AutoModelForCausalLM.from_pretrained(
                    model_name,
                    torch_dtype=torch.float32,
                    device_map="cpu"
                )
                
                self.generator = pipeline(
                    "text-generation",
                    model=self.model,
                    tokenizer=self.tokenizer,
                    device=-1  # force CPU
                )
                
                elapsed = time.time() - start_time
                logger.info(f"Local LLM loaded successfully in {elapsed:.2f}s.")
                self.is_local_loaded = True
            except Exception as e:
                logger.error(f"Failed to load local LLM: {e}")
                raise RuntimeError(f"Failed to load local LLM: {e}") from e
        else:
            logger.info("Forced fallback mode active.")

    def _format_historical_logs(self, historical_logs: list) -> str:
        """Format FAISS vector search hits into clean human-readable bullets."""
        if not historical_logs:
            return "No highly similar past failure cases found in vector database."
        
        formatted_items = []
        for idx, hit in enumerate(historical_logs[:2]):
            raw_text = hit.get('log', '')
            # Clean out technical prefixes and timestamps
            cleaned = re.sub(r'\b\d{4}-\d{2}-\d{2}(?:\s+\d{2}:\d{2}:\d{2})?\b', '', raw_text)
            cleaned = re.sub(r'Timestamp:\s*\|\s*', '', cleaned)
            cleaned = re.sub(r'^On\s*,\s*', '', cleaned)
            cleaned = cleaned.strip().strip('|').strip()
            
            dist = hit.get('l2_distance', 0.5)
            similarity = max(50, min(99, int((1.0 - (dist / 2.0)) * 100)))
            formatted_items.append(f"- **Match {idx+1} ({similarity}% similarity):** {cleaned}")
            
        return "\n".join(formatted_items)

    def _generate_rule_based_report(self, machine_id: int, failure_prob: float, rul: float, 
                                    historical_logs: list, sensor_values: dict = None) -> str:
        """Deterministic structured markdown report adhering strictly to the exact format."""
        # Threshold logic: 0-15% Healthy, 16-40% At Risk, 41%+ Critical
        if failure_prob > 0.40:
            risk_status = "Critical"
        elif failure_prob > 0.15:
            risk_status = "At Risk"
        else:
            risk_status = "Healthy"
            
        volt = sensor_values.get('volt', 170.5) if sensor_values else 170.5
        rot = sensor_values.get('rotate', 440.2) if sensor_values else 440.2
        press = sensor_values.get('pressure', 100.1) if sensor_values else 100.1
        vib = sensor_values.get('vibration', 40.2) if sensor_values else 40.2

        if risk_status == "Critical":
            summary = f"Machine {machine_id} is in critical condition and requires immediate operational intervention."
            analysis = (
                f"Elevated telemetry readings (Voltage={volt}, Vibration={vib} mm/s, Pressure={press} kPa) indicate active component distress. "
                f"With a {failure_prob * 100:.1f}% failure probability, continued operation will lead to an unplanned breakdown during the current shift."
            )
            actions = (
                "1. Immediately halt current job and switch production to designated backup machine.\n"
                "2. Dispatch maintenance technician for physical vibration and bearing assembly inspection.\n"
                "3. Perform electrical load testing and verify pressure relief valve functionality."
            )
        elif risk_status == "At Risk":
            summary = f"Machine {machine_id} is showing early degradation symptoms across operational telemetry channels."
            analysis = (
                f"Current telemetry values (Voltage={volt}, Rotation={rot} RPM, Vibration={vib} mm/s) indicate moderate wear beyond nominal baseline tolerances. "
                f"A failure probability of {failure_prob * 100:.1f}% indicates that proactive maintenance is needed to prevent escalation to critical failure."
            )
            actions = (
                "1. Schedule targeted component inspection during the next scheduled shift change.\n"
                "2. Increase telemetry polling frequency and set real-time threshold alert triggers.\n"
                "3. Inspect mechanical lubrication levels and check drive belt tension."
            )
        else:
            summary = f"Machine {machine_id} is operating in a stable, healthy state within expected operational parameters."
            analysis = (
                f"Sensor telemetry (Voltage={volt}, Rotation={rot} RPM, Pressure={press} kPa, Vibration={vib} mm/s) remains strictly within normal baseline ranges. "
                f"The low failure probability of {failure_prob * 100:.1f}% confirms steady-state production health with no immediate maintenance required."
            )
            actions = (
                "1. Continue standard production schedule without interruption.\n"
                "2. Maintain routine automated sensor logging and health telemetry monitoring.\n"
                "3. Proceed with standard scheduled maintenance intervals."
            )

        hist_context = self._format_historical_logs(historical_logs)

        report = f"""**Machine Status:** {risk_status}
**Failure Probability:** {failure_prob * 100:.1f}%

**Summary:**
{summary}

**Analysis:**
{analysis}

**Historical Context (FAISS Vector Database):**
{hist_context}

**Recommended Actions:**
{actions}"""
        return report

    def generate_explanation(self, machine_id: int, failure_prob: float, rul: float, 
                             historical_logs: list, sensor_values: dict = None) -> str:
        """Construct the prompt and generate the markdown maintenance explanation report."""
        if not self.is_local_loaded:
            return self._generate_rule_based_report(machine_id, failure_prob, rul, historical_logs, sensor_values)
            
        # Threshold logic: 0-15% Healthy, 16-40% At Risk, 41%+ Critical
        if failure_prob > 0.40:
            risk_status = "Critical"
        elif failure_prob > 0.15:
            risk_status = "At Risk"
        else:
            risk_status = "Healthy"
            
        if sensor_values:
            sensors_str = (
                f"Voltage={sensor_values.get('volt', '170.5')}, "
                f"Rotation={sensor_values.get('rotate', '440.2')}, "
                f"Pressure={sensor_values.get('pressure', '100.1')}, "
                f"Vibration={sensor_values.get('vibration', '40.2')}"
            )
        else:
            sensors_str = "Voltage=170.5, Rotation=440.2, Pressure=100.1, Vibration=40.2"

        system_prompt = (
            "You are a factory maintenance assistant. Write a short professional diagnostic report.\n"
            "Do not repeat sentences. Do not use square brackets. Write complete sentences.\n\n"
            "Format:\n"
            "**Machine Status:** Healthy, At Risk, or Critical\n"
            "**Failure Probability:** percentage\n\n"
            "**Summary:**\n"
            "One clear sentence about machine condition.\n\n"
            "**Analysis:**\n"
            "Two sentences explaining the sensor readings and failure probability.\n\n"
            "**Recommended Actions:**\n"
            "1. Action one\n"
            "2. Action two\n"
            "3. Action three"
        )

        user_prompt = (
            f"Machine ID: {machine_id}\n"
            f"Machine Status: {risk_status}\n"
            f"Failure Probability: {failure_prob * 100:.1f}%\n"
            f"Sensor Readings: {sensors_str}\n\n"
            "Write the diagnostic report now."
        )

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        try:
            logger.info(f"Generating explanation report for Machine {machine_id} using local LLM...")
            start_time = time.time()
            
            prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            outputs = self.generator(
                prompt,
                max_new_tokens=180,
                temperature=0.1,
                top_p=0.8,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id
            )
            
            generated_text = outputs[0]["generated_text"]
            if "<|im_start|>assistant" in generated_text:
                report = generated_text.split("<|im_start|>assistant")[-1].strip()
            else:
                report = generated_text[len(prompt):].strip()
                
            cleaned = clean_report(report)
            
            # Validation: If the 135M model degenerated or parroted instructions, return the validated structure
            if "**Summary:**" not in cleaned or "**Analysis:**" not in cleaned or "[" in cleaned or len(cleaned.split("\n")) < 5:
                logger.warning("LLM generated incomplete or bracket-polluted text. Using validated structure.")
                return self._generate_rule_based_report(machine_id, failure_prob, rul, historical_logs, sensor_values)
                
            # Append FAISS vector database context to the report
            hist_context = self._format_historical_logs(historical_logs)
            if "**Historical Context" not in cleaned:
                # Insert before recommended actions
                if "**Recommended Actions:**" in cleaned:
                    parts = cleaned.split("**Recommended Actions:**")
                    cleaned = f"{parts[0].strip()}\n\n**Historical Context (FAISS Vector Database):**\n{hist_context}\n\n**Recommended Actions:**\n{parts[1].strip()}"
                else:
                    cleaned += f"\n\n**Historical Context (FAISS Vector Database):**\n{hist_context}"
                    
            elapsed = time.time() - start_time
            logger.info(f"Explanation report generated in {elapsed:.2f}s.")
            return cleaned
            
        except Exception as e:
            logger.error(f"Error during LLM generation: {e}.")
            return self._generate_rule_based_report(machine_id, failure_prob, rul, historical_logs, sensor_values)


if __name__ == '__main__':
    explainer = LLMExplainer(force_fallback=True)
    mock_hits = [
        {"log": "Machine ID: 100 (Model model4, Age 5) triggered warning/error code error2. Speed/vibration telemetry registered anomalous spikes.", "l2_distance": 0.32},
        {"log": "Machine ID: 100 (Model model4, Age 5) triggered warning/error code error3. Speed/vibration telemetry registered anomalous spikes.", "l2_distance": 0.45}
    ]
    report = explainer.generate_explanation(
        machine_id=101,
        failure_prob=0.30,
        rul=70.0,
        historical_logs=mock_hits,
        sensor_values={'volt': 175.2, 'rotate': 460.0, 'pressure': 102.4, 'vibration': 42.1}
    )
    print("\n--- Test Report with FAISS Historical Context: ---")
    print(report)