Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| from typing import List, Dict, Optional | |
| from pymongo import MongoClient, errors | |
| from pymongo.collection import Collection | |
| class MongoDBHandler: | |
| """ | |
| Handles interactions with MongoDB for the Gemma-3 RAG system. | |
| Follows PythonMind principles: Simple, Modular, and Robust. | |
| """ | |
| def __init__(self, uri: str = "mongodb://localhost:27017/", db_name: str = "rag_db", collection_name: str = "chunks"): | |
| self.uri = uri | |
| self.db_name = db_name | |
| self.collection_name = collection_name | |
| self.client: Optional[MongoClient] = None | |
| self.db = None | |
| self.collection: Optional[Collection] = None | |
| def connect(self) -> bool: | |
| """Establishes connection to MongoDB.""" | |
| try: | |
| self.client = MongoClient(self.uri, serverSelectionTimeoutMS=5000) | |
| # Trigger a server selection to verify connection | |
| self.client.server_info() | |
| self.db = self.client[self.db_name] | |
| self.collection = self.db[self.collection_name] | |
| print(f"[MONGO] Connected to {self.uri}, Database: {self.db_name}, Collection: {self.collection_name}") | |
| return True | |
| except errors.ServerSelectionTimeoutError as e: | |
| print(f"[MONGO ERROR] Could not connect to MongoDB: {e}") | |
| return False | |
| except Exception as e: | |
| print(f"[MONGO ERROR] An unexpected error occurred: {e}") | |
| return False | |
| def insert_chunk(self, content: str, metadata: Dict) -> bool: | |
| """Inserts a single chunk into the collection.""" | |
| if self.collection is None: | |
| if not self.connect(): return False | |
| try: | |
| document = { | |
| "content": content, | |
| "metadata": metadata, | |
| "timestamp": time.time() | |
| } | |
| self.collection.insert_one(document) | |
| return True | |
| except Exception as e: | |
| print(f"[MONGO ERROR] Insertion failed: {e}") | |
| return False | |
| def find_relevant(self, query_dict: Dict, limit: int = 5) -> List[Dict]: | |
| """Performs a standard query search.""" | |
| if self.collection is None: | |
| if not self.connect(): return [] | |
| try: | |
| results = self.collection.find(query_dict).limit(limit) | |
| return list(results) | |
| except Exception as e: | |
| print(f"[MONGO ERROR] Search failed: {e}") | |
| return [] | |
| def clear(self) -> bool: | |
| """Clears the collection.""" | |
| if self.collection is None: | |
| if not self.connect(): return False | |
| try: | |
| self.collection.delete_many({}) | |
| return True | |
| except Exception as e: | |
| print(f"[MONGO ERROR] Clear failed: {e}") | |
| return False | |
| def get_stats(self) -> Dict: | |
| """Returns collection statistics.""" | |
| if self.collection is None: | |
| if not self.connect(): return {"error": "Not connected"} | |
| try: | |
| count = self.collection.count_documents({}) | |
| return {"count": count, "collection": self.collection_name, "db": self.db_name} | |
| except Exception as e: | |
| return {"error": str(e)} | |
| def close(self): | |
| """Closes the connection.""" | |
| if self.client: | |
| self.client.close() | |