mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-22 07:31:05 +08:00
649 lines
26 KiB
Python
649 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
import html
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from io import BytesIO
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterable, Optional
|
|
|
|
import pdfplumber
|
|
import requests
|
|
from PIL import Image
|
|
|
|
from common.constants import MAXIMUM_PAGE_NUMBER
|
|
|
|
try:
|
|
from deepdoc.parser.pdf_parser import RAGFlowPdfParser
|
|
except Exception:
|
|
|
|
class RAGFlowPdfParser:
|
|
pass
|
|
|
|
|
|
from deepdoc.parser.utils import extract_pdf_outlines
|
|
|
|
|
|
class OpenDataLoaderContentType(str, Enum):
|
|
IMAGE = "image"
|
|
TABLE = "table"
|
|
TEXT = "text"
|
|
EQUATION = "equation"
|
|
|
|
|
|
@dataclass
|
|
class _BBox:
|
|
page_no: int
|
|
x0: float
|
|
y0: float
|
|
x1: float
|
|
y1: float
|
|
|
|
|
|
_TEXT_TYPES = {"heading", "title", "paragraph", "text", "list", "list_item", "caption"}
|
|
_TABLE_TYPES = {"table"}
|
|
_IMAGE_TYPES = {"image", "picture", "figure"}
|
|
_FORMULA_TYPES = {"formula", "equation"}
|
|
|
|
|
|
def _as_float(v) -> Optional[float]:
|
|
try:
|
|
return float(v)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _bbox_from_element(el: dict) -> Optional[_BBox]:
|
|
bb = el.get("bounding box") or el.get("bounding_box") or el.get("bbox")
|
|
pn = el.get("page number")
|
|
if pn is None:
|
|
pn = el.get("page_number")
|
|
if pn is None:
|
|
pn = el.get("page")
|
|
if bb is None or pn is None:
|
|
return None
|
|
if not isinstance(bb, (list, tuple)) or len(bb) < 4:
|
|
return None
|
|
coords = [_as_float(x) for x in bb[:4]]
|
|
if any(c is None for c in coords):
|
|
return None
|
|
try:
|
|
page_no = int(pn)
|
|
except Exception:
|
|
return None
|
|
# OpenDataLoader emits [left, bottom, right, top] in PDF points.
|
|
left, bottom, right, top = coords
|
|
x0, x1 = min(left, right), max(left, right)
|
|
y0, y1 = min(bottom, top), max(bottom, top)
|
|
return _BBox(page_no=page_no, x0=x0, y0=y0, x1=x1, y1=y1)
|
|
|
|
|
|
def _iter_elements(node: Any) -> Iterable[dict]:
|
|
if isinstance(node, dict):
|
|
if "type" in node and ("content" in node or "text" in node or "cells" in node or "table_cells" in node):
|
|
yield node
|
|
for v in node.values():
|
|
yield from _iter_elements(v)
|
|
elif isinstance(node, list):
|
|
for item in node:
|
|
yield from _iter_elements(item)
|
|
|
|
|
|
def _element_text(el: dict) -> str:
|
|
content = el.get("content")
|
|
if isinstance(content, str):
|
|
return content
|
|
text = el.get("text")
|
|
if isinstance(text, str):
|
|
return text
|
|
# tables may expose cells; join row-wise if needed
|
|
cells = el.get("cells")
|
|
if isinstance(cells, list):
|
|
rows: dict[int, list[str]] = {}
|
|
for c in cells:
|
|
if not isinstance(c, dict):
|
|
continue
|
|
row = c.get("row") or c.get("row_index") or 0
|
|
rows.setdefault(int(row), []).append(str(c.get("content") or c.get("text") or ""))
|
|
return "\n".join(" | ".join(v) for _, v in sorted(rows.items()))
|
|
return ""
|
|
|
|
|
|
def _element_html(el: dict) -> str:
|
|
for key in ("html", "html_content"):
|
|
v = el.get(key)
|
|
if isinstance(v, str) and v.strip():
|
|
return v
|
|
# Docling tables: rebuild HTML from table_cells
|
|
cells = el.get("table_cells")
|
|
if isinstance(cells, list) and cells:
|
|
num_rows = el.get("num_rows") or 0
|
|
num_cols = el.get("num_cols") or 0
|
|
if num_rows > 0 and num_cols > 0:
|
|
grid: list[list[dict]] = [[{} for _ in range(num_cols)] for _ in range(num_rows)]
|
|
for c in cells:
|
|
if not isinstance(c, dict):
|
|
continue
|
|
r = c.get("start_row_offset_idx", 0)
|
|
col = c.get("start_col_offset_idx", 0)
|
|
row_span = max(1, c.get("row_span", 1) or 1)
|
|
col_span = max(1, c.get("col_span", 1) or 1)
|
|
text = html.escape(c.get("text", ""))
|
|
if 0 <= r < num_rows and 0 <= col < num_cols:
|
|
grid[r][col] = {"text": text, "row_span": row_span, "col_span": col_span}
|
|
# Mark spanned positions as occupied so they are skipped
|
|
for dr in range(row_span):
|
|
for dc in range(col_span):
|
|
if dr == 0 and dc == 0:
|
|
continue
|
|
rr, cc = r + dr, col + dc
|
|
if 0 <= rr < num_rows and 0 <= cc < num_cols:
|
|
grid[rr][cc] = {"_skip": True}
|
|
rows_html = []
|
|
for row in grid:
|
|
cols = []
|
|
for cell_data in row:
|
|
if cell_data.get("_skip"):
|
|
continue
|
|
if not cell_data:
|
|
continue
|
|
attrs = ""
|
|
if cell_data.get("row_span", 1) > 1:
|
|
attrs += f' rowspan="{cell_data["row_span"]}"'
|
|
if cell_data.get("col_span", 1) > 1:
|
|
attrs += f' colspan="{cell_data["col_span"]}"'
|
|
cols.append(f"<td{attrs}>{cell_data['text']}</td>")
|
|
rows_html.append("<tr>" + "".join(cols) + "</tr>")
|
|
return "<table>" + "".join(rows_html) + "</table>"
|
|
return ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DoclingDocument → intermediate format adapter
|
|
# ---------------------------------------------------------------------------
|
|
# The opendataloader-pdf-hybrid service (/v1/convert/file) returns a
|
|
# DoclingDocument JSON whose structure differs from the original OpenDataLoader
|
|
# format. This adapter normalises it so _transfer_from_json can consume it.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_DOCLING_TEXT_LABELS = {
|
|
"paragraph",
|
|
"text",
|
|
"section_header",
|
|
"title",
|
|
"heading",
|
|
"list_item",
|
|
"code",
|
|
"caption",
|
|
"footnote",
|
|
"page_header",
|
|
"page_footer",
|
|
}
|
|
_DOCLING_TABLE_LABELS = {"table"}
|
|
_DOCLING_PICTURE_LABELS = {"picture", "figure"}
|
|
_DOCLING_FORMULA_LABELS = {"formula"}
|
|
|
|
|
|
def _normalize_docling_response(doc: dict, page_heights: dict[int, float] | None = None) -> list[dict]:
|
|
"""Convert a DoclingDocument JSON into the intermediate element format
|
|
that _iter_elements / _transfer_from_json / _bbox_from_element expect.
|
|
|
|
Each returned element dict has:
|
|
- "type": element type string (table, picture, paragraph, section_header, …)
|
|
- "text": or "content" for text content
|
|
- "page_number": page number
|
|
- "bounding_box": [left, bottom, right, top] flat array in PDF points
|
|
- "html": rebuilt HTML for tables (optional)
|
|
- "cells": / "table_cells": for tables (optional)
|
|
|
|
Parameters:
|
|
doc: the DoclingDocument JSON dict.
|
|
page_heights: optional mapping from page_no (1-based) to page height
|
|
in PDF points. When provided, Y coordinates from
|
|
Docling's TOPLEFT origin are flipped to bottom-left
|
|
origin. When absent the raw values are used (may
|
|
produce incorrect positions).
|
|
"""
|
|
if page_heights is None:
|
|
page_heights = {}
|
|
elements: list[dict] = []
|
|
|
|
# --- texts ---
|
|
for item in doc.get("texts") or []:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
lbl = (item.get("label") or "").lower()
|
|
if lbl in {"page_header", "page_footer"}:
|
|
continue
|
|
el_type = lbl # preserve original label (section_header, paragraph, etc.)
|
|
text = item.get("text", "")
|
|
page_no, bbox_arr = _extract_docling_prov(item, page_heights)
|
|
el: dict[str, Any] = {"type": el_type, "text": text}
|
|
if page_no is not None:
|
|
el["page_number"] = page_no
|
|
if bbox_arr is not None:
|
|
el["bounding_box"] = bbox_arr
|
|
elements.append(el)
|
|
|
|
# --- tables ---
|
|
for item in doc.get("tables") or []:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
el_type = "table"
|
|
page_no, bbox_arr = _extract_docling_prov(item, page_heights)
|
|
el: dict[str, Any] = {"type": el_type, "table_cells": item.get("table_cells", []), "num_rows": item.get("num_rows", 0), "num_cols": item.get("num_cols", 0)}
|
|
if page_no is not None:
|
|
el["page_number"] = page_no
|
|
if bbox_arr is not None:
|
|
el["bounding_box"] = bbox_arr
|
|
elements.append(el)
|
|
|
|
# --- pictures ---
|
|
for item in doc.get("pictures") or []:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
el_type = "picture"
|
|
# Try to get caption from embedded data or referenced captions
|
|
caption = ""
|
|
# Some picture items may have a "data" field with description
|
|
data = item.get("data")
|
|
if isinstance(data, dict):
|
|
caption = data.get("description", "")
|
|
page_no, bbox_arr = _extract_docling_prov(item, page_heights)
|
|
el: dict[str, Any] = {"type": el_type, "content": caption}
|
|
if page_no is not None:
|
|
el["page_number"] = page_no
|
|
if bbox_arr is not None:
|
|
el["bounding_box"] = bbox_arr
|
|
elements.append(el)
|
|
|
|
return elements
|
|
|
|
|
|
def _extract_docling_prov(item: dict, page_heights: dict[int, float] | None = None) -> tuple[Optional[int], Optional[list[float]]]:
|
|
"""Extract page number and bounding box from a Docling item's prov list.
|
|
|
|
Docling prov format: [{"page_no": int, "bbox": {"l": float, "t": float,
|
|
"r": float, "b": float, "coord_origin": "TOPLEFT"}}]
|
|
|
|
Returns (page_no, [left, bottom, right, top]) in PDF points with
|
|
bottom-left origin (matching OpenDataLoader convention).
|
|
|
|
When a page_heights mapping is provided, Docling's TOPLEFT Y coordinates
|
|
are flipped using the page height. Without it the raw values are
|
|
returned as-is (which places the bounding box at the wrong vertical
|
|
position).
|
|
"""
|
|
if page_heights is None:
|
|
page_heights = {}
|
|
prov_list = item.get("prov")
|
|
if not isinstance(prov_list, list) or not prov_list:
|
|
return None, None
|
|
# Use the first provenance entry
|
|
prov = prov_list[0]
|
|
page_no = prov.get("page_no")
|
|
if page_no is None:
|
|
return None, None
|
|
bbox = prov.get("bbox")
|
|
if not isinstance(bbox, dict):
|
|
return None, None
|
|
left = _as_float(bbox.get("l"))
|
|
top_from_top = _as_float(bbox.get("t"))
|
|
right = _as_float(bbox.get("r"))
|
|
bottom_from_top = _as_float(bbox.get("b"))
|
|
if any(v is None for v in (left, top_from_top, right, bottom_from_top)):
|
|
return None, None
|
|
# Docling uses top-left origin (t = distance from top of page).
|
|
# Convert to bottom-left origin using page height when available.
|
|
page_height = page_heights.get(int(page_no))
|
|
if page_height is not None:
|
|
bottom_to_bottom = page_height - bottom_from_top
|
|
top_to_bottom = page_height - top_from_top
|
|
else:
|
|
# Fallback: raw values (incorrect positions but better than nothing)
|
|
bottom_to_bottom = bottom_from_top
|
|
top_to_bottom = top_from_top
|
|
return int(page_no), [left, bottom_to_bottom, right, top_to_bottom]
|
|
|
|
|
|
class OpenDataLoaderParser(RAGFlowPdfParser):
|
|
def __init__(self):
|
|
self.logger = logging.getLogger(self.__class__.__name__)
|
|
self.page_images: list[Image.Image] = []
|
|
self.page_from = 0
|
|
self.page_to = 10_000
|
|
self.outlines = []
|
|
self.api_url = os.environ.get("OPENDATALOADER_APISERVER", "").rstrip("/")
|
|
self.api_key = os.environ.get("OPENDATALOADER_API_KEY", "").strip()
|
|
try:
|
|
self.timeout = int(os.environ.get("OPENDATALOADER_TIMEOUT", "600") or "600")
|
|
except ValueError:
|
|
self.logger.warning("[OpenDataLoader] Invalid OPENDATALOADER_TIMEOUT, falling back to 600s")
|
|
self.timeout = 600
|
|
|
|
def check_installation(self) -> bool:
|
|
"""Return True when the OpenDataLoader service is reachable."""
|
|
if not self.api_url:
|
|
self.logger.warning("[OpenDataLoader] OPENDATALOADER_APISERVER is not set. Start the opendataloader service and set the env var.")
|
|
return False
|
|
try:
|
|
headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
|
|
resp = requests.get(f"{self.api_url}/health", timeout=5, headers=headers)
|
|
if resp.status_code == 200:
|
|
return True
|
|
self.logger.warning(f"[OpenDataLoader] Health check returned {resp.status_code}: {resp.text[:200]}")
|
|
return False
|
|
except Exception as exc:
|
|
self.logger.warning(f"[OpenDataLoader] Health check failed: {exc}")
|
|
return False
|
|
|
|
def __images__(self, fnm, zoomin: int = 1, page_from=0, page_to=MAXIMUM_PAGE_NUMBER, callback=None):
|
|
self.page_from = page_from
|
|
self.page_to = page_to
|
|
bytes_io = None
|
|
try:
|
|
if not isinstance(fnm, (str, PathLike)):
|
|
bytes_io = fnm if isinstance(fnm, BytesIO) else BytesIO(fnm)
|
|
opener = pdfplumber.open(fnm) if isinstance(fnm, (str, PathLike)) else pdfplumber.open(bytes_io)
|
|
with opener as pdf:
|
|
pages = pdf.pages[page_from:page_to]
|
|
self.page_images = [p.to_image(resolution=72 * zoomin, antialias=True).original for p in pages]
|
|
except Exception as e:
|
|
self.page_images = []
|
|
self.logger.exception(e)
|
|
finally:
|
|
if bytes_io:
|
|
bytes_io.close()
|
|
|
|
def _make_line_tag(self, bbox: _BBox) -> str:
|
|
if bbox is None:
|
|
return ""
|
|
# Guard: only emit a crop tag when the page was actually rendered.
|
|
if not self.page_images or bbox.page_no <= 0 or len(self.page_images) < bbox.page_no:
|
|
return ""
|
|
x0, x1 = bbox.x0, bbox.x1
|
|
# OpenDataLoader bbox uses PDF coordinate space (origin bottom-left).
|
|
# Convert to image-space (origin top-left) by subtracting from page height.
|
|
_, page_height = self.page_images[bbox.page_no - 1].size
|
|
top = page_height - bbox.y1
|
|
bott = page_height - bbox.y0
|
|
return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##".format(bbox.page_no, x0, x1, top, bott)
|
|
|
|
@staticmethod
|
|
def extract_positions(txt: str) -> list[tuple[list[int], float, float, float, float]]:
|
|
poss = []
|
|
for tag in re.findall(r"@@[0-9-]+\t[0-9.\t]+##", txt):
|
|
pn, left, right, top, bottom = tag.strip("#").strip("@").split("\t")
|
|
left, right, top, bottom = float(left), float(right), float(top), float(bottom)
|
|
poss.append(([int(p) - 1 for p in pn.split("-")], left, right, top, bottom))
|
|
return poss
|
|
|
|
def crop(self, text: str, ZM: int = 1, need_position: bool = False):
|
|
if not self.page_images:
|
|
return (None, None) if need_position else None
|
|
imgs = []
|
|
poss = self.extract_positions(text)
|
|
if not poss:
|
|
return (None, None) if need_position else None
|
|
# Drop positions whose page indices fall outside the rendered range.
|
|
max_page = len(self.page_images) - 1
|
|
poss = [p for p in poss if all(0 <= pn <= max_page for pn in p[0])]
|
|
if not poss:
|
|
return (None, None) if need_position else None
|
|
GAP = 6
|
|
pos = poss[0]
|
|
poss.insert(0, ([pos[0][0]], pos[1], pos[2], max(0, pos[3] - 120), max(pos[3] - GAP, 0)))
|
|
pos = poss[-1]
|
|
poss.append(([pos[0][-1]], pos[1], pos[2], min(self.page_images[pos[0][-1]].size[1], pos[4] + GAP), min(self.page_images[pos[0][-1]].size[1], pos[4] + 120)))
|
|
positions = []
|
|
for ii, (pns, left, right, top, bottom) in enumerate(poss):
|
|
if bottom <= top:
|
|
bottom = top + 4
|
|
img0 = self.page_images[pns[0]]
|
|
x0, y0, x1, y1 = int(left), int(top), int(right), int(min(bottom, img0.size[1]))
|
|
crop0 = img0.crop((x0, y0, x1, y1))
|
|
imgs.append(crop0)
|
|
if 0 < ii < len(poss) - 1:
|
|
positions.append((pns[0] + self.page_from, x0, x1, y0, y1))
|
|
remain_bottom = bottom - img0.size[1]
|
|
for pn in pns[1:]:
|
|
if remain_bottom <= 0:
|
|
break
|
|
page = self.page_images[pn]
|
|
x0, y0, x1, y1 = int(left), 0, int(right), int(min(remain_bottom, page.size[1]))
|
|
cimgp = page.crop((x0, y0, x1, y1))
|
|
imgs.append(cimgp)
|
|
if 0 < ii < len(poss) - 1:
|
|
positions.append((pn + self.page_from, x0, x1, y0, y1))
|
|
remain_bottom -= page.size[1]
|
|
if not imgs:
|
|
return (None, None) if need_position else None
|
|
height = sum(i.size[1] + GAP for i in imgs)
|
|
width = max(i.size[0] for i in imgs)
|
|
pic = Image.new("RGB", (width, int(height)), (245, 245, 245))
|
|
h = 0
|
|
for ii, img in enumerate(imgs):
|
|
if ii == 0 or ii + 1 == len(imgs):
|
|
img = img.convert("RGBA")
|
|
overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
|
|
overlay.putalpha(128)
|
|
img = Image.alpha_composite(img, overlay).convert("RGB")
|
|
pic.paste(img, (0, int(h)))
|
|
h += img.size[1] + GAP
|
|
return (pic, positions) if need_position else pic
|
|
|
|
def _cropout_region(self, bbox: _BBox, zoomin: int = 1):
|
|
if not self.page_images:
|
|
return None, ""
|
|
idx = (bbox.page_no - 1) - self.page_from
|
|
if idx < 0 or idx >= len(self.page_images):
|
|
return None, ""
|
|
page_img = self.page_images[idx]
|
|
W, H = page_img.size
|
|
x0 = max(0.0, min(float(bbox.x0), W - 1))
|
|
y0 = max(0.0, min(float(H - bbox.y1), H - 1))
|
|
x1 = max(x0 + 1.0, min(float(bbox.x1), W))
|
|
y1 = max(y0 + 1.0, min(float(H - bbox.y0), H))
|
|
try:
|
|
crop = page_img.crop((int(x0), int(y0), int(x1), int(y1))).convert("RGB")
|
|
except Exception:
|
|
return None, ""
|
|
pos = (bbox.page_no - 1 if bbox.page_no > 0 else 0, x0, x1, y0, y1)
|
|
return crop, [pos]
|
|
|
|
def _classify(self, el_type: str) -> str:
|
|
t = (el_type or "").lower()
|
|
if t in _TABLE_TYPES:
|
|
return OpenDataLoaderContentType.TABLE.value
|
|
if t in _IMAGE_TYPES:
|
|
return OpenDataLoaderContentType.IMAGE.value
|
|
if t in _FORMULA_TYPES:
|
|
return OpenDataLoaderContentType.EQUATION.value
|
|
# Preserve the original structural type (heading, title, paragraph,
|
|
# list, caption, …) so downstream parsers can apply heading/title heuristics.
|
|
return t if t else OpenDataLoaderContentType.TEXT.value
|
|
|
|
def _transfer_from_json(self, root: Any, parse_method: str):
|
|
sections: list[tuple[str, ...]] = []
|
|
tables: list = []
|
|
for el in _iter_elements(root):
|
|
el_type = self._classify(el.get("type", ""))
|
|
bbox = _bbox_from_element(el)
|
|
tag = self._make_line_tag(bbox) if bbox else ""
|
|
|
|
if el_type == OpenDataLoaderContentType.TABLE.value:
|
|
html = _element_html(el) or _element_text(el)
|
|
img = None
|
|
positions = ""
|
|
if bbox:
|
|
img, positions = self._cropout_region(bbox)
|
|
tables.append(((img, html), positions if positions else ""))
|
|
continue
|
|
|
|
if el_type == OpenDataLoaderContentType.IMAGE.value:
|
|
img = None
|
|
positions = ""
|
|
if bbox:
|
|
img, positions = self._cropout_region(bbox)
|
|
caption = _element_text(el)
|
|
tables.append(((img, [caption] if caption else [""]), positions if positions else ""))
|
|
continue
|
|
|
|
text = _element_text(el).strip()
|
|
if not text:
|
|
continue
|
|
if parse_method in {"manual", "pipeline"}:
|
|
sections.append((text, el_type, tag))
|
|
elif parse_method == "paper":
|
|
sections.append((text + tag, el_type))
|
|
else:
|
|
sections.append((text, tag))
|
|
return sections, tables
|
|
|
|
@staticmethod
|
|
def _sections_from_markdown(md: str, parse_method: str) -> list[tuple[str, ...]]:
|
|
txt = (md or "").strip()
|
|
if not txt:
|
|
return []
|
|
if parse_method in {"manual", "pipeline"}:
|
|
return [(txt, OpenDataLoaderContentType.TEXT.value, "")]
|
|
if parse_method == "paper":
|
|
return [(txt, OpenDataLoaderContentType.TEXT.value)]
|
|
return [(txt, "")]
|
|
|
|
def parse_pdf(
|
|
self,
|
|
filepath: str | PathLike[str],
|
|
binary: BytesIO | bytes | None = None,
|
|
callback: Optional[Callable] = None,
|
|
*,
|
|
parse_method: str = "raw",
|
|
hybrid: Optional[str] = None,
|
|
image_output: Optional[str] = None,
|
|
sanitize: Optional[bool] = None,
|
|
):
|
|
self.outlines = extract_pdf_outlines(binary if binary is not None else filepath)
|
|
|
|
if not self.api_url:
|
|
raise RuntimeError("[OpenDataLoader] OPENDATALOADER_APISERVER is not configured. Please start the opendataloader service and set the env var.")
|
|
|
|
# Render page images locally — used by _make_line_tag() and crop().
|
|
# The image rendering stays on the RAGFlow host; only the Java conversion
|
|
# runs inside the opendataloader service container.
|
|
try:
|
|
if binary is not None:
|
|
src = BytesIO(binary) if isinstance(binary, (bytes, bytearray)) else binary
|
|
self.__images__(src, zoomin=1)
|
|
else:
|
|
self.__images__(str(filepath), zoomin=1)
|
|
except Exception as e:
|
|
self.logger.warning(f"[OpenDataLoader] render pages failed: {e}")
|
|
|
|
# Build page-height lookup for Docling TOPLEFT → bottom-left
|
|
# coordinate conversion. Pages are 1-based in Docling.
|
|
page_heights: dict[int, float] = {}
|
|
for i, img in enumerate(self.page_images):
|
|
if img is not None:
|
|
page_heights[i + 1] = float(img.size[1])
|
|
|
|
# Read PDF bytes for the multipart upload
|
|
if binary is not None:
|
|
pdf_bytes = binary if isinstance(binary, (bytes, bytearray)) else binary.getvalue()
|
|
else:
|
|
with open(filepath, "rb") as fh:
|
|
pdf_bytes = fh.read()
|
|
|
|
filename = Path(str(filepath)).name or "input.pdf"
|
|
|
|
if callback:
|
|
callback(0.1, f"[OpenDataLoader] Sending '{filename}' to service")
|
|
|
|
# Docling Fast Server endpoint: POST /v1/convert/file
|
|
# - files: PDF file (multipart/form-data)
|
|
# - page_ranges: optional, e.g. "1-5"
|
|
form_data: dict[str, str] = {"to_formats": '["json", "md"]'}
|
|
if hybrid:
|
|
form_data["hybrid"] = hybrid
|
|
if image_output:
|
|
form_data["image_output"] = image_output
|
|
if sanitize is not None:
|
|
form_data["sanitize"] = "true" if sanitize else "false"
|
|
|
|
headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
|
|
last_exc: Exception | None = None
|
|
for attempt in range(1, 4):
|
|
try:
|
|
self.logger.info(f"[OpenDataLoader] POST {self.api_url}/v1/convert/file for '{filename}' (attempt {attempt})")
|
|
resp = requests.post(
|
|
url=f"{self.api_url}/v1/convert/file",
|
|
files={"files": (filename, pdf_bytes, "application/pdf")},
|
|
data=form_data,
|
|
headers=headers,
|
|
timeout=self.timeout,
|
|
)
|
|
resp.raise_for_status()
|
|
result = resp.json()
|
|
break
|
|
except Exception as exc:
|
|
last_exc = exc
|
|
self.logger.warning(f"[OpenDataLoader] attempt {attempt} failed: {exc}")
|
|
else:
|
|
raise RuntimeError(f"[OpenDataLoader] service call failed after 3 attempts: {last_exc}") from last_exc
|
|
|
|
if callback:
|
|
callback(0.7, "[OpenDataLoader] Processing response")
|
|
|
|
# /v1/convert/file response structure (nested):
|
|
# {"status": ..., "document": {"json_content": {DoclingDocument}}, "processing_time": ...}
|
|
json_doc = None
|
|
md_text = None
|
|
if isinstance(result, dict):
|
|
doc = result.get("document")
|
|
inner = None
|
|
if isinstance(doc, dict):
|
|
inner = doc.get("json_content")
|
|
if isinstance(inner, str):
|
|
try:
|
|
inner = json.loads(inner)
|
|
except (json.JSONDecodeError, TypeError):
|
|
self.logger.warning("[OpenDataLoader] json_content is a string but not valid JSON")
|
|
inner = None
|
|
|
|
# Use DoclingDocument if found at inner, otherwise fall back
|
|
if isinstance(inner, dict) and ("texts" in inner or "tables" in inner or "pictures" in inner):
|
|
json_doc = _normalize_docling_response(inner, page_heights)
|
|
elif isinstance(doc, dict) and ("texts" in doc or "tables" in doc or "pictures" in doc):
|
|
json_doc = _normalize_docling_response(doc, page_heights)
|
|
elif isinstance(result, dict) and ("texts" in result or "tables" in result or "pictures" in result):
|
|
json_doc = _normalize_docling_response(result, page_heights)
|
|
else:
|
|
# Legacy / flat format
|
|
json_doc = result.get("json_doc")
|
|
md_text = result.get("md_text")
|
|
elif isinstance(result, str):
|
|
md_text = result
|
|
|
|
sections: list[tuple[str, ...]] = []
|
|
tables: list = []
|
|
if json_doc is not None:
|
|
sections, tables = self._transfer_from_json(json_doc, parse_method=parse_method)
|
|
self.logger.info(f"[OpenDataLoader] Extracted {len(sections)} sections, {len(tables)} tables")
|
|
if not sections and md_text:
|
|
sections = self._sections_from_markdown(md_text, parse_method=parse_method)
|
|
|
|
if callback:
|
|
callback(1.0, f"[OpenDataLoader] Done. Sections: {len(sections)}, Tables: {len(tables)}")
|
|
|
|
return sections, tables
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO)
|
|
parser = OpenDataLoaderParser()
|
|
print("OpenDataLoader service reachable:", parser.check_installation())
|