"""Marker/OCR integration for document processing.""" from __future__ import annotations import json import asyncio import io import re from pathlib import Path from typing import Any import httpx from app.config import get_settings import app.db.database as _database_mod class MarkerProcessor: """Handles OCR via Marker API with polygon output.""" def __init__(self): self.settings = get_settings() async def process_pdf(self, file_data: bytes, filename: str) -> dict[str, Any]: """Send PDF to Marker API for OCR with polygon extraction.""" async with httpx.AsyncClient(timeout=300) as client: resp = await client.post( f"{self.settings.marker_api_url}/convert", files={"file": (filename, file_data, "application/pdf")}, ) resp.raise_for_status() return resp.json() async def process_pdf_url(self, pdf_url: str, filename: str) -> dict[str, Any]: """Process PDF from URL via Marker API.""" async with httpx.AsyncClient(timeout=300) as client: resp = await client.post( f"{self.settings.marker_api_url}/convert", json={"pdf_url": pdf_url}, ) resp.raise_for_status() return resp.json() async def parse_marker_json(self, marker_output: dict, doc_id: str) -> list[dict]: """Parse Marker JSON output into vectorized chunks with polygon data.""" pages = marker_output.get("pages", []) chunks = [] for page in pages: page_num = page.get("meta", {}).get("page_num", page.get("page", 0)) text_lines = page.get("text_lines", []) for block_idx, tl in enumerate(text_lines): content = tl.get("text", "") polygon = tl.get("bbox") or tl.get("polygon") block_type = tl.get("type", "text") if not content or not isinstance(content, str) or not content.strip(): continue chunks.append({ "content": content.strip(), "page_num": page_num, "block_index": block_idx, "polygon": polygon, "chunk_type": block_type, }) if chunks: await _database_mod.db.batch_chunk(doc_id, chunks) return chunks async def process_document_file( self, file_data: bytes, filename: str, doc_id: str ) -> dict[str, Any]: """Process document file and return structured result.""" ext = Path(filename).suffix.lower() if ext == ".pdf": result = await self.process_pdf(file_data, filename) pages_data = result.get("pages", []) chunks = await self.parse_marker_json({"pages": pages_data}, doc_id) return { "success": result.get("success", True), "doc_id": doc_id, "page_count": result.get("page_count", len(pages_data)), "chunks": len(chunks), "ocr_model": result.get("ocr_model", "deepseek-ocr"), } elif ext in (".txt", ".pdf", ".md"): text = file_data.decode("utf-8", errors="replace") paragraphs = re.split(r'\n\s*\n', text) chunks = [] for i, para in enumerate(paragraphs): if len(para.strip()) > 20: chunks.append({ "content": para.strip(), "page_num": 0, "block_index": i, "polygon": None, "chunk_type": "text", }) if chunks: await _database_mod.db.batch_chunk(doc_id, chunks) return {"success": True, "doc_id": doc_id, "chunks": len(chunks), "page_count": 1} return {"success": False, "error": f"Unsupported file type: {ext}"}