Update rag.py

rag.py

@@ -1,30 +1,30 @@
 import json
-import csv
-from pathlib import Path
 from datetime import datetime
 from sentence_transformers import SentenceTransformer, util
 from groq import Groq
 from dotenv import load_dotenv
 import os
+from datasets import load_dataset, Dataset, DatasetDict
 import pandas as pd
 
 # Load environment variables
 load_dotenv()
 
-# Initialize Groq client
+# Initialize clients
 groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
-
-# Load models and dataset
 similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
 
-# Load dataset
+# Configuration
+HF_DATASET_REPO = "midrees2806/unmatched_queries"  # Your dataset repo
+HF_TOKEN = os.getenv("HF_TOKEN")  # From Space secrets
+
+# --- Dataset Loading ---
 try:
     with open('dataset.json', 'r') as f:
         dataset = json.load(f)
-    # Validate dataset structure
     if not all(isinstance(item, dict) and 'input' in item and 'response' in item for item in dataset):
         raise ValueError("Invalid dataset structure")
-except (json.JSONDecodeError, ValueError, FileNotFoundError) as e:
+except Exception as e:
     print(f"Error loading dataset: {e}")
     dataset = []
 
@@ -33,88 +33,82 @@ dataset_questions = [item.get("input", "").lower().strip() for item in dataset]
 dataset_answers = [item.get("response", "") for item in dataset]
 dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)
 
-# Initialize unmatched queries CSV
-def init_unmatched_queries_file():
-    csv_file = Path('unmatched_queries.csv')
-    if not csv_file.exists():
-        with open(csv_file, 'w', newline='', encoding='utf-8') as f:
-            writer = csv.writer(f)
-            writer.writerow(['Unmatched Queries', 'Timestamp'])
-
-init_unmatched_queries_file()
+# --- Unmatched Queries Handler ---
+def manage_unmatched_queries(query: str):
+    """Save unmatched queries to HF Dataset with error handling"""
+    try:
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        
+        # Load existing dataset or create new
+        try:
+            ds = load_dataset(HF_DATASET_REPO, token=HF_TOKEN)
+            df = ds["train"].to_pandas()
+        except:
+            df = pd.DataFrame(columns=["Query", "Timestamp", "Processed"])
+        
+        # Append new query (avoid duplicates)
+        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)
+            
+            # Push to Hub
+            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}")
 
-def query_groq_llm(prompt, model_name="llama3-70b-8192"):
+# --- Enhanced LLM Query ---
+def query_llm(prompt: str, model: str = "llama3-70b-8192") -> str:
     try:
-        chat_completion = groq_client.chat.completions.create(
-            messages=[{
-                "role": "user",
-                "content": prompt
-            }],
-            model=model_name,
+        response = groq_client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=model,
             temperature=0.7,
-            max_tokens=500
+            max_tokens=1024,
+            top_p=0.9
         )
-        return chat_completion.choices[0].message.content.strip()
+        return response.choices[0].message.content.strip()
     except Exception as e:
-        print(f"Error querying Groq API: {e}")
+        print(f"LLM Error: {e}")
         return None
 
-def save_unmatched_query(query):
-    try:
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        with open('unmatched_queries.csv', 'a', newline='', encoding='utf-8') as f:
-            writer = csv.writer(f)
-            writer.writerow([query, timestamp])
-    except Exception as e:
-        print(f"Error saving unmatched query: {e}")
-
-def get_best_answer(user_input):
-    user_input_lower = user_input.lower().strip()
+# --- Main Chat Function ---
+def get_best_answer(user_input: str) -> str:
+    user_input = user_input.strip()
+    lower_input = user_input.lower()
 
-    # Handle fee-related questions
-    if any(keyword in user_input_lower for keyword in ["fee", "fees", "charges", "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"
-        )
+    # 1. Handle special cases
+    if any(kw in lower_input for kw in ["fee", "fees", "tuition"]):
+        return ("💰 Fee information:\n"
+                "Please visit: https://ue.edu.pk/allfeestructure.php\n"
+                "For personalized help, contact [email protected]")
 
-    # Similarity matching
-    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()
+    # 2. Semantic similarity search
+    query_embedding = similarity_model.encode(lower_input, convert_to_tensor=True)
+    scores = util.pytorch_cos_sim(query_embedding, dataset_embeddings)[0]
+    best_idx = scores.argmax().item()
+    best_score = scores[best_idx].item()
 
-    # Save unmatched queries
+    # 3. Save unmatched queries (threshold = 0.65)
     if best_score < 0.65:
-        save_unmatched_query(user_input)
+        manage_unmatched_queries(user_input)
 
+    # 4. Generate response
     if best_score >= 0.65:
-        original_answer = dataset_answers[best_match_idx]
-        prompt = f"""As an official assistant for University of Education Lahore, provide a clear response:
+        context = dataset_answers[best_idx]
+        prompt = f"""University Assistant Task:
         Question: {user_input}
-        Original Answer: {original_answer}
-        Improved Answer:"""
+        Context: {context}
+        Generate a helpful, accurate response using the context. If unsure, say "Please contact [email protected]" """
     else:
-        prompt = f"""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.
+        prompt = f"""As an official University of Education assistant, answer:
         Question: {user_input}
-        Official Answer:"""
-
-    llm_response = query_groq_llm(prompt)
-
-    if llm_response:
-        for marker in ["Improved Answer:", "Official Answer:"]:
-            if marker in llm_response:
-                response = llm_response.split(marker)[-1].strip()
-                break
-        else:
-            response = llm_response
-    else:
-        response = dataset_answers[best_match_idx] if best_score >= 0.65 else """For official information:
-        📞 +92-42-99262231-33
-        ✉️ [email protected]
-        🌐 ue.edu.pk"""
-
-    return response
+        Guidelines:
+        - Be polite and professional
+        - Direct to official channels if uncertain
+        - Keep responses under 3 sentences"""
+
+    response = query_llm(prompt)
+    return response or """For official assistance:
+    📞 +92-42-99262231-33
+    ✉️ [email protected]"""