82 lines
3.4 KiB
Python
82 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import fitz
|
|
|
|
from backend.app.services.markdown_builder import title_from_path
|
|
from backend.app.services.processors.base import ProcessedContent
|
|
from backend.app.services.providers.mistral_client import MistralClient
|
|
|
|
|
|
class PdfProcessor:
|
|
def __init__(self, use_ocr: bool = False, page_ranges: str | None = None):
|
|
self.use_ocr = use_ocr
|
|
self.page_ranges = page_ranges
|
|
|
|
def process(self, path: Path) -> ProcessedContent:
|
|
if self.use_ocr:
|
|
return self._process_ocr(path)
|
|
doc = fitz.open(path)
|
|
pages: list[str] = []
|
|
for index in self._selected_pages(len(doc)):
|
|
page = doc[index - 1]
|
|
text = page.get_text("text").strip()
|
|
if text:
|
|
pages.append(f"## Pagina {index}\n\n{text}")
|
|
body = "\n\n".join(pages).strip()
|
|
if len(body) < 80:
|
|
raise RuntimeError("El PDF tiene poco texto extraible; procesalo con OCR")
|
|
return ProcessedContent(title=title_from_path(path), body=body, processor="pdf-text")
|
|
|
|
def _process_ocr(self, path: Path) -> ProcessedContent:
|
|
doc = fitz.open(path)
|
|
client = MistralClient()
|
|
parts: list[str] = []
|
|
temp_paths: list[Path] = []
|
|
try:
|
|
for index in self._selected_pages(len(doc)):
|
|
page = doc[index - 1]
|
|
pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
|
|
image_path = path.with_name(f"{path.stem}-page-{index}.jpg")
|
|
pix.save(image_path)
|
|
temp_paths.append(image_path)
|
|
parts.append(f"## Pagina {index}\n\n{client.ocr_image(image_path)}")
|
|
finally:
|
|
for temp_path in temp_paths:
|
|
temp_path.unlink(missing_ok=True)
|
|
return ProcessedContent(title=title_from_path(path), body="\n\n".join(parts), processor="mistral-ocr")
|
|
|
|
def _selected_pages(self, total_pages: int) -> list[int]:
|
|
if total_pages < 1:
|
|
raise RuntimeError("El PDF no tiene paginas")
|
|
if not self.page_ranges or not self.page_ranges.strip():
|
|
return list(range(1, total_pages + 1))
|
|
|
|
selected: list[int] = []
|
|
for raw_part in self.page_ranges.split(","):
|
|
part = raw_part.strip()
|
|
if not part:
|
|
continue
|
|
if "-" in part:
|
|
start_text, end_text = [piece.strip() for piece in part.split("-", 1)]
|
|
if not start_text.isdigit() or not end_text.isdigit():
|
|
raise RuntimeError(f"Rango de paginas invalido: {part}")
|
|
start = int(start_text)
|
|
end = int(end_text)
|
|
if start > end:
|
|
raise RuntimeError(f"Rango de paginas invertido: {part}")
|
|
selected.extend(range(start, end + 1))
|
|
else:
|
|
if not part.isdigit():
|
|
raise RuntimeError(f"Pagina invalida: {part}")
|
|
selected.append(int(part))
|
|
|
|
unique_pages = sorted(set(selected))
|
|
invalid = [page for page in unique_pages if page < 1 or page > total_pages]
|
|
if invalid:
|
|
raise RuntimeError(f"Paginas fuera de rango: {', '.join(map(str, invalid))}. El PDF tiene {total_pages} paginas")
|
|
if not unique_pages:
|
|
raise RuntimeError("No se seleccionaron paginas para procesar")
|
|
return unique_pages
|