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/crud/rag/open_ai.py b/backend/app/crud/rag/open_ai.py index 0df73ab39..fdc2f8d56 100644 --- a/backend/app/crud/rag/open_ai.py +++ b/backend/app/crud/rag/open_ai.py @@ -1,10 +1,20 @@ 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 tenacity import ( + RetryCallState, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) from app.models import Document, ProviderType @@ -12,6 +22,26 @@ OPENAI_PROVIDER = ProviderType.openai.value +# 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 = {} @@ -85,7 +115,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 +131,100 @@ 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: + """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 + 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, @@ -115,11 +239,7 @@ def update( ) try: - batch = self.client.vector_stores.file_batches.upload_and_poll( - vector_store_id=vector_store_id, - files=[], - file_ids=[doc.file_id[OPENAI_PROVIDER] for doc in docs], - ) + 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}): " @@ -215,30 +335,9 @@ 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, - 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}") - 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 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 234443a00..4e1b79b9c 100644 --- a/backend/app/services/collections/create_collection.py +++ b/backend/app/services/collections/create_collection.py @@ -403,9 +403,11 @@ 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 - ] + 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, @@ -507,7 +509,10 @@ 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") logger.warning( "[create_collection.execute_batch_job] Collection Creation Timed Out | {'collection_job_id': '%s', 'error': '%s'}", diff --git a/backend/app/services/collections/providers/registry.py b/backend/app/services/collections/providers/registry.py index 7c40d8e5c..cf1729d6b 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_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=0, + timeout=OPENAI_TIMEOUT_SECONDS, + ) elif provider == LLMProvider.GOOGLE_AISTUDIO: if "api_key" not in credentials: raise ValueError( 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..c8cb77ea0 100644 --- a/backend/app/services/documents/helpers.py +++ b/backend/app/services/documents/helpers.py @@ -1,3 +1,4 @@ +import logging from typing import Optional, Tuple, Iterable, Union from uuid import UUID @@ -11,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, @@ -23,6 +28,39 @@ ) +logger = logging.getLogger(__name__) + + +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. 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/crud/rag/test_open_ai.py b/backend/app/tests/crud/rag/test_open_ai.py index d8c85c227..38c40940d 100644 --- a/backend/app/tests/crud/rag/test_open_ai.py +++ b/backend/app/tests/crud/rag/test_open_ai.py @@ -4,12 +4,18 @@ OpenAI exception type). """ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import openai import pytest +from tenacity import stop_after_attempt, wait_none -from app.crud.rag.open_ai import 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 @pytest.fixture @@ -22,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} @@ -37,10 +54,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.""" - counts = MagicMock(completed=completed, failed=failed) - return MagicMock(id=batch_id, file_counts=counts) +REAL_BATCH_ID = "vsfb_real123" +CORRUPT_POLL_ID = "vs_corrupt999" + + +def _wire_batch(mock_client, *, completed: int, failed: int) -> None: + """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, in_progress=0) + mock_client.vector_stores.file_batches.retrieve.return_value = MagicMock( + id=CORRUPT_POLL_ID, status="completed", file_counts=counts + ) def _failed_file(file_id: str, error_message: str | None): @@ -53,30 +79,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.retrieve.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 +120,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 +154,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") ] @@ -119,22 +162,20 @@ 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.""" - mock_client.vector_stores.file_batches.upload_and_poll.return_value = ( - _batch_result(completed=0, failed=2) - ) + 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) 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,9 +271,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 = ( - 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) @@ -243,7 +282,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,8 +295,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.upload_and_poll.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: @@ -273,3 +312,113 @@ class TestOpenAIVectorStoreCrudInit: def test_none_client_raises(self): with pytest.raises(ValueError): OpenAIVectorStoreCrud(client=None) + + +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) + + +class TestUpdateTerminalStatus: + """A batch ending 'cancelled'/'failed' with no per-file failures must raise, not + slip through the failed-count check as success.""" + + 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) -> None: + 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 + # _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" + + +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 1c07c100f..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) @@ -380,11 +391,15 @@ 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() 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 - return 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: @@ -398,24 +413,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( - "rate limit" - ) + client.vector_stores.file_batches.create.side_effect = OpenAIError("rate limit") crud = OpenAIVectorStoreCrud(client) with pytest.raises(InterruptedError, match="rate limit"): @@ -430,9 +441,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 +453,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 +466,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/services/documents/test_helpers.py b/backend/app/tests/services/documents/test_helpers.py index 9a4b25344..7eaddf9c0 100644 --- a/backend/app/tests/services/documents/test_helpers.py +++ b/backend/app/tests/services/documents/test_helpers.py @@ -1,8 +1,28 @@ +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 ( + calculate_file_size, + validate_upload, +) +from app.services.documents.validator import ( + SNIFF_HEAD_BYTES, + SNIFF_TAIL_BYTES, + DocumentValidationError, + _check_csv, + _check_doc, + _check_docx, + _check_json, + _check_pdf, + _check_xls, + _check_xlsx, + _decode_utf8, + _read_edges, + validate_document_content, +) def make_upload_file(content: bytes, size: int | None = None) -> UploadFile: @@ -58,3 +78,252 @@ 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 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: + _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 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..7f680057b 100644 --- a/backend/app/tests/utils/openai.py +++ b/backend/app/tests/utils/openai.py @@ -116,13 +116,16 @@ 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.status = "completed" 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 + # _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 348fe8e7f..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) @@ -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 4307e6b30..c6b8462df 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -35,3 +35,11 @@ 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. +- 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.