Spaces:
Runtime error
Runtime error
File size: 2,622 Bytes
1e8bb26 | 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 | """MongoDB metadata store for uploaded documents.
If MongoDB is unreachable (e.g. Atlas IP allow-list not configured yet), the
store gracefully falls back to an in-process dictionary so the rest of the app
keeps working. Fallback data is not persisted across restarts.
"""
from __future__ import annotations
from datetime import datetime, timezone
from functools import lru_cache
from typing import List, Optional
from pymongo import MongoClient, DESCENDING
from .config import get_settings
# In-memory fallback store (used only when Mongo can't be reached).
_MEM: dict[str, dict] = {}
_USING_FALLBACK = False
@lru_cache
def _client() -> MongoClient:
settings = get_settings()
return MongoClient(
settings.mongodb_uri,
appname="doc-intelligence-rag",
serverSelectionTimeoutMS=4000,
)
def _collection():
settings = get_settings()
return _client()[settings.mongodb_db]["documents"]
def _fallback(reason: str = "") -> None:
global _USING_FALLBACK
if not _USING_FALLBACK:
_USING_FALLBACK = True
def using_fallback() -> bool:
return _USING_FALLBACK
def save_document(
document_id: str,
filename: str,
content_type: str,
num_chunks: int,
num_chars: int,
ocr_used: bool,
) -> dict:
doc = {
"_id": document_id,
"filename": filename,
"content_type": content_type or "application/octet-stream",
"num_chunks": num_chunks,
"num_chars": num_chars,
"ocr_used": ocr_used,
"created_at": datetime.now(timezone.utc).isoformat(),
}
try:
_collection().insert_one(dict(doc))
except Exception as exc: # noqa: BLE001
_fallback(str(exc))
_MEM[document_id] = doc
return doc
def list_documents() -> List[dict]:
try:
return list(_collection().find().sort("created_at", DESCENDING))
except Exception as exc: # noqa: BLE001
_fallback(str(exc))
return sorted(_MEM.values(), key=lambda d: d["created_at"], reverse=True)
def get_document(document_id: str) -> Optional[dict]:
try:
return _collection().find_one({"_id": document_id})
except Exception as exc: # noqa: BLE001
_fallback(str(exc))
return _MEM.get(document_id)
def delete_document(document_id: str) -> int:
try:
result = _collection().delete_one({"_id": document_id})
return result.deleted_count
except Exception as exc: # noqa: BLE001
_fallback(str(exc))
return 1 if _MEM.pop(document_id, None) else 0
def ping() -> bool:
_client().admin.command("ping")
return True
|