From 4fcec8039b857de62236be07171d4a295921483a Mon Sep 17 00:00:00 2001 From: Vadim Malanov Date: Thu, 25 Jun 2026 11:43:29 +0300 Subject: [PATCH] chore: pin document recognition engine --- README.md | 6 +- app/ingestion/ENGINE_SOURCE.md | 53 +++ app/ingestion/docling_extractor.py | 386 +----------------- app/ingestion/ocr.py | 89 +--- docs/ADR-shared-core.md | 12 +- pyproject.toml | 4 + ...est_document_recognition_engine_adapter.py | 63 +++ 7 files changed, 154 insertions(+), 459 deletions(-) create mode 100644 app/ingestion/ENGINE_SOURCE.md create mode 100644 tests/test_document_recognition_engine_adapter.py diff --git a/README.md b/README.md index 1bf33d4..ef3965a 100644 --- a/README.md +++ b/README.md @@ -22,8 +22,8 @@ FastAPI /search ◀── BGE Reranker ◀── RRF merge ◀────── | Component | Tech | |------------------|------------------------------------------| -| OCR | OCRmyPDF + Tesseract (rus + eng) | -| Extraction | Docling (layout, tables, figures) | +| OCR | `teamhub-document-recognition-engine` over OCRmyPDF + Tesseract (rus + eng) | +| Extraction | `teamhub-document-recognition-engine` over Docling (layout, tables, figures) | | Object storage | MinIO (S3-compatible) | | Relational store | PostgreSQL 16 | | Lexical search | OpenSearch 2.x (BM25 + ru/en analyzers) | @@ -196,7 +196,7 @@ docker compose exec postgres psql -U legacyhub -d legacyhub \ See [`.env.example`](.env.example) for the full list. Key ones: -- `OCR_LANGUAGES` - Tesseract language packs (default `rus+eng`). +- `OCR_LANGUAGES` - Tesseract language packs passed to the shared recognition engine (default `rus+eng`). - `OCR_ENABLED` - set `false` to skip OCR completely. - `DOCLING_OCR_ENABLED` - prefer OCRmyPDF; only enable if you do not run OCRmyPDF. - `EMBEDDING_DEVICE` / `RERANKER_DEVICE` - `cpu`, `cuda`, or `mps`. diff --git a/app/ingestion/ENGINE_SOURCE.md b/app/ingestion/ENGINE_SOURCE.md new file mode 100644 index 0000000..108ec27 --- /dev/null +++ b/app/ingestion/ENGINE_SOURCE.md @@ -0,0 +1,53 @@ +# Engine Source + +Engine: `teamhub-document-recognition-engine` + +Package: `teamhub-document-recognition-engine` + +Version: `0.1.0` + +Source repository: `C:\Users\manag\TeamHUB\Engines` + +Source path: +`engines/teamhub-document-recognition-engine` + +Source reference status: release-tagged dependency pin. + +Source tag: `teamhub-document-recognition-engine-v0.1.0` + +Tag SHA: `efa29f9a27f890a40dd9081773973d8b4ed03509` + +Tag target commit: `0c9a1b5b8d6879bd44aeb792114431537c14bb2c` + +Dependency file: `pyproject.toml` + +Dependency spec: +`teamhub-document-recognition-engine @ git+https://github.com/GoodOneFather/TeamHUB_Engines.git@teamhub-document-recognition-engine-v0.1.0#subdirectory=engines/teamhub-document-recognition-engine` + +Adapter paths: + +- `app/ingestion/ocr.py` +- `app/ingestion/docling_extractor.py` + +Copied artifacts: none. LegacyHUB imports the package and keeps only thin +adapters. + +Rollback path: restore the previous local OCR/Docling implementation in the two +adapter files and remove the dependency entry from `pyproject.toml`. + +Removal condition: replace this git-tag dependency with a registry package +when a TeamHUB package registry is available. + +Focused verification: + +```powershell +python -m pytest tests/test_document_recognition_engine_adapter.py -q +``` + +Verified against: + +```powershell +C:\Users\manag\.cache\codex-runtimes\codex-primary-runtime\dependencies\python\python.exe -m pytest tests\test_document_recognition_engine_adapter.py -q +``` + +Verified on: `2026-06-25` diff --git a/app/ingestion/docling_extractor.py b/app/ingestion/docling_extractor.py index 1ad9cd2..b4afa4d 100644 --- a/app/ingestion/docling_extractor.py +++ b/app/ingestion/docling_extractor.py @@ -1,384 +1,20 @@ -"""Docling structured extraction. - -Docling produces a hierarchical document model with reading order, layout, tables -and figures. We export both Markdown and a JSON representation, then walk the -JSON to emit normalized blocks (title, heading, paragraph, list, table caption, -figure caption) for downstream chunking. - -The extractor is intentionally defensive: Docling's exact Python API has -shifted across releases. We probe for the safest exporter methods and fall -back to ``str(document)`` only as a last resort. -""" +"""LegacyHUB adapter for the shared Docling recognition engine.""" from __future__ import annotations -import json -from dataclasses import dataclass, field from pathlib import Path -from typing import Any + +from teamhub_document_recognition_engine import ( + ExtractedBlock, + ExtractedFigure, + ExtractedPage, + ExtractedTable, + ExtractionResult, +) +from teamhub_document_recognition_engine import extract as _engine_extract from app.config import settings -from app.logging_config import get_logger - -logger = get_logger(__name__) - - -@dataclass -class ExtractedBlock: - page_number: int - block_type: str - text: str - block_id: str | None = None - extra: dict[str, Any] = field(default_factory=dict) - - -@dataclass -class ExtractedTable: - page_number: int - table_index: int - markdown: str - csv_text: str | None = None - json_data: dict[str, Any] | None = None - block_id: str | None = None - - -@dataclass -class ExtractedFigure: - page_number: int - figure_index: int - caption: str | None - block_id: str | None = None - image_bytes: bytes | None = None - image_ext: str = "png" - - -@dataclass -class ExtractedPage: - page_number: int - text: str - has_tables: bool = False - has_figures: bool = False - has_handwriting: bool = False - ocr_confidence: float | None = None - - -@dataclass -class ExtractionResult: - markdown: str - json_payload: dict[str, Any] - blocks: list[ExtractedBlock] - tables: list[ExtractedTable] - figures: list[ExtractedFigure] - pages: list[ExtractedPage] def extract(pdf_path: Path) -> ExtractionResult: - """Run Docling on ``pdf_path`` and return a normalized result.""" - from docling.datamodel.base_models import InputFormat - from docling.datamodel.pipeline_options import PdfPipelineOptions - from docling.document_converter import DocumentConverter, PdfFormatOption - - pipeline_options = PdfPipelineOptions() - # We let OCRmyPDF do the heavy OCR; Docling OCR is opt-in. - pipeline_options.do_ocr = settings.docling_ocr_enabled - pipeline_options.do_table_structure = True - try: - pipeline_options.table_structure_options.do_cell_matching = True - except Exception: # noqa: BLE001 - older docling versions lack this - pass - try: - pipeline_options.generate_page_images = True - except Exception: # noqa: BLE001 - pass - - converter = DocumentConverter( - format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)} - ) - - logger.info("docling.start", input=str(pdf_path)) - conv = converter.convert(str(pdf_path)) - doc = conv.document - - markdown = _safe_export_markdown(doc) - json_payload = _safe_export_dict(doc) - - blocks = _walk_blocks(json_payload) - tables = _walk_tables(doc, json_payload) - figures = _walk_figures(doc, json_payload) - pages = _walk_pages(json_payload, blocks, tables, figures) - - logger.info( - "docling.done", - pages=len(pages), - blocks=len(blocks), - tables=len(tables), - figures=len(figures), - ) - return ExtractionResult( - markdown=markdown, - json_payload=json_payload, - blocks=blocks, - tables=tables, - figures=figures, - pages=pages, - ) - - -# ---------------- Internal helpers ---------------- - -def _safe_export_markdown(doc: Any) -> str: - for attr in ("export_to_markdown", "to_markdown"): - fn = getattr(doc, attr, None) - if callable(fn): - try: - return fn() - except Exception: # noqa: BLE001 - continue - return str(doc) - - -def _safe_export_dict(doc: Any) -> dict[str, Any]: - for attr in ("export_to_dict", "model_dump", "dict"): - fn = getattr(doc, attr, None) - if callable(fn): - try: - data = fn() - if isinstance(data, dict): - return data - except Exception: # noqa: BLE001 - continue - # Last resort: serialize via JSON round-trip - try: - return json.loads(getattr(doc, "model_dump_json", lambda: "{}")()) - except Exception: # noqa: BLE001 - return {} - - -_DOCLING_LABEL_TO_BLOCK = { - "title": "title", - "section_header": "heading", - "section-header": "heading", - "subtitle": "heading", - "page_header": "heading", - "header": "heading", - "list_item": "list", - "list-item": "list", - "list": "list", - "paragraph": "paragraph", - "text": "paragraph", - "caption": "figure_caption", - "figure": "figure_caption", - "table": "table", - "footnote": "paragraph", -} - - -def _walk_blocks(payload: dict[str, Any]) -> list[ExtractedBlock]: - """Flatten Docling's text items into ordered blocks per page.""" - blocks: list[ExtractedBlock] = [] - items = ( - payload.get("texts") - or payload.get("text_items") - or payload.get("body", {}).get("text_items", []) - or [] - ) - if not isinstance(items, list): - return blocks - - for item in items: - if not isinstance(item, dict): - continue - label = (item.get("label") or item.get("category") or "paragraph").lower() - text = (item.get("text") or "").strip() - if not text: - continue - block_type = _DOCLING_LABEL_TO_BLOCK.get(label, "paragraph") - page = _page_of(item) - blocks.append( - ExtractedBlock( - page_number=page, - block_type=block_type, - text=text, - block_id=item.get("self_ref") or item.get("id"), - extra={"label": label}, - ) - ) - return blocks - - -def _walk_tables(doc: Any, payload: dict[str, Any]) -> list[ExtractedTable]: - tables: list[ExtractedTable] = [] - raw_tables = payload.get("tables") or [] - for idx, t in enumerate(raw_tables): - if not isinstance(t, dict): - continue - page = _page_of(t) - md = _table_markdown(doc, t, idx) - csv_text = _table_csv(t) - tables.append( - ExtractedTable( - page_number=page, - table_index=idx, - markdown=md, - csv_text=csv_text, - json_data=t, - block_id=t.get("self_ref") or t.get("id"), - ) - ) - return tables - - -def _walk_figures(doc: Any, payload: dict[str, Any]) -> list[ExtractedFigure]: - figures: list[ExtractedFigure] = [] - raw_figures = payload.get("pictures") or payload.get("figures") or [] - for idx, f in enumerate(raw_figures): - if not isinstance(f, dict): - continue - page = _page_of(f) - caption = (f.get("caption") or "").strip() or None - figures.append( - ExtractedFigure( - page_number=page, - figure_index=idx, - caption=caption, - block_id=f.get("self_ref") or f.get("id"), - ) - ) - return figures - - -def _walk_pages( - payload: dict[str, Any], - blocks: list[ExtractedBlock], - tables: list[ExtractedTable], - figures: list[ExtractedFigure], -) -> list[ExtractedPage]: - pages_meta = payload.get("pages") or {} - page_numbers: set[int] = set() - if isinstance(pages_meta, dict): - for k in pages_meta.keys(): - try: - page_numbers.add(int(k)) - except (ValueError, TypeError): - continue - elif isinstance(pages_meta, list): - for p in pages_meta: - if isinstance(p, dict): - pn = p.get("page_no") or p.get("page") or p.get("number") - if isinstance(pn, int): - page_numbers.add(pn) - - for b in blocks: - page_numbers.add(b.page_number) - for t in tables: - page_numbers.add(t.page_number) - for f in figures: - page_numbers.add(f.page_number) - page_numbers.discard(0) - if not page_numbers: - page_numbers = {1} - - by_page_text: dict[int, list[str]] = {pn: [] for pn in page_numbers} - for b in blocks: - by_page_text.setdefault(b.page_number, []).append(b.text) - - has_tables_set = {t.page_number for t in tables} - has_figures_set = {f.page_number for f in figures} - - return [ - ExtractedPage( - page_number=pn, - text="\n\n".join(by_page_text.get(pn, [])), - has_tables=pn in has_tables_set, - has_figures=pn in has_figures_set, - ) - for pn in sorted(page_numbers) - ] - - -def _page_of(item: dict[str, Any]) -> int: - prov = item.get("prov") or item.get("provenance") - if isinstance(prov, list) and prov: - first = prov[0] - if isinstance(first, dict): - pn = first.get("page_no") or first.get("page") or first.get("page_number") - if isinstance(pn, int): - return pn - pn = item.get("page_no") or item.get("page") or item.get("page_number") - if isinstance(pn, int): - return pn - return 1 - - -def _table_markdown(doc: Any, raw: dict[str, Any], idx: int) -> str: - # Try Docling's own export first (per-table). - try: - export = getattr(doc, "export_table_to_markdown", None) - if callable(export): - return export(idx) - except Exception: # noqa: BLE001 - pass - - grid = raw.get("data") or raw.get("table_cells") or raw.get("grid") - if isinstance(grid, list) and grid and isinstance(grid[0], list): - return _grid_to_markdown(grid) - cells = raw.get("table_cells") - if isinstance(cells, list): - return _cells_to_markdown(cells) - return "" - - -def _grid_to_markdown(grid: list[list[Any]]) -> str: - if not grid: - return "" - - def _cell(c: Any) -> str: - if isinstance(c, dict): - return str(c.get("text") or c.get("value") or "").replace("|", "\\|").strip() - return str(c).replace("|", "\\|").strip() - - header = grid[0] - body = grid[1:] if len(grid) > 1 else [] - cols = len(header) - out = ["| " + " | ".join(_cell(c) for c in header) + " |"] - out.append("| " + " | ".join(["---"] * cols) + " |") - for row in body: - cells = [_cell(c) for c in row] - if len(cells) < cols: - cells += [""] * (cols - len(cells)) - out.append("| " + " | ".join(cells[:cols]) + " |") - return "\n".join(out) - - -def _cells_to_markdown(cells: list[Any]) -> str: - rows: dict[int, dict[int, str]] = {} - for c in cells: - if not isinstance(c, dict): - continue - r = c.get("start_row_offset_idx", c.get("row", 0)) or 0 - col = c.get("start_col_offset_idx", c.get("col", 0)) or 0 - rows.setdefault(r, {})[col] = (c.get("text") or "").replace("|", "\\|").strip() - if not rows: - return "" - max_col = max((max(r.keys()) for r in rows.values()), default=0) - grid = [] - for r_idx in sorted(rows): - row = [rows[r_idx].get(c, "") for c in range(max_col + 1)] - grid.append(row) - return _grid_to_markdown(grid) - - -def _table_csv(raw: dict[str, Any]) -> str | None: - grid = raw.get("data") or raw.get("grid") - if not (isinstance(grid, list) and grid and isinstance(grid[0], list)): - return None - import csv - import io - - buf = io.StringIO() - writer = csv.writer(buf) - for row in grid: - writer.writerow([ - (c.get("text") if isinstance(c, dict) else c) or "" for c in row - ]) - return buf.getvalue() + return _engine_extract(pdf_path, docling_ocr_enabled=settings.docling_ocr_enabled) diff --git a/app/ingestion/ocr.py b/app/ingestion/ocr.py index 821ed51..6b32e3a 100644 --- a/app/ingestion/ocr.py +++ b/app/ingestion/ocr.py @@ -1,87 +1,22 @@ -"""OCRmyPDF integration with Tesseract. - -We treat OCR as best-effort: if the input PDF already has a text layer (or OCR is -disabled by config), we skip OCR and use the original PDF. On failure, the -caller is expected to mark the document ``OCR_FAILED`` and continue without it. -""" +"""LegacyHUB adapter for the shared document recognition engine.""" from __future__ import annotations -from dataclasses import dataclass from pathlib import Path -import ocrmypdf +from teamhub_document_recognition_engine import OcrResult +from teamhub_document_recognition_engine import run_ocr as _engine_run_ocr from app.config import settings -from app.logging_config import get_logger -from app.utils.pdf import has_searchable_text - -logger = get_logger(__name__) - - -@dataclass -class OcrResult: - output_path: Path - skipped: bool - reason: str - languages: str def run_ocr(input_pdf: Path, output_pdf: Path, languages: str | None = None) -> OcrResult: - """Run OCRmyPDF. - - - If ``OCR_ENABLED`` is false: copy the input as the output and skip. - - If the input already has searchable text: skip OCR but still produce - ``output_pdf`` (a hard-link / copy to keep downstream code simple). - - On unexpected exceptions: re-raise (caller handles status update). - """ - langs = languages or settings.ocr_languages - - if not settings.ocr_enabled: - return _skip(input_pdf, output_pdf, langs, "ocr_disabled") - - if has_searchable_text(input_pdf): - return _skip(input_pdf, output_pdf, langs, "already_searchable") - - output_pdf.parent.mkdir(parents=True, exist_ok=True) - logger.info("ocr.start", input=str(input_pdf), output=str(output_pdf), languages=langs) - - try: - ocrmypdf.ocr( - input_file=str(input_pdf), - output_file=str(output_pdf), - language=langs, - skip_text=False, - redo_ocr=False, - force_ocr=False, - deskew=settings.ocr_deskew, - clean=settings.ocr_clean, - optimize=settings.ocr_optimize, - progress_bar=False, - jobs=1, - output_type="pdf", - # tolerate already-OCR pages where present - skip_big=200.0, - ) - except ocrmypdf.exceptions.PriorOcrFoundError: - logger.info("ocr.skip.prior_ocr", input=str(input_pdf)) - return _skip(input_pdf, output_pdf, langs, "prior_ocr_found") - except ocrmypdf.exceptions.DigitalSignatureError: - logger.warning("ocr.skip.signed_pdf", input=str(input_pdf)) - return _skip(input_pdf, output_pdf, langs, "digitally_signed") - except ocrmypdf.exceptions.EncryptedPdfError as exc: - logger.warning("ocr.encrypted", input=str(input_pdf), error=str(exc)) - raise - except ocrmypdf.exceptions.MissingDependencyError as exc: - logger.error("ocr.missing_dependency", error=str(exc)) - raise - - logger.info("ocr.done", output=str(output_pdf)) - return OcrResult(output_path=output_pdf, skipped=False, reason="ocr_completed", languages=langs) - - -def _skip(input_pdf: Path, output_pdf: Path, langs: str, reason: str) -> OcrResult: - output_pdf.parent.mkdir(parents=True, exist_ok=True) - if not output_pdf.exists() or output_pdf.resolve() != input_pdf.resolve(): - output_pdf.write_bytes(input_pdf.read_bytes()) - return OcrResult(output_path=output_pdf, skipped=True, reason=reason, languages=langs) + return _engine_run_ocr( + input_pdf, + output_pdf, + languages=languages or settings.ocr_languages, + enabled=settings.ocr_enabled, + deskew=settings.ocr_deskew, + clean=settings.ocr_clean, + optimize=settings.ocr_optimize, + ) diff --git a/docs/ADR-shared-core.md b/docs/ADR-shared-core.md index 9506bca..73a032b 100644 --- a/docs/ADR-shared-core.md +++ b/docs/ADR-shared-core.md @@ -1,6 +1,6 @@ # ADR: Extract OCR / Markdown / Search into a shared `teamhub_core` engine -Status: accepted (staged) · 2026-06-15 · scope: `legacy-knowledge-indexer` + `TeamHUB/Engines` +Status: accepted (staged, recognition engine extracted) · 2026-06-15 · scope: `legacy-knowledge-indexer` + `TeamHUB/Engines` ## Контекст @@ -42,12 +42,16 @@ versioned dependency / git submodule (D10). Извлечение делаетс - **`obs`-ядро готово.** `app/common/json_logger.py` написан без app-specific импортов и переносится в ядро как есть. - **Плейтменеджмент секретов вынесен из дефолтов** (см. `_enforce_secret_policy`). +- **Первый физический engine вынесен.** OCRmyPDF/Tesseract + Docling recognition + теперь живут в `TeamHUB/Engines/engines/teamhub-document-recognition-engine`; + LegacyHUB оставляет `app/ingestion/ocr.py` и + `app/ingestion/docling_extractor.py` как тонкие адаптеры и source metadata в + `app/ingestion/ENGINE_SOURCE.md`. ## Остаточные шаги (staged, координируются с QMS/MailHUB) -1. Создать `Engines/engines/teamhub_core/` с README (public API/CLI, inputs/outputs, - profile/config model, команда проверки, release notes) и версионным тегом - (`teamhub-core-v0.1.0`). +1. Поддерживать LegacyHUB pin на `teamhub-document-recognition-engine-v0.1.0` + до следующего совместимого release. 2. Перенести `contracts` (типы поиска/цитирования) из `app/api/schemas` в `teamhub_core.contracts`; `app/api/schemas` ре-экспортирует их для совместимости. 3. Перенести `search`/`ingest`/`storage`/`obs` в ядро; `run_search` принимает diff --git a/pyproject.toml b/pyproject.toml index 1165cf6..0e25224 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ # Ingestion - pin Docling tight since its DocumentConverter API # still moves between minor releases; lift the upper bound only # after a smoke test on a staging corpus. + "teamhub-document-recognition-engine @ git+https://github.com/GoodOneFather/TeamHUB_Engines.git@teamhub-document-recognition-engine-v0.1.0#subdirectory=engines/teamhub-document-recognition-engine", "ocrmypdf>=16.4.0,<17", "pikepdf>=9.0.0,<10", "pypdf>=4.3.0,<6", @@ -84,6 +85,9 @@ legacyhub-smoke = "scripts.smoke_test:main" [tool.hatch.build.targets.wheel] packages = ["app", "scripts"] +[tool.hatch.metadata] +allow-direct-references = true + [tool.ruff] line-length = 100 target-version = "py311" diff --git a/tests/test_document_recognition_engine_adapter.py b/tests/test_document_recognition_engine_adapter.py new file mode 100644 index 0000000..8e10752 --- /dev/null +++ b/tests/test_document_recognition_engine_adapter.py @@ -0,0 +1,63 @@ +from pathlib import Path + + +def test_legacy_ocr_adapter_reexports_engine_model_and_settings(monkeypatch, tmp_path: Path) -> None: + from teamhub_document_recognition_engine import OcrResult as EngineOcrResult + + import app.ingestion.ocr as legacy_ocr + + calls = {} + + def fake_run_ocr(input_pdf, output_pdf, **kwargs): + calls.update(kwargs) + return EngineOcrResult(output_path=output_pdf, skipped=True, reason="fake", languages=kwargs["languages"]) + + monkeypatch.setattr(legacy_ocr, "_engine_run_ocr", fake_run_ocr) + monkeypatch.setattr(legacy_ocr.settings, "ocr_enabled", False) + monkeypatch.setattr(legacy_ocr.settings, "ocr_deskew", False) + monkeypatch.setattr(legacy_ocr.settings, "ocr_clean", False) + monkeypatch.setattr(legacy_ocr.settings, "ocr_optimize", 0) + + result = legacy_ocr.run_ocr(tmp_path / "in.pdf", tmp_path / "out.pdf", languages="eng") + + assert legacy_ocr.OcrResult is EngineOcrResult + assert isinstance(result, EngineOcrResult) + assert calls == { + "languages": "eng", + "enabled": False, + "deskew": False, + "clean": False, + "optimize": 0, + } + + +def test_legacy_docling_adapter_reexports_engine_models_and_settings(monkeypatch, tmp_path: Path) -> None: + from teamhub_document_recognition_engine import ( + ExtractedBlock as EngineExtractedBlock, + ExtractionResult as EngineExtractionResult, + ) + + import app.ingestion.docling_extractor as legacy_docling + + calls = {} + + def fake_extract(pdf_path, **kwargs): + calls.update(kwargs) + return EngineExtractionResult( + markdown="", + json_payload={}, + blocks=[], + tables=[], + figures=[], + pages=[], + ) + + monkeypatch.setattr(legacy_docling, "_engine_extract", fake_extract) + monkeypatch.setattr(legacy_docling.settings, "docling_ocr_enabled", True) + + result = legacy_docling.extract(tmp_path / "ocr.pdf") + + assert legacy_docling.ExtractedBlock is EngineExtractedBlock + assert legacy_docling.ExtractionResult is EngineExtractionResult + assert isinstance(result, EngineExtractionResult) + assert calls == {"docling_ocr_enabled": True}