Update rag.py

rag.py

@@ -18,6 +18,10 @@ groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 # Load models and dataset
 similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
 
+# Configuration
+HF_DATASET_REPO = "midrees2806/unmatched_queries"  # Your dataset repo
+HF_TOKEN = os.getenv("HF_TOKEN")  # From Space secrets
+
 # Load dataset (automatically using the path)
 with open('dataset.json', 'r') as f:
     dataset = json.load(f)
@@ -27,6 +31,32 @@ 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)
 
+# --- 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}")
+
+# --- Enhanced LLM Query ---        
+
 def query_groq_llm(prompt, model_name="llama3-70b-8192"):
     try:
         chat_completion = groq_client.chat.completions.create(