Update app.py

app.py

@@ -389,11 +389,10 @@ def rag_pipeline(question):
         mlflow.log_dict(error_info, "artifacts/pipeline_errors.json")
         return f"Error: {str(e)}"
 
-# ----------- 7. Gradio UI with Enhanced Tracking -----------
+# ----------- 7. Gradio UI with Streaming Response -----------
 with gr.Blocks() as demo:
     gr.Markdown("# 🛠 Maintenance AI Assistant")
     
-    # Track additional UI metrics
     usage_counter = gr.State(value=0)
     session_start = gr.State(value=datetime.now().isoformat())
     
@@ -402,36 +401,62 @@ with gr.Blocks() as demo:
     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):
-        """Wrapper to track usage metrics with feedback"""
+    def track_usage_stream(question, count, session_start, feedback=None):
+        """Generator for streaming response token by token"""
         count += 1
         
-        # Start tracking context
         with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True):
             mlflow.log_param("question", question)
             mlflow.log_param("session_start", session_start)
             
-            # Get response
-            response = rag_pipeline(question)
+            # Retrieve context first
+            retrieved_context = get_retrieved_context(question)
+            mlflow.log_text(retrieved_context, "artifacts/retrieved_context.txt")
+            
+            # Stream response in chunks
+            def response_generator():
+                # Here, simulate streaming by splitting into words or chunks
+                response_full = llm_chain.invoke({
+                    "context": retrieved_context,
+                    "question": question
+                })["text"].strip()
+                
+                if "Answer:" in response_full:
+                    response_full = response_full.split("Answer:", 1)[-1].strip()
+                
+                # Yield response word by word (or token by token)
+                words = response_full.split()
+                current_text = ""
+                for word in words:
+                    current_text += word + " "
+                    yield current_text
+                
+                # After streaming is done, log evaluation
+                evaluation_metrics = evaluator.evaluate_all(
+                    question=question,
+                    response=response_full,
+                    reference=retrieved_context
+                )
+                for metric_name, metric_value in evaluation_metrics.items():
+                    mlflow.log_metric(metric_name, metric_value)
             
             # Log feedback if provided
             if feedback:
                 mlflow.log_param("user_feedback", feedback)
                 mlflow.log_metric("helpful_responses", 1 if feedback == "Helpful" else 0)
             
-            # Update metrics
             mlflow.log_metric("total_queries", count)
             
-            return response, count, session_start
+            return response_generator(), count, session_start
     
     ask_button.click(
-        track_usage, 
-        inputs=[question_input, usage_counter, session_start], 
+        track_usage_stream,
+        inputs=[question_input, usage_counter, session_start, feedback],
         outputs=[answer_output, usage_counter, session_start]
     )
     
     feedback.change(
-        lambda feedback, question, count, session_start: track_usage(question, count, session_start, feedback),
+        lambda feedback, question, count, session_start: track_usage_stream(question, count, session_start, feedback),
         inputs=[feedback, question_input, usage_counter, session_start],
         outputs=[answer_output, usage_counter, session_start]
     )