Update app.py

app.py

@@ -1,4 +1,3 @@
-##
 import torch
 import mauve
 from sacrebleu import corpus_bleu
@@ -258,52 +257,158 @@ Answer:
 
 llm_chain = LLMChain(llm=model, prompt=prompt)
 
-# ------------------ RAG Streaming Pipeline ------------------
-def rag_pipeline_stream(question):
+# ------------------ RAG Pipeline ------------------
+def get_rag_response(question):
+    """Get the complete RAG response without streaming"""
     retrieved_context = get_retrieved_context(question)
     full_response = llm_chain.invoke({
         "context": retrieved_context,
         "question": question
     })["text"].strip()
+    
     if "Answer:" in full_response:
         full_response = full_response.split("Answer:", 1)[-1].strip()
     
+    return full_response, retrieved_context
+
+def rag_pipeline_stream(question):
+    """Streaming version of RAG pipeline"""
+    full_response, _ = get_rag_response(question)
+    
     # Stream word by word
+    words = full_response.split()
     current_text = ""
-    for word in full_response.split():
+    for word in words:
         current_text += word + " "
         yield current_text
-
-    # Evaluate after streaming
-    evaluator.evaluate_all(question, full_response, retrieved_context)
+        time.sleep(0.05)  # Adjust speed as needed
 
 # ------------------ Gradio UI ------------------
-with gr.Blocks() as demo:
-    gr.Markdown("# 🛠 Maintenance AI Assistant")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🛠 Maintenance AI Assistant
+    *Your intelligent companion for maintenance queries and troubleshooting*
+    """)
+    
     usage_counter = gr.State(value=0)
     session_start = gr.State(value=datetime.now().isoformat())
+    current_response = gr.State(value="")  # Store current response for evaluation
     
-    question_input = gr.Textbox(label="Ask your maintenance question")
-    answer_output = gr.Textbox(label="AI Response")
-    ask_button = gr.Button("Get Answer")
-    feedback = gr.Radio(["Helpful", "Not Helpful"], label="Was this response helpful?")
-    
-    def track_usage(question, count, session_start, feedback=None):
+    with gr.Row():
+        with gr.Column(scale=1):
+            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
+            )
+            ask_button = gr.Button("Get Answer 🚀", variant="primary")
+            
+            feedback = gr.Radio(
+                ["Helpful", "Not Helpful"], 
+                label="Was this response helpful?",
+                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,
+                interactive=False,
+                show_copy_button=True
+            )
+            
+            with gr.Row():
+                clear_btn = gr.Button("Clear Chat 🗑️")
+                evaluate_btn = gr.Button("Show Metrics 📈", variant="secondary")
+
+    def track_usage(question, count, session_start, feedback_value=None):
+        """Track usage and get response"""
         count += 1
         with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True):
             mlflow.log_param("question", question)
             mlflow.log_param("session_start", session_start)
-            if feedback:
-                mlflow.log_param("user_feedback", feedback)
-                mlflow.log_metric("helpful_responses", 1 if feedback == "Helpful" else 0)
+            mlflow.log_param("user_feedback", feedback_value or "No feedback")
+            
+            if feedback_value:
+                mlflow.log_metric("helpful_responses", 1 if feedback_value == "Helpful" else 0)
+            
             mlflow.log_metric("total_queries", count)
-            return rag_pipeline_stream(question), count, session_start
+            
+            # Get response and context
+            response, context = get_rag_response(question)
+            
+            # Log response metrics
+            mlflow.log_metric("response_length", len(response))
+            mlflow.log_metric("response_tokens", len(response.split()))
+            
+            return response, count, session_start, response
 
+    def evaluate_response(question, response):
+        """Evaluate the response and return metrics"""
+        if not question or not response:
+            return gr.update(value={}, visible=False)
+            
+        try:
+            context = get_retrieved_context(question)
+            metrics = evaluator.evaluate_all(question, response, context)
+            
+            # Log metrics to MLflow
+            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}")
+            return gr.update(value={"error": str(e)}, visible=True)
+
+    def clear_chat():
+        """Clear the chat interface"""
+        return "", "", gr.update(visible=False)
+
+    # Main interaction flow
     ask_button.click(
-        track_usage,
+        fn=lambda: ("", gr.update(visible=False)),  # Clear previous metrics
+        outputs=[answer_output, metrics_output]
+    ).then(
+        fn=rag_pipeline_stream,
+        inputs=[question_input],
+        outputs=[answer_output]
+    ).then(
+        fn=track_usage,
         inputs=[question_input, usage_counter, session_start, feedback],
-        outputs=[answer_output, usage_counter, session_start]
+        outputs=[answer_output, usage_counter, session_start, current_response]
+    )
+
+    # Evaluation flow
+    evaluate_btn.click(
+        fn=evaluate_response,
+        inputs=[question_input, current_response],
+        outputs=[metrics_output]
+    )
+
+    # Clear chat
+    clear_btn.click(
+        fn=clear_chat,
+        outputs=[question_input, answer_output, metrics_output]
+    )
+
+    # Feedback handling
+    feedback.change(
+        fn=lambda feedback_val: mlflow.log_metric("user_feedback_score", 1 if feedback_val == "Helpful" else 0),
+        inputs=[feedback],
+        outputs=[]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        show_error=True
+    )