Update app.py

app.py

@@ -3,7 +3,7 @@ import mauve
 from sacrebleu import corpus_bleu
 from rouge_score import rouge_scorer
 from bert_score import score
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline, AutoTokenizer, AutoModelForCausalLM
 import re
 from mauve import compute_mauve
 import os
@@ -54,7 +54,6 @@ class RAGEvaluator:
         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
         return model, tokenizer
 
-    # BLEU, ROUGE
     def evaluate_bleu_rouge(self, candidates, references):
         try:
             bleu_score = corpus_bleu(candidates, [references]).score
@@ -68,7 +67,6 @@ class RAGEvaluator:
             print(f"BLEU/ROUGE evaluation failed: {e}")
             return 0, 0, 0, 0
 
-    # BERT Score
     def evaluate_bert_score(self, candidates, references):
         try:
             P, R, F1 = score(candidates, references, lang="en", model_type='bert-base-multilingual-cased')
@@ -77,7 +75,6 @@ class RAGEvaluator:
             print(f"BERT score evaluation failed: {e}")
             return 0, 0, 0
 
-    # Perplexity
     def evaluate_perplexity(self, text):
         try:
             encodings = self.gpt2_tokenizer(text, return_tensors='pt')
@@ -99,9 +96,8 @@ class RAGEvaluator:
             return ppl.item()
         except Exception as e:
             print(f"Perplexity evaluation failed: {e}")
-            return 1000.0  # High perplexity indicates error
+            return 1000.0
 
-    # Diversity
     def evaluate_diversity(self, texts):
         try:
             all_tokens = []
@@ -116,7 +112,6 @@ class RAGEvaluator:
             print(f"Diversity evaluation failed: {e}")
             return 0
 
-    # Racial bias
     def evaluate_racial_bias(self, text):
         try:
             results = self.bias_pipeline([text], candidate_labels=["hate speech", "not hate speech"])
@@ -126,7 +121,6 @@ class RAGEvaluator:
             print(f"Bias evaluation failed: {e}")
             return 0
 
-    # METEOR
     def evaluate_meteor(self, candidates, references):
         try:
             meteor_scores = []
@@ -146,7 +140,6 @@ class RAGEvaluator:
             print(f"METEOR evaluation failed: {e}")
             return 0
 
-    # CHRF
     def evaluate_chrf(self, candidates, references):
         try:
             chrf_scores = []
@@ -169,7 +162,6 @@ class RAGEvaluator:
             print(f"CHRF evaluation failed: {e}")
             return 0
 
-    # Readability
     def evaluate_readability(self, text):
         try:
             words = re.findall(r'\b\w+\b', text.lower())
@@ -185,7 +177,6 @@ class RAGEvaluator:
             print(f"Readability evaluation failed: {e}")
             return 0, 0
 
-    # MAUVE
     def evaluate_mauve(self, reference_texts, generated_texts):
         try:
             out = compute_mauve(
@@ -276,16 +267,98 @@ def get_retrieved_context(query: str, top_k=3):
         if mlflow.active_run():
             mlflow.log_metric("retrieved_chunks", len(results['matches']))
         
-        return "\n".join([m['metadata']['text'] for m in results['matches']])
+        context_texts = [m['metadata']['text'] for m in results['matches']]
+        return "\n".join(context_texts) if context_texts else "No relevant context found."
     except Exception as e:
         print(f"Context retrieval failed: {e}")
         return f"Context retrieval error: {str(e)}"
 
-# ------------------ Custom LLM ------------------
-class LitServeLLM(LLM):
-    endpoint_url: str
+# ------------------ Fallback LLM Models ------------------
+class FallbackLLM:
+    def __init__(self):
+        self.models_loaded = False
+        self.pipeline = None
+        self.load_fallback_models()
+    
+    def load_fallback_models(self):
+        """Load local models as fallback"""
+        try:
+            # Use a smaller model for fallback
+            self.pipeline = pipeline(
+                "text-generation",
+                model="microsoft/DialoGPT-small",
+                tokenizer="microsoft/DialoGPT-small",
+                max_length=150,
+                do_sample=True,
+                temperature=0.7
+            )
+            self.models_loaded = True
+            print("Fallback model loaded successfully")
+        except Exception as e:
+            print(f"Fallback model loading failed: {e}")
+            self.models_loaded = False
+    
+    def generate_response(self, context, question):
+        if not self.models_loaded:
+            return "I'm currently experiencing technical difficulties. Please try again later."
+        
+        try:
+            prompt = f"""
+Based on the following context, please provide a concise answer to the question.
 
+Context: {context}
+
+Question: {question}
+
+Answer: """
+            
+            response = self.pipeline(
+                prompt,
+                max_new_tokens=100,
+                num_return_sequences=1,
+                pad_token_id=50256
+            )
+            
+            if response and len(response) > 0:
+                full_response = response[0]['generated_text']
+                # Extract only the answer part
+                if "Answer:" in full_response:
+                    answer = full_response.split("Answer:")[-1].strip()
+                    return answer
+                return full_response.strip()
+            else:
+                return "I couldn't generate a response at the moment. Please try again."
+                
+        except Exception as e:
+            print(f"Fallback model generation failed: {e}")
+            return "I'm having trouble generating a response. Please try again later."
+
+# Initialize fallback LLM
+fallback_llm = FallbackLLM()
+
+# ------------------ Custom LLM with Fallback ------------------
+class RobustLitServeLLM(LLM):
+    endpoint_url: str
+    use_fallback: bool = True
+    
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        # Try the primary endpoint first
+        primary_success, primary_response = self._try_primary_endpoint(prompt)
+        
+        if primary_success:
+            return primary_response
+        
+        # If primary fails and fallback is enabled, use fallback
+        if self.use_fallback:
+            print("Using fallback LLM due to primary endpoint failure")
+            # Extract context and question from prompt
+            context, question = self._extract_context_question(prompt)
+            return fallback_llm.generate_response(context, question)
+        else:
+            return "I apologize, but the AI service is currently unavailable. Please try again later."
+    
+    def _try_primary_endpoint(self, prompt: str):
+        """Try to get response from primary endpoint"""
         try:
             payload = {"prompt": prompt}
             start_time = time.time()
@@ -298,43 +371,52 @@ class LitServeLLM(LLM):
                 data = response.json()
                 if mlflow.active_run():
                     mlflow.log_metric("response_tokens", len(data.get("response", "").split()))
-                return data.get("response", "").strip()
+                return True, data.get("response", "").strip()
             else:
                 if mlflow.active_run():
                     mlflow.log_metric("request_errors", 1)
-                error_msg = f"Request failed with status {response.status_code}"
-                print(f"LLM Error: {error_msg}")
-                return f"I apologize, but I'm currently experiencing technical difficulties. Error: {error_msg}"
+                print(f"Primary endpoint failed with status: {response.status_code}")
+                return False, ""
                 
-        except requests.exceptions.Timeout:
-            error_msg = "Request timeout - service unavailable"
-            print(f"LLM Timeout: {error_msg}")
-            return f"I apologize, but the service is currently unavailable. Please try again later."
         except Exception as e:
-            error_msg = f"Connection error: {str(e)}"
-            print(f"LLM Connection Error: {error_msg}")
-            return f"I apologize, but I'm having trouble connecting to the service. Please try again later."
+            print(f"Primary endpoint error: {e}")
+            return False, ""
+    
+    def _extract_context_question(self, prompt: str):
+        """Extract context and question from the prompt template"""
+        try:
+            if "Context:" in prompt and "Question:" in prompt:
+                context_part = prompt.split("Context:")[1].split("Question:")[0].strip()
+                question_part = prompt.split("Question:")[1].split("Answer:")[0].strip()
+                return context_part, question_part
+            return "", prompt
+        except:
+            return "", prompt
 
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
-        return {"endpoint_url": self.endpoint_url}
+        return {"endpoint_url": self.endpoint_url, "use_fallback": self.use_fallback}
 
     @property
     def _llm_type(self) -> str:
-        return "litserve_llm"
+        return "robust_litserve_llm"
 
-# Initialize model with fallback
+# Initialize the robust model
 try:
-    model = LitServeLLM(endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict")
+    model = RobustLitServeLLM(
+        endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict",
+        use_fallback=True
+    )
+    print("Robust LLM initialized successfully")
 except Exception as e:
-    print(f"Model initialization failed: {e}")
+    print(f"Robust LLM initialization failed: {e}")
     model = None
 
+# ------------------ Prompt Template ------------------
 prompt = PromptTemplate(
     input_variables=["context", "question"],
     template="""
-You are a smart assistant. Based on the provided context, answer the question in 1–2 lines only.
-If the context has more details, summarize it concisely.
+You are a smart maintenance assistant. Based on the provided context, answer the question concisely in 1-2 lines.
 
 Context:
 {context}
@@ -345,63 +427,71 @@ Answer:
 """
 )
 
-# Initialize chain with error handling
+# Initialize LLM chain
 try:
     if model:
         llm_chain = LLMChain(llm=model, prompt=prompt)
+        print("LLM chain initialized successfully")
     else:
         llm_chain = None
-        print("LLM chain not initialized due to model failure")
+        print("LLM chain not initialized - no model available")
 except Exception as e:
     print(f"LLM chain initialization failed: {e}")
     llm_chain = None
 
 # ------------------ RAG Pipeline ------------------
 def get_rag_response(question):
-    """Get the complete RAG response with error handling"""
+    """Get the complete RAG response with robust error handling"""
     try:
-        # Get context
+        if not question.strip():
+            return "Please enter a valid question.", ""
+        
+        # Get context from Pinecone
         retrieved_context = get_retrieved_context(question)
         
-        # If LLM chain is not available, return fallback response
-        if not llm_chain:
-            fallback_response = "I'm currently experiencing technical difficulties. Please try again later or contact support."
+        # If we have an LLM chain, use it
+        if llm_chain:
+            result = llm_chain.invoke({
+                "context": retrieved_context,
+                "question": question
+            })
+            full_response = result["text"].strip()
+            
+            # Clean up the response
+            if "Answer:" in full_response:
+                full_response = full_response.split("Answer:")[-1].strip()
+            
+            return full_response, retrieved_context
+        else:
+            # Use direct fallback
+            fallback_response = fallback_llm.generate_response(retrieved_context, question)
             return fallback_response, retrieved_context
-        
-        # Get response from LLM
-        result = llm_chain.invoke({
-            "context": retrieved_context,
-            "question": question
-        })
-        
-        full_response = result["text"].strip()
-        
-        if "Answer:" in full_response:
-            full_response = full_response.split("Answer:", 1)[-1].strip()
-        
-        return full_response, retrieved_context
-        
+            
     except Exception as e:
-        error_msg = f"Error generating response: {str(e)}"
-        print(f"RAG pipeline error: {error_msg}")
-        return f"I apologize, but I encountered an error while processing your request. Please try again. Error: {str(e)}", "Error retrieving context"
+        error_msg = f"Error in RAG pipeline: {str(e)}"
+        print(error_msg)
+        # Final fallback - simple response based on context
+        if "context" in locals() and retrieved_context:
+            return f"Based on available information: I found relevant maintenance data, but encountered an issue processing it. Context available: {len(retrieved_context)} characters.", retrieved_context
+        else:
+            return "I apologize, but I'm experiencing technical difficulties. Please try again later or contact support.", "No context retrieved"
 
 def rag_pipeline_stream(question):
     """Streaming version of RAG pipeline"""
     try:
         full_response, _ = get_rag_response(question)
         
-        # Stream word by word
+        # Stream word by word for better UX
         words = full_response.split()
         current_text = ""
         for word in words:
             current_text += word + " "
             yield current_text
-            time.sleep(0.05)  # Adjust speed as needed
+            time.sleep(0.03)  # Faster streaming
             
     except Exception as e:
         error_msg = f"Error in streaming: {str(e)}"
-        print(f"Streaming error: {error_msg}")
+        print(error_msg)
         yield "I apologize, but I encountered an error while generating the response."
 
 # ------------------ Gradio UI ------------------
@@ -409,6 +499,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
     gr.Markdown("""
     # 🛠 Maintenance AI Assistant
     *Your intelligent companion for maintenance queries and troubleshooting*
+    
+    **Note**: This system uses multiple fallback mechanisms to ensure reliability.
     """)
     
     usage_counter = gr.State(value=0)
@@ -420,10 +512,14 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
             gr.Markdown("### 💬 Chat Interface")
             question_input = gr.Textbox(
                 label="Ask your maintenance question",
-                placeholder="e.g., How do I troubleshoot a leaking valve?",
-                lines=2
+                placeholder="e.g., How do I troubleshoot a leaking valve? What's the maintenance schedule for pumps?",
+                lines=3
             )
-            ask_button = gr.Button("Get Answer 🚀", variant="primary")
+            ask_button = gr.Button("Get Answer 🚀", variant="primary", size="lg")
+            
+            with gr.Row():
+                clear_btn = gr.Button("Clear Chat 🗑️")
+                evaluate_btn = gr.Button("Show Metrics 📈", variant="secondary")
             
             feedback = gr.Radio(
                 ["Helpful", "Not Helpful"], 
@@ -431,21 +527,22 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
                 info="Your feedback helps improve the system"
             )
             
-            gr.Markdown("### 📊 Evaluation Metrics")
-            metrics_output = gr.JSON(label="Quality Metrics", visible=False)
-            
         with gr.Column(scale=1):
             gr.Markdown("### 🤖 AI Response")
             answer_output = gr.Textbox(
                 label="Response", 
-                lines=6,
+                lines=8,
                 interactive=False,
-                show_copy_button=True
+                show_copy_button=True,
+                autoscroll=True
             )
             
-            with gr.Row():
-                clear_btn = gr.Button("Clear Chat 🗑️")
-                evaluate_btn = gr.Button("Show Metrics 📈", variant="secondary")
+            gr.Markdown("### 📊 Evaluation Metrics")
+            metrics_output = gr.JSON(
+                label="Quality Metrics", 
+                visible=False,
+                show_label=True
+            )
 
     def track_usage(question, count, session_start, feedback_value=None):
         """Track usage and get response"""
@@ -455,8 +552,9 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
         count += 1
         
         try:
-            with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True) if mlflow_tracking_uri else dummy_context():
-                if mlflow_tracking_uri:
+            # Only use MLflow if properly configured
+            if mlflow_tracking_uri:
+                with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True):
                     mlflow.log_param("question", question)
                     mlflow.log_param("session_start", session_start)
                     mlflow.log_param("user_feedback", feedback_value or "No feedback")
@@ -465,36 +563,38 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
                         mlflow.log_metric("helpful_responses", 1 if feedback_value == "Helpful" else 0)
                     
                     mlflow.log_metric("total_queries", count)
-                
-                # Get response and context
-                response, context = get_rag_response(question)
-                
-                if mlflow_tracking_uri:
+                    
+                    # Get response and context
+                    response, context = get_rag_response(question)
+                    
                     mlflow.log_metric("response_length", len(response))
                     mlflow.log_metric("response_tokens", len(response.split()))
-                
+                    mlflow.log_metric("context_length", len(context))
+                    
+                    return response, count, session_start, response
+            else:
+                # Without MLflow
+                response, context = get_rag_response(question)
                 return response, count, session_start, response
                 
         except Exception as e:
             print(f"Tracking error: {e}")
-            error_msg = f"System error: {str(e)}"
-            return error_msg, count, session_start, error_msg
+            error_response = f"I encountered a system error. Please try again. Error: {str(e)}"
+            return error_response, count, session_start, error_response
 
     def evaluate_response(question, response):
         """Evaluate the response and return metrics"""
-        if not question or not response or "error" in response.lower() or "apologize" in response.lower():
-            return gr.update(value={"info": "Evaluation skipped due to error response"}, visible=True)
+        if not question or not response:
+            return gr.update(value={"info": "No question or response to evaluate"}, visible=True)
+        
+        # Skip evaluation for error responses
+        if any(error_word in response.lower() for error_word in ["error", "apologize", "unavailable", "technical"]):
+            return gr.update(value={"info": "Evaluation skipped for error response"}, visible=True)
             
         try:
             context = get_retrieved_context(question)
             metrics = evaluator.evaluate_all(question, response, context)
             
-            # Log metrics to MLflow if available
-            if mlflow_tracking_uri and mlflow.active_run():
-                for metric_name, metric_value in metrics.items():
-                    if isinstance(metric_value, (int, float)):
-                        mlflow.log_metric(metric_name, metric_value)
-            
             return gr.update(value=metrics, visible=True)
         except Exception as e:
             print(f"Evaluation error: {e}")
@@ -504,16 +604,9 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
         """Clear the chat interface"""
         return "", "", gr.update(visible=False)
 
-    # Dummy context manager for when MLflow is not available
-    class dummy_context:
-        def __enter__(self):
-            return self
-        def __exit__(self, *args):
-            pass
-
     # Main interaction flow
     ask_button.click(
-        fn=lambda: ("", gr.update(visible=False)),  # Clear previous metrics
+        fn=lambda: ("", gr.update(visible=False)),  # Clear previous metrics and response
         outputs=[answer_output, metrics_output]
     ).then(
         fn=rag_pipeline_stream,
@@ -538,14 +631,13 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
         outputs=[question_input, answer_output, metrics_output]
     )
 
-    # Feedback handling
+    # Handle feedback
     def handle_feedback(feedback_val):
         try:
             if mlflow_tracking_uri and mlflow.active_run():
                 mlflow.log_metric("user_feedback_score", 1 if feedback_val == "Helpful" else 0)
         except:
-            pass
-        return
+            pass  # Silently fail if feedback logging doesn't work
 
     feedback.change(
         fn=handle_feedback,
@@ -554,9 +646,14 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo
     )
 
 if __name__ == "__main__":
+    print("🚀 Starting Maintenance AI Assistant...")
+    print("✅ System initialized with fallback mechanisms")
+    print("🌐 Web interface available at http://0.0.0.0:7860")
+    
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,  # Disable sharing to avoid the warning
-        show_error=True
+        share=False,
+        show_error=True,
+        debug=False
     )