chore: import upstream snapshot with attribution

application/parser/document_reader.py

@@ -0,0 +1,508 @@
+"""In-process document parsing for the ``read_document`` tool, run on the Celery parsing worker.
+
+``parse_document_bytes`` turns untrusted document bytes into a bounded, shaped
+result (markdown/text/structured/chunks) using the BACKEND parsers (Docling by
+default). It applies the same untrusted-content safeguards as uploads — an
+extension whitelist, a byte cap, ``safe_filename`` staging into a temp file, and
+temp cleanup — so a hostile filename or document is treated as inert data.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import os
+import tempfile
+import zipfile
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional
+
+from application.core.settings import settings
+from application.parser.file.bulk import get_default_file_extractor
+from application.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS
+from application.utils import safe_filename
+
+logger = logging.getLogger(__name__)
+
+# Default cap for the LLM-facing VIEW of the text (applied in ``bound_parse_payload``,
+# NOT during parsing) so a huge document can't flood context; the full result is still
+# persisted as a ``data`` artifact. When the text exceeds the cap a head+tail window keeps
+# both the document's beginning AND end (e.g. totals/signatures) within the byte budget.
+_TEXT_MAX_BYTES = 8000
+_MAX_TABLES_RETURNED = 20
+_MAX_TABLE_ROWS = 50
+_MAX_CELL_CHARS = 200
+# Caps applied to the bounded view that rides back through the Redis result
+# backend (the full result still lives in the persisted artifact).
+_MAX_CHUNKS_RETURNED = 50
+_MAX_PAGE_SELECTIONS = 10_000
+_MAX_PAGE_SELECTOR_TOKENS = 20_000
+
+_VALID_OUTPUTS = ("markdown", "text", "structured", "chunks")
+_VALID_OCR = ("auto", "on", "off")
+_VALID_ENGINES = ("auto", "docling", "fast")
+
+
+def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str:
+    """Bound text to a head+tail byte window so a large file can't flood context."""
+    cap = int(max_bytes or _TEXT_MAX_BYTES)
+    if cap <= 0:
+        return text
+    encoded = text.encode("utf-8")
+    if len(encoded) <= cap:
+        return text
+    head = cap // 2
+    tail = cap - head
+    dropped = len(encoded) - head - tail
+    head_text = encoded[:head].decode("utf-8", errors="ignore")
+    tail_text = encoded[-tail:].decode("utf-8", errors="ignore")
+    return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}"
+
+
+def bound_parse_payload(payload: Dict[str, Any], max_chars: Optional[int] = None) -> Dict[str, Any]:
+    """Bound every shape of a parse payload so the Redis-backed result stays small.
+
+    This is where ALL view-bounding happens: parsing now returns the FULL content and the
+    persisted ``data`` artifact keeps it, while the view ridden back through the Redis result
+    backend is bounded here. ``content`` is capped to ``max_chars`` when given, else re-windowed
+    to a head+tail byte window; ``chunks`` is capped in count and per-chunk length. ``structured``
+    is left as-is: it rides back so json_schema validation in the tool can run against it, and it
+    is already bounded by the input byte cap plus the table caps (``_compact_table`` /
+    ``summary``). ``payload['truncated']`` is set when the content view actually cut. The dict is
+    mutated in place.
+    """
+    content = payload.get("content")
+    if isinstance(content, str):
+        if max_chars and int(max_chars) > 0:
+            capped = content[: int(max_chars)]
+        else:
+            capped = truncate_text_head_tail(content)
+        payload["content"] = capped
+        payload["truncated"] = capped != content
+
+    chunks = payload.get("chunks")
+    if isinstance(chunks, list):
+        bounded = [
+            truncate_text_head_tail(chunk) if isinstance(chunk, str) else chunk
+            for chunk in chunks[:_MAX_CHUNKS_RETURNED]
+        ]
+        if len(chunks) > _MAX_CHUNKS_RETURNED:
+            payload["chunks_truncated"] = True
+            payload["total_chunks"] = len(chunks)
+        payload["chunks"] = bounded
+
+    return payload
+
+
+def _max_input_bytes() -> int:
+    """Return the size cap for a parsed document (its own setting, else the sandbox cap)."""
+    explicit = int(getattr(settings, "DOCUMENT_PARSE_MAX_BYTES", 0) or 0)
+    if explicit > 0:
+        return explicit
+    return int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024))
+
+
+_ZIP_CONTAINER_EXTENSIONS = frozenset({".docx", ".xlsx", ".pptx", ".epub"})
+
+
+def _reject_zip_bomb(data: bytes, suffix: str) -> Optional[str]:
+    """Return an error string if a zip-based document declares an implausible expansion, else None."""
+    if suffix not in _ZIP_CONTAINER_EXTENSIONS:
+        return None
+    max_entries = int(getattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000))
+    cap = int(getattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024))
+    try:
+        with zipfile.ZipFile(io.BytesIO(data)) as zf:
+            infos = zf.infolist()
+            if len(infos) > max_entries:
+                return f"document archive has too many entries ({len(infos)} > {max_entries})."
+            total = 0
+            for info in infos:
+                total += info.file_size
+                if total > cap:
+                    return f"document decompresses to too much data: exceeds the {cap}-byte cap."
+    except zipfile.BadZipFile:
+        # Not a readable zip; the format-specific parser will surface a clean error.
+        return None
+    return None
+
+
+def _resolve_ocr_enabled(ocr: str) -> bool:
+    """Resolve the OCR flag from the ``ocr`` arg and the deployment setting."""
+    if ocr == "on":
+        return True
+    if ocr == "off":
+        return False
+    return bool(getattr(settings, "DOCLING_OCR_ENABLED", False))
+
+
+def _pick_parser(suffix: str, *, ocr_enabled: bool, engine: str):
+    """Select the parser for ``suffix`` honoring the requested engine; None when unsupported."""
+    if engine == "fast":
+        legacy = _legacy_parser_for(suffix)
+        if legacy is not None:
+            return legacy
+    extractor = get_default_file_extractor(ocr_enabled=ocr_enabled)
+    return extractor.get(suffix)
+
+
+def _legacy_parser_for(suffix: str):
+    """Return a non-Docling parser for ``suffix`` (the ``fast`` engine), or None."""
+    from application.parser.file.docs_parser import DocxParser, PDFParser
+    from application.parser.file.html_parser import HTMLParser
+    from application.parser.file.markdown_parser import MarkdownParser
+    from application.parser.file.tabular_parser import ExcelParser, PandasCSVParser
+
+    legacy = {
+        ".pdf": PDFParser,
+        ".docx": DocxParser,
+        ".csv": PandasCSVParser,
+        ".xlsx": ExcelParser,
+        ".html": HTMLParser,
+        ".md": MarkdownParser,
+        ".mdx": MarkdownParser,
+    }
+    cls = legacy.get(suffix)
+    return cls() if cls is not None else None
+
+
+def _parse_to_text(parser: Any, path: Path) -> str:
+    """Run a parser and coerce its ``str | List[str]`` result to a single text blob."""
+    if not parser.parser_config_set:
+        parser.init_parser()
+    parsed = parser.parse_file(path, errors="ignore")
+    if isinstance(parsed, list):
+        return "\n\n".join(str(part) for part in parsed)
+    return str(parsed)
+
+
+def _is_docling_parser(parser: Any) -> bool:
+    """True when ``parser`` is Docling-backed (collecting tables would otherwise re-convert)."""
+    if parser is None:
+        return False
+    try:
+        from application.parser.file.docling_parser import DoclingParser
+    except Exception:
+        return False
+    return isinstance(parser, DoclingParser)
+
+
+def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]:
+    """Bound a single table's rows and cell sizes so one giant table can't bloat context."""
+
+    def _cell(value: Any) -> Any:
+        if isinstance(value, str) and len(value) > _MAX_CELL_CHARS:
+            return value[:_MAX_CELL_CHARS] + "...[truncated]"
+        return value
+
+    rows = table.get("rows")
+    if not isinstance(rows, list):
+        return table
+    capped = [[_cell(c) for c in row] if isinstance(row, list) else _cell(row) for row in rows[:_MAX_TABLE_ROWS]]
+    compact = dict(table)
+    compact["rows"] = capped
+    if len(rows) > _MAX_TABLE_ROWS:
+        compact["rows_truncated"] = True
+        compact["total_rows"] = len(rows)
+    return compact
+
+
+def _docling_structured(path: Path, *, ocr_enabled: bool, include_tables: bool, parser: Any = None) -> Dict[str, Any]:
+    """Convert a document with Docling and return markdown + structured dict + bounded tables.
+
+    When ``parser`` is the configured ``DoclingParser`` (the collapse-the-double-convert
+    path), reuse ITS converter + export so OCR/pipeline options and the export fallback
+    are honored and the content matches the legacy single-parse output exactly; otherwise
+    fall back to a vanilla converter.
+    """
+    if _is_docling_parser(parser):
+        if getattr(parser, "_converter", None) is None:
+            parser._init_parser()
+        doc = parser._converter.convert(str(path)).document
+        markdown = parser._export_content(doc)
+    else:
+        from docling.document_converter import DocumentConverter
+
+        converter = DocumentConverter()
+        doc = converter.convert(str(path)).document
+        markdown = doc.export_to_markdown()
+    structured = doc.export_to_dict()
+    tables: List[Dict[str, Any]] = []
+    if include_tables:
+        for tbl in getattr(doc, "tables", []) or []:
+            try:
+                df = tbl.export_to_dataframe()
+                tables.append({"columns": [str(c) for c in df.columns], "rows": df.astype(str).values.tolist()})
+            except Exception:
+                try:
+                    tables.append({"markdown": tbl.export_to_markdown()})
+                except Exception:
+                    continue
+            if len(tables) >= _MAX_TABLES_RETURNED:
+                break
+    page_count = len(getattr(doc, "pages", {}) or {})
+    return {"markdown": markdown, "structured": structured, "tables": tables, "page_count": page_count}
+
+
+def _structure_summary(structured: Any) -> Dict[str, int]:
+    """Summarize the Docling structured dict by top-level element counts (keeps context compact)."""
+    if not isinstance(structured, dict):
+        return {}
+    counts: Dict[str, int] = {}
+    for key in ("texts", "tables", "pictures", "groups", "pages"):
+        value = structured.get(key)
+        if isinstance(value, (list, dict)):
+            counts[key] = len(value)
+    return counts
+
+
+def _apply_pages(text: str, pages: Any) -> str:
+    """Best-effort page-range slice on a page-delimited markdown blob (``\\f`` separated)."""
+    if not pages or "\f" not in text:
+        return text
+    total = text.count("\f") + 1
+    selected = _selected_page_indices(pages, total)
+    if not selected:
+        return text
+
+    # Splitting a hostile form-feed blob could allocate millions of strings.
+    # Walk boundaries and retain only the bounded set of requested pages.
+    wanted = set(selected)
+    found: Dict[int, tuple[int, int]] = {}
+    page_index = 0
+    start = 0
+    while wanted:
+        end = text.find("\f", start)
+        if end < 0:
+            end = len(text)
+        if page_index in wanted:
+            found[page_index] = (start, end)
+            wanted.remove(page_index)
+        if end == len(text):
+            break
+        page_index += 1
+        start = end + 1
+
+    # Preserve requested order and ordinary duplicates, but never let repeated
+    # selectors amplify the output beyond the source text's resident size.
+    output = io.StringIO()
+    output_chars = 0
+    wrote_page = False
+    for index in selected:
+        span = found.get(index)
+        if span is None:
+            continue
+        page_start, page_end = span
+        added = (page_end - page_start) + (1 if wrote_page else 0)
+        if output_chars + added > len(text):
+            break
+        if wrote_page:
+            output.write("\f")
+        output.write(text[page_start:page_end])
+        output_chars += added
+        wrote_page = True
+    return output.getvalue() if wrote_page else text
+
+
+def _iter_page_tokens(pages: Any) -> Iterator[Any]:
+    """Yield bounded selector tokens without materializing a comma-split list."""
+    if isinstance(pages, list):
+        for position, token in enumerate(pages):
+            if position >= _MAX_PAGE_SELECTOR_TOKENS:
+                break
+            yield token
+        return
+
+    raw = str(pages)
+    start = 0
+    emitted = 0
+    while emitted < _MAX_PAGE_SELECTOR_TOKENS:
+        end = raw.find(",", start)
+        if end < 0:
+            yield raw[start:]
+            return
+        yield raw[start:end]
+        emitted += 1
+        start = end + 1
+
+
+def _selected_page_indices(pages: Any, total: int) -> List[int]:
+    """Parse ``pages`` into a bounded list of valid 0-based page occurrences."""
+    if total <= 0:
+        return []
+
+    indices: List[int] = []
+    for token in _iter_page_tokens(pages):
+        if len(indices) >= _MAX_PAGE_SELECTIONS:
+            break
+        token = str(token).strip()
+        if "-" in token:
+            try:
+                lo, hi = (int(p) for p in token.split("-", 1))
+            except ValueError:
+                continue
+            start = max(lo - 1, 0)
+            stop = min(hi, total, start + (_MAX_PAGE_SELECTIONS - len(indices)))
+            indices.extend(range(start, stop))
+        else:
+            try:
+                index = int(token) - 1
+            except ValueError:
+                continue
+            if 0 <= index < total:
+                indices.append(index)
+    return indices
+
+
+def _to_chunks(text: str, max_chars: Optional[int]) -> List[str]:
+    """Chunk parsed text via the ingestion chunker; bounded and JSON-safe for the result."""
+    from application.parser.chunking_creator import ChunkerCreator
+    from application.parser.schema.base import Document
+
+    chunker = ChunkerCreator.create_chunker("classic_chunk")
+    chunks = chunker.chunk([Document(text=text)])
+    cap = int(max_chars or 0)
+    out: List[str] = []
+    for chunk in chunks:
+        body = getattr(chunk, "text", str(chunk))
+        out.append(body[:cap] if cap > 0 else body)
+        if len(out) >= 200:
+            break
+    return out
+
+
+def parse_document_bytes(
+    data: bytes,
+    filename: str,
+    *,
+    output: str = "markdown",
+    ocr: str = "auto",
+    pages: Any = None,
+    engine: str = "auto",
+    max_chars: Optional[int] = None,
+    include_tables: bool = True,
+) -> Dict[str, Any]:
+    """Parse untrusted document bytes into a bounded shaped result; whitelist + size + cleanup guarded."""
+    if output not in _VALID_OUTPUTS:
+        return {"error": f"unsupported output '{output}'; expected one of {_VALID_OUTPUTS}."}
+    if ocr not in _VALID_OCR:
+        return {"error": f"unsupported ocr '{ocr}'; expected one of {_VALID_OCR}."}
+    if engine not in _VALID_ENGINES:
+        return {"error": f"unsupported engine '{engine}'; expected one of {_VALID_ENGINES}."}
+
+    safe_name = safe_filename(filename) or "document"
+    suffix = os.path.splitext(safe_name)[1].lower()
+    if suffix not in SUPPORTED_SOURCE_EXTENSIONS:
+        return {"error": f"unsupported file type '{suffix or filename}'."}
+
+    cap = _max_input_bytes()
+    if len(data) > cap:
+        return {"error": f"input document is too large: {len(data)} bytes exceeds the {cap}-byte cap."}
+
+    bomb = _reject_zip_bomb(data, suffix)
+    if bomb is not None:
+        return {"error": bomb}
+
+    ocr_enabled = _resolve_ocr_enabled(ocr)
+    tmp_dir = tempfile.mkdtemp(prefix="docparse-")
+    tmp_path = Path(tmp_dir) / safe_name
+    try:
+        tmp_path.write_bytes(data)
+        return _shape(tmp_path, suffix, output, ocr_enabled, engine, pages, max_chars, include_tables)
+    except Exception as exc:
+        logger.exception("parse_document_bytes: parsing failed")
+        return {"error": f"parsing failed: {type(exc).__name__}: {exc}"}
+    finally:
+        try:
+            tmp_path.unlink(missing_ok=True)
+            os.rmdir(tmp_dir)
+        except OSError:
+            logger.warning("parse_document_bytes: temp cleanup failed for %s", tmp_dir, exc_info=True)
+
+
+def _shape(
+    path: Path,
+    suffix: str,
+    output: str,
+    ocr_enabled: bool,
+    engine: str,
+    pages: Any,
+    max_chars: Optional[int],
+    include_tables: bool,
+) -> Dict[str, Any]:
+    """Run the selected parser/engine and shape the result per ``output``; bounded throughout.
+
+    ``output='structured'`` always uses Docling regardless of ``engine`` — the ``fast``
+    engine is markdown/text only and cannot produce the structured dict.
+    """
+    if output == "structured":
+        try:
+            extracted = _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=include_tables)
+        except Exception as exc:
+            return {"error": f"structured parsing requires Docling: {type(exc).__name__}: {exc}"}
+        bounded, truncated = _bounded(extracted["markdown"])
+        return {
+            "output": "structured",
+            "content": bounded,
+            "truncated": truncated,
+            "tables": [_compact_table(t) for t in extracted["tables"]],
+            "structured": extracted["structured"],
+            "summary": _structure_summary(extracted["structured"]),
+            "page_count": extracted["page_count"],
+        }
+
+    parser = _pick_parser(suffix, ocr_enabled=ocr_enabled, engine=engine)
+    wants_tables = include_tables and engine != "fast" and output != "chunks"
+
+    # A Docling-backed parser already converts the whole document to produce its text.
+    # When tables are also requested, reuse that single conversion for both the markdown
+    # content and the tables instead of converting a second time just to collect tables
+    # (Docling/torch conversion dominates the cost, so a re-convert ~doubles latency).
+    if wants_tables and _is_docling_parser(parser):
+        try:
+            extracted = _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=True, parser=parser)
+            text = extracted["markdown"]
+            tables: List[Dict[str, Any]] = [_compact_table(t) for t in extracted["tables"]]
+        except Exception:
+            text, tables = _parse_to_text(parser, path), []
+        text = _apply_pages(text, pages)
+        bounded, truncated = _bounded(text)
+        payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
+        if tables:
+            payload["tables"] = tables
+        return payload
+
+    if parser is None:
+        # A whitelisted extension with no dedicated parser (e.g. .txt) reads as plain
+        # text, matching SimpleDirectoryReader's standard-read fallback.
+        text = path.read_text(errors="ignore")
+    else:
+        text = _parse_to_text(parser, path)
+    text = _apply_pages(text, pages)
+
+    if output == "chunks":
+        return {"output": "chunks", "chunks": _to_chunks(text, max_chars), "truncated": False}
+
+    tables: List[Dict[str, Any]] = []
+    if wants_tables:
+        try:
+            tables = [
+                _compact_table(t)
+                for t in _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=True)["tables"]
+            ]
+        except Exception:
+            tables = []
+    bounded, truncated = _bounded(text)
+    payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
+    if tables:
+        payload["tables"] = tables
+    return payload
+
+
+def _bounded(text: str) -> tuple[str, bool]:
+    """Return the FULL extracted text (never truncated here); the view is bounded in ``bound_parse_payload``.
+
+    Parsing keeps the complete text so the persisted ``data`` artifact is the full parse;
+    ``max_chars`` and the default head+tail window now bound only the LLM-facing view.
+    """
+    return text, False