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 1/4] 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 2/4] 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 3/4] 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 4/4] 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.