Spaces:
Running
Running
| import os | |
| import json | |
| import firebase_admin | |
| from firebase_admin import credentials, firestore | |
| class FirebaseManager: | |
| """ | |
| Manages connection to Firebase Firestore for caching deepfake detection results | |
| and tracking async task states. Falls back to in-memory dictionaries if Firebase is offline. | |
| Uses a single collection 'deepfake_analyses' for all documents. | |
| """ | |
| def __init__(self): | |
| self.db = None | |
| self.local_cache = {} | |
| cred_json_str = os.environ.get("FIREBASE_CREDENTIALS_JSON") | |
| cred_file_path = os.environ.get("FIREBASE_SERVICE_ACCOUNT_KEY") | |
| cred = None | |
| try: | |
| if cred_json_str: | |
| cred_info = json.loads(cred_json_str) | |
| cred = credentials.Certificate(cred_info) | |
| elif cred_file_path and os.path.exists(cred_file_path): | |
| cred = credentials.Certificate(cred_file_path) | |
| if not firebase_admin._apps: | |
| if cred: | |
| firebase_admin.initialize_app(cred) | |
| else: | |
| print("No Firebase credentials provided. Skipping Firebase initialization.") | |
| raise ValueError("Missing Firebase credentials") | |
| self.db = firestore.client() | |
| print("Connected to Firebase Firestore successfully!") | |
| except Exception as e: | |
| print(f"Warning: Firebase is not reachable ({e}). Falling back to in-memory cache.") | |
| self.db = None | |
| def get_cached_result(self, file_hash: str) -> dict | None: | |
| if self.db: | |
| try: | |
| doc_ref = self.db.collection("deepfake_analyses").document(file_hash) | |
| doc = doc_ref.get() | |
| if doc.exists: | |
| return doc.to_dict() | |
| except Exception as e: | |
| print(f"Firebase get cache error: {e}") | |
| return self.local_cache.get(file_hash) | |
| def cache_result(self, file_hash: str, doc_data: dict) -> None: | |
| self.local_cache[file_hash] = doc_data | |
| if self.db: | |
| try: | |
| doc_ref = self.db.collection("deepfake_analyses").document(file_hash) | |
| doc_ref.set(doc_data, merge=True) | |
| except Exception as e: | |
| print(f"Firebase set cache error: {e}") | |
| def get_task_status(self, task_id: str) -> dict | None: | |
| if self.db: | |
| try: | |
| docs = self.db.collection("deepfake_analyses").where("task_id", "==", task_id).limit(1).get() | |
| if docs: | |
| return docs[0].to_dict() | |
| except Exception as e: | |
| print(f"Firebase get task error: {e}") | |
| # Local search | |
| for k, v in self.local_cache.items(): | |
| if v.get("task_id") == task_id: | |
| return v | |
| return None | |
| def update_task_status(self, file_hash: str, status: dict) -> None: | |
| existing = self.local_cache.get(file_hash, {}) | |
| existing.update(status) | |
| self.local_cache[file_hash] = existing | |
| if self.db: | |
| try: | |
| doc_ref = self.db.collection("deepfake_analyses").document(file_hash) | |
| doc_ref.set(status, merge=True) | |
| except Exception as e: | |
| print(f"Firebase set task error: {e}") | |
| def get_history(self) -> list[dict]: | |
| history = [] | |
| if self.db: | |
| try: | |
| docs = self.db.collection("deepfake_analyses").stream() | |
| for doc in docs: | |
| res = doc.to_dict() | |
| if res and "result" in res: | |
| history.append(self._strip_heavy_payload(res)) | |
| history.sort(key=lambda x: x.get("timestamp", 0), reverse=True) | |
| return history | |
| except Exception as e: | |
| print(f"Firebase get history error: {e}") | |
| for file_hash, res in self.local_cache.items(): | |
| if "result" in res: | |
| history.append(self._strip_heavy_payload(res)) | |
| history.sort(key=lambda x: x.get("timestamp", 0), reverse=True) | |
| return history | |
| def _strip_heavy_payload(self, res: dict) -> dict: | |
| light_res = res.copy() | |
| if "image" in light_res: | |
| light_res["image"] = None | |
| if "result" in light_res and isinstance(light_res["result"], dict): | |
| if "frames" in light_res["result"]: | |
| for frame in light_res["result"]["frames"]: | |
| if "image" in frame: | |
| frame["image"] = None | |
| if "faces" in frame: | |
| for face in frame["faces"]: | |
| if "crop_b64" in face: | |
| face["crop_b64"] = None | |
| return light_res | |
| def clear_cache(self) -> None: | |
| self.local_cache.clear() | |
| if self.db: | |
| try: | |
| results_ref = self.db.collection("deepfake_analyses") | |
| docs = results_ref.list_documents() | |
| for doc in docs: | |
| doc.delete() | |
| print("Firebase cache cleared.") | |
| except Exception as e: | |
| print(f"Firebase clear cache error: {e}") | |