Update rag.py

rag.py

@@ -1,13 +1,15 @@
 import json
-from sentence_transformers import SentenceTransformer, util
-from groq import Groq
-from datetime import datetime
+import glob
 import os
+import random
 import pandas as pd
-from datasets import load_dataset, Dataset
+from datetime import datetime
 from dotenv import load_dotenv
-import random
-import glob
+
+# Core AI Libraries
+from sentence_transformers import SentenceTransformer, util
+from groq import Groq
+from datasets import load_dataset, Dataset
 
 # Load environment variables
 load_dotenv()
@@ -25,16 +27,7 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 # Greeting list
 GREETINGS = [
     "hi", "hello", "hey", "good morning", "good afternoon", "good evening",
-    "assalam o alaikum", "salam", "aoa", "hi there",
-    "hey there", "greetings"
-]
-
-# Fixed rephrased unmatched query responses
-UNMATCHED_RESPONSES = [
-    "Thank you for your query. We’ve forwarded it to our support team and it will be added soon. In the meantime, you can visit the University of Education official website or reach out via the contact details below.\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk",
-    "We’ve noted your question and it’s in queue for inclusion. For now, please check the University of Education website or contact the administration directly.\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk",
-    "Your query has been recorded. We’ll update the system with relevant information shortly. Meanwhile, you can visit UE's official site or reach out using the details below:\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk",
-    "We appreciate your question. It has been forwarded for further processing. Until it’s available here, feel free to visit the official UE website or use the contact options:\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk"
+    "assalam o alaikum", "salam", "aoa", "hi there", "hey there", "greetings"
 ]
 
 # Load multiple JSON datasets
@@ -48,10 +41,6 @@ try:
                 for item in data:
                     if isinstance(item, dict) and 'Question' in item and 'Answer' in item:
                         dataset.append(item)
-                    else:
-                        print(f"Invalid entry in {file_path}: {item}")
-            else:
-                print(f"File {file_path} does not contain a list.")
 except Exception as e:
     print(f"Error loading datasets: {e}")
 
@@ -60,8 +49,8 @@ dataset_questions = [item.get("Question", "").lower().strip() for item in datase
 dataset_answers = [item.get("Answer", "") for item in dataset]
 dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)
 
-# Save unmatched queries to Hugging Face
 def manage_unmatched_queries(query: str):
+    """Logs unknown queries to Hugging Face Hub."""
     try:
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         try:
@@ -69,6 +58,7 @@ def manage_unmatched_queries(query: str):
             df = ds["train"].to_pandas()
         except:
             df = pd.DataFrame(columns=["Query", "Timestamp", "Processed"])
+        
         if query not in df["Query"].values:
             new_entry = {"Query": query, "Timestamp": timestamp, "Processed": False}
             df = pd.concat([df, pd.DataFrame([new_entry])], ignore_index=True)
@@ -77,67 +67,65 @@ def manage_unmatched_queries(query: str):
     except Exception as e:
         print(f"Failed to save query: {e}")
 
-# Query Groq LLM
-def query_groq_llm(prompt, model_name="llama3-70b-8192"):
+def query_groq_llm(prompt, system_message):
+    """Utility to call Groq Llama 3 API."""
     try:
         chat_completion = groq_client.chat.completions.create(
-            messages=[{
-                "role": "user",
-                "content": prompt
-            }],
-            model=model_name,
-            temperature=0.7,
-            max_tokens=500
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": prompt}
+            ],
+            model="llama3-70b-8192",
+            temperature=0.6,
+            max_tokens=800
         )
         return chat_completion.choices[0].message.content.strip()
     except Exception as e:
         print(f"Error querying Groq API: {e}")
         return ""
 
-# Main logic function to be called from Gradio
 def get_best_answer(user_input):
     if not user_input.strip():
         return "Please enter a valid question."
 
     user_input_lower = user_input.lower().strip()
 
+    # 1. Length/Greeting Check
     if len(user_input_lower.split()) < 3 and not any(greet in user_input_lower for greet in GREETINGS):
-        return "Please ask your question properly with at least 3 words."
+        return "Please ask your question properly with at least 3 words so I can assist you better."
 
+    # 2. Hardcoded Keyword Check (Fees)
     if any(keyword in user_input_lower for keyword in ["fee structure", "fees structure", "semester fees", "semester fee"]):
         return (
-            "💰 For the most complete and up-to-date fee details for your program at the University of Education Lahore, please visit the official fee structure page.\n"
-            "This webpage offers a detailed overview of the fee structure, providing you with essential information to support your academic journey at our institution.\n"
-            "🔗 https://drive.google.com/file/d/1B30FKoP6GrkS9pQk10PWKCwcjco5E9Cc/view"
+            "💰 **Official Fee Information**\n\n"
+            "For the most complete and up-to-date fee details for your program at the University of Education Lahore, please visit the official link:\n"
+            "🔗 [View Fee Structure](https://drive.google.com/file/d/1B30FKoP6GrkS9pQk10PWKCwcjco5E9Cc/view)"
         )
 
+    # 3. Vector Similarity Search
     user_embedding = similarity_model.encode(user_input_lower, convert_to_tensor=True)
     similarities = util.pytorch_cos_sim(user_embedding, dataset_embeddings)[0]
     best_match_idx = similarities.argmax().item()
     best_score = similarities[best_match_idx].item()
 
-    if best_score < 0.65:
+    # 4. DECISION LOGIC
+    if best_score >= 0.65:
+        # PATH A: High similarity (Rephrase Dataset)
+        original_answer = dataset_answers[best_match_idx]
+        system_message = "You are the UOE AI Assistant. Your job is to rephrase the provided official answer to make it more attractive, using bold text, bullet points, and clear headings. Do NOT change the facts."
+        prompt = f"User Question: {user_input}\nOfficial Context: {original_answer}\n\nRephrase this beautifully:"
+    else:
+        # PATH B: Low similarity (LLM General Knowledge + Logging)
         manage_unmatched_queries(user_input)
-        return random.choice(UNMATCHED_RESPONSES)
-
-    original_answer = dataset_answers[best_match_idx]
-    prompt = f"""Name is UOE AI Assistant! You are an official assistant for the University of Education Lahore.
-Rephrase the following official answer clearly and professionally. 
-Use structured formatting (like headings, bullet points, or numbered lists) where appropriate.
-DO NOT add any new or extra information. ONLY rephrase and improve the clarity and formatting of the original answer.
-### Question:
-{user_input}
-### Original Answer:
-{original_answer}
-### Rephrased Answer:
-"""
-
-    llm_response = query_groq_llm(prompt)
+        system_message = "You are the UOE AI Assistant for the University of Education Lahore. Answer the user's question based on your general knowledge of the university. If you are unsure about specific dates or costs, tell them to contact info@ue.edu.pk."
+        prompt = f"The following question about University of Education (UE) was asked: {user_input}\n\nPlease provide a helpful and professional response:"
+
+    # 5. Execute LLM Call
+    llm_response = query_groq_llm(prompt, system_message)
 
+    # 6. Fallback Strategy
     if llm_response:
-        for marker in ["Improved Answer:", "Official Answer:", "Rephrased Answer:"]:
-            if marker in llm_response:
-                return llm_response.split(marker)[-1].strip()
         return llm_response
     else:
-        return dataset_answers[best_match_idx]
+        # If LLM fails, return raw dataset answer if score was high, else return contact info
+        return dataset_answers[best_match_idx] if best_score >= 0.65 else "I'm having trouble connecting to the server. Please contact info@ue.edu.pk."