From 27a317080266480f4ab9ac96c5f00b8c5291c2a8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 19 May 2026 08:27:13 +0000 Subject: [PATCH] fix: configure pdf page limit --- apps/api/.env.example | 1 + apps/api/app/api/v1/routes/retrieval.py | 10 +- apps/worker/.env.example | 1 + .../orchestration/parse_session.py | 9 +- .../contract/test_parse_task_contract.py | 147 ++++++++++++++++++ .../shared/core/config/storage.py | 5 + 6 files changed, 165 insertions(+), 8 deletions(-) diff --git a/apps/api/.env.example b/apps/api/.env.example index 0139bfd4f..609cc7cb6 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -90,6 +90,7 @@ ARK_API_KEY= # File handling defaults SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md MAX_FILE_SIZE=104857600 +MAX_PDF_PAGE_LIMIT=600 # Required for specific features: webhooks and callbacks WEBHOOK_MASTER_KEY= diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 700049fed..f814ec3ef 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -17,6 +17,10 @@ router = APIRouter(tags=["Retrieval"]) +def _is_none(value: object) -> bool: + return value is None + + class ExcludeSection(BaseModel): document_id: str section_path: str @@ -83,9 +87,9 @@ class RetrievalQueryResponse(BaseModel): answer_text: str | None = None referenced_chunks: list[dict] = Field(default_factory=list) results: list[dict] = Field(default_factory=list) - evidence_text: str | None = None - stop_reason: str | None = None - failure_reason: str | None = None + evidence_text: str | None = Field(default=None, exclude_if=_is_none) + stop_reason: str | None = Field(default=None, exclude_if=_is_none) + failure_reason: str | None = Field(default=None, exclude_if=_is_none) @router.post("/query", response_model=RetrievalQueryResponse) diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 9b3ad9208..bf5f8251e 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -115,6 +115,7 @@ ILOVEAPI_SECRET_KEY= # File handling defaults SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md MAX_FILE_SIZE=104857600 +MAX_PDF_PAGE_LIMIT=600 # Legacy parser compatibility fields. # ALL_DF_COLS=content,path,type,length,keywords,summary,know_id,tokens,connectto,addtime,page_nums diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py index 3fe956f1a..7e1141a84 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -16,8 +16,6 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -PDF_PAGE_LIMIT = 600 - @dataclass(frozen=True) class ParseSession: @@ -111,16 +109,17 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: f"ℹ️ VLM rejected atlas for {parse_input.filename}, routing as generic" ) - if profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT: + pdf_page_limit = settings.MAX_PDF_PAGE_LIMIT + if profile.file_type == "pdf" and profile.page_count > pdf_page_limit: raise ValidationException( user_message=( - f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. " + f"Document too large: {profile.page_count} pages exceeds the {pdf_page_limit}-page limit. " "Please split the document and upload in smaller batches." ), violations=[ { "field": "page_count", - "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}", + "description": f"PDF has {profile.page_count} pages, limit is {pdf_page_limit}", } ], ) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 1ca733f56..7fd6f5684 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -23,6 +23,17 @@ _SAMPLE_PDF_PATH: Path = _FIXTURES_ROOT / "sample_3pages.pdf" +def _write_blank_pdf(file_path: Path, page_count: int) -> None: + from pypdf import PdfWriter + + writer = PdfWriter() + for _ in range(page_count): + writer.add_blank_page(width=72, height=72) + + with file_path.open("wb") as pdf_file: + writer.write(pdf_file) + + def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]: job_metadata: dict[str, Any] = { "namespace": "worker-contract", @@ -1480,3 +1491,139 @@ def fake_download_s3_file_to_temp( ("start_processing", "running"), ("mark_failed", "failed"), ] + + +def test_should_reject_pdf_when_page_count_exceeds_configured_limit( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("BILLING_ENABLED", "false") + ( + document_ingestion_tasks, + _parse_service, + parse_job_service, + engine, + sync_job_info_service_cls, + sync_job_metadata_service_cls, + sync_redis_service_factory, + ) = _load_parse_task_modules() + settings = _load_worker_settings() + + user_id: str = f"worker-user-{uuid4().hex[:12]}" + job_id: str = f"job_pdf_page_limit_{uuid4().hex[:12]}" + source_file_name: str = "oversized-contract.pdf" + s3_key: str = f"uploads/{job_id}.pdf" + max_pdf_page_limit: int = 1 + actual_page_count: int = 2 + pdf_path = tmp_path / source_file_name + _write_blank_pdf(pdf_path, actual_page_count) + + with engine.begin() as connection: + insert_contract_user(connection, user_id=user_id) + job_metadata = _build_pending_file_job_metadata(source_file_name) + insert_contract_job( + connection, + job_id=job_id, + user_id=user_id, + status="pending", + source_type="file", + s3_key=s3_key, + webhook_enabled=False, + job_metadata=job_metadata, + billing_status="pending", + ) + + redis_service = _save_worker_task_cache( + job_id=job_id, + user_id=user_id, + s3_key=s3_key, + metadata=job_metadata, + sync_job_info_service_cls=sync_job_info_service_cls, + sync_job_metadata_service_cls=sync_job_metadata_service_cls, + sync_redis_service_factory=sync_redis_service_factory, + ) + + _bind_parse_task_to_current_module( + monkeypatch, + document_ingestion_tasks=document_ingestion_tasks, + ) + monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) + monkeypatch.setattr(settings, "BILLING_ENABLED", False) + monkeypatch.setattr(settings, "MAX_PDF_PAGE_LIMIT", max_pdf_page_limit) + + def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: + return { + "exists": storage_key == s3_key, + "size": pdf_path.stat().st_size, + } + + _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) + + def fake_download_s3_file_to_temp( + storage_key: str, + file_ext: str, + temp_dir: str, + ) -> str: + assert storage_key == s3_key + assert file_ext == ".pdf" + downloaded_path = Path(temp_dir) / f"downloaded{file_ext}" + shutil.copy2(pdf_path, downloaded_path) + return str(downloaded_path) + + monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp) + monkeypatch.setattr( + parse_job_service, + "get_result_storage", + lambda: (_ for _ in ()).throw( + AssertionError("result storage should not run after page-limit rejection") + ), + ) + + result = document_ingestion_tasks.parse_task.apply( + args=[job_id, user_id, "document_ingestion"], + throw=False, + ) + + assert result.status == "FAILURE" + assert _find_task_workspaces(tmp_path, job_id) == [] + + metadata = sync_job_metadata_service_cls(redis_service).get_metadata(job_id) + assert metadata is not None + assert metadata["page_count"] == actual_page_count + assert metadata["billing_status"] == "skipped" + + with engine.begin() as connection: + job_row = ( + connection.execute( + text( + """ + SELECT status, billing_status, page_count, credits_charged, + error_code, error_message, job_metadata + FROM jobs + WHERE job_id = :job_id + """ + ), + {"job_id": job_id}, + ) + .mappings() + .one() + ) + + assert job_row["status"] == "failed" + assert job_row["billing_status"] == "skipped" + assert job_row["page_count"] == actual_page_count + assert job_row["credits_charged"] == 0 + assert job_row["error_code"] == "INVALID_ARGUMENT" + assert ( + job_row["error_message"] + == "Document too large: 2 pages exceeds the 1-page limit. Please split the document and upload in smaller batches." + ) + assert job_row["job_metadata"]["error_details"] == { + "violations": [ + { + "field": "page_count", + "description": "PDF has 2 pages, limit is 1", + } + ] + } diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index a7237481f..556879f7f 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -55,6 +55,11 @@ class StorageConfig(BaseModel): MAX_FILE_SIZE: int = Field( default=104857600, description="Maximum file size in bytes" ) + MAX_PDF_PAGE_LIMIT: int = Field( + default=600, + ge=1, + description="Maximum allowed PDF page count before parsing is rejected", + ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md", description="Supported file extensions",