from __future__ import annotations import io import os import re import tempfile import time import threading from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from pathlib import Path from typing import Iterable DIRECT_TEXT_MIN_CHARS = 30 SHORT_PDF_PAGE_LIMIT = 20 TROCR_CONFIDENCE_THRESHOLD = 0.78 DEFAULT_GEMINI_MODELS = ( "gemini-2.5-flash-lite", "gemini-2.5-flash", ) @dataclass class PageResult: page_number: int engine: str text: str confidence: float | None = None @dataclass class ExtractionResult: text: str route: str page_count: int direct_text_found: bool pages: list[PageResult] = field(default_factory=list) warnings: list[str] = field(default_factory=list) class PdfTextExtractor: """Implements the requested PDF-to-text decision flow.""" def __init__( self, gemini_model: str | None = None, trocr_model: str = "microsoft/trocr-base-printed", short_pdf_page_limit: int = SHORT_PDF_PAGE_LIMIT, trocr_confidence_threshold: float = TROCR_CONFIDENCE_THRESHOLD, ) -> None: self.trocr_model_name = trocr_model self.short_pdf_page_limit = short_pdf_page_limit self.trocr_confidence_threshold = trocr_confidence_threshold self._gemini_client = None self._gemini_clients = {} self._gemini_key_cursor = 0 self._gemini_key_lock = threading.Lock() self._gemini_quota_blocked_until: dict[str, float] = {} self._trocr_processor = None self._trocr_model = None self._load_local_env() self.gemini_models = self._gemini_model_candidates(gemini_model) self.gemini_max_retries = min( self._env_int("OCR_GEMINI_MAX_RETRIES", "GEMINI_MAX_RETRIES", default=1, minimum=1), 4, ) self.gemini_retry_delay = min( self._env_float("OCR_GEMINI_RETRY_DELAY", "GEMINI_RETRY_DELAY", default=0.5, minimum=0.0), 3.0, ) self.gemini_timeout_ms = int( min( self._env_float("OCR_GEMINI_TIMEOUT", "GEMINI_TIMEOUT", default=20.0, minimum=5.0), 120.0, ) * 1000 ) self.gemini_key_quota_cooldown = min( self._env_float("GEMINI_KEY_QUOTA_COOLDOWN", default=300.0, minimum=0.0), 86400.0, ) self.gemini_concurrency = min( self._env_int("OCR_GEMINI_CONCURRENCY", default=3, minimum=1), 5, ) self.render_scale = self._env_float("OCR_RENDER_SCALE", default=1.25, minimum=1.0) self.fast_ocr = os.getenv("OCR_FAST_MODE", "1") != "0" def extract(self, pdf_path: str | Path) -> ExtractionResult: path = Path(pdf_path) document = self._open_document(path) try: page_count = document.page_count direct_text = self._extract_direct_text(document) if self._has_enough_text(direct_text): return ExtractionResult( text=self._clean_text_only(direct_text), route="PyMuPDF direct text extraction", page_count=page_count, direct_text_found=True, ) images = self._convert_pages_to_preprocessed_images(document) if page_count <= self.short_pdf_page_limit: pages = [] warnings = [] for index, image_bytes in images: text, warning = self._safe_gemini_ocr(image_bytes, index) if warning: warnings.append(warning) pages.append( PageResult( page_number=index, engine="Gemini Vision OCR" if text else "Skipped non-text page", text=text, ) ) if not pages: warnings.append("No readable text pages were processed.") result = self._finalize_ocr_result( route="PyMuPDF images -> OpenCV preprocess -> Gemini Vision OCR", page_count=page_count, direct_text_found=False, pages=pages, ) result.warnings.extend(warnings) return result pages = [] warnings = [] for page_number, image_bytes in images: try: trocr_text, confidence = self._trocr_ocr(image_bytes) except Exception as exc: trocr_text = "" confidence = 0.0 warnings.append(f"Page {page_number} skipped: TrOCR failed ({exc}).") if confidence < self.trocr_confidence_threshold: gemini_text, warning = self._safe_gemini_ocr(image_bytes, page_number) if warning: warnings.append(warning) pages.append( PageResult( page_number=page_number, engine="TrOCR low confidence -> Gemini Vision OCR" if gemini_text else "Skipped non-text page", text=gemini_text, confidence=confidence, ) ) else: trocr_text = self._clean_text_only(trocr_text) pages.append( PageResult( page_number=page_number, engine="TrOCR" if trocr_text else "Skipped non-text page", text=trocr_text, confidence=confidence, ) ) if not pages: warnings.append("No pages were processed.") result = self._finalize_ocr_result( route="PyMuPDF images -> OpenCV preprocess -> TrOCR -> Gemini low-confidence fallback", page_count=page_count, direct_text_found=False, pages=pages, ) result.warnings.extend(warnings) return result finally: document.close() def stream_extract(self, pdf_path: str | Path) -> Iterable[dict]: path = Path(pdf_path) document = self._open_document(path) try: page_count = document.page_count yield {"type": "start", "page_count": page_count} short_pdf_ocr_jobs = [] for page_number, page in enumerate(document, start=1): direct_text = self._clean_text_only(page.get_text("text")) if self._has_enough_text(direct_text, min_chars=5): yield { "type": "page", "page_number": page_number, "page_count": page_count, "engine": "PyMuPDF direct text extraction", "route": "PyMuPDF direct text extraction", "direct_text_found": True, "confidence": None, "text": direct_text, } continue try: image_bytes = self._convert_page_to_preprocessed_image(page) except Exception as exc: yield self._skipped_page_event(page_number, page_count, f"Image preprocessing failed: {exc}") continue if page_count <= self.short_pdf_page_limit: short_pdf_ocr_jobs.append((page_number, image_bytes)) continue try: trocr_text, confidence = self._trocr_ocr(image_bytes) except Exception as exc: trocr_text = "" confidence = 0.0 yield self._skipped_page_event(page_number, page_count, f"TrOCR failed: {exc}", confidence) continue if confidence < self.trocr_confidence_threshold: ocr_text, warning = self._safe_gemini_ocr(image_bytes, page_number) yield { "type": "page", "page_number": page_number, "page_count": page_count, "engine": "TrOCR low confidence -> Gemini Vision OCR" if ocr_text else "Skipped non-text page", "route": "PyMuPDF image -> OpenCV preprocess -> TrOCR -> Gemini fallback", "direct_text_found": False, "confidence": confidence, "text": ocr_text, "warning": warning, } else: trocr_text = self._clean_text_only(trocr_text) if not self._has_enough_text(trocr_text, min_chars=2): yield self._skipped_page_event(page_number, page_count, "No readable text found.", confidence) continue yield { "type": "page", "page_number": page_number, "page_count": page_count, "engine": "TrOCR", "route": "PyMuPDF image -> OpenCV preprocess -> TrOCR", "direct_text_found": False, "confidence": confidence, "text": self._clean_text(trocr_text), } if short_pdf_ocr_jobs: yield from self._stream_gemini_ocr_jobs(short_pdf_ocr_jobs, page_count) yield {"type": "done", "page_count": page_count} finally: document.close() def _open_document(self, pdf_path: Path): try: import fitz except ImportError as exc: raise RuntimeError("PyMuPDF is required. Install it with: pip install pymupdf") from exc try: document = fitz.open(pdf_path) if document.is_encrypted: raise RuntimeError("This PDF is password-protected. Please upload an unlocked PDF.") if document.page_count == 0: raise RuntimeError("This PDF has no pages.") return document except RuntimeError: raise except Exception as exc: raise RuntimeError(f"Could not open PDF. It may be corrupted or unsupported: {exc}") from exc def _extract_direct_text(self, document) -> str: chunks = [] for page in document: chunks.append(page.get_text("text")) return "\n\n".join(chunks) def _has_enough_text(self, text: str, min_chars: int = DIRECT_TEXT_MIN_CHARS) -> bool: normalized = re.sub(r"\s+", "", text or "") return len(normalized) >= min_chars def _convert_pages_to_preprocessed_images(self, document) -> list[tuple[int, bytes]]: images = [] for index, page in enumerate(document, start=1): try: images.append((index, self._convert_page_to_preprocessed_image(page))) except Exception: continue return images def _convert_page_to_preprocessed_image(self, page) -> bytes: try: import cv2 import fitz import numpy as np from PIL import Image except ImportError as exc: raise RuntimeError( "OCR image preprocessing needs opencv-python, numpy, Pillow, and PyMuPDF." ) from exc matrix = fitz.Matrix(self.render_scale, self.render_scale) pixmap = page.get_pixmap(matrix=matrix, alpha=False) pil_image = Image.open(io.BytesIO(pixmap.tobytes("png"))).convert("RGB") array = np.array(pil_image) gray = cv2.cvtColor(array, cv2.COLOR_RGB2GRAY) if self.fast_ocr: _, thresholded = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) else: denoised = cv2.fastNlMeansDenoising(gray, h=10) thresholded = cv2.adaptiveThreshold( denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 31, 11, ) output = Image.fromarray(thresholded).convert("RGB") buffer = io.BytesIO() output.save(buffer, format="JPEG", quality=72, optimize=True) return buffer.getvalue() def _gemini_ocr(self, image_bytes: bytes) -> str: try: from google import genai from google.genai import types except ImportError as exc: raise RuntimeError("Gemini OCR requires: pip install google-genai") from exc api_keys = self._gemini_api_keys() if not api_keys: raise RuntimeError("Set GEMINI_API_KEY before using Gemini Vision OCR.") prompt = ( "Extract all readable text from this preprocessed PDF page image. " "Preserve natural reading order, headings, bullet points, tables as plain text, " "and do not add commentary. Ignore photos, diagrams, icons, borders, handwriting-like noise, " "and decorative/non-text visual content. Do not describe images. If there is no readable text, " "return an empty response." ) last_error: Exception | None = None attempted_keys = 0 quota_blocked_keys = 0 for api_key in self._gemini_key_attempt_order(api_keys): attempted_keys += 1 client = self._gemini_clients.get(api_key) if client is None: client = genai.Client( api_key=api_key, http_options=types.HttpOptions(timeout=self.gemini_timeout_ms), ) self._gemini_clients[api_key] = client quota_hit = False for model in self.gemini_models: for attempt in range(1, self.gemini_max_retries + 1): try: response = client.models.generate_content( model=model, contents=[ types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg"), prompt, ], ) self._advance_gemini_key_cursor(api_keys, api_key) return (getattr(response, "text", None) or "").strip() except Exception as exc: last_error = exc if self._is_quota_error(exc): self._mark_gemini_key_quota_exhausted(api_keys, api_key) quota_blocked_keys += 1 quota_hit = True break if not self._is_retryable_gemini_error(exc): break if attempt < self.gemini_max_retries: time.sleep(self.gemini_retry_delay * attempt) if quota_hit: break models = ", ".join(self.gemini_models) raise RuntimeError( "Gemini OCR is temporarily unavailable after retries. " f"Tried {attempted_keys} active API key(s) out of {len(api_keys)} configured key(s); " f"{quota_blocked_keys} key(s) returned quota/rate-limit errors. Models tried: {models}. " f"Last error: {last_error}" ) def _safe_gemini_ocr(self, image_bytes: bytes, page_number: int) -> tuple[str, str | None]: try: text = self._clean_text_only(self._gemini_ocr(image_bytes)) except Exception as exc: return "", f"Page {page_number} skipped: OCR failed ({exc})." if not self._has_enough_text(text, min_chars=2): return "", f"Page {page_number} skipped: no readable text found." return text, None def _skipped_page_event( self, page_number: int, page_count: int, warning: str, confidence: float | None = None, ) -> dict: return { "type": "page", "page_number": page_number, "page_count": page_count, "engine": "Skipped non-text page", "route": "Skipped non-text content", "direct_text_found": False, "confidence": confidence, "text": "", "warning": warning, } def _stream_gemini_ocr_jobs(self, jobs: list[tuple[int, bytes]], page_count: int) -> Iterable[dict]: worker_count = min(self.gemini_concurrency, len(jobs)) if worker_count <= 1: for page_number, image_bytes in jobs: yield self._gemini_page_event(page_number, page_count, image_bytes) return with ThreadPoolExecutor(max_workers=worker_count) as executor: futures = { executor.submit(self._gemini_page_event, page_number, page_count, image_bytes): page_number for page_number, image_bytes in jobs } for future in as_completed(futures): try: yield future.result() except Exception as exc: yield self._skipped_page_event(futures[future], page_count, f"OCR failed: {exc}") def _gemini_page_event(self, page_number: int, page_count: int, image_bytes: bytes) -> dict: ocr_text, warning = self._safe_gemini_ocr(image_bytes, page_number) return { "type": "page", "page_number": page_number, "page_count": page_count, "engine": "Gemini Vision OCR" if ocr_text else "Skipped non-text page", "route": "PyMuPDF image -> OpenCV preprocess -> Gemini Vision OCR", "direct_text_found": False, "confidence": None, "text": ocr_text, "warning": warning, } def _gemini_api_keys(self) -> list[str]: keys = [] for key in sorted(os.environ, key=self._gemini_env_key_sort): value = os.environ[key] if self._is_gemini_key_name(key): self._add_gemini_key_values(keys, value) for file_name in (".env", "env"): env_path = Path(file_name) if not env_path.exists(): continue for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines(): stripped = line.strip() if not stripped or stripped.startswith("#") or "=" not in stripped: continue key, value = stripped.split("=", 1) key = key.strip() value = value.strip().strip('"').strip("'") if self._is_gemini_key_name(key): self._add_gemini_key_values(keys, value) return keys def _gemini_key_attempt_order(self, api_keys: list[str]) -> list[str]: with self._gemini_key_lock: now = time.time() active_keys = [ api_key for api_key in api_keys if self._gemini_quota_blocked_until.get(api_key, 0.0) <= now ] if not active_keys: active_keys = api_keys active_key_set = set(active_keys) start = self._gemini_key_cursor % len(api_keys) ordered_keys = api_keys[start:] + api_keys[:start] return [api_key for api_key in ordered_keys if api_key in active_key_set] def _advance_gemini_key_cursor(self, api_keys: list[str], api_key: str) -> None: with self._gemini_key_lock: try: self._gemini_key_cursor = (api_keys.index(api_key) + 1) % len(api_keys) except ValueError: self._gemini_key_cursor = 0 def _mark_gemini_key_quota_exhausted(self, api_keys: list[str], api_key: str) -> None: with self._gemini_key_lock: self._gemini_quota_blocked_until[api_key] = time.time() + self.gemini_key_quota_cooldown try: self._gemini_key_cursor = (api_keys.index(api_key) + 1) % len(api_keys) except ValueError: self._gemini_key_cursor = 0 def _is_gemini_key_name(self, key: str) -> bool: return ( key == "GEMINI_API_KEY" or key == "GOOGLE_API_KEY" or key == "GEMINI_API_KEYS" or key.startswith("GEMINI_API_KEY_") ) def _gemini_env_key_sort(self, key: str) -> tuple[int, int, str]: match = re.fullmatch(r"GEMINI_API_KEY_(\d+)", key) if match: return (0, int(match.group(1)), key) if key == "GEMINI_API_KEY": return (1, 0, key) if key == "GEMINI_API_KEYS": return (2, 0, key) if key == "GOOGLE_API_KEY": return (3, 0, key) return (4, 0, key) def _add_gemini_key_values(self, keys: list[str], value: str) -> None: for api_key in (part.strip() for part in value.split(",")): if api_key and api_key not in keys: keys.append(api_key) def _gemini_model_candidates(self, configured_model: str | None) -> list[str]: configured_models = os.getenv("GEMINI_MODELS") or configured_model if configured_models: models = [ model.strip() for model in configured_models.split(",") if model.strip() ] if models: return models return list(DEFAULT_GEMINI_MODELS) def _is_retryable_gemini_error(self, error: Exception) -> bool: message = str(error).lower() retryable_markers = ( "503", "unavailable", "high demand", "429", "resource_exhausted", "rate limit", "quota", "deadline", "timeout", ) return any(marker in message for marker in retryable_markers) def _is_quota_error(self, error: Exception | None) -> bool: if error is None: return False message = str(error).lower() return "429" in message or "resource_exhausted" in message or "quota" in message def _load_local_env(self) -> None: for file_name in (".env", "env"): env_path = Path(file_name) if not env_path.exists(): continue for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines(): stripped = line.strip() if not stripped or stripped.startswith("#") or "=" not in stripped: continue key, value = stripped.split("=", 1) key = key.strip() value = value.strip().strip('"').strip("'") if key and key not in os.environ: os.environ[key] = value def _env_int(self, key: str, fallback_key: str | None = None, default: int = 0, minimum: int = 0) -> int: raw_value = os.getenv(key) if raw_value is None and fallback_key: raw_value = os.getenv(fallback_key) if raw_value is None: return default try: value = int(raw_value) except ValueError: return default return max(value, minimum) def _env_float(self, key: str, fallback_key: str | None = None, default: float = 0.0, minimum: float = 0.0) -> float: raw_value = os.getenv(key) if raw_value is None and fallback_key: raw_value = os.getenv(fallback_key) if raw_value is None: return default try: value = float(raw_value) except ValueError: return default return max(value, minimum) def _trocr_ocr(self, image_bytes: bytes) -> tuple[str, float]: try: import torch from PIL import Image from transformers import TrOCRProcessor, VisionEncoderDecoderModel except ImportError as exc: raise RuntimeError( "TrOCR requires: pip install torch transformers Pillow" ) from exc if self._trocr_processor is None or self._trocr_model is None: self._trocr_processor = TrOCRProcessor.from_pretrained(self.trocr_model_name) self._trocr_model = VisionEncoderDecoderModel.from_pretrained(self.trocr_model_name) image = Image.open(io.BytesIO(image_bytes)).convert("RGB") pixel_values = self._trocr_processor(images=image, return_tensors="pt").pixel_values with torch.no_grad(): generated = self._trocr_model.generate( pixel_values, max_new_tokens=512, output_scores=True, return_dict_in_generate=True, ) text = self._trocr_processor.batch_decode( generated.sequences, skip_special_tokens=True, )[0].strip() confidence = self._score_trocr_confidence(generated.scores) return text, confidence def _score_trocr_confidence(self, scores: Iterable) -> float: try: import torch except ImportError: return 0.0 confidences = [] for score in scores: probabilities = torch.softmax(score, dim=-1) confidences.append(float(probabilities.max())) if not confidences: return 0.0 return sum(confidences) / len(confidences) def _finalize_ocr_result( self, route: str, page_count: int, direct_text_found: bool, pages: list[PageResult], ) -> ExtractionResult: merged_text = "\n\n".join( f"--- Page {page.page_number} ---\n{page.text.strip()}" for page in pages if page.text.strip() ) return ExtractionResult( text=self._clean_text(merged_text), route=route, page_count=page_count, direct_text_found=direct_text_found, pages=pages, ) def _clean_text(self, text: str) -> str: text = text.replace("\x00", "") text = re.sub(r"[ \t]+", " ", text) text = re.sub(r"\n{3,}", "\n\n", text) text = "\n".join(line.rstrip() for line in text.splitlines()) return text.strip() def _clean_text_only(self, text: str) -> str: cleaned = self._clean_text(text) lines = [] for line in cleaned.splitlines(): stripped = line.strip() if not stripped: lines.append("") continue if self._looks_like_visual_description(stripped): continue lines.append(stripped) text_only = self._clean_text("\n".join(lines)) return text_only if self._has_real_text_signal(text_only) else "" def _looks_like_visual_description(self, line: str) -> bool: normalized = re.sub(r"\s+", " ", line.lower()).strip() visual_description_patterns = ( r"^(the|this|an|a)\s+(image|photo|picture|diagram|figure|chart|graph|illustration|logo|icon)\s+", r"^(the|this)\s+page\s+(contains|shows|appears|has|is)\s+", r"^(it|this)\s+(shows|appears|looks like|contains|depicts)\s+", r"^(i can see|there is|there are)\s+", r"\b(no readable text|no text|cannot extract|not able to extract)\b", r"\b(image shows|picture shows|diagram shows|photo shows|chart shows|graph shows)\b", ) if any(re.search(pattern, normalized) for pattern in visual_description_patterns): return True visual_words = { "image", "photo", "picture", "diagram", "illustration", "visual", "icon", "logo", "background", "border", "shape", "graphic", } tokens = re.findall(r"[a-z0-9]+", normalized) if not tokens: return True visual_count = sum(1 for token in tokens if token in visual_words) return visual_count >= 2 and len(tokens) <= 18 def _has_real_text_signal(self, text: str) -> bool: normalized = re.sub(r"\s+", " ", text or "").strip() if not normalized: return False alnum_count = sum(char.isalnum() for char in normalized) alpha_count = sum(char.isalpha() for char in normalized) return alnum_count >= 2 and alpha_count >= 2 def extract_pdf_to_text(pdf_path: str | Path) -> ExtractionResult: return PdfTextExtractor().extract(pdf_path) def extract_uploaded_pdf(file_name: str, file_bytes: bytes) -> ExtractionResult: suffix = Path(file_name).suffix or ".pdf" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: temp_file.write(file_bytes) temp_path = Path(temp_file.name) try: return extract_pdf_to_text(temp_path) finally: temp_path.unlink(missing_ok=True) def stream_uploaded_pdf(file_name: str, file_bytes: bytes) -> Iterable[dict]: suffix = Path(file_name).suffix or ".pdf" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: temp_file.write(file_bytes) temp_path = Path(temp_file.name) try: yield from PdfTextExtractor().stream_extract(temp_path) finally: temp_path.unlink(missing_ok=True)