mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 08:16:38 +00:00
refactor: rename the application package to docsgpt
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
This commit is contained in:
1 parent
3e7f1f91b4
commit
574f96341e
984 files changed
+9847
-9682
No files matched your search
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,417 @@
|
||||
"""anydoc parser.
|
||||
|
||||
Converts documents to GitHub-Flavored Markdown with `firecrawl-anydoc
|
||||
<https://github.com/firecrawl/anydoc>`_: a single Rust extension with no ML
|
||||
models and no Python dependencies. Measured against docling on a 20-document
|
||||
corpus it was 20-160x faster per file (1.2 s vs 193 s on a 150-page annual
|
||||
report), imported in 6 ms instead of 3.9 s, and peaked at 107 MB RSS instead
|
||||
of 2.3 GB, with equivalent output on clean office files, CSV and text-layer
|
||||
PDFs.
|
||||
|
||||
anydoc never OCRs. It *detects* a scanned or image-only PDF and raises
|
||||
``UnsupportedError`` ("OCR is required"), which is the routing point: a
|
||||
parser given a ``fallback_parser`` (docling when installed, the native OCR
|
||||
parser when OCR is on without docling, otherwise the legacy parser for the
|
||||
suffix) delegates there; without one the file fails loudly as
|
||||
``DocumentParseError`` rather than being stored empty.
|
||||
"""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.parser.file.base_parser import (
|
||||
BaseParser,
|
||||
DocumentParseError,
|
||||
delegate_parse,
|
||||
module_available,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# A fallback parse (scanned-PDF delegation or trust-check re-parse) yielding
|
||||
# fewer stripped characters than this is an empty document in the making,
|
||||
# not content.
|
||||
_MIN_SCAN_FALLBACK_CHARS = 50
|
||||
|
||||
# Suffixes anydoc 0.2.3 converts, verified with ``anydoc.format_from_extension``.
|
||||
# ``.epub`` is deliberately absent: it stays on ``EpubParser``.
|
||||
ANYDOC_SUFFIXES: Tuple[str, ...] = (
|
||||
# Word processing
|
||||
".pdf",
|
||||
".docx",
|
||||
".docm",
|
||||
".doc",
|
||||
".odt",
|
||||
".rtf",
|
||||
# Presentations
|
||||
".pptx",
|
||||
".pptm",
|
||||
".ppsx",
|
||||
".ppsm",
|
||||
".ppt",
|
||||
".pps",
|
||||
".pot",
|
||||
".odp",
|
||||
# Spreadsheets / tabular
|
||||
".xlsx",
|
||||
".xlsm",
|
||||
".xlsb",
|
||||
".xls",
|
||||
".ods",
|
||||
".csv",
|
||||
)
|
||||
|
||||
|
||||
# The five formats every engine has its own parser for (docling / legacy);
|
||||
# under the anydoc engine they get that parser as the fallback. Everything
|
||||
# else in ``ANYDOC_SUFFIXES`` is anydoc-only ("gained") and is mapped to a
|
||||
# fallback-less ``AnydocParser`` in every parser map.
|
||||
_CORE_SUFFIXES = frozenset({".pdf", ".docx", ".pptx", ".xlsx", ".csv"})
|
||||
|
||||
# Spreadsheets anydoc may refuse on its fixed, non-configurable limits
|
||||
# (2M XML nodes per part, 128 MiB decompressed per part — a ~265k-row sheet
|
||||
# trips them) that the plain tabular parsers read fine in a few hundred MB.
|
||||
# For these a ResourceLimitError is a reason to delegate, not to fail.
|
||||
_TABULAR_SUFFIXES = frozenset({".csv", ".xlsx", ".xlsm", ".xlsb", ".xls", ".ods"})
|
||||
ANYDOC_GAINED_SUFFIXES: Tuple[str, ...] = tuple(
|
||||
suffix for suffix in ANYDOC_SUFFIXES if suffix not in _CORE_SUFFIXES
|
||||
)
|
||||
|
||||
|
||||
def anydoc_available() -> bool:
|
||||
"""Whether the ``anydoc`` extension can be imported."""
|
||||
return module_available("anydoc")
|
||||
|
||||
|
||||
def _result_chars(result: Union[str, List[str]]) -> int:
|
||||
"""Stripped character count of a parser result (str or list of rows)."""
|
||||
if isinstance(result, list):
|
||||
return sum(len(str(part).strip()) for part in result)
|
||||
return len(str(result).strip())
|
||||
|
||||
|
||||
def _is_docling_backed(parser: Optional[BaseParser]) -> bool:
|
||||
"""True when ``parser`` is a docling parser (worth a trust-check re-parse).
|
||||
|
||||
Looks through ``NativeOcrPdfParser``, which under ``OCR_BACKEND=native``
|
||||
wraps the docling PDF parser as its ``text_parser`` and hands text-layer
|
||||
PDFs — the only kind the trust check flags — straight to it.
|
||||
"""
|
||||
if parser is None:
|
||||
return False
|
||||
try:
|
||||
from docsgpt.parser.file.docling_parser import DoclingParser
|
||||
except ImportError:
|
||||
return False
|
||||
if isinstance(parser, DoclingParser):
|
||||
return True
|
||||
return isinstance(getattr(parser, "text_parser", None), DoclingParser)
|
||||
|
||||
|
||||
class AnydocParser(BaseParser):
|
||||
"""Markdown conversion via ``anydoc.to_markdown`` with a typed fallback.
|
||||
|
||||
Attributes:
|
||||
fallback_parser: Parser used when anydoc cannot convert the file
|
||||
(scanned PDF, unsupported or malformed input). ``None`` means
|
||||
such files raise ``DocumentParseError``.
|
||||
last_engine: Name of the engine that produced the most recent parse:
|
||||
``"anydoc"``, or the fallback parser's ``last_engine`` / class
|
||||
name when it was delegated to. Mirrors ``PdfiumTextParser`` so
|
||||
the attachment worker records what actually ran.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fallback_parser: Optional[BaseParser] = None,
|
||||
parser_config: Optional[Dict] = None,
|
||||
) -> None:
|
||||
super().__init__(parser_config)
|
||||
self.fallback_parser = fallback_parser
|
||||
self.last_engine: Optional[str] = None
|
||||
# (path, problems) from the most recent parse whose output the trust
|
||||
# check flagged but kept — surfaced via ``get_file_metadata`` as
|
||||
# ``parse_warnings`` so the document carries its own caveat.
|
||||
self._last_warnings: Optional[Tuple[Path, List[str]]] = None
|
||||
# (path, count) of scanned pages OCR'd into the most recent parse.
|
||||
self._last_ocr_pages: Optional[Tuple[Path, int]] = None
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
# Import for real rather than trusting ``find_spec``: a wheel whose
|
||||
# native extension fails to load (glibc/arch mismatch) has a spec but
|
||||
# no importable module, and that must fail here — where the ingest
|
||||
# task treats it as a setup error — not as a bare ImportError from
|
||||
# ``parse_file`` mid-batch, which ``load_data`` does not catch.
|
||||
try:
|
||||
import anydoc # noqa: F401
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"firecrawl-anydoc is required for AnydocParser. "
|
||||
"Install it with: pip install firecrawl-anydoc"
|
||||
) from exc
|
||||
fallback = self.fallback_parser
|
||||
return {"fallback_parser": type(fallback).__name__ if fallback else None}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Convert ``file`` to Markdown.
|
||||
|
||||
Args:
|
||||
file: Path to the document.
|
||||
errors: Decoding error policy; anydoc decodes internally, so this
|
||||
is only forwarded to the fallback parser.
|
||||
|
||||
Returns:
|
||||
The document as Markdown.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: When neither anydoc nor the fallback could
|
||||
produce text for the file, or anydoc hit a resource limit.
|
||||
"""
|
||||
import anydoc
|
||||
|
||||
path = Path(file)
|
||||
self._last_warnings = None
|
||||
self._last_ocr_pages = None
|
||||
try:
|
||||
content = anydoc.to_markdown(str(path))
|
||||
except anydoc.ResourceLimitError as exc:
|
||||
# anydoc refused because the file is too expensive (declared
|
||||
# decompression size, nesting depth, node count). For office
|
||||
# documents a heavier engine would spend exactly what was just
|
||||
# refused, so this is terminal. Spreadsheets are the exception:
|
||||
# anydoc's cell/node limits sit below what ordinary business
|
||||
# exports reach, and the tabular fallback (docling's gate hands
|
||||
# oversized sheets to pandas/openpyxl) is lighter, not heavier.
|
||||
if path.suffix.lower() in _TABULAR_SUFFIXES and self.fallback_parser is not None:
|
||||
return self._delegate(path, errors, f"{type(exc).__name__}: {exc}")
|
||||
self.last_engine = None
|
||||
raise DocumentParseError(
|
||||
f"Failed to parse {path.name}: {type(exc).__name__}: {exc}"
|
||||
) from exc
|
||||
except anydoc.ConvertError as exc:
|
||||
# Typed, per-document failures another engine may still handle:
|
||||
# UnsupportedError (scanned PDF / unknown format), MalformedError,
|
||||
# EncryptedError, MissingPartError.
|
||||
ocr_needed = isinstance(exc, anydoc.UnsupportedError) and "OCR" in str(exc)
|
||||
return self._delegate(
|
||||
path, errors, f"{type(exc).__name__}: {exc}", ocr_needed=ocr_needed
|
||||
)
|
||||
except OSError as exc:
|
||||
raise DocumentParseError(
|
||||
f"Failed to parse {path.name}: the file could not be read."
|
||||
) from exc
|
||||
except Exception as exc:
|
||||
logger.error(f"anydoc failed on {path.name}: {exc}", exc_info=True)
|
||||
raise DocumentParseError(
|
||||
f"Failed to parse {path.name} with anydoc: {exc}"
|
||||
) from exc
|
||||
|
||||
if not content or not content.strip():
|
||||
return self._delegate(path, errors, "anydoc produced no text")
|
||||
|
||||
self.last_engine = "anydoc"
|
||||
if path.suffix.lower() == ".pdf":
|
||||
content = self._finish_pdf(path, content, errors)
|
||||
return content
|
||||
|
||||
def _finish_pdf(self, path: Path, content: str, errors: str) -> Union[str, List[str]]:
|
||||
"""Post-process a successful anydoc PDF conversion.
|
||||
|
||||
Three PDF-only steps:
|
||||
|
||||
1. Trust check (``PDF_TRUST_CHECK``): the two classes where anydoc
|
||||
drops text *silently* — Type0 fonts without ToUnicode, and
|
||||
CJK-declaring PDFs with CJK-less output. A flagged file re-parses
|
||||
on the docling fallback when one is wired; otherwise the anydoc
|
||||
output is kept and the problems ride along as ``parse_warnings``
|
||||
metadata.
|
||||
2. Mixed documents: anydoc only refuses a PDF when *every* page lacks
|
||||
text, so a text document with scanned pages in it converts fine
|
||||
and silently loses those pages. When the fallback parser can OCR
|
||||
(it is wired with OCR on), the pages with no text layer are OCR'd
|
||||
through it and appended (``_ocr_scanned_pages``).
|
||||
3. ``ANYDOC_TABLEIZE`` (off by default): rewrite dot-leader /
|
||||
whitespace table runs as GFM tables. Never applied to a docling
|
||||
re-parse — docling emits real tables already.
|
||||
"""
|
||||
problems = self._trust_problems(path, content)
|
||||
if problems:
|
||||
rerouted = self._reroute_flagged(path, errors, problems)
|
||||
if rerouted is not None:
|
||||
return rerouted
|
||||
self._last_warnings = (path, problems)
|
||||
logger.warning(
|
||||
f"anydoc output for {path.name} failed the PDF trust check "
|
||||
f"({'; '.join(problems)}); keeping it with parse_warnings"
|
||||
)
|
||||
content = self._ocr_scanned_pages(path, content)
|
||||
if settings.ANYDOC_TABLEIZE:
|
||||
from docsgpt.parser.file.tableize import tableize
|
||||
|
||||
content = tableize(content)
|
||||
return content
|
||||
|
||||
def _trust_problems(self, path: Path, content: str) -> List[str]:
|
||||
"""Trust-check findings for ``content``; [] when disabled, clean, or the check errors."""
|
||||
if not settings.PDF_TRUST_CHECK:
|
||||
return []
|
||||
try:
|
||||
from docsgpt.parser.file.pdf_trust import verify_pdf_file
|
||||
|
||||
return verify_pdf_file(path, content)
|
||||
except Exception:
|
||||
# The check exists to catch silent loss; it must never turn a
|
||||
# successful parse into a failure.
|
||||
logger.warning(
|
||||
f"PDF trust check errored on {path.name}; trusting the output",
|
||||
exc_info=True,
|
||||
)
|
||||
return []
|
||||
|
||||
def _reroute_flagged(
|
||||
self, path: Path, errors: str, problems: List[str]
|
||||
) -> Optional[Union[str, List[str]]]:
|
||||
"""Re-parse a trust-flagged PDF with a docling-backed fallback.
|
||||
|
||||
Returns the fallback's output, or None when there is no docling
|
||||
fallback, it fails, or it comes back near-empty — the caller then
|
||||
keeps the anydoc output.
|
||||
Only docling is worth the re-parse: it ships Adobe's predefined
|
||||
CMaps, which is exactly what the flagged font class needs; the
|
||||
legacy pypdf parser does not.
|
||||
"""
|
||||
fallback = self.fallback_parser
|
||||
if not _is_docling_backed(fallback):
|
||||
return None
|
||||
logger.warning(
|
||||
f"anydoc output for {path.name} failed the PDF trust check "
|
||||
f"({'; '.join(problems)}); re-parsing with {type(fallback).__name__}"
|
||||
)
|
||||
try:
|
||||
result = delegate_parse(fallback, path, errors)
|
||||
except DocumentParseError:
|
||||
logger.warning(
|
||||
f"docling re-parse of trust-flagged {path.name} failed; "
|
||||
"keeping the anydoc output",
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
if _result_chars(result) < _MIN_SCAN_FALLBACK_CHARS:
|
||||
# A docling pipeline dropout ('' / '<!-- image -->') returns
|
||||
# without raising; adopting it would swap anydoc's real text for
|
||||
# an empty document.
|
||||
logger.warning(
|
||||
f"docling re-parse of trust-flagged {path.name} returned almost "
|
||||
"no text; keeping the anydoc output"
|
||||
)
|
||||
return None
|
||||
self.last_engine = getattr(fallback, "last_engine", None) or type(fallback).__name__
|
||||
return result
|
||||
|
||||
def _ocr_scanned_pages(self, path: Path, content: str) -> str:
|
||||
"""Append OCR text for the pages of ``path`` that have no text layer.
|
||||
|
||||
Runs only when the fallback parser both has OCR on and exposes
|
||||
``ocr_pages`` (the native OCR PDF parser, or the docling PDF parser).
|
||||
With OCR off the fallback is the plain pypdf parser and mixed
|
||||
documents keep today's behaviour: text pages only. Pages are OCR'd
|
||||
one at a time so a failure on one page costs that page, not the OCR
|
||||
of every other; the anydoc text is always kept — a partial document
|
||||
beats a rejected upload, and the loud-failure path stays reserved
|
||||
for fully scanned files. anydoc's PDF markdown carries no page
|
||||
separators, so the OCR'd pages are appended after the text pages
|
||||
rather than spliced into position.
|
||||
"""
|
||||
fallback = self.fallback_parser
|
||||
if not (getattr(fallback, "ocr_enabled", False) and hasattr(fallback, "ocr_pages")):
|
||||
return content
|
||||
from docsgpt.parser.file.ocr_parser import scanned_page_indices
|
||||
|
||||
indices = scanned_page_indices(path)
|
||||
if not indices:
|
||||
return content
|
||||
logger.info(
|
||||
"%s has %d page(s) without a text layer among its text pages; OCR-ing them with %s",
|
||||
path.name,
|
||||
len(indices),
|
||||
type(fallback).__name__,
|
||||
)
|
||||
try:
|
||||
if not fallback.parser_config_set:
|
||||
fallback.init_parser()
|
||||
except Exception: # noqa: BLE001 - keep the text pages rather than fail the file
|
||||
logger.warning(
|
||||
"OCR of the scanned pages of %s failed; indexing the text pages only",
|
||||
path.name,
|
||||
exc_info=True,
|
||||
)
|
||||
return content
|
||||
extra: List[str] = []
|
||||
for index in indices:
|
||||
try:
|
||||
text = fallback.ocr_pages(path, [index]).get(index, "")
|
||||
except Exception: # noqa: BLE001 - one page's failure must not drop the others
|
||||
logger.warning(
|
||||
"OCR of page %d of %s failed; skipping that page",
|
||||
index + 1,
|
||||
path.name,
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
if text and text.strip():
|
||||
extra.append(text)
|
||||
if not extra:
|
||||
return content
|
||||
self._last_ocr_pages = (path, len(extra))
|
||||
return content.rstrip() + "\n\n" + "\n\n".join(extra) + "\n"
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Surface trust-check findings and OCR'd page counts for the file just parsed."""
|
||||
metadata: Dict = {}
|
||||
last = self._last_warnings
|
||||
if last is not None and last[0] == Path(file):
|
||||
metadata["parse_warnings"] = list(last[1])
|
||||
ocr = self._last_ocr_pages
|
||||
if ocr is not None and ocr[0] == Path(file):
|
||||
metadata["ocr_pages"] = ocr[1]
|
||||
return metadata
|
||||
|
||||
def _delegate(
|
||||
self, path: Path, errors: str, reason: str, ocr_needed: bool = False
|
||||
) -> Union[str, List[str]]:
|
||||
"""Hand ``path`` to the fallback parser, or fail loudly without one.
|
||||
|
||||
``ocr_needed`` marks anydoc's scanned-PDF refusal ("OCR is required").
|
||||
The fallback still runs first — docling extracts text layers anydoc
|
||||
refuses (seen on degenerate CJK text layers) even with OCR off — but
|
||||
when it comes back near-empty the parse fails loudly, telling the
|
||||
user how to get OCR, instead of silently storing an empty document
|
||||
for a scan.
|
||||
"""
|
||||
fallback = self.fallback_parser
|
||||
if fallback is None:
|
||||
self.last_engine = None
|
||||
raise DocumentParseError(f"Failed to parse {path.name}: {reason}")
|
||||
logger.warning(
|
||||
f"anydoc could not convert {path.name} ({reason}); "
|
||||
f"falling back to {type(fallback).__name__}"
|
||||
)
|
||||
result = delegate_parse(fallback, path, errors)
|
||||
if ocr_needed and _result_chars(result) < _MIN_SCAN_FALLBACK_CHARS:
|
||||
self.last_engine = None
|
||||
if getattr(fallback, "ocr_enabled", False):
|
||||
hint = " even with OCR enabled."
|
||||
else:
|
||||
hint = (
|
||||
". Enable OCR to ingest scans: set OCR_ENABLED=true (and "
|
||||
"OCR_ATTACHMENTS_ENABLED for attachments). OCR runs on the "
|
||||
"tesseract binary or a DeepSeek-OCR endpoint (OCR_ENGINE), or "
|
||||
"through the optional docling extra when installed."
|
||||
)
|
||||
raise DocumentParseError(
|
||||
f"{path.name} appears to be a scanned PDF (no text layer), and "
|
||||
f"{type(fallback).__name__} extracted almost nothing{hint}"
|
||||
)
|
||||
self.last_engine = getattr(fallback, "last_engine", None) or type(fallback).__name__
|
||||
return result
|
||||
@@ -0,0 +1,48 @@
|
||||
from pathlib import Path
|
||||
from typing import Dict, Union
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
from docsgpt.stt.stt_creator import STTCreator
|
||||
from docsgpt.stt.upload_limits import enforce_audio_file_size_limit
|
||||
|
||||
|
||||
class AudioParser(BaseParser):
|
||||
def __init__(self, parser_config=None):
|
||||
super().__init__(parser_config=parser_config)
|
||||
self._transcript_metadata: Dict[str, Dict] = {}
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]:
|
||||
_ = errors
|
||||
try:
|
||||
enforce_audio_file_size_limit(file.stat().st_size)
|
||||
except OSError:
|
||||
pass
|
||||
stt = STTCreator.create_stt(settings.STT_PROVIDER)
|
||||
result = stt.transcribe(
|
||||
file,
|
||||
language=settings.STT_LANGUAGE,
|
||||
timestamps=settings.STT_ENABLE_TIMESTAMPS,
|
||||
diarize=settings.STT_ENABLE_DIARIZATION,
|
||||
)
|
||||
|
||||
transcript_metadata = {
|
||||
"transcript_duration_s": result.get("duration_s"),
|
||||
"transcript_language": result.get("language"),
|
||||
"transcript_provider": result.get("provider"),
|
||||
}
|
||||
if result.get("segments"):
|
||||
transcript_metadata["transcript_segments"] = result["segments"]
|
||||
|
||||
self._transcript_metadata[str(file)] = {
|
||||
key: value
|
||||
for key, value in transcript_metadata.items()
|
||||
if value not in (None, [], {})
|
||||
}
|
||||
return result.get("text", "")
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
return self._transcript_metadata.get(str(file), {})
|
||||
@@ -0,0 +1,19 @@
|
||||
"""Base reader class."""
|
||||
from abc import abstractmethod
|
||||
from typing import Any, List
|
||||
|
||||
from docsgpt.parser.schema.base import Document
|
||||
from docsgpt.vectorstore.document_class import Document as VectorDocument
|
||||
|
||||
|
||||
class BaseReader:
|
||||
"""Utilities for loading data from a directory."""
|
||||
|
||||
@abstractmethod
|
||||
def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
|
||||
"""Load data from the input directory."""
|
||||
|
||||
def load_vector_documents(self, **load_kwargs: Any) -> List[VectorDocument]:
|
||||
"""Load data in the vector-store document format."""
|
||||
docs = self.load_data(**load_kwargs)
|
||||
return [d.to_vector_format() for d in docs]
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Base parser and config class."""
|
||||
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DocumentParseError(Exception):
|
||||
"""A file could not be converted to text.
|
||||
|
||||
Raised instead of returning the failure message as the document body: a
|
||||
parser that hands back its own traceback produces an attachment/source
|
||||
whose "content" is an error string, which then reaches the LLM as if it
|
||||
were the document. Callers should treat this as a failed upload and tell
|
||||
the user, rather than storing anything.
|
||||
|
||||
Deterministic for a given input, so it is marked non-retryable on the
|
||||
Celery tasks that parse (alongside ``DataError`` and
|
||||
``AttachmentRejectedError``) — a poison file fails once instead of
|
||||
re-failing identically on every retry.
|
||||
"""
|
||||
|
||||
|
||||
class BaseParser:
|
||||
"""Base class for all parsers."""
|
||||
|
||||
def __init__(self, parser_config: Optional[Dict] = None):
|
||||
"""Init params."""
|
||||
self._parser_config = parser_config
|
||||
|
||||
def init_parser(self) -> None:
|
||||
"""Init parser and store it."""
|
||||
parser_config = self._init_parser()
|
||||
self._parser_config = parser_config
|
||||
|
||||
@property
|
||||
def parser_config_set(self) -> bool:
|
||||
"""Check if parser config is set."""
|
||||
return self._parser_config is not None
|
||||
|
||||
@property
|
||||
def parser_config(self) -> Dict:
|
||||
"""Check if parser config is set."""
|
||||
if self._parser_config is None:
|
||||
raise ValueError("Parser config not set.")
|
||||
return self._parser_config
|
||||
|
||||
@abstractmethod
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Initialize the parser with the config."""
|
||||
|
||||
@abstractmethod
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse file."""
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Return parser-specific metadata for the most recently parsed file."""
|
||||
_ = file
|
||||
return {}
|
||||
|
||||
|
||||
def module_available(name: str) -> bool:
|
||||
"""Whether top-level module ``name`` can be imported, without importing it.
|
||||
|
||||
Cheap enough to call while building a parser map. Honors ``sys.modules``
|
||||
first so tests that stub a module in (or block one out with ``None``)
|
||||
are seen the same way the import system would see them; a stub with no
|
||||
``__spec__`` counts as present rather than raising.
|
||||
|
||||
Args:
|
||||
name: Top-level module name, e.g. ``"docling"``.
|
||||
|
||||
Returns:
|
||||
True when an import of ``name`` would succeed.
|
||||
"""
|
||||
if name in sys.modules:
|
||||
return sys.modules[name] is not None
|
||||
try:
|
||||
return importlib.util.find_spec(name) is not None
|
||||
except (ImportError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def delegate_parse(
|
||||
parser: "BaseParser", file: Path, errors: str
|
||||
) -> Union[str, List[str]]:
|
||||
"""Run a fallback parser, normalizing its failures to ``DocumentParseError``.
|
||||
|
||||
Used wherever one parser hands a file to another (docling to the plain
|
||||
tabular parsers for oversized sheets, anydoc to docling/legacy parsers
|
||||
for files it cannot convert). Wrapping matters twice over: the Celery
|
||||
tasks that own these paths list ``DocumentParseError`` in
|
||||
``dont_autoretry_for``, so a deterministic content error fails once
|
||||
instead of retrying four times; and ``SimpleDirectoryReader.load_data``
|
||||
catches only ``DocumentParseError``, so one poison file is skipped into
|
||||
``failed_files`` instead of aborting an entire multi-file ingest.
|
||||
|
||||
Args:
|
||||
parser: The parser to delegate to; initialized here if needed.
|
||||
file: Path to the file being parsed.
|
||||
errors: Decoding error policy, forwarded to the parser.
|
||||
|
||||
Returns:
|
||||
The parsed text, or list of row strings.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: If the fallback parser fails for any reason.
|
||||
"""
|
||||
try:
|
||||
if not parser.parser_config_set:
|
||||
parser.init_parser()
|
||||
return parser.parse_file(file, errors)
|
||||
except DocumentParseError:
|
||||
raise
|
||||
except ImportError:
|
||||
# A fallback whose dependency is missing is a deployment problem, not
|
||||
# a bad document: let it reach the ingest task's setup-error path
|
||||
# instead of blaming every file with "could not be read".
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Fallback parse of {Path(file).name} with "
|
||||
f"{type(parser).__name__} failed: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise DocumentParseError(
|
||||
f"Failed to parse {Path(file).name}: the file could not be read."
|
||||
) from e
|
||||
@@ -0,0 +1,650 @@
|
||||
"""Simple reader that reads files of different formats from a directory."""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from docsgpt.parser.file.base import BaseReader
|
||||
from docsgpt.parser.file.base_parser import BaseParser, DocumentParseError, module_available
|
||||
from docsgpt.parser.file.docs_parser import DocxParser, PDFParser
|
||||
from docsgpt.parser.file.epub_parser import EpubParser
|
||||
from docsgpt.parser.file.html_parser import HTMLParser
|
||||
from docsgpt.parser.file.markdown_parser import MarkdownParser
|
||||
from docsgpt.parser.file.rst_parser import RstParser
|
||||
from docsgpt.parser.file.tabular_parser import PandasCSVParser, ExcelParser
|
||||
from docsgpt.parser.file.json_parser import JSONParser
|
||||
from docsgpt.parser.file.pptx_parser import PPTXParser
|
||||
from docsgpt.parser.file.image_parser import ImageParser
|
||||
from docsgpt.parser.file.audio_parser import AudioParser
|
||||
from docsgpt.parser.schema.base import Document
|
||||
from docsgpt.stt.constants import SUPPORTED_AUDIO_EXTENSIONS
|
||||
from docsgpt.utils import num_tokens_from_string
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.core.optional_deps import install_hint
|
||||
|
||||
|
||||
def _build_audio_parser_mapping() -> Dict[str, BaseParser]:
|
||||
return {extension: AudioParser() for extension in SUPPORTED_AUDIO_EXTENSIONS}
|
||||
|
||||
|
||||
def _wrap_pdf_fast_path(pdf_parser: BaseParser) -> BaseParser:
|
||||
"""Put the pypdfium2 text-layer parser in front of ``pdf_parser``.
|
||||
|
||||
Args:
|
||||
pdf_parser: Parser to fall back to for PDFs with no usable text layer.
|
||||
|
||||
Returns:
|
||||
BaseParser: The fast-path parser, or ``pdf_parser`` unchanged when
|
||||
pypdfium2 is unavailable.
|
||||
"""
|
||||
try:
|
||||
from docsgpt.parser.file.pdfium_parser import PdfiumTextParser
|
||||
except ImportError:
|
||||
logging.warning(
|
||||
"pypdfium2 is not installed; PDF attachments will use %s",
|
||||
type(pdf_parser).__name__,
|
||||
)
|
||||
return pdf_parser
|
||||
return PdfiumTextParser(
|
||||
fallback_parser=pdf_parser,
|
||||
min_median_chars=settings.ATTACHMENT_PDF_TEXT_MIN_MEDIAN_CHARS,
|
||||
)
|
||||
|
||||
|
||||
def _gained_format_entries() -> Dict[str, BaseParser]:
|
||||
"""Anydoc-only formats (legacy/macro Office, OpenDocument, RTF) for every map.
|
||||
|
||||
anydoc is a core dependency, so these suffixes are parseable under both
|
||||
engines. When anydoc is somehow missing the entries are omitted and
|
||||
``SimpleDirectoryReader.load_data`` rejects such files with a
|
||||
``DocumentParseError`` naming the install, rather than reading OLE/zip
|
||||
bytes as text.
|
||||
"""
|
||||
from docsgpt.parser.file.anydoc_parser import (
|
||||
ANYDOC_GAINED_SUFFIXES,
|
||||
AnydocParser,
|
||||
anydoc_available,
|
||||
)
|
||||
|
||||
if not anydoc_available():
|
||||
return {}
|
||||
return {suffix: AnydocParser() for suffix in ANYDOC_GAINED_SUFFIXES}
|
||||
|
||||
|
||||
def _native_ocr_parsers(
|
||||
text_parser: Optional[BaseParser] = None,
|
||||
) -> Optional[Tuple[BaseParser, Callable[[], BaseParser]]]:
|
||||
"""PDF parser and image-parser factory for the native OCR backend.
|
||||
|
||||
Args:
|
||||
text_parser: Parser for PDFs whose every page has a text layer (the
|
||||
docling PDF parser under the docling engine); None reads text
|
||||
layers with pypdfium2.
|
||||
|
||||
Returns:
|
||||
``(pdf_parser, image_parser_factory)``, or None when pypdfium2 or
|
||||
Pillow is missing — both are core dependencies, so that is a broken
|
||||
install, logged as such.
|
||||
"""
|
||||
from docsgpt.parser.file.ocr_parser import (
|
||||
NativeOcrImageParser,
|
||||
NativeOcrPdfParser,
|
||||
native_ocr_available,
|
||||
)
|
||||
|
||||
if not native_ocr_available():
|
||||
logging.error(
|
||||
"OCR is enabled but pypdfium2/Pillow are not importable, so the native "
|
||||
"OCR backend cannot run; scanned PDFs and images will not be OCR'd"
|
||||
)
|
||||
return None
|
||||
return NativeOcrPdfParser(text_parser=text_parser), NativeOcrImageParser
|
||||
|
||||
|
||||
def _image_entries(factory: Callable[[], BaseParser], suffixes) -> Dict[str, BaseParser]:
|
||||
"""One parser instance per image suffix."""
|
||||
return {suffix: factory() for suffix in suffixes}
|
||||
|
||||
|
||||
_LEGACY_IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg")
|
||||
|
||||
|
||||
def _legacy_file_extractor(pdf_text_fast_path: bool = False, ocr_enabled: bool = False) -> Dict[str, BaseParser]:
|
||||
"""Parser map that needs neither docling nor anydoc.
|
||||
|
||||
Args:
|
||||
pdf_text_fast_path: Put the pypdfium2 text-layer parser in front of
|
||||
the PDF parser.
|
||||
ocr_enabled: Use the native OCR parsers (tesseract / DeepSeek-OCR via
|
||||
pypdfium2 + Pillow) for PDFs and images instead of the plain
|
||||
pypdf / remote-service parsers.
|
||||
|
||||
Returns:
|
||||
Dict[str, BaseParser]: Parser keyed by lower-case file suffix.
|
||||
"""
|
||||
from docsgpt.parser.file.ocr_parser import IMAGE_SUFFIXES
|
||||
|
||||
pdf_parser: BaseParser = PDFParser()
|
||||
images: Dict[str, BaseParser] = _image_entries(ImageParser, _LEGACY_IMAGE_SUFFIXES)
|
||||
if ocr_enabled:
|
||||
native = _native_ocr_parsers()
|
||||
if native is not None:
|
||||
pdf_parser, image_factory = native
|
||||
images = _image_entries(image_factory, sorted(IMAGE_SUFFIXES))
|
||||
if pdf_text_fast_path:
|
||||
pdf_parser = _wrap_pdf_fast_path(pdf_parser)
|
||||
return {
|
||||
".pdf": pdf_parser,
|
||||
".docx": DocxParser(),
|
||||
".csv": PandasCSVParser(),
|
||||
".xlsx": ExcelParser(),
|
||||
".epub": EpubParser(),
|
||||
".md": MarkdownParser(),
|
||||
".rst": RstParser(),
|
||||
".html": HTMLParser(),
|
||||
".xhtml": HTMLParser(),
|
||||
".mdx": MarkdownParser(),
|
||||
".json": JSONParser(),
|
||||
".pptx": PPTXParser(),
|
||||
**images,
|
||||
**_build_audio_parser_mapping(),
|
||||
**_gained_format_entries(),
|
||||
}
|
||||
|
||||
|
||||
def _docling_file_extractor(
|
||||
ocr_enabled: bool,
|
||||
pdf_text_fast_path: bool = False,
|
||||
*,
|
||||
missing_log_level: int = logging.WARNING,
|
||||
) -> Dict[str, BaseParser]:
|
||||
"""docling parser map, or the legacy map when docling is not installed.
|
||||
|
||||
With OCR on, ``OCR_BACKEND`` decides who OCRs: docling's own hybrid
|
||||
pipeline, or the native parsers (``ocr_parser.py``) with the docling PDF
|
||||
parser kept — OCR off — for PDFs that have a text layer on every page.
|
||||
Without docling installed the legacy map takes over, itself using the
|
||||
native OCR parsers when OCR is on.
|
||||
|
||||
The availability check has to happen here, not at parse time:
|
||||
``DoclingParser`` only imports docling inside ``_init_parser``, and the
|
||||
``ImportError`` that raises is not a ``DocumentParseError``, so a map
|
||||
that handed out docling parsers on a docling-less install would abort
|
||||
the whole ingest on its first file instead of degrading.
|
||||
|
||||
Args:
|
||||
ocr_enabled: Enable OCR in the docling PDF/image parsers.
|
||||
pdf_text_fast_path: Put the pypdfium2 text-layer parser in front of
|
||||
the PDF parser (attachments only; see ``get_default_file_extractor``).
|
||||
missing_log_level: Level at which a missing docling install is logged;
|
||||
informational when docling is only the *fallback* engine.
|
||||
|
||||
Returns:
|
||||
Dict[str, BaseParser]: Parser keyed by lower-case file suffix.
|
||||
"""
|
||||
try:
|
||||
if not module_available("docling"):
|
||||
raise ImportError("docling is not installed")
|
||||
from docsgpt.parser.file.docling_parser import (
|
||||
DoclingPDFParser,
|
||||
DoclingDocxParser,
|
||||
DoclingPPTXParser,
|
||||
DoclingXLSXParser,
|
||||
DoclingHTMLParser,
|
||||
DoclingImageParser,
|
||||
DoclingCSVParser,
|
||||
DoclingAsciiDocParser,
|
||||
DoclingVTTParser,
|
||||
DoclingXMLParser,
|
||||
)
|
||||
except ImportError:
|
||||
logging.log(
|
||||
missing_log_level,
|
||||
"docling is not installed. Using standard parsers%s. For layout-model "
|
||||
"parsing, install the docling extra: %s",
|
||||
" with native OCR" if ocr_enabled else "",
|
||||
install_hint("docling"),
|
||||
)
|
||||
return _legacy_file_extractor(pdf_text_fast_path, ocr_enabled=ocr_enabled)
|
||||
|
||||
from docsgpt.parser.file.ocr_parser import IMAGE_SUFFIXES, resolve_ocr_backend
|
||||
|
||||
native = None
|
||||
if ocr_enabled and resolve_ocr_backend() == "native":
|
||||
native = _native_ocr_parsers(text_parser=DoclingPDFParser(ocr_enabled=False))
|
||||
if native is not None:
|
||||
pdf_parser, image_factory = native
|
||||
images = _image_entries(image_factory, sorted(IMAGE_SUFFIXES))
|
||||
else:
|
||||
pdf_parser = DoclingPDFParser(ocr_enabled=ocr_enabled)
|
||||
# Images (with OCR) - only use Docling when OCR is enabled
|
||||
image_factory = (
|
||||
(lambda: DoclingImageParser(ocr_enabled=True)) if ocr_enabled else ImageParser
|
||||
)
|
||||
images = _image_entries(image_factory, sorted(IMAGE_SUFFIXES))
|
||||
if pdf_text_fast_path:
|
||||
pdf_parser = _wrap_pdf_fast_path(pdf_parser)
|
||||
return {
|
||||
# Documents
|
||||
".pdf": pdf_parser,
|
||||
".docx": DoclingDocxParser(),
|
||||
".pptx": DoclingPPTXParser(),
|
||||
".xlsx": DoclingXLSXParser(),
|
||||
# Web formats
|
||||
".html": DoclingHTMLParser(),
|
||||
".xhtml": DoclingHTMLParser(),
|
||||
# Data formats
|
||||
".csv": DoclingCSVParser(),
|
||||
".json": JSONParser(), # Keep JSON parser (specialized handling)
|
||||
# Text/markup formats
|
||||
".md": MarkdownParser(), # Keep markdown parser (specialized handling)
|
||||
".mdx": MarkdownParser(),
|
||||
".rst": RstParser(),
|
||||
".adoc": DoclingAsciiDocParser(),
|
||||
".asciidoc": DoclingAsciiDocParser(),
|
||||
# Images: OCR'd by the selected backend when OCR is on, else ImageParser
|
||||
**images,
|
||||
# Media/subtitles
|
||||
".vtt": DoclingVTTParser(),
|
||||
**_build_audio_parser_mapping(),
|
||||
# Specialized XML formats
|
||||
".xml": DoclingXMLParser(),
|
||||
# Formats docling doesn't support - use standard parsers
|
||||
".epub": EpubParser(),
|
||||
# Formats only anydoc reads (legacy/macro Office, OpenDocument, RTF)
|
||||
**_gained_format_entries(),
|
||||
}
|
||||
|
||||
|
||||
def _anydoc_file_extractor(ocr_enabled: bool, pdf_text_fast_path: bool = False) -> Dict[str, BaseParser]:
|
||||
"""anydoc parser map, layered over the docling (or legacy) map.
|
||||
|
||||
anydoc takes every format it converts; each of those parsers gets the
|
||||
base map's parser for the same suffix as ``fallback_parser``, so a
|
||||
scanned PDF still reaches OCR — docling's when it is installed and
|
||||
selected, the native OCR parser otherwise — and the legacy parser when
|
||||
OCR is off. docling, when installed, also keeps serving what anydoc
|
||||
cannot: ``.adoc``/``.vtt``/``.xml``. Images go to whichever OCR backend
|
||||
is active. HTML goes to the markdownify-based ``HTMLMarkdownParser``.
|
||||
|
||||
``pdf_text_fast_path`` is deliberately ignored here: anydoc already reads
|
||||
the text layer in milliseconds *and* keeps headings/tables, and does its
|
||||
own scanned-PDF detection, so putting pypdfium2 in front of it would
|
||||
only lose structure.
|
||||
|
||||
Args:
|
||||
ocr_enabled: Enable OCR in the docling fallback/image parsers.
|
||||
pdf_text_fast_path: Honored only when anydoc is not installed and the
|
||||
map degrades to the base engine.
|
||||
|
||||
Returns:
|
||||
Dict[str, BaseParser]: Parser keyed by lower-case file suffix.
|
||||
"""
|
||||
from docsgpt.parser.file.anydoc_parser import (
|
||||
ANYDOC_SUFFIXES,
|
||||
AnydocParser,
|
||||
anydoc_available,
|
||||
)
|
||||
from docsgpt.parser.file.html_parser import HTMLMarkdownParser
|
||||
|
||||
if not anydoc_available():
|
||||
logging.warning(
|
||||
"firecrawl-anydoc is not installed; DOC_PARSER_ENGINE=anydoc is using "
|
||||
"the %s parsers instead. Install with: pip install firecrawl-anydoc",
|
||||
"docling" if module_available("docling") else "standard",
|
||||
)
|
||||
return _docling_file_extractor(ocr_enabled, pdf_text_fast_path)
|
||||
|
||||
base = _docling_file_extractor(
|
||||
ocr_enabled, pdf_text_fast_path=False, missing_log_level=logging.INFO
|
||||
)
|
||||
extractor = dict(base)
|
||||
for suffix in ANYDOC_SUFFIXES:
|
||||
fallback = base.get(suffix)
|
||||
if isinstance(fallback, AnydocParser):
|
||||
# A gained-format entry from the base map is already a
|
||||
# fallback-less AnydocParser — anydoc delegating to anydoc would
|
||||
# just repeat the same failure.
|
||||
continue
|
||||
extractor[suffix] = AnydocParser(fallback_parser=fallback)
|
||||
extractor[".html"] = HTMLMarkdownParser()
|
||||
extractor[".xhtml"] = HTMLMarkdownParser()
|
||||
return extractor
|
||||
|
||||
|
||||
def get_default_file_extractor(
|
||||
ocr_enabled: Optional[bool] = None,
|
||||
pdf_text_fast_path: bool = False,
|
||||
engine: Optional[str] = None,
|
||||
) -> Dict[str, BaseParser]:
|
||||
"""Get the default file extractor.
|
||||
|
||||
The engine is ``settings.DOC_PARSER_ENGINE`` unless overridden:
|
||||
|
||||
* ``anydoc`` (default): firecrawl-anydoc for office/PDF/tabular formats,
|
||||
markdownify for HTML, the active OCR backend for images and as the
|
||||
fallback for scanned PDFs; docling (if installed) for the rest anydoc
|
||||
rejects, legacy parsers otherwise.
|
||||
* ``docling``: the docling map, with the legacy parsers when docling is
|
||||
not installed.
|
||||
|
||||
Args:
|
||||
ocr_enabled: Enable OCR for scanned PDFs and images, through the
|
||||
backend ``OCR_BACKEND`` resolves to. Defaults to
|
||||
``settings.OCR_ENABLED``.
|
||||
pdf_text_fast_path: Read PDFs via their embedded text layer
|
||||
(pypdfium2) instead of docling, falling back to docling for scans.
|
||||
Off by default: it trades docling's structural markdown for speed,
|
||||
which suits attachments (read into a prompt) but not source
|
||||
ingestion (chunked and embedded for retrieval). A no-op under the
|
||||
anydoc engine, which is already that fast and keeps structure.
|
||||
engine: ``"anydoc"`` or ``"docling"``; ``None`` reads the setting.
|
||||
|
||||
Returns:
|
||||
Dict[str, BaseParser]: Parser keyed by lower-case file suffix.
|
||||
"""
|
||||
if ocr_enabled is None:
|
||||
ocr_enabled = settings.OCR_ENABLED
|
||||
selected = (engine or getattr(settings, "DOC_PARSER_ENGINE", None) or "anydoc")
|
||||
selected = str(selected).strip().lower()
|
||||
if selected == "docling":
|
||||
return _docling_file_extractor(ocr_enabled, pdf_text_fast_path)
|
||||
if selected != "anydoc":
|
||||
logging.warning("Unknown DOC_PARSER_ENGINE %r; using anydoc", selected)
|
||||
return _anydoc_file_extractor(ocr_enabled, pdf_text_fast_path)
|
||||
|
||||
|
||||
# For backwards compatibility
|
||||
DEFAULT_FILE_EXTRACTOR: Dict[str, BaseParser] = get_default_file_extractor()
|
||||
|
||||
|
||||
class SimpleDirectoryReader(BaseReader):
|
||||
"""Simple directory reader.
|
||||
|
||||
Can read files into separate documents, or concatenates
|
||||
files into one document text.
|
||||
|
||||
Args:
|
||||
input_dir (str): Path to the directory.
|
||||
input_files (List): List of file paths to read (Optional; overrides input_dir)
|
||||
exclude_hidden (bool): Whether to exclude hidden files (dotfiles).
|
||||
errors (str): how encoding and decoding errors are to be handled,
|
||||
see https://docs.python.org/3/library/functions.html#open
|
||||
recursive (bool): Whether to recursively search in subdirectories.
|
||||
False by default.
|
||||
required_exts (Optional[List[str]]): List of required extensions.
|
||||
Default is None.
|
||||
file_extractor (Optional[Dict[str, BaseParser]]): A mapping of file
|
||||
extension to a BaseParser class that specifies how to convert that file
|
||||
to text. See DEFAULT_FILE_EXTRACTOR.
|
||||
num_files_limit (Optional[int]): Maximum number of files to read.
|
||||
Default is None.
|
||||
file_metadata (Optional[Callable[str, Dict]]): A function that takes
|
||||
in a filename and returns a Dict of metadata for the Document.
|
||||
Default is None.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dir: Optional[str] = None,
|
||||
input_files: Optional[List] = None,
|
||||
exclude_hidden: bool = True,
|
||||
errors: str = "ignore",
|
||||
recursive: bool = True,
|
||||
required_exts: Optional[List[str]] = None,
|
||||
file_extractor: Optional[Dict[str, BaseParser]] = None,
|
||||
num_files_limit: Optional[int] = None,
|
||||
file_metadata: Optional[Callable[[str], Dict]] = None,
|
||||
) -> None:
|
||||
"""Initialize with parameters."""
|
||||
super().__init__()
|
||||
|
||||
if not input_dir and not input_files:
|
||||
raise ValueError("Must provide either `input_dir` or `input_files`.")
|
||||
|
||||
self.errors = errors
|
||||
|
||||
self.recursive = recursive
|
||||
self.exclude_hidden = exclude_hidden
|
||||
# Normalize extensions to lowercase for case-insensitive matching
|
||||
self.required_exts = (
|
||||
[ext.lower() for ext in required_exts] if required_exts else None
|
||||
)
|
||||
self.num_files_limit = num_files_limit
|
||||
|
||||
if input_files:
|
||||
self.input_files = []
|
||||
for path in input_files:
|
||||
print(path)
|
||||
input_file = Path(path)
|
||||
self.input_files.append(input_file)
|
||||
elif input_dir:
|
||||
self.input_dir = Path(input_dir)
|
||||
self.input_files = self._add_files(self.input_dir)
|
||||
|
||||
self.file_extractor = file_extractor or DEFAULT_FILE_EXTRACTOR
|
||||
self.file_metadata = file_metadata
|
||||
# (path, message) per file skipped by ``load_data`` as unparseable.
|
||||
self.failed_files: List[Tuple[Path, str]] = []
|
||||
|
||||
def _add_files(self, input_dir: Path) -> List[Path]:
|
||||
"""Add files."""
|
||||
input_files = sorted(input_dir.iterdir())
|
||||
new_input_files = []
|
||||
dirs_to_explore = []
|
||||
for input_file in input_files:
|
||||
if input_file.is_dir():
|
||||
if self.recursive:
|
||||
dirs_to_explore.append(input_file)
|
||||
elif self.exclude_hidden and input_file.name.startswith("."):
|
||||
continue
|
||||
elif (
|
||||
self.required_exts is not None
|
||||
and input_file.suffix.lower() not in self.required_exts
|
||||
):
|
||||
continue
|
||||
else:
|
||||
new_input_files.append(input_file)
|
||||
|
||||
for dir_to_explore in dirs_to_explore:
|
||||
sub_input_files = self._add_files(dir_to_explore)
|
||||
new_input_files.extend(sub_input_files)
|
||||
|
||||
if self.num_files_limit is not None and self.num_files_limit > 0:
|
||||
new_input_files = new_input_files[0: self.num_files_limit]
|
||||
|
||||
# print total number of files added
|
||||
logging.debug(
|
||||
f"> [SimpleDirectoryReader] Total files added: {len(new_input_files)}"
|
||||
)
|
||||
|
||||
return new_input_files
|
||||
|
||||
def load_data(
|
||||
self,
|
||||
concatenate: bool = False,
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None,
|
||||
) -> List[Document]:
|
||||
"""Load data from the input directory.
|
||||
|
||||
Args:
|
||||
concatenate (bool): whether to concatenate all files into one document.
|
||||
If set to True, file metadata is ignored.
|
||||
False by default.
|
||||
progress_callback (Optional[Callable[[int, int], None]]): Called
|
||||
after each file is parsed with ``(files_done, total_files)``.
|
||||
Lets callers surface parse/OCR progress before embedding
|
||||
begins. Exceptions raised by the callback are swallowed so
|
||||
progress reporting can never fail ingestion.
|
||||
|
||||
Returns:
|
||||
List[Document]: A list of documents.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: if no input file could be parsed. Individual
|
||||
unreadable files are skipped and recorded in ``failed_files``
|
||||
so one corrupt document does not cost the caller the rest of
|
||||
the batch; a single-file read (the attachment path) still
|
||||
raises, since skipping there would only defer the failure to
|
||||
an empty result.
|
||||
"""
|
||||
data: Union[str, List[str]] = ""
|
||||
data_list: List[str] = []
|
||||
metadata_list = []
|
||||
self.file_token_counts = {}
|
||||
self.failed_files = []
|
||||
|
||||
total_files = len(self.input_files)
|
||||
|
||||
def report_progress(files_done: int) -> None:
|
||||
if progress_callback is None:
|
||||
return
|
||||
try:
|
||||
progress_callback(files_done, total_files)
|
||||
except Exception:
|
||||
logging.warning("load_data progress callback failed", exc_info=True)
|
||||
|
||||
for file_index, input_file in enumerate(self.input_files):
|
||||
suffix_lower = input_file.suffix.lower()
|
||||
parser_metadata = {}
|
||||
try:
|
||||
if suffix_lower in self.file_extractor:
|
||||
parser = self.file_extractor[suffix_lower]
|
||||
if not parser.parser_config_set:
|
||||
parser.init_parser()
|
||||
data = parser.parse_file(input_file, errors=self.errors)
|
||||
parser_metadata = parser.get_file_metadata(input_file)
|
||||
else:
|
||||
from docsgpt.parser.file.anydoc_parser import ANYDOC_GAINED_SUFFIXES
|
||||
|
||||
if suffix_lower in ANYDOC_GAINED_SUFFIXES:
|
||||
# Binary office formats only anydoc reads; without it
|
||||
# the standard read would index OLE/zip bytes as text.
|
||||
raise DocumentParseError(
|
||||
f"No parser is available for {suffix_lower} files: "
|
||||
"firecrawl-anydoc is not installed (pip install firecrawl-anydoc)"
|
||||
)
|
||||
# do standard read
|
||||
with open(input_file, "r", errors=self.errors) as f:
|
||||
data = f.read()
|
||||
except DocumentParseError as e:
|
||||
logging.warning(f"Skipping unreadable file {input_file.name}: {e}")
|
||||
self.failed_files.append((input_file, str(e)))
|
||||
report_progress(file_index + 1)
|
||||
continue
|
||||
|
||||
# Calculate token count for this file
|
||||
if isinstance(data, List):
|
||||
file_tokens = sum(num_tokens_from_string(str(d)) for d in data)
|
||||
else:
|
||||
file_tokens = num_tokens_from_string(str(data))
|
||||
|
||||
full_path = str(input_file.resolve())
|
||||
self.file_token_counts[full_path] = file_tokens
|
||||
|
||||
base_metadata = {
|
||||
'title': input_file.name,
|
||||
'token_count': file_tokens,
|
||||
}
|
||||
if parser_metadata:
|
||||
base_metadata.update(parser_metadata)
|
||||
|
||||
if hasattr(self, 'input_dir'):
|
||||
try:
|
||||
relative_path = str(input_file.relative_to(self.input_dir))
|
||||
base_metadata['source'] = relative_path
|
||||
except ValueError:
|
||||
base_metadata['source'] = str(input_file)
|
||||
else:
|
||||
base_metadata['source'] = str(input_file)
|
||||
|
||||
if self.file_metadata is not None:
|
||||
custom_metadata = self.file_metadata(input_file.name)
|
||||
base_metadata.update(custom_metadata)
|
||||
|
||||
if isinstance(data, List):
|
||||
# Extend data_list with each item in the data list
|
||||
data_list.extend([str(d) for d in data])
|
||||
# copy(): chunking writes token_count into this dict in
|
||||
# place, so a shared reference gives every chunk the last
|
||||
# chunk's count.
|
||||
metadata_list.extend([base_metadata.copy() for _ in data])
|
||||
else:
|
||||
data_list.append(str(data))
|
||||
metadata_list.append(base_metadata.copy())
|
||||
|
||||
report_progress(file_index + 1)
|
||||
|
||||
# Every file failed: there is nothing to ingest, so this is a failed
|
||||
# read rather than an empty one. Callers (the attachment worker, the
|
||||
# ingest tasks) treat it as terminal and tell the user.
|
||||
if self.failed_files and not data_list:
|
||||
if len(self.failed_files) == 1:
|
||||
# Single-file read (the attachment path): the parser's own
|
||||
# message reaches the user verbatim, so don't wrap it in
|
||||
# "None of the 1 file(s)…".
|
||||
raise DocumentParseError(self.failed_files[0][1])
|
||||
names = ", ".join(p.name for p, _ in self.failed_files[:5])
|
||||
raise DocumentParseError(
|
||||
f"None of the {len(self.failed_files)} files could be parsed: "
|
||||
f"{names}{'…' if len(self.failed_files) > 5 else ''} "
|
||||
f"({self.failed_files[0][1]})"
|
||||
)
|
||||
|
||||
# Build directory structure if input_dir is provided
|
||||
if hasattr(self, 'input_dir'):
|
||||
self.directory_structure = self.build_directory_structure(self.input_dir)
|
||||
logging.info("Directory structure built successfully")
|
||||
else:
|
||||
self.directory_structure = {}
|
||||
|
||||
if concatenate:
|
||||
return [Document("\n".join(data_list))]
|
||||
elif self.file_metadata is not None:
|
||||
return [Document(d, extra_info=m) for d, m in zip(data_list, metadata_list)]
|
||||
else:
|
||||
return [Document(d) for d in data_list]
|
||||
|
||||
def build_directory_structure(self, base_path):
|
||||
"""Build a dictionary representing the directory structure.
|
||||
|
||||
Args:
|
||||
base_path: The base path to start building the structure from.
|
||||
|
||||
Returns:
|
||||
dict: A nested dictionary representing the directory structure.
|
||||
"""
|
||||
import mimetypes
|
||||
|
||||
def build_tree(path):
|
||||
"""Helper function to recursively build the directory tree."""
|
||||
result = {}
|
||||
|
||||
for item in path.iterdir():
|
||||
if self.exclude_hidden and item.name.startswith('.'):
|
||||
continue
|
||||
|
||||
if item.is_dir():
|
||||
subtree = build_tree(item)
|
||||
if subtree:
|
||||
result[item.name] = subtree
|
||||
else:
|
||||
if self.required_exts is not None and item.suffix.lower() not in self.required_exts:
|
||||
continue
|
||||
|
||||
full_path = str(item.resolve())
|
||||
file_size_bytes = item.stat().st_size
|
||||
mime_type = mimetypes.guess_type(item.name)[0] or "application/octet-stream"
|
||||
|
||||
file_info = {
|
||||
"type": mime_type,
|
||||
"size_bytes": file_size_bytes
|
||||
}
|
||||
|
||||
if hasattr(self, 'file_token_counts') and full_path in self.file_token_counts:
|
||||
file_info["token_count"] = self.file_token_counts[full_path]
|
||||
|
||||
result[item.name] = file_info
|
||||
|
||||
return result
|
||||
|
||||
return build_tree(Path(base_path))
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Shared file-extension constants for parsing and ingestion flows."""
|
||||
|
||||
import os
|
||||
|
||||
from docsgpt.parser.file.anydoc_parser import ANYDOC_GAINED_SUFFIXES
|
||||
from docsgpt.stt.constants import SUPPORTED_AUDIO_EXTENSIONS
|
||||
|
||||
|
||||
SUPPORTED_SOURCE_DOCUMENT_EXTENSIONS = (
|
||||
".rst",
|
||||
".md",
|
||||
".pdf",
|
||||
".txt",
|
||||
".docx",
|
||||
".csv",
|
||||
".epub",
|
||||
".html",
|
||||
".mdx",
|
||||
".json",
|
||||
".xlsx",
|
||||
".pptx",
|
||||
# Read by the HTML parsers on every engine.
|
||||
".xhtml",
|
||||
# Read by the anydoc engine (legacy/macro Office, OpenDocument, RTF).
|
||||
# Parseable regardless of DOC_PARSER_ENGINE: anydoc is a core dependency,
|
||||
# and both parser maps route these suffixes to it.
|
||||
*ANYDOC_GAINED_SUFFIXES,
|
||||
)
|
||||
|
||||
SUPPORTED_SOURCE_IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg")
|
||||
|
||||
SUPPORTED_SOURCE_EXTENSIONS = (
|
||||
*SUPPORTED_SOURCE_DOCUMENT_EXTENSIONS,
|
||||
*SUPPORTED_SOURCE_IMAGE_EXTENSIONS,
|
||||
*SUPPORTED_AUDIO_EXTENSIONS,
|
||||
)
|
||||
|
||||
# Suffixes the attachment path has a dedicated parser for — exactly the keys
|
||||
# of ``get_default_file_extractor()``. Kept as a literal (importing ``bulk``
|
||||
# here would drag docling into the API process), with
|
||||
# ``tests/parser/file/test_constants.py`` asserting the two agree.
|
||||
#
|
||||
# This is *not* the whole attachment allow-list: a suffix with no parser is
|
||||
# read by ``SimpleDirectoryReader``'s plain-text fallthrough, which is right
|
||||
# for a .py or a .log and catastrophic for a video — the reader "extracts"
|
||||
# megabytes of binary garbage, truncates it, and stores it with
|
||||
# ``extraction.status == "ok"``. So unparsed suffixes are admitted on
|
||||
# content instead (``upload_limits.enforce_parseable_attachment``): text
|
||||
# passes, binary is refused. Zip is deliberately absent — source ingestion
|
||||
# extracts archives, the attachment path does not, and a zip fails the
|
||||
# content check like any other binary.
|
||||
#
|
||||
# Mirrored in ``frontend/src/constants/fileUpload.ts``; update both together.
|
||||
ATTACHMENT_PARSER_EXTENSIONS = frozenset(
|
||||
{
|
||||
*SUPPORTED_SOURCE_EXTENSIONS,
|
||||
".xhtml",
|
||||
".adoc",
|
||||
".asciidoc",
|
||||
".tiff",
|
||||
".tif",
|
||||
".bmp",
|
||||
".webp",
|
||||
".vtt",
|
||||
".xml",
|
||||
}
|
||||
# .txt has no parser of its own — it *is* the plain-text fallthrough. It
|
||||
# must be sniffed like any other unparsed suffix, or renaming a video to
|
||||
# notes.txt walks straight back into the bug this gate exists for.
|
||||
- {".txt"}
|
||||
)
|
||||
|
||||
|
||||
def attachment_extension(filename: str | None) -> str:
|
||||
"""Return the lower-cased extension of ``filename`` including the dot, or ``""``.
|
||||
|
||||
Args:
|
||||
filename: A bare filename or path; ``None`` and empty strings yield ``""``.
|
||||
|
||||
Returns:
|
||||
The last suffix in lower case (``".pdf"``), or ``""`` when there is none.
|
||||
"""
|
||||
if not filename:
|
||||
return ""
|
||||
return os.path.splitext(os.path.basename(str(filename)))[1].lower()
|
||||
|
||||
|
||||
def has_attachment_parser(filename: str | None) -> bool:
|
||||
"""Return whether an attachment's suffix has a dedicated parser.
|
||||
|
||||
A False result does not mean the file is refused: it means nothing but the
|
||||
plain-text fallthrough will read it, so it has to earn its place on
|
||||
content. The decision is by extension only, never by the mime type a
|
||||
browser reports (mobile pickers ignore ``accept`` and lie).
|
||||
|
||||
Args:
|
||||
filename: The upload's filename.
|
||||
|
||||
Returns:
|
||||
True when the suffix is in ``ATTACHMENT_PARSER_EXTENSIONS``.
|
||||
"""
|
||||
return attachment_extension(filename) in ATTACHMENT_PARSER_EXTENSIONS
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,70 @@
|
||||
"""Docs parser.
|
||||
|
||||
Contains parsers for docx, pdf files.
|
||||
|
||||
"""
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
from docsgpt.core.settings import settings
|
||||
import requests
|
||||
|
||||
class PDFParser(BaseParser):
|
||||
"""PDF parser."""
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> str:
|
||||
"""Parse file."""
|
||||
if settings.PARSE_PDF_AS_IMAGE:
|
||||
doc2md_service = "https://llm.arc53.com/doc2md"
|
||||
# alternatively you can use local vision capable LLM
|
||||
with open(file, "rb") as file_loaded:
|
||||
files = {'file': file_loaded}
|
||||
response = requests.post(doc2md_service, files=files, timeout=100)
|
||||
data = response.json()["markdown"]
|
||||
return data
|
||||
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
except ImportError:
|
||||
raise ValueError("pypdf is required to read PDF files.")
|
||||
text_list = []
|
||||
with open(file, "rb") as fp:
|
||||
# Create a PDF object
|
||||
pdf = PdfReader(fp)
|
||||
|
||||
# Get the number of pages in the PDF document
|
||||
num_pages = len(pdf.pages)
|
||||
|
||||
# Iterate over every page
|
||||
for page_index in range(num_pages):
|
||||
# Extract the text from the page
|
||||
page = pdf.pages[page_index]
|
||||
page_text = page.extract_text()
|
||||
text_list.append(page_text)
|
||||
text = "\n".join(text_list)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
class DocxParser(BaseParser):
|
||||
"""Docx parser."""
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> str:
|
||||
"""Parse file."""
|
||||
try:
|
||||
import docx2txt
|
||||
except ImportError:
|
||||
raise ValueError("docx2txt is required to read Microsoft Word files.")
|
||||
|
||||
text = docx2txt.process(file)
|
||||
|
||||
return text
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Epub parser.
|
||||
|
||||
Contains parsers for epub files.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
|
||||
|
||||
class EpubParser(BaseParser):
|
||||
"""Epub Parser."""
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> str:
|
||||
"""Parse file."""
|
||||
try:
|
||||
from fast_ebook import epub
|
||||
except ImportError:
|
||||
raise ValueError("`fast-ebook` is required to read Epub files.")
|
||||
|
||||
book = epub.read_epub(file)
|
||||
text = book.to_markdown()
|
||||
return text
|
||||
@@ -0,0 +1,238 @@
|
||||
"""HTML parsers.
|
||||
|
||||
``HTMLParser`` returns the visible text of a page (the ``fast`` engine and
|
||||
the legacy no-docling map rely on that plain-text contract).
|
||||
``HTMLMarkdownParser`` returns Markdown instead and is what the anydoc
|
||||
engine maps ``.html``/``.xhtml`` to: anydoc has no HTML support, and
|
||||
markdownify — already a dependency of the web crawler — keeps links, GFM
|
||||
tables, headings and byte-exact code blocks at a few ms per page.
|
||||
"""
|
||||
import logging
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
from docsgpt.utils import truncate_to_line_boundary
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The crawler's conventions (``crawler_markdown.py`` / ``read_webpage.py``),
|
||||
# so file and web ingestion produce the same Markdown shape.
|
||||
MARKDOWNIFY_OPTIONS = {"heading_style": "ATX", "newline_style": "BACKSLASH"}
|
||||
|
||||
# Elements whose text is never document content. ``title`` is reported via
|
||||
# ``get_file_metadata`` instead of leaking in as a stray first line.
|
||||
_DROP_TAGS = ("title", "script", "style", "noscript", "template")
|
||||
# Attributes whose ``data:`` URIs would otherwise land in the Markdown as
|
||||
# link/image targets: one inline image is megabytes of base64 "text" for the
|
||||
# chunker and embedder (a 5 MB data-URI ``<img>`` measured 7 MB of output).
|
||||
_URI_ATTRIBUTES = ("src", "href", "srcset", "poster", "data")
|
||||
|
||||
|
||||
def html_to_markdown(html: Union[str, bytes]) -> str:
|
||||
"""Convert an HTML/XHTML document to Markdown.
|
||||
|
||||
Args:
|
||||
html: The markup. Pass the raw bytes when you have them: BeautifulSoup
|
||||
then decodes by the document's own BOM / ``<meta charset>``
|
||||
instead of the process locale, so a windows-1250 page keeps its
|
||||
diacritics.
|
||||
|
||||
Returns:
|
||||
Markdown with runs of blank lines collapsed to one.
|
||||
"""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
return soup_to_markdown(BeautifulSoup(html, "html.parser"))
|
||||
|
||||
|
||||
def soup_to_markdown(soup) -> str:
|
||||
"""Convert a parsed ``BeautifulSoup`` tree to Markdown.
|
||||
|
||||
Destructive: the non-content elements in ``_DROP_TAGS`` (including
|
||||
``<title>``) are removed from ``soup``; read the title first.
|
||||
|
||||
Args:
|
||||
soup: The parsed document.
|
||||
|
||||
Returns:
|
||||
Markdown with runs of blank lines collapsed to one.
|
||||
"""
|
||||
from bs4 import CData, Declaration, ProcessingInstruction
|
||||
from markdownify import MarkdownConverter
|
||||
|
||||
for tag in soup.find_all(_DROP_TAGS):
|
||||
tag.decompose()
|
||||
# ``<?xml ...?>`` (every XHTML file), CDATA and stray declarations are
|
||||
# not text; markdownify would emit them as the document's first line.
|
||||
for node in soup.find_all(string=lambda s: isinstance(s, (ProcessingInstruction, CData, Declaration))):
|
||||
node.extract()
|
||||
for tag in soup.find_all(True):
|
||||
for attribute in _URI_ATTRIBUTES:
|
||||
value = tag.get(attribute)
|
||||
if isinstance(value, list):
|
||||
value = " ".join(value)
|
||||
if isinstance(value, str) and "data:" in value.lower():
|
||||
del tag[attribute]
|
||||
markdown = MarkdownConverter(**MARKDOWNIFY_OPTIONS).convert_soup(soup)
|
||||
return re.sub(r"\n{3,}", "\n\n", markdown).strip()
|
||||
|
||||
|
||||
def _soup_title(soup) -> Optional[str]:
|
||||
"""The document's ``<title>`` text, or None."""
|
||||
if soup.title and soup.title.string:
|
||||
return str(soup.title.string)
|
||||
return None
|
||||
|
||||
|
||||
def _trim_torn_utf8_tail(data: bytes) -> bytes:
|
||||
"""Drop a trailing partial multi-byte UTF-8 sequence left by a byte cut.
|
||||
|
||||
A cut that lands inside a multi-byte character matters more than one
|
||||
lost glyph: BeautifulSoup decodes strictly and, when the declared UTF-8
|
||||
fails on the torn tail, retries as windows-1252 — and if every other
|
||||
non-ASCII byte happens to be cp1252-valid (German umlauts are), the
|
||||
*whole* document comes back as mojibake. Only a UTF-8 tail is
|
||||
recognised; a single-byte page loses at most the final byte.
|
||||
|
||||
Args:
|
||||
data: A head window whose last bytes may be a torn sequence.
|
||||
|
||||
Returns:
|
||||
``data`` minus any trailing incomplete UTF-8 sequence.
|
||||
"""
|
||||
end = len(data)
|
||||
trailing = 0
|
||||
while end > 0 and trailing < 3 and 0x80 <= data[end - 1] < 0xC0:
|
||||
end -= 1
|
||||
trailing += 1
|
||||
if end == 0:
|
||||
return data
|
||||
lead = data[end - 1]
|
||||
if lead < 0xC0:
|
||||
return data
|
||||
need = 2 if lead < 0xE0 else 3 if lead < 0xF0 else 4
|
||||
if trailing + 1 < need:
|
||||
return data[: end - 1]
|
||||
return data
|
||||
|
||||
|
||||
def read_markup_head(file: Path, max_bytes: int) -> bytes:
|
||||
"""Read a markup file, head-truncated to ``max_bytes`` on a line boundary.
|
||||
|
||||
HTML has no lightweight fallback parser to bound it by other means (the
|
||||
way oversized CSV/XLSX go to the plain tabular parsers), and the soup +
|
||||
markdownify tree costs ~50x the input — 30 MB measured at 1.6 GB RSS,
|
||||
against a 100 MB upload cap. Parsing only the head keeps one upload from
|
||||
taking the ingest worker down; the cut lands on a line boundary when one
|
||||
is reasonably close, and the lenient HTML parser copes with the torn tail.
|
||||
|
||||
Args:
|
||||
file: Path to the markup file.
|
||||
max_bytes: Size gate; ``<= 0`` reads the whole file.
|
||||
|
||||
Returns:
|
||||
The file's bytes, or its first ``max_bytes`` when it is larger.
|
||||
"""
|
||||
with open(file, "rb") as fh:
|
||||
if max_bytes <= 0:
|
||||
return fh.read()
|
||||
head = fh.read(max_bytes + 1)
|
||||
if len(head) <= max_bytes:
|
||||
return head
|
||||
logger.warning(
|
||||
f"Markup {Path(file).name} exceeds MARKUP_MAX_BYTES ({max_bytes}); "
|
||||
f"parsing the first {max_bytes} bytes to bound memory"
|
||||
)
|
||||
head = head[:max_bytes]
|
||||
if head[:2] in (b"\xff\xfe", b"\xfe\xff"):
|
||||
# UTF-16 (BOM-declared): a byte-level line cut lands between the two
|
||||
# bytes of a code unit and the strict decode then fails outright, so
|
||||
# keep an even byte count and let the lenient parser take the torn
|
||||
# tail. (Only 2-byte units matter here; a torn surrogate pair is one
|
||||
# lost character.)
|
||||
return head[: len(head) - (len(head) % 2)]
|
||||
return _trim_torn_utf8_tail(truncate_to_line_boundary(head))
|
||||
|
||||
|
||||
class HTMLParser(BaseParser):
|
||||
"""HTML parser."""
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]:
|
||||
"""Extract the visible text of an HTML file.
|
||||
|
||||
Returns the text as a string, matching the other file parsers — the
|
||||
title is recovered separately by ``get_file_metadata``.
|
||||
"""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
with open(file, "r", errors=errors) as f:
|
||||
soup = BeautifulSoup(f, "html.parser")
|
||||
return soup.get_text("\n")
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Return the document title, when the markup carries one."""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
try:
|
||||
with open(file, "r", errors="ignore") as f:
|
||||
soup = BeautifulSoup(f, "html.parser")
|
||||
except OSError:
|
||||
return {}
|
||||
if soup.title and soup.title.string:
|
||||
return {"title": str(soup.title.string)}
|
||||
return {}
|
||||
|
||||
|
||||
class HTMLMarkdownParser(HTMLParser):
|
||||
"""HTML/XHTML to Markdown, for chunking and retrieval.
|
||||
|
||||
Differs from ``HTMLParser`` in three ways: the body comes back as
|
||||
Markdown; the file is read as bytes so the document's own charset
|
||||
declaration wins over the process locale; and input is head-truncated at
|
||||
``MARKUP_MAX_BYTES`` (the docling HTML path has the same gate under
|
||||
``DOCLING_MARKUP_MAX_BYTES``).
|
||||
"""
|
||||
|
||||
def __init__(self, parser_config: Optional[Dict] = None) -> None:
|
||||
super().__init__(parser_config)
|
||||
# (path, title) of the most recent ``parse_file``, so the metadata
|
||||
# call that follows it does not build the soup a second time.
|
||||
self._last_title: Optional[Tuple[Path, Optional[str]]] = None
|
||||
|
||||
@staticmethod
|
||||
def _read(file: Path) -> bytes:
|
||||
return read_markup_head(Path(file), int(settings.MARKUP_MAX_BYTES))
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]:
|
||||
"""Convert an HTML file to Markdown.
|
||||
|
||||
``errors`` is accepted for the ``BaseParser`` contract but unused:
|
||||
decoding is done by BeautifulSoup from the raw bytes.
|
||||
"""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
_ = errors
|
||||
soup = BeautifulSoup(self._read(file), "html.parser")
|
||||
self._last_title = (Path(file), _soup_title(soup))
|
||||
return soup_to_markdown(soup)
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Return the document title, decoded the same way as the body."""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
last = self._last_title
|
||||
if last is not None and last[0] == Path(file):
|
||||
title = last[1]
|
||||
else:
|
||||
try:
|
||||
title = _soup_title(BeautifulSoup(self._read(file), "html.parser"))
|
||||
except OSError:
|
||||
return {}
|
||||
return {"title": title} if title else {}
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Image parser.
|
||||
|
||||
Contains parser for .png, .jpg, .jpeg files.
|
||||
|
||||
"""
|
||||
from pathlib import Path
|
||||
import requests
|
||||
from typing import Dict, Union
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
|
||||
class ImageParser(BaseParser):
|
||||
"""Image parser."""
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]:
|
||||
if settings.PARSE_IMAGE_REMOTE:
|
||||
doc2md_service = "https://llm.arc53.com/doc2md"
|
||||
# alternatively you can use local vision capable LLM
|
||||
with open(file, "rb") as file_loaded:
|
||||
files = {'file': file_loaded}
|
||||
response = requests.post(doc2md_service, files=files, timeout=100)
|
||||
data = response.json()["markdown"]
|
||||
else:
|
||||
data = ""
|
||||
return data
|
||||
@@ -0,0 +1,57 @@
|
||||
import json
|
||||
from typing import Any, Dict, List, Union
|
||||
from pathlib import Path
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
|
||||
class JSONParser(BaseParser):
|
||||
r"""JSON (.json) parser.
|
||||
|
||||
Parses JSON files into a list of strings or a concatenated document.
|
||||
It handles both JSON objects (dictionaries) and arrays (lists).
|
||||
|
||||
Args:
|
||||
concat_rows (bool): Whether to concatenate all rows into one document.
|
||||
If set to False, a Document will be created for each item in the JSON.
|
||||
True by default.
|
||||
|
||||
row_joiner (str): Separator to use for joining each row.
|
||||
Only used when `concat_rows=True`.
|
||||
Set to "\n" by default.
|
||||
|
||||
json_config (dict): Options for parsing JSON. Can be used to specify options like
|
||||
custom decoding or formatting. Set to empty dict by default.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
concat_rows: bool = True,
|
||||
row_joiner: str = "\n",
|
||||
json_config: dict = {},
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._concat_rows = concat_rows
|
||||
self._row_joiner = row_joiner
|
||||
self._json_config = json_config
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse JSON file."""
|
||||
|
||||
with open(file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f, **self._json_config)
|
||||
|
||||
if isinstance(data, dict):
|
||||
data = [data]
|
||||
|
||||
if self._concat_rows:
|
||||
return self._row_joiner.join([str(item) for item in data])
|
||||
else:
|
||||
return data
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Markdown parser.
|
||||
|
||||
Contains parser for md files.
|
||||
|
||||
"""
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
from docsgpt.utils import num_tokens_from_string
|
||||
|
||||
|
||||
class MarkdownParser(BaseParser):
|
||||
"""Markdown parser.
|
||||
|
||||
Extract text from markdown files.
|
||||
Returns dictionary with keys as headers and values as the text between headers.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
remove_hyperlinks: bool = True,
|
||||
remove_images: bool = True,
|
||||
max_tokens: int = 2048,
|
||||
# remove_tables: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._remove_hyperlinks = remove_hyperlinks
|
||||
self._remove_images = remove_images
|
||||
self._max_tokens = max_tokens
|
||||
# self._remove_tables = remove_tables
|
||||
|
||||
def tups_chunk_append(self, tups: List[Tuple[Optional[str], str]], current_header: Optional[str],
|
||||
current_text: str):
|
||||
"""Append to tups chunk."""
|
||||
num_tokens = num_tokens_from_string(current_text)
|
||||
if num_tokens > self._max_tokens:
|
||||
chunks = [current_text[i:i + self._max_tokens] for i in range(0, len(current_text), self._max_tokens)]
|
||||
for chunk in chunks:
|
||||
tups.append((current_header, chunk))
|
||||
else:
|
||||
tups.append((current_header, current_text))
|
||||
return tups
|
||||
|
||||
def markdown_to_tups(self, markdown_text: str) -> List[Tuple[Optional[str], str]]:
|
||||
"""Convert a markdown file to a dictionary.
|
||||
|
||||
The keys are the headers and the values are the text under each header.
|
||||
|
||||
"""
|
||||
markdown_tups: List[Tuple[Optional[str], str]] = []
|
||||
lines = markdown_text.split("\n")
|
||||
|
||||
current_header = None
|
||||
current_text = ""
|
||||
|
||||
for line in lines:
|
||||
header_match = re.match(r"^#+\s", line)
|
||||
if header_match:
|
||||
if current_header is not None:
|
||||
if current_text == "" or None:
|
||||
continue
|
||||
markdown_tups = self.tups_chunk_append(markdown_tups, current_header, current_text)
|
||||
|
||||
current_header = line
|
||||
current_text = ""
|
||||
else:
|
||||
current_text += line + "\n"
|
||||
markdown_tups = self.tups_chunk_append(markdown_tups, current_header, current_text)
|
||||
|
||||
if current_header is not None:
|
||||
# pass linting, assert keys are defined
|
||||
markdown_tups = [
|
||||
(re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
|
||||
for key, value in markdown_tups
|
||||
]
|
||||
else:
|
||||
markdown_tups = [
|
||||
(key, re.sub("\n", "", value)) for key, value in markdown_tups
|
||||
]
|
||||
|
||||
return markdown_tups
|
||||
|
||||
def remove_images(self, content: str) -> str:
|
||||
"""Get a dictionary of a markdown file from its path."""
|
||||
pattern = r"!{1}\[\[(.*)\]\]"
|
||||
content = re.sub(pattern, "", content)
|
||||
return content
|
||||
|
||||
# def remove_tables(self, content: str) -> List[List[str]]:
|
||||
# """Convert markdown tables to nested lists."""
|
||||
# table_rows_pattern = r"((\r?\n){2}|^)([^\r\n]*\|[^\r\n]*(\r?\n)?)+(?=(\r?\n){2}|$)"
|
||||
# table_cells_pattern = r"([^\|\r\n]*)\|"
|
||||
#
|
||||
# table_rows = re.findall(table_rows_pattern, content, re.MULTILINE)
|
||||
# table_lists = []
|
||||
# for row in table_rows:
|
||||
# cells = re.findall(table_cells_pattern, row[2])
|
||||
# cells = [cell.strip() for cell in cells if cell.strip()]
|
||||
# table_lists.append(cells)
|
||||
# return str(table_lists)
|
||||
|
||||
def remove_hyperlinks(self, content: str) -> str:
|
||||
"""Get a dictionary of a markdown file from its path."""
|
||||
pattern = r"\[(.*?)\]\((.*?)\)"
|
||||
content = re.sub(pattern, r"\1", content)
|
||||
return content
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Initialize the parser with the config."""
|
||||
return {}
|
||||
|
||||
def parse_tups(
|
||||
self, filepath: Path, errors: str = "ignore"
|
||||
) -> List[Tuple[Optional[str], str]]:
|
||||
"""Parse file into tuples."""
|
||||
with open(filepath, "r") as f:
|
||||
content = f.read()
|
||||
if self._remove_hyperlinks:
|
||||
content = self.remove_hyperlinks(content)
|
||||
if self._remove_images:
|
||||
content = self.remove_images(content)
|
||||
# if self._remove_tables:
|
||||
# content = self.remove_tables(content)
|
||||
markdown_tups = self.markdown_to_tups(content)
|
||||
return markdown_tups
|
||||
|
||||
def parse_file(
|
||||
self, filepath: Path, errors: str = "ignore"
|
||||
) -> Union[str, List[str]]:
|
||||
"""Parse file into string."""
|
||||
tups = self.parse_tups(filepath, errors=errors)
|
||||
results = []
|
||||
# TODO: don't include headers right now
|
||||
for header, value in tups:
|
||||
if header is None:
|
||||
results.append(value)
|
||||
else:
|
||||
results.append(f"\n\n{header}\n{value}")
|
||||
return results
|
||||
@@ -0,0 +1,874 @@
|
||||
"""Native OCR parsers: scanned PDFs and images without docling.
|
||||
|
||||
OCR in DocsGPT has two backends, selected by ``OCR_BACKEND``:
|
||||
|
||||
* ``docling`` — the layout-model pipeline in ``docling_parser.py``. Hybrid
|
||||
OCR (only the bitmap regions of a page), reading-order and table
|
||||
structure recovery, and five engines (tesseract / auto / ocrmac /
|
||||
rapidocr / deepseek). Costs the optional docling install (torch, ONNX
|
||||
models, gigabytes of image) and seconds to minutes per file.
|
||||
* ``native`` (this module) — page rendering with pypdfium2 and Pillow, both
|
||||
core dependencies, feeding one of two engines directly: the system
|
||||
``tesseract`` binary, or a DeepSeek-OCR model behind an OpenAI-compatible
|
||||
endpoint (Ollama / vLLM). No ML models load in the worker. Pages that
|
||||
carry a text layer are read through pypdfium2 and never OCR'd; pages
|
||||
without one are rendered and OCR'd. What it lacks against docling is the
|
||||
layout model: multi-column scans rely on tesseract's own page
|
||||
segmentation, and tesseract yields tables as plain lines (DeepSeek-OCR
|
||||
emits Markdown tables itself).
|
||||
|
||||
``auto`` picks docling when it is installed and native otherwise, so a
|
||||
deployment that never installs the docling extra gets working OCR from the
|
||||
tesseract binary alone, and one that does install it keeps today's
|
||||
behaviour unchanged.
|
||||
"""
|
||||
import base64
|
||||
import functools
|
||||
import io
|
||||
import logging
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Protocol, Tuple, Union
|
||||
|
||||
from docsgpt.parser.file.base_parser import (
|
||||
BaseParser,
|
||||
DocumentParseError,
|
||||
delegate_parse,
|
||||
module_available,
|
||||
)
|
||||
|
||||
from docsgpt.core.optional_deps import install_hint
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Every engine ``OCR_ENGINE`` accepts. ``auto``, ``ocrmac`` and ``rapidocr``
|
||||
# exist only inside docling; the native backend maps them to tesseract.
|
||||
VALID_OCR_ENGINES: Tuple[str, ...] = ("tesseract", "auto", "ocrmac", "rapidocr", "deepseek")
|
||||
NATIVE_OCR_ENGINES: Tuple[str, ...] = ("tesseract", "deepseek")
|
||||
VALID_OCR_BACKENDS: Tuple[str, ...] = ("auto", "docling", "native")
|
||||
|
||||
IMAGE_SUFFIXES = frozenset({".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".webp"})
|
||||
|
||||
# Characters in a page's text layer at or above which the page is read from
|
||||
# that layer instead of being OCR'd. Same separation ``PdfiumTextParser``
|
||||
# relies on: scanned pages sit at 0-17 chars, text pages in the hundreds+.
|
||||
DEFAULT_TEXT_LAYER_MIN_CHARS = 32
|
||||
# Chars-per-page floor for the near-empty guard when the setting is unusable.
|
||||
_DEFAULT_MIN_CHARS_PER_PAGE = 20
|
||||
_DEFAULT_RENDER_DPI = 200
|
||||
_MIN_RENDER_DPI, _MAX_RENDER_DPI = 72, 600
|
||||
# Pixel budget for one rendered page. pypdfium2 allocates width*height*4
|
||||
# bytes up front with no cap of its own, so a 14400x14400 pt page (the PDF
|
||||
# maximum; a 191-byte file) at 200 dpi would ask for 40000x40000 px = 6 GiB.
|
||||
# 40 MP is ~10x a Letter page at 200 dpi and still OCRs an A0 poster at a
|
||||
# usable ~125 dpi; anything larger renders at the scale that fits.
|
||||
_MAX_RENDER_PIXELS = 40_000_000
|
||||
_TESSERACT_TIMEOUT_SECONDS = 300
|
||||
# tesseract reports a missing language pack on stderr and, when at least one
|
||||
# other requested pack loads, exits 0 and silently OCRs with what it has —
|
||||
# so the exit code alone cannot catch OCR_LANGS=eng+chi_sim without chi_sim.
|
||||
_TESSERACT_LANG_ERROR_RE = re.compile(r"Error opening data file|Failed loading language")
|
||||
_TESSERACT_LANG_RE = re.compile(r"^[A-Za-z0-9_/\-]+$")
|
||||
# docling's DeepSeek-OCR prompt minus its ``<|grounding|>`` prefix: grounding
|
||||
# makes the model wrap every element in ref/det tags with bounding boxes,
|
||||
# which docling parses back into a layout tree. Plain Markdown is what the
|
||||
# ingestion pipeline stores, so ask for that directly.
|
||||
DEEPSEEK_PROMPT = "Convert the document to markdown."
|
||||
_DEEPSEEK_MAX_TOKENS = 4096
|
||||
# Defensive cleanup should a served model still emit grounding markup.
|
||||
_DEEPSEEK_DET_BLOCK_RE = re.compile(r"<\|det\|>.*?<\|/det\|>", re.DOTALL)
|
||||
_DEEPSEEK_TAG_RE = re.compile(r"<\|/?ref\|>|<\|grounding\|>|<\|end▁of▁sentence\|>")
|
||||
# tesseract's CJK models emit one space between every character ("互相 保密 协议")
|
||||
# because they segment by glyph. Chinese and Japanese are written without
|
||||
# word spaces, so collapse whitespace between two CJK characters (or CJK
|
||||
# punctuation); spaces next to Latin text, digits and line breaks stay.
|
||||
_CJK_CHAR = r"[\u3000-\u303f\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff\uff00-\uffef]"
|
||||
_CJK_SPACE_RE = re.compile(rf"(?<={_CJK_CHAR})[ \t]+(?={_CJK_CHAR})")
|
||||
|
||||
|
||||
class OcrUnavailableError(DocumentParseError):
|
||||
"""The selected OCR engine cannot run on this host (missing binary or endpoint)."""
|
||||
|
||||
|
||||
class OcrEngine(Protocol):
|
||||
"""Anything that turns one page image into text."""
|
||||
|
||||
name: str
|
||||
|
||||
def ocr_image(self, image) -> str: # pragma: no cover - protocol
|
||||
"""Return the text (or Markdown) recognised in ``image`` (a PIL image)."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def native_ocr_available() -> bool:
|
||||
"""Whether the native backend's rendering stack (pypdfium2 + Pillow) is importable."""
|
||||
return module_available("pypdfium2") and module_available("PIL")
|
||||
|
||||
|
||||
def resolve_ocr_backend(requested: Optional[str] = None) -> str:
|
||||
"""Resolve ``OCR_BACKEND`` to the backend that will actually run.
|
||||
|
||||
Args:
|
||||
requested: ``auto`` | ``docling`` | ``native``, or None to read the setting.
|
||||
|
||||
Returns:
|
||||
``"docling"`` or ``"native"``. ``auto`` prefers docling when it is
|
||||
installed; ``docling`` without the install degrades to native with a
|
||||
warning rather than leaving OCR off.
|
||||
"""
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
backend = str(requested or getattr(settings, "OCR_BACKEND", None) or "auto").strip().lower()
|
||||
if backend not in VALID_OCR_BACKENDS:
|
||||
logger.warning(f"Unknown OCR_BACKEND {backend!r}; using auto")
|
||||
backend = "auto"
|
||||
docling_installed = module_available("docling")
|
||||
if backend == "docling" and not docling_installed:
|
||||
logger.warning(
|
||||
"OCR_BACKEND=docling but docling is not installed (%s); using the native OCR backend",
|
||||
install_hint("docling"),
|
||||
)
|
||||
return "native"
|
||||
if backend == "auto":
|
||||
return "docling" if docling_installed else "native"
|
||||
return backend
|
||||
|
||||
|
||||
def resolve_native_ocr_engine(requested: Optional[str] = None) -> str:
|
||||
"""Map ``OCR_ENGINE`` onto an engine the native backend implements.
|
||||
|
||||
Args:
|
||||
requested: Engine name, or None to read ``settings.OCR_ENGINE``.
|
||||
|
||||
Returns:
|
||||
``"tesseract"`` or ``"deepseek"``. docling-only engines (``auto``,
|
||||
``ocrmac``, ``rapidocr``) and unknown names become tesseract with a
|
||||
warning, so switching a deployment off docling never silently
|
||||
disables OCR.
|
||||
"""
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
engine = str(requested or getattr(settings, "OCR_ENGINE", None) or "tesseract").strip().lower()
|
||||
if engine in NATIVE_OCR_ENGINES:
|
||||
return engine
|
||||
if engine in VALID_OCR_ENGINES:
|
||||
logger.warning(
|
||||
f"OCR_ENGINE={engine!r} is only available through the docling backend; "
|
||||
"the native backend is using tesseract"
|
||||
)
|
||||
else:
|
||||
logger.warning(f"Unknown OCR_ENGINE {engine!r}; the native backend is using tesseract")
|
||||
return "tesseract"
|
||||
|
||||
|
||||
def ocr_min_chars_per_page() -> int:
|
||||
"""Chars-per-page floor for the near-empty OCR guard (``OCR_MIN_CHARS_PER_PAGE``); 0 disables it."""
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
try:
|
||||
return int(getattr(settings, "OCR_MIN_CHARS_PER_PAGE", _DEFAULT_MIN_CHARS_PER_PAGE))
|
||||
except (TypeError, ValueError):
|
||||
return _DEFAULT_MIN_CHARS_PER_PAGE
|
||||
|
||||
|
||||
def render_dpi() -> int:
|
||||
"""Rendering resolution for pages that need OCR (``OCR_RENDER_DPI``), clamped to a sane range."""
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
try:
|
||||
dpi = int(getattr(settings, "OCR_RENDER_DPI", _DEFAULT_RENDER_DPI))
|
||||
except (TypeError, ValueError):
|
||||
dpi = _DEFAULT_RENDER_DPI
|
||||
return max(_MIN_RENDER_DPI, min(_MAX_RENDER_DPI, dpi))
|
||||
|
||||
|
||||
def _has_alpha(image) -> bool:
|
||||
return "A" in image.getbands() or (image.mode == "P" and "transparency" in image.info)
|
||||
|
||||
|
||||
def _png_bytes(image) -> bytes:
|
||||
"""Encode a PIL image as PNG, flattening modes the engines cannot take.
|
||||
|
||||
Transparency is composited onto white: a plain ``convert("RGB")`` drops
|
||||
the alpha channel and leaves transparent pixels black, which turns dark
|
||||
text on a transparent background into an all-black page that OCRs to
|
||||
nothing.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
if _has_alpha(image):
|
||||
rgba = image.convert("RGBA")
|
||||
canvas = Image.new("RGB", rgba.size, "white")
|
||||
canvas.paste(rgba, mask=rgba.getchannel("A"))
|
||||
image = canvas
|
||||
elif image.mode not in ("RGB", "L"):
|
||||
image = image.convert("RGB")
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def fit_to_pixel_budget(image):
|
||||
"""Downscale a decoded image to ``_MAX_RENDER_PIXELS`` before OCR.
|
||||
|
||||
The same budget the PDF renderer applies: a 28 KB PNG can decode to
|
||||
81 MP (Pillow's bomb guard only trips at ~178 MP), and every extra
|
||||
pixel is paid again in the PNG re-encode and inside tesseract.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
width, height = image.size
|
||||
scale = render_scale(width, height, 72)
|
||||
if scale >= 1.0:
|
||||
return image
|
||||
if image.mode not in ("RGB", "L", "RGBA", "LA"):
|
||||
image = image.convert("RGBA" if _has_alpha(image) else "RGB")
|
||||
logger.warning(
|
||||
"Image is %dx%d px; downscaling to %.0f%% to stay within %d MP before OCR",
|
||||
width,
|
||||
height,
|
||||
scale * 100,
|
||||
_MAX_RENDER_PIXELS // 1_000_000,
|
||||
)
|
||||
return image.resize(
|
||||
(max(1, int(width * scale)), max(1, int(height * scale))), Image.Resampling.BILINEAR
|
||||
)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def tesseract_languages() -> Optional[frozenset]:
|
||||
"""Language packs the ``tesseract`` binary reports (``--list-langs``), or None when unknown.
|
||||
|
||||
Cached for the process: packs are installed with the image, not at runtime.
|
||||
"""
|
||||
if shutil.which("tesseract") is None:
|
||||
return None
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
["tesseract", "--list-langs"], capture_output=True, timeout=30, check=False
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return None
|
||||
output = completed.stdout.decode("utf-8", "replace") + "\n" + completed.stderr.decode("utf-8", "replace")
|
||||
langs = frozenset(
|
||||
line.strip() for line in output.splitlines() if _TESSERACT_LANG_RE.match(line.strip())
|
||||
)
|
||||
return langs or None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Engines
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TesseractEngine:
|
||||
"""OCR through the system ``tesseract`` binary (stdin -> stdout, no Python wrapper).
|
||||
|
||||
Attributes:
|
||||
languages: Tesseract language codes; None reads ``OCR_LANGS`` per call.
|
||||
psm: Page segmentation mode (3 = fully automatic, tesseract's default).
|
||||
timeout: Seconds allowed per page.
|
||||
"""
|
||||
|
||||
name = "tesseract"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
languages: Optional[List[str]] = None,
|
||||
psm: int = 3,
|
||||
timeout: float = _TESSERACT_TIMEOUT_SECONDS,
|
||||
) -> None:
|
||||
self.languages = languages
|
||||
self.psm = psm
|
||||
self.timeout = timeout
|
||||
|
||||
@staticmethod
|
||||
def available() -> bool:
|
||||
"""Whether a ``tesseract`` binary is on PATH."""
|
||||
return shutil.which("tesseract") is not None
|
||||
|
||||
def _language_arg(self) -> str:
|
||||
if self.languages:
|
||||
langs = [lang.strip() for lang in self.languages if lang and lang.strip()]
|
||||
else:
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
configured = str(getattr(settings, "OCR_LANGS", "") or "eng")
|
||||
langs = [lang.strip() for lang in configured.split("+") if lang.strip()]
|
||||
return "+".join(langs) or "eng"
|
||||
|
||||
def command(self) -> List[str]:
|
||||
"""The tesseract command line, for logging and tests."""
|
||||
return ["tesseract", "stdin", "stdout", "-l", self._language_arg(), "--psm", str(self.psm)]
|
||||
|
||||
def ocr_image(self, image) -> str:
|
||||
"""OCR one PIL image.
|
||||
|
||||
Raises:
|
||||
OcrUnavailableError: No tesseract binary on PATH.
|
||||
DocumentParseError: tesseract failed or timed out.
|
||||
"""
|
||||
if not self.available():
|
||||
raise OcrUnavailableError(
|
||||
"OCR_ENGINE=tesseract but no tesseract binary is on PATH. Install "
|
||||
"tesseract-ocr plus language packs (Docker: build with "
|
||||
"INSTALL_TESSERACT=true; local: apt/brew install tesseract-ocr), "
|
||||
"or set OCR_ENGINE=deepseek / install the docling extra."
|
||||
)
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
self.command(),
|
||||
input=_png_bytes(image),
|
||||
capture_output=True,
|
||||
timeout=self.timeout,
|
||||
check=False,
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
raise DocumentParseError(f"tesseract timed out after {self.timeout:.0f}s on a page") from exc
|
||||
except OSError as exc:
|
||||
raise DocumentParseError(f"tesseract could not be started: {exc}") from exc
|
||||
stderr = completed.stderr.decode("utf-8", "replace").strip()
|
||||
if _TESSERACT_LANG_ERROR_RE.search(stderr):
|
||||
# Exit code 0 here means tesseract dropped the missing pack and
|
||||
# OCR'd with the rest: Chinese scans would come out as garbage
|
||||
# with no signal. Fail the file and name the fix instead.
|
||||
raise OcrUnavailableError(
|
||||
f"tesseract could not load a language pack for OCR_LANGS={self._language_arg()!r}: "
|
||||
f"{stderr[:300]}. Install the tesseract-ocr-<lang> package for every language listed."
|
||||
)
|
||||
if completed.returncode != 0:
|
||||
raise DocumentParseError(f"tesseract failed (exit {completed.returncode}): {stderr[:300]}")
|
||||
return collapse_cjk_spaces(completed.stdout.decode("utf-8", "replace")).strip()
|
||||
|
||||
|
||||
def collapse_cjk_spaces(text: str) -> str:
|
||||
"""Remove the per-glyph spaces tesseract inserts inside CJK runs."""
|
||||
return _CJK_SPACE_RE.sub("", text)
|
||||
|
||||
|
||||
def clean_deepseek_output(text: str) -> str:
|
||||
"""Strip DeepSeek-OCR grounding markup, keeping the referenced text."""
|
||||
text = _DEEPSEEK_DET_BLOCK_RE.sub("", text)
|
||||
return _DEEPSEEK_TAG_RE.sub("", text).strip()
|
||||
|
||||
|
||||
class DeepseekOcrEngine:
|
||||
"""DeepSeek-OCR over an OpenAI-compatible chat-completions endpoint.
|
||||
|
||||
One request per page, sequentially. The image travels as a base64 data
|
||||
URL, which both Ollama and vLLM accept. Pages are sent one at a time
|
||||
because the model server, not the worker, is the bottleneck; docling's
|
||||
default of four concurrent requests and a 90 s timeout is what broke on
|
||||
a laptop-hosted Ollama in testing.
|
||||
|
||||
Attributes:
|
||||
url: Chat-completions URL (``OCR_DEEPSEEK_URL``).
|
||||
model: Model name at that endpoint (``OCR_DEEPSEEK_MODEL``).
|
||||
timeout: Seconds per page (``OCR_DEEPSEEK_TIMEOUT``).
|
||||
prompt: Instruction sent with every page image.
|
||||
"""
|
||||
|
||||
name = "deepseek"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
timeout: Optional[float] = None,
|
||||
prompt: str = DEEPSEEK_PROMPT,
|
||||
max_tokens: int = _DEEPSEEK_MAX_TOKENS,
|
||||
) -> None:
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
self.url = url or settings.OCR_DEEPSEEK_URL
|
||||
self.model = model or settings.OCR_DEEPSEEK_MODEL
|
||||
self.timeout = float(timeout if timeout is not None else getattr(settings, "OCR_DEEPSEEK_TIMEOUT", 300))
|
||||
self.prompt = prompt
|
||||
self.max_tokens = max_tokens
|
||||
|
||||
def payload(self, image) -> Dict:
|
||||
"""The chat-completions request body for one page image."""
|
||||
data_url = "data:image/png;base64," + base64.b64encode(_png_bytes(image)).decode("ascii")
|
||||
return {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
{"type": "text", "text": self.prompt},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": self.max_tokens,
|
||||
"temperature": 0,
|
||||
}
|
||||
|
||||
def ocr_image(self, image) -> str:
|
||||
"""OCR one PIL image through the endpoint.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: The request failed, timed out, or the
|
||||
response was not a chat completion.
|
||||
"""
|
||||
import requests
|
||||
|
||||
try:
|
||||
response = requests.post(self.url, json=self.payload(image), timeout=self.timeout)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
except requests.exceptions.JSONDecodeError as exc:
|
||||
# Subclasses RequestException too, so it must be caught first or a
|
||||
# proxy's HTML error page is reported as a connection failure.
|
||||
raise DocumentParseError(f"DeepSeek-OCR endpoint {self.url} returned a non-JSON body") from exc
|
||||
except requests.RequestException as exc:
|
||||
raise DocumentParseError(
|
||||
f"DeepSeek-OCR request to {self.url} failed: {exc}. Check OCR_DEEPSEEK_URL "
|
||||
f"and that model {self.model!r} is served (e.g. `ollama pull {self.model}`)."
|
||||
) from exc
|
||||
try:
|
||||
content = body["choices"][0]["message"]["content"]
|
||||
except (KeyError, IndexError, TypeError) as exc:
|
||||
raise DocumentParseError(
|
||||
f"DeepSeek-OCR endpoint {self.url} returned no chat completion: {str(body)[:200]}"
|
||||
) from exc
|
||||
if isinstance(content, list): # some servers return content parts
|
||||
content = "".join(part.get("text", "") for part in content if isinstance(part, dict))
|
||||
return clean_deepseek_output(str(content or ""))
|
||||
|
||||
|
||||
def build_native_ocr_engine(engine: Optional[str] = None, languages: Optional[List[str]] = None) -> OcrEngine:
|
||||
"""Instantiate the engine ``OCR_ENGINE`` (or ``engine``) resolves to for the native backend."""
|
||||
resolved = resolve_native_ocr_engine(engine)
|
||||
if resolved == "deepseek":
|
||||
return DeepseekOcrEngine()
|
||||
return TesseractEngine(languages=languages)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parsers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def render_scale(page_width_pt: float, page_height_pt: float, dpi: int) -> float:
|
||||
"""Render scale for a page of the given size at ``dpi``, capped by ``_MAX_RENDER_PIXELS``.
|
||||
|
||||
The cap is what stands between an oversized (or hostile) MediaBox and a
|
||||
multi-gigabyte bitmap allocation in the ingest worker.
|
||||
"""
|
||||
scale = dpi / 72.0
|
||||
area = max(1.0, float(page_width_pt)) * max(1.0, float(page_height_pt))
|
||||
max_scale = (_MAX_RENDER_PIXELS / area) ** 0.5
|
||||
return min(scale, max_scale)
|
||||
|
||||
|
||||
def _render_page(page, dpi: int):
|
||||
"""Render a pypdfium2 page to a PIL image at ``dpi``, within the pixel budget."""
|
||||
width_pt, height_pt = page.get_size()
|
||||
scale = render_scale(width_pt, height_pt, dpi)
|
||||
if scale < dpi / 72.0:
|
||||
logger.warning(
|
||||
"Page is %.0fx%.0f pt; rendering at %.0f dpi instead of %d to stay within %d MP",
|
||||
width_pt,
|
||||
height_pt,
|
||||
scale * 72.0,
|
||||
dpi,
|
||||
_MAX_RENDER_PIXELS // 1_000_000,
|
||||
)
|
||||
return page.render(scale=scale).to_pil()
|
||||
|
||||
|
||||
def text_layer_counts(path: Path) -> List[int]:
|
||||
"""Characters in each page's text layer (pypdfium2), in page order.
|
||||
|
||||
Raises:
|
||||
Exception: Propagated from pypdfium2 when the file cannot be opened.
|
||||
"""
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
pdf = pdfium.PdfDocument(str(path))
|
||||
counts: List[int] = []
|
||||
try:
|
||||
for index in range(len(pdf)):
|
||||
page = pdf[index]
|
||||
try:
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
counts.append(max(0, textpage.count_chars()))
|
||||
finally:
|
||||
textpage.close()
|
||||
finally:
|
||||
page.close()
|
||||
finally:
|
||||
pdf.close()
|
||||
return counts
|
||||
|
||||
|
||||
def _page_has_image(page) -> bool:
|
||||
"""Whether the page draws at least one image XObject (something OCR could read)."""
|
||||
import pypdfium2.raw as pdfium_c
|
||||
|
||||
try:
|
||||
# pypdfium2's default max_depth=2 misses images inside nested Form
|
||||
# XObjects, which print drivers and InDesign produce routinely.
|
||||
objects = page.get_objects(filter=(pdfium_c.FPDF_PAGEOBJ_IMAGE,), max_depth=16)
|
||||
return next(iter(objects), None) is not None
|
||||
except Exception: # noqa: BLE001 - treat an unreadable page as image-bearing
|
||||
return True
|
||||
|
||||
|
||||
def scanned_page_indices(path: Path, min_chars: int = DEFAULT_TEXT_LAYER_MIN_CHARS) -> List[int]:
|
||||
"""0-based indices of the pages that look scanned: no text layer, or a thin one over an image.
|
||||
|
||||
The probe behind mixed-document handling: a text-layer converter (anydoc)
|
||||
reads the text pages and silently skips these, so whoever owns OCR has to
|
||||
fill them in. A page with no text at all always qualifies. A page with a
|
||||
thin text layer (under ``min_chars``: a Bates stamp on a scan, a page
|
||||
number under a full-page figure) qualifies only when it also draws an
|
||||
image — a cover page reading just "Annual Report 2025" has nothing to OCR,
|
||||
and OCR-ing it would duplicate text the converter already read. Returns []
|
||||
when the file cannot be probed — the caller then keeps what it has rather
|
||||
than failing a successful parse.
|
||||
"""
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
floor = max(1, int(min_chars))
|
||||
indices: List[int] = []
|
||||
try:
|
||||
pdf = pdfium.PdfDocument(str(path))
|
||||
try:
|
||||
for index in range(len(pdf)):
|
||||
page = pdf[index]
|
||||
try:
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
count = max(0, textpage.count_chars())
|
||||
finally:
|
||||
textpage.close()
|
||||
if count == 0 or (count < floor and _page_has_image(page)):
|
||||
indices.append(index)
|
||||
finally:
|
||||
page.close()
|
||||
finally:
|
||||
pdf.close()
|
||||
except Exception: # noqa: BLE001 - a probe must never fail the parse
|
||||
logger.warning("Could not probe %s for scanned pages", Path(path).name, exc_info=True)
|
||||
return []
|
||||
return indices
|
||||
|
||||
|
||||
def _check_near_empty(name: str, engine: str, content: str, pages: int, ocr_pages: int) -> None:
|
||||
"""Loud failure for OCR that produced nothing, a warning for text-sparse documents.
|
||||
|
||||
Mirrors the docling parser's dropout guard: failing only on positive
|
||||
evidence — several pages that OCR'd to literally nothing — because a
|
||||
photo, logo or chart is genuinely text-poor, and ``DocumentParseError``
|
||||
rejects the upload permanently.
|
||||
"""
|
||||
if ocr_pages == 0:
|
||||
return
|
||||
floor = ocr_min_chars_per_page()
|
||||
chars = len(content.strip())
|
||||
if floor <= 0 or chars >= floor * pages:
|
||||
return
|
||||
if chars == 0 and pages > 1:
|
||||
raise DocumentParseError(
|
||||
f"OCR ({engine}) produced no text over {pages} pages of {name}; "
|
||||
"document not indexed. Check the engine's languages (OCR_LANGS) and "
|
||||
"that the scan is legible."
|
||||
)
|
||||
logger.warning(
|
||||
"%s OCR'd (%s) to only %d chars over %d page(s); indexing it as text-sparse",
|
||||
name,
|
||||
engine,
|
||||
chars,
|
||||
pages,
|
||||
)
|
||||
|
||||
|
||||
class NativeOcrPdfParser(BaseParser):
|
||||
"""PDF parser that OCRs pages without a text layer and reads the rest directly.
|
||||
|
||||
Per page: a text layer of at least ``min_text_chars`` characters is read
|
||||
with pypdfium2; anything thinner is rendered and sent to the engine. So
|
||||
a scan, a born-digital PDF and a mix of both all come out right without
|
||||
any prior classification.
|
||||
|
||||
Attributes:
|
||||
engine: The ``OcrEngine``; built from ``OCR_ENGINE`` on first use when None.
|
||||
text_parser: Optional parser for documents whose *every* page has a
|
||||
text layer — used under ``DOC_PARSER_ENGINE=docling`` to keep
|
||||
docling's structured Markdown for text PDFs while scans take the
|
||||
native OCR path. None reads text layers with pypdfium2.
|
||||
min_text_chars: Text-layer characters per page below which the page is OCR'd.
|
||||
ocr_enabled: Always True; lets callers that inspect their fallback
|
||||
parser (``AnydocParser``) phrase their errors correctly.
|
||||
last_engine: Engine name behind the most recent parse.
|
||||
"""
|
||||
|
||||
ocr_enabled = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine: Optional[OcrEngine] = None,
|
||||
text_parser: Optional[BaseParser] = None,
|
||||
min_text_chars: int = DEFAULT_TEXT_LAYER_MIN_CHARS,
|
||||
parser_config: Optional[Dict] = None,
|
||||
) -> None:
|
||||
super().__init__(parser_config)
|
||||
self.engine = engine
|
||||
self.text_parser = text_parser
|
||||
self.min_text_chars = max(1, int(min_text_chars))
|
||||
self.last_engine: Optional[str] = None
|
||||
self._last_metadata: Dict = {}
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
try:
|
||||
import pypdfium2 # noqa: F401
|
||||
from PIL import Image # noqa: F401
|
||||
except ImportError as exc:
|
||||
raise ImportError("pypdfium2 and Pillow are required for NativeOcrPdfParser") from exc
|
||||
if self.engine is None:
|
||||
self.engine = build_native_ocr_engine()
|
||||
return {
|
||||
"engine": self.engine.name,
|
||||
"text_parser": type(self.text_parser).__name__ if self.text_parser else None,
|
||||
"min_text_chars": self.min_text_chars,
|
||||
}
|
||||
|
||||
def _ensure_engine(self) -> OcrEngine:
|
||||
if self.engine is None:
|
||||
self.engine = build_native_ocr_engine()
|
||||
return self.engine
|
||||
|
||||
def _text_layer_counts(self, pdf) -> List[int]:
|
||||
counts: List[int] = []
|
||||
for index in range(len(pdf)):
|
||||
page = pdf[index]
|
||||
try:
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
counts.append(max(0, textpage.count_chars()))
|
||||
finally:
|
||||
textpage.close()
|
||||
finally:
|
||||
page.close()
|
||||
return counts
|
||||
|
||||
def ocr_pages(self, file: Path, indices: List[int]) -> Dict[int, str]:
|
||||
"""OCR only the given 0-based pages of ``file``, ignoring their text layers.
|
||||
|
||||
Used by ``AnydocParser`` for mixed documents: anydoc has already read
|
||||
the text pages, so only the scanned ones come here.
|
||||
|
||||
Returns:
|
||||
Page index -> recognised text, in the order requested.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: The file cannot be opened or the engine failed.
|
||||
"""
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
path = Path(file)
|
||||
engine = self._ensure_engine()
|
||||
try:
|
||||
pdf = pdfium.PdfDocument(str(path))
|
||||
except Exception as exc:
|
||||
raise DocumentParseError(f"Failed to open {path.name} with pypdfium2: {exc}") from exc
|
||||
texts: Dict[int, str] = {}
|
||||
try:
|
||||
dpi = render_dpi()
|
||||
for index in indices:
|
||||
if index < 0 or index >= len(pdf):
|
||||
continue
|
||||
page = pdf[index]
|
||||
try:
|
||||
texts[index] = (engine.ocr_image(_render_page(page, dpi)) or "").strip()
|
||||
finally:
|
||||
page.close()
|
||||
except DocumentParseError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise DocumentParseError(f"Failed to OCR pages of {path.name} ({engine.name}): {exc}") from exc
|
||||
finally:
|
||||
pdf.close()
|
||||
self.last_engine = engine.name
|
||||
return texts
|
||||
|
||||
def text_layer_delegate(self, file: Path) -> Optional[BaseParser]:
|
||||
"""The ``text_parser`` that ``parse_file`` would hand ``file`` to, or None.
|
||||
|
||||
None when there is no text parser, the file cannot be opened, or any
|
||||
page falls below ``min_text_chars`` (those are OCR'd here). Callers
|
||||
that want a by-product of the delegate's conversion (docling tables)
|
||||
can run it directly instead of converting the document twice.
|
||||
"""
|
||||
if self.text_parser is None:
|
||||
return None
|
||||
try:
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
pdf = pdfium.PdfDocument(str(Path(file)))
|
||||
except Exception: # noqa: BLE001 - a probe; parse_file reports the real error
|
||||
return None
|
||||
try:
|
||||
counts = self._text_layer_counts(pdf)
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
finally:
|
||||
pdf.close()
|
||||
if counts and all(count >= self.min_text_chars for count in counts):
|
||||
return self.text_parser
|
||||
return None
|
||||
|
||||
def _delegate_text(self, path: Path, errors: str) -> Union[str, List[str]]:
|
||||
text_parser = self.text_parser
|
||||
assert text_parser is not None
|
||||
logger.info(
|
||||
"%s has a text layer on every page; parsing with %s instead of OCR",
|
||||
path.name,
|
||||
type(text_parser).__name__,
|
||||
)
|
||||
result = delegate_parse(text_parser, path, errors)
|
||||
self.last_engine = getattr(text_parser, "last_engine", None) or type(text_parser).__name__
|
||||
self._last_metadata = {**text_parser.get_file_metadata(path), "parse_engine": self.last_engine}
|
||||
return result
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse a PDF, OCR-ing only the pages that need it.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: The file cannot be opened, the engine failed,
|
||||
or several pages OCR'd to nothing.
|
||||
"""
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
path = Path(file)
|
||||
engine = self._ensure_engine()
|
||||
try:
|
||||
pdf = pdfium.PdfDocument(str(path))
|
||||
except Exception as exc:
|
||||
raise DocumentParseError(f"Failed to open {path.name} with pypdfium2: {exc}") from exc
|
||||
|
||||
pages_text: List[str] = []
|
||||
ocr_pages = 0
|
||||
try:
|
||||
page_count = len(pdf)
|
||||
if page_count == 0:
|
||||
raise DocumentParseError(f"{path.name} has no pages")
|
||||
counts = self._text_layer_counts(pdf)
|
||||
if self.text_parser is not None and all(count >= self.min_text_chars for count in counts):
|
||||
pdf.close()
|
||||
pdf = None
|
||||
return self._delegate_text(path, errors)
|
||||
dpi = render_dpi()
|
||||
for index in range(page_count):
|
||||
page = pdf[index]
|
||||
try:
|
||||
if counts[index] >= self.min_text_chars:
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
text = textpage.get_text_bounded()
|
||||
finally:
|
||||
textpage.close()
|
||||
else:
|
||||
text = engine.ocr_image(_render_page(page, dpi))
|
||||
ocr_pages += 1
|
||||
finally:
|
||||
page.close()
|
||||
pages_text.append((text or "").strip())
|
||||
except DocumentParseError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise DocumentParseError(f"Failed to parse {path.name} with native OCR ({engine.name}): {exc}") from exc
|
||||
finally:
|
||||
if pdf is not None:
|
||||
pdf.close()
|
||||
|
||||
content = "\n\n".join(pages_text)
|
||||
_check_near_empty(path.name, engine.name, content, page_count, ocr_pages)
|
||||
self.last_engine = engine.name
|
||||
self._last_metadata = {"parse_engine": engine.name, "pdf_pages": page_count, "ocr_pages": ocr_pages}
|
||||
logger.info(
|
||||
"Parsed %s with native OCR (%s): %d/%d page(s) OCR'd, %d chars",
|
||||
path.name,
|
||||
engine.name,
|
||||
ocr_pages,
|
||||
page_count,
|
||||
len(content),
|
||||
)
|
||||
return content
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Engine and page counts for the most recently parsed file."""
|
||||
_ = file
|
||||
return dict(self._last_metadata)
|
||||
|
||||
|
||||
class NativeOcrImageParser(BaseParser):
|
||||
"""Image parser that OCRs every frame (multi-page TIFFs included) with the native engine."""
|
||||
|
||||
ocr_enabled = True
|
||||
|
||||
def __init__(self, engine: Optional[OcrEngine] = None, parser_config: Optional[Dict] = None) -> None:
|
||||
super().__init__(parser_config)
|
||||
self.engine = engine
|
||||
self.last_engine: Optional[str] = None
|
||||
self._last_metadata: Dict = {}
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
try:
|
||||
from PIL import Image # noqa: F401
|
||||
except ImportError as exc:
|
||||
raise ImportError("Pillow is required for NativeOcrImageParser") from exc
|
||||
if self.engine is None:
|
||||
self.engine = build_native_ocr_engine()
|
||||
return {"engine": self.engine.name}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""OCR an image file.
|
||||
|
||||
Only TIFF frames are pages (multi-page faxes and scans); every other
|
||||
multi-frame format is an animation (GIF, WebP, APNG) whose frames
|
||||
repeat one picture, so only the first is OCR'd — a 50-frame WebP
|
||||
would otherwise cost 50 tesseract runs and then be rejected as an
|
||||
empty multi-page document.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: The image cannot be decoded, the engine failed,
|
||||
or a multi-frame image OCR'd to nothing.
|
||||
"""
|
||||
from PIL import Image, ImageSequence
|
||||
|
||||
_ = errors
|
||||
path = Path(file)
|
||||
if self.engine is None:
|
||||
self.engine = build_native_ocr_engine()
|
||||
engine = self.engine
|
||||
texts: List[str] = []
|
||||
try:
|
||||
with Image.open(path) as image:
|
||||
frames = ImageSequence.Iterator(image) if path.suffix.lower() in (".tif", ".tiff") else [image]
|
||||
for frame in frames:
|
||||
texts.append((engine.ocr_image(fit_to_pixel_budget(frame)) or "").strip())
|
||||
except DocumentParseError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise DocumentParseError(f"Failed to OCR {path.name} ({engine.name}): {exc}") from exc
|
||||
|
||||
content = "\n\n".join(texts)
|
||||
_check_near_empty(path.name, engine.name, content, max(1, len(texts)), len(texts))
|
||||
self.last_engine = engine.name
|
||||
self._last_metadata = {"parse_engine": engine.name, "ocr_pages": len(texts)}
|
||||
return content
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Engine and frame count for the most recently parsed file."""
|
||||
_ = file
|
||||
return dict(self._last_metadata)
|
||||
@@ -0,0 +1,51 @@
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from openapi_parser import parse
|
||||
|
||||
try:
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
except ModuleNotFoundError:
|
||||
from base_parser import BaseParser
|
||||
|
||||
|
||||
class OpenAPI3Parser(BaseParser):
|
||||
def init_parser(self) -> None:
|
||||
return super().init_parser()
|
||||
|
||||
def get_base_urls(self, urls):
|
||||
base_urls = []
|
||||
for i in urls:
|
||||
parsed_url = urlparse(i)
|
||||
base_url = parsed_url.scheme + "://" + parsed_url.netloc
|
||||
if base_url not in base_urls:
|
||||
base_urls.append(base_url)
|
||||
return base_urls
|
||||
|
||||
def get_info_from_paths(self, path):
|
||||
info = ""
|
||||
if path.operations:
|
||||
for operation in path.operations:
|
||||
info += (
|
||||
f"\n{operation.method.value}="
|
||||
f"{operation.responses[0].description}"
|
||||
)
|
||||
return info
|
||||
|
||||
def parse_file(self, file_path):
|
||||
data = parse(file_path)
|
||||
results = ""
|
||||
base_urls = self.get_base_urls(link.url for link in data.servers)
|
||||
base_urls = ",".join([base_url for base_url in base_urls])
|
||||
results += f"Base URL:{base_urls}\n"
|
||||
i = 1
|
||||
for path in data.paths:
|
||||
info = self.get_info_from_paths(path)
|
||||
results += (
|
||||
f"Path{i}: {path.url}\n"
|
||||
f"description: {path.description}\n"
|
||||
f"parameters: {path.parameters}\nmethods: {info}\n"
|
||||
)
|
||||
i += 1
|
||||
with open("results.txt", "w") as f:
|
||||
f.write(results)
|
||||
return results
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Trust checks for PDF text extraction.
|
||||
|
||||
anydoc silently drops text from fonts it cannot map to Unicode — seen with
|
||||
non-embedded composite (Type0/CID) fonts relying on predefined CMaps, where
|
||||
the EN/ZH NDA corpus file loses its whole Chinese column with no error, and
|
||||
with Adobe-CNS1 fonts in HK legal PDFs where a valid ToUnicode exists but
|
||||
extraction still fails. These dependency-free checks scan the PDF's raw
|
||||
objects (including Flate-compressed object streams) so the pipeline can
|
||||
route such files to a heavier parser, or at least mark the output as
|
||||
unverified, instead of trusting silent partial text.
|
||||
|
||||
Two stages:
|
||||
|
||||
* ``check_pdf_fonts`` — pre-flight on the bytes alone: composite fonts
|
||||
without an embedded ToUnicode CMap, and whether the PDF declares CJK font
|
||||
resources at all.
|
||||
* ``verify_extraction`` — pre-flight plus the cross-check the font scan
|
||||
cannot do: a PDF that declares CJK fonts whose extracted text contains
|
||||
almost no CJK characters is not to be trusted.
|
||||
|
||||
A flag means "verify or route to a fallback", not "the text is wrong": tools
|
||||
shipping Adobe's predefined CMaps (docling's parser does) often extract such
|
||||
files fine. On the benchmark corpus the checks caught both known
|
||||
silent-drop cases with zero false positives on the other 14 PDFs, and cost
|
||||
~30 ms per scanned MB (92 ms on a 3 MB, 150-page annual report).
|
||||
"""
|
||||
import re
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterator, List, Optional, Union
|
||||
|
||||
_OBJ_HEADER = re.compile(rb"\d+\s+\d+\s+obj")
|
||||
_TYPE0 = re.compile(rb"/Subtype\s*/Type0")
|
||||
# The CJK expectation is keyed on CIDSystemInfo /Ordering ALONE — the
|
||||
# authoritative "this PDF maps text through a CJK character collection"
|
||||
# signal, present in every known silent-drop case. Matching CJK font *names*
|
||||
# (SimSun, MS-Gothic, MSungHK, ...) was tried and rejected: Word-exported
|
||||
# Latin PDFs routinely embed an MS-Gothic subset for a stray full-width
|
||||
# character (a 150-page English annual report in the benchmark corpus does),
|
||||
# and substring matching also catches Latin faces like FranklinGothic —
|
||||
# either way English-only documents would be rerouted to the heavy engine.
|
||||
_CJK_ORDERING = re.compile(rb"/Ordering\s*\((GB1|CNS1|Japan1|Japan2|KR|Korea1)\)")
|
||||
# Bytes kept around a Type0 marker found inside a decompressed object
|
||||
# stream: object streams hold many dicts with no obj/endobj markers, and a
|
||||
# font dict's keys sit close to its /Subtype entry.
|
||||
_WINDOW_BEFORE, _WINDOW_AFTER = 200, 800
|
||||
|
||||
# Extracted text with fewer CJK characters than this, from a PDF that
|
||||
# declares CJK fonts, is treated as a silent drop.
|
||||
_MIN_CJK_CHARS = 10
|
||||
_CJK_PROBLEM_PREFIX = "PDF declares CJK fonts"
|
||||
|
||||
_CJK_RANGES = (
|
||||
("一", "鿿"), # CJK Unified Ideographs
|
||||
("", "ヿ"), # Hiragana + Katakana
|
||||
("가", ""), # Hangul syllables
|
||||
)
|
||||
|
||||
|
||||
# Per-stream decompression cap. Flate reaches ~1000:1, so uncapped inflation
|
||||
# of a crafted (or merely image-heavy) PDF inside the upload cap could balloon
|
||||
# to gigabytes and OOM the ingest worker. Font dicts and object streams — the
|
||||
# only things the signals live in — are far smaller than this.
|
||||
_STREAM_INFLATE_CAP = 4_000_000
|
||||
|
||||
|
||||
def _stream_bodies(raw: bytes) -> Iterator[bytes]:
|
||||
"""The raw bytes between each ``stream`` / ``endstream`` keyword pair.
|
||||
|
||||
A ``find`` walk rather than a regex: ``stream\\r?\\n(.*?)\\r?\\nendstream``
|
||||
misses every stream whose data runs straight into ``endstream`` with no
|
||||
EOL (common in the wild), and when it misses one its lazy group scans on
|
||||
to the *next* stream's terminator — quadratic on large files (a 10 MB
|
||||
file measured 32 s, matching a fifth of its streams) and a merged,
|
||||
undecodable body for the rest.
|
||||
"""
|
||||
pos = 0
|
||||
while True:
|
||||
pos = raw.find(b"stream", pos)
|
||||
if pos < 0:
|
||||
return
|
||||
if raw[max(0, pos - 3):pos] == b"end":
|
||||
pos += 6
|
||||
continue
|
||||
start = pos + 6
|
||||
if raw[start:start + 2] == b"\r\n":
|
||||
start += 2
|
||||
elif raw[start:start + 1] in (b"\n", b"\r"):
|
||||
start += 1
|
||||
end = raw.find(b"endstream", start)
|
||||
if end < 0:
|
||||
return
|
||||
body = raw[start:end]
|
||||
if body.endswith(b"\r\n"):
|
||||
body = body[:-2]
|
||||
elif body.endswith((b"\n", b"\r")):
|
||||
body = body[:-1]
|
||||
pos = end + 9
|
||||
yield body
|
||||
|
||||
|
||||
def _object_bodies(raw: bytes) -> Iterator[bytes]:
|
||||
"""The bytes between each ``N G obj`` header and the next ``endobj``.
|
||||
|
||||
A forward walk for the same reason as ``_stream_bodies``: the regex
|
||||
``\\d+\\s+\\d+\\s+obj(.*?)endobj`` rescans to end-of-file for every header
|
||||
that has no ``endobj`` after it, so a malformed (or crafted) PDF with many
|
||||
such headers costs headers x size — 2000 orphan headers in 1 MB measured
|
||||
13.5 s, on the default upload path. Advancing past each ``endobj`` keeps
|
||||
the walk linear and yields the same bodies.
|
||||
"""
|
||||
pos = 0
|
||||
while True:
|
||||
header = _OBJ_HEADER.search(raw, pos)
|
||||
if header is None:
|
||||
return
|
||||
end = raw.find(b"endobj", header.end())
|
||||
if end < 0:
|
||||
return
|
||||
yield raw[header.end():end]
|
||||
pos = end + 6
|
||||
|
||||
|
||||
def _decompressed_streams(raw: bytes) -> Iterator[bytes]:
|
||||
"""Each Flate stream in ``raw`` that decompresses, one at a time, capped."""
|
||||
for body in _stream_bodies(raw):
|
||||
try:
|
||||
inflated = zlib.decompressobj().decompress(body, _STREAM_INFLATE_CAP)
|
||||
except zlib.error:
|
||||
continue
|
||||
yield inflated
|
||||
|
||||
|
||||
def _cjk_chars(text: str, up_to: int) -> int:
|
||||
"""Count CJK characters in ``text``, stopping once ``up_to`` is reached."""
|
||||
count = 0
|
||||
for char in text:
|
||||
for low, high in _CJK_RANGES:
|
||||
if low <= char <= high:
|
||||
count += 1
|
||||
break
|
||||
if count >= up_to:
|
||||
break
|
||||
return count
|
||||
|
||||
|
||||
def check_pdf_fonts(data: bytes) -> Dict[str, Union[int, bool]]:
|
||||
"""Pre-flight font scan of a PDF's bytes.
|
||||
|
||||
Args:
|
||||
data: The complete PDF file contents.
|
||||
|
||||
Returns:
|
||||
Dict with ``type0`` (composite fonts seen), ``type0_no_tounicode``
|
||||
(those without an embedded ToUnicode CMap), ``expects_cjk`` (the PDF
|
||||
declares CJK font resources), ``has_fonts``, and ``flagged`` — True
|
||||
when extraction should not be trusted unverified.
|
||||
"""
|
||||
type0 = bare = 0
|
||||
expects_cjk = bool(_CJK_ORDERING.search(data))
|
||||
has_fonts = b"/Font" in data
|
||||
|
||||
def count_font(unit: bytes) -> None:
|
||||
"""One analysis unit holding a Type0 font dict."""
|
||||
nonlocal type0, bare
|
||||
type0 += 1
|
||||
if b"/ToUnicode" not in unit:
|
||||
bare += 1
|
||||
|
||||
for body in _object_bodies(data):
|
||||
if _TYPE0.search(body):
|
||||
count_font(body)
|
||||
# Each decompressed stream is scanned for every signal in one pass and
|
||||
# then dropped; nothing inflated is retained past its loop iteration.
|
||||
for stream in _decompressed_streams(data):
|
||||
expects_cjk = expects_cjk or bool(_CJK_ORDERING.search(stream))
|
||||
has_fonts = has_fonts or b"/Font" in stream
|
||||
for type0_match in _TYPE0.finditer(stream):
|
||||
start = type0_match.start()
|
||||
count_font(stream[max(0, start - _WINDOW_BEFORE): start + _WINDOW_AFTER])
|
||||
return {
|
||||
"type0": type0,
|
||||
"type0_no_tounicode": bare,
|
||||
"expects_cjk": expects_cjk,
|
||||
"has_fonts": has_fonts,
|
||||
"flagged": bare > 0 or not has_fonts,
|
||||
}
|
||||
|
||||
|
||||
def verify_extraction(data: bytes, markdown: str) -> List[str]:
|
||||
"""Reasons the extracted ``markdown`` of the PDF ``data`` shouldn't be trusted.
|
||||
|
||||
Combines the pre-flight font scan with the post-conversion cross-check it
|
||||
cannot do alone: fonts with a valid ToUnicode that the converter
|
||||
nevertheless failed to extract still show up as missing CJK output.
|
||||
|
||||
Args:
|
||||
data: The complete PDF file contents.
|
||||
markdown: The text a converter extracted from it.
|
||||
|
||||
Returns:
|
||||
Human-readable problem strings; empty when the output looks sound.
|
||||
"""
|
||||
problems: List[str] = []
|
||||
result = check_pdf_fonts(data)
|
||||
if result["type0_no_tounicode"]:
|
||||
problems.append(
|
||||
f"{result['type0_no_tounicode']} composite (Type0) font(s) carry no "
|
||||
"ToUnicode map; extracted text may silently omit glyphs"
|
||||
)
|
||||
elif not result["has_fonts"]:
|
||||
problems.append(
|
||||
"no font resources detected; text extraction cannot be verified"
|
||||
)
|
||||
if result["expects_cjk"]:
|
||||
cjk = _cjk_chars(markdown, _MIN_CJK_CHARS)
|
||||
if cjk < _MIN_CJK_CHARS:
|
||||
problems.append(
|
||||
f"{_CJK_PROBLEM_PREFIX} but the extracted text contains only "
|
||||
f"{cjk} CJK character(s)"
|
||||
)
|
||||
return problems
|
||||
|
||||
|
||||
def _text_layer_cjk_chars(file: Path, up_to: int) -> Optional[int]:
|
||||
"""CJK characters in the text layer as pypdfium2 reads it, stopping at ``up_to``.
|
||||
|
||||
pdfium ships Adobe's predefined CMaps, so it is an independent reference
|
||||
for whether the document *has* CJK text to lose. None when pypdfium2 is
|
||||
unavailable or the file cannot be read that way.
|
||||
"""
|
||||
try:
|
||||
import pypdfium2 as pdfium
|
||||
except ImportError:
|
||||
return None
|
||||
try:
|
||||
pdf = pdfium.PdfDocument(str(file))
|
||||
except Exception: # noqa: BLE001 - a reference probe, not a parse
|
||||
return None
|
||||
count = 0
|
||||
try:
|
||||
for index in range(len(pdf)):
|
||||
page = pdf[index]
|
||||
try:
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
text = textpage.get_text_bounded()
|
||||
finally:
|
||||
textpage.close()
|
||||
finally:
|
||||
page.close()
|
||||
count += _cjk_chars(text, up_to - count)
|
||||
if count >= up_to:
|
||||
break
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
finally:
|
||||
pdf.close()
|
||||
return count
|
||||
|
||||
|
||||
def verify_pdf_file(file: Path, markdown: str) -> List[str]:
|
||||
"""``verify_extraction`` for a file on disk; unreadable files trust the output.
|
||||
|
||||
The CJK cross-check is confirmed against pdfium's own text layer: a
|
||||
stray ``/Ordering (Japan1)`` in an otherwise Latin document (a 48-page
|
||||
Quartz export in the corpus carries one among 258 Identity fonts) is not
|
||||
a dropped Chinese column, and must not send the file to a full docling
|
||||
re-parse.
|
||||
"""
|
||||
try:
|
||||
data = Path(file).read_bytes()
|
||||
except OSError:
|
||||
return []
|
||||
problems = verify_extraction(data, markdown)
|
||||
if any(problem.startswith(_CJK_PROBLEM_PREFIX) for problem in problems):
|
||||
reference = _text_layer_cjk_chars(Path(file), _MIN_CJK_CHARS)
|
||||
if reference is not None and reference < _MIN_CJK_CHARS:
|
||||
problems = [p for p in problems if not p.startswith(_CJK_PROBLEM_PREFIX)]
|
||||
return problems
|
||||
@@ -0,0 +1,261 @@
|
||||
"""Direct PDF text-layer extraction with pypdfium2, for the attachment path.
|
||||
|
||||
Attachments are read straight into a chat prompt; sources are chunked and
|
||||
embedded for retrieval. That difference is what justifies two parsers. Docling
|
||||
produces structured markdown — headings, table markup, image placeholders —
|
||||
which is worth its cost when the output feeds chunking, and mostly wasted when
|
||||
the output is concatenated into a prompt.
|
||||
|
||||
The cost is not marginal. Measured on eight real uploads drawn from the p90 of
|
||||
production attachment-ingest latency (docling 2.119.0, OCR off, warm process):
|
||||
|
||||
docling total 155.6 s pypdfium2 total 0.33 s (~470x)
|
||||
token yield within +/-15% per file, three files yielding *more* text
|
||||
|
||||
What is lost is structure: docling emitted 20-187 headings and 44-88 table rows
|
||||
per file where this parser emits none, and text arrives in content-stream order
|
||||
rather than reading order. Acceptable for prompt context; not acceptable for
|
||||
retrieval, which is why `get_default_file_extractor` keeps docling by default
|
||||
and only swaps in this parser when a caller opts in.
|
||||
|
||||
A scanned PDF has no text layer, so this parser would return an empty document
|
||||
for exactly the files that need docling most. It therefore probes first and
|
||||
delegates to `fallback_parser` when the text layer is too thin to be real,
|
||||
leaving scanned-PDF behaviour unchanged.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser, DocumentParseError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Pages sampled by the text-layer probe. Sampling a spread rather than the
|
||||
# first N matters: scanned documents routinely carry a text cover page, and
|
||||
# born-digital ones sometimes open with a full-page image.
|
||||
DEFAULT_SAMPLE_PAGES = 16
|
||||
|
||||
# Median characters per sampled page below which the file is treated as having
|
||||
# no usable text layer. Measured separation on the production corpus was stark:
|
||||
# scanned files sat at 0 and 17 characters per page, text-layer files at
|
||||
# 433-6834, so anything in the low hundreds separates them. 32 is deliberately
|
||||
# conservative — it keeps sparse-but-real documents on docling.
|
||||
DEFAULT_MIN_MEDIAN_CHARS = 32
|
||||
|
||||
|
||||
class PdfiumTextParser(BaseParser):
|
||||
"""Extract a PDF's embedded text layer with pypdfium2.
|
||||
|
||||
Falls back to ``fallback_parser`` when the file has no usable text layer
|
||||
(a scan), cannot be opened, or when pypdfium2 raises. The fallback is the
|
||||
parser this one displaces, so delegating restores the previous behaviour
|
||||
exactly rather than degrading it.
|
||||
|
||||
Attributes:
|
||||
fallback_parser: Parser used when the text layer is unusable.
|
||||
min_median_chars: Median chars/page required to use the fast path.
|
||||
sample_pages: Maximum number of pages inspected by the probe.
|
||||
last_engine: Name of the engine that produced the most recent parse,
|
||||
``"pypdfium2"`` or the fallback parser's class name.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fallback_parser: Optional[BaseParser] = None,
|
||||
min_median_chars: int = DEFAULT_MIN_MEDIAN_CHARS,
|
||||
sample_pages: int = DEFAULT_SAMPLE_PAGES,
|
||||
parser_config: Optional[Dict] = None,
|
||||
) -> None:
|
||||
"""Initialize the parser.
|
||||
|
||||
Args:
|
||||
fallback_parser: Parser to delegate to when the text layer is
|
||||
unusable. When ``None`` such files raise ``DocumentParseError``.
|
||||
min_median_chars: Median characters per sampled page required to
|
||||
take the fast path.
|
||||
sample_pages: Maximum pages the probe inspects.
|
||||
parser_config: Optional pre-built parser config.
|
||||
"""
|
||||
super().__init__(parser_config)
|
||||
self.fallback_parser = fallback_parser
|
||||
self.min_median_chars = min_median_chars
|
||||
self.sample_pages = max(1, sample_pages)
|
||||
self.last_engine: Optional[str] = None
|
||||
self._last_metadata: Dict[str, Any] = {}
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Initialize the parser config.
|
||||
|
||||
Returns:
|
||||
Dict: Static config describing the probe thresholds.
|
||||
"""
|
||||
return {
|
||||
"engine": "pypdfium2",
|
||||
"min_median_chars": self.min_median_chars,
|
||||
"sample_pages": self.sample_pages,
|
||||
}
|
||||
|
||||
def _sample_indices(self, page_count: int) -> List[int]:
|
||||
"""Return a spread of page indices to probe."""
|
||||
if page_count <= self.sample_pages:
|
||||
return list(range(page_count))
|
||||
step = (page_count - 1) / (self.sample_pages - 1)
|
||||
return sorted({round(i * step) for i in range(self.sample_pages)})
|
||||
|
||||
def _probe_text_layer(self, file: Path) -> Dict[str, Any]:
|
||||
"""Measure how much extractable text the PDF carries.
|
||||
|
||||
Args:
|
||||
file: Path to the PDF.
|
||||
|
||||
Returns:
|
||||
Dict: ``page_count`` and ``median_chars`` across sampled pages.
|
||||
|
||||
Raises:
|
||||
Exception: Propagated from pypdfium2 when the file cannot be read.
|
||||
"""
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
pdf = pdfium.PdfDocument(str(file))
|
||||
try:
|
||||
page_count = len(pdf)
|
||||
counts = []
|
||||
for index in self._sample_indices(page_count):
|
||||
page = pdf[index]
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
counts.append(textpage.count_chars())
|
||||
finally:
|
||||
textpage.close()
|
||||
page.close()
|
||||
finally:
|
||||
pdf.close()
|
||||
return {
|
||||
"page_count": page_count,
|
||||
"median_chars": int(statistics.median(counts)) if counts else 0,
|
||||
}
|
||||
|
||||
def _extract_text(self, file: Path) -> str:
|
||||
"""Extract every page's text layer.
|
||||
|
||||
Args:
|
||||
file: Path to the PDF.
|
||||
|
||||
Returns:
|
||||
str: Page texts joined by blank lines.
|
||||
"""
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
pdf = pdfium.PdfDocument(str(file))
|
||||
pages: List[str] = []
|
||||
try:
|
||||
for index in range(len(pdf)):
|
||||
page = pdf[index]
|
||||
textpage = page.get_textpage()
|
||||
try:
|
||||
pages.append(textpage.get_text_range())
|
||||
finally:
|
||||
textpage.close()
|
||||
page.close()
|
||||
finally:
|
||||
pdf.close()
|
||||
return "\n\n".join(pages)
|
||||
|
||||
def _delegate(self, file: Path, errors: str, reason: str) -> Union[str, List[str]]:
|
||||
"""Hand the file to the fallback parser.
|
||||
|
||||
Args:
|
||||
file: Path to the PDF.
|
||||
errors: Decoding leniency passed through to the fallback.
|
||||
reason: Why the fast path was skipped, for the log line.
|
||||
|
||||
Returns:
|
||||
The fallback parser's output.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: When no fallback parser is configured.
|
||||
"""
|
||||
if self.fallback_parser is None:
|
||||
raise DocumentParseError(
|
||||
f"{file.name} has no usable PDF text layer ({reason}) and no "
|
||||
"fallback parser is configured."
|
||||
)
|
||||
logger.info(
|
||||
"PDF text-layer fast path skipped for %s (%s); using %s",
|
||||
file.name,
|
||||
reason,
|
||||
type(self.fallback_parser).__name__,
|
||||
)
|
||||
if not self.fallback_parser.parser_config_set:
|
||||
self.fallback_parser.init_parser()
|
||||
self.last_engine = type(self.fallback_parser).__name__
|
||||
result = self.fallback_parser.parse_file(file, errors=errors)
|
||||
self._last_metadata = {
|
||||
**self.fallback_parser.get_file_metadata(file),
|
||||
"parse_engine": self.last_engine,
|
||||
}
|
||||
return result
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse a PDF, preferring its embedded text layer.
|
||||
|
||||
Args:
|
||||
file: Path to the PDF.
|
||||
errors: Decoding leniency, passed through to the fallback parser.
|
||||
|
||||
Returns:
|
||||
The extracted text, or the fallback parser's output.
|
||||
|
||||
Raises:
|
||||
DocumentParseError: When the text layer is unusable and no fallback
|
||||
parser is configured.
|
||||
"""
|
||||
file = Path(file)
|
||||
try:
|
||||
probe = self._probe_text_layer(file)
|
||||
except Exception as exc: # noqa: BLE001 - any read failure defers to docling
|
||||
return self._delegate(file, errors, f"unreadable by pypdfium2: {exc}")
|
||||
|
||||
if probe["page_count"] == 0:
|
||||
return self._delegate(file, errors, "no pages")
|
||||
if probe["median_chars"] < self.min_median_chars:
|
||||
return self._delegate(
|
||||
file,
|
||||
errors,
|
||||
f"median {probe['median_chars']} chars/page below "
|
||||
f"{self.min_median_chars}",
|
||||
)
|
||||
|
||||
try:
|
||||
text = self._extract_text(file)
|
||||
except Exception as exc: # noqa: BLE001 - as above
|
||||
return self._delegate(file, errors, f"extraction failed: {exc}")
|
||||
|
||||
self.last_engine = "pypdfium2"
|
||||
self._last_metadata = {
|
||||
"parse_engine": "pypdfium2",
|
||||
"pdf_pages": probe["page_count"],
|
||||
}
|
||||
logger.info(
|
||||
"Parsed %s via pypdfium2 text layer: %d pages, %d chars",
|
||||
file.name,
|
||||
probe["page_count"],
|
||||
len(text),
|
||||
)
|
||||
return text
|
||||
|
||||
def get_file_metadata(self, file: Path) -> Dict:
|
||||
"""Return metadata for the most recently parsed file.
|
||||
|
||||
Args:
|
||||
file: Path to the PDF (unused; state is per-parse).
|
||||
|
||||
Returns:
|
||||
Dict: Includes ``parse_engine`` so a stored attachment records
|
||||
which parser actually produced its content.
|
||||
"""
|
||||
_ = file
|
||||
return dict(self._last_metadata)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""PPT parser.
|
||||
Contains parsers for presentation (.pptx) files to extract slide text.
|
||||
"""
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
|
||||
class PPTXParser(BaseParser):
|
||||
r"""PPTX (.pptx) parser for extracting text from PowerPoint slides.
|
||||
Args:
|
||||
concat_slides (bool): Specifies whether to concatenate all slide text into one document.
|
||||
- If True, slide texts will be joined together as a single string.
|
||||
- If False, each slide's text will be stored as a separate entry in a list.
|
||||
Set to True by default.
|
||||
slide_separator (str): Separator used to join slides' text content.
|
||||
Only used when `concat_slides=True`. Default is "\n".
|
||||
Refer to https://python-pptx.readthedocs.io/en/latest/ for more information.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
concat_slides: bool = True,
|
||||
slide_separator: str = "\n",
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._concat_slides = concat_slides
|
||||
self._slide_separator = slide_separator
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
r"""
|
||||
Parse a .pptx file and extract text from each slide.
|
||||
Args:
|
||||
file (Path): Path to the .pptx file.
|
||||
errors (str): Error handling policy ('ignore' by default).
|
||||
Returns:
|
||||
Union[str, List[str]]: Concatenated text if concat_slides is True,
|
||||
otherwise a list of slide texts.
|
||||
"""
|
||||
|
||||
try:
|
||||
from pptx import Presentation
|
||||
except ImportError:
|
||||
raise ImportError("pptx module is required to read .PPTX files.")
|
||||
|
||||
try:
|
||||
presentation = Presentation(file)
|
||||
slide_texts=[]
|
||||
|
||||
# Iterate over each slide in the presentation
|
||||
for slide in presentation.slides:
|
||||
slide_text=""
|
||||
|
||||
# Iterate over each shape in the slide
|
||||
for shape in slide.shapes:
|
||||
# Check if the shape has a 'text' attribute and append that to the slide_text
|
||||
if hasattr(shape,"text"):
|
||||
slide_text+=shape.text
|
||||
|
||||
slide_texts.append(slide_text.strip())
|
||||
|
||||
if self._concat_slides:
|
||||
return self._slide_separator.join(slide_texts)
|
||||
else:
|
||||
return slide_texts
|
||||
|
||||
except Exception as e:
|
||||
raise e
|
||||
@@ -0,0 +1,201 @@
|
||||
"""reStructuredText parser.
|
||||
|
||||
Contains parser for md files.
|
||||
|
||||
"""
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
|
||||
|
||||
class RstParser(BaseParser):
|
||||
"""reStructuredText parser.
|
||||
|
||||
Extract text from .rst files.
|
||||
Returns dictionary with keys as headers and values as the text between headers.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
remove_hyperlinks: bool = True,
|
||||
remove_images: bool = True,
|
||||
remove_table_excess: bool = True,
|
||||
remove_interpreters: bool = True,
|
||||
remove_directives: bool = True,
|
||||
remove_whitespaces_excess: bool = True,
|
||||
# Be careful with remove_characters_excess, might cause data loss
|
||||
remove_characters_excess: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._remove_hyperlinks = remove_hyperlinks
|
||||
self._remove_images = remove_images
|
||||
self._remove_table_excess = remove_table_excess
|
||||
self._remove_interpreters = remove_interpreters
|
||||
self._remove_directives = remove_directives
|
||||
self._remove_whitespaces_excess = remove_whitespaces_excess
|
||||
self._remove_characters_excess = remove_characters_excess
|
||||
|
||||
def rst_to_tups(self, rst_text: str) -> List[Tuple[Optional[str], str]]:
|
||||
"""Convert a reStructuredText file to a dictionary.
|
||||
|
||||
The keys are the headers and the values are the text under each header.
|
||||
|
||||
"""
|
||||
rst_tups: List[Tuple[Optional[str], str]] = []
|
||||
lines = rst_text.split("\n")
|
||||
|
||||
current_header = None
|
||||
current_text = ""
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
header_match = re.match(r"^[^\S\n]*[-=]+[^\S\n]*$", line)
|
||||
if header_match and i > 0 and (
|
||||
len(lines[i - 1].strip()) == len(header_match.group().strip()) or lines[i - 2] == lines[i - 2]):
|
||||
if current_header is not None:
|
||||
if current_text == "" or None:
|
||||
continue
|
||||
# removes the next heading from current Document
|
||||
if current_text.endswith(lines[i - 1] + "\n"):
|
||||
current_text = current_text[:len(current_text) - len(lines[i - 1] + "\n")]
|
||||
rst_tups.append((current_header, current_text))
|
||||
|
||||
current_header = lines[i - 1]
|
||||
current_text = ""
|
||||
else:
|
||||
current_text += line + "\n"
|
||||
|
||||
rst_tups.append((current_header, current_text))
|
||||
|
||||
# TODO: Format for rst
|
||||
#
|
||||
# if current_header is not None:
|
||||
# # pass linting, assert keys are defined
|
||||
# rst_tups = [
|
||||
# (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
|
||||
# for key, value in rst_tups
|
||||
# ]
|
||||
# else:
|
||||
# rst_tups = [
|
||||
# (key, re.sub("\n", "", value)) for key, value in rst_tups
|
||||
# ]
|
||||
|
||||
if current_header is None:
|
||||
rst_tups = [
|
||||
(key, re.sub("\n", "", value)) for key, value in rst_tups
|
||||
]
|
||||
return rst_tups
|
||||
|
||||
def chunk_by_token_count(self, text: str, max_tokens: int = 100) -> List[str]:
|
||||
"""Chunk text by token count."""
|
||||
|
||||
avg_token_length = 5
|
||||
|
||||
chunk_size = max_tokens * avg_token_length
|
||||
|
||||
chunks = []
|
||||
for i in range(0, len(text), chunk_size):
|
||||
chunk = text[i:i+chunk_size]
|
||||
if i + chunk_size < len(text):
|
||||
last_space = chunk.rfind(' ')
|
||||
if last_space != -1:
|
||||
chunk = chunk[:last_space]
|
||||
|
||||
chunks.append(chunk.strip())
|
||||
|
||||
return chunks
|
||||
|
||||
def remove_images(self, content: str) -> str:
|
||||
pattern = r"\.\. image:: (.*)"
|
||||
content = re.sub(pattern, "", content)
|
||||
return content
|
||||
|
||||
def remove_hyperlinks(self, content: str) -> str:
|
||||
pattern = r"`(.*?) <(.*?)>`_"
|
||||
content = re.sub(pattern, r"\1", content)
|
||||
return content
|
||||
|
||||
def remove_directives(self, content: str) -> str:
|
||||
"""Removes reStructuredText Directives"""
|
||||
pattern = r"`\.\.([^:]+)::"
|
||||
content = re.sub(pattern, "", content)
|
||||
return content
|
||||
|
||||
def remove_interpreters(self, content: str) -> str:
|
||||
"""Removes reStructuredText Interpreted Text Roles"""
|
||||
pattern = r":(\w+):"
|
||||
content = re.sub(pattern, "", content)
|
||||
return content
|
||||
|
||||
def remove_table_excess(self, content: str) -> str:
|
||||
"""Pattern to remove grid table separators"""
|
||||
pattern = r"^\+[-]+\+[-]+\+$"
|
||||
content = re.sub(pattern, "", content, flags=re.MULTILINE)
|
||||
return content
|
||||
|
||||
def remove_whitespaces_excess(self, content: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
|
||||
"""Pattern to match 2 or more consecutive whitespaces"""
|
||||
pattern = r"\s{2,}"
|
||||
content = [(key, re.sub(pattern, " ", value)) for key, value in content]
|
||||
return content
|
||||
|
||||
def remove_characters_excess(self, content: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
|
||||
"""Pattern to match 2 or more consecutive characters"""
|
||||
pattern = r"(\S)\1{2,}"
|
||||
content = [(key, re.sub(pattern, r"\1\1\1", value, flags=re.MULTILINE)) for key, value in content]
|
||||
return content
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Initialize the parser with the config."""
|
||||
return {}
|
||||
|
||||
def parse_tups(
|
||||
self, filepath: Path, errors: str = "ignore",max_tokens: Optional[int] = 1000
|
||||
) -> List[Tuple[Optional[str], str]]:
|
||||
"""Parse file into tuples."""
|
||||
with open(filepath, "r") as f:
|
||||
content = f.read()
|
||||
if self._remove_hyperlinks:
|
||||
content = self.remove_hyperlinks(content)
|
||||
if self._remove_images:
|
||||
content = self.remove_images(content)
|
||||
if self._remove_table_excess:
|
||||
content = self.remove_table_excess(content)
|
||||
if self._remove_directives:
|
||||
content = self.remove_directives(content)
|
||||
if self._remove_interpreters:
|
||||
content = self.remove_interpreters(content)
|
||||
rst_tups = self.rst_to_tups(content)
|
||||
if self._remove_whitespaces_excess:
|
||||
rst_tups = self.remove_whitespaces_excess(rst_tups)
|
||||
if self._remove_characters_excess:
|
||||
rst_tups = self.remove_characters_excess(rst_tups)
|
||||
|
||||
# Apply chunking if max_tokens is provided
|
||||
if max_tokens is not None:
|
||||
chunked_tups = []
|
||||
for header, text in rst_tups:
|
||||
chunks = self.chunk_by_token_count(text, max_tokens)
|
||||
for idx, chunk in enumerate(chunks):
|
||||
chunked_tups.append((f"{header} - Chunk {idx + 1}", chunk))
|
||||
return chunked_tups
|
||||
return rst_tups
|
||||
|
||||
def parse_file(
|
||||
self, filepath: Path, errors: str = "ignore"
|
||||
) -> Union[str, List[str]]:
|
||||
"""Parse file into string."""
|
||||
tups = self.parse_tups(filepath, errors=errors)
|
||||
results = []
|
||||
# TODO: don't include headers right now
|
||||
for header, value in tups:
|
||||
if header is None:
|
||||
results.append(value)
|
||||
else:
|
||||
results.append(f"\n\n{header}\n{value}")
|
||||
return results
|
||||
@@ -0,0 +1,105 @@
|
||||
"""Reconstruct typographic tables in anydoc's PDF Markdown output.
|
||||
|
||||
anydoc converts PDFs straight to Markdown with no document model, so tables
|
||||
drawn with dot leaders or bare whitespace alignment (financial statements,
|
||||
tables of contents) come out as flat text lines — values intact, structure
|
||||
lost. This post-processor detects runs of such lines and rewrites them as
|
||||
GFM tables: the lightweight alternative to a table-structure model for this
|
||||
layout family. (The robust fix is upstream in anydoc's Rust PDF code, where
|
||||
glyph x-positions exist; this is the recoverable-downstream version.)
|
||||
|
||||
Deliberately conservative — a run converts only when it has at least
|
||||
``min_rows`` consecutive lines that each parse as ``label [leaders]
|
||||
numeric-columns`` with the *same* column count; anything else passes through
|
||||
untouched. Gated by ``ANYDOC_TABLEIZE`` (off by default) and applied only to
|
||||
anydoc's own PDF output, never to docling's.
|
||||
"""
|
||||
import re
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
# label ...... 1,234 (56) — dot-leader row
|
||||
_LEADER = re.compile(r"^(.*?)\s*\.{3,}\s*(.+)$")
|
||||
# numeric-ish token: $ 1,234 / (30) / 83,431 / 12.5% / —
|
||||
_NUM_TOKEN = re.compile(r"\$?\(?-?\d[\d,.]*\)?%?|—")
|
||||
# row without leaders: label text, then a trailing run of numeric tokens
|
||||
_TRAILING = re.compile(r"^(.*?[^\s\d,.])\s+([\d$(].*)$")
|
||||
|
||||
|
||||
def _merge_currency(tokens: List[str]) -> List[str]:
|
||||
"""Join a free-standing ``$`` onto the number that follows it."""
|
||||
out: List[str] = []
|
||||
for token in tokens:
|
||||
if out and out[-1] == "$":
|
||||
out[-1] = "$" + token
|
||||
else:
|
||||
out.append(token)
|
||||
return out
|
||||
|
||||
|
||||
def _parse_row(line: str) -> Optional[Tuple[str, List[str]]]:
|
||||
"""``(label, values)`` when the line looks like a typographic table row, else None.
|
||||
|
||||
Without dot leaders a row needs at least two numeric columns: a single
|
||||
trailing number is what "Chapter 1", "Footnote 2", "ISO 9001" and
|
||||
"Version 3.0" look like, and three of those in a row are a list, not a
|
||||
table.
|
||||
"""
|
||||
leader = _LEADER.match(line)
|
||||
match = leader or _TRAILING.match(line)
|
||||
if not match:
|
||||
return None
|
||||
label, rest = match.group(1).strip(" ."), match.group(2)
|
||||
values = _merge_currency(rest.split())
|
||||
if not values or not all(_NUM_TOKEN.fullmatch(v.lstrip("$")) for v in values):
|
||||
return None
|
||||
if leader is None and len(values) < 2:
|
||||
return None
|
||||
if not label or _NUM_TOKEN.fullmatch(label):
|
||||
return None
|
||||
return label, values
|
||||
|
||||
|
||||
def tableize(markdown: str, min_rows: int = 3) -> str:
|
||||
"""Rewrite runs of typographic table rows in ``markdown`` as GFM tables.
|
||||
|
||||
Args:
|
||||
markdown: Converter output to post-process.
|
||||
min_rows: Minimum consecutive, same-width rows for a run to convert.
|
||||
|
||||
Returns:
|
||||
``markdown`` with qualifying runs rewritten; everything else verbatim.
|
||||
"""
|
||||
out: List[str] = []
|
||||
run: List[Tuple[str, List[str], str]] = [] # (label, values, original line)
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal run
|
||||
widths = {len(values) for _, values, _ in run}
|
||||
if len(run) >= min_rows and len(widths) == 1:
|
||||
ncols = widths.pop()
|
||||
out.append("")
|
||||
# GFM needs a header row; an empty one keeps invented labels out
|
||||
# of the indexed text (they would otherwise be embedded as content).
|
||||
out.append("|" + " |" * (ncols + 1))
|
||||
out.append("|" + " --- |" * (ncols + 1))
|
||||
for label, values, _ in run:
|
||||
# Only the label can carry a '|' (values are numeric tokens);
|
||||
# unescaped it would add a cell and mis-column the row.
|
||||
cells = [label.replace("|", "\\|")] + values
|
||||
out.append("| " + " | ".join(cells) + " |")
|
||||
out.append("")
|
||||
else:
|
||||
out.extend(original for _, _, original in run)
|
||||
run = []
|
||||
|
||||
for line in markdown.splitlines():
|
||||
parsed = _parse_row(line.strip()) if line.strip() else None
|
||||
if parsed:
|
||||
run.append((*parsed, line))
|
||||
else:
|
||||
if run:
|
||||
flush()
|
||||
out.append(line)
|
||||
if run:
|
||||
flush()
|
||||
return "\n".join(out)
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Tabular parser.
|
||||
|
||||
Contains parsers for tabular data files.
|
||||
|
||||
"""
|
||||
import datetime
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
from docsgpt.parser.file.base_parser import BaseParser
|
||||
|
||||
|
||||
def cell_to_text(value: Any) -> str:
|
||||
"""Render one spreadsheet/CSV cell as text that is always safe to join.
|
||||
|
||||
``Series.astype(str)`` cannot be trusted for this. Up to pandas 2.x it
|
||||
rendered missing values as the string ``"nan"``; pandas 3.0 *preserves*
|
||||
them, so ``.tolist()`` hands back Python ``float`` NaNs and ``str.join``
|
||||
raises ``TypeError: sequence item N: expected str instance, float found``.
|
||||
A single blank cell anywhere in a sheet was therefore enough to destroy
|
||||
an entire upload.
|
||||
|
||||
Missing values (blank cells, ``NaT``, Excel error values, and formulas
|
||||
with no cached result) become the empty string. Dates and times are
|
||||
rendered ISO-8601 rather than pandas' repr. Integral floats drop the
|
||||
``.0`` that pandas adds when a blank upcasts an integer column — an ID
|
||||
column with one gap otherwise turns every id into ``1001.0``. NUL bytes
|
||||
are stripped here because they crash the downstream Postgres text write.
|
||||
|
||||
Args:
|
||||
value: A single cell value, of any type pandas or openpyxl may yield.
|
||||
|
||||
Returns:
|
||||
str: The cell rendered as text; ``""`` for missing values.
|
||||
"""
|
||||
if value is None:
|
||||
return ""
|
||||
# Handled before the pandas lookup below: it is by far the common case
|
||||
# (every blank cell) and needs no import.
|
||||
if isinstance(value, float) and math.isnan(value):
|
||||
return ""
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
# pandas is optional here — plain float NaN is already covered above,
|
||||
# and pd.NA/NaT can only reach us via pandas in the first place.
|
||||
pd = None
|
||||
if pd is not None:
|
||||
try:
|
||||
if pd.isna(value):
|
||||
return ""
|
||||
except (TypeError, ValueError):
|
||||
# pd.isna raises on array-like cells (lists, ndarrays). Those are
|
||||
# real values rather than missing data, so fall through and
|
||||
# stringify them below.
|
||||
pass
|
||||
if isinstance(value, (datetime.datetime, datetime.date, datetime.time)):
|
||||
return value.isoformat()
|
||||
if isinstance(value, float) and value.is_integer():
|
||||
return str(int(value))
|
||||
text = str(value)
|
||||
return text.replace("\x00", "") if "\x00" in text else text
|
||||
|
||||
|
||||
def _row_to_texts(row: Any) -> List[str]:
|
||||
"""Render a pandas row as a list of join-safe strings."""
|
||||
return [cell_to_text(v) for v in row.tolist()]
|
||||
|
||||
|
||||
class CSVParser(BaseParser):
|
||||
"""CSV parser.
|
||||
|
||||
Args:
|
||||
concat_rows (bool): whether to concatenate all rows into one document.
|
||||
If set to False, a Document will be created for each row.
|
||||
True by default.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, *args: Any, concat_rows: bool = True, **kwargs: Any) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._concat_rows = concat_rows
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse file.
|
||||
|
||||
Returns:
|
||||
Union[str, List[str]]: a string or a List of strings.
|
||||
|
||||
"""
|
||||
try:
|
||||
import csv
|
||||
except ImportError:
|
||||
raise ValueError("csv module is required to read CSV files.")
|
||||
text_list = []
|
||||
with open(file, "r") as fp:
|
||||
csv_reader = csv.reader(fp)
|
||||
for row in csv_reader:
|
||||
text_list.append(", ".join(row))
|
||||
if self._concat_rows:
|
||||
return "\n".join(text_list)
|
||||
else:
|
||||
return text_list
|
||||
|
||||
|
||||
class PandasCSVParser(BaseParser):
|
||||
r"""Pandas-based CSV parser.
|
||||
|
||||
Parses CSVs using the separator detection from Pandas `read_csv`function.
|
||||
If special parameters are required, use the `pandas_config` dict.
|
||||
|
||||
Args:
|
||||
concat_rows (bool): whether to concatenate all rows into one document.
|
||||
If set to False, a Document will be created for each row.
|
||||
True by default.
|
||||
|
||||
col_joiner (str): Separator to use for joining cols per row.
|
||||
Set to ", " by default.
|
||||
|
||||
row_joiner (str): Separator to use for joining each row.
|
||||
Only used when `concat_rows=True`.
|
||||
Set to "\n" by default.
|
||||
|
||||
pandas_config (dict): Options for the `pandas.read_csv` function call.
|
||||
Refer to https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html
|
||||
for more information.
|
||||
Set to empty dict by default, this means pandas will try to figure
|
||||
out the separators, table head, etc. on its own.
|
||||
|
||||
header_period (int): Controls how headers are included in output:
|
||||
- 0: Headers only at the beginning
|
||||
- 1: Headers in every row
|
||||
- N > 1: Headers every N rows
|
||||
|
||||
header_prefix (str): Prefix for header rows. Default is "HEADERS: ".
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
concat_rows: bool = True,
|
||||
col_joiner: str = ", ",
|
||||
row_joiner: str = "\n",
|
||||
pandas_config: dict = {},
|
||||
header_period: int = 20,
|
||||
header_prefix: str = "HEADERS: ",
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._concat_rows = concat_rows
|
||||
self._col_joiner = col_joiner
|
||||
self._row_joiner = row_joiner
|
||||
self._pandas_config = pandas_config
|
||||
self._header_period = header_period
|
||||
self._header_prefix = header_prefix
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse file."""
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
raise ValueError("pandas module is required to read CSV files.")
|
||||
|
||||
df = pd.read_csv(file, **self._pandas_config)
|
||||
headers = [cell_to_text(h) for h in df.columns.tolist()]
|
||||
header_row = f"{self._header_prefix}{self._col_joiner.join(headers)}"
|
||||
|
||||
if not self._concat_rows:
|
||||
return df.apply(
|
||||
lambda row: (self._col_joiner).join(_row_to_texts(row)), axis=1
|
||||
).tolist()
|
||||
|
||||
text_list = []
|
||||
if self._header_period != 1:
|
||||
text_list.append(header_row)
|
||||
|
||||
for i, row in df.iterrows():
|
||||
if (self._header_period > 1 and i > 0 and i % self._header_period == 0):
|
||||
text_list.append(header_row)
|
||||
text_list.append(self._col_joiner.join(_row_to_texts(row)))
|
||||
if self._header_period == 1 and i < len(df) - 1:
|
||||
text_list.append(header_row)
|
||||
|
||||
return self._row_joiner.join(text_list)
|
||||
|
||||
|
||||
class ExcelParser(BaseParser):
|
||||
r"""Excel (.xlsx) parser.
|
||||
|
||||
Parses Excel files using Pandas `read_excel` function.
|
||||
If special parameters are required, use the `pandas_config` dict.
|
||||
|
||||
Args:
|
||||
concat_rows (bool): whether to concatenate all rows into one document.
|
||||
If set to False, a Document will be created for each row.
|
||||
True by default.
|
||||
|
||||
col_joiner (str): Separator to use for joining cols per row.
|
||||
Set to ", " by default.
|
||||
|
||||
row_joiner (str): Separator to use for joining each row.
|
||||
Only used when `concat_rows=True`.
|
||||
Set to "\n" by default.
|
||||
|
||||
pandas_config (dict): Options for the `pandas.read_excel` function call.
|
||||
Refer to https://pandas.pydata.org/docs/reference/api/pandas.read_excel.html
|
||||
for more information.
|
||||
Set to empty dict by default, this means pandas will try to figure
|
||||
out the table structure on its own.
|
||||
|
||||
header_period (int): Controls how headers are included in output:
|
||||
- 0: Headers only at the beginning (default)
|
||||
- 1: Headers in every row
|
||||
- N > 1: Headers every N rows
|
||||
|
||||
header_prefix (str): Prefix for header rows. Default is "HEADERS: ".
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
concat_rows: bool = True,
|
||||
col_joiner: str = ", ",
|
||||
row_joiner: str = "\n",
|
||||
pandas_config: dict = {},
|
||||
header_period: int = 20,
|
||||
header_prefix: str = "HEADERS: ",
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""Init params."""
|
||||
super().__init__(*args, **kwargs)
|
||||
self._concat_rows = concat_rows
|
||||
self._col_joiner = col_joiner
|
||||
self._row_joiner = row_joiner
|
||||
self._pandas_config = pandas_config
|
||||
self._header_period = header_period
|
||||
self._header_prefix = header_prefix
|
||||
|
||||
def _init_parser(self) -> Dict:
|
||||
"""Init parser."""
|
||||
return {}
|
||||
|
||||
def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
|
||||
"""Parse file."""
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
raise ValueError("pandas module is required to read Excel files.")
|
||||
|
||||
df = pd.read_excel(file, **self._pandas_config)
|
||||
headers = [cell_to_text(h) for h in df.columns.tolist()]
|
||||
header_row = f"{self._header_prefix}{self._col_joiner.join(headers)}"
|
||||
|
||||
if not self._concat_rows:
|
||||
return df.apply(
|
||||
lambda row: (self._col_joiner).join(_row_to_texts(row)), axis=1
|
||||
).tolist()
|
||||
|
||||
text_list = []
|
||||
if self._header_period != 1:
|
||||
text_list.append(header_row)
|
||||
|
||||
for i, row in df.iterrows():
|
||||
if (self._header_period > 1 and i > 0 and i % self._header_period == 0):
|
||||
text_list.append(header_row)
|
||||
text_list.append(self._col_joiner.join(_row_to_texts(row)))
|
||||
if self._header_period == 1 and i < len(df) - 1:
|
||||
text_list.append(header_row)
|
||||
return self._row_joiner.join(text_list)
|
||||
Reference in new issue
Block a user