| import os |
| import re |
| import time |
| from typing import Dict, List, Optional |
|
|
| try: |
| import pymongo |
| HAS_PYMONGO = True |
| except ImportError: |
| HAS_PYMONGO = False |
|
|
| |
| MONGODB_URI = os.environ.get("MONGODB_URI", "").strip() |
| DB_NAME = os.environ.get("MONGODB_DB_NAME", "llm_xray_db") |
| COLLECTION_NAME = "audit_certificates" |
|
|
| _mongo_client = None |
|
|
|
|
| def get_db_collection(retries: int = 3, delay: float = 2.0): |
| """Initializes and returns the MongoDB collection using HF Secrets. |
| Retries the initial connection — HF Space cold starts often need a |
| few extra seconds for DNS/TLS to Atlas before serverSelection succeeds. |
| """ |
| global _mongo_client |
| if not HAS_PYMONGO: |
| return None |
| if not MONGODB_URI: |
| return None |
|
|
| for attempt in range(1, retries + 1): |
| try: |
| if _mongo_client is None: |
| _mongo_client = pymongo.MongoClient( |
| MONGODB_URI, |
| serverSelectionTimeoutMS=8000, |
| connectTimeoutMS=8000, |
| ) |
| |
| |
| |
| _mongo_client.admin.command("ping") |
| return _mongo_client[DB_NAME][COLLECTION_NAME] |
| except Exception as e: |
| print(f"[MongoDB] Connection attempt {attempt}/{retries} warning: {e}", flush=True) |
| _mongo_client = None |
| if attempt < retries: |
| time.sleep(delay) |
| return None |
|
|
|
|
| def mongo_save_certificate(cert: Dict) -> bool: |
| """Persists an audit certificate permanently to MongoDB Atlas.""" |
| col = get_db_collection() |
| if col is None: |
| return False |
| try: |
| config = cert.get("config", {}) |
| model_name = config.get("model_name", "unknown").strip() |
| model_sha = str(config.get("model_sha", "main")).strip() |
| doc_id = f"{model_name}__{model_sha}".replace("/", "__").lower() |
|
|
| doc = dict(cert) |
| doc["_id"] = doc_id |
| doc["model_name"] = model_name |
| doc["model_sha"] = model_sha |
| doc["updated_at"] = cert.get("audited_at", "") |
|
|
| col.replace_one({"_id": doc_id}, doc, upsert=True) |
| print(f"[MongoDB] ✅ Successfully persisted '{model_name}' to Atlas collection", flush=True) |
| return True |
| except Exception as e: |
| print(f"[MongoDB] ⚠️ Error saving certificate: {e}", flush=True) |
| return False |
|
|
|
|
| def mongo_get_certificate(clean_model_name: str) -> Optional[Dict]: |
| """Retrieves an audit certificate by model name from MongoDB.""" |
| col = get_db_collection() |
| if col is None or not clean_model_name: |
| return None |
| try: |
| norm = clean_model_name.strip() |
| doc_id = norm.replace("/", "__").lower() |
| doc = col.find_one({ |
| "$or": [ |
| {"_id": {"$regex": f"^{doc_id}", "$options": "i"}}, |
| {"config.model_name": {"$regex": f"^{re.escape(norm)}$", "$options": "i"}}, |
| {"model_name": {"$regex": f"^{re.escape(norm)}$", "$options": "i"}} |
| ] |
| }) |
| if doc: |
| doc.pop("_id", None) |
| return doc |
| except Exception as e: |
| print(f"[MongoDB] ⚠️ Error retrieving '{clean_model_name}': {e}", flush=True) |
| return None |
|
|
|
|
| def mongo_get_all_certificates() -> List[Dict]: |
| """Recalls all saved certificates.""" |
| col = get_db_collection() |
| if col is None: |
| return [] |
| try: |
| docs = list(col.find({})) |
| for d in docs: |
| d.pop("_id", None) |
| return docs |
| except Exception as e: |
| print(f"[MongoDB] ⚠️ Error loading certificates: {e}", flush=True) |
| return [] |
