File size: 3,475 Bytes
d7d27a3 f39e265 9d8bbd6 f39e265 d7d27a3 7adb4f3 9d8bbd6 f39e265 9d8bbd6 7adb4f3 9d8bbd6 7adb4f3 9d8bbd6 d7d27a3 9d8bbd6 | 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 101 102 103 104 105 106 107 108 109 110 111 112 113 | import logging
from functools import lru_cache
from pathlib import Path
SEPARATOR = "\n\n" + "-" * 60 + "\n\n"
SUPPORTED_UPLOAD_EXTENSIONS = {".pdf", ".docx", ".md", ".txt"}
logger = logging.getLogger(__name__)
@lru_cache(maxsize=1)
def _get_ocr_engine():
from rapidocr import RapidOCR
return RapidOCR()
def _convert_pdf_with_ocr(path):
"""Render a PDF page-by-page and extract text with RapidOCR."""
import pypdfium2
ocr = _get_ocr_engine()
pdf = pypdfium2.PdfDocument(str(path))
pages = []
try:
for page_number in range(len(pdf)):
page = pdf[page_number]
try:
image = page.render(scale=2).to_numpy()
result = ocr(image)
if result is not None and result.txts:
pages.append("\n".join(text for text in result.txts if text.strip()))
finally:
page.close()
finally:
pdf.close()
return "\n\n".join(pages)
def _convert_pdf(path):
"""Prefer MarkItDown and use OCR when extraction fails or returns no text."""
try:
from markitdown import MarkItDown
content = MarkItDown().convert(str(path)).markdown
except Exception as exc:
markitdown_error = exc
logger.warning(
"MarkItDown failed for PDF %s; trying RapidOCR: %s",
path.name,
exc,
)
else:
if content and content.strip():
return content
logger.info("MarkItDown returned no text for PDF %s; trying RapidOCR.", path.name)
markitdown_error = None
try:
return _convert_pdf_with_ocr(path)
except Exception as ocr_exc:
markitdown_detail = (
f"MarkItDown failed: {markitdown_error}"
if markitdown_error
else "MarkItDown returned no readable text"
)
raise RuntimeError(
f"{markitdown_detail}; RapidOCR could not read PDF "
f"'{path.name}': {ocr_exc}"
) from ocr_exc
def format_documents(documents):
"""Combine named Markdown strings into the profile context."""
sections = [
f"# Source: {Path(name).stem}\n\n{content.strip()}"
for name, content in documents
if content and content.strip()
]
if not sections:
raise ValueError("The uploaded documents contain no readable text.")
return SEPARATOR.join(sections)
def convert_uploaded_documents(files):
"""Convert supported uploaded files to one Markdown profile in memory."""
documents = []
converter = None
for file_path in files or []:
path = Path(file_path)
extension = path.suffix.lower()
if extension not in SUPPORTED_UPLOAD_EXTENSIONS:
supported = ", ".join(sorted(SUPPORTED_UPLOAD_EXTENSIONS))
raise ValueError(
f"Unsupported file type '{extension or '(none)'}' for {path.name}. "
f"Upload files with these extensions: {supported}."
)
if extension in {".md", ".txt"}:
content = path.read_text(encoding="utf-8-sig")
elif extension == ".pdf":
content = _convert_pdf(path)
else:
if converter is None:
from markitdown import MarkItDown
converter = MarkItDown()
content = converter.convert(str(path)).markdown
documents.append((path.name, content))
return format_documents(documents)
|