Sharada25/dhammaai
1
1"""2MongoDB Database Helper for Vipassana Guide AI3Handles all database operations for chat history and reinforcement learning4"""5 6from pymongo import MongoClient7from pymongo.errors import ConnectionFailure, OperationFailure8import os9from dotenv import load_dotenv10from datetime import datetime11from typing import List, Dict, Optional12from bson.objectid import ObjectId13from urllib.parse import quote_plus14import ssl15import certifi16 17load_dotenv()18 19class MongoDBManager:20 """Manages MongoDB Atlas connections and operations"""21 22 def __init__(self):23 """Initialize MongoDB connection"""24 self.client = None25 self.db = None26 self._connect()27 28 def _connect(self):29 """Connect to MongoDB Atlas"""30 try:31 # Get MongoDB connection string from environment32 mongo_uri = os.getenv('MONGODB_URI')33 34 if not mongo_uri:35 print("[ERROR] MONGODB_URI not found in environment variables")36 print("[INFO] Make sure to set MONGODB_URI in HuggingFace Spaces secrets")37 raise ValueError("MONGODB_URI is required")38 39 # Check if URI needs manual encoding (if it contains @ in password)40 # Format should be: mongodb+srv://username:password@host/41 print(f"[INFO] Processing MongoDB URI...")42 43 # If the URI doesn't already have %40 encoded, it might fail44 # Try to use it as-is first, pymongo should handle it45 mongo_uri_to_use = mongo_uri46 47 # Log connection attempt (without showing password)48 if '://' in mongo_uri:49 display_uri = mongo_uri.split('://')[0] + '://***:***@' + mongo_uri.split('@')[-1]50 print(f"[INFO] MongoDB connection string (censored): {display_uri}")51 52 # Connect to MongoDB Atlas with proper error handling53 print(f"[INFO] Connecting to MongoDB Atlas (timeout: 20s)...")54 55 # MongoDB Atlas mongodb+srv:// URI automatically enables TLS56 # Using tlsAllowInvalidCertificates=True to bypass certificate verification in restricted Docker environments57 # This is acceptable for HF Spaces since the connection is still encrypted via TLS58 print(f"[INFO] Connecting with TLS (certificate verification disabled)...")59 self.client = MongoClient(60 mongo_uri_to_use,61 serverSelectionTimeoutMS=20000, # Increased to 20s62 connectTimeoutMS=20000,63 socketTimeoutMS=20000,64 retryWrites=False, # Disable retry writes in case of connection issues65 tls=True, # Explicitly enable TLS66 tlsAllowInvalidCertificates=True, # Allow invalid/self-signed certificates67 tlsAllowInvalidHostnames=True # Allow hostname mismatches68 )69 70 # Test connection with proper error messaging71 print(f"[INFO] Testing MongoDB connection with ping command...")72 self.client.admin.command('ping')73 print(f"[OK] ✓ MongoDB connection test successful!")74 75 # Get database name from env or use default76 db_name = os.getenv('MONGODB_DATABASE', 'vipassana_chat')77 self.db = self.client[db_name]78 79 print(f"[OK] ✓ MongoDB Atlas connected successfully to database: {db_name}")80 81 # Create indexes for better performance82 self._create_indexes()83 84 except ConnectionFailure as e:85 print(f"[ERROR] Connection failed to MongoDB Atlas")86 print(f"[ERROR] This is likely a network issue on your current machine (DNS/firewall)")87 print(f"[ERROR] But the app will still work - MongoDB is optional")88 print(f"[ERROR] Error: {str(e)[:200]}")89 print(f"[WARNING] Continuing without MongoDB - chat history will use local fallback")90 except ValueError as e:91 print(f"[ERROR] ✗ Configuration error: {e}")92 print(f"[ERROR] Make sure MONGODB_URI is set in environment variables")93 raise94 except Exception as e:95 print(f"[ERROR] ✗ Unexpected error connecting to MongoDB")96 print(f"[ERROR] Error type: {type(e).__name__}")97 print(f"[ERROR] Error message: {str(e)}")98 print(f"[ERROR] Full traceback: {repr(e)}")99 raise100 101 def _create_indexes(self):102 """Create indexes for better query performance"""103 try:104 # Conversations collection indexes105 self.db.conversations.create_index("created_at")106 self.db.conversations.create_index("updated_at")107 108 # Messages collection indexes109 self.db.messages.create_index("conversation_id")110 self.db.messages.create_index("timestamp")111 self.db.messages.create_index("feedback")112 113 # RAG performance collection indexes114 self.db.rag_performance.create_index("chunk_id")115 self.db.rag_performance.create_index("feedback_score")116 117 except Exception as e:118 print(f"Warning: Could not create indexes: {e}")119 120 # ===== CONVERSATION OPERATIONS =====121 122 def save_conversation(self, conversation_id: str, title: str, user_id: str = None) -> bool:123 """124 Create or update a conversation125 Returns True if successful126 """127 try:128 conversation = {129 "_id": conversation_id,130 "user_id": user_id,131 "title": title,132 "created_at": datetime.utcnow(),133 "updated_at": datetime.utcnow()134 }135 136 # Upsert: update if exists, insert if not137 self.db.conversations.update_one(138 {"_id": conversation_id},139 {140 "$set": {"title": title, "updated_at": datetime.utcnow()},141 "$setOnInsert": {"created_at": datetime.utcnow(), "user_id": user_id}142 },143 upsert=True144 )145 return True146 except Exception as e:147 print(f"Error saving conversation: {e}")148 return False149 150 def get_conversation(self, conversation_id: str) -> Optional[Dict]:151 """Get a conversation by ID with all its messages"""152 try:153 # Get conversation details154 conversation = self.db.conversations.find_one({"_id": conversation_id})155 156 if not conversation:157 return None158 159 # Get all messages for this conversation160 messages = list(self.db.messages.find(161 {"conversation_id": conversation_id}162 ).sort("timestamp", 1))163 164 # Convert ObjectId to string for JSON serialization165 for msg in messages:166 if "_id" in msg:167 msg["id"] = str(msg["_id"])168 del msg["_id"]169 170 conversation["messages"] = messages171 return conversation172 173 except Exception as e:174 print(f"Error getting conversation: {e}")175 return None176 177 def get_all_conversations(self, user_id: str = None, limit: int = 100) -> List[Dict]:178 """Get all conversations, optionally filtered by user_id"""179 try:180 query = {}181 if user_id:182 query["user_id"] = user_id183 184 conversations = list(self.db.conversations.find(query)185 .sort("updated_at", -1)186 .limit(limit))187 188 return conversations189 190 except Exception as e:191 print(f"Error getting conversations: {e}")192 return []193 194 def delete_conversation(self, conversation_id: str) -> bool:195 """Delete a conversation and all its messages"""196 try:197 # Delete all messages198 self.db.messages.delete_many({"conversation_id": conversation_id})199 200 # Delete conversation201 self.db.conversations.delete_one({"_id": conversation_id})202 203 return True204 except Exception as e:205 print(f"Error deleting conversation: {e}")206 return False207 208 # ===== MESSAGE OPERATIONS =====209 210 def save_message(self, conversation_id: str, role: str, content: str,211 rag_chunks: Optional[List[Dict]] = None) -> Optional[str]:212 """213 Save a message to the database214 Returns the message ID if successful215 """216 try:217 message = {218 "conversation_id": conversation_id,219 "role": role,220 "content": content,221 "timestamp": datetime.utcnow(),222 "feedback": "neutral",223 "rating": None,224 "feedback_comment": None,225 "rag_chunks_used": rag_chunks226 }227 228 result = self.db.messages.insert_one(message)229 230 # Update conversation's updated_at231 self.db.conversations.update_one(232 {"_id": conversation_id},233 {"$set": {"updated_at": datetime.utcnow()}}234 )235 236 return str(result.inserted_id)237 238 except Exception as e:239 print(f"Error saving message: {e}")240 return None241 242 def update_message_feedback(self, message_id: str, feedback: str,243 rating: Optional[int] = None,244 comment: Optional[str] = None) -> bool:245 """246 Update feedback for a message247 feedback: 'helpful', 'not_helpful', or 'neutral'248 rating: 1-5 stars (optional)249 """250 try:251 update_data = {252 "feedback": feedback,253 "rating": rating,254 "feedback_comment": comment255 }256 257 self.db.messages.update_one(258 {"_id": ObjectId(message_id)},259 {"$set": update_data}260 )261 262 # Also update RAG performance tracking263 if feedback in ['helpful', 'not_helpful']:264 self._update_rag_performance(message_id, feedback == 'helpful', rating)265 266 return True267 268 except Exception as e:269 print(f"Error updating message feedback: {e}")270 return False271 272 def _update_rag_performance(self, message_id: str, was_helpful: bool, rating: Optional[int]):273 """Track RAG chunk performance for reinforcement learning"""274 try:275 # Get the message with RAG chunks276 message = self.db.messages.find_one({"_id": ObjectId(message_id)})277 278 if not message or not message.get('rag_chunks_used'):279 return280 281 chunks = message['rag_chunks_used']282 283 # Calculate feedback score284 feedback_score = 1.0 if was_helpful else -0.5285 if rating:286 feedback_score *= (rating / 3.0)287 288 # Insert performance records for each chunk289 for chunk in chunks:290 perf_record = {291 "chunk_id": chunk.get('id', ''),292 "chunk_source": chunk.get('source', ''),293 "query": message['content'][:500],294 "message_id": message_id,295 "feedback_score": feedback_score,296 "was_helpful": was_helpful,297 "created_at": datetime.utcnow()298 }299 300 self.db.rag_performance.insert_one(perf_record)301 302 except Exception as e:303 print(f"Error updating RAG performance: {e}")304 305 # ===== REINFORCEMENT LEARNING QUERIES =====306 307 def get_chunk_performance(self, min_uses: int = 5) -> List[Dict]:308 """309 Get performance statistics for RAG chunks310 Used for reinforcement learning optimization311 """312 try:313 pipeline = [314 {315 "$group": {316 "_id": {317 "chunk_id": "$chunk_id",318 "chunk_source": "$chunk_source"319 },320 "times_used": {"$sum": 1},321 "avg_feedback_score": {"$avg": "$feedback_score"},322 "helpful_count": {323 "$sum": {"$cond": ["$was_helpful", 1, 0]}324 },325 "not_helpful_count": {326 "$sum": {"$cond": ["$was_helpful", 0, 1]}327 }328 }329 },330 {331 "$match": {332 "times_used": {"$gte": min_uses}333 }334 },335 {336 "$project": {337 "chunk_id": "$_id.chunk_id",338 "chunk_source": "$_id.chunk_source",339 "times_used": 1,340 "avg_feedback_score": 1,341 "helpful_count": 1,342 "not_helpful_count": 1,343 "success_rate": {344 "$multiply": [345 {"$divide": ["$helpful_count", "$times_used"]},346 100347 ]348 }349 }350 },351 {352 "$sort": {"avg_feedback_score": -1}353 }354 ]355 356 results = list(self.db.rag_performance.aggregate(pipeline))357 return results358 359 except Exception as e:360 print(f"Error getting chunk performance: {e}")361 return []362 363 def get_low_rated_conversations(self, threshold: float = 3.0, limit: int = 20) -> List[Dict]:364 """365 Get conversations with low ratings for analysis366 """367 try:368 pipeline = [369 {370 "$group": {371 "_id": "$conversation_id",372 "avg_rating": {"$avg": "$rating"},373 "message_count": {"$sum": 1},374 "helpful_count": {375 "$sum": {"$cond": [{"$eq": ["$feedback", "helpful"]}, 1, 0]}376 },377 "not_helpful_count": {378 "$sum": {"$cond": [{"$eq": ["$feedback", "not_helpful"]}, 1, 0]}379 }380 }381 },382 {383 "$match": {384 "avg_rating": {"$lt": threshold, "$ne": None}385 }386 },387 {388 "$sort": {"avg_rating": 1}389 },390 {391 "$limit": limit392 }393 ]394 395 results = list(self.db.messages.aggregate(pipeline))396 397 # Get conversation titles398 for result in results:399 conv = self.db.conversations.find_one({"_id": result["_id"]})400 result["title"] = conv.get("title", "Untitled") if conv else "Untitled"401 result["conversation_id"] = result["_id"]402 403 return results404 405 except Exception as e:406 print(f"Error getting low rated conversations: {e}")407 return []408 409 def get_analytics_summary(self) -> Dict:410 """Get overall analytics summary for dashboard"""411 try:412 # Total conversations413 total_conversations = self.db.conversations.count_documents({})414 415 # Total messages416 total_messages = self.db.messages.count_documents({})417 418 # Average rating419 pipeline = [420 {"$match": {"rating": {"$ne": None}}},421 {"$group": {"_id": None, "avg_rating": {"$avg": "$rating"}}}422 ]423 rating_result = list(self.db.messages.aggregate(pipeline))424 overall_avg_rating = rating_result[0]["avg_rating"] if rating_result else None425 426 # Feedback counts427 total_helpful = self.db.messages.count_documents({"feedback": "helpful"})428 total_not_helpful = self.db.messages.count_documents({"feedback": "not_helpful"})429 430 return {431 "total_conversations": total_conversations,432 "total_messages": total_messages,433 "overall_avg_rating": overall_avg_rating,434 "total_helpful": total_helpful,435 "total_not_helpful": total_not_helpful436 }437 438 except Exception as e:439 print(f"Error getting analytics summary: {e}")440 return {}441 442 # ===== EXPORT FUNCTIONS =====443 444 def export_to_dict(self) -> List[Dict]:445 """Export all data for Excel/CSV export"""446 try:447 pipeline = [448 {449 "$lookup": {450 "from": "conversations",451 "localField": "conversation_id",452 "foreignField": "_id",453 "as": "conversation"454 }455 },456 {457 "$unwind": {458 "path": "$conversation",459 "preserveNullAndEmptyArrays": True460 }461 },462 {463 "$project": {464 "conversation_id": "$conversation_id",465 "title": "$conversation.title",466 "conversation_created": "$conversation.created_at",467 "message_id": {"$toString": "$_id"},468 "role": 1,469 "content": 1,470 "timestamp": 1,471 "feedback": 1,472 "rating": 1473 }474 },475 {476 "$sort": {"conversation_created": -1, "timestamp": 1}477 }478 ]479 480 data = list(self.db.messages.aggregate(pipeline))481 return data482 483 except Exception as e:484 print(f"Error exporting data: {e}")485 return []486 487 def test_connection(self) -> bool:488 """Test if database connection works"""489 try:490 self.client.admin.command('ping')491 print("[OK] Database connection test successful")492 return True493 except Exception as e:494 print(f"[ERROR] Database connection test failed: {e}")495 return False496 497 498# Singleton instance499db_manager = None500 501def get_db_manager() -> MongoDBManager:502 """Get or create database manager instance"""503 global db_manager504 if db_manager is None:505 db_manager = MongoDBManager()506 return db_manager507 508 509if __name__ == "__main__":510 # Test the database connection511 print("Testing MongoDB Atlas connection...")512 try:513 db = MongoDBManager()514 if db.test_connection():515 print("\n[OK] MongoDB Atlas is ready to use!")516 print("\nTesting basic operations...")517 518 # Test saving a conversation519 test_conv_id = "test_123"520 db.save_conversation(test_conv_id, "Test Conversation")521 print("[OK] Conversation saved")522 523 # Test saving messages524 msg_id = db.save_message(test_conv_id, "user", "Test question?")525 print(f"[OK] Message saved with ID: {msg_id}")526 527 # Test retrieving conversation528 conv = db.get_conversation(test_conv_id)529 print(f"[OK] Retrieved conversation: {conv['title']}")530 531 # Clean up test data532 db.delete_conversation(test_conv_id)533 print("[OK] Test data cleaned up")534 535 print("\n[OK] All tests passed! MongoDB Atlas is working correctly!")536 except Exception as e:537 print(f"\n[ERROR] MongoDB Atlas connection failed: {e}")538 print("\nPlease check:")539 print("1. Your MONGODB_URI in .env file")540 print("2. MongoDB Atlas network access (IP whitelist)")541 print("3. Database user credentials")542