feat: batch large documents with page_range and progress reporting
Use Docling's native page_range parameter to split large PDFs into sequential batches, preventing memory exhaustion and timeouts. Progress is reported via existing polling mechanism. Closes #56
This commit is contained in:
Pier-Jean Malandrino
parent
18f3b91af4
commit
fc866ce229
16 changed files with 521 additions and 13 deletions
|
|
@ -39,6 +39,8 @@ def _to_response(job) -> AnalysisResponse:
|
|||
chunks_json=job.chunks_json,
|
||||
has_document_json=job.document_json is not None,
|
||||
error_message=job.error_message,
|
||||
progress_current=job.progress_current,
|
||||
progress_total=job.progress_total,
|
||||
started_at=str(job.started_at) if job.started_at else None,
|
||||
completed_at=str(job.completed_at) if job.completed_at else None,
|
||||
created_at=str(job.created_at),
|
||||
|
|
|
|||
|
|
@ -47,6 +47,8 @@ class AnalysisResponse(_CamelModel):
|
|||
chunks_json: str | None = None
|
||||
has_document_json: bool = False
|
||||
error_message: str | None = None
|
||||
progress_current: int | None = None
|
||||
progress_total: int | None = None
|
||||
started_at: str | datetime | None = None
|
||||
completed_at: str | datetime | None = None
|
||||
created_at: str | datetime
|
||||
|
|
|
|||
|
|
@ -45,6 +45,8 @@ class AnalysisJob:
|
|||
document_json: str | None = None
|
||||
chunks_json: str | None = None
|
||||
error_message: str | None = None
|
||||
progress_current: int | None = None
|
||||
progress_total: int | None = None
|
||||
started_at: datetime | None = None
|
||||
completed_at: datetime | None = None
|
||||
created_at: datetime = field(default_factory=_utcnow)
|
||||
|
|
@ -74,6 +76,11 @@ class AnalysisJob:
|
|||
self.chunks_json = chunks_json
|
||||
self.completed_at = _utcnow()
|
||||
|
||||
def update_progress(self, current: int, total: int) -> None:
|
||||
"""Update batch progress counters."""
|
||||
self.progress_current = current
|
||||
self.progress_total = total
|
||||
|
||||
def mark_failed(self, error: str) -> None:
|
||||
"""Transition to FAILED with an error message."""
|
||||
self.status = AnalysisStatus.FAILED
|
||||
|
|
|
|||
|
|
@ -28,6 +28,8 @@ class DocumentConverter(Protocol):
|
|||
self,
|
||||
file_path: str,
|
||||
options: ConversionOptions,
|
||||
*,
|
||||
page_range: tuple[int, int] | None = None,
|
||||
) -> ConversionResult: ...
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -216,7 +216,12 @@ def _process_content_item(
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResult:
|
||||
def _convert_sync(
|
||||
file_path: str,
|
||||
options: ConversionOptions,
|
||||
*,
|
||||
page_range: tuple[int, int] | None = None,
|
||||
) -> ConversionResult:
|
||||
acquired = _converter_lock.acquire(timeout=settings.lock_timeout)
|
||||
if not acquired:
|
||||
raise TimeoutError(
|
||||
|
|
@ -230,6 +235,8 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul
|
|||
kwargs["max_num_pages"] = settings.max_page_count
|
||||
if settings.max_file_size > 0:
|
||||
kwargs["max_file_size"] = settings.max_file_size
|
||||
if page_range is not None:
|
||||
kwargs["page_range"] = page_range
|
||||
result = conv.convert(file_path, **kwargs)
|
||||
finally:
|
||||
_converter_lock.release()
|
||||
|
|
@ -275,5 +282,7 @@ class LocalConverter:
|
|||
self,
|
||||
file_path: str,
|
||||
options: ConversionOptions,
|
||||
*,
|
||||
page_range: tuple[int, int] | None = None,
|
||||
) -> ConversionResult:
|
||||
return await asyncio.to_thread(_convert_sync, file_path, options)
|
||||
return await asyncio.to_thread(_convert_sync, file_path, options, page_range=page_range)
|
||||
|
|
|
|||
|
|
@ -76,12 +76,14 @@ class ServeConverter:
|
|||
self,
|
||||
file_path: str,
|
||||
options: ConversionOptions,
|
||||
*,
|
||||
page_range: tuple[int, int] | None = None,
|
||||
) -> ConversionResult:
|
||||
"""Convert a document by uploading it to Docling Serve."""
|
||||
path = Path(file_path)
|
||||
content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
|
||||
|
||||
form_data = _build_form_data(options)
|
||||
form_data = _build_form_data(options, page_range=page_range)
|
||||
url = f"{self._base_url}{_API_PREFIX}/convert/file"
|
||||
|
||||
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
||||
|
|
@ -112,13 +114,17 @@ class ServeConverter:
|
|||
return False
|
||||
|
||||
|
||||
def _build_form_data(options: ConversionOptions) -> dict[str, str | list[str]]:
|
||||
def _build_form_data(
|
||||
options: ConversionOptions,
|
||||
*,
|
||||
page_range: tuple[int, int] | None = None,
|
||||
) -> dict[str, str | list[str]]:
|
||||
"""Build form fields matching Docling Serve's multipart form contract.
|
||||
|
||||
Array fields (to_formats) are sent as lists — httpx encodes them as
|
||||
repeated form keys (to_formats=md&to_formats=html&to_formats=json).
|
||||
"""
|
||||
return {
|
||||
data: dict[str, str | list[str]] = {
|
||||
"to_formats": ["md", "html", "json"],
|
||||
"do_ocr": str(options.do_ocr).lower(),
|
||||
"do_table_structure": str(options.do_table_structure).lower(),
|
||||
|
|
@ -131,6 +137,9 @@ def _build_form_data(options: ConversionOptions) -> dict[str, str | list[str]]:
|
|||
"generate_page_images": str(options.generate_page_images).lower(),
|
||||
"images_scale": str(options.images_scale),
|
||||
}
|
||||
if page_range is not None:
|
||||
data["page_range"] = f"{page_range[0]}-{page_range[1]}"
|
||||
return data
|
||||
|
||||
|
||||
def _parse_response(data: dict) -> ConversionResult:
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ class Settings:
|
|||
default_table_mode: str = "accurate" # "accurate" or "fast"
|
||||
max_page_count: int = 0 # 0 = unlimited (upload validation)
|
||||
max_file_size: int = 0 # 0 = unlimited (Docling-level, bytes)
|
||||
batch_page_size: int = 0 # 0 = disabled, > 0 = pages per batch
|
||||
upload_dir: str = "./uploads"
|
||||
db_path: str = "./data/docling_studio.db"
|
||||
cors_origins: list[str] = field(
|
||||
|
|
@ -42,6 +43,8 @@ class Settings:
|
|||
errors.append(f"max_page_count must be >= 0 (got {self.max_page_count})")
|
||||
if self.max_file_size < 0:
|
||||
errors.append(f"max_file_size must be >= 0 (got {self.max_file_size})")
|
||||
if self.batch_page_size < 0:
|
||||
errors.append(f"batch_page_size must be >= 0 (got {self.batch_page_size})")
|
||||
if self.default_table_mode not in ("accurate", "fast"):
|
||||
errors.append(
|
||||
f"default_table_mode must be 'accurate' or 'fast' (got '{self.default_table_mode}')"
|
||||
|
|
@ -78,6 +81,7 @@ class Settings:
|
|||
default_table_mode=os.environ.get("DEFAULT_TABLE_MODE", "accurate"),
|
||||
max_page_count=int(os.environ.get("MAX_PAGE_COUNT", "0")),
|
||||
max_file_size=int(os.environ.get("MAX_FILE_SIZE", "0")),
|
||||
batch_page_size=int(os.environ.get("BATCH_PAGE_SIZE", "0")),
|
||||
upload_dir=os.environ.get("UPLOAD_DIR", "./uploads"),
|
||||
db_path=os.environ.get("DB_PATH", "./data/docling_studio.db"),
|
||||
cors_origins=[o.strip() for o in cors_raw.split(",")],
|
||||
|
|
|
|||
|
|
@ -27,6 +27,8 @@ def _row_to_job(row) -> AnalysisJob:
|
|||
document_json=row["document_json"] if "document_json" in keys else None,
|
||||
chunks_json=row["chunks_json"] if "chunks_json" in keys else None,
|
||||
error_message=row["error_message"],
|
||||
progress_current=row["progress_current"] if "progress_current" in keys else None,
|
||||
progress_total=row["progress_total"] if "progress_total" in keys else None,
|
||||
started_at=_parse_dt(row["started_at"]),
|
||||
completed_at=_parse_dt(row["completed_at"]),
|
||||
created_at=_parse_dt(row["created_at"]) or datetime.now(UTC),
|
||||
|
|
@ -78,7 +80,8 @@ async def update_status(job: AnalysisJob) -> None:
|
|||
"""UPDATE analysis_jobs
|
||||
SET status = ?, content_markdown = ?, content_html = ?,
|
||||
pages_json = ?, document_json = ?, chunks_json = ?,
|
||||
error_message = ?, started_at = ?, completed_at = ?
|
||||
error_message = ?, progress_current = ?, progress_total = ?,
|
||||
started_at = ?, completed_at = ?
|
||||
WHERE id = ?""",
|
||||
(
|
||||
job.status.value,
|
||||
|
|
@ -88,6 +91,8 @@ async def update_status(job: AnalysisJob) -> None:
|
|||
job.document_json,
|
||||
job.chunks_json,
|
||||
job.error_message,
|
||||
job.progress_current,
|
||||
job.progress_total,
|
||||
str(job.started_at) if job.started_at else None,
|
||||
str(job.completed_at) if job.completed_at else None,
|
||||
job.id,
|
||||
|
|
@ -96,6 +101,16 @@ async def update_status(job: AnalysisJob) -> None:
|
|||
await db.commit()
|
||||
|
||||
|
||||
async def update_progress(job_id: str, current: int, total: int) -> None:
|
||||
"""Update only the progress columns for a running analysis."""
|
||||
async with get_connection() as db:
|
||||
await db.execute(
|
||||
"UPDATE analysis_jobs SET progress_current = ?, progress_total = ? WHERE id = ?",
|
||||
(current, total, job_id),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
|
||||
async def update_chunks(job_id: str, chunks_json: str) -> bool:
|
||||
"""Update only the chunks_json column for a completed analysis."""
|
||||
async with get_connection() as db:
|
||||
|
|
|
|||
|
|
@ -50,6 +50,8 @@ CREATE INDEX IF NOT EXISTS idx_documents_created_at ON documents(created_at);
|
|||
_MIGRATIONS = [
|
||||
("document_json", "ALTER TABLE analysis_jobs ADD COLUMN document_json TEXT"),
|
||||
("chunks_json", "ALTER TABLE analysis_jobs ADD COLUMN chunks_json TEXT"),
|
||||
("progress_current", "ALTER TABLE analysis_jobs ADD COLUMN progress_current INTEGER"),
|
||||
("progress_total", "ALTER TABLE analysis_jobs ADD COLUMN progress_total INTEGER"),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,11 +10,21 @@ import asyncio
|
|||
import functools
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
from dataclasses import asdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
from domain.models import AnalysisJob, AnalysisStatus
|
||||
from domain.value_objects import ChunkingOptions, ChunkResult, ConversionOptions, ConversionResult
|
||||
from domain.value_objects import (
|
||||
ChunkingOptions,
|
||||
ChunkResult,
|
||||
ConversionOptions,
|
||||
ConversionResult,
|
||||
PageDetail,
|
||||
)
|
||||
from infra.settings import settings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -38,6 +48,67 @@ def _chunk_to_dict(c: ChunkResult) -> dict:
|
|||
# Maximum number of concurrent analysis jobs to prevent resource exhaustion.
|
||||
_DEFAULT_MAX_CONCURRENT = 3
|
||||
|
||||
# Regex to extract <body> content from Docling's well-formed HTML output.
|
||||
_BODY_RE = re.compile(r"<body[^>]*>(.*)</body>", re.DOTALL | re.IGNORECASE)
|
||||
|
||||
|
||||
def _count_pdf_pages(file_path: str) -> int:
|
||||
"""Count pages in a PDF. Returns 0 if the file is not a valid PDF."""
|
||||
try:
|
||||
pdf = pdfium.PdfDocument(file_path)
|
||||
count = len(pdf)
|
||||
pdf.close()
|
||||
return count
|
||||
except Exception:
|
||||
logger.debug("Cannot open %s as PDF, batching disabled", file_path)
|
||||
return 0
|
||||
|
||||
|
||||
def _extract_html_body(html: str) -> str:
|
||||
"""Extract content between <body> tags.
|
||||
|
||||
Docling produces well-formed HTML — regex is safe for this controlled output.
|
||||
Returns raw html as fallback if no <body> tag is found.
|
||||
"""
|
||||
match = _BODY_RE.search(html)
|
||||
return match.group(1).strip() if match else html
|
||||
|
||||
|
||||
def _merge_results(results: list[ConversionResult]) -> ConversionResult:
|
||||
"""Merge multiple batch ConversionResults into a single consolidated result.
|
||||
|
||||
document_json is intentionally set to None: merging DoclingDocument's internal
|
||||
tree structure across batches is error-prone. Re-chunking is disabled for
|
||||
batched conversions (robustness decision for 0.3.1).
|
||||
"""
|
||||
if not results:
|
||||
return ConversionResult(page_count=0, content_markdown="", content_html="", pages=[])
|
||||
|
||||
all_pages: list[PageDetail] = []
|
||||
all_md: list[str] = []
|
||||
html_bodies: list[str] = []
|
||||
total_skipped = 0
|
||||
|
||||
for r in results:
|
||||
all_pages.extend(r.pages)
|
||||
all_md.append(r.content_markdown)
|
||||
html_bodies.append(_extract_html_body(r.content_html))
|
||||
total_skipped += r.skipped_items
|
||||
|
||||
merged_body = "\n".join(html_bodies)
|
||||
merged_html = (
|
||||
f'<!DOCTYPE html><html><head><meta charset="utf-8"></head><body>{merged_body}</body></html>'
|
||||
)
|
||||
|
||||
return ConversionResult(
|
||||
page_count=sum(r.page_count for r in results),
|
||||
content_markdown="\n\n".join(all_md),
|
||||
content_html=merged_html,
|
||||
pages=all_pages,
|
||||
skipped_items=total_skipped,
|
||||
document_json=None,
|
||||
)
|
||||
|
||||
|
||||
class AnalysisService:
|
||||
"""Orchestrates document analysis using an injected converter."""
|
||||
|
|
@ -121,6 +192,66 @@ class AnalysisService:
|
|||
|
||||
return chunks
|
||||
|
||||
async def _run_batched_conversion(
|
||||
self,
|
||||
job_id: str,
|
||||
file_path: str,
|
||||
options: ConversionOptions,
|
||||
total_pages: int,
|
||||
batch_size: int,
|
||||
) -> ConversionResult | None:
|
||||
"""Convert a document in batches using page_range.
|
||||
|
||||
Returns None if the job was deleted mid-batch (caller should abort).
|
||||
Raises on batch failure (fail-fast: entire job fails).
|
||||
"""
|
||||
num_batches = math.ceil(total_pages / batch_size)
|
||||
await analysis_repo.update_progress(job_id, 0, total_pages)
|
||||
logger.info(
|
||||
"Batched conversion: %d pages in %d batches of %d for job %s",
|
||||
total_pages,
|
||||
num_batches,
|
||||
batch_size,
|
||||
job_id,
|
||||
)
|
||||
|
||||
results: list[ConversionResult] = []
|
||||
for batch_idx in range(num_batches):
|
||||
start = batch_idx * batch_size + 1
|
||||
end = min(start + batch_size - 1, total_pages)
|
||||
|
||||
if not await analysis_repo.find_by_id(job_id):
|
||||
logger.info(
|
||||
"Job %s deleted during batch %d/%d, aborting",
|
||||
job_id,
|
||||
batch_idx + 1,
|
||||
num_batches,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
batch_result = await asyncio.wait_for(
|
||||
self._converter.convert(file_path, options, page_range=(start, end)),
|
||||
timeout=self._conversion_timeout,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
f"Batch {batch_idx + 1}/{num_batches} (pages {start}-{end}) failed: {exc}"
|
||||
) from exc
|
||||
|
||||
results.append(batch_result)
|
||||
await analysis_repo.update_progress(job_id, end, total_pages)
|
||||
logger.info(
|
||||
"Batch %d/%d done (pages %d-%d) for job %s",
|
||||
batch_idx + 1,
|
||||
num_batches,
|
||||
start,
|
||||
end,
|
||||
job_id,
|
||||
)
|
||||
|
||||
return _merge_results(results)
|
||||
|
||||
def _on_task_done(self, task: asyncio.Task, *, job_id: str) -> None:
|
||||
"""Cleanup running tasks and delegate to module-level handler."""
|
||||
self._running_tasks.pop(job_id, None)
|
||||
|
|
@ -168,10 +299,24 @@ class AnalysisService:
|
|||
opts_dict = {**opts_dict, "table_mode": settings.default_table_mode}
|
||||
options = ConversionOptions(**opts_dict)
|
||||
|
||||
result: ConversionResult = await asyncio.wait_for(
|
||||
self._converter.convert(file_path, options),
|
||||
timeout=self._conversion_timeout,
|
||||
)
|
||||
total_pages = _count_pdf_pages(file_path)
|
||||
batch_size = settings.batch_page_size
|
||||
|
||||
if batch_size > 0 and total_pages > batch_size:
|
||||
result = await self._run_batched_conversion(
|
||||
job_id,
|
||||
file_path,
|
||||
options,
|
||||
total_pages,
|
||||
batch_size,
|
||||
)
|
||||
if result is None:
|
||||
return # job was deleted mid-batch
|
||||
else:
|
||||
result = await asyncio.wait_for(
|
||||
self._converter.convert(file_path, options),
|
||||
timeout=self._conversion_timeout,
|
||||
)
|
||||
|
||||
pages_json = json.dumps([asdict(p) for p in result.pages])
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Tests for AnalysisService — callbacks, concurrency, and orchestration."""
|
||||
"""Tests for AnalysisService — callbacks, concurrency, orchestration, and batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -8,7 +8,14 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from services.analysis_service import AnalysisService, _on_task_done
|
||||
from domain.value_objects import ConversionResult, PageDetail
|
||||
from services.analysis_service import (
|
||||
AnalysisService,
|
||||
_count_pdf_pages,
|
||||
_extract_html_body,
|
||||
_merge_results,
|
||||
_on_task_done,
|
||||
)
|
||||
|
||||
|
||||
class TestOnTaskDone:
|
||||
|
|
@ -186,3 +193,224 @@ class TestAnalysisServiceConcurrency:
|
|||
# Both should have completed
|
||||
assert call_order.count("start") == 2
|
||||
assert call_order.count("end") == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch helper tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCountPdfPages:
|
||||
def test_valid_pdf(self, tmp_path):
|
||||
"""A real (minimal) PDF should return its page count."""
|
||||
# Create a minimal valid 1-page PDF using pypdfium2
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
pdf = pdfium.PdfDocument.new()
|
||||
pdf.new_page(612, 792)
|
||||
path = tmp_path / "test.pdf"
|
||||
pdf.save(str(path))
|
||||
pdf.close()
|
||||
|
||||
assert _count_pdf_pages(str(path)) == 1
|
||||
|
||||
def test_non_pdf_file(self, tmp_path):
|
||||
"""A non-PDF file should return 0."""
|
||||
path = tmp_path / "test.docx"
|
||||
path.write_bytes(b"PK\x03\x04 not a pdf")
|
||||
assert _count_pdf_pages(str(path)) == 0
|
||||
|
||||
def test_nonexistent_file(self):
|
||||
"""A nonexistent file should return 0."""
|
||||
assert _count_pdf_pages("/no/such/file.pdf") == 0
|
||||
|
||||
def test_empty_file(self, tmp_path):
|
||||
"""An empty file should return 0."""
|
||||
path = tmp_path / "empty.pdf"
|
||||
path.write_bytes(b"")
|
||||
assert _count_pdf_pages(str(path)) == 0
|
||||
|
||||
|
||||
class TestExtractHtmlBody:
|
||||
def test_extracts_body(self):
|
||||
html = '<html><head></head><body class="x"><p>Hello</p></body></html>'
|
||||
assert _extract_html_body(html) == "<p>Hello</p>"
|
||||
|
||||
def test_no_body_tag_returns_raw(self):
|
||||
html = "<p>No body tag</p>"
|
||||
assert _extract_html_body(html) == html
|
||||
|
||||
def test_empty_body(self):
|
||||
html = "<html><body></body></html>"
|
||||
assert _extract_html_body(html) == ""
|
||||
|
||||
|
||||
class TestMergeResults:
|
||||
def test_empty_list(self):
|
||||
result = _merge_results([])
|
||||
assert result.page_count == 0
|
||||
assert result.content_markdown == ""
|
||||
assert result.pages == []
|
||||
assert result.document_json is None
|
||||
|
||||
def test_single_result_passthrough(self):
|
||||
r = ConversionResult(
|
||||
page_count=3,
|
||||
content_markdown="# Page 1",
|
||||
content_html="<html><body><p>Page 1</p></body></html>",
|
||||
pages=[PageDetail(page_number=1, width=612, height=792)],
|
||||
document_json='{"pages": {}}',
|
||||
)
|
||||
merged = _merge_results([r])
|
||||
assert merged.page_count == 3
|
||||
assert merged.content_markdown == "# Page 1"
|
||||
assert merged.pages == [PageDetail(page_number=1, width=612, height=792)]
|
||||
assert merged.document_json is None # intentionally dropped
|
||||
|
||||
def test_merges_multiple_results(self):
|
||||
r1 = ConversionResult(
|
||||
page_count=2,
|
||||
content_markdown="# Batch 1",
|
||||
content_html="<html><body><p>B1</p></body></html>",
|
||||
pages=[
|
||||
PageDetail(page_number=1, width=612, height=792),
|
||||
PageDetail(page_number=2, width=612, height=792),
|
||||
],
|
||||
skipped_items=1,
|
||||
)
|
||||
r2 = ConversionResult(
|
||||
page_count=2,
|
||||
content_markdown="# Batch 2",
|
||||
content_html="<html><body><p>B2</p></body></html>",
|
||||
pages=[
|
||||
PageDetail(page_number=3, width=612, height=792),
|
||||
PageDetail(page_number=4, width=612, height=792),
|
||||
],
|
||||
skipped_items=2,
|
||||
)
|
||||
merged = _merge_results([r1, r2])
|
||||
assert merged.page_count == 4
|
||||
assert merged.content_markdown == "# Batch 1\n\n# Batch 2"
|
||||
assert len(merged.pages) == 4
|
||||
assert merged.pages[0].page_number == 1
|
||||
assert merged.pages[3].page_number == 4
|
||||
assert merged.skipped_items == 3
|
||||
assert merged.document_json is None
|
||||
assert "<p>B1</p>" in merged.content_html
|
||||
assert "<p>B2</p>" in merged.content_html
|
||||
# Valid HTML structure
|
||||
assert merged.content_html.startswith("<!DOCTYPE html>")
|
||||
assert "</body></html>" in merged.content_html
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batched conversion integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBatchedConversion:
|
||||
@pytest.mark.asyncio
|
||||
async def test_batched_conversion_produces_merged_result(self):
|
||||
"""When batch_page_size is set and document exceeds it, results are merged."""
|
||||
converter = AsyncMock()
|
||||
|
||||
# Simulate 2 batches: pages 1-5 and 6-8
|
||||
converter.convert.side_effect = [
|
||||
ConversionResult(
|
||||
page_count=5,
|
||||
content_markdown="# Batch 1",
|
||||
content_html="<html><body><p>B1</p></body></html>",
|
||||
pages=[PageDetail(page_number=i, width=612, height=792) for i in range(1, 6)],
|
||||
),
|
||||
ConversionResult(
|
||||
page_count=3,
|
||||
content_markdown="# Batch 2",
|
||||
content_html="<html><body><p>B2</p></body></html>",
|
||||
pages=[PageDetail(page_number=i, width=612, height=792) for i in range(6, 9)],
|
||||
),
|
||||
]
|
||||
|
||||
service = AnalysisService(converter=converter, conversion_timeout=60)
|
||||
|
||||
with patch("services.analysis_service.analysis_repo") as mock_repo:
|
||||
mock_repo.find_by_id = AsyncMock(return_value=MagicMock())
|
||||
mock_repo.update_progress = AsyncMock()
|
||||
|
||||
result = await service._run_batched_conversion(
|
||||
"job-1",
|
||||
"/fake.pdf",
|
||||
MagicMock(),
|
||||
total_pages=8,
|
||||
batch_size=5,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.page_count == 8
|
||||
assert len(result.pages) == 8
|
||||
assert result.document_json is None
|
||||
assert converter.convert.call_count == 2
|
||||
|
||||
# Verify page_range was passed correctly
|
||||
call1_kwargs = converter.convert.call_args_list[0].kwargs
|
||||
call2_kwargs = converter.convert.call_args_list[1].kwargs
|
||||
assert call1_kwargs["page_range"] == (1, 5)
|
||||
assert call2_kwargs["page_range"] == (6, 8)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_failure_raises_with_enriched_message(self):
|
||||
"""If a batch fails, RuntimeError is raised with batch info."""
|
||||
converter = AsyncMock()
|
||||
converter.convert.side_effect = [
|
||||
ConversionResult(
|
||||
page_count=5,
|
||||
content_markdown="ok",
|
||||
content_html="<html><body>ok</body></html>",
|
||||
pages=[PageDetail(page_number=i, width=612, height=792) for i in range(1, 6)],
|
||||
),
|
||||
RuntimeError("OOM"),
|
||||
]
|
||||
|
||||
service = AnalysisService(converter=converter, conversion_timeout=60)
|
||||
|
||||
with patch("services.analysis_service.analysis_repo") as mock_repo:
|
||||
mock_repo.find_by_id = AsyncMock(return_value=MagicMock())
|
||||
mock_repo.update_progress = AsyncMock()
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"Batch 2/2 \(pages 6-8\) failed: OOM"):
|
||||
await service._run_batched_conversion(
|
||||
"job-fail",
|
||||
"/fake.pdf",
|
||||
MagicMock(),
|
||||
total_pages=8,
|
||||
batch_size=5,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_job_deleted_mid_batch_returns_none(self):
|
||||
"""If the job is deleted between batches, conversion aborts with None."""
|
||||
converter = AsyncMock()
|
||||
converter.convert.return_value = ConversionResult(
|
||||
page_count=5,
|
||||
content_markdown="ok",
|
||||
content_html="<html><body>ok</body></html>",
|
||||
pages=[PageDetail(page_number=i, width=612, height=792) for i in range(1, 6)],
|
||||
)
|
||||
|
||||
service = AnalysisService(converter=converter, conversion_timeout=60)
|
||||
|
||||
# First check returns job, second returns None (deleted)
|
||||
with patch("services.analysis_service.analysis_repo") as mock_repo:
|
||||
mock_repo.find_by_id = AsyncMock(side_effect=[MagicMock(), None])
|
||||
mock_repo.update_progress = AsyncMock()
|
||||
|
||||
result = await service._run_batched_conversion(
|
||||
"job-del",
|
||||
"/fake.pdf",
|
||||
MagicMock(),
|
||||
total_pages=10,
|
||||
batch_size=5,
|
||||
)
|
||||
|
||||
assert result is None
|
||||
# Only first batch should have been converted
|
||||
assert converter.convert.call_count == 1
|
||||
|
|
|
|||
|
|
@ -56,6 +56,14 @@ class TestBuildFormData:
|
|||
assert data["images_scale"] == "2.0"
|
||||
assert data["include_images"] == "true"
|
||||
|
||||
def test_page_range_included_when_set(self):
|
||||
data = _build_form_data(ConversionOptions(), page_range=(11, 20))
|
||||
assert data["page_range"] == "11-20"
|
||||
|
||||
def test_page_range_absent_when_none(self):
|
||||
data = _build_form_data(ConversionOptions())
|
||||
assert "page_range" not in data
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests — response parsing
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ class TestSettingsDefaults:
|
|||
assert s.document_timeout == 120.0
|
||||
assert s.lock_timeout == 300
|
||||
assert s.max_page_count == 0
|
||||
assert s.batch_page_size == 0
|
||||
assert s.upload_dir == "./uploads"
|
||||
assert s.db_path == "./data/docling_studio.db"
|
||||
assert "http://localhost:3000" in s.cors_origins
|
||||
|
|
@ -67,6 +68,20 @@ class TestSettingsValidation:
|
|||
with pytest.raises(ValueError, match="max_file_size must be >= 0"):
|
||||
Settings(max_file_size=-1)
|
||||
|
||||
def test_negative_batch_page_size_rejected(self):
|
||||
import pytest
|
||||
|
||||
with pytest.raises(ValueError, match="batch_page_size must be >= 0"):
|
||||
Settings(batch_page_size=-1)
|
||||
|
||||
def test_zero_batch_page_size_accepted(self):
|
||||
s = Settings(batch_page_size=0)
|
||||
assert s.batch_page_size == 0
|
||||
|
||||
def test_positive_batch_page_size_accepted(self):
|
||||
s = Settings(batch_page_size=10)
|
||||
assert s.batch_page_size == 10
|
||||
|
||||
def test_zero_lock_timeout_rejected(self):
|
||||
import pytest
|
||||
|
||||
|
|
@ -114,6 +129,7 @@ class TestSettingsFromEnv:
|
|||
monkeypatch.setenv("DOCUMENT_TIMEOUT", "60.0")
|
||||
monkeypatch.setenv("LOCK_TIMEOUT", "600")
|
||||
monkeypatch.setenv("MAX_PAGE_COUNT", "20")
|
||||
monkeypatch.setenv("BATCH_PAGE_SIZE", "15")
|
||||
monkeypatch.setenv("UPLOAD_DIR", "/data/uploads")
|
||||
monkeypatch.setenv("DB_PATH", "/data/test.db")
|
||||
monkeypatch.setenv("CORS_ORIGINS", "http://a.com, http://b.com")
|
||||
|
|
@ -129,6 +145,7 @@ class TestSettingsFromEnv:
|
|||
assert s.lock_timeout == 600
|
||||
assert s.document_timeout == 60.0
|
||||
assert s.max_page_count == 20
|
||||
assert s.batch_page_size == 15
|
||||
assert s.upload_dir == "/data/uploads"
|
||||
assert s.db_path == "/data/test.db"
|
||||
assert s.cors_origins == ["http://a.com", "http://b.com"]
|
||||
|
|
|
|||
|
|
@ -38,6 +38,18 @@
|
|||
<div class="status-row">
|
||||
<span class="status-badge" :class="statusClass">
|
||||
{{ analysisStore.currentAnalysis.status }}
|
||||
<template
|
||||
v-if="
|
||||
analysisStore.currentAnalysis.status === 'RUNNING' &&
|
||||
analysisStore.currentAnalysis.progressTotal &&
|
||||
analysisStore.currentAnalysis.progressTotal > 0
|
||||
"
|
||||
>
|
||||
({{ analysisStore.currentAnalysis.progressCurrent ?? 0 }}/{{
|
||||
analysisStore.currentAnalysis.progressTotal
|
||||
}}
|
||||
pages)
|
||||
</template>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -105,6 +105,18 @@
|
|||
<div v-else-if="store.currentAnalysis?.status === 'RUNNING'" class="result-placeholder">
|
||||
<div class="spinner-large" />
|
||||
<span>{{ t('studio.analysisRunning') }}</span>
|
||||
<div
|
||||
v-if="store.currentAnalysis.progressTotal && store.currentAnalysis.progressTotal > 0"
|
||||
class="batch-progress"
|
||||
>
|
||||
<div class="progress-bar">
|
||||
<div class="progress-fill" :style="{ width: progressPercent + '%' }" />
|
||||
</div>
|
||||
<span class="progress-text">
|
||||
Pages {{ store.currentAnalysis.progressCurrent ?? 0 }} /
|
||||
{{ store.currentAnalysis.progressTotal }}
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
<div v-else-if="store.currentAnalysis?.status === 'FAILED'" class="result-placeholder error">
|
||||
<svg viewBox="0 0 20 20" fill="currentColor" class="error-icon">
|
||||
|
|
@ -171,6 +183,12 @@ const tabs = computed(() => [
|
|||
|
||||
const totalPages = computed(() => store.currentPages.length)
|
||||
|
||||
const progressPercent = computed(() => {
|
||||
const a = store.currentAnalysis
|
||||
if (!a?.progressTotal || a.progressTotal <= 0) return 0
|
||||
return Math.min(100, Math.round(((a.progressCurrent ?? 0) / a.progressTotal) * 100))
|
||||
})
|
||||
|
||||
const currentPageData = computed(() => {
|
||||
return store.currentPages.find((p) => p.page_number === props.currentPage) || null
|
||||
})
|
||||
|
|
@ -505,4 +523,30 @@ async function copyElement(idx: number, content: string) {
|
|||
transform: rotate(360deg);
|
||||
}
|
||||
}
|
||||
|
||||
.batch-progress {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
margin-top: 8px;
|
||||
width: 200px;
|
||||
}
|
||||
.progress-bar {
|
||||
width: 100%;
|
||||
height: 6px;
|
||||
background: var(--border-light);
|
||||
border-radius: 3px;
|
||||
overflow: hidden;
|
||||
}
|
||||
.progress-fill {
|
||||
height: 100%;
|
||||
background: var(--accent);
|
||||
border-radius: 3px;
|
||||
transition: width 0.3s ease;
|
||||
}
|
||||
.progress-text {
|
||||
font-size: 0.8rem;
|
||||
color: var(--text-secondary);
|
||||
}
|
||||
</style>
|
||||
|
|
|
|||
|
|
@ -33,6 +33,8 @@ export interface Analysis {
|
|||
chunksJson: string | null
|
||||
hasDocumentJson: boolean
|
||||
errorMessage: string | null
|
||||
progressCurrent: number | null
|
||||
progressTotal: number | null
|
||||
startedAt: string | null
|
||||
completedAt: string | null
|
||||
createdAt: string
|
||||
|
|
|
|||
Loading…
Reference in a new issue