Update rag.py

rag.py

@@ -1,48 +1,50 @@
+import os
 import json
+import requests
+import pandas as pd
+from dotenv import load_dotenv
+from datetime import datetime
 from sentence_transformers import SentenceTransformer, util
 from groq import Groq
-from datetime import datetime
-import os
-import pandas as pd
 from datasets import load_dataset, Dataset
-from dotenv import load_dotenv
 
-# Load environment variables
+# ✅ Load environment variables from .env
 load_dotenv()
 
-# Initialize Groq client
-groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+# ✅ API Keys
+HF_TOKEN = os.getenv("HF_TOKEN")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
-# Load similarity model
-similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
+# ✅ Initialize Groq client
+groq_client = Groq(api_key=GROQ_API_KEY)
 
-# Config
+# ✅ Hugging Face Dataset Repo
 HF_DATASET_REPO = "midrees2806/unmatched_queries"
-HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Greeting list
+# ✅ Sentence Transformer model for semantic similarity
+similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
+
+# ✅ Greeting keywords
 GREETINGS = [
     "hi", "hello", "hey", "good morning", "good afternoon", "good evening",
-    "assalam o alaikum", "salam", "aoa", "hi there",
-    "hey there", "greetings"
+    "assalam o alaikum", "salam", "aoa", "hi there", "hey there", "greetings"
 ]
 
-# Load local dataset
+# ✅ Load dataset
 try:
     with open('dataset.json', 'r') as f:
         dataset = json.load(f)
-    if not all(isinstance(item, dict) and 'input' in item and 'response' in item for item in dataset):
-        raise ValueError("Invalid dataset structure")
+    assert all('input' in d and 'response' in d for d in dataset), "Invalid dataset format"
 except Exception as e:
-    print(f"Error loading dataset: {e}")
+    print(f"[ERROR] Loading dataset: {e}")
     dataset = []
 
-# Precompute embeddings
-dataset_questions = [item.get("input", "").lower().strip() for item in dataset]
-dataset_answers = [item.get("response", "") for item in dataset]
+# ✅ Prepare embeddings
+dataset_questions = [d["input"].lower().strip() for d in dataset]
+dataset_answers = [d["response"] for d in dataset]
 dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)
 
-# Save unmatched queries to Hugging Face
+# ✅ Function: Save unmatched queries to Hugging Face Hub
 def manage_unmatched_queries(query: str):
     try:
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
@@ -57,95 +59,93 @@ def manage_unmatched_queries(query: str):
             updated_ds = Dataset.from_pandas(df)
             updated_ds.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN)
     except Exception as e:
-        print(f"Failed to save query: {e}")
+        print(f"[ERROR] Logging unmatched query: {e}")
 
-# Query Groq LLM
-def query_groq_llm(prompt, model_name="llama3-70b-8192"):
+# ✅ Function: Call Groq LLM
+def query_groq_llm(prompt: str, model_name="llama3-70b-8192") -> str:
     try:
-        chat_completion = groq_client.chat.completions.create(
-            messages=[{
-                "role": "user",
-                "content": prompt
-            }],
+        completion = groq_client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
             model=model_name,
             temperature=0.7,
             max_tokens=500
         )
-        return chat_completion.choices[0].message.content.strip()
+        return completion.choices[0].message.content.strip()
     except Exception as e:
-        print(f"Error querying Groq API: {e}")
+        print(f"[ERROR] Groq LLM call failed: {e}")
         return ""
 
-# Main logic function to be called from Gradio
-def get_best_answer(user_input):
+# ✅ Main RAG logic
+def get_best_answer(user_input: str) -> str:
     if not user_input.strip():
-        return "Please enter a valid question."
+        return "⚠️ Please enter a valid question."
 
     user_input_lower = user_input.lower().strip()
 
-    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."
+    # Handle short or vague questions
+    if len(user_input_lower.split()) < 3 and not any(g in user_input_lower for g in GREETINGS):
+        return "🔎 Please provide more details or ask a complete question (at least 3 words)."
 
+    # Handle greetings
     if any(greet in user_input_lower for greet in GREETINGS):
-        greeting_response = query_groq_llm(
-            f"You are an official assistant for University of Education Lahore. "
-            f"Respond to this greeting in a friendly and professional manner: {user_input}"
-        )
-        return greeting_response if greeting_response else "Hello! How can I assist you today?"
+        prompt = f"You are an official assistant for University of Education Lahore. Respond to this greeting in a professional and friendly tone: {user_input}"
+        return query_groq_llm(prompt) or "👋 Hello! How can I assist you today?"
 
+    # Handle direct FAQ (e.g., fee structure)
     if any(keyword in user_input_lower for keyword in ["fee structure", "fees structure", "semester fees", "semester fee"]):
         return (
-            "💰 For complete and up-to-date fee details for this program, we recommend visiting the official University of Education fee structure page.\n"
-            "You'll find comprehensive information regarding tuition, admission charges, and other applicable fees there.\n"
-            "🔗 https://ue.edu.pk/allfeestructure.php"
+            "💰 For complete and up-to-date fee details for this program, please visit:\n"
+            "🔗 https://ue.edu.pk/allfeestructure.php\n"
+            "It contains all relevant information including tuition, admission, and semester-wise fees."
         )
 
+    # Semantic search for best matching question
     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:
-        manage_unmatched_queries(user_input)
+    # Threshold to determine match quality
+    SIMILARITY_THRESHOLD = 0.65
 
-    if best_score >= 0.65:
+    if best_score >= SIMILARITY_THRESHOLD:
         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.
-
+        prompt = f"""You are UOE AI Assistant! As an official assistant for University of Education Lahore, rephrase the following answer clearly and professionally.
+Use bullet points or headings if helpful. Do NOT add extra information.
 ### Question:
 {user_input}
-
 ### Original Answer:
 {original_answer}
-
-### Rephrased Answer:
-"""
+### Rephrased Answer:"""
     else:
-        prompt = f"""Name is UOE AI Assistant! As an official assistant for University of Education Lahore, provide a helpful response:
-Include relevant details about university policies.
-If unsure, direct to official channels.
-
+        manage_unmatched_queries(user_input)
+        prompt = f"""You are UOE AI Assistant. As an official assistant for University of Education Lahore, provide a helpful response to this query.
+If unsure, direct the user to the official university contact options.
 ### Question:
 {user_input}
+### Official Answer:"""
 
-### Official Answer:
-"""
+    # Get the response from LLM
+    response = query_groq_llm(prompt)
 
-    llm_response = query_groq_llm(prompt)
-
-    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
+    if response:
+        for marker in ["Rephrased Answer:", "Official Answer:", "Improved Answer:"]:
+            if marker in response:
+                return response.split(marker)[-1].strip()
+        return response  # if no marker found
     else:
-        return dataset_answers[best_match_idx] if best_score >= 0.65 else (
+        return dataset_answers[best_match_idx] if best_score >= SIMILARITY_THRESHOLD else (
             "For official information:\n"
             "📞 +92-42-99262231-33\n"
             "✉️ [email protected]\n"
             "🌐 https://ue.edu.pk"
         )
+
+# ✅ Example (for direct testing)
+if __name__ == "__main__":
+    while True:
+        user_input = input("\n🧑‍🎓 You: ")
+        if user_input.lower() in ["exit", "quit"]:
+            break
+        answer = get_best_answer(user_input)
+        print(f"\n🤖 UOE Assistant:\n{answer}")