From aea2dfabdaaf98a754aa1d35c95abc5c9a2fe553 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 24 Jul 2026 02:44:40 +0530 Subject: [PATCH 01/10] fix(collections): Use file batch ID from create() when listing failed files --- backend/app/crud/rag/open_ai.py | 19 +++-- backend/app/tests/crud/rag/test_open_ai.py | 72 ++++++++++++------- .../providers/test_openai_provider.py | 32 ++++----- backend/app/tests/utils/llm_provider.py | 7 +- backend/app/tests/utils/openai.py | 9 +-- docs/wiki/modules/knowledge-base.md | 1 + 6 files changed, 85 insertions(+), 55 deletions(-) diff --git a/backend/app/crud/rag/open_ai.py b/backend/app/crud/rag/open_ai.py index 0df73ab39..6d80e12c9 100644 --- a/backend/app/crud/rag/open_ai.py +++ b/backend/app/crud/rag/open_ai.py @@ -115,11 +115,16 @@ def update( ) try: - batch = self.client.vector_stores.file_batches.upload_and_poll( + created = self.client.vector_stores.file_batches.create( vector_store_id=vector_store_id, - files=[], file_ids=[doc.file_id[OPENAI_PROVIDER] for doc in docs], ) + # poll()'s return deserializes a vector-store body, so its .id is the vs_ id; + # capture the real vsfb_ batch id from create() before polling. + batch_id = created.id + batch = self.client.vector_stores.file_batches.poll( + batch_id, vector_store_id=vector_store_id + ) except openai.RateLimitError as e: error_message = ( f"[OPENAI] Rate limit exceeded (code: {e.status_code}): " @@ -215,14 +220,14 @@ def update( logger.info( f"[OpenAIVectorStoreCrud.update] Batch complete | " - f"{{'vector_store_id': '{vector_store_id}', " + f"{{'vector_store_id': '{vector_store_id}', 'batch_id': '{batch_id}', " f"'completed': {batch.file_counts.completed}, 'failed': {batch.file_counts.failed}}}" ) if batch.file_counts.failed > 0: try: failed_files = self.client.vector_stores.file_batches.list_files( vector_store_id=vector_store_id, - batch_id=batch.id, + batch_id=batch_id, filter="failed", ) doc_by_file_id = {d.file_id[OPENAI_PROVIDER]: d for d in docs} @@ -232,11 +237,15 @@ def update( label = d.fname if d else f.id msg = f.last_error.message if f.last_error else "no error detail" parts.append(f"{label}: {msg}") + logger.error( + f"[OpenAIVectorStoreCrud.update] Files failed to index | " + f"{{'batch_id': '{batch_id}', 'failed_files': '{', '.join(parts)}'}}" + ) raise RuntimeError("; ".join(parts)) except OpenAIError as err: logger.warning( f"[OpenAIVectorStoreCrud.update] Could not fetch per-file errors | " - f"{{'batch_id': '{batch.id}', 'error': '{str(err)}'}}" + f"{{'batch_id': '{batch_id}', 'error': '{str(err)}'}}" ) raise diff --git a/backend/app/tests/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index d8c85c227..fb5693331 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -37,10 +37,19 @@ def docs(): ] -def _batch_result(*, completed: int, failed: int, batch_id: str = "batch_abc"): - """Mock the return of vector_stores.file_batches.upload_and_poll.""" +REAL_BATCH_ID = "vsfb_real123" +CORRUPT_POLL_ID = "vs_corrupt999" + + +def _wire_batch(mock_client, *, completed: int, failed: int) -> None: + """Stub create() with the real vsfb_ id and poll() with the corrupt vs_ id the SDK returns.""" + mock_client.vector_stores.file_batches.create.return_value = MagicMock( + id=REAL_BATCH_ID + ) counts = MagicMock(completed=completed, failed=failed) - return MagicMock(id=batch_id, file_counts=counts) + mock_client.vector_stores.file_batches.poll.return_value = MagicMock( + id=CORRUPT_POLL_ID, file_counts=counts + ) def _failed_file(file_id: str, error_message: str | None): @@ -53,30 +62,35 @@ def _failed_file(file_id: str, error_message: str | None): class TestOpenAIVectorStoreCrudUpdateSuccess: def test_completes_when_all_files_complete(self, crud, mock_client, docs): - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - _batch_result(completed=2, failed=0) - ) + _wire_batch(mock_client, completed=2, failed=0) crud.update("vs_1", docs) - _, kwargs = mock_client.vector_stores.file_batches.upload_and_poll.call_args + _, kwargs = mock_client.vector_stores.file_batches.create.call_args assert kwargs["vector_store_id"] == "vs_1" assert kwargs["file_ids"] == ["file-1", "file-2"] # list_files should not have been called on the happy path mock_client.vector_stores.file_batches.list_files.assert_not_called() + def test_polls_with_the_create_batch_id(self, crud, mock_client, docs): + _wire_batch(mock_client, completed=2, failed=0) + + crud.update("vs_1", docs) + + args, kwargs = mock_client.vector_stores.file_batches.poll.call_args + assert args[0] == REAL_BATCH_ID + assert kwargs["vector_store_id"] == "vs_1" + def test_skips_upload_when_no_docs(self, crud, mock_client): crud.update("vs_1", []) - mock_client.vector_stores.file_batches.upload_and_poll.assert_not_called() + mock_client.vector_stores.file_batches.create.assert_not_called() class TestOpenAIVectorStoreCrudUpdatePartialFailure: """Failed files -> RuntimeError with per-file reasons labelled by fname.""" def test_includes_failed_fnames_and_messages(self, crud, mock_client, docs): - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - _batch_result(completed=1, failed=1) - ) + _wire_batch(mock_client, completed=1, failed=1) mock_client.vector_stores.file_batches.list_files.return_value = [ _failed_file("file-1", "Unsupported file type"), _failed_file("file-2", "File too large"), @@ -89,15 +103,29 @@ def test_includes_failed_fnames_and_messages(self, crud, mock_client, docs): assert "file1.pdf: Unsupported file type" in msg assert "file2.pdf: File too large" in msg + def test_looks_up_failures_with_create_batch_id_not_poll_id( + self, crud, mock_client, docs + ): + """Regression: poll()'s return .id is the vs_ id, which list_files rejects.""" + _wire_batch(mock_client, completed=1, failed=1) + mock_client.vector_stores.file_batches.list_files.return_value = [ + _failed_file("file-1", "Unsupported file type") + ] + + with pytest.raises(RuntimeError): + crud.update("vs_1", docs) + + _, kwargs = mock_client.vector_stores.file_batches.list_files.call_args + assert kwargs["batch_id"] == REAL_BATCH_ID + assert kwargs["batch_id"] != CORRUPT_POLL_ID + def test_reports_no_error_detail_when_last_error_missing( self, crud, mock_client, docs ): """A failed file with no `last_error` shouldn't drop out of the summary — it gets 'no error detail' so the user sees that something was wrong with that file even if OpenAI didn't tell us what.""" - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - _batch_result(completed=1, failed=1) - ) + _wire_batch(mock_client, completed=1, failed=1) mock_client.vector_stores.file_batches.list_files.return_value = [ _failed_file("file-1", None) ] @@ -109,9 +137,7 @@ def test_falls_back_to_file_id_label_for_unknown_file( self, crud, mock_client, docs ): """A failed file ID not matching any doc is labelled by its file ID.""" - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - _batch_result(completed=1, failed=1) - ) + _wire_batch(mock_client, completed=1, failed=1) mock_client.vector_stores.file_batches.list_files.return_value = [ _failed_file("file-unknown", "parse error") ] @@ -122,9 +148,7 @@ def test_falls_back_to_file_id_label_for_unknown_file( def test_reraises_when_list_files_errors(self, crud, mock_client, docs): """If the follow-up list_files lookup itself raises, the OpenAI error propagates instead of masking the real upload problem.""" - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - _batch_result(completed=0, failed=2) - ) + _wire_batch(mock_client, completed=0, failed=2) mock_client.vector_stores.file_batches.list_files.side_effect = ( openai.OpenAIError("list failed") ) @@ -134,7 +158,7 @@ def test_reraises_when_list_files_errors(self, crud, mock_client, docs): class TestOpenAIVectorStoreCrudUpdateOpenAIExceptions: - """`upload_and_poll` raising each specific OpenAI exception type maps to + """`create` raising each specific OpenAI exception type maps to `InterruptedError` with a category-prefixed message that includes the upstream status code and a remediation hint. @@ -230,7 +254,7 @@ def test_specific_openai_exception_maps_to_category_prefix( expected_status, original_message, ): - mock_client.vector_stores.file_batches.upload_and_poll.side_effect = ( + mock_client.vector_stores.file_batches.create.side_effect = ( exception_factory() ) @@ -243,7 +267,7 @@ def test_specific_openai_exception_maps_to_category_prefix( def test_api_timeout_error(self, crud, mock_client, docs): """APITimeoutError doesn't expose .message — handler interpolates str(e).""" - mock_client.vector_stores.file_batches.upload_and_poll.side_effect = ( + mock_client.vector_stores.file_batches.create.side_effect = ( openai.APITimeoutError(request=MagicMock()) ) @@ -256,7 +280,7 @@ def test_generic_openai_error_falls_through(self, crud, mock_client, docs): bottom-most `except openai.OpenAIError` block — prefixed with the generic "OpenAI error" tag but still carrying the original message. """ - mock_client.vector_stores.file_batches.upload_and_poll.side_effect = ( + mock_client.vector_stores.file_batches.create.side_effect = ( openai.OpenAIError("something else") ) diff --git a/backend/app/tests/services/collections/providers/test_openai_provider.py b/backend/app/tests/services/collections/providers/test_openai_provider.py index 1c07c100f..7f75c1bd9 100644 --- a/backend/app/tests/services/collections/providers/test_openai_provider.py +++ b/backend/app/tests/services/collections/providers/test_openai_provider.py @@ -380,11 +380,13 @@ def test_upload_files_first_failure_stops_remaining_docs() -> None: # --------------------------------------------------------------------------- -def _make_batch(completed: int, failed: int) -> MagicMock: - batch = MagicMock() +def _wire_batch(client: MagicMock, completed: int, failed: int) -> None: + """create() yields the real vsfb_ id; poll()'s return carries the corrupt vs_ id.""" + client.vector_stores.file_batches.create.return_value = MagicMock(id="vsfb_real") + batch = MagicMock(id="vs_corrupt") batch.file_counts.completed = completed batch.file_counts.failed = failed - return batch + client.vector_stores.file_batches.poll.return_value = batch def _make_openai_doc(file_id: str = "file-abc", fname: str = "doc.pdf") -> MagicMock: @@ -398,22 +400,20 @@ def test_vector_store_update_skips_when_no_docs() -> None: client = MagicMock() crud = OpenAIVectorStoreCrud(client) crud.update("vs_123", []) - client.vector_stores.file_batches.upload_and_poll.assert_not_called() + client.vector_stores.file_batches.create.assert_not_called() def test_vector_store_update_succeeds_with_no_failures() -> None: client = MagicMock() - client.vector_stores.file_batches.upload_and_poll.return_value = _make_batch( - completed=3, failed=0 - ) + _wire_batch(client, completed=3, failed=0) crud = OpenAIVectorStoreCrud(client) crud.update("vs_123", [_make_openai_doc() for _ in range(3)]) - client.vector_stores.file_batches.upload_and_poll.assert_called_once() + client.vector_stores.file_batches.create.assert_called_once() def test_vector_store_update_raises_on_openai_error() -> None: client = MagicMock() - client.vector_stores.file_batches.upload_and_poll.side_effect = OpenAIError( + client.vector_stores.file_batches.create.side_effect = OpenAIError( "rate limit" ) crud = OpenAIVectorStoreCrud(client) @@ -430,9 +430,7 @@ def _make_failed_file(message: str) -> MagicMock: def test_vector_store_update_raises_on_partial_failures() -> None: client = MagicMock() - client.vector_stores.file_batches.upload_and_poll.return_value = _make_batch( - completed=2, failed=1 - ) + _wire_batch(client, completed=2, failed=1) client.vector_stores.file_batches.list_files.return_value = [ _make_failed_file("unsupported file type") ] @@ -444,9 +442,7 @@ def test_vector_store_update_raises_on_partial_failures() -> None: def test_vector_store_update_raises_on_all_failures() -> None: client = MagicMock() - client.vector_stores.file_batches.upload_and_poll.return_value = _make_batch( - completed=0, failed=2 - ) + _wire_batch(client, completed=0, failed=2) client.vector_stores.file_batches.list_files.return_value = [ _make_failed_file("invalid pdf"), _make_failed_file("parse error"), @@ -459,15 +455,13 @@ def test_vector_store_update_raises_on_all_failures() -> None: def test_vector_store_update_passes_file_ids_to_openai() -> None: client = MagicMock() - client.vector_stores.file_batches.upload_and_poll.return_value = _make_batch( - completed=2, failed=0 - ) + _wire_batch(client, completed=2, failed=0) crud = OpenAIVectorStoreCrud(client) docs = [_make_openai_doc("file-1"), _make_openai_doc("file-2")] crud.update("vs_abc", docs) - _, kwargs = client.vector_stores.file_batches.upload_and_poll.call_args + _, kwargs = client.vector_stores.file_batches.create.call_args assert kwargs["vector_store_id"] == "vs_abc" assert kwargs["file_ids"] == ["file-1", "file-2"] diff --git a/backend/app/tests/utils/llm_provider.py b/backend/app/tests/utils/llm_provider.py index 564f27984..df135141c 100644 --- a/backend/app/tests/utils/llm_provider.py +++ b/backend/app/tests/utils/llm_provider.py @@ -116,11 +116,12 @@ def get_mock_openai_client_with_vector_store() -> MagicMock: mock_client.vector_stores.create.return_value = mock_vector_store mock_file_batch = MagicMock() + mock_file_batch.id = "vsfb_mock" mock_file_batch.file_counts.completed = 2 mock_file_batch.file_counts.total = 2 - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - mock_file_batch - ) + mock_file_batch.file_counts.failed = 0 + mock_client.vector_stores.file_batches.create.return_value = mock_file_batch + mock_client.vector_stores.file_batches.poll.return_value = mock_file_batch mock_client.vector_stores.files.list.return_value = {"data": []} diff --git a/backend/app/tests/utils/openai.py b/backend/app/tests/utils/openai.py index 4cceab033..093f45ae6 100644 --- a/backend/app/tests/utils/openai.py +++ b/backend/app/tests/utils/openai.py @@ -116,13 +116,14 @@ def get_mock_openai_client_with_vector_store() -> MagicMock: mock_vector_store.id = "mock_vector_store_id" mock_client.vector_stores.create.return_value = mock_vector_store - # File upload + polling + # File batch creation + polling mock_file_batch = MagicMock() + mock_file_batch.id = "vsfb_mock" mock_file_batch.file_counts.completed = 2 mock_file_batch.file_counts.total = 2 - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - mock_file_batch - ) + mock_file_batch.file_counts.failed = 0 + mock_client.vector_stores.file_batches.create.return_value = mock_file_batch + mock_client.vector_stores.file_batches.poll.return_value = mock_file_batch # File list mock_client.vector_stores.files.list.return_value = {"data": []} diff --git a/docs/wiki/modules/knowledge-base.md b/docs/wiki/modules/knowledge-base.md index 4307e6b30..1582afa74 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -35,3 +35,4 @@ All paths relative to `backend/app/`. ## Gotchas - Uploads de-duplicate by provider file ID (see deep dive §7). - Collections are immutable-ish: deletion semantics in deep dive §10. +- OpenAI file-batch id: the SDK's `file_batches.poll()` / `upload_and_poll()` final return deserializes a vector-store body, so its `.id` is the `vs_` id, not the `vsfb_` batch id. `crud/rag/open_ai.py` captures the batch id from `create()` before polling and uses it for `list_files`. Any failed file is a hard failure (whole vector store rolled back); partial indexing needs an add-documents endpoint first. From 6ca053f379b51e7f03927708995ca8f61d46dcfa Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Fri, 24 Jul 2026 02:55:56 +0530 Subject: [PATCH 02/10] fix(tests): Simplify exception handling in OpenAI vector store tests --- backend/app/tests/crud/rag/test_open_ai.py | 8 +++----- .../collections/providers/test_openai_provider.py | 4 +--- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/backend/app/tests/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index fb5693331..04558c28c 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -254,9 +254,7 @@ def test_specific_openai_exception_maps_to_category_prefix( expected_status, original_message, ): - mock_client.vector_stores.file_batches.create.side_effect = ( - exception_factory() - ) + mock_client.vector_stores.file_batches.create.side_effect = exception_factory() with pytest.raises(InterruptedError) as exc_info: crud.update("vs_1", docs) @@ -280,8 +278,8 @@ def test_generic_openai_error_falls_through(self, crud, mock_client, docs): bottom-most `except openai.OpenAIError` block — prefixed with the generic "OpenAI error" tag but still carrying the original message. """ - mock_client.vector_stores.file_batches.create.side_effect = ( - openai.OpenAIError("something else") + mock_client.vector_stores.file_batches.create.side_effect = openai.OpenAIError( + "something else" ) with pytest.raises(InterruptedError) as exc_info: diff --git a/backend/app/tests/services/collections/providers/test_openai_provider.py b/backend/app/tests/services/collections/providers/test_openai_provider.py index 7f75c1bd9..9911bf133 100644 --- a/backend/app/tests/services/collections/providers/test_openai_provider.py +++ b/backend/app/tests/services/collections/providers/test_openai_provider.py @@ -413,9 +413,7 @@ def test_vector_store_update_succeeds_with_no_failures() -> None: def test_vector_store_update_raises_on_openai_error() -> None: client = MagicMock() - client.vector_stores.file_batches.create.side_effect = OpenAIError( - "rate limit" - ) + client.vector_stores.file_batches.create.side_effect = OpenAIError("rate limit") crud = OpenAIVectorStoreCrud(client) with pytest.raises(InterruptedError, match="rate limit"): From 8c02d57a37f02b10d9dfa815a6bbe296cbf12b47 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Sun, 26 Jul 2026 09:40:53 +0530 Subject: [PATCH 03/10] fix(upload): Replace pre-transform validation with upload validation in document upload process --- backend/app/api/routes/documents.py | 6 +- backend/app/services/doctransform/registry.py | 4 + backend/app/services/documents/helpers.py | 307 +++++++++++++++++- 3 files changed, 313 insertions(+), 4 deletions(-) diff --git a/backend/app/api/routes/documents.py b/backend/app/api/routes/documents.py index 838216337..ae553abca 100644 --- a/backend/app/api/routes/documents.py +++ b/backend/app/api/routes/documents.py @@ -33,7 +33,7 @@ from app.services.documents.helpers import ( calculate_file_size, schedule_transformation, - pre_transform_validation, + validate_upload, build_document_schema, build_document_schemas, ) @@ -126,8 +126,8 @@ async def upload_doc( if callback_url: validate_callback_url(callback_url) - source_format, actual_transformer = pre_transform_validation( - src_filename=src.filename, + source_format, actual_transformer = validate_upload( + src=src, target_format=target_format, transformer=transformer, ) diff --git a/backend/app/services/doctransform/registry.py b/backend/app/services/doctransform/registry.py index e84be60fb..b0507ce3c 100644 --- a/backend/app/services/doctransform/registry.py +++ b/backend/app/services/doctransform/registry.py @@ -39,6 +39,8 @@ class TransformationError(Exception): ".markdown": "markdown", ".csv": "csv", ".json": "json", + ".xlsx": "xlsx", + ".xls": "xls", } # Map format names to file extensions @@ -51,6 +53,8 @@ class TransformationError(Exception): "markdown": ".md", "csv": ".csv", "json": ".json", + "xlsx": ".xlsx", + "xls": ".xls", } diff --git a/backend/app/services/documents/helpers.py b/backend/app/services/documents/helpers.py index 7d78b6160..f4818337f 100644 --- a/backend/app/services/documents/helpers.py +++ b/backend/app/services/documents/helpers.py @@ -1,4 +1,9 @@ -from typing import Optional, Tuple, Iterable, Union +import csv +import json +import logging +from codecs import getincrementaldecoder +from dataclasses import dataclass +from typing import Callable, Optional, Tuple, Iterable, Union from uuid import UUID from fastapi import HTTPException, UploadFile @@ -23,6 +28,306 @@ ) +logger = logging.getLogger(__name__) + +SNIFF_HEAD_BYTES = 4096 +SNIFF_TAIL_BYTES = 2048 + +PDF_EOF_MARKER = b"%%EOF" +PDF_ENCRYPT_MARKER = b"/Encrypt" +OOXML_CONTENT_TYPES = b"[Content_Types].xml" +OOXML_WORD_PART = b"word/" +OOXML_EXCEL_PART = b"xl/" +JSON_OPENERS = (b"{", b"[") +JSON_CLOSERS: dict[int, int] = {ord("{"): ord("}"), ord("["): ord("]")} +UTF8_BOM = b"\xef\xbb\xbf" + +# Past this size a full parse is no longer cheap enough for the request path; +# larger documents fall back to the O(1) truncation check in _check_json_tail. +JSON_FULL_PARSE_MAX_BYTES = 256 * 1024 + +# Set False if ragged rows turn out to be common in real uploads. +CSV_STRICT_COLUMN_COUNT = True +CSV_SAMPLE_MAX_ROWS = 20 + +MISLABELLED_BINARY_SIGNATURES: dict[bytes, str] = { + b"PK\x03\x04": "an Office/zip file (xlsx, docx)", + b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1": "a legacy Office file (xls, doc)", + b"%PDF-": "a PDF", +} + +UNNAMED_FILE_PLACEHOLDER = "" + + +CLIENT_VALIDATION_MESSAGE = ( + "Document '{filename}' failed parsing. The document appears to be corrupted " + "or invalid. Please re-upload a valid document." +) + + +class DocumentValidationError(Exception): + """Raised when a document fails the pre-upload sanity check.""" + + def __init__(self, filename: str, reason: str) -> None: + self.filename = filename + self.reason = reason + super().__init__(f"Document '{filename}' failed validation: {reason}") + + @property + def client_message(self) -> str: + """Client-safe wording. `reason` stays internal - it is log-only.""" + return CLIENT_VALIDATION_MESSAGE.format(filename=self.filename) + + +def _decode_utf8(filename: str, sample: bytes) -> str: + if b"\x00" in sample: + raise DocumentValidationError( + filename, + "parsing error - NUL byte found in a text-based file; " + "the file is binary, corrupt, or has the wrong extension", + ) + try: + # Buffers a multi-byte char split by the sample boundary instead of + # reporting it as a decode failure. + return getincrementaldecoder("utf-8")().decode(sample) + except UnicodeDecodeError as e: + raise DocumentValidationError( + filename, + f"parsing error - content is not valid UTF-8 at byte offset {e.start}; " + "re-export the file as UTF-8", + ) from e + + +def _check_pdf(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + if PDF_EOF_MARKER not in tail: + raise DocumentValidationError( + filename, + "parsing error - PDF is missing its end-of-file marker; " + "the upload is truncated or incomplete", + ) + if PDF_ENCRYPT_MARKER in tail: + raise DocumentValidationError( + filename, + "parsing error - PDF is password protected or encrypted and cannot be read", + ) + + +def _check_ooxml(filename: str, head: bytes, expected_part: bytes) -> None: + if OOXML_CONTENT_TYPES not in head: + raise DocumentValidationError( + filename, + "parsing error - file is a zip archive but not a valid Office document " + "(no OOXML content-types part)", + ) + + other_part = ( + OOXML_EXCEL_PART if expected_part == OOXML_WORD_PART else OOXML_WORD_PART + ) + # Only rejects on positive evidence of the wrong package type; neither + # marker landing in the sample leaves the file alone. + if other_part in head and expected_part not in head: + raise DocumentValidationError( + filename, + f"parsing error - file is a '{other_part.decode()}' Office package, " + "not the format its extension claims", + ) + + +def _check_docx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ooxml(filename, head, OOXML_WORD_PART) + + +def _check_xlsx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ooxml(filename, head, OOXML_EXCEL_PART) + + +def _check_csv(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + for signature, description in MISLABELLED_BINARY_SIGNATURES.items(): + if head.startswith(signature): + raise DocumentValidationError( + filename, + f"parsing error - file is {description} renamed to a CSV, not text; " + "export it as CSV instead of renaming it", + ) + + text = _decode_utf8(filename, head) + + lines = text.splitlines() + if len(head) == SNIFF_HEAD_BYTES and len(lines) > 1: + # Only a head-sized read can have cut its last line in half. + lines = lines[:-1] + + try: + rows = [row for row in csv.reader(lines[:CSV_SAMPLE_MAX_ROWS]) if row] + except csv.Error as e: + raise DocumentValidationError( + filename, f"parsing error - malformed CSV: {e}" + ) from e + + if not rows: + raise DocumentValidationError( + filename, "parsing error - no readable CSV rows found" + ) + + if CSV_STRICT_COLUMN_COUNT: + expected_columns = len(rows[0]) + for line_number, row in enumerate(rows[1:], start=2): + if len(row) != expected_columns: + raise DocumentValidationError( + filename, + f"parsing error - inconsistent column count at row {line_number} " + f"({len(row)} fields, expected {expected_columns})", + ) + + +def _check_json_tail(filename: str, opener: int, tail: bytes) -> None: + """Truncation check for JSON too large to parse: the opening bracket must be + matched by the last non-whitespace byte.""" + closing = tail.rstrip() + if not closing or closing[-1] != JSON_CLOSERS[opener]: + raise DocumentValidationError( + filename, + "parsing error - JSON is not closed correctly; " + "the upload is truncated or incomplete", + ) + + +def _check_json(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + body = head.lstrip(UTF8_BOM).lstrip() + if not body.startswith(JSON_OPENERS): + raise DocumentValidationError( + filename, "parsing error - JSON document must start with '{' or '['" + ) + + stream = file.file + stream.seek(0, 2) + size_bytes = stream.tell() + stream.seek(0) + + if size_bytes > JSON_FULL_PARSE_MAX_BYTES: + logger.info( + f"[_check_json] Too large for a full parse, checking closure only | " + f"filename: {filename} | size_bytes: {size_bytes}" + ) + _check_json_tail(filename, body[0], tail) + return + + try: + # The hook discards every object as it is parsed: the whole document is + # still validated, but no result graph is built, cutting peak memory + # roughly 5x versus a plain json.loads. + json.loads(stream.read(), object_pairs_hook=lambda pairs: None) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + raise DocumentValidationError( + filename, f"parsing error - invalid JSON: {e}" + ) from e + finally: + stream.seek(0) + + +@dataclass(frozen=True) +class FormatSpec: + signatures: tuple[bytes, ...] = () + needs_tail: bool = False + checker: Callable[[str, bytes, bytes, UploadFile], None] | None = None + + +# Deliberately partial: text, markdown and html are omitted because any +# decodable byte string is a valid document in those formats. +FORMAT_SPECS: dict[str, FormatSpec] = { + "pdf": FormatSpec(signatures=(b"%PDF-",), needs_tail=True, checker=_check_pdf), + # OOXML is a zip; empty/spanned archives use the other two headers. + "docx": FormatSpec( + signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), checker=_check_docx + ), + "xlsx": FormatSpec( + signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), checker=_check_xlsx + ), + # OLE2, shared with .doc/.ppt/.msg - telling them apart needs a real parse. + "xls": FormatSpec(signatures=(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",)), + "csv": FormatSpec(checker=_check_csv), + "json": FormatSpec(needs_tail=True, checker=_check_json), +} + + +def _read_edges(file: UploadFile, needs_tail: bool) -> Tuple[bytes, bytes]: + """Sample the head (and optionally tail) of the upload, leaving it rewound.""" + stream = file.file + + stream.seek(0) + head = stream.read(SNIFF_HEAD_BYTES) + + tail = b"" + if needs_tail: + stream.seek(0, 2) + stream.seek(max(0, stream.tell() - SNIFF_TAIL_BYTES)) + tail = stream.read(SNIFF_TAIL_BYTES) + + stream.seek(0) + return head, tail + + +def validate_document_content(*, file: UploadFile, source_format: str) -> None: + """ + Sanity-check an upload's bytes before it is persisted. Only the first + SNIFF_HEAD_BYTES (plus SNIFF_TAIL_BYTES for PDF) are read, so the cost is + constant in file size - small JSON, which is fully parsed, is the exception. + + Raises: + DocumentValidationError: naming the offending file and the reason. + """ + filename = (file.filename or "").strip() or UNNAMED_FILE_PLACEHOLDER + + spec = FORMAT_SPECS.get(source_format) + head, tail = _read_edges(file, needs_tail=bool(spec and spec.needs_tail)) + + if not head: + raise DocumentValidationError(filename, "the file is empty") + + if spec is None: + return + + if spec.signatures and not head.startswith(spec.signatures): + raise DocumentValidationError( + filename, + f"parsing error - content does not match the {source_format} file " + "signature; the file is corrupt or has the wrong extension", + ) + + if spec.checker: + spec.checker(filename, head, tail, file) + + +def validate_upload( + *, + src: UploadFile, + target_format: str | None, + transformer: str | None, +) -> Tuple[str, str | None]: + """ + Full pre-storage gate: extension and transformer validation plus a content + sanity check. Returns (source_format, actual_transformer_or_none). + + Raises: HTTPException(400) on client errors. + """ + source_format, actual_transformer = pre_transform_validation( + src_filename=src.filename, + target_format=target_format, + transformer=transformer, + ) + + try: + validate_document_content(file=src, source_format=source_format) + except DocumentValidationError as e: + logger.warning( + f"[validate_upload] Document failed sanity check | " + f"filename: {e.filename} | format: {source_format} | reason: {e.reason}" + ) + raise HTTPException(status_code=400, detail=e.client_message) + + return source_format, actual_transformer + + def calculate_file_size(file: UploadFile) -> float: """ Calculate the size of an uploaded file in kilobytes. From d391ad26dddaf826d33944cefa193670e96bc514 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:18:17 +0530 Subject: [PATCH 04/10] fix(openai): Enhance OpenAI vector store handling with improved batch processing and error management --- backend/app/crud/rag/open_ai.py | 72 +++++- .../services/collections/create_collection.py | 88 ++++++- .../collections/providers/registry.py | 7 +- backend/app/services/documents/helpers.py | 44 ++-- backend/app/tests/crud/rag/test_open_ai.py | 10 +- .../providers/test_openai_provider.py | 6 +- .../collections/test_create_collection.py | 223 +++++++++++++++++- .../kaapi-knowledge-base-ARCHITECTURE.md | 94 +++++--- docs/wiki/modules/knowledge-base.md | 7 + 9 files changed, 464 insertions(+), 87 deletions(-) diff --git a/backend/app/crud/rag/open_ai.py b/backend/app/crud/rag/open_ai.py index 6d80e12c9..57a8409dd 100644 --- a/backend/app/crud/rag/open_ai.py +++ b/backend/app/crud/rag/open_ai.py @@ -1,9 +1,12 @@ import json import logging +import time import functools as ft import openai from openai import OpenAI, OpenAIError +from openai.types import VectorStore +from openai.types.vector_stores import VectorStoreFileBatch from pydantic import BaseModel from app.models import Document, ProviderType @@ -12,6 +15,15 @@ OPENAI_PROVIDER = ProviderType.openai.value +OPENAI_MAX_RETRIES = 2 + +# Per socket operation, not per call: a steady upload runs as long as it needs and +# this only fires on a dead connection. Sized so all attempts fit one Celery task. +OPENAI_TIMEOUT_SECONDS = 90 + +BATCH_POLL_INTERVAL_SECONDS = 2 +BATCH_POLL_TIMEOUT_SECONDS = 240 + def vs_ls(client: OpenAI, vector_store_id: str): kwargs = {} @@ -85,7 +97,7 @@ def __init__(self, client): class OpenAIVectorStoreCrud(OpenAICrud): - def create(self): + def create(self) -> VectorStore: logger.info( f"[OpenAIVectorStoreCrud.create] Creating vector store | {{'action': 'create'}}" ) @@ -101,6 +113,52 @@ def read(self, vector_store_id: str): ) yield from vs_ls(self.client, vector_store_id) + def _create_file_batch(self, vector_store_id: str, file_ids: list[str]) -> str: + """Returns the vsfb_ id. poll()'s return deserializes a vector-store body, + so its .id is the vs_ id - take the batch id from create().""" + created = self.client.vector_stores.file_batches.create( + vector_store_id=vector_store_id, + file_ids=file_ids, + ) + return created.id + + def _retrieve_file_batch( + self, batch_id: str, vector_store_id: str + ) -> VectorStoreFileBatch: + return self.client.vector_stores.file_batches.retrieve( + batch_id, vector_store_id=vector_store_id + ) + + def _poll_file_batch( + self, batch_id: str, vector_store_id: str + ) -> VectorStoreFileBatch: + """Wait for indexing to finish. The SDK's own poll() loops forever.""" + deadline = time.monotonic() + BATCH_POLL_TIMEOUT_SECONDS + + while True: + batch = self._retrieve_file_batch(batch_id, vector_store_id) + if batch.status != "in_progress": + return batch + + if time.monotonic() >= deadline: + error_message = ( + f"[KAAPI] Timed out (code: batch-poll-timeout) after " + f"{BATCH_POLL_TIMEOUT_SECONDS}s waiting for OpenAI to finish indexing " + f"{batch.file_counts.in_progress} file(s). Retry the " + f"collection, or split it into smaller batches." + ) + logger.error( + f"[OpenAIVectorStoreCrud._poll_file_batch] {error_message} | " + f"vector_store_id={vector_store_id}, batch_id={batch_id}, " + f"completed={batch.file_counts.completed}, " + f"in_progress={batch.file_counts.in_progress}" + ) + raise InterruptedError(error_message) + + time.sleep( + min(BATCH_POLL_INTERVAL_SECONDS, max(0.0, deadline - time.monotonic())) + ) + def update( self, vector_store_id: str, @@ -115,16 +173,10 @@ def update( ) try: - created = self.client.vector_stores.file_batches.create( - vector_store_id=vector_store_id, - file_ids=[doc.file_id[OPENAI_PROVIDER] for doc in docs], - ) - # poll()'s return deserializes a vector-store body, so its .id is the vs_ id; - # capture the real vsfb_ batch id from create() before polling. - batch_id = created.id - batch = self.client.vector_stores.file_batches.poll( - batch_id, vector_store_id=vector_store_id + batch_id = self._create_file_batch( + vector_store_id, [doc.file_id[OPENAI_PROVIDER] for doc in docs] ) + batch = self._poll_file_batch(batch_id, vector_store_id) except openai.RateLimitError as e: error_message = ( f"[OPENAI] Rate limit exceeded (code: {e.status_code}): " diff --git a/backend/app/services/collections/create_collection.py b/backend/app/services/collections/create_collection.py index 234443a00..f810fe4d2 100644 --- a/backend/app/services/collections/create_collection.py +++ b/backend/app/services/collections/create_collection.py @@ -4,7 +4,7 @@ from sqlmodel import Session from gevent import Timeout -from celery.exceptions import SoftTimeLimitExceeded +from celery.exceptions import Retry, SoftTimeLimitExceeded from opentelemetry import trace from asgi_correlation_id import correlation_id @@ -40,6 +40,48 @@ tracer = trace.get_tracer(__name__) +def _can_retry_batch(task_instance) -> bool: + """Whether a timed-out batch has attempts left. + + Worth retrying because upload_files persists each file ID as it uploads, so + every attempt resumes where the last stopped. `max_retries=None` means retry + forever, which would never fail the job or fire its callback; treat as none. + """ + if task_instance is None or task_instance.max_retries is None: + return False + return task_instance.request.retries < task_instance.max_retries + + +def _note_batch_retry( + project_id: int, + job_id: str, + batch_number: int, + attempt: int, + max_attempts: int, +) -> None: + """Record a pending retry on the job row. + + The job stays PROCESSING between attempts and pending_monitor only watches + PENDING, so a retry that never lands would be invisible. + """ + try: + with Session(engine) as session: + CollectionJobCrud(session, project_id).update( + UUID(job_id), + CollectionJobUpdate( + error_message=( + f"Batch {batch_number} timed out; " + f"retry {attempt}/{max_attempts} queued" + ) + ), + ) + except Exception: + logger.warning( + "[create_collection._note_batch_retry] Could not record retry | job_id=%s", + job_id, + ) + + def start_job( db: Session, request: CreationRequest, @@ -403,9 +445,13 @@ def execute_batch_job( collection_job_crud = CollectionJobCrud(session, project_id) collection_job = collection_job_crud.read_one(job_uuid) already_uploaded = collection_job.documents_uploaded or [] - now_uploaded = already_uploaded + [ - str(d) for d in all_doc_ids_this_batch - ] + # A re-run past this point re-appends its own IDs, and read_each + # rejects a list whose duplicates collapse in its IN. + now_uploaded = list( + dict.fromkeys( + already_uploaded + [str(d) for d in all_doc_ids_this_batch] + ) + ) collection_job = collection_job_crud.update( job_uuid, @@ -509,6 +555,40 @@ def execute_batch_job( except (Timeout, SoftTimeLimitExceeded) as err: timeout_err = TimeoutError("Task exceeded soft time limit") + + if _can_retry_batch(task_instance): + attempt = task_instance.request.retries + 1 + logger.warning( + "[create_collection.execute_batch_job] Batch timed out, retrying | " + "{'collection_job_id': '%s', 'batch_number': %d, 'attempt': '%d/%d'}", + job_id, + batch_number, + attempt, + task_instance.max_retries, + ) + span.set_attribute("collection.batch_retry", attempt) + _note_batch_retry( + project_id, + job_id, + batch_number, + attempt, + task_instance.max_retries, + ) + try: + # Only Retry means it reached the broker. retry() also raises + # Reject on a failed publish, or exc outside a worker - those + # must fail the job, not strand it in PROCESSING. + raise task_instance.retry(exc=timeout_err) + except Retry: + raise + except Exception: + logger.exception( + "[create_collection.execute_batch_job] Retry could not be scheduled, " + "failing the job | {'collection_job_id': '%s', 'batch_number': %d}", + job_id, + batch_number, + ) + logger.warning( "[create_collection.execute_batch_job] Collection Creation Timed Out | {'collection_job_id': '%s', 'error': '%s'}", job_id, diff --git a/backend/app/services/collections/providers/registry.py b/backend/app/services/collections/providers/registry.py index 7c40d8e5c..a4f9ed2da 100644 --- a/backend/app/services/collections/providers/registry.py +++ b/backend/app/services/collections/providers/registry.py @@ -5,6 +5,7 @@ from openai import OpenAI from app.crud import get_provider_credential +from app.crud.rag.open_ai import OPENAI_MAX_RETRIES, OPENAI_TIMEOUT_SECONDS from app.services.collections.providers.base import BaseProvider from app.services.collections.providers.gemini import GeminiAIStudioProvider from app.services.collections.providers.openai import OpenAIProvider @@ -63,7 +64,11 @@ def get_llm_provider( if provider == LLMProvider.OPENAI: if "api_key" not in credentials: raise ValueError("OpenAI credentials not configured for this project.") - client = OpenAI(api_key=credentials["api_key"]) + client = OpenAI( + api_key=credentials["api_key"], + max_retries=OPENAI_MAX_RETRIES, + timeout=OPENAI_TIMEOUT_SECONDS, + ) elif provider == LLMProvider.GOOGLE_AISTUDIO: if "api_key" not in credentials: raise ValueError( diff --git a/backend/app/services/documents/helpers.py b/backend/app/services/documents/helpers.py index f4818337f..27ed2a800 100644 --- a/backend/app/services/documents/helpers.py +++ b/backend/app/services/documents/helpers.py @@ -42,17 +42,16 @@ JSON_CLOSERS: dict[int, int] = {ord("{"): ord("}"), ord("["): ord("]")} UTF8_BOM = b"\xef\xbb\xbf" -# Past this size a full parse is no longer cheap enough for the request path; -# larger documents fall back to the O(1) truncation check in _check_json_tail. JSON_FULL_PARSE_MAX_BYTES = 256 * 1024 -# Set False if ragged rows turn out to be common in real uploads. CSV_STRICT_COLUMN_COUNT = True -CSV_SAMPLE_MAX_ROWS = 20 +CSV_SAMPLE_MAX_ROWS = 200 + +OLE2_SIGNATURE = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" MISLABELLED_BINARY_SIGNATURES: dict[bytes, str] = { b"PK\x03\x04": "an Office/zip file (xlsx, docx)", - b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1": "a legacy Office file (xls, doc)", + OLE2_SIGNATURE: "a legacy Office file (xls, doc)", b"%PDF-": "a PDF", } @@ -87,8 +86,6 @@ def _decode_utf8(filename: str, sample: bytes) -> str: "the file is binary, corrupt, or has the wrong extension", ) try: - # Buffers a multi-byte char split by the sample boundary instead of - # reporting it as a decode failure. return getincrementaldecoder("utf-8")().decode(sample) except UnicodeDecodeError as e: raise DocumentValidationError( @@ -112,8 +109,8 @@ def _check_pdf(filename: str, head: bytes, tail: bytes, file: UploadFile) -> Non ) -def _check_ooxml(filename: str, head: bytes, expected_part: bytes) -> None: - if OOXML_CONTENT_TYPES not in head: +def _check_ooxml(filename: str, sample: bytes, expected_part: bytes) -> None: + if OOXML_CONTENT_TYPES not in sample: raise DocumentValidationError( filename, "parsing error - file is a zip archive but not a valid Office document " @@ -123,9 +120,7 @@ def _check_ooxml(filename: str, head: bytes, expected_part: bytes) -> None: other_part = ( OOXML_EXCEL_PART if expected_part == OOXML_WORD_PART else OOXML_WORD_PART ) - # Only rejects on positive evidence of the wrong package type; neither - # marker landing in the sample leaves the file alone. - if other_part in head and expected_part not in head: + if other_part in sample and expected_part not in sample: raise DocumentValidationError( filename, f"parsing error - file is a '{other_part.decode()}' Office package, " @@ -134,11 +129,13 @@ def _check_ooxml(filename: str, head: bytes, expected_part: bytes) -> None: def _check_docx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - _check_ooxml(filename, head, OOXML_WORD_PART) + """Head+tail: the zip central directory lists all entry names and sits at the + end; the head holds only the first entry.""" + _check_ooxml(filename, head + tail, OOXML_WORD_PART) def _check_xlsx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - _check_ooxml(filename, head, OOXML_EXCEL_PART) + _check_ooxml(filename, head + tail, OOXML_EXCEL_PART) def _check_csv(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: @@ -154,7 +151,6 @@ def _check_csv(filename: str, head: bytes, tail: bytes, file: UploadFile) -> Non lines = text.splitlines() if len(head) == SNIFF_HEAD_BYTES and len(lines) > 1: - # Only a head-sized read can have cut its last line in half. lines = lines[:-1] try: @@ -213,9 +209,6 @@ def _check_json(filename: str, head: bytes, tail: bytes, file: UploadFile) -> No return try: - # The hook discards every object as it is parsed: the whole document is - # still validated, but no result graph is built, cutting peak memory - # roughly 5x versus a plain json.loads. json.loads(stream.read(), object_pairs_hook=lambda pairs: None) except (json.JSONDecodeError, UnicodeDecodeError) as e: raise DocumentValidationError( @@ -232,19 +225,20 @@ class FormatSpec: checker: Callable[[str, bytes, bytes, UploadFile], None] | None = None -# Deliberately partial: text, markdown and html are omitted because any -# decodable byte string is a valid document in those formats. FORMAT_SPECS: dict[str, FormatSpec] = { "pdf": FormatSpec(signatures=(b"%PDF-",), needs_tail=True, checker=_check_pdf), - # OOXML is a zip; empty/spanned archives use the other two headers. "docx": FormatSpec( - signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), checker=_check_docx + signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), + needs_tail=True, + checker=_check_docx, ), "xlsx": FormatSpec( - signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), checker=_check_xlsx + signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), + needs_tail=True, + checker=_check_xlsx, ), - # OLE2, shared with .doc/.ppt/.msg - telling them apart needs a real parse. - "xls": FormatSpec(signatures=(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",)), + "xls": FormatSpec(signatures=(OLE2_SIGNATURE,)), + "doc": FormatSpec(signatures=(OLE2_SIGNATURE,)), "csv": FormatSpec(checker=_check_csv), "json": FormatSpec(needs_tail=True, checker=_check_json), } diff --git a/backend/app/tests/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index 04558c28c..4d3d3de2b 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -42,13 +42,13 @@ def docs(): def _wire_batch(mock_client, *, completed: int, failed: int) -> None: - """Stub create() with the real vsfb_ id and poll() with the corrupt vs_ id the SDK returns.""" + """create() gives the real vsfb_ id; retrieve() gives the corrupt vs_ id.""" mock_client.vector_stores.file_batches.create.return_value = MagicMock( id=REAL_BATCH_ID ) - counts = MagicMock(completed=completed, failed=failed) - mock_client.vector_stores.file_batches.poll.return_value = MagicMock( - id=CORRUPT_POLL_ID, file_counts=counts + counts = MagicMock(completed=completed, failed=failed, in_progress=0) + mock_client.vector_stores.file_batches.retrieve.return_value = MagicMock( + id=CORRUPT_POLL_ID, status="completed", file_counts=counts ) @@ -77,7 +77,7 @@ def test_polls_with_the_create_batch_id(self, crud, mock_client, docs): crud.update("vs_1", docs) - args, kwargs = mock_client.vector_stores.file_batches.poll.call_args + args, kwargs = mock_client.vector_stores.file_batches.retrieve.call_args assert args[0] == REAL_BATCH_ID assert kwargs["vector_store_id"] == "vs_1" diff --git a/backend/app/tests/services/collections/providers/test_openai_provider.py b/backend/app/tests/services/collections/providers/test_openai_provider.py index 9911bf133..17e5e8306 100644 --- a/backend/app/tests/services/collections/providers/test_openai_provider.py +++ b/backend/app/tests/services/collections/providers/test_openai_provider.py @@ -381,12 +381,14 @@ def test_upload_files_first_failure_stops_remaining_docs() -> None: def _wire_batch(client: MagicMock, completed: int, failed: int) -> None: - """create() yields the real vsfb_ id; poll()'s return carries the corrupt vs_ id.""" + """create() gives the real vsfb_ id; retrieve() gives the corrupt vs_ id.""" client.vector_stores.file_batches.create.return_value = MagicMock(id="vsfb_real") batch = MagicMock(id="vs_corrupt") + batch.status = "completed" batch.file_counts.completed = completed batch.file_counts.failed = failed - client.vector_stores.file_batches.poll.return_value = batch + batch.file_counts.in_progress = 0 + client.vector_stores.file_batches.retrieve.return_value = batch def _make_openai_doc(file_id: str = "file-abc", fname: str = "doc.pdf") -> MagicMock: diff --git a/backend/app/tests/services/collections/test_create_collection.py b/backend/app/tests/services/collections/test_create_collection.py index e2b6b3475..112de9bcb 100644 --- a/backend/app/tests/services/collections/test_create_collection.py +++ b/backend/app/tests/services/collections/test_create_collection.py @@ -4,7 +4,7 @@ import uuid from uuid import UUID, uuid4 -from celery.exceptions import SoftTimeLimitExceeded +from celery.exceptions import Reject, Retry, SoftTimeLimitExceeded from gevent import Timeout import pytest @@ -703,6 +703,227 @@ def test_execute_batch_job_soft_time_limit_marks_failed_and_reraises( assert "soft time limit" in (updated_job.error_message or "") +def _celery_task_instance(retries: int, max_retries: int | None = 3) -> MagicMock: + """Bound-task stand-in; production always passes `task_instance=self`.""" + task = MagicMock() + task.request.retries = retries + task.max_retries = max_retries + task.retry.side_effect = Retry("requeued") + return task + + +@pytest.mark.parametrize("timeout_exc", [Timeout(300), SoftTimeLimitExceeded()]) +@patch("app.services.collections.create_collection.get_cloud_storage") +@patch("app.services.collections.create_collection.get_llm_provider") +def test_execute_batch_job_timeout_retries_and_keeps_vector_store( + mock_get_provider: MagicMock, + mock_get_storage: MagicMock, + timeout_exc: BaseException, + db: Session, +) -> None: + project = get_project(db) + store = DocumentStore(db=db, project_id=project.id) + doc = store.put() + + mock_provider = get_mock_provider("vs_123", "openai vector store") + mock_provider.create.side_effect = timeout_exc + mock_get_provider.return_value = mock_provider + + job = get_collection_job( + db, + project, + action_type=CollectionActionType.CREATE, + status=CollectionJobStatus.PROCESSING, + ) + request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) + task = _celery_task_instance(retries=0) + + patcher = _patch_session(db) + try: + with pytest.raises(Retry): + execute_batch_job( + request=request.model_dump(mode="json"), + project_id=project.id, + organization_id=project.organization_id, + task_id=str(uuid4()), + job_id=str(job.id), + task_instance=task, + vector_store_id="vs_123", + batch_number=1, + batch_doc_ids=[str(doc.id)], + remaining_batches=[], + ) + finally: + patcher.stop() + + task.retry.assert_called_once() + mock_provider.delete.assert_not_called() + + updated_job = CollectionJobCrud(db, project.id).read_one(job.id) + assert updated_job.status == CollectionJobStatus.PROCESSING + assert "retry 1/3 queued" in (updated_job.error_message or "") + + +@pytest.mark.parametrize( + "retry_effect", + [ + pytest.param(Reject(OSError("broker down"), False), id="broker_publish_fails"), + pytest.param(TimeoutError("soft limit"), id="called_outside_worker"), + ], +) +@patch("app.services.collections.create_collection.get_cloud_storage") +@patch("app.services.collections.create_collection.get_llm_provider") +def test_execute_batch_job_unschedulable_retry_fails_job( + mock_get_provider: MagicMock, + mock_get_storage: MagicMock, + retry_effect: BaseException, + db: Session, +) -> None: + """Reject (publish failed) and a bare exc (called outside a worker) both mean + nothing was queued, so the job must fail rather than strand in PROCESSING.""" + project = get_project(db) + store = DocumentStore(db=db, project_id=project.id) + doc = store.put() + + mock_provider = get_mock_provider("vs_123", "openai vector store") + mock_provider.create.side_effect = SoftTimeLimitExceeded() + mock_get_provider.return_value = mock_provider + + job = get_collection_job( + db, + project, + action_type=CollectionActionType.CREATE, + status=CollectionJobStatus.PROCESSING, + ) + request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) + task = _celery_task_instance(retries=0) + task.retry.side_effect = retry_effect + + patcher = _patch_session(db) + try: + with pytest.raises(SoftTimeLimitExceeded): + execute_batch_job( + request=request.model_dump(mode="json"), + project_id=project.id, + organization_id=project.organization_id, + task_id=str(uuid4()), + job_id=str(job.id), + task_instance=task, + vector_store_id="vs_123", + batch_number=1, + batch_doc_ids=[str(doc.id)], + remaining_batches=[], + ) + finally: + patcher.stop() + + mock_provider.delete.assert_called_once() + + updated_job = CollectionJobCrud(db, project.id).read_one(job.id) + assert updated_job.status == CollectionJobStatus.FAILED + + +@patch("app.services.collections.create_collection.get_cloud_storage") +@patch("app.services.collections.create_collection.get_llm_provider") +def test_execute_batch_job_infinite_max_retries_does_not_crash( + mock_get_provider: MagicMock, + mock_get_storage: MagicMock, + db: Session, +) -> None: + """`max_retries=None` is retry-forever; comparing against it raises TypeError, + which would skip failure handling entirely.""" + project = get_project(db) + store = DocumentStore(db=db, project_id=project.id) + doc = store.put() + + mock_provider = get_mock_provider("vs_123", "openai vector store") + mock_provider.create.side_effect = SoftTimeLimitExceeded() + mock_get_provider.return_value = mock_provider + + job = get_collection_job( + db, + project, + action_type=CollectionActionType.CREATE, + status=CollectionJobStatus.PROCESSING, + ) + request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) + task = _celery_task_instance(retries=0, max_retries=None) + + patcher = _patch_session(db) + try: + with pytest.raises(SoftTimeLimitExceeded): + execute_batch_job( + request=request.model_dump(mode="json"), + project_id=project.id, + organization_id=project.organization_id, + task_id=str(uuid4()), + job_id=str(job.id), + task_instance=task, + vector_store_id="vs_123", + batch_number=1, + batch_doc_ids=[str(doc.id)], + remaining_batches=[], + ) + finally: + patcher.stop() + + task.retry.assert_not_called() + assert ( + CollectionJobCrud(db, project.id).read_one(job.id).status + == CollectionJobStatus.FAILED + ) + + +@patch("app.services.collections.create_collection.get_cloud_storage") +@patch("app.services.collections.create_collection.get_llm_provider") +def test_execute_batch_job_timeout_fails_once_retries_exhausted( + mock_get_provider: MagicMock, + mock_get_storage: MagicMock, + db: Session, +) -> None: + project = get_project(db) + store = DocumentStore(db=db, project_id=project.id) + doc = store.put() + + mock_provider = get_mock_provider("vs_123", "openai vector store") + mock_provider.create.side_effect = SoftTimeLimitExceeded() + mock_get_provider.return_value = mock_provider + + job = get_collection_job( + db, + project, + action_type=CollectionActionType.CREATE, + status=CollectionJobStatus.PROCESSING, + ) + request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) + task = _celery_task_instance(retries=3) + + patcher = _patch_session(db) + try: + with pytest.raises(SoftTimeLimitExceeded): + execute_batch_job( + request=request.model_dump(mode="json"), + project_id=project.id, + organization_id=project.organization_id, + task_id=str(uuid4()), + job_id=str(job.id), + task_instance=task, + vector_store_id="vs_123", + batch_number=1, + batch_doc_ids=[str(doc.id)], + remaining_batches=[], + ) + finally: + patcher.stop() + + task.retry.assert_not_called() + mock_provider.delete.assert_called_once() + + updated_job = CollectionJobCrud(db, project.id).read_one(job.id) + assert updated_job.status == CollectionJobStatus.FAILED + assert "soft time limit" in (updated_job.error_message or "") + + @patch("app.services.collections.create_collection.get_cloud_storage") @patch("app.services.collections.create_collection.send_callback") @patch("app.services.collections.create_collection.get_llm_provider") diff --git a/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md b/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md index 348fe8e7f..c09176279 100644 --- a/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md +++ b/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md @@ -289,20 +289,17 @@ heavy work is chunked: 1. Load the requested `document` rows; resolve the provider + project credentials. 2. Move the job to `PROCESSING`. -3. **Upload all files** to the provider via `provider.upload_files()` — this is - where the de-dup optimization lives: docs that already carry an - `openai_file_id` are skipped; new ones are uploaded and their file IDs - persisted (§7). -4. Compute `total_size_mb` and call `batch_documents()` to produce the batch plan. +3. Compute `total_size_mb` and call `batch_documents()` to produce the batch plan. +4. **Create the vector store** via `provider.create_vector_store()` and record it + on `result` immediately, so a later failure can always tear it down. 5. Persist batch metadata on the job (`total_batches`, `current_batch_number=0`, `documents_uploaded=[]`). -6. Enqueue **batch 1** with `vector_store_id=None` and the remaining batches as a - tail list. +6. Enqueue **batch 1** with the resolved `vector_store_id` and the remaining + batches as a tail list. -> Note the asymmetry: *file upload to the provider is not itself batched* — it -> happens for all docs in this single setup task. Only the **vector-store attach** -> is batched across Phase-2 tasks. See §11 for the residual timeout risk this -> leaves. +> Setup does **no** file uploading. Both the upload and the vector-store attach +> happen per batch in Phase 2, so each is bounded by the batch caps (≤200 docs / +> ≤30 MB). ### Phase 2 — `execute_batch_job` (one task per batch, self-chaining) @@ -311,16 +308,16 @@ heavy work is chunked: For each batch the task: 1. Resolves the provider and reads this batch's document rows. -2. Calls `provider.create(batch_docs, vector_store_id)`: - - On batch 1, `vector_store_id is None` → the vector store is **created**, then - this batch's file IDs are attached. - - On later batches, the existing `vector_store_id` (threaded through the task - args) is reused and the batch's file IDs are attached to it. -3. **Checkpoints** progress to the job row (`current_batch_number`, +2. Calls `provider.upload_files(storage, batch_docs, project_id)` — ships the + bytes of any doc that does not already carry a provider file ID (§7). +3. Calls `provider.create(batch_docs, vector_store_id)` — attaches this batch's + file IDs to the vector store created during setup, then waits for indexing + under a deadline (§12). +4. **Checkpoints** progress to the job row (`current_batch_number`, `documents_uploaded += this batch`). -4. If batches remain → enqueue the next batch task (passing the resolved +5. If batches remain → enqueue the next batch task (passing the resolved `vector_store_id` and the shrinking `remaining_batches` tail) and **return**. -5. On the **final** batch → finalize: +6. On the **final** batch → finalize: - Create the `Collection` row (`llm_service_id` = vector store ID, `llm_service_name`, provider, name, description). - Link every uploaded doc to it via `document_collection`. @@ -620,26 +617,41 @@ store, the remote vector store can be left orphaned: sends a failure callback. No `Collection` row is created (it only appears on the final batch), and no further batches are queued. - Cleanup is guarded by `if provider is not None and result is not None: - provider.delete(result)` in `_handle_job_failure`. But `result` is only set - **after** `provider.create()` returns. If `provider.create()` itself raises - (the common failure — e.g. an attach/parse error), `result` stays `None`, so - the **cleanup is skipped** and the vector store built by previous batches is - left dangling on the provider. + provider.delete(result)` in `_handle_job_failure`. **Resolved:** `result` is now + assigned as soon as the vector store exists — immediately after + `create_vector_store()` in setup, and at the top of the `try` in each batch task + — so a failure inside `upload_files()` or `create()` still tears the store down. - Uploaded provider **files persist** regardless (their IDs are saved on the `document` rows) — this is intentional for reuse, not a leak. -- **TODO:** track the resolved `vector_store_id` independently of `result` so a - mid-chain failure can always tear down the in-progress vector store. - -### 11.5 File upload to the provider is not batched - -- Batching protects the **vector-store attach** step, but `execute_setup_job` - uploads **all** new files to the provider in a single task before batching - begins. For a collection with many large *new* (not-yet-uploaded) documents, - that one setup task can still approach the soft time limit. -- The reuse optimization (§7) mitigates this for repeat documents, but a - first-time bulk upload is still a single-task operation. -- **TODO:** consider batching the upload phase too, or moving uploads into the - per-batch tasks. +- Residual gap: if `create_vector_store()` creates the store remotely but the + response is lost, the retry creates a second store and the first is orphaned. + +### 11.5 Upload and attach are both batched — timing bounds + +**Resolved:** uploads moved out of `execute_setup_job` into `execute_batch_job`, +so both phases are bounded by the batch caps. Measured worst case for a full +200-doc / 30 MB batch: **122 s at 300 ms RTT** (62 s at 150 ms, 22 s at 50 ms); +pure code overhead is 0.3 s, the rest is round-trip latency. That fits inside +`CELERY_TASK_SOFT_TIME_LIMIT = 300s`. + +The remaining unbounded call was the **indexing wait**: the SDK's +`file_batches.poll()` is a `while True` with no deadline, so a stuck batch could +outlast the soft limit and kill the task mid-write. `OpenAIVectorStoreCrud` now +polls with its own deadline instead (§12). + +Do **not** solve batch timing with a self-requeueing "continuation" task. Celery +runs with `task_acks_late=True` and `task_reject_on_worker_lost=True`, so a +*crashed* task is redelivered — but a task that returns normally and relies on a +freshly enqueued message has no such guarantee. If that message is lost the job +sits in `PROCESSING` forever, and `crud/job/pending_monitor.py` only counts +`status == PENDING` and only alerts, never recovers. + +The timeout retry (§12) shares that residual risk — `task_instance.retry()` +publishes a new message too — so it is bounded rather than unbounded: at most +`CELERY_TASK_MAX_RETRIES` attempts, and each one is recorded on the job's +`error_message` (`"Batch N timed out; retry k/3 queued"`) so a retry that never +lands is at least diagnosable from the job row. Widening `pending_monitor` to +watch stalled `PROCESSING` jobs is still open. ### 11.6 Two divergent delete paths @@ -658,8 +670,12 @@ store, the remote vector store can be left orphaned: |---|---| | **Async** | `POST /collections` returns `job_id` immediately; work runs on the Celery `low_priority` (priority 1) queue. | | **Batching** | One task per batch (≤30 MB / ≤200 docs), self-chaining; progress checkpointed to the job row across tasks. | -| **Timeout** | `SoftTimeLimitExceeded` → job `FAILED` ("Task exceeded soft time limit") + failure callback, then re-raised. | -| **Provider / attach error** | `OpenAIVectorStoreCrud.update` raises if any file fails to attach → job `FAILED`; see §11.4 for the orphan-cleanup gap. | +| **Timeout** | `execute_batch_job` **retries the batch** via `task_instance.retry()` while attempts remain (`CELERY_TASK_MAX_RETRIES = 3`, delay 180 s). The retry deliberately skips `_handle_job_failure`, so the vector store survives and the retried task reuses the same `vector_store_id`. Only once attempts are exhausted does the job go `FAILED` ("Task exceeded soft time limit") with a failure callback and vector-store teardown. `execute_setup_job` does **not** retry — a second `create_vector_store()` would orphan the first (§11.4). | +| **Why retry converges** | `upload_files` persists each document's provider file ID in its own session the moment that document uploads, so attempt N+1 skips everything attempt N finished. A 200-doc batch against a provider degraded to ~70 uploads/attempt completes on attempt 3 (measured). This is why a timeout is worth retrying rather than failing: each attempt does strictly less work. | +| **Indexing wait** | `OpenAIVectorStoreCrud._poll_file_batch` polls `file_batches.retrieve` every `BATCH_POLL_INTERVAL_SECONDS` (2 s) until the batch leaves `in_progress`, bounded by `BATCH_POLL_TIMEOUT_SECONDS` (240 s). Expiry raises `InterruptedError` tagged `(code: batch-poll-timeout)`. One constant, no coupling to the caller — overrunning the soft limit is survivable now that the batch retries. | +| **What `OPENAI_TIMEOUT_SECONDS` actually bounds** | A **stall**, not a duration. httpx applies its timeout per socket operation — there is no total-request deadline in httpx at all — so a large upload that streams steadily takes as long as it takes. Verified: a 10 MB upload running 12 s wall-clock under a 2 s write timeout **succeeds**; only a receiver that stops draining trips `WriteTimeout`. So 90 s means "90 consecutive seconds of zero bytes moving", i.e. a dead connection. It imposes no throughput floor on documents. Kept under a third of the soft limit so all attempts against one dead socket still fit in a task, letting the batch recover without spending a Celery retry. | +| **Transient provider errors** | Handled by the **OpenAI SDK's own retries** — `max_retries=OPENAI_MAX_RETRIES` (2) and `timeout=OPENAI_TIMEOUT_SECONDS` (90 s), both set in `providers/registry.py`. The timeout is what makes the ceiling provable: without it the SDK default is a 600 s read timeout, so one hung call is 3 × 600 = 1800 s — six times the soft limit, with zero progress made and nothing for a retry to build on. The SDK retries connection errors, 408, 409, 429 and 5xx, honours `retry-after`, and rewinds the upload stream between attempts. Deterministic errors (400/401/404/422) are not retried. Do **not** stack a second retry layer on top: two layers multiply attempts (3×3 = 9 requests, measured) and the outer layer's longer backoff dominates — a custom tenacity layer measured **4.5× slower** on an intermittent-failure batch (402 s vs 90 s projected over 200 docs). | +| **Provider / attach error** | `OpenAIVectorStoreCrud.update` raises if any file fails to attach → job `FAILED`; failed files are never retried, since an unsupported or corrupt file fails identically every time. | | **Upload then DB failure** | The just-uploaded provider file is deleted to avoid an orphan; the job fails. | | **Credentials missing** | `get_llm_provider` raises `ValueError` → clean job failure. | | **Callback delivery** | Best-effort; HTTPS-only, SSRF-validated, optional HMAC signing. A failed callback does not change job status (still pollable). | diff --git a/docs/wiki/modules/knowledge-base.md b/docs/wiki/modules/knowledge-base.md index 1582afa74..f43995399 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -36,3 +36,10 @@ All paths relative to `backend/app/`. - Uploads de-duplicate by provider file ID (see deep dive §7). - Collections are immutable-ish: deletion semantics in deep dive §10. - OpenAI file-batch id: the SDK's `file_batches.poll()` / `upload_and_poll()` final return deserializes a vector-store body, so its `.id` is the `vs_` id, not the `vsfb_` batch id. `crud/rag/open_ai.py` captures the batch id from `create()` before polling and uses it for `list_files`. Any failed file is a hard failure (whole vector store rolled back); partial indexing needs an add-documents endpoint first. +- The SDK's `file_batches.poll()` never times out. `_poll_file_batch` polls `retrieve` itself under a fixed `BATCH_POLL_TIMEOUT_SECONDS` deadline. It used to take that deadline from the caller via a `task_budget` contextvar; that was deleted once the batch retry landed, because overrunning the soft limit stopped being fatal. Don't reintroduce caller coupling here. +- Retries are the SDK's own (`max_retries=OPENAI_MAX_RETRIES`, `timeout=OPENAI_TIMEOUT_SECONDS` in `providers/registry.py`). It covers connection errors/408/409/429/5xx, honours `retry-after`, and rewinds the upload stream between attempts (verified). Do not add a second retry layer: stacking gives 3x3=9 requests per call, and a custom tenacity layer measured 4.5x slower than the SDK on an intermittent-failure batch. +- `OPENAI_TIMEOUT_SECONDS` (90s) is a **stall detector, not an upload deadline**. httpx has no total-request timeout — connect/read/write/pool are each per-socket-operation — so a big document that streams steadily uploads fine however long it takes; the timeout only fires after 90 consecutive seconds of zero bytes. Verified: a 10MB body taking 12s under a 2s write timeout succeeds, and only a stalled receiver raises `WriteTimeout`. Don't raise it out of fear that large files get cut off — they don't. Its actual job is bounding how long we sit on a *dead* socket (SDK default is 600s, so 3 attempts = 1800s of nothing). +- A batch that hits the Celery soft limit is **retried** (`task_instance.retry()`, max 3), not failed. The retry skips `_handle_job_failure` so the vector store survives and is reused. This converges because `upload_files` persists each doc's file ID as it uploads, so each attempt does strictly less work — measured: 200 docs against a provider managing ~70 uploads/attempt finishes on attempt 3. Setup does not retry (a second `create_vector_store()` orphans the first). Note `celery.exceptions.Retry` subclasses `Exception`; it escapes only because it is raised from inside an `except` clause, which sibling `except Exception` handlers do not catch. +- `documents_uploaded` is deduped on append — a batch retried past its checkpoint re-adds its own IDs, and `DocumentCrud.read_each` raises when duplicates collapse in its `IN` clause. +- Do not "fix" batch timing with a self-requeueing continuation task — a task that returns normally is not redelivered, and a lost continuation strands the job in `PROCESSING`, which nothing monitors or recovers. Deep dive §11.5. +- Uploads happen per batch in `execute_batch_job`, not in setup; setup only plans batches and creates the vector store. From 8966baa668b107befe2226e1c15b91b419e2a974 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:18:27 +0530 Subject: [PATCH 05/10] test: Enhance OpenAI vector store tests with batch timeout handling and document validation checks --- backend/app/tests/crud/rag/test_open_ai.py | 43 +++- .../tests/services/documents/test_helpers.py | 238 +++++++++++++++++- 2 files changed, 278 insertions(+), 3 deletions(-) diff --git a/backend/app/tests/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index 4d3d3de2b..b8ad62af7 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -4,12 +4,13 @@ OpenAI exception type). """ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import openai import pytest from app.crud.rag.open_ai import OpenAIVectorStoreCrud +from app.tests.utils.openai import get_mock_openai_client_with_vector_store @pytest.fixture @@ -295,3 +296,43 @@ class TestOpenAIVectorStoreCrudInit: def test_none_client_raises(self): with pytest.raises(ValueError): OpenAIVectorStoreCrud(client=None) + + +class TestPollFileBatchTimeout: + """When the batch never leaves 'in_progress', the poll loop must give up at the + deadline rather than spin forever like the SDK's own poll().""" + + def test_raises_interrupted_error_past_deadline(self, crud, mock_client): + with patch("app.crud.rag.open_ai.time") as mock_time: + # monotonic: deadline setup, first below-deadline check, its sleep arg, + # then a value past the deadline that trips the timeout. + mock_time.monotonic.side_effect = [0, 0, 0, 300] + mock_time.sleep = MagicMock() + mock_client.vector_stores.file_batches.retrieve.return_value = MagicMock( + status="in_progress", + file_counts=MagicMock(in_progress=1, completed=0), + ) + + with pytest.raises(InterruptedError) as exc_info: + crud._poll_file_batch("vsfb_x", "vs_1") + + assert "batch-poll-timeout" in str(exc_info.value) + assert mock_time.sleep.called + + +class TestGetMockOpenAIClientWithVectorStore: + """Contract test for the repo test fixture: exercises the whole helper and + pins the wiring the callers that use it depend on.""" + + def test_wiring_contract(self): + client = get_mock_openai_client_with_vector_store() + + assert client.vector_stores.create.return_value.id == "mock_vector_store_id" + + batch = client.vector_stores.file_batches.create.return_value + assert batch.id == "vsfb_mock" + assert batch.file_counts.failed == 0 + assert batch.file_counts.completed == 2 + assert client.vector_stores.file_batches.poll.return_value is batch + + assert client.beta.assistants.create.return_value.id == "mock_assistant_id" diff --git a/backend/app/tests/services/documents/test_helpers.py b/backend/app/tests/services/documents/test_helpers.py index 9a4b25344..6e4414410 100644 --- a/backend/app/tests/services/documents/test_helpers.py +++ b/backend/app/tests/services/documents/test_helpers.py @@ -1,8 +1,24 @@ +import csv from io import BytesIO -from fastapi import UploadFile +import pytest +from fastapi import HTTPException, UploadFile -from app.services.documents.helpers import calculate_file_size +from app.services.documents.helpers import ( + SNIFF_HEAD_BYTES, + SNIFF_TAIL_BYTES, + DocumentValidationError, + _check_csv, + _check_docx, + _check_json, + _check_pdf, + _check_xlsx, + _decode_utf8, + _read_edges, + calculate_file_size, + validate_document_content, + validate_upload, +) def make_upload_file(content: bytes, size: int | None = None) -> UploadFile: @@ -58,3 +74,221 @@ def test_empty_file_via_seek(self) -> None: """Returns 0 for an empty file when size is None.""" file = make_upload_file(b"", size=None) assert calculate_file_size(file) == 0 + + +class TestDocumentValidationErrorClientMessage: + def test_client_message_hides_internal_reason(self) -> None: + """The client-facing wording must never echo the internal `reason`, which + can describe byte offsets / binary sniffing details we don't expose.""" + err = DocumentValidationError("report.pdf", "NUL byte found at offset 12") + assert "report.pdf" in err.client_message + assert "NUL byte found at offset 12" not in err.client_message + assert "corrupted" in err.client_message + + +class TestDecodeUtf8: + def test_nul_byte_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _decode_utf8("f.csv", b"abc\x00def") + assert "NUL byte" in exc.value.reason + + def test_invalid_utf8_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _decode_utf8("f.csv", b"\xff\xfe") + assert "not valid UTF-8" in exc.value.reason + + def test_valid_utf8_returns_decoded_str(self) -> None: + assert _decode_utf8("f.csv", b"hello, world") == "hello, world" + + +class TestCheckPdf: + def test_missing_eof_marker_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_pdf("f.pdf", b"%PDF-1.4", b"no marker", make_upload_file(b"")) + assert "end-of-file marker" in exc.value.reason + + def test_encrypted_pdf_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_pdf( + "f.pdf", b"%PDF-1.4", b"/Encrypt 4 0 R %%EOF", make_upload_file(b"") + ) + assert "password protected" in exc.value.reason + + def test_valid_pdf_passes(self) -> None: + _check_pdf("f.pdf", b"%PDF-1.4", b"trailer\n%%EOF", make_upload_file(b"")) + + +class TestCheckOoxml: + def test_docx_missing_content_types_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_docx("f.docx", b"PK\x03\x04junk", b"", make_upload_file(b"")) + assert "no OOXML content-types part" in exc.value.reason + + def test_docx_that_is_really_xlsx_raises(self) -> None: + sample = b"[Content_Types].xml ... xl/worksheets/sheet1.xml" + with pytest.raises(DocumentValidationError) as exc: + _check_docx("f.docx", sample, b"", make_upload_file(b"")) + assert "not the format its extension claims" in exc.value.reason + + def test_valid_docx_passes(self) -> None: + sample = b"[Content_Types].xml ... word/document.xml" + _check_docx("f.docx", sample, b"", make_upload_file(b"")) + + def test_xlsx_that_is_really_docx_raises(self) -> None: + sample = b"[Content_Types].xml ... word/document.xml" + with pytest.raises(DocumentValidationError) as exc: + _check_xlsx("f.xlsx", sample, b"", make_upload_file(b"")) + assert "not the format its extension claims" in exc.value.reason + + def test_valid_xlsx_passes(self) -> None: + sample = b"[Content_Types].xml ... xl/worksheets/sheet1.xml" + _check_xlsx("f.xlsx", sample, b"", make_upload_file(b"")) + + +class TestCheckCsv: + def test_binary_signature_renamed_to_csv_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_csv("f.csv", b"%PDF-1.4 pretending", b"", make_upload_file(b"")) + assert "renamed to a CSV" in exc.value.reason + + def test_nul_byte_raises_via_decode(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_csv("f.csv", b"a,b\x00,c\n", b"", make_upload_file(b"")) + assert "NUL byte" in exc.value.reason + + def test_no_rows_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_csv("f.csv", b"", b"", make_upload_file(b"")) + assert "no readable CSV rows" in exc.value.reason + + def test_malformed_csv_raises(self) -> None: + """A single field larger than csv's field-size limit makes the reader + raise csv.Error, which surfaces as a validation failure.""" + head = b"x" * (csv.field_size_limit() + 1) + with pytest.raises(DocumentValidationError) as exc: + _check_csv("f.csv", head, b"", make_upload_file(b"")) + assert "malformed CSV" in exc.value.reason + + def test_inconsistent_column_count_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_csv("f.csv", b"a,b,c\n1,2\n", b"", make_upload_file(b"")) + assert "inconsistent column count at row 2" in exc.value.reason + + def test_clean_csv_passes(self) -> None: + _check_csv("f.csv", b"a,b,c\n1,2,3\n4,5,6\n", b"", make_upload_file(b"")) + + def test_full_head_drops_trailing_partial_line(self) -> None: + """At exactly SNIFF_HEAD_BYTES the final row is likely cut mid-line, so the + sniffer drops it instead of flagging a bogus column-count mismatch.""" + head = (b"a,b,c\n" * 700)[:SNIFF_HEAD_BYTES] + assert len(head) == SNIFF_HEAD_BYTES + _check_csv("f.csv", head, b"", make_upload_file(b"")) + + +class TestCheckJson: + def test_body_not_starting_with_bracket_raises(self) -> None: + content = b"plain text, not json" + with pytest.raises(DocumentValidationError) as exc: + _check_json("f.json", content, b"", make_upload_file(content)) + assert "must start with" in exc.value.reason + + def test_small_valid_json_passes(self) -> None: + content = b'{"a": 1, "b": [2, 3]}' + _check_json("f.json", content, content, make_upload_file(content)) + + def test_small_invalid_json_raises(self) -> None: + content = b'{"a": 1, oops}' + with pytest.raises(DocumentValidationError) as exc: + _check_json("f.json", content, content, make_upload_file(content)) + assert "invalid JSON" in exc.value.reason + + def test_large_json_closed_correctly_passes(self) -> None: + """Over the full-parse ceiling only the tail is checked for a matching + closer, so a properly closed large array passes without parsing.""" + content = b"[" + b" " * 300000 + b"]" + head = content[:SNIFF_HEAD_BYTES] + tail = content[-SNIFF_TAIL_BYTES:] + _check_json("f.json", head, tail, make_upload_file(content)) + + def test_large_json_not_closed_raises(self) -> None: + content = b"[" + b" " * 300000 + head = content[:SNIFF_HEAD_BYTES] + tail = content[-SNIFF_TAIL_BYTES:] + with pytest.raises(DocumentValidationError) as exc: + _check_json("f.json", head, tail, make_upload_file(content)) + assert "not closed correctly" in exc.value.reason + + +class TestReadEdges: + def test_head_only_when_tail_not_needed(self) -> None: + file = make_upload_file(b"hello world") + head, tail = _read_edges(file, needs_tail=False) + assert head == b"hello world" + assert tail == b"" + assert file.file.tell() == 0 + + def test_head_and_tail_when_needed(self) -> None: + content = b"START" + b"x" * 10000 + b"END" + file = make_upload_file(content) + head, tail = _read_edges(file, needs_tail=True) + assert head == content[:SNIFF_HEAD_BYTES] + assert tail == content[-SNIFF_TAIL_BYTES:] + assert file.file.tell() == 0 + + +class TestValidateDocumentContent: + def test_empty_file_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + validate_document_content(file=make_upload_file(b""), source_format="csv") + assert "the file is empty" in exc.value.reason + + def test_unknown_format_skips_validation(self) -> None: + """A format with no FormatSpec is accepted as-is, even binary junk.""" + result = validate_document_content( + file=make_upload_file(b"\x00\x01 arbitrary"), source_format="txt" + ) + assert result is None + + def test_signature_mismatch_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + validate_document_content( + file=make_upload_file(b"not a pdf %%EOF"), source_format="pdf" + ) + assert "does not match the pdf file signature" in exc.value.reason + + def test_valid_pdf_dispatches_to_checker_and_passes(self) -> None: + content = b"%PDF-1.4\n" + b"body " * 100 + b"trailer\n%%EOF\n" + validate_document_content( + file=make_upload_file(content), source_format="pdf" + ) + + +class TestValidateUpload: + def test_valid_csv_returns_format_and_transformer(self, monkeypatch) -> None: + monkeypatch.setattr( + "app.services.documents.helpers.pre_transform_validation", + lambda **kwargs: ("csv", None), + ) + result = validate_upload( + src=make_upload_file(b"a,b,c\n1,2,3\n"), + target_format=None, + transformer=None, + ) + assert result == ("csv", None) + + def test_invalid_content_raises_http_400_with_client_message( + self, monkeypatch + ) -> None: + monkeypatch.setattr( + "app.services.documents.helpers.pre_transform_validation", + lambda **kwargs: ("csv", None), + ) + with pytest.raises(HTTPException) as exc: + validate_upload( + src=make_upload_file(b"a,b\x00,c\n"), + target_format=None, + transformer=None, + ) + assert exc.value.status_code == 400 + assert "NUL byte" not in exc.value.detail + assert "corrupted" in exc.value.detail From a86653cabadfb95a7e0d1cc236fb9a85c2418e03 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:18:41 +0530 Subject: [PATCH 06/10] refactor(tests): Simplify validate_document_content call in PDF validation test --- backend/app/tests/services/documents/test_helpers.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/backend/app/tests/services/documents/test_helpers.py b/backend/app/tests/services/documents/test_helpers.py index 6e4414410..960875c28 100644 --- a/backend/app/tests/services/documents/test_helpers.py +++ b/backend/app/tests/services/documents/test_helpers.py @@ -258,9 +258,7 @@ def test_signature_mismatch_raises(self) -> None: def test_valid_pdf_dispatches_to_checker_and_passes(self) -> None: content = b"%PDF-1.4\n" + b"body " * 100 + b"trailer\n%%EOF\n" - validate_document_content( - file=make_upload_file(content), source_format="pdf" - ) + validate_document_content(file=make_upload_file(content), source_format="pdf") class TestValidateUpload: From 50ce4c9f8651d65ad02eb7f5116b2b03c8410c8e Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 30 Jul 2026 15:05:01 +0530 Subject: [PATCH 07/10] fix(openai): Adjust OpenAI timeout settings and improve batch polling logic for better error handling --- backend/app/crud/rag/open_ai.py | 46 +++++++-------- backend/app/services/documents/helpers.py | 47 ++++++++++++++- backend/app/tests/crud/rag/test_open_ai.py | 59 ++++++++++++------- .../tests/services/documents/test_helpers.py | 35 +++++++++++ backend/app/tests/utils/openai.py | 4 +- .../kaapi-knowledge-base-ARCHITECTURE.md | 10 ++-- docs/wiki/modules/knowledge-base.md | 4 +- 7 files changed, 149 insertions(+), 56 deletions(-) diff --git a/backend/app/crud/rag/open_ai.py b/backend/app/crud/rag/open_ai.py index 57a8409dd..cd105890d 100644 --- a/backend/app/crud/rag/open_ai.py +++ b/backend/app/crud/rag/open_ai.py @@ -17,12 +17,10 @@ OPENAI_MAX_RETRIES = 2 -# Per socket operation, not per call: a steady upload runs as long as it needs and -# this only fires on a dead connection. Sized so all attempts fit one Celery task. -OPENAI_TIMEOUT_SECONDS = 90 +# Under the Celery soft time limit so a hung call plus its retries can't eat the window. +OPENAI_TIMEOUT_SECONDS = 30 BATCH_POLL_INTERVAL_SECONDS = 2 -BATCH_POLL_TIMEOUT_SECONDS = 240 def vs_ls(client: OpenAI, vector_store_id: str): @@ -132,32 +130,12 @@ def _retrieve_file_batch( def _poll_file_batch( self, batch_id: str, vector_store_id: str ) -> VectorStoreFileBatch: - """Wait for indexing to finish. The SDK's own poll() loops forever.""" - deadline = time.monotonic() + BATCH_POLL_TIMEOUT_SECONDS - + """Poll until indexing finishes; the Celery soft time limit is the deadline.""" while True: batch = self._retrieve_file_batch(batch_id, vector_store_id) if batch.status != "in_progress": return batch - - if time.monotonic() >= deadline: - error_message = ( - f"[KAAPI] Timed out (code: batch-poll-timeout) after " - f"{BATCH_POLL_TIMEOUT_SECONDS}s waiting for OpenAI to finish indexing " - f"{batch.file_counts.in_progress} file(s). Retry the " - f"collection, or split it into smaller batches." - ) - logger.error( - f"[OpenAIVectorStoreCrud._poll_file_batch] {error_message} | " - f"vector_store_id={vector_store_id}, batch_id={batch_id}, " - f"completed={batch.file_counts.completed}, " - f"in_progress={batch.file_counts.in_progress}" - ) - raise InterruptedError(error_message) - - time.sleep( - min(BATCH_POLL_INTERVAL_SECONDS, max(0.0, deadline - time.monotonic())) - ) + time.sleep(BATCH_POLL_INTERVAL_SECONDS) def update( self, @@ -301,6 +279,22 @@ def update( ) raise + # Only 'completed' is success; 'cancelled'/'failed' with no per-file failures + # would otherwise slip through the failed-count check above. + if batch.status != "completed": + error_message = ( + f"[OPENAI] Vector store indexing did not complete " + f"(status: {batch.status}). Retry the collection." + ) + logger.error( + f"[OpenAIVectorStoreCrud.update] {error_message} | " + f"vector_store_id={vector_store_id}, batch_id={batch_id}, " + f"status={batch.status}, completed={batch.file_counts.completed}, " + f"failed={batch.file_counts.failed}, " + f"cancelled={batch.file_counts.cancelled}" + ) + raise RuntimeError(error_message) + def delete(self, vector_store_id: str, retries: int = 3): if retries < 1: try: diff --git a/backend/app/services/documents/helpers.py b/backend/app/services/documents/helpers.py index 27ed2a800..a200c21d9 100644 --- a/backend/app/services/documents/helpers.py +++ b/backend/app/services/documents/helpers.py @@ -49,6 +49,13 @@ OLE2_SIGNATURE = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" +# The OLE2 signature is a generic Compound File Binary container shared by .xls and +# .doc, so it can't tell them apart. The CFB directory stores per-stream names as +# UTF-16LE: an .xls carries a Workbook (BIFF8) or Book (BIFF5) stream, a .doc carries +# WordDocument. +OLE2_EXCEL_STREAMS = ("Workbook".encode("utf-16-le"), "Book".encode("utf-16-le")) +OLE2_WORD_STREAMS = ("WordDocument".encode("utf-16-le"),) + MISLABELLED_BINARY_SIGNATURES: dict[bytes, str] = { b"PK\x03\x04": "an Office/zip file (xlsx, docx)", OLE2_SIGNATURE: "a legacy Office file (xls, doc)", @@ -138,6 +145,38 @@ def _check_xlsx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> No _check_ooxml(filename, head + tail, OOXML_EXCEL_PART) +def _check_ole2_stream( + filename: str, + sample: bytes, + wanted: tuple[bytes, ...], + rivals: tuple[bytes, ...], + rival_label: str, +) -> None: + """Reject an OLE2 file whose extension claims one format but whose CFB directory + holds the rival's stream. Inconclusive samples pass — the directory can sit past + the sampled edges, and rejecting a valid file is worse than the loose check.""" + if any(name in sample for name in wanted): + return + if any(name in sample for name in rivals): + raise DocumentValidationError( + filename, + f"parsing error - file is a {rival_label} document, " + "not the format its extension claims", + ) + + +def _check_xls(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ole2_stream( + filename, head + tail, OLE2_EXCEL_STREAMS, OLE2_WORD_STREAMS, "Word (.doc)" + ) + + +def _check_doc(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ole2_stream( + filename, head + tail, OLE2_WORD_STREAMS, OLE2_EXCEL_STREAMS, "Excel (.xls)" + ) + + def _check_csv(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: for signature, description in MISLABELLED_BINARY_SIGNATURES.items(): if head.startswith(signature): @@ -237,8 +276,12 @@ class FormatSpec: needs_tail=True, checker=_check_xlsx, ), - "xls": FormatSpec(signatures=(OLE2_SIGNATURE,)), - "doc": FormatSpec(signatures=(OLE2_SIGNATURE,)), + "xls": FormatSpec( + signatures=(OLE2_SIGNATURE,), needs_tail=True, checker=_check_xls + ), + "doc": FormatSpec( + signatures=(OLE2_SIGNATURE,), needs_tail=True, checker=_check_doc + ), "csv": FormatSpec(checker=_check_csv), "json": FormatSpec(needs_tail=True, checker=_check_json), } diff --git a/backend/app/tests/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index b8ad62af7..03e3a97dd 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -9,7 +9,7 @@ import openai import pytest -from app.crud.rag.open_ai import OpenAIVectorStoreCrud +from app.crud.rag.open_ai import BATCH_POLL_INTERVAL_SECONDS, OpenAIVectorStoreCrud from app.tests.utils.openai import get_mock_openai_client_with_vector_store @@ -298,33 +298,51 @@ def test_none_client_raises(self): OpenAIVectorStoreCrud(client=None) -class TestPollFileBatchTimeout: - """When the batch never leaves 'in_progress', the poll loop must give up at the - deadline rather than spin forever like the SDK's own poll().""" +class TestPollFileBatch: + """Poll loop sleeps while in_progress and returns once the batch settles. A batch + that never finishes is bounded by the Celery soft time limit, not an internal one. + """ + + def test_polls_until_not_in_progress( + self, crud: OpenAIVectorStoreCrud, mock_client: MagicMock + ) -> None: + done = MagicMock(status="completed") + mock_client.vector_stores.file_batches.retrieve.side_effect = [ + MagicMock(status="in_progress"), + done, + ] + + with patch("app.crud.rag.open_ai.time.sleep") as mock_sleep: + result = crud._poll_file_batch("vsfb_x", "vs_1") + + assert result is done + mock_sleep.assert_called_once_with(BATCH_POLL_INTERVAL_SECONDS) - def test_raises_interrupted_error_past_deadline(self, crud, mock_client): - with patch("app.crud.rag.open_ai.time") as mock_time: - # monotonic: deadline setup, first below-deadline check, its sleep arg, - # then a value past the deadline that trips the timeout. - mock_time.monotonic.side_effect = [0, 0, 0, 300] - mock_time.sleep = MagicMock() - mock_client.vector_stores.file_batches.retrieve.return_value = MagicMock( - status="in_progress", - file_counts=MagicMock(in_progress=1, completed=0), - ) - with pytest.raises(InterruptedError) as exc_info: - crud._poll_file_batch("vsfb_x", "vs_1") +class TestUpdateTerminalStatus: + """A batch ending 'cancelled'/'failed' with no per-file failures must raise, not + slip through the failed-count check as success.""" - assert "batch-poll-timeout" in str(exc_info.value) - assert mock_time.sleep.called + def test_cancelled_batch_raises( + self, crud: OpenAIVectorStoreCrud, mock_client: MagicMock, docs: list[MagicMock] + ) -> None: + mock_client.vector_stores.file_batches.create.return_value = MagicMock( + id=REAL_BATCH_ID + ) + counts = MagicMock(completed=0, failed=0, cancelled=2, in_progress=0) + mock_client.vector_stores.file_batches.retrieve.return_value = MagicMock( + status="cancelled", file_counts=counts + ) + + with pytest.raises(RuntimeError, match="cancelled"): + crud.update("vs_1", docs) class TestGetMockOpenAIClientWithVectorStore: """Contract test for the repo test fixture: exercises the whole helper and pins the wiring the callers that use it depend on.""" - def test_wiring_contract(self): + def test_wiring_contract(self) -> None: client = get_mock_openai_client_with_vector_store() assert client.vector_stores.create.return_value.id == "mock_vector_store_id" @@ -333,6 +351,7 @@ def test_wiring_contract(self): assert batch.id == "vsfb_mock" assert batch.file_counts.failed == 0 assert batch.file_counts.completed == 2 - assert client.vector_stores.file_batches.poll.return_value is batch + # _poll_file_batch polls retrieve, not poll — pin the endpoint production uses. + assert client.vector_stores.file_batches.retrieve.return_value is batch assert client.beta.assistants.create.return_value.id == "mock_assistant_id" diff --git a/backend/app/tests/services/documents/test_helpers.py b/backend/app/tests/services/documents/test_helpers.py index 960875c28..524cb5926 100644 --- a/backend/app/tests/services/documents/test_helpers.py +++ b/backend/app/tests/services/documents/test_helpers.py @@ -9,9 +9,11 @@ SNIFF_TAIL_BYTES, DocumentValidationError, _check_csv, + _check_doc, _check_docx, _check_json, _check_pdf, + _check_xls, _check_xlsx, _decode_utf8, _read_edges, @@ -145,6 +147,39 @@ def test_valid_xlsx_passes(self) -> None: _check_xlsx("f.xlsx", sample, b"", make_upload_file(b"")) +class TestCheckOle2: + """OLE2 is a generic container shared by .xls and .doc; the checker distinguishes + them by the UTF-16LE stream name in the CFB directory.""" + + WORKBOOK = "Workbook".encode("utf-16-le") + WORD = "WordDocument".encode("utf-16-le") + + def test_doc_mislabelled_as_xls_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_xls("f.xls", self.WORD, b"", make_upload_file(b"")) + assert "not the format its extension claims" in exc.value.reason + + def test_valid_xls_passes(self) -> None: + _check_xls("f.xls", self.WORKBOOK, b"", make_upload_file(b"")) + + def test_xls_mislabelled_as_doc_raises(self) -> None: + with pytest.raises(DocumentValidationError) as exc: + _check_doc("f.doc", self.WORKBOOK, b"", make_upload_file(b"")) + assert "not the format its extension claims" in exc.value.reason + + def test_valid_doc_passes(self) -> None: + _check_doc("f.doc", self.WORD, b"", make_upload_file(b"")) + + def test_inconclusive_sample_passes(self) -> None: + """Directory can sit past the sampled edges; don't reject on absence.""" + _check_xls( + "f.xls", + b"\xd0\xcf\x11\xe0 no stream names here", + b"", + make_upload_file(b""), + ) + + class TestCheckCsv: def test_binary_signature_renamed_to_csv_raises(self) -> None: with pytest.raises(DocumentValidationError) as exc: diff --git a/backend/app/tests/utils/openai.py b/backend/app/tests/utils/openai.py index 093f45ae6..7f680057b 100644 --- a/backend/app/tests/utils/openai.py +++ b/backend/app/tests/utils/openai.py @@ -119,11 +119,13 @@ def get_mock_openai_client_with_vector_store() -> MagicMock: # File batch creation + polling mock_file_batch = MagicMock() mock_file_batch.id = "vsfb_mock" + mock_file_batch.status = "completed" mock_file_batch.file_counts.completed = 2 mock_file_batch.file_counts.total = 2 mock_file_batch.file_counts.failed = 0 mock_client.vector_stores.file_batches.create.return_value = mock_file_batch - mock_client.vector_stores.file_batches.poll.return_value = mock_file_batch + # _poll_file_batch polls retrieve; wire it so callers of the real update() path work. + mock_client.vector_stores.file_batches.retrieve.return_value = mock_file_batch # File list mock_client.vector_stores.files.list.return_value = {"data": []} diff --git a/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md b/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md index c09176279..8837e985a 100644 --- a/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md +++ b/docs/architecture/kaapi-knowledge-base-ARCHITECTURE.md @@ -255,16 +255,16 @@ sequenceDiagram BR->>SU: deliver setup task SU->>DB: read documents · resolve provider + credentials · job → PROCESSING - SU->>Prov: upload_files() — upload EVERY not-yet-uploaded doc, persist file IDs - SU->>DB: persist openai_file_id per doc SU->>SU: batch_documents() → plan N batches - SU->>DB: job ← total_size_mb, total_batches=N, current_batch_number=0, documents_uploaded=[] - SU->>BR: enqueue batch 1 (vector_store_id=None, remaining_batches=[2..N]) + SU->>Prov: create_vector_store() + SU->>DB: job ← total_size_mb, total_batches=N, current_batch_number=0, documents_uploaded=[], knowledge_base_id + SU->>BR: enqueue batch 1 (vector_store_id, remaining_batches=[2..N]) loop each batch k = 1..N BR->>BA: deliver batch task (batch k, vector_store_id) + BA->>Prov: upload_files(batch k docs) — persist file IDs BA->>Prov: provider.create(batch_docs, vector_store_id) - Note over Prov: k=1 creates the vector store, every batch attaches its file IDs via file_batches.upload_and_poll + Note over Prov: every batch attaches its file IDs to the vector store via file_batches BA->>DB: checkpoint: current_batch_number=k, documents_uploaded += batch k alt batches remain BA->>BR: enqueue batch k+1 (vector_store_id threaded through) diff --git a/docs/wiki/modules/knowledge-base.md b/docs/wiki/modules/knowledge-base.md index f43995399..312d16ccd 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -36,9 +36,9 @@ All paths relative to `backend/app/`. - Uploads de-duplicate by provider file ID (see deep dive §7). - Collections are immutable-ish: deletion semantics in deep dive §10. - OpenAI file-batch id: the SDK's `file_batches.poll()` / `upload_and_poll()` final return deserializes a vector-store body, so its `.id` is the `vs_` id, not the `vsfb_` batch id. `crud/rag/open_ai.py` captures the batch id from `create()` before polling and uses it for `list_files`. Any failed file is a hard failure (whole vector store rolled back); partial indexing needs an add-documents endpoint first. -- The SDK's `file_batches.poll()` never times out. `_poll_file_batch` polls `retrieve` itself under a fixed `BATCH_POLL_TIMEOUT_SECONDS` deadline. It used to take that deadline from the caller via a `task_budget` contextvar; that was deleted once the batch retry landed, because overrunning the soft limit stopped being fatal. Don't reintroduce caller coupling here. +- The SDK's `file_batches.poll()` never times out. `_poll_file_batch` polls `retrieve` in a loop with no internal deadline — the Celery soft time limit bounds it, and its `SoftTimeLimitExceeded` fails/retries the task. An earlier version took a deadline from the caller via a `task_budget` `ContextVar` (and a fixed `BATCH_POLL_TIMEOUT_SECONDS`); both were deleted once the batch retry landed, because overrunning the soft limit stopped being fatal. Don't reintroduce caller coupling here. - Retries are the SDK's own (`max_retries=OPENAI_MAX_RETRIES`, `timeout=OPENAI_TIMEOUT_SECONDS` in `providers/registry.py`). It covers connection errors/408/409/429/5xx, honours `retry-after`, and rewinds the upload stream between attempts (verified). Do not add a second retry layer: stacking gives 3x3=9 requests per call, and a custom tenacity layer measured 4.5x slower than the SDK on an intermittent-failure batch. -- `OPENAI_TIMEOUT_SECONDS` (90s) is a **stall detector, not an upload deadline**. httpx has no total-request timeout — connect/read/write/pool are each per-socket-operation — so a big document that streams steadily uploads fine however long it takes; the timeout only fires after 90 consecutive seconds of zero bytes. Verified: a 10MB body taking 12s under a 2s write timeout succeeds, and only a stalled receiver raises `WriteTimeout`. Don't raise it out of fear that large files get cut off — they don't. Its actual job is bounding how long we sit on a *dead* socket (SDK default is 600s, so 3 attempts = 1800s of nothing). +- `OPENAI_TIMEOUT_SECONDS` (30s) is a **stall detector, not an upload deadline**. httpx has no total-request timeout — connect/read/write/pool are each per-socket-operation — so a big document that streams steadily uploads fine however long it takes; the timeout only fires after 30 consecutive seconds of zero bytes. Kept well under the Celery soft limit so a hung call plus its SDK retries can't consume the whole task window. Verified: a 10MB body taking 12s under a 2s write timeout succeeds, and only a stalled receiver raises `WriteTimeout`. Don't raise it out of fear that large files get cut off — they don't. Its actual job is bounding how long we sit on a *dead* socket (SDK default is 600s, so 3 attempts = 1800s of nothing). - A batch that hits the Celery soft limit is **retried** (`task_instance.retry()`, max 3), not failed. The retry skips `_handle_job_failure` so the vector store survives and is reused. This converges because `upload_files` persists each doc's file ID as it uploads, so each attempt does strictly less work — measured: 200 docs against a provider managing ~70 uploads/attempt finishes on attempt 3. Setup does not retry (a second `create_vector_store()` orphans the first). Note `celery.exceptions.Retry` subclasses `Exception`; it escapes only because it is raised from inside an `except` clause, which sibling `except Exception` handlers do not catch. - `documents_uploaded` is deduped on append — a batch retried past its checkpoint re-adds its own IDs, and `DocumentCrud.read_each` raises when duplicates collapse in its `IN` clause. - Do not "fix" batch timing with a self-requeueing continuation task — a task that returns normally is not redelivered, and a lost continuation strands the job in `PROCESSING`, which nothing monitors or recovers. Deep dive §11.5. From 9fdb45c43d8ac4d4881e1f086c7f40023816667a Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 30 Jul 2026 15:21:01 +0530 Subject: [PATCH 08/10] fix(collections): Remove unused batch retry logging and update test assertions for job status --- backend/app/crud/rag/open_ai.py | 8 ++-- .../services/collections/create_collection.py | 37 ------------------- .../collections/test_create_collection.py | 1 - 3 files changed, 3 insertions(+), 43 deletions(-) diff --git a/backend/app/crud/rag/open_ai.py b/backend/app/crud/rag/open_ai.py index cd105890d..70fd0472c 100644 --- a/backend/app/crud/rag/open_ai.py +++ b/backend/app/crud/rag/open_ai.py @@ -279,8 +279,8 @@ def update( ) raise - # Only 'completed' is success; 'cancelled'/'failed' with no per-file failures - # would otherwise slip through the failed-count check above. + # Only 'completed' is success; a 'cancelled'/'failed' batch with no per-file + # failures slips past the failed-count check above. if batch.status != "completed": error_message = ( f"[OPENAI] Vector store indexing did not complete " @@ -289,9 +289,7 @@ def update( logger.error( f"[OpenAIVectorStoreCrud.update] {error_message} | " f"vector_store_id={vector_store_id}, batch_id={batch_id}, " - f"status={batch.status}, completed={batch.file_counts.completed}, " - f"failed={batch.file_counts.failed}, " - f"cancelled={batch.file_counts.cancelled}" + f"status={batch.status}" ) raise RuntimeError(error_message) diff --git a/backend/app/services/collections/create_collection.py b/backend/app/services/collections/create_collection.py index f810fe4d2..56f379b2d 100644 --- a/backend/app/services/collections/create_collection.py +++ b/backend/app/services/collections/create_collection.py @@ -52,36 +52,6 @@ def _can_retry_batch(task_instance) -> bool: return task_instance.request.retries < task_instance.max_retries -def _note_batch_retry( - project_id: int, - job_id: str, - batch_number: int, - attempt: int, - max_attempts: int, -) -> None: - """Record a pending retry on the job row. - - The job stays PROCESSING between attempts and pending_monitor only watches - PENDING, so a retry that never lands would be invisible. - """ - try: - with Session(engine) as session: - CollectionJobCrud(session, project_id).update( - UUID(job_id), - CollectionJobUpdate( - error_message=( - f"Batch {batch_number} timed out; " - f"retry {attempt}/{max_attempts} queued" - ) - ), - ) - except Exception: - logger.warning( - "[create_collection._note_batch_retry] Could not record retry | job_id=%s", - job_id, - ) - - def start_job( db: Session, request: CreationRequest, @@ -567,13 +537,6 @@ def execute_batch_job( task_instance.max_retries, ) span.set_attribute("collection.batch_retry", attempt) - _note_batch_retry( - project_id, - job_id, - batch_number, - attempt, - task_instance.max_retries, - ) try: # Only Retry means it reached the broker. retry() also raises # Reject on a failed publish, or exc outside a worker - those diff --git a/backend/app/tests/services/collections/test_create_collection.py b/backend/app/tests/services/collections/test_create_collection.py index 112de9bcb..fc2c7f6e1 100644 --- a/backend/app/tests/services/collections/test_create_collection.py +++ b/backend/app/tests/services/collections/test_create_collection.py @@ -761,7 +761,6 @@ def test_execute_batch_job_timeout_retries_and_keeps_vector_store( updated_job = CollectionJobCrud(db, project.id).read_one(job.id) assert updated_job.status == CollectionJobStatus.PROCESSING - assert "retry 1/3 queued" in (updated_job.error_message or "") @pytest.mark.parametrize( From 80c7180e6685010570ed9061a498536b2919eb41 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:48:49 +0530 Subject: [PATCH 09/10] fix(openai): Implement tenacity for batch create+index retries and improve error handling --- backend/app/crud/rag/open_ai.py | 138 +++++++---- .../services/collections/create_collection.py | 46 +--- .../collections/providers/registry.py | 5 +- backend/app/tests/crud/rag/test_open_ai.py | 77 +++++- .../providers/test_openai_provider.py | 11 + .../collections/test_create_collection.py | 222 +----------------- docs/wiki/modules/knowledge-base.md | 10 +- 7 files changed, 189 insertions(+), 320 deletions(-) diff --git a/backend/app/crud/rag/open_ai.py b/backend/app/crud/rag/open_ai.py index 70fd0472c..fdc2f8d56 100644 --- a/backend/app/crud/rag/open_ai.py +++ b/backend/app/crud/rag/open_ai.py @@ -8,6 +8,13 @@ from openai.types import VectorStore from openai.types.vector_stores import VectorStoreFileBatch from pydantic import BaseModel +from tenacity import ( + RetryCallState, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) from app.models import Document, ProviderType @@ -15,13 +22,26 @@ OPENAI_PROVIDER = ProviderType.openai.value -OPENAI_MAX_RETRIES = 2 - -# Under the Celery soft time limit so a hung call plus its retries can't eat the window. +# Under the Celery soft time limit so a hung call can't eat the whole task window. +# SDK-level retries are off (registry.py: max_retries=0); tenacity is the sole +# retry layer, wrapping batch create+index. OPENAI_TIMEOUT_SECONDS = 30 BATCH_POLL_INTERVAL_SECONDS = 2 +# Retry batch create+index on any OpenAI/indexing failure, exponential backoff +# (~2s, 4s, 8s), all inside one Celery soft-time-limit window. +BATCH_INDEX_MAX_ATTEMPTS = 4 +BATCH_RETRY_BACKOFF_BASE_SECONDS = 2 + + +def _log_batch_retry(retry_state: RetryCallState) -> None: + logger.warning( + f"[OpenAIVectorStoreCrud._create_and_index_batch] Batch attempt failed, retrying | " + f"attempt={retry_state.attempt_number}, " + f"error={retry_state.outcome.exception() if retry_state.outcome else None}" + ) + def vs_ls(client: OpenAI, vector_store_id: str): kwargs = {} @@ -137,6 +157,74 @@ def _poll_file_batch( return batch time.sleep(BATCH_POLL_INTERVAL_SECONDS) + def _raise_if_batch_incomplete( + self, + batch: VectorStoreFileBatch, + batch_id: str, + vector_store_id: str, + docs: list[Document], + ) -> None: + """Raise on any indexing failure so the batch attempt is retried.""" + if batch.file_counts.failed > 0: + try: + failed_files = self.client.vector_stores.file_batches.list_files( + vector_store_id=vector_store_id, + batch_id=batch_id, + filter="failed", + ) + doc_by_file_id = {d.file_id[OPENAI_PROVIDER]: d for d in docs} + parts = [] + for f in failed_files: + d = doc_by_file_id.get(f.id) + label = d.fname if d else f.id + msg = f.last_error.message if f.last_error else "no error detail" + parts.append(f"{label}: {msg}") + logger.error( + f"[OpenAIVectorStoreCrud._raise_if_batch_incomplete] Files failed to index | " + f"{{'batch_id': '{batch_id}', 'failed_files': '{', '.join(parts)}'}}" + ) + raise RuntimeError("; ".join(parts)) + except OpenAIError as err: + logger.warning( + f"[OpenAIVectorStoreCrud._raise_if_batch_incomplete] Could not fetch per-file errors | " + f"{{'batch_id': '{batch_id}', 'error': '{str(err)}'}}" + ) + raise + + # Only 'completed' is success; a 'cancelled'/'failed' batch with no per-file + # failures slips past the failed-count check above. + if batch.status != "completed": + error_message = ( + f"[OPENAI] Vector store indexing did not complete " + f"(status: {batch.status}). Retry the collection." + ) + logger.error( + f"[OpenAIVectorStoreCrud._raise_if_batch_incomplete] {error_message} | " + f"vector_store_id={vector_store_id}, batch_id={batch_id}, " + f"status={batch.status}" + ) + raise RuntimeError(error_message) + + @retry( + reraise=True, + stop=stop_after_attempt(BATCH_INDEX_MAX_ATTEMPTS), + wait=wait_exponential(multiplier=BATCH_RETRY_BACKOFF_BASE_SECONDS), + retry=retry_if_exception_type((OpenAIError, RuntimeError)), + before_sleep=_log_batch_retry, + ) + def _create_and_index_batch( + self, vector_store_id: str, docs: list[Document] + ) -> tuple[VectorStoreFileBatch, str]: + """Create the file batch, wait for indexing, verify it completed. Retried as + a unit on any OpenAI/indexing failure; SoftTimeLimitExceeded is not retried + (not an OpenAIError/RuntimeError) so it aborts the task inside the window.""" + batch_id = self._create_file_batch( + vector_store_id, [doc.file_id[OPENAI_PROVIDER] for doc in docs] + ) + batch = self._poll_file_batch(batch_id, vector_store_id) + self._raise_if_batch_incomplete(batch, batch_id, vector_store_id, docs) + return batch, batch_id + def update( self, vector_store_id: str, @@ -151,10 +239,7 @@ def update( ) try: - batch_id = self._create_file_batch( - vector_store_id, [doc.file_id[OPENAI_PROVIDER] for doc in docs] - ) - batch = self._poll_file_batch(batch_id, vector_store_id) + batch, batch_id = self._create_and_index_batch(vector_store_id, docs) except openai.RateLimitError as e: error_message = ( f"[OPENAI] Rate limit exceeded (code: {e.status_code}): " @@ -253,45 +338,6 @@ def update( f"{{'vector_store_id': '{vector_store_id}', 'batch_id': '{batch_id}', " f"'completed': {batch.file_counts.completed}, 'failed': {batch.file_counts.failed}}}" ) - if batch.file_counts.failed > 0: - try: - failed_files = self.client.vector_stores.file_batches.list_files( - vector_store_id=vector_store_id, - batch_id=batch_id, - filter="failed", - ) - doc_by_file_id = {d.file_id[OPENAI_PROVIDER]: d for d in docs} - parts = [] - for f in failed_files: - d = doc_by_file_id.get(f.id) - label = d.fname if d else f.id - msg = f.last_error.message if f.last_error else "no error detail" - parts.append(f"{label}: {msg}") - logger.error( - f"[OpenAIVectorStoreCrud.update] Files failed to index | " - f"{{'batch_id': '{batch_id}', 'failed_files': '{', '.join(parts)}'}}" - ) - raise RuntimeError("; ".join(parts)) - except OpenAIError as err: - logger.warning( - f"[OpenAIVectorStoreCrud.update] Could not fetch per-file errors | " - f"{{'batch_id': '{batch_id}', 'error': '{str(err)}'}}" - ) - raise - - # Only 'completed' is success; a 'cancelled'/'failed' batch with no per-file - # failures slips past the failed-count check above. - if batch.status != "completed": - error_message = ( - f"[OPENAI] Vector store indexing did not complete " - f"(status: {batch.status}). Retry the collection." - ) - logger.error( - f"[OpenAIVectorStoreCrud.update] {error_message} | " - f"vector_store_id={vector_store_id}, batch_id={batch_id}, " - f"status={batch.status}" - ) - raise RuntimeError(error_message) def delete(self, vector_store_id: str, retries: int = 3): if retries < 1: diff --git a/backend/app/services/collections/create_collection.py b/backend/app/services/collections/create_collection.py index 56f379b2d..4dd8edf61 100644 --- a/backend/app/services/collections/create_collection.py +++ b/backend/app/services/collections/create_collection.py @@ -4,7 +4,7 @@ from sqlmodel import Session from gevent import Timeout -from celery.exceptions import Retry, SoftTimeLimitExceeded +from celery.exceptions import SoftTimeLimitExceeded from opentelemetry import trace from asgi_correlation_id import correlation_id @@ -40,18 +40,6 @@ tracer = trace.get_tracer(__name__) -def _can_retry_batch(task_instance) -> bool: - """Whether a timed-out batch has attempts left. - - Worth retrying because upload_files persists each file ID as it uploads, so - every attempt resumes where the last stopped. `max_retries=None` means retry - forever, which would never fail the job or fire its callback; treat as none. - """ - if task_instance is None or task_instance.max_retries is None: - return False - return task_instance.request.retries < task_instance.max_retries - - def start_job( db: Session, request: CreationRequest, @@ -523,35 +511,11 @@ def execute_batch_job( webhook_secret=webhook_secret, ) - except (Timeout, SoftTimeLimitExceeded) as err: + except (Timeout, SoftTimeLimitExceeded): + # Batch-level retries happen in-task (tenacity in OpenAIVectorStoreCrud), + # so hitting the soft time limit means the window is spent — fail, don't + # re-queue. timeout_err = TimeoutError("Task exceeded soft time limit") - - if _can_retry_batch(task_instance): - attempt = task_instance.request.retries + 1 - logger.warning( - "[create_collection.execute_batch_job] Batch timed out, retrying | " - "{'collection_job_id': '%s', 'batch_number': %d, 'attempt': '%d/%d'}", - job_id, - batch_number, - attempt, - task_instance.max_retries, - ) - span.set_attribute("collection.batch_retry", attempt) - try: - # Only Retry means it reached the broker. retry() also raises - # Reject on a failed publish, or exc outside a worker - those - # must fail the job, not strand it in PROCESSING. - raise task_instance.retry(exc=timeout_err) - except Retry: - raise - except Exception: - logger.exception( - "[create_collection.execute_batch_job] Retry could not be scheduled, " - "failing the job | {'collection_job_id': '%s', 'batch_number': %d}", - job_id, - batch_number, - ) - logger.warning( "[create_collection.execute_batch_job] Collection Creation Timed Out | {'collection_job_id': '%s', 'error': '%s'}", job_id, diff --git a/backend/app/services/collections/providers/registry.py b/backend/app/services/collections/providers/registry.py index a4f9ed2da..1d23050f9 100644 --- a/backend/app/services/collections/providers/registry.py +++ b/backend/app/services/collections/providers/registry.py @@ -5,7 +5,7 @@ from openai import OpenAI from app.crud import get_provider_credential -from app.crud.rag.open_ai import OPENAI_MAX_RETRIES, OPENAI_TIMEOUT_SECONDS +from app.crud.rag.open_ai import OPENAI_TIMEOUT_SECONDS from app.services.collections.providers.base import BaseProvider from app.services.collections.providers.gemini import GeminiAIStudioProvider from app.services.collections.providers.openai import OpenAIProvider @@ -66,7 +66,8 @@ def get_llm_provider( raise ValueError("OpenAI credentials not configured for this project.") client = OpenAI( api_key=credentials["api_key"], - max_retries=OPENAI_MAX_RETRIES, + # SDK retries off; tenacity in OpenAIVectorStoreCrud is the sole retry layer. + max_retries=0, timeout=OPENAI_TIMEOUT_SECONDS, ) elif provider == LLMProvider.GOOGLE_AISTUDIO: diff --git a/backend/app/tests/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index 03e3a97dd..38c40940d 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -8,8 +8,13 @@ import openai import pytest +from tenacity import stop_after_attempt, wait_none -from app.crud.rag.open_ai import BATCH_POLL_INTERVAL_SECONDS, OpenAIVectorStoreCrud +from app.crud.rag.open_ai import ( + BATCH_INDEX_MAX_ATTEMPTS, + BATCH_POLL_INTERVAL_SECONDS, + OpenAIVectorStoreCrud, +) from app.tests.utils.openai import get_mock_openai_client_with_vector_store @@ -23,6 +28,17 @@ def crud(mock_client): return OpenAIVectorStoreCrud(client=mock_client) +@pytest.fixture(autouse=True) +def _single_attempt_batch_retry(): + """_create_and_index_batch is tenacity-retried; default tests to a single + attempt with no backoff. TestBatchRetry re-enables retries explicitly.""" + retrying = OpenAIVectorStoreCrud._create_and_index_batch.retry + stop, wait = retrying.stop, retrying.wait + retrying.stop, retrying.wait = stop_after_attempt(1), wait_none() + yield + retrying.stop, retrying.wait = stop, wait + + def _make_doc(file_id: str, fname: str) -> MagicMock: doc = MagicMock() doc.file_id = {"openai": file_id} @@ -146,15 +162,15 @@ def test_falls_back_to_file_id_label_for_unknown_file( with pytest.raises(RuntimeError, match="file-unknown: parse error"): crud.update("vs_1", docs) - def test_reraises_when_list_files_errors(self, crud, mock_client, docs): - """If the follow-up list_files lookup itself raises, the OpenAI error - propagates instead of masking the real upload problem.""" + def test_surfaces_when_list_files_errors(self, crud, mock_client, docs): + """If the follow-up list_files lookup itself raises, the OpenAI error is + surfaced (mapped to InterruptedError) instead of being masked as success.""" _wire_batch(mock_client, completed=0, failed=2) mock_client.vector_stores.file_batches.list_files.side_effect = ( openai.OpenAIError("list failed") ) - with pytest.raises(openai.OpenAIError, match="list failed"): + with pytest.raises(InterruptedError, match="list failed"): crud.update("vs_1", docs) @@ -355,3 +371,54 @@ def test_wiring_contract(self) -> None: assert client.vector_stores.file_batches.retrieve.return_value is batch assert client.beta.assistants.create.return_value.id == "mock_assistant_id" + + +class TestBatchRetry: + """_create_and_index_batch retries the whole create+poll+validate unit on any + OpenAI/indexing failure, up to BATCH_INDEX_MAX_ATTEMPTS (tenacity).""" + + @staticmethod + def _enable_retries() -> None: + retrying = OpenAIVectorStoreCrud._create_and_index_batch.retry + retrying.stop = stop_after_attempt(BATCH_INDEX_MAX_ATTEMPTS) + retrying.wait = wait_none() + + def test_retries_then_succeeds( + self, crud: OpenAIVectorStoreCrud, mock_client: MagicMock + ) -> None: + self._enable_retries() + mock_client.vector_stores.file_batches.create.return_value = MagicMock( + id=REAL_BATCH_ID + ) + completed = MagicMock( + status="completed", + file_counts=MagicMock(completed=1, failed=0, in_progress=0), + ) + mock_client.vector_stores.file_batches.retrieve.side_effect = [ + openai.APIConnectionError(request=MagicMock()), + completed, + ] + + crud.update("vs_1", [_make_doc("file-1", "f1.pdf")]) + + assert mock_client.vector_stores.file_batches.create.call_count == 2 + + def test_gives_up_after_max_attempts( + self, crud: OpenAIVectorStoreCrud, mock_client: MagicMock + ) -> None: + self._enable_retries() + mock_client.vector_stores.file_batches.create.return_value = MagicMock( + id=REAL_BATCH_ID + ) + mock_client.vector_stores.file_batches.retrieve.return_value = MagicMock( + status="cancelled", + file_counts=MagicMock(completed=0, failed=0, cancelled=1, in_progress=0), + ) + + with pytest.raises(RuntimeError): + crud.update("vs_1", [_make_doc("file-1", "f1.pdf")]) + + assert ( + mock_client.vector_stores.file_batches.create.call_count + == BATCH_INDEX_MAX_ATTEMPTS + ) diff --git a/backend/app/tests/services/collections/providers/test_openai_provider.py b/backend/app/tests/services/collections/providers/test_openai_provider.py index 17e5e8306..d31fb0ae3 100644 --- a/backend/app/tests/services/collections/providers/test_openai_provider.py +++ b/backend/app/tests/services/collections/providers/test_openai_provider.py @@ -5,6 +5,7 @@ import pytest from openai import OpenAIError +from tenacity import stop_after_attempt, wait_none from app.crud.rag.open_ai import OpenAIVectorStoreCrud from app.services.collections.providers.openai import OpenAIProvider @@ -16,6 +17,16 @@ ) +@pytest.fixture(autouse=True) +def _single_attempt_batch_retry(): + """update() retries batch create+index via tenacity; run one attempt, no backoff.""" + retrying = OpenAIVectorStoreCrud._create_and_index_batch.retry + stop, wait = retrying.stop, retrying.wait + retrying.stop, retrying.wait = stop_after_attempt(1), wait_none() + yield + retrying.stop, retrying.wait = stop, wait + + def test_create_vector_store_returns_id() -> None: client = get_mock_openai_client_with_vector_store() provider = OpenAIProvider(client=client) diff --git a/backend/app/tests/services/collections/test_create_collection.py b/backend/app/tests/services/collections/test_create_collection.py index fc2c7f6e1..e2b6b3475 100644 --- a/backend/app/tests/services/collections/test_create_collection.py +++ b/backend/app/tests/services/collections/test_create_collection.py @@ -4,7 +4,7 @@ import uuid from uuid import UUID, uuid4 -from celery.exceptions import Reject, Retry, SoftTimeLimitExceeded +from celery.exceptions import SoftTimeLimitExceeded from gevent import Timeout import pytest @@ -703,226 +703,6 @@ def test_execute_batch_job_soft_time_limit_marks_failed_and_reraises( assert "soft time limit" in (updated_job.error_message or "") -def _celery_task_instance(retries: int, max_retries: int | None = 3) -> MagicMock: - """Bound-task stand-in; production always passes `task_instance=self`.""" - task = MagicMock() - task.request.retries = retries - task.max_retries = max_retries - task.retry.side_effect = Retry("requeued") - return task - - -@pytest.mark.parametrize("timeout_exc", [Timeout(300), SoftTimeLimitExceeded()]) -@patch("app.services.collections.create_collection.get_cloud_storage") -@patch("app.services.collections.create_collection.get_llm_provider") -def test_execute_batch_job_timeout_retries_and_keeps_vector_store( - mock_get_provider: MagicMock, - mock_get_storage: MagicMock, - timeout_exc: BaseException, - db: Session, -) -> None: - project = get_project(db) - store = DocumentStore(db=db, project_id=project.id) - doc = store.put() - - mock_provider = get_mock_provider("vs_123", "openai vector store") - mock_provider.create.side_effect = timeout_exc - mock_get_provider.return_value = mock_provider - - job = get_collection_job( - db, - project, - action_type=CollectionActionType.CREATE, - status=CollectionJobStatus.PROCESSING, - ) - request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) - task = _celery_task_instance(retries=0) - - patcher = _patch_session(db) - try: - with pytest.raises(Retry): - execute_batch_job( - request=request.model_dump(mode="json"), - project_id=project.id, - organization_id=project.organization_id, - task_id=str(uuid4()), - job_id=str(job.id), - task_instance=task, - vector_store_id="vs_123", - batch_number=1, - batch_doc_ids=[str(doc.id)], - remaining_batches=[], - ) - finally: - patcher.stop() - - task.retry.assert_called_once() - mock_provider.delete.assert_not_called() - - updated_job = CollectionJobCrud(db, project.id).read_one(job.id) - assert updated_job.status == CollectionJobStatus.PROCESSING - - -@pytest.mark.parametrize( - "retry_effect", - [ - pytest.param(Reject(OSError("broker down"), False), id="broker_publish_fails"), - pytest.param(TimeoutError("soft limit"), id="called_outside_worker"), - ], -) -@patch("app.services.collections.create_collection.get_cloud_storage") -@patch("app.services.collections.create_collection.get_llm_provider") -def test_execute_batch_job_unschedulable_retry_fails_job( - mock_get_provider: MagicMock, - mock_get_storage: MagicMock, - retry_effect: BaseException, - db: Session, -) -> None: - """Reject (publish failed) and a bare exc (called outside a worker) both mean - nothing was queued, so the job must fail rather than strand in PROCESSING.""" - project = get_project(db) - store = DocumentStore(db=db, project_id=project.id) - doc = store.put() - - mock_provider = get_mock_provider("vs_123", "openai vector store") - mock_provider.create.side_effect = SoftTimeLimitExceeded() - mock_get_provider.return_value = mock_provider - - job = get_collection_job( - db, - project, - action_type=CollectionActionType.CREATE, - status=CollectionJobStatus.PROCESSING, - ) - request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) - task = _celery_task_instance(retries=0) - task.retry.side_effect = retry_effect - - patcher = _patch_session(db) - try: - with pytest.raises(SoftTimeLimitExceeded): - execute_batch_job( - request=request.model_dump(mode="json"), - project_id=project.id, - organization_id=project.organization_id, - task_id=str(uuid4()), - job_id=str(job.id), - task_instance=task, - vector_store_id="vs_123", - batch_number=1, - batch_doc_ids=[str(doc.id)], - remaining_batches=[], - ) - finally: - patcher.stop() - - mock_provider.delete.assert_called_once() - - updated_job = CollectionJobCrud(db, project.id).read_one(job.id) - assert updated_job.status == CollectionJobStatus.FAILED - - -@patch("app.services.collections.create_collection.get_cloud_storage") -@patch("app.services.collections.create_collection.get_llm_provider") -def test_execute_batch_job_infinite_max_retries_does_not_crash( - mock_get_provider: MagicMock, - mock_get_storage: MagicMock, - db: Session, -) -> None: - """`max_retries=None` is retry-forever; comparing against it raises TypeError, - which would skip failure handling entirely.""" - project = get_project(db) - store = DocumentStore(db=db, project_id=project.id) - doc = store.put() - - mock_provider = get_mock_provider("vs_123", "openai vector store") - mock_provider.create.side_effect = SoftTimeLimitExceeded() - mock_get_provider.return_value = mock_provider - - job = get_collection_job( - db, - project, - action_type=CollectionActionType.CREATE, - status=CollectionJobStatus.PROCESSING, - ) - request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) - task = _celery_task_instance(retries=0, max_retries=None) - - patcher = _patch_session(db) - try: - with pytest.raises(SoftTimeLimitExceeded): - execute_batch_job( - request=request.model_dump(mode="json"), - project_id=project.id, - organization_id=project.organization_id, - task_id=str(uuid4()), - job_id=str(job.id), - task_instance=task, - vector_store_id="vs_123", - batch_number=1, - batch_doc_ids=[str(doc.id)], - remaining_batches=[], - ) - finally: - patcher.stop() - - task.retry.assert_not_called() - assert ( - CollectionJobCrud(db, project.id).read_one(job.id).status - == CollectionJobStatus.FAILED - ) - - -@patch("app.services.collections.create_collection.get_cloud_storage") -@patch("app.services.collections.create_collection.get_llm_provider") -def test_execute_batch_job_timeout_fails_once_retries_exhausted( - mock_get_provider: MagicMock, - mock_get_storage: MagicMock, - db: Session, -) -> None: - project = get_project(db) - store = DocumentStore(db=db, project_id=project.id) - doc = store.put() - - mock_provider = get_mock_provider("vs_123", "openai vector store") - mock_provider.create.side_effect = SoftTimeLimitExceeded() - mock_get_provider.return_value = mock_provider - - job = get_collection_job( - db, - project, - action_type=CollectionActionType.CREATE, - status=CollectionJobStatus.PROCESSING, - ) - request = CreationRequest(documents=[doc.id], provider="openai", callback_url=None) - task = _celery_task_instance(retries=3) - - patcher = _patch_session(db) - try: - with pytest.raises(SoftTimeLimitExceeded): - execute_batch_job( - request=request.model_dump(mode="json"), - project_id=project.id, - organization_id=project.organization_id, - task_id=str(uuid4()), - job_id=str(job.id), - task_instance=task, - vector_store_id="vs_123", - batch_number=1, - batch_doc_ids=[str(doc.id)], - remaining_batches=[], - ) - finally: - patcher.stop() - - task.retry.assert_not_called() - mock_provider.delete.assert_called_once() - - updated_job = CollectionJobCrud(db, project.id).read_one(job.id) - assert updated_job.status == CollectionJobStatus.FAILED - assert "soft time limit" in (updated_job.error_message or "") - - @patch("app.services.collections.create_collection.get_cloud_storage") @patch("app.services.collections.create_collection.send_callback") @patch("app.services.collections.create_collection.get_llm_provider") diff --git a/docs/wiki/modules/knowledge-base.md b/docs/wiki/modules/knowledge-base.md index 312d16ccd..c6b8462df 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -36,10 +36,10 @@ All paths relative to `backend/app/`. - Uploads de-duplicate by provider file ID (see deep dive §7). - Collections are immutable-ish: deletion semantics in deep dive §10. - OpenAI file-batch id: the SDK's `file_batches.poll()` / `upload_and_poll()` final return deserializes a vector-store body, so its `.id` is the `vs_` id, not the `vsfb_` batch id. `crud/rag/open_ai.py` captures the batch id from `create()` before polling and uses it for `list_files`. Any failed file is a hard failure (whole vector store rolled back); partial indexing needs an add-documents endpoint first. -- The SDK's `file_batches.poll()` never times out. `_poll_file_batch` polls `retrieve` in a loop with no internal deadline — the Celery soft time limit bounds it, and its `SoftTimeLimitExceeded` fails/retries the task. An earlier version took a deadline from the caller via a `task_budget` `ContextVar` (and a fixed `BATCH_POLL_TIMEOUT_SECONDS`); both were deleted once the batch retry landed, because overrunning the soft limit stopped being fatal. Don't reintroduce caller coupling here. -- Retries are the SDK's own (`max_retries=OPENAI_MAX_RETRIES`, `timeout=OPENAI_TIMEOUT_SECONDS` in `providers/registry.py`). It covers connection errors/408/409/429/5xx, honours `retry-after`, and rewinds the upload stream between attempts (verified). Do not add a second retry layer: stacking gives 3x3=9 requests per call, and a custom tenacity layer measured 4.5x slower than the SDK on an intermittent-failure batch. -- `OPENAI_TIMEOUT_SECONDS` (30s) is a **stall detector, not an upload deadline**. httpx has no total-request timeout — connect/read/write/pool are each per-socket-operation — so a big document that streams steadily uploads fine however long it takes; the timeout only fires after 30 consecutive seconds of zero bytes. Kept well under the Celery soft limit so a hung call plus its SDK retries can't consume the whole task window. Verified: a 10MB body taking 12s under a 2s write timeout succeeds, and only a stalled receiver raises `WriteTimeout`. Don't raise it out of fear that large files get cut off — they don't. Its actual job is bounding how long we sit on a *dead* socket (SDK default is 600s, so 3 attempts = 1800s of nothing). -- A batch that hits the Celery soft limit is **retried** (`task_instance.retry()`, max 3), not failed. The retry skips `_handle_job_failure` so the vector store survives and is reused. This converges because `upload_files` persists each doc's file ID as it uploads, so each attempt does strictly less work — measured: 200 docs against a provider managing ~70 uploads/attempt finishes on attempt 3. Setup does not retry (a second `create_vector_store()` orphans the first). Note `celery.exceptions.Retry` subclasses `Exception`; it escapes only because it is raised from inside an `except` clause, which sibling `except Exception` handlers do not catch. -- `documents_uploaded` is deduped on append — a batch retried past its checkpoint re-adds its own IDs, and `DocumentCrud.read_each` raises when duplicates collapse in its `IN` clause. +- The SDK's `file_batches.poll()` never times out. `_poll_file_batch` polls `retrieve` in a loop with no internal deadline — the Celery soft time limit bounds it, and its `SoftTimeLimitExceeded` aborts the task. An earlier version took a deadline from the caller via a `task_budget` `ContextVar` (and a fixed `BATCH_POLL_TIMEOUT_SECONDS`); both were deleted. Don't reintroduce caller coupling here. +- Retries are a **single tenacity layer** wrapping `_create_and_index_batch` (create + poll + validate) in `crud/rag/open_ai.py`: `stop_after_attempt(BATCH_INDEX_MAX_ATTEMPTS)` (3 retries) with exponential backoff (~2s/4s/8s), retrying on `OpenAIError` or `RuntimeError` (indexing error, failed files, or non-`completed` status). SDK-level retries are **off** (`max_retries=0` in `providers/registry.py`) so nothing stacks. `SoftTimeLimitExceeded` is deliberately *not* retried (neither `OpenAIError` nor `RuntimeError`), so a spent window aborts immediately. History: a prior tenacity attempt that *stacked on top of* SDK retries measured 4.5x slower (3×3=9 requests/call) — the fix was to make tenacity the sole layer, not to drop it. `upload_files` runs once per task (outside the retry); only the create+attach+index is retried, so retries re-attach already-uploaded file IDs rather than re-uploading. +- `OPENAI_TIMEOUT_SECONDS` (30s) is a **stall detector, not an upload deadline**. httpx has no total-request timeout — connect/read/write/pool are each per-socket-operation — so a big document that streams steadily uploads fine however long it takes; the timeout only fires after 30 consecutive seconds of zero bytes. With SDK retries off, one hung call is capped near 30s, and tenacity's backoff keeps the whole batch inside the soft-limit window. Verified: a 10MB body taking 12s under a 2s write timeout succeeds, and only a stalled receiver raises `WriteTimeout`. Don't raise it out of fear that large files get cut off — they don't. +- A batch failure (indexing error, failed files, cancelled/failed status, timeout) is retried **in-task** by the tenacity layer above — 3 retries, exponential backoff, all inside one Celery soft-time-limit window. There is **no** Celery-level re-queue: if the batch (or its retries) can't finish within the window, `SoftTimeLimitExceeded` fires and the job is marked FAILED (`_handle_job_failure`), not re-queued. Trade-off: a collection whose batch genuinely needs more than one window fails rather than resuming across windows — size batches to fit. Setup does not retry either (a second `create_vector_store()` orphans the first). +- `documents_uploaded` is deduped on append — a task re-run (e.g. redelivered after worker loss, since `acks_late` is on) re-adds its own IDs, and `DocumentCrud.read_each` raises when duplicates collapse in its `IN` clause. - Do not "fix" batch timing with a self-requeueing continuation task — a task that returns normally is not redelivered, and a lost continuation strands the job in `PROCESSING`, which nothing monitors or recovers. Deep dive §11.5. - Uploads happen per batch in `execute_batch_job`, not in setup; setup only plans batches and creates the vector store. From 09b8196b7cf6b2c114a0e4841d3a7be910cdd951 Mon Sep 17 00:00:00 2001 From: Prashant Vasudevan <71649489+vprashrex@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:59:48 +0530 Subject: [PATCH 10/10] feat(documents): Add document validation logic and helper functions for file uploads --- .../services/collections/create_collection.py | 2 - .../collections/providers/registry.py | 1 - backend/app/services/documents/helpers.py | 314 +---------------- backend/app/services/documents/validator.py | 315 ++++++++++++++++++ .../tests/services/documents/test_helpers.py | 6 +- 5 files changed, 324 insertions(+), 314 deletions(-) create mode 100644 backend/app/services/documents/validator.py diff --git a/backend/app/services/collections/create_collection.py b/backend/app/services/collections/create_collection.py index 4dd8edf61..4e1b79b9c 100644 --- a/backend/app/services/collections/create_collection.py +++ b/backend/app/services/collections/create_collection.py @@ -403,8 +403,6 @@ def execute_batch_job( collection_job_crud = CollectionJobCrud(session, project_id) collection_job = collection_job_crud.read_one(job_uuid) already_uploaded = collection_job.documents_uploaded or [] - # A re-run past this point re-appends its own IDs, and read_each - # rejects a list whose duplicates collapse in its IN. now_uploaded = list( dict.fromkeys( already_uploaded + [str(d) for d in all_doc_ids_this_batch] diff --git a/backend/app/services/collections/providers/registry.py b/backend/app/services/collections/providers/registry.py index 1d23050f9..cf1729d6b 100644 --- a/backend/app/services/collections/providers/registry.py +++ b/backend/app/services/collections/providers/registry.py @@ -66,7 +66,6 @@ def get_llm_provider( raise ValueError("OpenAI credentials not configured for this project.") client = OpenAI( api_key=credentials["api_key"], - # SDK retries off; tenacity in OpenAIVectorStoreCrud is the sole retry layer. max_retries=0, timeout=OPENAI_TIMEOUT_SECONDS, ) diff --git a/backend/app/services/documents/helpers.py b/backend/app/services/documents/helpers.py index a200c21d9..c8cb77ea0 100644 --- a/backend/app/services/documents/helpers.py +++ b/backend/app/services/documents/helpers.py @@ -1,9 +1,5 @@ -import csv -import json import logging -from codecs import getincrementaldecoder -from dataclasses import dataclass -from typing import Callable, Optional, Tuple, Iterable, Union +from typing import Optional, Tuple, Iterable, Union from uuid import UUID from fastapi import HTTPException, UploadFile @@ -16,6 +12,10 @@ ) from app.crud import DocTransformationJobCrud, DocumentCrud from app.services.doctransform import job as transformation_job +from app.services.documents.validator import ( + DocumentValidationError, + validate_document_content, +) from app.models import ( DocTransformJobCreate, TransformationStatus, @@ -30,310 +30,6 @@ logger = logging.getLogger(__name__) -SNIFF_HEAD_BYTES = 4096 -SNIFF_TAIL_BYTES = 2048 - -PDF_EOF_MARKER = b"%%EOF" -PDF_ENCRYPT_MARKER = b"/Encrypt" -OOXML_CONTENT_TYPES = b"[Content_Types].xml" -OOXML_WORD_PART = b"word/" -OOXML_EXCEL_PART = b"xl/" -JSON_OPENERS = (b"{", b"[") -JSON_CLOSERS: dict[int, int] = {ord("{"): ord("}"), ord("["): ord("]")} -UTF8_BOM = b"\xef\xbb\xbf" - -JSON_FULL_PARSE_MAX_BYTES = 256 * 1024 - -CSV_STRICT_COLUMN_COUNT = True -CSV_SAMPLE_MAX_ROWS = 200 - -OLE2_SIGNATURE = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" - -# The OLE2 signature is a generic Compound File Binary container shared by .xls and -# .doc, so it can't tell them apart. The CFB directory stores per-stream names as -# UTF-16LE: an .xls carries a Workbook (BIFF8) or Book (BIFF5) stream, a .doc carries -# WordDocument. -OLE2_EXCEL_STREAMS = ("Workbook".encode("utf-16-le"), "Book".encode("utf-16-le")) -OLE2_WORD_STREAMS = ("WordDocument".encode("utf-16-le"),) - -MISLABELLED_BINARY_SIGNATURES: dict[bytes, str] = { - b"PK\x03\x04": "an Office/zip file (xlsx, docx)", - OLE2_SIGNATURE: "a legacy Office file (xls, doc)", - b"%PDF-": "a PDF", -} - -UNNAMED_FILE_PLACEHOLDER = "" - - -CLIENT_VALIDATION_MESSAGE = ( - "Document '{filename}' failed parsing. The document appears to be corrupted " - "or invalid. Please re-upload a valid document." -) - - -class DocumentValidationError(Exception): - """Raised when a document fails the pre-upload sanity check.""" - - def __init__(self, filename: str, reason: str) -> None: - self.filename = filename - self.reason = reason - super().__init__(f"Document '{filename}' failed validation: {reason}") - - @property - def client_message(self) -> str: - """Client-safe wording. `reason` stays internal - it is log-only.""" - return CLIENT_VALIDATION_MESSAGE.format(filename=self.filename) - - -def _decode_utf8(filename: str, sample: bytes) -> str: - if b"\x00" in sample: - raise DocumentValidationError( - filename, - "parsing error - NUL byte found in a text-based file; " - "the file is binary, corrupt, or has the wrong extension", - ) - try: - return getincrementaldecoder("utf-8")().decode(sample) - except UnicodeDecodeError as e: - raise DocumentValidationError( - filename, - f"parsing error - content is not valid UTF-8 at byte offset {e.start}; " - "re-export the file as UTF-8", - ) from e - - -def _check_pdf(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - if PDF_EOF_MARKER not in tail: - raise DocumentValidationError( - filename, - "parsing error - PDF is missing its end-of-file marker; " - "the upload is truncated or incomplete", - ) - if PDF_ENCRYPT_MARKER in tail: - raise DocumentValidationError( - filename, - "parsing error - PDF is password protected or encrypted and cannot be read", - ) - - -def _check_ooxml(filename: str, sample: bytes, expected_part: bytes) -> None: - if OOXML_CONTENT_TYPES not in sample: - raise DocumentValidationError( - filename, - "parsing error - file is a zip archive but not a valid Office document " - "(no OOXML content-types part)", - ) - - other_part = ( - OOXML_EXCEL_PART if expected_part == OOXML_WORD_PART else OOXML_WORD_PART - ) - if other_part in sample and expected_part not in sample: - raise DocumentValidationError( - filename, - f"parsing error - file is a '{other_part.decode()}' Office package, " - "not the format its extension claims", - ) - - -def _check_docx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - """Head+tail: the zip central directory lists all entry names and sits at the - end; the head holds only the first entry.""" - _check_ooxml(filename, head + tail, OOXML_WORD_PART) - - -def _check_xlsx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - _check_ooxml(filename, head + tail, OOXML_EXCEL_PART) - - -def _check_ole2_stream( - filename: str, - sample: bytes, - wanted: tuple[bytes, ...], - rivals: tuple[bytes, ...], - rival_label: str, -) -> None: - """Reject an OLE2 file whose extension claims one format but whose CFB directory - holds the rival's stream. Inconclusive samples pass — the directory can sit past - the sampled edges, and rejecting a valid file is worse than the loose check.""" - if any(name in sample for name in wanted): - return - if any(name in sample for name in rivals): - raise DocumentValidationError( - filename, - f"parsing error - file is a {rival_label} document, " - "not the format its extension claims", - ) - - -def _check_xls(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - _check_ole2_stream( - filename, head + tail, OLE2_EXCEL_STREAMS, OLE2_WORD_STREAMS, "Word (.doc)" - ) - - -def _check_doc(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - _check_ole2_stream( - filename, head + tail, OLE2_WORD_STREAMS, OLE2_EXCEL_STREAMS, "Excel (.xls)" - ) - - -def _check_csv(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - for signature, description in MISLABELLED_BINARY_SIGNATURES.items(): - if head.startswith(signature): - raise DocumentValidationError( - filename, - f"parsing error - file is {description} renamed to a CSV, not text; " - "export it as CSV instead of renaming it", - ) - - text = _decode_utf8(filename, head) - - lines = text.splitlines() - if len(head) == SNIFF_HEAD_BYTES and len(lines) > 1: - lines = lines[:-1] - - try: - rows = [row for row in csv.reader(lines[:CSV_SAMPLE_MAX_ROWS]) if row] - except csv.Error as e: - raise DocumentValidationError( - filename, f"parsing error - malformed CSV: {e}" - ) from e - - if not rows: - raise DocumentValidationError( - filename, "parsing error - no readable CSV rows found" - ) - - if CSV_STRICT_COLUMN_COUNT: - expected_columns = len(rows[0]) - for line_number, row in enumerate(rows[1:], start=2): - if len(row) != expected_columns: - raise DocumentValidationError( - filename, - f"parsing error - inconsistent column count at row {line_number} " - f"({len(row)} fields, expected {expected_columns})", - ) - - -def _check_json_tail(filename: str, opener: int, tail: bytes) -> None: - """Truncation check for JSON too large to parse: the opening bracket must be - matched by the last non-whitespace byte.""" - closing = tail.rstrip() - if not closing or closing[-1] != JSON_CLOSERS[opener]: - raise DocumentValidationError( - filename, - "parsing error - JSON is not closed correctly; " - "the upload is truncated or incomplete", - ) - - -def _check_json(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: - body = head.lstrip(UTF8_BOM).lstrip() - if not body.startswith(JSON_OPENERS): - raise DocumentValidationError( - filename, "parsing error - JSON document must start with '{' or '['" - ) - - stream = file.file - stream.seek(0, 2) - size_bytes = stream.tell() - stream.seek(0) - - if size_bytes > JSON_FULL_PARSE_MAX_BYTES: - logger.info( - f"[_check_json] Too large for a full parse, checking closure only | " - f"filename: {filename} | size_bytes: {size_bytes}" - ) - _check_json_tail(filename, body[0], tail) - return - - try: - json.loads(stream.read(), object_pairs_hook=lambda pairs: None) - except (json.JSONDecodeError, UnicodeDecodeError) as e: - raise DocumentValidationError( - filename, f"parsing error - invalid JSON: {e}" - ) from e - finally: - stream.seek(0) - - -@dataclass(frozen=True) -class FormatSpec: - signatures: tuple[bytes, ...] = () - needs_tail: bool = False - checker: Callable[[str, bytes, bytes, UploadFile], None] | None = None - - -FORMAT_SPECS: dict[str, FormatSpec] = { - "pdf": FormatSpec(signatures=(b"%PDF-",), needs_tail=True, checker=_check_pdf), - "docx": FormatSpec( - signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), - needs_tail=True, - checker=_check_docx, - ), - "xlsx": FormatSpec( - signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), - needs_tail=True, - checker=_check_xlsx, - ), - "xls": FormatSpec( - signatures=(OLE2_SIGNATURE,), needs_tail=True, checker=_check_xls - ), - "doc": FormatSpec( - signatures=(OLE2_SIGNATURE,), needs_tail=True, checker=_check_doc - ), - "csv": FormatSpec(checker=_check_csv), - "json": FormatSpec(needs_tail=True, checker=_check_json), -} - - -def _read_edges(file: UploadFile, needs_tail: bool) -> Tuple[bytes, bytes]: - """Sample the head (and optionally tail) of the upload, leaving it rewound.""" - stream = file.file - - stream.seek(0) - head = stream.read(SNIFF_HEAD_BYTES) - - tail = b"" - if needs_tail: - stream.seek(0, 2) - stream.seek(max(0, stream.tell() - SNIFF_TAIL_BYTES)) - tail = stream.read(SNIFF_TAIL_BYTES) - - stream.seek(0) - return head, tail - - -def validate_document_content(*, file: UploadFile, source_format: str) -> None: - """ - Sanity-check an upload's bytes before it is persisted. Only the first - SNIFF_HEAD_BYTES (plus SNIFF_TAIL_BYTES for PDF) are read, so the cost is - constant in file size - small JSON, which is fully parsed, is the exception. - - Raises: - DocumentValidationError: naming the offending file and the reason. - """ - filename = (file.filename or "").strip() or UNNAMED_FILE_PLACEHOLDER - - spec = FORMAT_SPECS.get(source_format) - head, tail = _read_edges(file, needs_tail=bool(spec and spec.needs_tail)) - - if not head: - raise DocumentValidationError(filename, "the file is empty") - - if spec is None: - return - - if spec.signatures and not head.startswith(spec.signatures): - raise DocumentValidationError( - filename, - f"parsing error - content does not match the {source_format} file " - "signature; the file is corrupt or has the wrong extension", - ) - - if spec.checker: - spec.checker(filename, head, tail, file) - def validate_upload( *, diff --git a/backend/app/services/documents/validator.py b/backend/app/services/documents/validator.py new file mode 100644 index 000000000..d9942335f --- /dev/null +++ b/backend/app/services/documents/validator.py @@ -0,0 +1,315 @@ +import csv +import json +import logging +from codecs import getincrementaldecoder +from dataclasses import dataclass +from typing import Callable, Tuple + +from fastapi import UploadFile + + +logger = logging.getLogger(__name__) + +SNIFF_HEAD_BYTES = 4096 +SNIFF_TAIL_BYTES = 2048 + +PDF_EOF_MARKER = b"%%EOF" +PDF_ENCRYPT_MARKER = b"/Encrypt" +OOXML_CONTENT_TYPES = b"[Content_Types].xml" +OOXML_WORD_PART = b"word/" +OOXML_EXCEL_PART = b"xl/" +JSON_OPENERS = (b"{", b"[") +JSON_CLOSERS: dict[int, int] = {ord("{"): ord("}"), ord("["): ord("]")} +UTF8_BOM = b"\xef\xbb\xbf" + +JSON_FULL_PARSE_MAX_BYTES = 256 * 1024 + +CSV_STRICT_COLUMN_COUNT = True +CSV_SAMPLE_MAX_ROWS = 200 + +OLE2_SIGNATURE = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" + +# The OLE2 signature is a generic Compound File Binary container shared by .xls and +# .doc, so it can't tell them apart. The CFB directory stores per-stream names as +# UTF-16LE: an .xls carries a Workbook (BIFF8) or Book (BIFF5) stream, a .doc carries +# WordDocument. +OLE2_EXCEL_STREAMS = ("Workbook".encode("utf-16-le"), "Book".encode("utf-16-le")) +OLE2_WORD_STREAMS = ("WordDocument".encode("utf-16-le"),) + +MISLABELLED_BINARY_SIGNATURES: dict[bytes, str] = { + b"PK\x03\x04": "an Office/zip file (xlsx, docx)", + OLE2_SIGNATURE: "a legacy Office file (xls, doc)", + b"%PDF-": "a PDF", +} + +UNNAMED_FILE_PLACEHOLDER = "" + + +CLIENT_VALIDATION_MESSAGE = ( + "Document '{filename}' failed parsing. The document appears to be corrupted " + "or invalid. Please re-upload a valid document." +) + + +class DocumentValidationError(Exception): + """Raised when a document fails the pre-upload sanity check.""" + + def __init__(self, filename: str, reason: str) -> None: + self.filename = filename + self.reason = reason + super().__init__(f"Document '{filename}' failed validation: {reason}") + + @property + def client_message(self) -> str: + """Client-safe wording. `reason` stays internal - it is log-only.""" + return CLIENT_VALIDATION_MESSAGE.format(filename=self.filename) + + +def _decode_utf8(filename: str, sample: bytes) -> str: + if b"\x00" in sample: + raise DocumentValidationError( + filename, + "parsing error - NUL byte found in a text-based file; " + "the file is binary, corrupt, or has the wrong extension", + ) + try: + return getincrementaldecoder("utf-8")().decode(sample) + except UnicodeDecodeError as e: + raise DocumentValidationError( + filename, + f"parsing error - content is not valid UTF-8 at byte offset {e.start}; " + "re-export the file as UTF-8", + ) from e + + +def _check_pdf(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + if PDF_EOF_MARKER not in tail: + raise DocumentValidationError( + filename, + "parsing error - PDF is missing its end-of-file marker; " + "the upload is truncated or incomplete", + ) + if PDF_ENCRYPT_MARKER in tail: + raise DocumentValidationError( + filename, + "parsing error - PDF is password protected or encrypted and cannot be read", + ) + + +def _check_ooxml(filename: str, sample: bytes, expected_part: bytes) -> None: + if OOXML_CONTENT_TYPES not in sample: + raise DocumentValidationError( + filename, + "parsing error - file is a zip archive but not a valid Office document " + "(no OOXML content-types part)", + ) + + other_part = ( + OOXML_EXCEL_PART if expected_part == OOXML_WORD_PART else OOXML_WORD_PART + ) + if other_part in sample and expected_part not in sample: + raise DocumentValidationError( + filename, + f"parsing error - file is a '{other_part.decode()}' Office package, " + "not the format its extension claims", + ) + + +def _check_docx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + """Head+tail: the zip central directory lists all entry names and sits at the + end; the head holds only the first entry.""" + _check_ooxml(filename, head + tail, OOXML_WORD_PART) + + +def _check_xlsx(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ooxml(filename, head + tail, OOXML_EXCEL_PART) + + +def _check_ole2_stream( + filename: str, + sample: bytes, + wanted: tuple[bytes, ...], + rivals: tuple[bytes, ...], + rival_label: str, +) -> None: + """Reject an OLE2 file whose extension claims one format but whose CFB directory + holds the rival's stream. Inconclusive samples pass — the directory can sit past + the sampled edges, and rejecting a valid file is worse than the loose check.""" + if any(name in sample for name in wanted): + return + if any(name in sample for name in rivals): + raise DocumentValidationError( + filename, + f"parsing error - file is a {rival_label} document, " + "not the format its extension claims", + ) + + +def _check_xls(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ole2_stream( + filename, head + tail, OLE2_EXCEL_STREAMS, OLE2_WORD_STREAMS, "Word (.doc)" + ) + + +def _check_doc(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + _check_ole2_stream( + filename, head + tail, OLE2_WORD_STREAMS, OLE2_EXCEL_STREAMS, "Excel (.xls)" + ) + + +def _check_csv(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + for signature, description in MISLABELLED_BINARY_SIGNATURES.items(): + if head.startswith(signature): + raise DocumentValidationError( + filename, + f"parsing error - file is {description} renamed to a CSV, not text; " + "export it as CSV instead of renaming it", + ) + + text = _decode_utf8(filename, head) + + lines = text.splitlines() + if len(head) == SNIFF_HEAD_BYTES and len(lines) > 1: + lines = lines[:-1] + + try: + rows = [row for row in csv.reader(lines[:CSV_SAMPLE_MAX_ROWS]) if row] + except csv.Error as e: + raise DocumentValidationError( + filename, f"parsing error - malformed CSV: {e}" + ) from e + + if not rows: + raise DocumentValidationError( + filename, "parsing error - no readable CSV rows found" + ) + + if CSV_STRICT_COLUMN_COUNT: + expected_columns = len(rows[0]) + for line_number, row in enumerate(rows[1:], start=2): + if len(row) != expected_columns: + raise DocumentValidationError( + filename, + f"parsing error - inconsistent column count at row {line_number} " + f"({len(row)} fields, expected {expected_columns})", + ) + + +def _check_json_tail(filename: str, opener: int, tail: bytes) -> None: + """Truncation check for JSON too large to parse: the opening bracket must be + matched by the last non-whitespace byte.""" + closing = tail.rstrip() + if not closing or closing[-1] != JSON_CLOSERS[opener]: + raise DocumentValidationError( + filename, + "parsing error - JSON is not closed correctly; " + "the upload is truncated or incomplete", + ) + + +def _check_json(filename: str, head: bytes, tail: bytes, file: UploadFile) -> None: + body = head.lstrip(UTF8_BOM).lstrip() + if not body.startswith(JSON_OPENERS): + raise DocumentValidationError( + filename, "parsing error - JSON document must start with '{' or '['" + ) + + stream = file.file + stream.seek(0, 2) + size_bytes = stream.tell() + stream.seek(0) + + if size_bytes > JSON_FULL_PARSE_MAX_BYTES: + logger.info( + f"[_check_json] Too large for a full parse, checking closure only | " + f"filename: {filename} | size_bytes: {size_bytes}" + ) + _check_json_tail(filename, body[0], tail) + return + + try: + json.loads(stream.read(), object_pairs_hook=lambda pairs: None) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + raise DocumentValidationError( + filename, f"parsing error - invalid JSON: {e}" + ) from e + finally: + stream.seek(0) + + +@dataclass(frozen=True) +class FormatSpec: + signatures: tuple[bytes, ...] = () + needs_tail: bool = False + checker: Callable[[str, bytes, bytes, UploadFile], None] | None = None + + +FORMAT_SPECS: dict[str, FormatSpec] = { + "pdf": FormatSpec(signatures=(b"%PDF-",), needs_tail=True, checker=_check_pdf), + "docx": FormatSpec( + signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), + needs_tail=True, + checker=_check_docx, + ), + "xlsx": FormatSpec( + signatures=(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), + needs_tail=True, + checker=_check_xlsx, + ), + "xls": FormatSpec( + signatures=(OLE2_SIGNATURE,), needs_tail=True, checker=_check_xls + ), + "doc": FormatSpec( + signatures=(OLE2_SIGNATURE,), needs_tail=True, checker=_check_doc + ), + "csv": FormatSpec(checker=_check_csv), + "json": FormatSpec(needs_tail=True, checker=_check_json), +} + + +def _read_edges(file: UploadFile, needs_tail: bool) -> Tuple[bytes, bytes]: + """Sample the head (and optionally tail) of the upload, leaving it rewound.""" + stream = file.file + + stream.seek(0) + head = stream.read(SNIFF_HEAD_BYTES) + + tail = b"" + if needs_tail: + stream.seek(0, 2) + stream.seek(max(0, stream.tell() - SNIFF_TAIL_BYTES)) + tail = stream.read(SNIFF_TAIL_BYTES) + + stream.seek(0) + return head, tail + + +def validate_document_content(*, file: UploadFile, source_format: str) -> None: + """ + Sanity-check an upload's bytes before it is persisted. Only the first + SNIFF_HEAD_BYTES (plus SNIFF_TAIL_BYTES for PDF) are read, so the cost is + constant in file size - small JSON, which is fully parsed, is the exception. + + Raises: + DocumentValidationError: naming the offending file and the reason. + """ + filename = (file.filename or "").strip() or UNNAMED_FILE_PLACEHOLDER + + spec = FORMAT_SPECS.get(source_format) + head, tail = _read_edges(file, needs_tail=bool(spec and spec.needs_tail)) + + if not head: + raise DocumentValidationError(filename, "the file is empty") + + if spec is None: + return + + if spec.signatures and not head.startswith(spec.signatures): + raise DocumentValidationError( + filename, + f"parsing error - content does not match the {source_format} file " + "signature; the file is corrupt or has the wrong extension", + ) + + if spec.checker: + spec.checker(filename, head, tail, file) diff --git a/backend/app/tests/services/documents/test_helpers.py b/backend/app/tests/services/documents/test_helpers.py index 524cb5926..7eaddf9c0 100644 --- a/backend/app/tests/services/documents/test_helpers.py +++ b/backend/app/tests/services/documents/test_helpers.py @@ -5,6 +5,10 @@ from fastapi import HTTPException, UploadFile from app.services.documents.helpers import ( + calculate_file_size, + validate_upload, +) +from app.services.documents.validator import ( SNIFF_HEAD_BYTES, SNIFF_TAIL_BYTES, DocumentValidationError, @@ -17,9 +21,7 @@ _check_xlsx, _decode_utf8, _read_edges, - calculate_file_size, validate_document_content, - validate_upload, )