File size: 3,858 Bytes
422468d 275f2ec 9ad792d 422468d 275f2ec 422468d 275f2ec 9ad792d 422468d 275f2ec 422468d 9ad792d 422468d 275f2ec 422468d 275f2ec 422468d 275f2ec 422468d 275f2ec 422468d d7f0565 422468d 275f2ec 422468d 275f2ec 422468d d7f0565 422468d 275f2ec 422468d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import os
import re
import time
from typing import Dict, List, Optional
try:
import pymongo
HAS_PYMONGO = True
except ImportError:
HAS_PYMONGO = False
# Read connection string strictly from Hugging Face Space Secrets
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,
)
# Force a real round-trip here (pymongo is lazy otherwise),
# so a slow cold-start connection is retried at this single
# choke point rather than surfacing as "0 certificates".
_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 [] |