Update app.py

app.py

@@ -27,22 +27,27 @@ pinecone_api_key = os.environ.get("PINECONE_API_KEY")
 mlflow_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI")
 
 # ------------------ DagsHub & MLflow Setup ------------------
-dagshub.init(
-    repo_owner='prathamesh.khade20',
-    repo_name='Maintenance_AI_website',
-    mlflow=True
-)
-
-mlflow.set_tracking_uri(mlflow_tracking_uri)
-mlflow.set_experiment("Maintenance-RAG-Chatbot")
-mlflow.langchain.autolog()
+try:
+    dagshub.init(
+        repo_owner='prathamesh.khade20',
+        repo_name='Maintenance_AI_website',
+        mlflow=True
+    )
+    mlflow.set_tracking_uri(mlflow_tracking_uri)
+    mlflow.set_experiment("Maintenance-RAG-Chatbot")
+    mlflow.langchain.autolog()
+except Exception as e:
+    print(f"MLflow/DagsHub initialization failed: {e}")
 
 # ------------------ RAG Evaluator ------------------
 class RAGEvaluator:
     def __init__(self):
-        self.gpt2_model, self.gpt2_tokenizer = self.load_gpt2_model()
-        self.bias_pipeline = pipeline("zero-shot-classification", model="Hate-speech-CNERG/dehatebert-mono-english")
-        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        try:
+            self.gpt2_model, self.gpt2_tokenizer = self.load_gpt2_model()
+            self.bias_pipeline = pipeline("zero-shot-classification", model="Hate-speech-CNERG/dehatebert-mono-english")
+            self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        except Exception as e:
+            print(f"Evaluator initialization failed: {e}")
         
     def load_gpt2_model(self):
         model = GPT2LMHeadModel.from_pretrained('gpt2')
@@ -51,187 +56,264 @@ class RAGEvaluator:
 
     # BLEU, ROUGE
     def evaluate_bleu_rouge(self, candidates, references):
-        bleu_score = corpus_bleu(candidates, [references]).score
-        scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)
-        rouge_scores = [scorer.score(ref, cand) for ref, cand in zip(references, candidates)]
-        rouge1 = sum([score['rouge1'].fmeasure for score in rouge_scores]) / len(rouge_scores)
-        rouge2 = sum([score['rouge2'].fmeasure for score in rouge_scores]) / len(rouge_scores)
-        rougeL = sum([score['rougeL'].fmeasure for score in rouge_scores]) / len(rouge_scores)
-        return bleu_score, rouge1, rouge2, rougeL
+        try:
+            bleu_score = corpus_bleu(candidates, [references]).score
+            scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)
+            rouge_scores = [scorer.score(ref, cand) for ref, cand in zip(references, candidates)]
+            rouge1 = sum([score['rouge1'].fmeasure for score in rouge_scores]) / len(rouge_scores)
+            rouge2 = sum([score['rouge2'].fmeasure for score in rouge_scores]) / len(rouge_scores)
+            rougeL = sum([score['rougeL'].fmeasure for score in rouge_scores]) / len(rouge_scores)
+            return bleu_score, rouge1, rouge2, rougeL
+        except Exception as e:
+            print(f"BLEU/ROUGE evaluation failed: {e}")
+            return 0, 0, 0, 0
 
     # BERT Score
     def evaluate_bert_score(self, candidates, references):
-        P, R, F1 = score(candidates, references, lang="en", model_type='bert-base-multilingual-cased')
-        return P.mean().item(), R.mean().item(), F1.mean().item()
+        try:
+            P, R, F1 = score(candidates, references, lang="en", model_type='bert-base-multilingual-cased')
+            return P.mean().item(), R.mean().item(), F1.mean().item()
+        except Exception as e:
+            print(f"BERT score evaluation failed: {e}")
+            return 0, 0, 0
 
     # Perplexity
     def evaluate_perplexity(self, text):
-        encodings = self.gpt2_tokenizer(text, return_tensors='pt')
-        max_length = self.gpt2_model.config.n_positions
-        stride = 512
-        lls = []
-        for i in range(0, encodings.input_ids.size(1), stride):
-            begin_loc = max(i + stride - max_length, 0)
-            end_loc = min(i + stride, encodings.input_ids.size(1))
-            trg_len = end_loc - i
-            input_ids = encodings.input_ids[:, begin_loc:end_loc]
-            target_ids = input_ids.clone()
-            target_ids[:, :-trg_len] = -100
-            with torch.no_grad():
-                outputs = self.gpt2_model(input_ids, labels=target_ids)
-                log_likelihood = outputs[0] * trg_len
-            lls.append(log_likelihood)
-        ppl = torch.exp(torch.stack(lls).sum() / end_loc)
-        return ppl.item()
+        try:
+            encodings = self.gpt2_tokenizer(text, return_tensors='pt')
+            max_length = self.gpt2_model.config.n_positions
+            stride = 512
+            lls = []
+            for i in range(0, encodings.input_ids.size(1), stride):
+                begin_loc = max(i + stride - max_length, 0)
+                end_loc = min(i + stride, encodings.input_ids.size(1))
+                trg_len = end_loc - i
+                input_ids = encodings.input_ids[:, begin_loc:end_loc]
+                target_ids = input_ids.clone()
+                target_ids[:, :-trg_len] = -100
+                with torch.no_grad():
+                    outputs = self.gpt2_model(input_ids, labels=target_ids)
+                    log_likelihood = outputs[0] * trg_len
+                lls.append(log_likelihood)
+            ppl = torch.exp(torch.stack(lls).sum() / end_loc)
+            return ppl.item()
+        except Exception as e:
+            print(f"Perplexity evaluation failed: {e}")
+            return 1000.0  # High perplexity indicates error
 
     # Diversity
     def evaluate_diversity(self, texts):
-        all_tokens = []
-        for text in texts:
-            tokens = self.tokenizer.tokenize(text)
-            all_tokens.extend(tokens)
-        unique_bigrams = set()
-        for i in range(len(all_tokens) - 1):
-            unique_bigrams.add((all_tokens[i], all_tokens[i+1]))
-        return len(unique_bigrams) / len(all_tokens) if all_tokens else 0
+        try:
+            all_tokens = []
+            for text in texts:
+                tokens = self.tokenizer.tokenize(text)
+                all_tokens.extend(tokens)
+            unique_bigrams = set()
+            for i in range(len(all_tokens) - 1):
+                unique_bigrams.add((all_tokens[i], all_tokens[i+1]))
+            return len(unique_bigrams) / len(all_tokens) if all_tokens else 0
+        except Exception as e:
+            print(f"Diversity evaluation failed: {e}")
+            return 0
 
     # Racial bias
     def evaluate_racial_bias(self, text):
-        results = self.bias_pipeline([text], candidate_labels=["hate speech", "not hate speech"])
-        bias_score = results[0]['scores'][results[0]['labels'].index('hate speech')]
-        return bias_score
+        try:
+            results = self.bias_pipeline([text], candidate_labels=["hate speech", "not hate speech"])
+            bias_score = results[0]['scores'][results[0]['labels'].index('hate speech')]
+            return bias_score
+        except Exception as e:
+            print(f"Bias evaluation failed: {e}")
+            return 0
 
     # METEOR
     def evaluate_meteor(self, candidates, references):
-        meteor_scores = []
-        for ref, cand in zip(references, candidates):
-            ref_tokens = self.tokenizer.tokenize(ref)
-            cand_tokens = self.tokenizer.tokenize(cand)
-            common_tokens = set(ref_tokens) & set(cand_tokens)
-            precision = len(common_tokens) / len(cand_tokens) if cand_tokens else 0
-            recall = len(common_tokens) / len(ref_tokens) if ref_tokens else 0
-            if precision + recall == 0:
-                f_score = 0
-            else:
-                f_score = (10 * precision * recall) / (9 * precision + recall)
-            meteor_scores.append(f_score)
-        return sum(meteor_scores) / len(meteor_scores) if meteor_scores else 0
+        try:
+            meteor_scores = []
+            for ref, cand in zip(references, candidates):
+                ref_tokens = self.tokenizer.tokenize(ref)
+                cand_tokens = self.tokenizer.tokenize(cand)
+                common_tokens = set(ref_tokens) & set(cand_tokens)
+                precision = len(common_tokens) / len(cand_tokens) if cand_tokens else 0
+                recall = len(common_tokens) / len(ref_tokens) if ref_tokens else 0
+                if precision + recall == 0:
+                    f_score = 0
+                else:
+                    f_score = (10 * precision * recall) / (9 * precision + recall)
+                meteor_scores.append(f_score)
+            return sum(meteor_scores) / len(meteor_scores) if meteor_scores else 0
+        except Exception as e:
+            print(f"METEOR evaluation failed: {e}")
+            return 0
 
     # CHRF
     def evaluate_chrf(self, candidates, references):
-        chrf_scores = []
-        for ref, cand in zip(references, candidates):
-            ref_chars = list(ref)
-            cand_chars = list(cand)
-            ref_ngrams = set()
-            cand_ngrams = set()
-            for i in range(len(ref_chars) - 5):
-                ref_ngrams.add(tuple(ref_chars[i:i+6]))
-            for i in range(len(cand_chars) - 5):
-                cand_ngrams.add(tuple(cand_chars[i:i+6]))
-            common_ngrams = ref_ngrams & cand_ngrams
-            precision = len(common_ngrams) / len(cand_ngrams) if cand_ngrams else 0
-            recall = len(common_ngrams) / len(ref_ngrams) if ref_ngrams else 0
-            chrf_score = 2 * precision * recall / (precision + recall) if precision + recall else 0
-            chrf_scores.append(chrf_score)
-        return sum(chrf_scores) / len(chrf_scores) if chrf_scores else 0
+        try:
+            chrf_scores = []
+            for ref, cand in zip(references, candidates):
+                ref_chars = list(ref)
+                cand_chars = list(cand)
+                ref_ngrams = set()
+                cand_ngrams = set()
+                for i in range(len(ref_chars) - 5):
+                    ref_ngrams.add(tuple(ref_chars[i:i+6]))
+                for i in range(len(cand_chars) - 5):
+                    cand_ngrams.add(tuple(cand_chars[i:i+6]))
+                common_ngrams = ref_ngrams & cand_ngrams
+                precision = len(common_ngrams) / len(cand_ngrams) if cand_ngrams else 0
+                recall = len(common_ngrams) / len(ref_ngrams) if ref_ngrams else 0
+                chrf_score = 2 * precision * recall / (precision + recall) if precision + recall else 0
+                chrf_scores.append(chrf_score)
+            return sum(chrf_scores) / len(chrf_scores) if chrf_scores else 0
+        except Exception as e:
+            print(f"CHRF evaluation failed: {e}")
+            return 0
 
     # Readability
     def evaluate_readability(self, text):
-        words = re.findall(r'\b\w+\b', text.lower())
-        sentences = re.split(r'[.!?]+', text)
-        num_words = len(words)
-        num_sentences = len([s for s in sentences if s.strip()])
-        avg_word_length = sum(len(word) for word in words) / num_words if num_words else 0
-        words_per_sentence = num_words / num_sentences if num_sentences else 0
-        flesch_ease = 206.835 - (1.015 * words_per_sentence) - (84.6 * avg_word_length)
-        flesch_grade = (0.39 * words_per_sentence) + (11.8 * avg_word_length) - 15.59
-        return flesch_ease, flesch_grade
+        try:
+            words = re.findall(r'\b\w+\b', text.lower())
+            sentences = re.split(r'[.!?]+', text)
+            num_words = len(words)
+            num_sentences = len([s for s in sentences if s.strip()])
+            avg_word_length = sum(len(word) for word in words) / num_words if num_words else 0
+            words_per_sentence = num_words / num_sentences if num_sentences else 0
+            flesch_ease = 206.835 - (1.015 * words_per_sentence) - (84.6 * avg_word_length)
+            flesch_grade = (0.39 * words_per_sentence) + (11.8 * avg_word_length) - 15.59
+            return flesch_ease, flesch_grade
+        except Exception as e:
+            print(f"Readability evaluation failed: {e}")
+            return 0, 0
 
     # MAUVE
     def evaluate_mauve(self, reference_texts, generated_texts):
-        out = compute_mauve(
-            p_text=reference_texts,
-            q_text=generated_texts,
-            device_id=0,
-            max_text_length=1024,
-            verbose=False
-        )
-        return out.mauve
+        try:
+            out = compute_mauve(
+                p_text=reference_texts,
+                q_text=generated_texts,
+                device_id=0,
+                max_text_length=1024,
+                verbose=False
+            )
+            return out.mauve
+        except Exception as e:
+            print(f"MAUVE evaluation failed: {e}")
+            return 0.0
 
     def evaluate_all(self, question, response, reference):
-        candidates = [response]
-        references = [reference]
-        bleu, rouge1, rouge2, rougeL = self.evaluate_bleu_rouge(candidates, references)
-        bert_p, bert_r, bert_f1 = self.evaluate_bert_score(candidates, references)
-        perplexity = self.evaluate_perplexity(response)
-        diversity = self.evaluate_diversity(candidates)
-        racial_bias = self.evaluate_racial_bias(response)
-        meteor = self.evaluate_meteor(candidates, references)
-        chrf = self.evaluate_chrf(candidates, references)
-        flesch_ease, flesch_grade = self.evaluate_readability(response)
-        mauve_score = self.evaluate_mauve(references, candidates) if len(references) > 1 else 0.0
-        return {
-            "BLEU": bleu,
-            "ROUGE-1": rouge1,
-            "ROUGE-2": rouge2,
-            "ROUGE-L": rougeL,
-            "BERT_Precision": bert_p,
-            "BERT_Recall": bert_r,
-            "BERT_F1": bert_f1,
-            "Perplexity": perplexity,
-            "Diversity": diversity,
-            "Racial_Bias": racial_bias,
-            "MAUVE": mauve_score,
-            "METEOR": meteor,
-            "CHRF": chrf,
-            "Flesch_Reading_Ease": flesch_ease,
-            "Flesch_Kincaid_Grade": flesch_grade,
-        }
+        try:
+            candidates = [response]
+            references = [reference]
+            bleu, rouge1, rouge2, rougeL = self.evaluate_bleu_rouge(candidates, references)
+            bert_p, bert_r, bert_f1 = self.evaluate_bert_score(candidates, references)
+            perplexity = self.evaluate_perplexity(response)
+            diversity = self.evaluate_diversity(candidates)
+            racial_bias = self.evaluate_racial_bias(response)
+            meteor = self.evaluate_meteor(candidates, references)
+            chrf = self.evaluate_chrf(candidates, references)
+            flesch_ease, flesch_grade = self.evaluate_readability(response)
+            mauve_score = self.evaluate_mauve(references, candidates) if len(references) > 1 else 0.0
+            
+            return {
+                "BLEU": bleu,
+                "ROUGE-1": rouge1,
+                "ROUGE-2": rouge2,
+                "ROUGE-L": rougeL,
+                "BERT_Precision": bert_p,
+                "BERT_Recall": bert_r,
+                "BERT_F1": bert_f1,
+                "Perplexity": perplexity,
+                "Diversity": diversity,
+                "Racial_Bias": racial_bias,
+                "MAUVE": mauve_score,
+                "METEOR": meteor,
+                "CHRF": chrf,
+                "Flesch_Reading_Ease": flesch_ease,
+                "Flesch_Kincaid_Grade": flesch_grade,
+            }
+        except Exception as e:
+            print(f"Complete evaluation failed: {e}")
+            return {"error": str(e)}
 
 # Initialize evaluator
 evaluator = RAGEvaluator()
 
 # ------------------ Pinecone ------------------
 def init_pinecone():
-    pc = Pinecone(api_key=pinecone_api_key)
-    return pc.Index("rag-granite-index")
+    try:
+        pc = Pinecone(api_key=pinecone_api_key)
+        return pc.Index("rag-granite-index")
+    except Exception as e:
+        print(f"Pinecone initialization failed: {e}")
+        return None
 
 index = init_pinecone()
 
 # ------------------ Embeddings ------------------
-embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+try:
+    embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+except Exception as e:
+    print(f"Embeddings initialization failed: {e}")
+    embeddings_model = None
 
 def get_retrieved_context(query: str, top_k=3):
-    start = time.time()
-    query_embedding = embeddings_model.embed_query(query)
-    mlflow.log_metric("embedding_latency", time.time() - start)
-    results = index.query(
-        namespace="rag-ns",
-        vector=query_embedding,
-        top_k=top_k,
-        include_metadata=True
-    )
-    mlflow.log_metric("retrieved_chunks", len(results['matches']))
-    return "\n".join([m['metadata']['text'] for m in results['matches']])
+    if not index or not embeddings_model:
+        return "No context available - system initialization failed"
+    
+    try:
+        start = time.time()
+        query_embedding = embeddings_model.embed_query(query)
+        if mlflow.active_run():
+            mlflow.log_metric("embedding_latency", time.time() - start)
+        
+        results = index.query(
+            namespace="rag-ns",
+            vector=query_embedding,
+            top_k=top_k,
+            include_metadata=True
+        )
+        
+        if mlflow.active_run():
+            mlflow.log_metric("retrieved_chunks", len(results['matches']))
+        
+        return "\n".join([m['metadata']['text'] for m in results['matches']])
+    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
 
-    @mlflow.trace
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
-        payload = {"prompt": prompt}
-        start_time = time.time()
-        response = requests.post(self.endpoint_url, json=payload)
-        mlflow.log_metric("lit_serve_latency", time.time() - start_time)
-        if response.status_code == 200:
-            data = response.json()
-            mlflow.log_metric("response_tokens", len(data.get("response", "").split()))
-            return data.get("response", "").strip()
-        else:
-            mlflow.log_metric("request_errors", 1)
-            raise ValueError(f"Request failed: {response.status_code}")
+        try:
+            payload = {"prompt": prompt}
+            start_time = time.time()
+            response = requests.post(self.endpoint_url, json=payload, timeout=30)
+            
+            if mlflow.active_run():
+                mlflow.log_metric("lit_serve_latency", time.time() - start_time)
+            
+            if response.status_code == 200:
+                data = response.json()
+                if mlflow.active_run():
+                    mlflow.log_metric("response_tokens", len(data.get("response", "").split()))
+                return 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}"
+                
+        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."
 
     @property
     def _identifying_params(self) -> Mapping[str, Any]:
@@ -241,50 +323,89 @@ class LitServeLLM(LLM):
     def _llm_type(self) -> str:
         return "litserve_llm"
 
-model = LitServeLLM(endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict")
+# Initialize model with fallback
+try:
+    model = LitServeLLM(endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict")
+except Exception as e:
+    print(f"Model initialization failed: {e}")
+    model = None
 
 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.
+
 Context:
 {context}
+
 Question: {question}
+
 Answer:
 """
 )
 
-llm_chain = LLMChain(llm=model, prompt=prompt)
+# Initialize chain with error handling
+try:
+    if model:
+        llm_chain = LLMChain(llm=model, prompt=prompt)
+    else:
+        llm_chain = None
+        print("LLM chain not initialized due to model failure")
+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 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
+    """Get the complete RAG response with error handling"""
+    try:
+        # Get context
+        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."
+            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"
 
 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 words:
-        current_text += word + " "
-        yield current_text
-        time.sleep(0.05)  # Adjust speed as needed
+    try:
+        full_response, _ = get_rag_response(question)
+        
+        # Stream word by word
+        words = full_response.split()
+        current_text = ""
+        for word in words:
+            current_text += word + " "
+            yield current_text
+            time.sleep(0.05)  # Adjust speed as needed
+            
+    except Exception as e:
+        error_msg = f"Error in streaming: {str(e)}"
+        print(f"Streaming error: {error_msg}")
+        yield "I apologize, but I encountered an error while generating the response."
 
 # ------------------ Gradio UI ------------------
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+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*
@@ -292,7 +413,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     usage_counter = gr.State(value=0)
     session_start = gr.State(value=datetime.now().isoformat())
-    current_response = gr.State(value="")  # Store current response for evaluation
+    current_response = gr.State(value="")
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -328,49 +449,68 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     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)
-            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)
-            
-            # Get response and context
-            response, context = get_rag_response(question)
+        if not question.strip():
+            return "Please enter a question.", count, session_start, ""
             
-            # Log response metrics
-            mlflow.log_metric("response_length", len(response))
-            mlflow.log_metric("response_tokens", len(response.split()))
-            
-            return response, count, session_start, response
+        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:
+                    mlflow.log_param("question", question)
+                    mlflow.log_param("session_start", session_start)
+                    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)
+                
+                # Get response and context
+                response, context = get_rag_response(question)
+                
+                if mlflow_tracking_uri:
+                    mlflow.log_metric("response_length", len(response))
+                    mlflow.log_metric("response_tokens", len(response.split()))
+                
+                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
 
     def evaluate_response(question, response):
         """Evaluate the response and return metrics"""
-        if not question or not response:
-            return gr.update(value={}, visible=False)
+        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)
             
         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)
+            # 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}")
-            return gr.update(value={"error": str(e)}, visible=True)
+            return gr.update(value={"error": f"Evaluation failed: {str(e)}"}, visible=True)
 
     def clear_chat():
         """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
@@ -399,8 +539,16 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
     # Feedback handling
+    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
+
     feedback.change(
-        fn=lambda feedback_val: mlflow.log_metric("user_feedback_score", 1 if feedback_val == "Helpful" else 0),
+        fn=handle_feedback,
         inputs=[feedback],
         outputs=[]
     )
@@ -409,6 +557,6 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=True,
+        share=False,  # Disable sharing to avoid the warning
         show_error=True
     )