Spaces:
Sleeping
Sleeping
| """Document validators — MIME detection and file safety checks. | |
| Validates uploaded documents before any processing: | |
| 1. File size limit | |
| 2. MIME type detection via magic bytes (content-first, not extension-first) | |
| 3. Supported MIME type check | |
| 4. Basic structural sanity | |
| Never calls AI models. | |
| Never crashes on corrupt input — returns structured errors. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import logging | |
| from ehr_intel.models import DocumentTypeEnum | |
| log = logging.getLogger(__name__) | |
| # --------------------------------------------------------------------------- | |
| # Supported MIME types | |
| # --------------------------------------------------------------------------- | |
| SUPPORTED_MIME_TYPES: frozenset[str] = frozenset({ | |
| "application/pdf", | |
| "image/png", | |
| "image/jpeg", | |
| "image/webp", | |
| "text/plain", | |
| "application/json", | |
| "application/fhir+json", | |
| }) | |
| # Magic byte signatures for content-based MIME detection | |
| _MAGIC_SIGNATURES: list[tuple[bytes, str]] = [ | |
| (b"%PDF", "application/pdf"), | |
| (b"\x89PNG\r\n\x1a\n", "image/png"), | |
| (b"\xff\xd8\xff", "image/jpeg"), | |
| (b"RIFF", "image/webp"), # bytes 8-11 must also be b"WEBP" | |
| (b"\xff\xd8", "image/jpeg"), # looser JPEG prefix fallback | |
| ] | |
| _EXTENSION_MIME_MAP: dict[str, str] = { | |
| ".pdf": "application/pdf", | |
| ".png": "image/png", | |
| ".jpg": "image/jpeg", | |
| ".jpeg": "image/jpeg", | |
| ".webp": "image/webp", | |
| ".txt": "text/plain", | |
| ".json": "application/json", | |
| } | |
| class ValidationError(Exception): | |
| """Raised when a document fails validation.""" | |
| def __init__(self, message: str, code: str = "VALIDATION_ERROR") -> None: | |
| super().__init__(message) | |
| self.code = code | |
| self.message = message | |
| class ValidationResult: | |
| """Result of document validation.""" | |
| def __init__( | |
| self, | |
| *, | |
| valid: bool, | |
| mime_type: str, | |
| document_type: DocumentTypeEnum, | |
| file_size_bytes: int, | |
| sha256: str, | |
| error: str | None = None, | |
| error_code: str | None = None, | |
| ) -> None: | |
| self.valid = valid | |
| self.mime_type = mime_type | |
| self.document_type = document_type | |
| self.file_size_bytes = file_size_bytes | |
| self.sha256 = sha256 | |
| self.error = error | |
| self.error_code = error_code | |
| def detect_mime_type( | |
| raw_bytes: bytes, | |
| filename: str | None = None, | |
| claimed_mime: str | None = None, | |
| ) -> str: | |
| """Detects MIME type using magic bytes first. | |
| Falls back to claimed MIME and then file extension. | |
| Content always wins over extension claim. | |
| Args: | |
| raw_bytes: Raw file bytes. May be empty. | |
| filename: Original filename. Used only as final fallback. | |
| claimed_mime: MIME type claimed by the caller. Not trusted blindly. | |
| Returns: | |
| Detected MIME type string. | |
| """ | |
| if raw_bytes: | |
| # --- Primary: magic byte sniffing --- | |
| for signature, mime in _MAGIC_SIGNATURES: | |
| if raw_bytes.startswith(signature): | |
| if signature == b"RIFF": | |
| # WebP: bytes 8-12 must be b"WEBP" | |
| if len(raw_bytes) >= 12 and raw_bytes[8:12] == b"WEBP": | |
| return "image/webp" | |
| continue | |
| return mime | |
| # --- Check for JSON / FHIR Bundle --- | |
| try: | |
| snippet = raw_bytes[:1024].decode("utf-8", errors="ignore").strip() | |
| if snippet.startswith("{") or snippet.startswith("["): | |
| parsed = json.loads(raw_bytes.decode("utf-8", errors="replace")) | |
| if isinstance(parsed, dict) and parsed.get("resourceType") == "Bundle": | |
| return "application/fhir+json" | |
| return "application/json" | |
| except (json.JSONDecodeError, UnicodeDecodeError): | |
| pass | |
| # --- Check for plain text (UTF-8 decodable, no binary nulls) --- | |
| try: | |
| text = raw_bytes[:512].decode("utf-8") | |
| if "\x00" not in text: | |
| # Looks like plain text | |
| if claimed_mime == "text/plain" or ( | |
| filename and filename.lower().endswith(".txt") | |
| ): | |
| return "text/plain" | |
| except UnicodeDecodeError: | |
| pass | |
| # --- Secondary: claimed MIME --- | |
| if claimed_mime: | |
| normalized = claimed_mime.lower().strip() | |
| # Normalize "image/jpg" → "image/jpeg" | |
| if normalized == "image/jpg": | |
| normalized = "image/jpeg" | |
| if normalized in SUPPORTED_MIME_TYPES or normalized == "text/plain": | |
| return normalized | |
| # --- Tertiary: extension --- | |
| if filename: | |
| fn_lower = filename.lower().strip() | |
| for ext, mime in _EXTENSION_MIME_MAP.items(): | |
| if fn_lower.endswith(ext): | |
| return mime | |
| return "application/octet-stream" | |
| def classify_document_type( | |
| *, | |
| mime_type: str, | |
| filename: str | None = None, | |
| text_content: str | None = None, | |
| ) -> DocumentTypeEnum: | |
| """Classifies clinical document type from available signals. | |
| This is a lightweight heuristic classification. Gemini will refine it later. | |
| Args: | |
| mime_type: Detected MIME type. | |
| filename: Original filename. | |
| text_content: Short text sample, if available (e.g. from PDF extraction). | |
| Returns: | |
| DocumentTypeEnum best guess. | |
| """ | |
| if mime_type == "application/fhir+json": | |
| return DocumentTypeEnum.FHIR_BUNDLE | |
| probe = " ".join(filter(None, [filename or "", text_content or ""])).lower() | |
| # FHIR Bundle text indicator | |
| if "resourcetype" in probe or '"bundle"' in probe: | |
| return DocumentTypeEnum.FHIR_BUNDLE | |
| # Prescription indicators | |
| rx_kws = ["prescription", "rx:", "sig:", "disp:", "dispense", "refill"] | |
| rx_soft = ["tab", "cap", "od", "bd", "tid", "qid", "po", "mg", "tablet", "capsule"] | |
| if any(kw in probe for kw in rx_kws) or sum(1 for kw in rx_soft if kw in probe) >= 3: | |
| if "discharge summary" not in probe: | |
| return DocumentTypeEnum.PRESCRIPTION | |
| # Doctor note indicators | |
| note_kws = ["progress note", "soap note", "subjective:", "objective:", "assessment", "consultation note"] | |
| if any(kw in probe for kw in note_kws): | |
| return DocumentTypeEnum.DOCTOR_NOTE | |
| # Discharge summary indicators | |
| ds_kws = ["discharge summary", "hospital course", "admission date", "discharge date", "discharge diagnosis"] | |
| if any(kw in probe for kw in ds_kws): | |
| return DocumentTypeEnum.DISCHARGE_SUMMARY | |
| # Lab report indicators | |
| lab_kws = ["lab report", "laboratory", "reference range", "specimen", "hemoglobin", "creatinine"] | |
| if any(kw in probe for kw in lab_kws): | |
| return DocumentTypeEnum.LAB_REPORT | |
| # Filename-based hints | |
| if filename: | |
| fn = filename.lower() | |
| if any(kw in fn for kw in ["presc", "rx_", "_rx"]): | |
| return DocumentTypeEnum.PRESCRIPTION | |
| if any(kw in fn for kw in ["discharge", "disch", "summary"]): | |
| return DocumentTypeEnum.DISCHARGE_SUMMARY | |
| if any(kw in fn for kw in ["note", "soap", "doctor"]): | |
| return DocumentTypeEnum.DOCTOR_NOTE | |
| if "lab" in fn or "report" in fn: | |
| return DocumentTypeEnum.LAB_REPORT | |
| if mime_type in ("image/png", "image/jpeg", "image/webp"): | |
| return DocumentTypeEnum.GENERAL_CLINICAL | |
| return DocumentTypeEnum.UNKNOWN | |
| def validate_document( | |
| raw_bytes: bytes, | |
| *, | |
| filename: str | None = None, | |
| claimed_mime: str | None = None, | |
| max_file_mb: int = 20, | |
| ) -> ValidationResult: | |
| """Full document validation pipeline. | |
| Args: | |
| raw_bytes: Document bytes (may be empty — will fail gracefully). | |
| filename: Original filename. | |
| claimed_mime: MIME type claimed by caller. | |
| max_file_mb: Maximum allowed file size in MB. | |
| Returns: | |
| ValidationResult — always returns (never raises), callers check .valid. | |
| """ | |
| # Compute hash immediately (needed for logging even on failure) | |
| sha256 = hashlib.sha256(raw_bytes).hexdigest() if raw_bytes else "empty" | |
| size = len(raw_bytes) | |
| # --- Empty file check --- | |
| if not raw_bytes or size == 0: | |
| return ValidationResult( | |
| valid=False, | |
| mime_type="application/octet-stream", | |
| document_type=DocumentTypeEnum.UNKNOWN, | |
| file_size_bytes=0, | |
| sha256=sha256, | |
| error="Empty file received.", | |
| error_code="EMPTY_FILE", | |
| ) | |
| # --- Size limit check --- | |
| max_bytes = max_file_mb * 1024 * 1024 | |
| if size > max_bytes: | |
| return ValidationResult( | |
| valid=False, | |
| mime_type="application/octet-stream", | |
| document_type=DocumentTypeEnum.UNKNOWN, | |
| file_size_bytes=size, | |
| sha256=sha256, | |
| error=f"File size {size / 1024 / 1024:.1f} MB exceeds limit of {max_file_mb} MB.", | |
| error_code="FILE_TOO_LARGE", | |
| ) | |
| # --- MIME detection --- | |
| mime = detect_mime_type(raw_bytes, filename=filename, claimed_mime=claimed_mime) | |
| # --- Supported MIME check --- | |
| if mime not in SUPPORTED_MIME_TYPES and mime != "application/octet-stream": | |
| # Allow some unlisted mimes through with a warning | |
| pass | |
| if mime == "application/octet-stream": | |
| return ValidationResult( | |
| valid=False, | |
| mime_type=mime, | |
| document_type=DocumentTypeEnum.UNKNOWN, | |
| file_size_bytes=size, | |
| sha256=sha256, | |
| error=( | |
| f"Unsupported or unrecognized file type. " | |
| f"Accepted: {sorted(SUPPORTED_MIME_TYPES)}." | |
| ), | |
| error_code="UNSUPPORTED_MIME", | |
| ) | |
| # --- Document type classification --- | |
| doc_type = classify_document_type(mime_type=mime, filename=filename) | |
| log.info( | |
| "Document validated: sha256=%s... size=%d mime=%s type=%s", | |
| sha256[:8], | |
| size, | |
| mime, | |
| doc_type, | |
| ) | |
| return ValidationResult( | |
| valid=True, | |
| mime_type=mime, | |
| document_type=doc_type, | |
| file_size_bytes=size, | |
| sha256=sha256, | |
| ) | |
| __all__ = [ | |
| "SUPPORTED_MIME_TYPES", | |
| "ValidationError", | |
| "ValidationResult", | |
| "classify_document_type", | |
| "detect_mime_type", | |
| "validate_document", | |
| ] | |