diff --git a/api-compose.yaml b/api-compose.yaml
index aa2f95b8..02b8f515 100644
--- a/api-compose.yaml
+++ b/api-compose.yaml
@@ -1,14 +1,34 @@
version: '3.8'
services:
+ extraction-init: !override
+ extends:
+ file: extraction-compose.yaml
+ service: extraction-init
+
+ extraction: !override
+ extends:
+ file: extraction-compose.yaml
+ service: extraction
+
rag_api:
build: .
+ networks: [default, extraction]
+ depends_on:
+ extraction:
+ condition: service_healthy
ports:
- "8000:8000"
volumes:
- ./uploads:/app/uploads
+ - extraction-credentials:/run/extraction:ro
env_file:
- .env
volumes:
+ extraction-credentials:
pgdata2:
+
+networks:
+ extraction:
+ internal: true
diff --git a/app/config.py b/app/config.py
index 3fc8560e..85d5d4d4 100644
--- a/app/config.py
+++ b/app/config.py
@@ -10,6 +10,7 @@
from starlette.middleware.base import BaseHTTPMiddleware
from app.services.vector_store.factory import get_vector_store
+from app.services.extraction_client import ExtractionSettings
load_dotenv(find_dotenv())
@@ -41,6 +42,8 @@ def get_env_variable(
return value
+EXTRACTION_SETTINGS = ExtractionSettings.from_env()
+
RAG_HOST = os.getenv("RAG_HOST", "0.0.0.0")
RAG_PORT = int(os.getenv("RAG_PORT", 8000))
diff --git a/app/routes/document_routes.py b/app/routes/document_routes.py
index c3cdb2ee..659ac588 100644
--- a/app/routes/document_routes.py
+++ b/app/routes/document_routes.py
@@ -7,8 +7,12 @@
import traceback
import aiofiles
import aiofiles.os
+from anyio import CancelScope
import asyncio
import time
+import inspect
+from contextvars import ContextVar
+from threading import Lock
from shutil import copyfileobj
from typing import List, Iterable, Optional, Union, TYPE_CHECKING
from concurrent.futures import ThreadPoolExecutor
@@ -25,7 +29,7 @@
)
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
-from functools import lru_cache
+from functools import lru_cache, wraps
if TYPE_CHECKING:
from app.services.vector_store.async_pg_vector import AsyncPgVector
@@ -44,8 +48,11 @@
EMBEDDING_MAX_QUEUE_SIZE,
PARALLEL_EXECUTION,
RAG_DISTANCE_THRESHOLD,
+ EXTRACTION_SETTINGS,
+ PDF_EXTRACT_IMAGES,
)
from app.scope import file_clause, files_clause, resolve_scope
+from app.services.extraction_client import ExtractionClient, ExtractionError
# Warn once at import time if the user set a threshold under Atlas, where
# the score direction is inverted (Atlas vectorSearchScore: higher = better)
@@ -94,6 +101,7 @@ def _apply_distance_threshold(documents):
from app.utils.async_utils import run_in_executor
from app.utils.document_loader import (
get_loader,
+ extraction_format,
clean_text,
process_documents,
cleanup_temp_encoding_file,
@@ -101,6 +109,7 @@ def _apply_distance_threshold(documents):
from app.utils.health import is_health_ok
router = APIRouter()
+extraction_client = ExtractionClient(EXTRACTION_SETTINGS)
_INGESTION_ATTEMPT_ID_KEY = "_rag_ingestion_attempt_id"
_INGESTION_ATTEMPT_STARTED_AT_NS_KEY = "_rag_ingestion_attempt_started_at_ns"
@@ -335,12 +344,168 @@ def _make_unique_temp_path(user_id: str, filename: str) -> Optional[str]:
return str(Path(RAG_UPLOAD_DIR, user_id, unique_name).resolve())
+_active_extraction_request = ContextVar("rag_extraction_request", default=None)
+
+
+class _ExtractionRequest:
+ def __init__(self):
+ self.disconnected = False
+ self.finish_storage = False
+ self.cleanup_started = False
+ self._mutation_lock = Lock()
+
+ def mark_disconnected(self):
+ with self._mutation_lock:
+ self.disconnected = True
+
+ def checkpoint(self):
+ if self.disconnected:
+ raise asyncio.CancelledError()
+
+ def begin_cleanup(self):
+ self.cleanup_started = True
+
+ def start_uncancellable_storage(self):
+ with self._mutation_lock:
+ self.checkpoint()
+ # These store APIs expose no safe attempt-scoped cancellation/undo boundary.
+ self.finish_storage = True
+
+
+def _with_extraction_request(file_argument, content_type="content_type"):
+ """Keep selected requests alive until cancellable work or started storage drains."""
+
+ def decorate(handler):
+ signature = inspect.signature(handler)
+
+ @wraps(handler)
+ async def wrapped(*args, **kwargs):
+ arguments = signature.bind(*args, **kwargs).arguments
+ file = arguments[file_argument]
+ format = extraction_format(file.filename, getattr(file, content_type))
+ if format is None:
+ return await handler(*args, **kwargs)
+
+ request = arguments["request"]
+ lifecycle = _ExtractionRequest()
+
+ async def wait_for_disconnect():
+ # FastAPI consumed the complete multipart/JSON body before the handler.
+ while (await request.receive())["type"] != "http.disconnect":
+ pass
+ lifecycle.mark_disconnected()
+
+ token = _active_extraction_request.set(lifecycle)
+ try:
+ operation = asyncio.create_task(handler(*args, **kwargs))
+ finally:
+ _active_extraction_request.reset(token)
+ disconnected = asyncio.create_task(wait_for_disconnect())
+ try:
+ await asyncio.wait(
+ (operation, disconnected), return_when=asyncio.FIRST_COMPLETED
+ )
+ if disconnected.done():
+ disconnected.result()
+ raise HTTPException(status_code=499, detail={"code": "CANCELLED"})
+ return operation.result()
+ finally:
+ if (
+ not operation.done()
+ and not operation.cancelling()
+ and not (lifecycle.finish_storage or lifecycle.cleanup_started)
+ ):
+ operation.cancel()
+ disconnected.cancel()
+ # Do not clean uploads or return while workers/rollback can still run.
+ await _drain_request_worker(
+ asyncio.gather(operation, disconnected, return_exceptions=True)
+ )
+
+ return wrapped
+
+ return decorate
+
+
+async def _drain_request_worker(future):
+ """Defer repeated cancellation until submitted work or cleanup has finished."""
+ cancellation = None
+ with CancelScope(shield=True):
+ while True:
+ try:
+ result = await asyncio.shield(future)
+ break
+ except asyncio.CancelledError as error:
+ lifecycle = _active_extraction_request.get()
+ if lifecycle is not None:
+ lifecycle.begin_cleanup()
+ # Executor futures can store CancelledError without being cancelled.
+ if future.cancelled() or (
+ future.done()
+ and isinstance(future.exception(), asyncio.CancelledError)
+ ):
+ raise
+ cancellation = error
+ except Exception:
+ if cancellation is not None:
+ raise cancellation
+ raise
+ if cancellation is not None:
+ raise cancellation
+ return result
+
+
+async def _run_sync_request_store(executor, store, documents, ids):
+ lifecycle = _active_extraction_request.get()
+
+ def insert():
+ if not lifecycle.finish_storage:
+ lifecycle.start_uncancellable_storage()
+ return store.add_documents(documents, ids=ids)
+
+ if not lifecycle.finish_storage:
+ lifecycle.checkpoint()
+ future = asyncio.get_running_loop().run_in_executor(executor, insert)
+ return await _drain_request_worker(future)
+
+
+async def _extract_for_request(request: Optional[Request], file_path: str, format: str):
+ lifecycle = _active_extraction_request.get()
+ if request is None or lifecycle is not None:
+ if lifecycle is not None:
+ lifecycle.checkpoint()
+ return await extraction_client.extract(file_path, format)
+
+ async def wait_for_disconnect():
+ # FastAPI has already consumed the multipart/JSON body before this route.
+ while (await request.receive())["type"] != "http.disconnect":
+ pass
+
+ disconnected = asyncio.create_task(wait_for_disconnect())
+ extraction = asyncio.create_task(extraction_client.extract(file_path, format))
+ try:
+ await asyncio.wait(
+ (disconnected, extraction), return_when=asyncio.FIRST_COMPLETED
+ )
+ if disconnected.done():
+ disconnected.result()
+ raise ExtractionError("CANCELLED", status_code=499)
+ return extraction.result()
+ finally:
+ disconnected.cancel()
+ extraction.cancel()
+ # Drain both tasks before cleanup or persistence, including ASGI cancellation.
+ with CancelScope(shield=True):
+ await asyncio.gather(disconnected, extraction, return_exceptions=True)
+
+
async def load_file_content(
filename: str,
content_type: str,
file_path: str,
executor,
raw_text: bool = False,
+ request: Optional[Request] = None,
) -> tuple:
"""Load file content using appropriate loader.
@@ -348,6 +513,30 @@ async def load_file_content(
the ``/text`` endpoint) so text-formatted files are not semantically
parsed.
"""
+ format = extraction_format(filename, content_type)
+ if format is not None:
+ try:
+ if format == "pdf" and PDF_EXTRACT_IMAGES:
+ raise ExtractionError("UNSUPPORTED_CONTENT")
+ segments = await _extract_for_request(request, file_path, format)
+ except ExtractionError as error:
+ raise HTTPException(
+ status_code=error.status_code,
+ detail={"code": error.code},
+ ) from None
+ # Service pages are one-based; RAG metadata remains zero-based.
+ documents = [
+ Document(
+ page_content=segment["text"],
+ metadata={
+ "source": file_path,
+ **({"page": segment["index"] - 1} if format == "pdf" else {}),
+ },
+ )
+ for segment in segments
+ ]
+ return documents, True, format
+
loader = None
try:
loader, known_type, file_ext = get_loader(
@@ -380,8 +569,18 @@ def extract_text_from_documents(documents: List[Document], file_ext: str) -> str
async def cleanup_temp_file_async(file_path: str) -> None:
"""Clean up temporary file asynchronously."""
+ lifecycle = _active_extraction_request.get()
+ if lifecycle is not None:
+ lifecycle.begin_cleanup()
try:
- await aiofiles.os.remove(file_path)
+ # Starlette cancellation scopes must not interrupt upload cleanup.
+ with CancelScope(shield=True):
+ if lifecycle is None:
+ await aiofiles.os.remove(file_path)
+ else:
+ await _drain_request_worker(
+ asyncio.create_task(aiofiles.os.remove(file_path))
+ )
except Exception as e:
logger.error(
"Failed to remove temporary file | Path: %s | Error: %s | Traceback: %s",
@@ -684,6 +883,9 @@ async def insert_batch(
) -> List[str]:
"""Track started inserts so rollback waits for executor-backed work."""
nonlocal cleanup_needed
+ lifecycle = _active_extraction_request.get()
+ if lifecycle is not None:
+ lifecycle.checkpoint()
insert_task = asyncio.create_task(
vector_store.aadd_documents(
batch_documents, ids=batch_ids, executor=executor
@@ -778,39 +980,7 @@ async def embedding_consumer():
producer_task = None
consumer_tasks: List[asyncio.Task] = []
- try:
- # Start producer and consumers concurrently
- producer_task = asyncio.create_task(batch_producer())
- consumer_tasks = [
- asyncio.create_task(embedding_consumer()) for _ in range(consumer_count)
- ]
-
- # Wait for all tasks to complete
- await asyncio.gather(producer_task, *consumer_tasks, return_exceptions=False)
-
- for batch_num in range(1, num_batches + 1):
- all_ids.extend(successful_batch_ids[batch_num])
-
- logger.info(
- "Async pipeline completed | user_id=%s | file_id=%s | inserted_ids=%d | %s",
- user_id,
- file_id,
- len(all_ids),
- get_process_memory_details(),
- )
-
- return all_ids
-
- except (Exception, asyncio.CancelledError) as e:
- failure_event.set()
- logger.error(
- "Pipeline failed | user_id=%s | file_id=%s | inserted_batches=%d | error=%s | %s",
- user_id,
- file_id,
- len(successful_batch_ids),
- e,
- get_process_memory_details(),
- )
+ async def cleanup_failed_pipeline():
if producer_task is not None or consumer_tasks:
# if one of the tasks is still running, cancel it
if producer_task is not None and not producer_task.done():
@@ -859,6 +1029,49 @@ async def embedding_consumer():
get_process_memory_details(),
)
+ try:
+ # Start producer and consumers concurrently
+ producer_task = asyncio.create_task(batch_producer())
+ consumer_tasks = [
+ asyncio.create_task(embedding_consumer()) for _ in range(consumer_count)
+ ]
+
+ # Wait for all tasks to complete
+ await asyncio.gather(producer_task, *consumer_tasks, return_exceptions=False)
+ lifecycle = _active_extraction_request.get()
+ if lifecycle is not None:
+ lifecycle.checkpoint()
+
+ for batch_num in range(1, num_batches + 1):
+ all_ids.extend(successful_batch_ids[batch_num])
+
+ logger.info(
+ "Async pipeline completed | user_id=%s | file_id=%s | inserted_ids=%d | %s",
+ user_id,
+ file_id,
+ len(all_ids),
+ get_process_memory_details(),
+ )
+
+ return all_ids
+
+ except (Exception, asyncio.CancelledError) as e:
+ failure_event.set()
+ logger.error(
+ "Pipeline failed | user_id=%s | file_id=%s | inserted_batches=%d | error=%s | %s",
+ user_id,
+ file_id,
+ len(successful_batch_ids),
+ e,
+ get_process_memory_details(),
+ )
+ lifecycle = _active_extraction_request.get()
+ if lifecycle is None:
+ await cleanup_failed_pipeline()
+ else:
+ lifecycle.begin_cleanup()
+ await _drain_request_worker(asyncio.create_task(cleanup_failed_pipeline()))
+
# Re-raise the original error
raise
@@ -916,21 +1129,27 @@ async def _process_documents_batched_sync(
try:
# Wrap sync call in executor to avoid blocking the event loop
- batch_result_ids = await loop.run_in_executor(
- executor,
- lambda docs=batch_documents, ids=batch_ids: vector_store.add_documents(
- docs, ids=ids
- ),
- )
+ if _active_extraction_request.get() is not None:
+ batch_result_ids = await _run_sync_request_store(
+ executor, vector_store, batch_documents, batch_ids
+ )
+ else:
+ batch_result_ids = await loop.run_in_executor(
+ executor,
+ lambda docs=batch_documents, ids=batch_ids: vector_store.add_documents(
+ docs, ids=ids
+ ),
+ )
all_ids.extend(batch_result_ids)
except Exception as batch_error:
logger.error("Batch %d failed: %s", batch_idx + 1, batch_error)
# Rollback entire file from vector store
- if (
- all_ids
- ): # any batch succeeded (i.e., any chunks for this file were inserted)
+ if all_ids and _active_extraction_request.get() is not None:
+ # Sync stores cannot safely distinguish this attempt from older rows.
+ logger.warning("Retaining committed sync batches for file %s", file_id)
+ elif all_ids:
logger.warning("Rolling back file %s due to batch failure", file_id)
try:
await loop.run_in_executor(
@@ -1003,7 +1222,10 @@ async def store_data_in_vector_db(
start_time = time.perf_counter()
# Run document preparation in executor to avoid blocking the event loop
loop = asyncio.get_running_loop()
- docs = await loop.run_in_executor(
+ lifecycle = _active_extraction_request.get()
+ if lifecycle is not None:
+ lifecycle.checkpoint()
+ preparation = loop.run_in_executor(
executor,
_prepare_documents_sync,
data,
@@ -1011,6 +1233,11 @@ async def store_data_in_vector_db(
user_id,
clean_content,
)
+ if lifecycle is None:
+ docs = await preparation
+ else:
+ docs = await _drain_request_worker(preparation)
+ lifecycle.checkpoint()
logger.info(
"Documents prepared | %s",
@@ -1030,10 +1257,16 @@ async def store_data_in_vector_db(
if EMBEDDING_BATCH_SIZE <= 0:
# synchronously embed the file and insert into vector store in one go
if isinstance(vector_store, AsyncPgVector):
+ if lifecycle is not None:
+ lifecycle.start_uncancellable_storage()
_tag_documents_for_ingestion(docs, file_id)
ids = await vector_store.aadd_documents(
docs, ids=[file_id] * len(docs), executor=executor
)
+ elif lifecycle is not None:
+ ids = await _run_sync_request_store(
+ executor, vector_store, docs, [file_id] * len(docs)
+ )
else:
ids = vector_store.add_documents(docs, ids=[file_id] * len(docs))
else:
@@ -1094,6 +1327,7 @@ async def store_data_in_vector_db(
@router.post("/local/embed")
+@_with_extraction_request("document", content_type="file_content_type")
async def embed_local_file(
document: StoreDocument, request: Request, entity_id: str = None
):
@@ -1114,7 +1348,6 @@ async def embed_local_file(
detail=ERROR_MESSAGES.FILE_NOT_FOUND,
)
- loader = None
try:
logger.info(
"Ingestion started | %s",
@@ -1127,12 +1360,12 @@ async def embed_local_file(
temp_file_path=file_path,
),
)
- loader, known_type, file_ext = get_loader(
- document.filename, document.file_content_type, file_path
- )
- loop = asyncio.get_running_loop()
- data = await loop.run_in_executor(
- request.app.state.thread_pool, lambda: list(loader.lazy_load())
+ data, known_type, file_ext = await load_file_content(
+ document.filename,
+ document.file_content_type,
+ file_path,
+ request.app.state.thread_pool,
+ request=request,
)
result = await store_data_in_vector_db(
@@ -1147,7 +1380,7 @@ async def embed_local_file(
temp_file_path=file_path,
)
- if result:
+ if result and "error" not in result:
return {
"status": True,
"file_id": document.file_id,
@@ -1188,13 +1421,10 @@ async def embed_local_file(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
- finally:
- # Clean up temporary UTF-8 file if it was created for encoding conversion
- if loader is not None:
- cleanup_temp_encoding_file(loader)
@router.post("/embed")
+@_with_extraction_request("file")
async def embed_file(
request: Request,
file_id: str = Form(...),
@@ -1239,6 +1469,7 @@ async def embed_file(
file.content_type,
validated_file_path,
request.app.state.thread_pool,
+ request=request,
)
logger.debug(
@@ -1261,23 +1492,11 @@ async def embed_file(
temp_file_path=validated_file_path,
)
- if not result:
- response_status = False
- response_message = "Failed to process/store the file data."
+ if not result or "error" in result:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to process/store the file data.",
)
- elif "error" in result:
- response_status = False
- response_message = "Failed to process/store the file data."
- if isinstance(result["error"], str):
- response_message = result["error"]
- else:
- raise HTTPException(
- status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
- detail="An unspecified error occurred.",
- )
except HTTPException as http_exc:
response_status = False
response_message = f"HTTP Exception: {http_exc.detail}"
@@ -1368,6 +1587,7 @@ async def load_document_context(request: Request, id: str, entity_id: str = None
@router.post("/embed-upload")
+@_with_extraction_request("uploaded_file")
async def embed_file_upload(
request: Request,
file_id: str = Form(...),
@@ -1409,6 +1629,7 @@ async def embed_file_upload(
uploaded_file.content_type,
validated_temp_file_path,
request.app.state.thread_pool,
+ request=request,
)
result = await store_data_in_vector_db(
@@ -1423,7 +1644,7 @@ async def embed_file_upload(
temp_file_path=validated_temp_file_path,
)
- if not result:
+ if not result or "error" in result:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to process/store the file data.",
@@ -1518,6 +1739,7 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody
@router.post("/text")
+@_with_extraction_request("file")
async def extract_text_from_file(
request: Request,
file_id: str = Form(...),
@@ -1547,6 +1769,7 @@ async def extract_text_from_file(
validated_temp_file_path,
request.app.state.thread_pool,
raw_text=True,
+ request=request,
)
# Extract text content from loaded documents
diff --git a/app/services/extraction_client.py b/app/services/extraction_client.py
new file mode 100644
index 00000000..2bbf5b8f
--- /dev/null
+++ b/app/services/extraction_client.py
@@ -0,0 +1,436 @@
+"""Bounded v1 extraction transport. Only caller-owned bytes and service auth cross HTTP."""
+
+import asyncio
+import hashlib
+import json
+import math
+import os
+import re
+import ssl
+import stat
+from dataclasses import dataclass
+from urllib.parse import urlsplit
+
+import aiofiles
+
+INPUT_BYTES = 10 * 1024 * 1024
+TEXT_BYTES = 1024 * 1024
+RESULT_BYTES = 3 * 1024 * 1024
+HEADER_BYTES = 4096
+CAPABILITY_BYTES = 8192
+CHUNK_BYTES = 64 * 1024
+OPERATION = "document.extract-text"
+ERROR_STATUS = {
+ "UNAUTHORIZED": 401,
+ "INVALID_REQUEST": 400,
+ "INPUT_LIMIT": 413,
+ "OUTPUT_LIMIT": 413,
+ "DECOMPRESSION_LIMIT": 413,
+ "STRUCTURE_LIMIT": 413,
+ "BUSY": 429,
+ "UNAVAILABLE": 503,
+ "DEADLINE": 504,
+ **dict.fromkeys(
+ (
+ "CANCELLED",
+ "INVALID_DOCUMENT",
+ "EMPTY_OUTPUT",
+ "ENCRYPTED_DOCUMENT",
+ "UNSUPPORTED_ENCODING",
+ "UNSUPPORTED_CONTENT",
+ "RESOURCE_LIMIT",
+ ),
+ 422,
+ ),
+}
+
+
+class ExtractionError(Exception):
+ def __init__(self, code: str, status_code: int = None):
+ self.code = code
+ self.status_code = status_code or ERROR_STATUS.get(code, 502)
+ super().__init__(code)
+
+
+def _endpoint(url):
+ try:
+ if (
+ not isinstance(url, str)
+ or not 1 <= len(url) <= 2048
+ or any(ord(char) <= 32 or ord(char) >= 127 for char in url)
+ or "?" in url
+ or "#" in url
+ or "\\" in url
+ ):
+ raise ValueError()
+ endpoint = urlsplit(url)
+ if (
+ endpoint.scheme not in {"http", "https"}
+ or not endpoint.hostname
+ or endpoint.username is not None
+ or endpoint.password is not None
+ or endpoint.path not in {"", "/"}
+ or endpoint.netloc.endswith(":")
+ or not 1
+ <= (endpoint.port or (443 if endpoint.scheme == "https" else 80))
+ <= 65535
+ or endpoint.port == 0
+ or not re.fullmatch(r"[a-zA-Z0-9.:-]+", endpoint.hostname)
+ ):
+ raise ValueError()
+ authority = (
+ f"[{endpoint.hostname}]" if ":" in endpoint.hostname else endpoint.hostname
+ )
+ if endpoint.port is not None:
+ authority += f":{endpoint.port}"
+ if endpoint.netloc.lower() != authority:
+ raise ValueError()
+ if endpoint.scheme == "http" and not (
+ endpoint.hostname in {"localhost", "127.0.0.1", "::1"}
+ or (endpoint.hostname == "extraction" and endpoint.port == 8000)
+ ):
+ raise ValueError()
+ return endpoint
+ except ValueError:
+ raise ValueError(
+ "RAG_EXTRACTION_URL must be an HTTPS origin or local HTTP extraction origin"
+ ) from None
+
+
+@dataclass(frozen=True)
+class ExtractionSettings:
+ url: str = "http://extraction:8000"
+ token_file: str = "/run/extraction/token"
+ deadline: float = 10.0
+
+ def __post_init__(self):
+ _endpoint(self.url)
+ if (
+ not isinstance(self.token_file, str)
+ or not os.path.isabs(self.token_file)
+ or "\x00" in self.token_file
+ ):
+ raise ValueError("RAG_EXTRACTION_TOKEN_FILE must be an absolute file path")
+ if (
+ type(self.deadline) not in (int, float)
+ or not math.isfinite(self.deadline)
+ or not 0 < self.deadline <= 10
+ ):
+ raise ValueError("RAG_EXTRACTION_DEADLINE_SECONDS must be in (0, 10]")
+
+ @classmethod
+ def from_env(cls):
+ for key in (
+ "RAG_EXTRACTION_ENABLED",
+ "RAG_EXTRACTION_FORMATS",
+ "RAG_EXTRACTION_SOCKET",
+ ):
+ if key in os.environ:
+ raise ValueError(
+ f"{key} is obsolete; authenticated PDF/DOCX service extraction is required"
+ )
+ try:
+ deadline = float(os.getenv("RAG_EXTRACTION_DEADLINE_SECONDS", "10"))
+ except ValueError:
+ raise ValueError(
+ "RAG_EXTRACTION_DEADLINE_SECONDS must be in (0, 10]"
+ ) from None
+ return cls(
+ url=os.getenv("RAG_EXTRACTION_URL", "http://extraction:8000"),
+ token_file=os.getenv("RAG_EXTRACTION_TOKEN_FILE", "/run/extraction/token"),
+ deadline=deadline,
+ )
+
+
+def _read_token(path):
+ descriptor = os.open(
+ path, os.O_RDONLY | os.O_NONBLOCK | os.O_NOFOLLOW | os.O_CLOEXEC
+ )
+ try:
+ metadata = os.fstat(descriptor)
+ if not stat.S_ISREG(metadata.st_mode) or metadata.st_size != 64:
+ raise ExtractionError("UNAVAILABLE")
+ token = os.read(descriptor, 65)
+ if not re.fullmatch(b"[0-9a-f]{64}", token):
+ raise ExtractionError("UNAVAILABLE")
+ return token.decode("ascii")
+ finally:
+ os.close(descriptor)
+
+
+def _object(pairs):
+ result = {}
+ for key, value in pairs:
+ if key in result:
+ raise ValueError("duplicate key")
+ result[key] = value
+ return result
+
+
+def _json(body: bytes):
+ try:
+ return json.loads(
+ body.decode("utf-8", "strict"),
+ object_pairs_hook=_object,
+ parse_constant=lambda _: (_ for _ in ()).throw(ValueError("constant")),
+ )
+ except (UnicodeError, ValueError, RecursionError):
+ raise ExtractionError("INVALID_RESPONSE") from None
+
+
+def _keys(value, expected):
+ return type(value) is dict and value.keys() == set(expected)
+
+
+def _integer(value, minimum, maximum):
+ return type(value) is int and minimum <= value <= maximum
+
+
+def validate_capabilities(value, format: str, size: int):
+ if not _keys(
+ value, ("version", "operation", "formats", "structure", "limitations", "limits")
+ ):
+ raise ExtractionError("INCOMPATIBLE_SERVICE")
+ if (
+ type(value["version"]) is not int
+ or value["version"] != 1
+ or value["operation"] != OPERATION
+ or type(value["formats"]) is not list
+ or not value["formats"]
+ or any(
+ type(item) is not str or item not in {"pdf", "docx"}
+ for item in value["formats"]
+ )
+ or len(set(value["formats"])) != len(value["formats"])
+ or format not in value["formats"]
+ ):
+ raise ExtractionError("INCOMPATIBLE_SERVICE")
+ for key in ("structure", "limitations"):
+ if not _keys(value[key], value["formats"]):
+ raise ExtractionError("INCOMPATIBLE_SERVICE")
+ for text in value[key].values():
+ try:
+ if (
+ type(text) is not str
+ or not text.strip()
+ or len(text.encode("utf-8")) > 2048
+ ):
+ raise ValueError()
+ except (ValueError, UnicodeError):
+ raise ExtractionError("INCOMPATIBLE_SERVICE") from None
+ limits = value["limits"]
+ maximums = {
+ "inputBytes": INPUT_BYTES,
+ "textBytes": TEXT_BYTES,
+ "resultBytes": RESULT_BYTES,
+ "deadlineMs": 10_000,
+ "maxSegments": 128,
+ "concurrency": 2,
+ "expandedBytes": 32 * 1024 * 1024,
+ "archiveEntries": 512,
+ "tableColumns": 512,
+ }
+ if not _keys(limits, maximums) or any(
+ not _integer(limits[key], 1, maximum) for key, maximum in maximums.items()
+ ):
+ raise ExtractionError("INCOMPATIBLE_SERVICE")
+ if size > limits["inputBytes"]:
+ raise ExtractionError("INPUT_LIMIT")
+ return limits
+
+
+def validate_result(value, format: str, digest: str, limits):
+ if not _keys(
+ value,
+ ("version", "operation", "format", "inputSha256", "textBytes", "segments"),
+ ):
+ raise ExtractionError("INVALID_RESPONSE")
+ if (
+ type(value["version"]) is not int
+ or value["version"] != 1
+ or value["operation"] != OPERATION
+ or value["format"] != format
+ or value["inputSha256"] != digest
+ or not _integer(value["textBytes"], 1, limits["textBytes"])
+ or type(value["segments"]) is not list
+ or not 1
+ <= len(value["segments"])
+ <= (limits["maxSegments"] if format == "pdf" else 1)
+ ):
+ raise ExtractionError("INVALID_RESPONSE")
+ size = 0
+ nonempty = False
+ for index, segment in enumerate(value["segments"], 1):
+ if (
+ not _keys(segment, ("kind", "index", "text"))
+ or segment["kind"] != ("page" if format == "pdf" else "document")
+ or type(segment["index"]) is not int
+ or segment["index"] != index
+ or type(segment["text"]) is not str
+ ):
+ raise ExtractionError("INVALID_RESPONSE")
+ try:
+ size += len(segment["text"].encode("utf-8", "strict"))
+ except UnicodeError:
+ raise ExtractionError("INVALID_RESPONSE") from None
+ nonempty |= bool(segment["text"].strip())
+ if size != value["textBytes"] or size > limits["textBytes"] or not nonempty:
+ raise ExtractionError("INVALID_RESPONSE")
+ return value["segments"]
+
+
+class ExtractionClient:
+ def __init__(self, settings: ExtractionSettings):
+ self.settings = settings
+ self._admitted = 0
+ self._endpoint = _endpoint(settings.url)
+ self._ssl = (
+ ssl.create_default_context() if self._endpoint.scheme == "https" else None
+ )
+
+ async def _request(
+ self,
+ method: str,
+ token: str,
+ *,
+ file=None,
+ size=0,
+ format=None,
+ maximum=CAPABILITY_BYTES,
+ ):
+ reader, writer = await asyncio.open_connection(
+ self._endpoint.hostname,
+ self._endpoint.port or (443 if self._ssl else 80),
+ ssl=self._ssl,
+ server_hostname=self._endpoint.hostname if self._ssl else None,
+ limit=HEADER_BYTES,
+ )
+ # At most one chunk may be queued without drain, including on cancellation.
+ writer.transport.set_write_buffer_limits(high=CHUNK_BYTES, low=CHUNK_BYTES // 2)
+ digest = hashlib.sha256()
+ try:
+ path = "/v1/extract-text" if file is not None else "/v1/capabilities"
+ headers = (
+ f"{method} {path} HTTP/1.1\r\nHost: {self._endpoint.netloc}\r\n"
+ f"Authorization: Bearer {token}\r\nConnection: close\r\n"
+ )
+ if file is not None:
+ headers += (
+ f"Content-Length: {size}\r\nContent-Type: application/octet-stream\r\n"
+ f"X-Extraction-Version: 1\r\nX-Extraction-Format: {format}\r\n"
+ )
+ writer.write((headers + "\r\n").encode("ascii"))
+ await writer.drain()
+ remaining = size
+ while remaining:
+ chunk = await file.read(min(CHUNK_BYTES, remaining))
+ if not chunk:
+ raise ExtractionError("INVALID_REQUEST")
+ digest.update(chunk)
+ writer.write(chunk)
+ await writer.drain()
+ remaining -= len(chunk)
+ if file is not None and await file.read(1):
+ raise ExtractionError("INPUT_LIMIT")
+ try:
+ header = await reader.readuntil(b"\r\n\r\n")
+ except (asyncio.LimitOverrunError, asyncio.IncompleteReadError):
+ raise ExtractionError("INVALID_RESPONSE") from None
+ if len(header) > HEADER_BYTES:
+ raise ExtractionError("INVALID_RESPONSE")
+ lines = header[:-4].split(b"\r\n")
+ if not re.fullmatch(rb"HTTP/1\.1 [1-5][0-9]{2} [\x20-\x7e]*", lines[0]):
+ raise ExtractionError("INVALID_RESPONSE")
+ status = int(lines[0][9:12])
+ fields = {}
+ for line in lines[1:]:
+ key, sep, val = line.partition(b":")
+ key = key.lower()
+ if not sep or not re.fullmatch(rb"[a-z0-9-]+", key) or key in fields:
+ raise ExtractionError("INVALID_RESPONSE")
+ if any(byte < 32 or byte > 126 for byte in val):
+ raise ExtractionError("INVALID_RESPONSE")
+ fields[key] = val.strip()
+ length = fields.get(b"content-length", b"")
+ if (
+ not re.fullmatch(rb"[1-9][0-9]{0,7}", length)
+ or int(length) > maximum
+ or fields.get(b"content-type") != b"application/json; charset=utf-8"
+ or fields.get(b"connection", b"").lower() != b"close"
+ or b"transfer-encoding" in fields
+ or b"content-encoding" in fields
+ ):
+ raise ExtractionError("INVALID_RESPONSE")
+ try:
+ body = await reader.readexactly(int(length))
+ if await reader.read(1):
+ raise ExtractionError("INVALID_RESPONSE")
+ except asyncio.IncompleteReadError:
+ raise ExtractionError("INVALID_RESPONSE") from None
+ value = _json(body)
+ if status != 200:
+ if (
+ _keys(value, ("error",))
+ and _keys(value["error"], ("code",))
+ and type(value["error"]["code"]) is str
+ and ERROR_STATUS.get(value["error"]["code"]) == status
+ ):
+ code = value["error"]["code"]
+ # Service authentication is not the RAG caller's authentication.
+ raise ExtractionError(
+ "UNAVAILABLE" if code == "UNAUTHORIZED" else code
+ )
+ raise ExtractionError("INVALID_RESPONSE")
+ return value, digest.hexdigest()
+ finally:
+ # Abort instead of flushing pending input on timeout/cancellation.
+ writer.transport.abort()
+ closed = asyncio.create_task(writer.wait_closed())
+ cancelled = False
+ while not closed.done():
+ try:
+ await asyncio.shield(closed)
+ except asyncio.CancelledError:
+ cancelled = True
+ except OSError:
+ break
+ try:
+ closed.result()
+ except OSError:
+ pass
+ if cancelled:
+ raise asyncio.CancelledError()
+
+ async def extract(self, file_path: str, format: str):
+ if format not in {"pdf", "docx"}:
+ raise ExtractionError("UNSUPPORTED_CONTENT")
+ if self._admitted >= 2:
+ raise ExtractionError("BUSY")
+ self._admitted += 1
+ try:
+ async with asyncio.timeout(self.settings.deadline):
+ async with aiofiles.open(file_path, "rb") as file:
+ stat = await asyncio.to_thread(os.fstat, file.fileno())
+ if not 0 < stat.st_size <= INPUT_BYTES:
+ raise ExtractionError("INPUT_LIMIT")
+ token = await asyncio.to_thread(
+ _read_token, self.settings.token_file
+ )
+ # No cache: authenticate and check the current service before upload.
+ capabilities, _ = await self._request("GET", token)
+ limits = validate_capabilities(capabilities, format, stat.st_size)
+ result, digest = await self._request(
+ "POST",
+ token,
+ file=file,
+ size=stat.st_size,
+ format=format,
+ maximum=limits["resultBytes"],
+ )
+ return validate_result(result, format, digest, limits)
+ except TimeoutError:
+ raise ExtractionError("DEADLINE") from None
+ except OSError:
+ raise ExtractionError("UNAVAILABLE") from None
+ finally:
+ self._admitted -= 1
diff --git a/app/utils/document_loader.py b/app/utils/document_loader.py
index 5a5576b5..93284d1f 100644
--- a/app/utils/document_loader.py
+++ b/app/utils/document_loader.py
@@ -4,15 +4,14 @@
import codecs
import tempfile
-from typing import Iterator, List, Optional
+from typing import List, Optional
import chardet
from langchain_core.documents import Document
-from app.config import known_source_ext, PDF_EXTRACT_IMAGES, CHUNK_OVERLAP, logger
+from app.config import known_source_ext, CHUNK_OVERLAP, logger
from langchain_community.document_loaders import (
TextLoader,
- PyPDFLoader,
CSVLoader,
Docx2txtLoader,
UnstructuredEPubLoader,
@@ -32,6 +31,18 @@
)
+def extraction_format(filename: str, content_type: str) -> Optional[str]:
+ """Identify migrated formats independently of text/Markdown MIME overrides."""
+ extension = filename.rsplit(".", 1)[-1].lower()
+ if extension == "pdf" or content_type == "application/pdf":
+ return "pdf"
+ if extension == "docx" or content_type == (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+ ):
+ return "docx"
+ return None
+
+
def detect_file_encoding(filepath: str) -> str:
"""
Detect the encoding of a file using BOM markers and chardet for broader support.
@@ -82,7 +93,7 @@ def get_loader(
filepath: str,
raw_text: bool = False,
):
- """Get the appropriate document loader based on file type and\or content type.
+ """Get the appropriate document loader based on file type and/or content type.
When ``raw_text`` is True, text-formatted files (e.g. Markdown) are loaded
verbatim with :class:`TextLoader` so their original formatting is
@@ -91,14 +102,14 @@ def get_loader(
(``raw_text=False``) so semantic loaders continue to strip formatting for
better vector search quality.
"""
+ if extraction_format(filename, file_content_type) is not None:
+ raise ValueError("PDF/DOCX require the asynchronous extraction service")
file_ext = filename.split(".")[-1].lower()
known_type = True
# File Content Type reference:
# ref.: https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/MIME_types/Common_types
- if file_ext == "pdf" or file_content_type == "application/pdf":
- loader = SafePyPDFLoader(filepath, extract_images=PDF_EXTRACT_IMAGES)
- elif file_ext == "csv" or file_content_type == "text/csv":
+ if file_ext == "csv" or file_content_type == "text/csv":
# Detect encoding for CSV files
encoding = detect_file_encoding(filepath)
@@ -158,10 +169,7 @@ def get_loader(
loader = UnstructuredMarkdownLoader(filepath)
elif file_ext == "epub" or file_content_type == "application/epub+zip":
loader = UnstructuredEPubLoader(filepath)
- elif file_ext in ["doc", "docx"] or file_content_type in [
- "application/msword",
- "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
- ]:
+ elif file_ext == "doc" or file_content_type == "application/msword":
loader = Docx2txtLoader(filepath)
elif file_ext in ["xls", "xlsx"] or file_content_type in [
"application/vnd.ms-excel",
@@ -241,52 +249,3 @@ def process_documents(documents: List[Document]) -> str:
processed_text += new_content
return processed_text.strip()
-
-
-class SafePyPDFLoader:
- """
- A wrapper around PyPDFLoader that handles image extraction failures gracefully.
- Falls back to text-only extraction when image extraction fails.
-
- This is a workaround for issues with PyPDFLoader that can occur when extracting images
- from PDFs, which can lead to KeyError exceptions if the PDF is malformed or has unsupported
- image formats. This class attempts to load the PDF with image extraction enabled, and if it
- fails due to a KeyError related to image filters, it falls back to loading the PDF
- without image extraction.
- ref.: https://github.com/langchain-ai/langchain/issues/26652
- """
-
- def __init__(self, filepath: str, extract_images: bool = False):
- self.filepath = filepath
- self.extract_images = extract_images
- self._temp_filepath = None # For compatibility with cleanup function
-
- def lazy_load(self) -> Iterator[Document]:
- """Lazy load PDF documents with automatic fallback on image extraction errors."""
- loader = PyPDFLoader(self.filepath, extract_images=self.extract_images)
-
- if not self.extract_images:
- # No image extraction: no fallback needed, stream directly
- yield from loader.lazy_load()
- return
-
- # extract_images=True: must collect eagerly so that a mid-stream
- # KeyError doesn't leave already-yielded pages duplicated by the
- # fallback (yield from + try/except would deliver partial + full).
- try:
- pages = list(loader.lazy_load())
- except KeyError as e:
- if "/Filter" in str(e):
- logger.warning(
- f"PDF image extraction failed for {self.filepath}, falling back to text-only: {e}"
- )
- fallback_loader = PyPDFLoader(self.filepath, extract_images=False)
- pages = list(fallback_loader.lazy_load())
- else:
- # Re-raise if it's a different error
- raise
- yield from pages
-
- def load(self) -> List[Document]:
- """Load PDF documents with automatic fallback on image extraction errors."""
- return list(self.lazy_load())
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 45f59cc8..37700811 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -1,4 +1,14 @@
services:
+ extraction-init: !override
+ extends:
+ file: extraction-compose.yaml
+ service: extraction-init
+
+ extraction: !override
+ extends:
+ file: extraction-compose.yaml
+ service: extraction
+
db:
image: ankane/pgvector:latest
environment:
@@ -12,6 +22,7 @@ services:
fastapi:
build: .
+ networks: [default, extraction]
environment:
- DB_HOST=db
- DB_PORT=5432
@@ -20,10 +31,19 @@ services:
- "8000:8000"
volumes:
- ./uploads:/app/uploads
+ - extraction-credentials:/run/extraction:ro
depends_on:
- - db
+ db:
+ condition: service_started
+ extraction:
+ condition: service_healthy
env_file:
- .env
volumes:
+ extraction-credentials:
pgdata2:
+
+networks:
+ extraction:
+ internal: true
diff --git a/extraction-compose.yaml b/extraction-compose.yaml
new file mode 100644
index 00000000..d4cbf9ce
--- /dev/null
+++ b/extraction-compose.yaml
@@ -0,0 +1,42 @@
+services:
+ extraction-init:
+ build: &extraction-build
+ context: https://github.com/LibreChat-AI/code-interpreter.git#c7bfe421243c66d4d5d9c6e3d08d276586ac47bf
+ dockerfile: packages/extraction/Dockerfile
+ target: runtime
+ command: ['node', 'dist/provision.js']
+ healthcheck:
+ disable: true
+ user: '10001:10001'
+ read_only: true
+ cap_drop: [ALL]
+ security_opt: [no-new-privileges:true]
+ network_mode: none
+ mem_limit: 128m
+ cpus: 0.25
+ pids_limit: 16
+ environment:
+ EXTRACTION_TOKEN_FILE: /run/extraction/token
+ volumes:
+ - extraction-credentials:/run/extraction
+
+ extraction:
+ build: *extraction-build
+ init: true
+ user: '10001:10001'
+ read_only: true
+ cap_drop: [ALL]
+ security_opt: [no-new-privileges:true]
+ networks: [extraction]
+ depends_on:
+ extraction-init:
+ condition: service_completed_successfully
+ environment:
+ EXTRACTION_TOKEN_FILE: /run/extraction/token
+ mem_limit: 768m
+ cpus: 2
+ pids_limit: 32
+ tmpfs:
+ - /jobs:rw,noexec,nosuid,nodev,size=96m,nr_inodes=1024,uid=10001,gid=10001,mode=0700
+ volumes:
+ - extraction-credentials:/run/extraction:ro
diff --git a/tests/extraction/conftest.py b/tests/extraction/conftest.py
new file mode 100644
index 00000000..4d7b06cb
--- /dev/null
+++ b/tests/extraction/conftest.py
@@ -0,0 +1,186 @@
+import asyncio
+import copy
+import hashlib
+import json
+import subprocess
+
+import pytest
+
+from app.services.extraction_client import ExtractionClient, ExtractionSettings
+
+CAPABILITIES = {
+ "version": 1,
+ "operation": "document.extract-text",
+ "formats": ["pdf", "docx"],
+ "structure": {
+ "pdf": "ordered pages, one-based index",
+ "docx": "document body, paragraphs and tables",
+ },
+ "limitations": {"pdf": "Text only; no OCR", "docx": "Body only"},
+ "limits": {
+ "inputBytes": 10 * 1024 * 1024,
+ "textBytes": 1024 * 1024,
+ "resultBytes": 3 * 1024 * 1024,
+ "deadlineMs": 10000,
+ "maxSegments": 128,
+ "concurrency": 2,
+ "expandedBytes": 32 * 1024 * 1024,
+ "archiveEntries": 512,
+ "tableColumns": 512,
+ },
+}
+
+
+class Processor:
+ def __init__(self, token_file):
+ self.url = None
+ self.token_file = str(token_file)
+ self.token = "ab" * 32
+ token_file.write_text(self.token)
+ self.capabilities = copy.deepcopy(CAPABILITIES)
+ self.requests = []
+ self.mutate = lambda value: value
+ self.raw = None
+ self.error = None
+ self.hold = False
+ self.hold_caps = False
+ self.caps_entered = asyncio.Event()
+ self.before_response = None
+ self.entered = asyncio.Event()
+ self.closed = asyncio.Event()
+ self.active = set()
+ self.texts = ["First\x00 page café", "Second page"]
+
+ def client(self, deadline=10):
+ return ExtractionClient(
+ ExtractionSettings(
+ url=self.url, token_file=self.token_file, deadline=deadline
+ )
+ )
+
+ async def handle(self, reader, writer):
+ task = asyncio.current_task()
+ self.active.add(task)
+ try:
+ headers = (await reader.readuntil(b"\r\n\r\n")).decode("ascii")
+ lines = headers.split("\r\n")
+ method, path, _ = lines[0].split()
+ fields = dict(line.split(": ", 1) for line in lines[1:] if line)
+ length = int(fields.get("Content-Length", "0"))
+ body = await reader.readexactly(length)
+ self.requests.append((method, path, fields, body))
+ if fields.get("Authorization") != f"Bearer {self.token}":
+ value = {"error": {"code": "UNAUTHORIZED"}}
+ status = 401
+ elif path == "/v1/capabilities":
+ self.caps_entered.set()
+ if self.hold_caps:
+ await reader.read()
+ self.closed.set()
+ return
+ value = self.capabilities
+ status = 200
+ else:
+ self.entered.set()
+ if self.hold:
+ await reader.read()
+ self.closed.set()
+ return
+ format = fields["X-Extraction-Format"]
+ texts = (
+ self.texts if format == "pdf" else ["Body paragraph\nLeft\tRight"]
+ )
+ value = self.mutate(
+ {
+ "version": 1,
+ "operation": "document.extract-text",
+ "format": format,
+ "inputSha256": hashlib.sha256(body).hexdigest(),
+ "textBytes": sum(len(text.encode("utf-8")) for text in texts),
+ "segments": [
+ {
+ "kind": "page" if format == "pdf" else "document",
+ "index": i,
+ "text": text,
+ }
+ for i, text in enumerate(texts, 1)
+ ],
+ }
+ )
+ status = 200
+ if self.error:
+ status, code = self.error
+ value = {"error": {"code": code}}
+ if self.raw is not None:
+ writer.write(self.raw)
+ await writer.drain()
+ return
+ if self.before_response is not None:
+ self.before_response()
+ wire = json.dumps(value).encode("utf-8")
+ writer.write(
+ f"HTTP/1.1 {status} Result\r\nContent-Type: application/json; charset=utf-8\r\n"
+ f"Content-Length: {len(wire)}\r\nConnection: close\r\n\r\n".encode(
+ "ascii"
+ )
+ + wire
+ )
+ await writer.drain()
+ except (ConnectionError, asyncio.IncompleteReadError):
+ pass
+ finally:
+ writer.close()
+ try:
+ await writer.wait_closed()
+ except ConnectionError:
+ pass
+ self.active.discard(task)
+
+
+@pytest.fixture
+async def processor(tmp_path):
+ service = Processor(tmp_path / "token")
+ server = await asyncio.start_server(service.handle, host="127.0.0.1", port=0)
+ service.url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}"
+ async with server:
+ yield service
+ for task in list(service.active):
+ task.cancel()
+ await asyncio.gather(*service.active, return_exceptions=True)
+
+
+@pytest.fixture
+def input_file(tmp_path):
+ file = tmp_path / "input"
+ file.write_bytes(b"caller-owned opaque document bytes")
+ return str(file)
+
+
+@pytest.fixture(scope="session")
+def tls_certificate(tmp_path_factory):
+ directory = tmp_path_factory.mktemp("extraction-tls")
+ certificate, key = directory / "certificate.pem", directory / "key.pem"
+ subprocess.run(
+ [
+ "openssl",
+ "req",
+ "-x509",
+ "-newkey",
+ "rsa:2048",
+ "-nodes",
+ "-days",
+ "1",
+ "-keyout",
+ str(key),
+ "-out",
+ str(certificate),
+ "-subj",
+ "/CN=localhost",
+ "-addext",
+ "subjectAltName=DNS:localhost",
+ ],
+ check=True,
+ capture_output=True,
+ timeout=15,
+ )
+ return certificate, key
diff --git a/tests/extraction/fixtures/.gitattributes b/tests/extraction/fixtures/.gitattributes
new file mode 100644
index 00000000..d5a03a01
--- /dev/null
+++ b/tests/extraction/fixtures/.gitattributes
@@ -0,0 +1,2 @@
+*.pdf binary
+*.docx binary
diff --git a/tests/extraction/fixtures/body.docx b/tests/extraction/fixtures/body.docx
new file mode 100644
index 00000000..77b5f53b
Binary files /dev/null and b/tests/extraction/fixtures/body.docx differ
diff --git a/tests/extraction/fixtures/pages.pdf b/tests/extraction/fixtures/pages.pdf
new file mode 100644
index 00000000..1e7acc24
Binary files /dev/null and b/tests/extraction/fixtures/pages.pdf differ
diff --git a/tests/extraction/test_client.py b/tests/extraction/test_client.py
new file mode 100644
index 00000000..5fbd403b
--- /dev/null
+++ b/tests/extraction/test_client.py
@@ -0,0 +1,283 @@
+import asyncio
+import json
+from pathlib import Path
+
+import pytest
+
+from app.services.extraction_client import (
+ ExtractionClient,
+ ExtractionError,
+ ExtractionSettings,
+ INPUT_BYTES,
+ TEXT_BYTES,
+)
+
+
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_raw_upload_and_valid_segments(processor, input_file, format):
+ segments = await processor.client().extract(input_file, format)
+ assert len(segments) == (2 if format == "pdf" else 1)
+ assert [segment["index"] for segment in segments] == list(
+ range(1, len(segments) + 1)
+ )
+ assert len(processor.requests) == 2
+ method, path, headers, body = processor.requests[1]
+ assert (method, path) == ("POST", "/v1/extract-text")
+ assert body == Path(input_file).read_bytes()
+ assert headers == {
+ "Host": processor.url.removeprefix("http://"),
+ "Authorization": f"Bearer {processor.token}",
+ "Connection": "close",
+ "Content-Length": str(len(body)),
+ "Content-Type": "application/octet-stream",
+ "X-Extraction-Version": "1",
+ "X-Extraction-Format": format,
+ }
+
+
+@pytest.mark.parametrize("size", [0, INPUT_BYTES + 1])
+async def test_input_bounds_before_any_connection(processor, input_file, size):
+ with open(input_file, "wb") as file:
+ file.truncate(size)
+ with pytest.raises(ExtractionError, match="INPUT_LIMIT"):
+ await processor.client().extract(input_file, "pdf")
+ assert not processor.requests
+
+
+@pytest.mark.parametrize(
+ "change",
+ [
+ lambda r: {**r, "version": True},
+ lambda r: {**r, "version": 2},
+ lambda r: {**r, "operation": "execute"},
+ lambda r: {**r, "format": "docx"},
+ lambda r: {**r, "inputSha256": "0" * 64},
+ lambda r: {**r, "textBytes": True},
+ lambda r: {**r, "textBytes": r["textBytes"] + 1},
+ lambda r: {**r, "textBytes": TEXT_BYTES + 1},
+ lambda r: {**r, "metadata": {"user_id": "attacker"}},
+ lambda r: {**r, "segments": []},
+ lambda r: {**r, "segments": r["segments"] * 65},
+ lambda r: {**r, "segments": list(reversed(r["segments"]))},
+ lambda r: {**r, "segments": [{**r["segments"][0], "index": True}]},
+ lambda r: {**r, "segments": [{**r["segments"][0], "index": 0}]},
+ lambda r: {**r, "segments": [{**r["segments"][0], "kind": "document"}]},
+ lambda r: {**r, "segments": [{**r["segments"][0], "text": "\ud800"}]},
+ lambda r: {**r, "segments": [{**r["segments"][0], "text": None}]},
+ lambda r: {**r, "segments": [{**r["segments"][0], "file_id": "forged"}]},
+ lambda r: {
+ **r,
+ "textBytes": 1,
+ "segments": [{"kind": "page", "index": 1, "text": " "}],
+ },
+ ],
+)
+async def test_strict_result_rejection(processor, input_file, change):
+ processor.mutate = change
+ with pytest.raises(ExtractionError, match="INVALID_RESPONSE"):
+ await processor.client().extract(input_file, "pdf")
+
+
+@pytest.mark.parametrize(
+ "change",
+ [
+ lambda c: c.update(version=True),
+ lambda c: c.update(operation="other"),
+ lambda c: c.update(formats=["docx"]),
+ lambda c: c.update(formats=["pdf", "pdf"]),
+ lambda c: c.update(structure={"pdf": "pages"}),
+ lambda c: c["limitations"].update(pdf="\ud800"),
+ lambda c: c["limits"].update(inputBytes=True),
+ lambda c: c["limits"].update(resultBytes=4 * 1024 * 1024),
+ lambda c: c["limits"].update(concurrency=0),
+ lambda c: c.update(extra="untrusted"),
+ ],
+)
+async def test_incompatible_capabilities_never_upload(processor, input_file, change):
+ change(processor.capabilities)
+ with pytest.raises(ExtractionError, match="INCOMPATIBLE_SERVICE"):
+ await processor.client().extract(input_file, "pdf")
+ assert len(processor.requests) == 1
+
+
+def response(body=b"{}", fields=None):
+ fields = fields or f"Content-Length: {len(body)}\r\n".encode()
+ return (
+ b"HTTP/1.1 200 OK\r\nContent-Type: application/json; charset=utf-8\r\n"
+ b"Connection: close\r\n" + fields + b"\r\n" + body
+ )
+
+
+@pytest.mark.parametrize(
+ "raw",
+ [
+ response(b"\xff"),
+ response(b'{"version":1,"version":1}'),
+ response(b'{"textBytes":NaN}'),
+ response(b"[]"),
+ response(b"x" * 20, b"Content-Length: 21\r\n"),
+ response(b"{}trailing", b"Content-Length: 2\r\n"),
+ response(b"{}", b"Content-Length: 3145729\r\n"),
+ response(b"{}", b"Content-Length: 2\r\nContent-Length: 2\r\n"),
+ response(b"{}", b"Content-Length: 2\r\nTransfer-Encoding: chunked\r\n"),
+ response(b"{}", b"Content-Length: 2\r\nContent-Encoding: gzip\r\n"),
+ response(b"{}", b"Content-Length: 2\r\nX-Pad: " + b"x" * 4096 + b"\r\n"),
+ b"HTTP/1.1 302 Redirect\r\nLocation: http://example.com\r\n\r\n",
+ ],
+)
+async def test_bounded_strict_http_and_json(processor, input_file, raw):
+ processor.raw = raw
+ with pytest.raises(ExtractionError, match="INVALID_RESPONSE"):
+ await processor.client().extract(input_file, "pdf")
+
+
+@pytest.mark.parametrize(
+ "status,code",
+ [
+ (400, "INVALID_REQUEST"),
+ (413, "INPUT_LIMIT"),
+ (413, "STRUCTURE_LIMIT"),
+ (422, "UNSUPPORTED_CONTENT"),
+ (429, "BUSY"),
+ (503, "UNAVAILABLE"),
+ (504, "DEADLINE"),
+ ],
+)
+async def test_safe_error_codes(processor, input_file, status, code):
+ processor.error = (status, code)
+ with pytest.raises(ExtractionError) as caught:
+ await processor.client().extract(input_file, "pdf")
+ assert (caught.value.code, caught.value.status_code) == (code, status)
+
+
+async def test_deadline_closes_transport(processor, input_file):
+ processor.hold = True
+ client = processor.client(deadline=0.05)
+ with pytest.raises(ExtractionError, match="DEADLINE"):
+ await client.extract(input_file, "pdf")
+ await asyncio.wait_for(processor.closed.wait(), 1)
+ assert client._admitted == 0
+
+
+async def test_cancellation_closes_transport_and_releases_admission(
+ processor, input_file
+):
+ processor.hold = True
+ client = processor.client()
+ task = asyncio.create_task(client.extract(input_file, "pdf"))
+ await asyncio.wait_for(processor.entered.wait(), 1)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ await asyncio.wait_for(processor.closed.wait(), 1)
+ assert client._admitted == 0
+
+
+async def test_client_admission_is_fail_fast(processor, input_file):
+ processor.hold = True
+ client = processor.client()
+ tasks = [asyncio.create_task(client.extract(input_file, "pdf")) for _ in range(2)]
+ while len([r for r in processor.requests if r[0] == "POST"]) < 2:
+ await asyncio.sleep(0)
+ try:
+ with pytest.raises(ExtractionError, match="BUSY"):
+ await client.extract(input_file, "pdf")
+ finally:
+ for task in tasks:
+ task.cancel()
+ await asyncio.gather(*tasks, return_exceptions=True)
+ assert client._admitted == 0
+
+
+async def test_capabilities_rechecked_after_service_change(processor, input_file):
+ client = processor.client()
+ await client.extract(input_file, "pdf")
+ processor.capabilities["version"] = 2
+ with pytest.raises(ExtractionError, match="INCOMPATIBLE_SERVICE"):
+ await client.extract(input_file, "pdf")
+ assert [r[0] for r in processor.requests] == ["GET", "POST", "GET"]
+
+
+async def test_outage_is_safe(input_file, processor, monkeypatch):
+ async def unavailable(*args, **kwargs):
+ raise ConnectionRefusedError("private endpoint")
+
+ monkeypatch.setattr(asyncio, "open_connection", unavailable)
+ client = processor.client()
+ with pytest.raises(ExtractionError, match="UNAVAILABLE"):
+ await client.extract(input_file, "pdf")
+
+
+def test_default_settings_require_service_without_environment(monkeypatch):
+ for key in (
+ "RAG_EXTRACTION_ENABLED",
+ "RAG_EXTRACTION_FORMATS",
+ "RAG_EXTRACTION_SOCKET",
+ "RAG_EXTRACTION_URL",
+ "RAG_EXTRACTION_TOKEN_FILE",
+ "RAG_EXTRACTION_DEADLINE_SECONDS",
+ ):
+ monkeypatch.delenv(key, raising=False)
+ settings = ExtractionSettings.from_env()
+ assert settings.url == "http://extraction:8000" and settings.deadline == 10
+ assert settings.token_file == "/run/extraction/token"
+ assert not hasattr(settings, "enabled") and not hasattr(settings, "formats")
+
+
+@pytest.mark.parametrize(
+ "key,value",
+ [
+ ("RAG_EXTRACTION_SOCKET", "/socket/old.sock"),
+ ("RAG_EXTRACTION_SOCKET", ""),
+ ("RAG_EXTRACTION_ENABLED", "false"),
+ ("RAG_EXTRACTION_ENABLED", "true"),
+ ("RAG_EXTRACTION_ENABLED", ""),
+ ("RAG_EXTRACTION_FORMATS", "docx"),
+ ("RAG_EXTRACTION_FORMATS", "pdf,docx"),
+ ("RAG_EXTRACTION_FORMATS", ""),
+ ],
+)
+def test_obsolete_opt_out_configuration_is_rejected(monkeypatch, key, value):
+ monkeypatch.setenv(key, value)
+ with pytest.raises(ValueError, match="obsolete"):
+ ExtractionSettings.from_env()
+
+
+@pytest.mark.parametrize(
+ "kwargs",
+ [
+ {"token_file": "relative"},
+ {"token_file": ""},
+ {"token_file": "/bad\x00path"},
+ {"deadline": 0},
+ {"deadline": 11},
+ {"deadline": True},
+ {"deadline": float("nan")},
+ {"deadline": float("inf")},
+ ],
+)
+def test_required_config_is_validated(kwargs):
+ with pytest.raises(ValueError):
+ ExtractionSettings(**kwargs)
+
+
+@pytest.mark.parametrize(
+ "key,value",
+ [
+ ("RAG_EXTRACTION_URL", "relative"),
+ ("RAG_EXTRACTION_TOKEN_FILE", "relative"),
+ ("RAG_EXTRACTION_DEADLINE_SECONDS", "stale"),
+ ],
+)
+def test_invalid_environment_never_disables_extraction(monkeypatch, key, value):
+ monkeypatch.setenv(key, value)
+ with pytest.raises(ValueError):
+ ExtractionSettings.from_env()
+
+
+@pytest.mark.parametrize(
+ "kwargs", [{"enabled": False}, {"formats": frozenset({"docx"})}]
+)
+def test_client_constructor_cannot_opt_out(kwargs):
+ with pytest.raises(TypeError):
+ ExtractionSettings(**kwargs)
diff --git a/tests/extraction/test_deployment.py b/tests/extraction/test_deployment.py
new file mode 100644
index 00000000..bd93ebac
--- /dev/null
+++ b/tests/extraction/test_deployment.py
@@ -0,0 +1,183 @@
+"""Contracts for the processor required by both standard RAG deployments."""
+
+import json
+import shutil
+import subprocess
+import uuid
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[2]
+PROCESSOR_REF = "c7bfe421243c66d4d5d9c6e3d08d276586ac47bf"
+PROCESSOR_CONTEXT = (
+ f"https://github.com/LibreChat-AI/code-interpreter.git#{PROCESSOR_REF}"
+)
+
+
+def stage_compose_files(directory):
+ for name in ("docker-compose.yaml", "api-compose.yaml", "extraction-compose.yaml"):
+ shutil.copyfile(ROOT / name, directory / name)
+ # Never read deployment credentials while rendering or qualifying the service.
+ (directory / ".env").write_text("")
+ return directory
+
+
+def compose_command(directory, project, *files):
+ command = ["docker", "compose", "--env-file", "/dev/null", "-p", project]
+ for file in files:
+ command.extend(["-f", str(directory / file)])
+ return command
+
+
+def render(directory, *files):
+ command = compose_command(directory, f"rag-contract-{uuid.uuid4().hex}", *files)
+ result = subprocess.run(
+ [*command, "config", "--no-env-resolution", "--format", "json"],
+ capture_output=True,
+ text=True,
+ timeout=15,
+ )
+ if result.returncode != 0:
+ raise AssertionError(
+ f"Compose render failed ({result.returncode}):\n{result.stderr}"
+ )
+ return json.loads(result.stdout)
+
+
+def assert_build(service):
+ assert service["build"]["context"] == PROCESSOR_CONTEXT
+ assert service["build"]["dockerfile"] == "packages/extraction/Dockerfile"
+ assert service["build"]["target"] == "runtime"
+
+
+def assert_credentials_mount(service, readonly):
+ mount = next(
+ item for item in service["volumes"] if item["target"] == "/run/extraction"
+ )
+ assert mount["type"] == "volume" and mount["source"] == "extraction-credentials"
+ assert mount.get("read_only", False) is readonly
+
+
+def assert_sandbox(service):
+ assert service["read_only"] is True
+ assert service["cap_drop"] == ["ALL"]
+ assert service["security_opt"] == ["no-new-privileges:true"]
+ for key in (
+ "env_file",
+ "ports",
+ "expose",
+ "profiles",
+ "privileged",
+ "cap_add",
+ "entrypoint",
+ "container_name",
+ "group_add",
+ ):
+ assert not service.get(key), key
+
+
+def assert_hardened_processor(service):
+ assert_build(service)
+ assert_sandbox(service)
+ assert service["init"] is True and service["user"] == "10001:10001"
+ assert set(service["networks"]) == {"extraction"}
+ assert "network_mode" not in service
+ assert (
+ service["depends_on"]["extraction-init"]["condition"]
+ == "service_completed_successfully"
+ )
+ assert not service.get("healthcheck")
+ assert service["environment"] == {"EXTRACTION_TOKEN_FILE": "/run/extraction/token"}
+ assert int(service["mem_limit"]) == 768 * 1024 * 1024
+ assert float(service["cpus"]) == 2 and int(service["pids_limit"]) == 32
+ assert service["tmpfs"] == [
+ "/jobs:rw,noexec,nosuid,nodev,size=96m,nr_inodes=1024,uid=10001,gid=10001,mode=0700"
+ ]
+ assert len(service["volumes"]) == 1
+ assert_credentials_mount(service, True)
+ assert not service.get("command")
+
+
+def assert_provisioning(config):
+ service = config["services"]["extraction-init"]
+ assert_build(service)
+ assert_sandbox(service)
+ assert service["build"] == config["services"]["extraction"]["build"]
+ assert service["command"] == ["node", "dist/provision.js"]
+ assert service["user"] == "10001:10001" and service["network_mode"] == "none"
+ assert service["healthcheck"] == {"disable": True}
+ assert int(service["mem_limit"]) == 128 * 1024 * 1024
+ assert float(service["cpus"]) == 0.25 and int(service["pids_limit"]) == 16
+ assert service["environment"] == {"EXTRACTION_TOKEN_FILE": "/run/extraction/token"}
+ assert len(service["volumes"]) == 1
+ assert_credentials_mount(service, False)
+ assert config["networks"]["extraction"]["internal"] is True
+ assert "extraction-credentials" in config["volumes"]
+ assert "extraction-socket" not in config["volumes"]
+
+
+def assert_caller(app):
+ assert not app.get("group_add") and not app.get("profiles")
+ assert set(app["networks"]) == {"default", "extraction"}
+ assert app["depends_on"]["extraction"]["condition"] == "service_healthy"
+ assert_credentials_mount(app, True)
+ assert (
+ not {
+ "RAG_EXTRACTION_ENABLED",
+ "RAG_EXTRACTION_FORMATS",
+ "RAG_EXTRACTION_SOCKET",
+ }
+ & app.get("environment", {}).keys()
+ )
+
+
+@pytest.mark.parametrize(
+ "recipe,caller",
+ [("docker-compose.yaml", "fastapi"), ("api-compose.yaml", "rag_api")],
+)
+def test_standard_deployment_provisions_hardened_processor(tmp_path, recipe, caller):
+ config = render(stage_compose_files(tmp_path), recipe)
+ assert_hardened_processor(config["services"]["extraction"])
+ assert_provisioning(config)
+ assert_caller(config["services"][caller])
+
+
+@pytest.mark.parametrize(
+ "recipes",
+ [
+ ("docker-compose.yaml", "api-compose.yaml"),
+ ("api-compose.yaml", "docker-compose.yaml"),
+ ],
+)
+def test_combined_recipes_share_provisioning_and_processor(tmp_path, recipes):
+ config = render(stage_compose_files(tmp_path), *recipes)
+ assert set(config["services"]) == {
+ "db",
+ "fastapi",
+ "rag_api",
+ "extraction-init",
+ "extraction",
+ }
+ assert_hardened_processor(config["services"]["extraction"])
+ assert_provisioning(config)
+ for caller in ("fastapi", "rag_api"):
+ assert_caller(config["services"][caller])
+
+
+def test_render_failure_includes_stderr_without_dumping_model(monkeypatch, tmp_path):
+ diagnostic = "services.extraction.security_opt items at 0 and 1 are equal"
+ monkeypatch.setattr(
+ subprocess,
+ "run",
+ lambda command, **kwargs: subprocess.CompletedProcess(
+ command,
+ 1,
+ stdout="do not dump the rendered application model",
+ stderr=diagnostic,
+ ),
+ )
+ with pytest.raises(AssertionError, match="security_opt") as caught:
+ render(tmp_path, "docker-compose.yaml", "api-compose.yaml")
+ assert diagnostic in str(caught.value)
+ assert "do not dump" not in str(caught.value)
diff --git a/tests/extraction/test_http.py b/tests/extraction/test_http.py
new file mode 100644
index 00000000..4fa0164d
--- /dev/null
+++ b/tests/extraction/test_http.py
@@ -0,0 +1,364 @@
+import asyncio
+import os
+import ssl
+import threading
+from pathlib import Path
+
+import pytest
+
+from app.services import extraction_client
+from app.services.extraction_client import ExtractionError, ExtractionSettings
+
+
+@pytest.mark.parametrize(
+ "url",
+ [
+ "http://extraction:8000",
+ "http://extraction:8000/",
+ "http://localhost:1234",
+ "http://127.0.0.1:1234",
+ "http://[::1]:1234",
+ "https://processor.example",
+ "https://processor.example:8443/",
+ "https://[2001:db8::1]:8443",
+ ],
+)
+def test_valid_origins(url):
+ assert ExtractionSettings(url=url).url == url
+
+
+@pytest.mark.parametrize(
+ "url",
+ [
+ None,
+ "",
+ "extraction:8000",
+ "unix:///socket/extraction.sock",
+ "file:///run/extraction/token",
+ "ftp://extraction:8000",
+ "http://extraction",
+ "http://extraction:8001",
+ "http://remote.example:8000",
+ "http://10.0.0.1:8000",
+ "http://192.168.1.2:8000",
+ "http://127.0.0.2:8000",
+ "http://localhost.evil:8000",
+ "https://user:secret@processor.example",
+ "https://secret@processor.example",
+ "https://processor.example/path",
+ "https://processor.example//",
+ "https://processor.example?token=secret",
+ "https://processor.example?",
+ "https://processor.example#secret",
+ "https://processor.example#",
+ "https://processor.example:0",
+ "https://processor.example:65536",
+ "https://processor.example:secret",
+ "https://processor.example:",
+ "https://[broken",
+ "https://[::1]ignored",
+ "https://[::1]:0443",
+ "https://[v1.localhost]",
+ "https://processor.example:0443",
+ "https://processor.example\\@localhost",
+ " https://processor.example",
+ "https://processor.example\n",
+ "https://processor.example\x00",
+ "https://processor.example\t:443",
+ "https://procéssor.example",
+ ],
+)
+def test_invalid_origin_is_safe(url):
+ with pytest.raises(ValueError) as caught:
+ ExtractionSettings(url=url)
+ assert "secret" not in str(caught.value)
+ assert "processor.example" not in str(caught.value)
+
+
+@pytest.mark.parametrize("phase", ["capabilities", "extract"])
+async def test_both_operations_require_service_auth(processor, input_file, phase):
+ if phase == "capabilities":
+ Path(processor.token_file).write_text("cd" * 32)
+ else:
+ processor.error = (401, "UNAUTHORIZED")
+ with pytest.raises(ExtractionError) as caught:
+ await processor.client().extract(input_file, "pdf")
+ assert (caught.value.code, caught.value.status_code) == ("UNAVAILABLE", 503)
+ assert len(processor.requests) == (1 if phase == "capabilities" else 2)
+
+
+@pytest.mark.parametrize(
+ "token",
+ [b"", b"a" * 63, b"a" * 65, b"a" * 64 + b"\n", b"A" * 64, b"g" * 64, b"\xff" * 64],
+)
+async def test_invalid_credentials_never_connect(processor, input_file, token):
+ Path(processor.token_file).write_bytes(token)
+ with pytest.raises(ExtractionError) as caught:
+ await processor.client().extract(input_file, "pdf")
+ assert (caught.value.code, caught.value.status_code) == ("UNAVAILABLE", 503)
+ assert not processor.requests
+
+
+@pytest.mark.parametrize(
+ "kind", ["missing", "directory", "fifo", "symlink", "oversized"]
+)
+async def test_nonregular_or_missing_credentials_fail_closed(
+ processor, input_file, kind
+):
+ path = Path(processor.token_file)
+ path.unlink()
+ if kind == "directory":
+ path.mkdir()
+ elif kind == "fifo":
+ os.mkfifo(path)
+ elif kind == "symlink":
+ target = path.with_name("other-token")
+ target.write_text(processor.token)
+ path.symlink_to(target)
+ elif kind == "oversized":
+ with path.open("wb") as file:
+ file.truncate(1024 * 1024 * 1024)
+ with pytest.raises(ExtractionError, match="UNAVAILABLE"):
+ await asyncio.wait_for(processor.client().extract(input_file, "pdf"), 1)
+ assert not processor.requests
+
+
+async def test_token_read_is_bounded_nonblocking_and_once_per_job(
+ processor, input_file, monkeypatch
+):
+ original_open, original_read = os.open, os.read
+ opens, reads = [], []
+
+ def capture_open(path, flags, *args, **kwargs):
+ descriptor = original_open(path, flags, *args, **kwargs)
+ if str(path) == processor.token_file:
+ opens.append((descriptor, flags))
+ return descriptor
+
+ def capture_read(descriptor, size):
+ if opens and descriptor == opens[-1][0]:
+ reads.append(size)
+ return original_read(descriptor, size)
+
+ monkeypatch.setattr(os, "open", capture_open)
+ monkeypatch.setattr(os, "read", capture_read)
+ await processor.client().extract(input_file, "pdf")
+ assert len(opens) == 1 and reads == [65]
+ assert opens[0][1] & os.O_NONBLOCK and opens[0][1] & os.O_NOFOLLOW
+ with pytest.raises(OSError):
+ os.fstat(opens[0][0])
+
+
+async def test_token_rotation_is_visible_on_next_job(processor, input_file):
+ client = processor.client()
+ await client.extract(input_file, "pdf")
+ replacement = Path(processor.token_file).with_name("rotated-token")
+ processor.token = "cd" * 32
+ replacement.write_text(processor.token)
+ replacement.replace(processor.token_file)
+ await client.extract(input_file, "pdf")
+ assert [request[2]["Authorization"] for request in processor.requests] == [
+ "Bearer " + "ab" * 32,
+ "Bearer " + "ab" * 32,
+ "Bearer " + "cd" * 32,
+ "Bearer " + "cd" * 32,
+ ]
+
+
+async def test_proxy_environment_is_not_used(processor, input_file, monkeypatch):
+ for key in (
+ "HTTP_PROXY",
+ "HTTPS_PROXY",
+ "ALL_PROXY",
+ "http_proxy",
+ "https_proxy",
+ "all_proxy",
+ ):
+ monkeypatch.setenv(key, "http://secret@127.0.0.1:1")
+ await processor.client().extract(input_file, "pdf")
+ assert len(processor.requests) == 2
+
+
+async def test_credential_read_is_inside_deadline(processor, input_file, monkeypatch):
+ entered, release, finished = threading.Event(), threading.Event(), threading.Event()
+ read = extraction_client._read_token
+
+ def delayed_read(path):
+ entered.set()
+ try:
+ release.wait(2)
+ return read(path)
+ finally:
+ finished.set()
+
+ monkeypatch.setattr(extraction_client, "_read_token", delayed_read)
+ client = processor.client(deadline=0.05)
+ try:
+ with pytest.raises(ExtractionError, match="DEADLINE"):
+ await client.extract(input_file, "pdf")
+ assert entered.is_set() and not processor.requests and client._admitted == 0
+ finally:
+ release.set()
+ assert await asyncio.to_thread(finished.wait, 1)
+
+
+async def test_connect_is_inside_deadline(processor, input_file, monkeypatch):
+ entered, closed = asyncio.Event(), asyncio.Event()
+
+ async def delayed_connect(*args, **kwargs):
+ entered.set()
+ try:
+ await asyncio.Future()
+ finally:
+ closed.set()
+
+ monkeypatch.setattr(asyncio, "open_connection", delayed_connect)
+ with pytest.raises(ExtractionError, match="DEADLINE"):
+ await processor.client(deadline=0.05).extract(input_file, "pdf")
+ assert entered.is_set() and closed.is_set() and not processor.requests
+
+
+async def test_capabilities_are_inside_deadline(processor, input_file):
+ processor.hold_caps = True
+ with pytest.raises(ExtractionError, match="DEADLINE"):
+ await processor.client(deadline=0.05).extract(input_file, "pdf")
+ await asyncio.wait_for(processor.closed.wait(), 1)
+ assert processor.caps_entered.is_set() and len(processor.requests) == 1
+
+
+@pytest.fixture
+async def tls_processor(processor, tls_certificate):
+ certificate, key = tls_certificate
+ context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+ context.load_cert_chain(certificate, key)
+ server = await asyncio.start_server(processor.handle, "127.0.0.1", 0, ssl=context)
+ processor.url = f"https://localhost:{server.sockets[0].getsockname()[1]}"
+ async with server:
+ yield processor
+
+
+async def test_https_uses_standard_trust_and_hostname(
+ tls_processor, tls_certificate, input_file, monkeypatch
+):
+ monkeypatch.setenv("SSL_CERT_FILE", str(tls_certificate[0]))
+ client = tls_processor.client()
+ assert client._ssl.verify_mode == ssl.CERT_REQUIRED and client._ssl.check_hostname
+ await client.extract(input_file, "pdf")
+ assert len(tls_processor.requests) == 2
+
+
+@pytest.mark.parametrize("failure", ["untrusted", "hostname"])
+async def test_https_never_skips_verification(
+ tls_processor, tls_certificate, input_file, monkeypatch, failure
+):
+ if failure == "hostname":
+ monkeypatch.setenv("SSL_CERT_FILE", str(tls_certificate[0]))
+ tls_processor.url = tls_processor.url.replace("localhost", "127.0.0.1")
+ with pytest.raises(ExtractionError) as caught:
+ await tls_processor.client().extract(input_file, "pdf")
+ assert (caught.value.code, caught.value.status_code) == ("UNAVAILABLE", 503)
+ assert not tls_processor.requests
+
+
+@pytest.mark.parametrize("cancel", [False, True])
+async def test_https_abort_closes_real_transport(
+ tls_processor, tls_certificate, input_file, monkeypatch, cancel
+):
+ monkeypatch.setenv("SSL_CERT_FILE", str(tls_certificate[0]))
+ tls_processor.hold = True
+ client = tls_processor.client(deadline=0.1 if not cancel else 10)
+ task = asyncio.create_task(client.extract(input_file, "pdf"))
+ await asyncio.wait_for(tls_processor.entered.wait(), 1)
+ if cancel:
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError if cancel else ExtractionError):
+ await task
+ await asyncio.wait_for(tls_processor.closed.wait(), 1)
+ assert client._admitted == 0
+
+
+async def test_tls_handshake_deadline_aborts_connection(processor, input_file):
+ connected, closed = asyncio.Event(), asyncio.Event()
+
+ async def never_tls(reader, writer):
+ connected.set()
+ try:
+ await reader.read()
+ closed.set()
+ finally:
+ writer.close()
+ await writer.wait_closed()
+
+ server = await asyncio.start_server(never_tls, "127.0.0.1", 0)
+ async with server:
+ processor.url = f"https://localhost:{server.sockets[0].getsockname()[1]}"
+ with pytest.raises(ExtractionError, match="DEADLINE"):
+ await processor.client(deadline=0.05).extract(input_file, "pdf")
+ assert connected.is_set()
+ await asyncio.wait_for(closed.wait(), 1)
+ assert not processor.requests
+
+
+async def test_capabilities_and_result_share_one_deadline(
+ processor, input_file, monkeypatch
+):
+ client = processor.client(deadline=0.05)
+ request = client._request
+ remaining = []
+
+ async def delay(method, *args, **kwargs):
+ if method == "POST":
+ remaining.append(asyncio.current_task())
+ await asyncio.sleep(0.05)
+ else:
+ await asyncio.sleep(0.03)
+ return await request(method, *args, **kwargs)
+
+ monkeypatch.setattr(client, "_request", delay)
+ with pytest.raises(ExtractionError, match="DEADLINE"):
+ await client.extract(input_file, "pdf")
+ assert remaining and len(processor.requests) == 1
+ assert client._admitted == 0
+
+
+async def test_remote_error_text_and_redirects_are_not_reflected(processor, input_file):
+ secret = b"private-service-details"
+ body = b'{"message":"private-service-details"}'
+ processor.raw = (
+ b"HTTP/1.1 401 Unauthorized\r\nContent-Type: application/json; charset=utf-8\r\n"
+ + f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n".encode()
+ + body
+ )
+ with pytest.raises(ExtractionError) as caught:
+ await processor.client().extract(input_file, "pdf")
+ assert secret.decode() not in str(caught.value)
+ processor.raw = (
+ b"HTTP/1.1 302 Redirect\r\nLocation: https://secret@private.example/\r\n"
+ b"Content-Type: application/json; charset=utf-8\r\nContent-Length: 2\r\n"
+ b"Connection: close\r\n\r\n{}"
+ )
+ with pytest.raises(ExtractionError, match="INVALID_RESPONSE"):
+ await processor.client().extract(input_file, "pdf")
+ assert len(processor.requests) == 4
+
+
+def test_invalid_deadline_environment_does_not_echo_value(monkeypatch):
+ monkeypatch.setenv("RAG_EXTRACTION_DEADLINE_SECONDS", "private-service-details")
+ with pytest.raises(ValueError) as caught:
+ ExtractionSettings.from_env()
+ assert "private-service-details" not in str(caught.value)
+
+
+@pytest.mark.parametrize(
+ "url", ["https://processor.example:8443", "http://localhost:2345"]
+)
+def test_service_environment_uses_only_service_configuration(monkeypatch, url):
+ monkeypatch.setenv("RAG_EXTRACTION_URL", url)
+ monkeypatch.setenv("RAG_EXTRACTION_TOKEN_FILE", "/run/extraction/custom-token")
+ monkeypatch.setenv("RAG_EXTRACTION_DEADLINE_SECONDS", "4.5")
+ monkeypatch.setenv("OPENAI_API_KEY", "unrelated-app-credential")
+ monkeypatch.setenv("JWT_SECRET", "unrelated-user-auth-secret")
+ settings = ExtractionSettings.from_env()
+ assert settings.url == url
+ assert settings.token_file == "/run/extraction/custom-token"
+ assert settings.deadline == 4.5
diff --git a/tests/extraction/test_processor.py b/tests/extraction/test_processor.py
new file mode 100644
index 00000000..ab76fbe5
--- /dev/null
+++ b/tests/extraction/test_processor.py
@@ -0,0 +1,280 @@
+"""Current Python client qualification against the default hardened deployment."""
+
+import json
+import subprocess
+import uuid
+from pathlib import Path
+
+import aiofiles
+import pytest
+
+from tests.extraction.test_deployment import (
+ assert_hardened_processor,
+ compose_command,
+ render,
+ stage_compose_files,
+)
+
+pytestmark = pytest.mark.integration
+CALLER_IMAGE = (
+ "python:3.13-slim-bookworm@"
+ "sha256:5024f48ba9441d4b13a95d3945abc6365538e3a31109833367a1923523c6efed"
+)
+GATEWAY = "/opt/extraction/gateway"
+
+
+def docker(*args, check=True):
+ result = subprocess.run(
+ ["docker", *args], capture_output=True, text=True, timeout=40
+ )
+ if check:
+ assert result.returncode == 0, result.stdout + result.stderr
+ return result
+
+
+def test_python_client_to_default_hardened_processor(tmp_path):
+ directory = stage_compose_files(tmp_path)
+ assert_hardened_processor(
+ render(directory, "docker-compose.yaml")["services"]["extraction"]
+ )
+ run_id = uuid.uuid4().hex
+ project = f"rag-extraction-{run_id}"
+ volume = f"{project}_extraction-credentials"
+ network = f"{project}_extraction"
+ caller = f"rag-extraction-caller-{run_id}"
+ compose = compose_command(directory, project, "docker-compose.yaml")
+ images = [f"{project}-extraction", f"{project}-extraction-init"]
+
+ def run_compose(*args, check=True, timeout=40):
+ result = subprocess.run(
+ [*compose, *args], capture_output=True, text=True, timeout=timeout
+ )
+ if check:
+ assert result.returncode == 0, result.stdout + result.stderr
+ return result
+
+ root = Path(__file__).resolve().parents[2]
+ client = root / "app/services/extraction_client.py"
+ fixtures = Path(__file__).parent / "fixtures"
+ library = Path(aiofiles.__file__).parent
+ try:
+ run_compose("build", "extraction-init", "extraction", timeout=240)
+ run_compose("up", "-d", "--wait", "--wait-timeout", "30", "extraction")
+ name = run_compose("ps", "-q", "extraction").stdout.strip()
+ image = docker("inspect", "--format", "{{.Config.Image}}", name).stdout.strip()
+ assert (
+ json.loads(docker("network", "inspect", network).stdout)[0]["Internal"]
+ is True
+ )
+ runtime = json.loads(docker("inspect", name).stdout)[0]
+ host = runtime["HostConfig"]
+ assert (
+ host["NetworkMode"] == network
+ and host["ReadonlyRootfs"]
+ and not host["Privileged"]
+ )
+ assert (
+ host["CapDrop"] == ["ALL"]
+ and "no-new-privileges:true" in host["SecurityOpt"]
+ )
+ assert "seccomp=unconfined" not in host["SecurityOpt"]
+ assert runtime["Config"]["User"] == "10001:10001"
+ assert runtime["Config"]["Entrypoint"] == [GATEWAY]
+ assert runtime["Config"]["Healthcheck"]["Test"] == [
+ "CMD",
+ GATEWAY,
+ "node",
+ "dist/health.js",
+ ]
+ init_name = run_compose("ps", "-a", "-q", "extraction-init").stdout.strip()
+ provisioner = json.loads(docker("inspect", init_name).stdout)[0]
+ assert provisioner["Config"]["User"] == "10001:10001"
+ assert provisioner["Config"]["Entrypoint"] == [GATEWAY]
+ assert provisioner["Config"]["Cmd"] == ["node", "dist/provision.js"]
+ assert provisioner["Config"]["Healthcheck"]["Test"] == ["NONE"]
+ # Compose adds service-specific image labels; the pinned payload must match.
+ provisioner_image = json.loads(
+ docker("image", "inspect", provisioner["Image"]).stdout
+ )[0]
+ runtime_image = json.loads(docker("image", "inspect", runtime["Image"]).stdout)[
+ 0
+ ]
+ assert provisioner_image["RootFS"] == runtime_image["RootFS"]
+ assert provisioner["State"]["ExitCode"] == 0
+ assert not host["PortBindings"]
+ assert set(runtime["NetworkSettings"]["Networks"]) == {network}
+ assert host["Memory"] == 768 * 1024 * 1024 and host["NanoCpus"] == 2 * 10**9
+ assert host["PidsLimit"] == 32 and set(host["Tmpfs"]) == {"/jobs"}
+ assert (
+ len(runtime["Mounts"]) == 1
+ and runtime["Mounts"][0]["Destination"] == "/run/extraction"
+ and runtime["Mounts"][0]["RW"] is False
+ )
+ script = """
+import asyncio
+import errno
+import os
+import stat
+import sys
+sys.path.insert(0, "/")
+from dataclasses import replace
+from pathlib import Path
+from client import ExtractionClient, ExtractionError, ExtractionSettings
+async def check():
+ settings = ExtractionSettings.from_env()
+ assert settings.url == "http://extraction:8000"
+ assert settings.token_file == "/run/extraction/token"
+ assert stat.S_IMODE(os.stat("/run/extraction").st_mode) == 0o755
+ assert stat.S_IMODE(os.stat(settings.token_file).st_mode) == 0o444
+ try:
+ open("/run/extraction/not-writable", "wb")
+ except OSError as error:
+ assert error.errno == errno.EROFS
+ else:
+ raise AssertionError("caller credential mount is writable")
+ for token_file in ("/jobs/missing-token", "/jobs/wrong-token"):
+ if token_file.endswith("wrong-token"):
+ token = Path(settings.token_file).read_text()
+ Path(token_file).write_text(("a" if token[0] != "a" else "b") + token[1:])
+ try:
+ await ExtractionClient(replace(settings, token_file=token_file)).extract("/fixtures/pages.pdf", "pdf")
+ except ExtractionError as error:
+ assert (error.code, error.status_code) == ("UNAVAILABLE", 503)
+ else:
+ raise AssertionError("missing/wrong service credentials accepted")
+ # Stop a real HTTP upload before the body completes, then reuse admission.
+ cancellation_input = Path("/jobs/cancel.pdf")
+ cancellation_input.write_bytes(b"%PDF-1.7\\n" + b"x" * 131072)
+ entered = asyncio.Event()
+ class SlowInput:
+ def __init__(self, file):
+ self.file = file
+ self.sent = False
+ async def read(self, size):
+ if not self.sent:
+ self.sent = True
+ return await self.file.read(1)
+ entered.set()
+ await asyncio.Future()
+ class CancellingClient(ExtractionClient):
+ async def _request(self, method, token, **kwargs):
+ if method == "POST":
+ kwargs["file"] = SlowInput(kwargs["file"])
+ return await super()._request(method, token, **kwargs)
+ cancelled_client = CancellingClient(settings)
+ task = asyncio.create_task(cancelled_client.extract(str(cancellation_input), "pdf"))
+ await asyncio.wait_for(entered.wait(), 2)
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError:
+ pass
+ else:
+ raise AssertionError("extraction did not propagate cancellation")
+ assert cancelled_client._admitted == 0
+ client = ExtractionClient(settings)
+ pdf = await client.extract("/fixtures/pages.pdf", "pdf")
+ assert pdf == [
+ {"kind": "page", "index": 1, "text": "First page text"},
+ {"kind": "page", "index": 2, "text": "Second page text"},
+ ], pdf
+ docx = await client.extract("/fixtures/body.docx", "docx")
+ assert len(docx) == 1 and docx[0]["kind"] == "document" and docx[0]["index"] == 1
+ text = docx[0]["text"]
+ assert text.index("First paragraph") < text.index("Left cell") < text.index("Right cell") < text.index("Last paragraph")
+asyncio.run(check())
+print("Python client to hardened PDF/DOCX processor passed")
+"""
+ # The trusted RAG caller uses Python's normal asyncio self-pipe, not the worker gateway.
+ caller_command = [
+ "run",
+ "--rm",
+ "--name",
+ caller,
+ "--init",
+ "--network",
+ network,
+ "--read-only",
+ "--user",
+ "65532:65532",
+ "--cap-drop",
+ "ALL",
+ "--security-opt",
+ "no-new-privileges:true",
+ "--mount",
+ f"type=volume,src={volume},dst=/run/extraction,readonly",
+ "--tmpfs",
+ "/jobs:rw,noexec,nosuid,nodev,size=1m,uid=65532,gid=65532,mode=0700",
+ "--mount",
+ f"type=bind,src={client},dst=/client.py,readonly",
+ "--mount",
+ f"type=bind,src={fixtures},dst=/fixtures,readonly",
+ "--mount",
+ f"type=bind,src={library},dst=/usr/local/lib/python3.13/site-packages/aiofiles,readonly",
+ CALLER_IMAGE,
+ "python",
+ "-B",
+ "-c",
+ script,
+ ]
+ result = docker(*caller_command)
+ assert "processor passed" in result.stdout
+ # Re-provisioning preserves credentials already loaded by the running service.
+ run_compose("run", "--rm", "--no-deps", "extraction-init")
+ docker("exec", name, GATEWAY, "node", "dist/health.js")
+ run_compose(
+ "up",
+ "-d",
+ "--wait",
+ "--wait-timeout",
+ "30",
+ "--scale",
+ "extraction=2",
+ "extraction",
+ )
+ replicas = run_compose("ps", "-q", "extraction").stdout.split()
+ assert len(replicas) == 2
+ for replica in replicas:
+ docker("exec", replica, GATEWAY, "node", "dist/health.js")
+ assert "processor passed" in docker(*caller_command).stdout
+ run_compose("restart", "extraction")
+ run_compose(
+ "up",
+ "-d",
+ "--wait",
+ "--wait-timeout",
+ "30",
+ "--scale",
+ "extraction=2",
+ "extraction",
+ )
+ assert "processor passed" in docker(*caller_command).stdout
+ # Exercise the shipped gateway rather than a test-only seccomp policy.
+ probe = """
+import errno
+import socket
+for operation in (lambda: socket.socket(socket.AF_UNIX), socket.socketpair):
+ try:
+ operation()
+ except OSError as error:
+ assert error.errno == errno.EPERM
+ else:
+ raise AssertionError("gateway permitted Unix sockets")
+"""
+ docker("exec", name, GATEWAY, "python", "-c", probe)
+ assert (
+ docker(
+ "exec",
+ name,
+ GATEWAY,
+ "node",
+ "-e",
+ 'if(require("node:fs").readdirSync("/jobs").length)process.exit(1)',
+ ).returncode
+ == 0
+ )
+ finally:
+ docker("rm", "-f", caller, check=False)
+ run_compose("down", "-v", "--remove-orphans", check=False)
+ for image in images:
+ docker("image", "rm", image, check=False)
diff --git a/tests/extraction/test_routes.py b/tests/extraction/test_routes.py
new file mode 100644
index 00000000..9fe87b3a
--- /dev/null
+++ b/tests/extraction/test_routes.py
@@ -0,0 +1,1320 @@
+import asyncio
+import copy
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from unittest.mock import Mock
+
+import httpx
+import jwt
+import pytest
+from fastapi import FastAPI
+from langchain_core.documents import Document
+
+from app.config import LogMiddleware
+from app.middleware import security_middleware
+from app.routes import document_routes as routes
+from app.services.extraction_client import ExtractionClient, ExtractionSettings
+from app.services.vector_store.async_pg_vector import AsyncPgVector
+
+ROUTES = ["/text", "/embed", "/embed-upload", "/local/embed"]
+MIMES = {
+ "txt": "text/plain",
+ "pdf": "application/pdf",
+ "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+}
+
+
+class RecordingStore(AsyncPgVector):
+ def __init__(self):
+ self._bind = None
+ self.calls = []
+ self.persisted = []
+ self.rollbacks = []
+ self.fail_on = None
+
+ async def aadd_documents(self, documents, ids=None, executor=None):
+ self.calls.append(copy.deepcopy(documents))
+ if len(self.calls) == self.fail_on:
+ raise RuntimeError("insertion failed")
+ self.persisted.extend(copy.deepcopy(documents))
+ return [f"row-{len(self.calls)}-{index}" for index in range(len(documents))]
+
+ async def delete_by_metadata(self, metadata, executor=None):
+ self.rollbacks.append(metadata)
+ self.persisted = [
+ doc
+ for doc in self.persisted
+ if not all(
+ doc.metadata.get(key) == value for key, value in metadata.items()
+ )
+ ]
+
+ async def get_filtered_ids(self, ids, owners=None, executor=None):
+ return [
+ file_id
+ for file_id in ids
+ if any(
+ doc.metadata["file_id"] == file_id and doc.metadata["user_id"] in owners
+ for doc in self.persisted
+ )
+ ]
+
+ async def get_documents_by_ids(self, ids, owners=None, executor=None):
+ return [
+ doc
+ for doc in reversed(self.persisted)
+ if doc.metadata["file_id"] in ids and doc.metadata["user_id"] in owners
+ ]
+
+
+@pytest.fixture
+async def api(processor, monkeypatch, tmp_path):
+ app = FastAPI()
+ app.add_middleware(LogMiddleware)
+ app.middleware("http")(security_middleware)
+ app.include_router(routes.router)
+ store = RecordingStore()
+ uploads = tmp_path / "uploads"
+ uploads.mkdir()
+ monkeypatch.setattr(routes, "RAG_UPLOAD_DIR", str(uploads))
+ monkeypatch.setattr(routes, "extraction_client", processor.client())
+ monkeypatch.setattr(routes, "PDF_EXTRACT_IMAGES", False)
+ monkeypatch.setattr(routes, "vector_store", store)
+ monkeypatch.setattr(routes, "PARALLEL_EXECUTION", 1)
+ monkeypatch.setenv("JWT_SECRET", "test-extraction-secret")
+ # The migrated path must never construct either legacy parser.
+ local_loader = routes.get_loader
+
+ def no_migrated_local_loader(filename, content_type, *args, **kwargs):
+ assert (
+ routes.extraction_format(filename, content_type) is None
+ ), "local migrated parser"
+ return local_loader(filename, content_type, *args, **kwargs)
+
+ monkeypatch.setattr(routes, "get_loader", no_migrated_local_loader)
+ monkeypatch.setattr(
+ "app.utils.document_loader.Docx2txtLoader",
+ Mock(side_effect=AssertionError("local DOCX parser")),
+ )
+ token = jwt.encode(
+ {
+ "id": "caller",
+ "exp": datetime.now(timezone.utc) + timedelta(minutes=5),
+ },
+ "test-extraction-secret",
+ algorithm="HS256",
+ )
+ with ThreadPoolExecutor(max_workers=2) as executor:
+ app.state.thread_pool = executor
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app=app),
+ base_url="http://rag",
+ headers={"Authorization": f"Bearer {token}"},
+ ) as client:
+ yield client, store, uploads
+
+
+async def post(
+ client,
+ uploads,
+ route,
+ format="pdf",
+ content=b"opaque bytes",
+ mime=None,
+ filename=None,
+ entity=None,
+):
+ filename = filename or f"report.{format}"
+ mime = mime or MIMES.get(format, "text/plain")
+ if route == "/local/embed":
+ local = uploads / filename
+ local.write_bytes(content)
+ return await client.post(
+ route,
+ params={"entity_id": entity} if entity else None,
+ json={
+ "filepath": filename,
+ "filename": filename,
+ "file_content_type": mime,
+ "file_id": "caller-file",
+ },
+ )
+ data = {"file_id": "caller-file"}
+ if entity:
+ data["entity_id"] = entity
+ return await client.post(
+ route,
+ data=data,
+ files={
+ "uploaded_file" if route == "/embed-upload" else "file": (
+ filename,
+ content,
+ mime,
+ ),
+ },
+ )
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_migrated_routes_parse_once_and_keep_caller_metadata(
+ api, processor, route, format
+):
+ client, store, uploads = api
+ response = await post(client, uploads, route, format)
+ assert response.status_code == 200, response.text
+ result = response.json()
+ assert result["file_id"] == "caller-file"
+ assert result["filename"] == f"report.{format}"
+ assert result["known_type"] is True
+ assert [r[0] for r in processor.requests] == ["GET", "POST"]
+ assert all(
+ request[2]["Authorization"] == f"Bearer {processor.token}"
+ and request[2]["Authorization"] != client.headers["Authorization"]
+ for request in processor.requests
+ )
+ if route == "/text":
+ assert result["text"] == (
+ "First page café\nSecond page"
+ if format == "pdf"
+ else "Body paragraph\nLeft\tRight"
+ )
+ assert not store.calls
+ else:
+ assert result["status"] is True
+ docs = store.persisted
+ assert docs
+ assert all(
+ doc.metadata["user_id"] == "caller"
+ and doc.metadata["file_id"] == "caller-file"
+ for doc in docs
+ )
+ assert all(str(uploads) in doc.metadata["source"] for doc in docs)
+ assert all("\x00" not in doc.page_content for doc in docs)
+ assert [doc.metadata["_rag_chunk_index"] for doc in docs] == list(
+ range(len(docs))
+ )
+ assert len({doc.metadata["_rag_ingestion_attempt_id"] for doc in docs}) == 1
+ if format == "pdf":
+ assert [doc.metadata["page"] for doc in docs] == [0, 1]
+ context = await client.get("/documents/caller-file/context")
+ assert context.status_code == 200
+ assert context.text.index("First page") < context.text.index("Second page")
+ assert "# PAGE 0" not in context.text
+ assert "# PAGE 1" in context.text # Existing zero-based context convention.
+ else:
+ assert all("page" not in doc.metadata for doc in docs)
+ remaining = list(uploads.rglob("*.*"))
+ assert len(remaining) == (1 if route == "/local/embed" else 0)
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize(
+ "failure,status",
+ [
+ ("unsupported", 422),
+ ("outage", 503),
+ ("incompatible", 502),
+ ("identity", 502),
+ ],
+)
+async def test_selected_format_failure_never_falls_back_or_persists(
+ api, processor, route, failure, status
+):
+ client, store, uploads = api
+ if failure == "unsupported":
+ processor.error = (422, "UNSUPPORTED_CONTENT")
+ elif failure == "outage":
+ routes.extraction_client = ExtractionClient(
+ ExtractionSettings(token_file="/run/extraction/missing-token")
+ )
+ elif failure == "incompatible":
+ processor.capabilities["version"] = 2
+ else:
+ processor.mutate = lambda value: {
+ **value,
+ "user_id": "forged",
+ "file_id": "forged",
+ "_rag_chunk_index": 99,
+ }
+ response = await post(client, uploads, route)
+ assert response.status_code == status, response.text
+ assert not store.calls and not store.persisted and not store.rollbacks
+ assert len(list(uploads.rglob("*.pdf"))) == (1 if route == "/local/embed" else 0)
+
+
+@pytest.mark.parametrize("route", ROUTES)
+async def test_unauthorized_request_cannot_reach_processor(api, processor, route):
+ client, store, uploads = api
+ client.headers.pop("Authorization")
+ response = await post(client, uploads, route)
+ assert response.status_code == 401
+ assert not processor.requests and not store.calls
+
+
+@pytest.mark.parametrize("route", ROUTES)
+async def test_cancellation_cleans_upload_and_never_stores(api, processor, route):
+ client, store, uploads = api
+ processor.hold = True
+ task = asyncio.create_task(post(client, uploads, route))
+ await asyncio.wait_for(processor.entered.wait(), 2)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ await asyncio.wait_for(processor.closed.wait(), 1)
+ assert not store.calls
+ assert len(list(uploads.rglob("*.pdf"))) == (1 if route == "/local/embed" else 0)
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+async def test_existing_attempt_rollback_and_no_success_on_storage_error(
+ api, processor, monkeypatch, route
+):
+ client, store, uploads = api
+ monkeypatch.setattr(routes, "EMBEDDING_BATCH_SIZE", 1)
+ store.fail_on = 2
+ old = Document(
+ page_content="earlier attempt",
+ metadata={
+ "file_id": "caller-file",
+ "user_id": "caller",
+ "_rag_ingestion_attempt_id": "old",
+ },
+ )
+ store.persisted.append(old)
+ response = await post(client, uploads, route)
+ assert response.status_code == 500, response.text
+ assert store.persisted == [old]
+ assert len(store.rollbacks) == 1
+ assert store.rollbacks[0]["file_id"] == "caller-file"
+ assert store.rollbacks[0]["_rag_ingestion_attempt_id"] != "old"
+ assert len([r for r in processor.requests if r[0] == "POST"]) == 1
+
+
+@pytest.mark.parametrize(
+ "filename,mime",
+ [
+ ("notes.md", "text/markdown"),
+ ("plain.txt", "text/plain"),
+ ("code.py", "text/x-python"),
+ ],
+)
+async def test_raw_text_parity_keeps_original_text_and_local_path(
+ api, processor, filename, mime
+):
+ client, store, uploads = api
+ content = "# Heading\n\n**bold** and café\n- item\n".encode()
+ response = await post(client, uploads, "/text", "txt", content, mime, filename)
+ assert response.status_code == 200
+ assert response.json()["text"] == content.decode().rstrip("\n")
+ assert response.json()["known_type"] is True
+ assert not processor.requests and not store.calls
+
+
+async def test_mime_selection_returns_canonical_pdf_and_cleans_text(api, processor):
+ client, store, uploads = api
+ response = await post(
+ client, uploads, "/text", mime="application/pdf", filename="opaque.bin"
+ )
+ assert response.status_code == 200
+ assert "\x00" not in response.json()["text"]
+ assert processor.requests[-1][2]["X-Extraction-Format"] == "pdf"
+
+
+async def test_selected_docx_cannot_be_hijacked_by_markdown_mime(api, processor):
+ client, store, uploads = api
+ response = await post(client, uploads, "/text", "docx", mime="text/markdown")
+ assert response.status_code == 200
+ assert processor.requests[-1][2]["X-Extraction-Format"] == "docx"
+
+
+async def test_requested_pdf_image_extraction_is_not_silently_downgraded(
+ api, processor, monkeypatch
+):
+ client, store, uploads = api
+ monkeypatch.setattr(routes, "PDF_EXTRACT_IMAGES", True)
+ response = await post(client, uploads, "/text")
+ assert response.status_code == 422
+ assert response.json()["detail"]["code"] == "UNSUPPORTED_CONTENT"
+ assert not processor.requests and not store.calls
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_no_extraction_configuration_still_requires_service(
+ api, processor, monkeypatch, route, format
+):
+ client, store, uploads = api
+ for key in (
+ "RAG_EXTRACTION_ENABLED",
+ "RAG_EXTRACTION_FORMATS",
+ "RAG_EXTRACTION_SOCKET",
+ "RAG_EXTRACTION_URL",
+ "RAG_EXTRACTION_TOKEN_FILE",
+ "RAG_EXTRACTION_DEADLINE_SECONDS",
+ ):
+ monkeypatch.delenv(key, raising=False)
+ monkeypatch.setenv("RAG_EXTRACTION_TOKEN_FILE", processor.token_file)
+ defaults = ExtractionSettings.from_env()
+ routes.extraction_client = ExtractionClient(defaults)
+ connect = asyncio.open_connection
+
+ async def disposable_endpoint(host, port, **kwargs):
+ assert (host, port) == ("extraction", 8000)
+ return await connect(
+ "127.0.0.1", int(processor.url.rsplit(":", 1)[1]), **kwargs
+ )
+
+ monkeypatch.setattr(asyncio, "open_connection", disposable_endpoint)
+ response = await post(client, uploads, route, format)
+ assert response.status_code == 200, response.text
+ assert [r[0] for r in processor.requests] == ["GET", "POST"]
+ assert response.json()["known_type"] is True
+
+
+class ASGIConnection:
+ """Deliver the complete upload, then a server disconnect without cancelling app."""
+
+ def __init__(self, app, request, body):
+ self.incoming = asyncio.Queue()
+ self.received = []
+ self.sent = []
+ chunks = [body[:17], body[17:53], body[53:]]
+ for index, chunk in enumerate(chunks):
+ self.incoming.put_nowait(
+ {
+ "type": "http.request",
+ "body": chunk,
+ "more_body": index < len(chunks) - 1,
+ }
+ )
+ self.scope = {
+ "type": "http",
+ "asgi": {"version": "3.0", "spec_version": "2.3"},
+ "http_version": "1.1",
+ "method": request.method,
+ "scheme": "http",
+ "path": request.url.path,
+ "raw_path": request.url.raw_path,
+ "query_string": request.url.query,
+ "root_path": "",
+ "headers": [(key.lower(), value) for key, value in request.headers.raw],
+ "client": ("127.0.0.1", 12345),
+ "server": ("rag", 80),
+ }
+ self.task = asyncio.create_task(app(self.scope, self.receive, self.send))
+
+ async def receive(self):
+ message = await self.incoming.get()
+ self.received.append(message["type"])
+ return message
+
+ async def send(self, message):
+ self.sent.append(message)
+
+ def disconnect(self):
+ self.incoming.put_nowait({"type": "http.disconnect"})
+
+ async def close(self):
+ if not self.task.done():
+ self.task.cancel()
+ await asyncio.gather(self.task, return_exceptions=True)
+
+
+def disconnect_request(client, uploads, route, format):
+ filename = f"report.{format}"
+ if route == "/local/embed":
+ (uploads / filename).write_bytes(b"opaque document bytes")
+ return client.build_request(
+ "POST",
+ route,
+ json={
+ "filepath": filename,
+ "filename": filename,
+ "file_content_type": MIMES[format],
+ "file_id": "caller-file",
+ },
+ )
+ return client.build_request(
+ "POST",
+ route,
+ data={"file_id": "caller-file"},
+ files={
+ "uploaded_file" if route == "/embed-upload" else "file": (
+ filename,
+ b"opaque document bytes",
+ MIMES[format],
+ ),
+ },
+ )
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_asgi_disconnect_stops_extraction_and_cleans_upload(
+ api, processor, route, format
+):
+ client, store, uploads = api
+ processor.hold = True
+ request = disconnect_request(client, uploads, route, format)
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ await asyncio.wait_for(processor.entered.wait(), 1)
+ assert connection.received == ["http.request"] * 3
+ connection.disconnect()
+ await asyncio.wait_for(processor.closed.wait(), 0.5)
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert "http.disconnect" in connection.received
+ assert connection.task.cancelling() == 0
+ assert not store.calls and not store.persisted and not store.rollbacks
+ assert not any(
+ message["type"] == "http.response.start" and 200 <= message["status"] < 300
+ for message in connection.sent
+ )
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ assert routes.extraction_client._admitted == 0
+ finally:
+ await connection.close()
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_asgi_disconnect_wins_over_processor_success(
+ api, processor, route, format
+):
+ client, store, uploads = api
+ request = disconnect_request(client, uploads, route, format)
+ body = await request.aread()
+ processor.before_response = lambda: connection.disconnect()
+ connection = ASGIConnection(client._transport.app, request, body)
+ try:
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert connection.received == ["http.request"] * 3 + ["http.disconnect"]
+ assert connection.task.cancelling() == 0
+ assert not store.calls and not store.persisted and not store.rollbacks
+ assert not any(
+ message["type"] == "http.response.start" and 200 <= message["status"] < 300
+ for message in connection.sent
+ )
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ assert routes.extraction_client._admitted == 0
+ finally:
+ await connection.close()
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+@pytest.mark.parametrize("phase", ["preparation", "insertion"])
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_disconnect_after_extraction_drains_and_preserves_old_rows(
+ api, processor, monkeypatch, route, phase, format
+):
+ import threading
+
+ client, store, uploads = api
+ monkeypatch.setattr(routes, "EMBEDDING_BATCH_SIZE", 1)
+ old = Document(
+ page_content="previous attempt",
+ metadata={
+ "file_id": "caller-file",
+ "user_id": "caller",
+ "_rag_ingestion_attempt_id": "old",
+ },
+ )
+ store.persisted.append(old)
+ events = []
+ preparation_entered = threading.Event()
+ preparation_release = threading.Event()
+ insert_entered = asyncio.Event()
+ insert_release = asyncio.Event()
+ prepare = routes._prepare_documents_sync
+ insert = store.aadd_documents
+ rollback = store.delete_by_metadata
+
+ def delayed_prepare(*args):
+ events.append("preparation:start")
+ preparation_entered.set()
+ assert preparation_release.wait(3), "test preparation gate timed out"
+ result = prepare(*args)
+ events.append("preparation:done")
+ return result
+
+ async def delayed_insert(*args, **kwargs):
+ events.append("insert:start")
+ insert_entered.set()
+ await insert_release.wait()
+ result = await insert(*args, **kwargs)
+ events.append("insert:commit")
+ return result
+
+ async def track_rollback(*args, **kwargs):
+ events.append("rollback")
+ return await rollback(*args, **kwargs)
+
+ if phase == "preparation":
+ monkeypatch.setattr(routes, "_prepare_documents_sync", delayed_prepare)
+ else:
+ monkeypatch.setattr(store, "aadd_documents", delayed_insert)
+ monkeypatch.setattr(store, "delete_by_metadata", track_rollback)
+ request = disconnect_request(client, uploads, route, format)
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ if phase == "preparation":
+ assert await asyncio.to_thread(preparation_entered.wait, 1)
+ else:
+ await asyncio.wait_for(insert_entered.wait(), 1)
+ assert [r[0] for r in processor.requests] == ["GET", "POST"]
+ events.append("disconnect:queued")
+ connection.disconnect()
+ await asyncio.sleep(0.05)
+ assert not connection.task.done()
+ assert not store.rollbacks
+ assert list(uploads.rglob(f"*.{format}"))
+ preparation_release.set()
+ insert_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ statuses = [
+ message["status"]
+ for message in connection.sent
+ if message["type"] == "http.response.start"
+ ]
+ assert store.persisted == [old]
+ assert statuses == [499]
+ assert connection.task.cancelling() == 0
+ if phase == "preparation":
+ assert not store.calls and not store.rollbacks
+ assert events.index("preparation:done") > events.index("disconnect:queued")
+ else:
+ assert events == [
+ "insert:start",
+ "disconnect:queued",
+ "insert:commit",
+ "rollback",
+ ]
+ assert store.rollbacks[0]["_rag_ingestion_attempt_id"] != "old"
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ finally:
+ preparation_release.set()
+ insert_release.set()
+ await connection.close()
+
+
+class SyncRecordingStore:
+ def __init__(self, entered, release, events, old):
+ self.entered = entered
+ self.release = release
+ self.events = events
+ self.calls = []
+ self.persisted = [old]
+ self.rollbacks = []
+
+ def add_documents(self, documents, ids=None):
+ self.calls.append(copy.deepcopy(documents))
+ self.events.append("insert:start")
+ self.entered.set()
+ assert self.release.wait(3), "test sync insertion gate timed out"
+ self.persisted.extend(copy.deepcopy(documents))
+ self.events.append("insert:commit")
+ return [
+ f"sync-row-{len(self.calls)}-{index}" for index in range(len(documents))
+ ]
+
+ def delete_scoped(self, *args, **kwargs):
+ self.rollbacks.append((args, kwargs))
+ raise AssertionError("disconnect must not delete previous attempts by file_id")
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+@pytest.mark.parametrize("mode", ["async_single", "sync_single", "sync_batch"])
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_disconnect_after_uncancellable_storage_starts_drains_without_undo(
+ api, processor, monkeypatch, route, mode, format
+):
+ import threading
+
+ client, store, uploads = api
+ monkeypatch.setattr(
+ routes, "EMBEDDING_BATCH_SIZE", 1 if mode == "sync_batch" else 0
+ )
+ old = Document(
+ page_content="previous attempt",
+ metadata={
+ "file_id": "caller-file",
+ "user_id": "caller",
+ "_rag_ingestion_attempt_id": "old",
+ },
+ )
+ events = []
+ sync_entered = threading.Event()
+ sync_release = threading.Event()
+ async_entered = asyncio.Event()
+ async_release = asyncio.Event()
+ if mode.startswith("sync"):
+ store = SyncRecordingStore(sync_entered, sync_release, events, old)
+ monkeypatch.setattr(routes, "vector_store", store)
+ else:
+ store.persisted.append(old)
+ insert = store.aadd_documents
+
+ async def delayed_insert(*args, **kwargs):
+ events.append("insert:start")
+ async_entered.set()
+ await async_release.wait()
+ result = await insert(*args, **kwargs)
+ events.append("insert:commit")
+ return result
+
+ monkeypatch.setattr(store, "aadd_documents", delayed_insert)
+ request = disconnect_request(client, uploads, route, format)
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ if mode.startswith("sync"):
+ assert await asyncio.to_thread(sync_entered.wait, 1)
+ else:
+ await asyncio.wait_for(async_entered.wait(), 1)
+ events.append("disconnect:queued")
+ connection.disconnect()
+ await asyncio.sleep(0.05)
+ assert "http.disconnect" in connection.received
+ assert not connection.task.done()
+ assert store.persisted == [old] and not store.rollbacks
+ assert list(uploads.rglob(f"*.{format}"))
+ sync_release.set()
+ async_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert old in store.persisted
+ assert len(store.persisted) == (3 if format == "pdf" else 2)
+ assert not store.rollbacks
+ assert events.index("insert:commit") > events.index("disconnect:queued")
+ expected_calls = 2 if mode == "sync_batch" and format == "pdf" else 1
+ assert len(store.calls) == expected_calls
+ assert [
+ m["status"] for m in connection.sent if m["type"] == "http.response.start"
+ ] == [499]
+ assert connection.task.cancelling() == 0
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ finally:
+ sync_release.set()
+ async_release.set()
+ await connection.close()
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+async def test_unmigrated_text_keeps_post_extraction_disconnect_behavior(
+ api, monkeypatch, route
+):
+ import threading
+
+ client, store, uploads = api
+ loader = Mock()
+ loader.lazy_load.return_value = [Document(page_content="legacy text")]
+ monkeypatch.setattr(routes, "get_loader", Mock(return_value=(loader, True, "txt")))
+ entered = threading.Event()
+ release = threading.Event()
+ prepare = routes._prepare_documents_sync
+
+ def delayed_prepare(*args):
+ entered.set()
+ assert release.wait(3), "test legacy preparation gate timed out"
+ return prepare(*args)
+
+ monkeypatch.setattr(routes, "_prepare_documents_sync", delayed_prepare)
+ request = disconnect_request(client, uploads, route, "txt")
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ assert await asyncio.to_thread(entered.wait, 1)
+ connection.disconnect()
+ await asyncio.sleep(0.05)
+ release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert "http.disconnect" not in connection.received
+ assert len(store.persisted) == 1
+ assert not store.rollbacks
+ assert [
+ m["status"] for m in connection.sent if m["type"] == "http.response.start"
+ ] == [200]
+ finally:
+ release.set()
+ await connection.close()
+
+
+@pytest.mark.parametrize("disconnect", [False, True])
+async def test_selected_sync_partial_failure_retains_old_and_committed_rows(
+ api, monkeypatch, disconnect
+):
+ import threading
+
+ client, _, uploads = api
+ monkeypatch.setattr(routes, "EMBEDDING_BATCH_SIZE", 1)
+ entered = threading.Event()
+ release = threading.Event()
+ events = []
+ old = Document(
+ page_content="old", metadata={"file_id": "caller-file", "user_id": "caller"}
+ )
+ store = SyncRecordingStore(entered, release, events, old)
+ insert = store.add_documents
+
+ def fail_second_batch(documents, ids=None):
+ if store.calls:
+ raise RuntimeError("second sync batch failed")
+ return insert(documents, ids)
+
+ monkeypatch.setattr(store, "add_documents", fail_second_batch)
+ monkeypatch.setattr(routes, "vector_store", store)
+ request = disconnect_request(client, uploads, "/embed", "pdf")
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ assert await asyncio.to_thread(entered.wait, 1)
+ if disconnect:
+ connection.disconnect()
+ await asyncio.sleep(0.05)
+ assert "http.disconnect" in connection.received
+ release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert len(store.persisted) == 2 and old in store.persisted
+ assert not store.rollbacks
+ assert [
+ m["status"] for m in connection.sent if m["type"] == "http.response.start"
+ ] == [499 if disconnect else 500]
+ finally:
+ release.set()
+ await connection.close()
+
+
+async def test_disconnect_during_upload_cleanup_waits_for_removal(api, monkeypatch):
+ client, _, uploads = api
+ entered = asyncio.Event()
+ release = asyncio.Event()
+ remove = routes.aiofiles.os.remove
+
+ async def delayed_remove(*args, **kwargs):
+ entered.set()
+ await release.wait()
+ await remove(*args, **kwargs)
+
+ monkeypatch.setattr(routes.aiofiles.os, "remove", delayed_remove)
+ request = disconnect_request(client, uploads, "/text", "pdf")
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ await asyncio.wait_for(entered.wait(), 1)
+ connection.disconnect()
+ await asyncio.sleep(0.05)
+ assert not connection.task.done()
+ release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert not list(uploads.rglob("*.pdf"))
+ assert [
+ m["status"] for m in connection.sent if m["type"] == "http.response.start"
+ ] == [499]
+ finally:
+ release.set()
+ await connection.close()
+
+
+@pytest.mark.parametrize("batch_size", [0, 1])
+async def test_disconnect_before_queued_sync_mutation_prevents_write(
+ api, monkeypatch, batch_size
+):
+ import threading
+
+ client, _, uploads = api
+ app = client._transport.app
+ original_executor = app.state.thread_pool
+ monkeypatch.setattr(routes, "EMBEDDING_BATCH_SIZE", batch_size)
+ blocker_entered = threading.Event()
+ blocker_release = threading.Event()
+ submitted = asyncio.Event()
+ old = Document(
+ page_content="old", metadata={"file_id": "caller-file", "user_id": "caller"}
+ )
+ store = SyncRecordingStore(threading.Event(), threading.Event(), [], old)
+ monkeypatch.setattr(routes, "vector_store", store)
+ prepare = routes._prepare_documents_sync
+ run_sync = routes._run_sync_request_store
+
+ def blocker():
+ blocker_entered.set()
+ assert blocker_release.wait(3), "test executor gate timed out"
+
+ def prepare_then_block_executor(*args):
+ docs = prepare(*args)
+ app.state.thread_pool.submit(blocker)
+ return docs
+
+ async def observe_submission(*args):
+ submitted.set()
+ return await run_sync(*args)
+
+ monkeypatch.setattr(routes, "_prepare_documents_sync", prepare_then_block_executor)
+ monkeypatch.setattr(routes, "_run_sync_request_store", observe_submission)
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ app.state.thread_pool = executor
+ request = disconnect_request(client, uploads, "/embed", "pdf")
+ connection = ASGIConnection(app, request, await request.aread())
+ try:
+ await asyncio.wait_for(submitted.wait(), 1)
+ assert await asyncio.to_thread(blocker_entered.wait, 1)
+ connection.disconnect()
+ await asyncio.sleep(0.05)
+ assert "http.disconnect" in connection.received
+ assert not connection.task.done()
+ assert not store.calls
+ blocker_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert store.persisted == [old]
+ assert not store.calls and not store.rollbacks
+ assert [
+ m["status"]
+ for m in connection.sent
+ if m["type"] == "http.response.start"
+ ] == [499]
+ finally:
+ blocker_release.set()
+ await connection.close()
+ app.state.thread_pool = original_executor
+
+
+@pytest.fixture
+def freeze_disconnect_wrapper(monkeypatch):
+ entered = asyncio.Event()
+ release = asyncio.Event()
+ original_wait = asyncio.wait
+ state = {}
+
+ async def frozen_wait(tasks, **kwargs):
+ result = await original_wait(tasks, **kwargs)
+ watcher = next(
+ (
+ task
+ for task in tasks
+ if "_with_extraction_request" in task.get_coro().__qualname__
+ and task.get_coro().__name__ == "wait_for_disconnect"
+ ),
+ None,
+ )
+ if watcher in result[0]:
+ state["operation"] = next(task for task in tasks if task is not watcher)
+ entered.set()
+ await release.wait()
+ return result
+
+ monkeypatch.setattr(asyncio, "wait", frozen_wait)
+ return entered, release, state
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+@pytest.mark.parametrize("repeat_cancel", [False, True])
+async def test_completed_batch_disconnect_cannot_cancel_queued_attempt_rollback(
+ api, monkeypatch, freeze_disconnect_wrapper, route, format, repeat_cancel
+):
+ import threading
+
+ client, store, uploads = api
+ app = client._transport.app
+ original_executor = app.state.thread_pool
+ monkeypatch.setattr(routes, "EMBEDDING_BATCH_SIZE", 500)
+ wrapper_entered, wrapper_release, state = freeze_disconnect_wrapper
+ insert_entered = asyncio.Event()
+ insert_release = asyncio.Event()
+ rollback_queued = asyncio.Event()
+ worker_entered = threading.Event()
+ worker_release = threading.Event()
+ events = []
+ queued_deletes = []
+ old = Document(
+ page_content="old attempt",
+ metadata={
+ "file_id": "caller-file",
+ "user_id": "caller",
+ "_rag_ingestion_attempt_id": "old",
+ },
+ )
+ store.persisted.append(old)
+ insert = store.aadd_documents
+
+ def occupy_worker():
+ worker_entered.set()
+ assert worker_release.wait(3), "test rollback executor gate timed out"
+
+ async def final_insert(*args, executor=None, **kwargs):
+ executor.submit(occupy_worker)
+ insert_entered.set()
+ await insert_release.wait()
+ result = await insert(*args, executor=executor, **kwargs)
+ events.append("insert:commit")
+ return result
+
+ def delete_current_attempt(metadata):
+ events.append("rollback:execute")
+ store.rollbacks.append(metadata)
+ store.persisted = [
+ doc
+ for doc in store.persisted
+ if not all(
+ doc.metadata.get(key) == value for key, value in metadata.items()
+ )
+ ]
+
+ async def queued_rollback(metadata, executor=None):
+ future = asyncio.get_running_loop().run_in_executor(
+ executor, delete_current_attempt, metadata
+ )
+ queued_deletes.append(future)
+ events.append("rollback:queued")
+ rollback_queued.set()
+ await future
+
+ monkeypatch.setattr(store, "aadd_documents", final_insert)
+ monkeypatch.setattr(store, "delete_by_metadata", queued_rollback)
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ app.state.thread_pool = executor
+ request = disconnect_request(client, uploads, route, format)
+ connection = ASGIConnection(app, request, await request.aread())
+ try:
+ await asyncio.wait_for(insert_entered.wait(), 1)
+ assert await asyncio.to_thread(worker_entered.wait, 1)
+ connection.disconnect()
+ await asyncio.wait_for(wrapper_entered.wait(), 1)
+ insert_release.set()
+ await asyncio.wait_for(rollback_queued.wait(), 1)
+ assert len(store.persisted) > 1 and not store.rollbacks
+ wrapper_release.set()
+ await asyncio.sleep(0.05)
+ assert state["operation"].cancelling() == 0
+ if repeat_cancel:
+ for _ in range(2):
+ state["operation"].cancel()
+ await asyncio.sleep(0)
+ await asyncio.sleep(0.05)
+ assert not connection.task.done(), events
+ assert not queued_deletes[0].cancelled(), events
+ assert list(uploads.rglob(f"*.{format}"))
+ worker_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert events == ["insert:commit", "rollback:queued", "rollback:execute"]
+ assert store.persisted == [old]
+ assert len(store.rollbacks) == 1
+ assert store.rollbacks[0]["_rag_ingestion_attempt_id"] != "old"
+ assert connection.task.cancelling() == 0
+ assert [
+ m["status"]
+ for m in connection.sent
+ if m["type"] == "http.response.start"
+ ] == [499]
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ finally:
+ wrapper_release.set()
+ insert_release.set()
+ worker_release.set()
+ await connection.close()
+ app.state.thread_pool = original_executor
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+@pytest.mark.parametrize("repeat_cancel", [False, True])
+async def test_disconnect_drain_owns_pending_insert_until_attempt_rollback(
+ api, monkeypatch, freeze_disconnect_wrapper, route, repeat_cancel
+):
+ import threading
+
+ client, store, uploads = api
+ app = client._transport.app
+ original_executor = app.state.thread_pool
+ monkeypatch.setattr(routes, "EMBEDDING_BATCH_SIZE", 1)
+ monkeypatch.setattr(routes, "PARALLEL_EXECUTION", 2)
+ wrapper_entered, wrapper_release, state = freeze_disconnect_wrapper
+ second_batch = asyncio.Event()
+ insert_draining = asyncio.Event()
+ worker_entered = threading.Event()
+ worker_release = threading.Event()
+ insert_tasks = []
+ insert_futures = []
+ events = []
+ old = Document(
+ page_content="old",
+ metadata={
+ "file_id": "caller-file",
+ "user_id": "caller",
+ "_rag_ingestion_attempt_id": "old",
+ },
+ )
+ store.persisted.append(old)
+ original_queue = asyncio.Queue
+ original_gather = asyncio.gather
+
+ class ControlledQueue(original_queue):
+ async def get(self):
+ item = await super().get()
+ if isinstance(item, tuple) and len(item) == 4 and item[2] >= 2:
+ await second_batch.wait()
+ return item
+
+ def observe_gather(*tasks, **kwargs):
+ if insert_tasks and insert_tasks[0] in tasks:
+ insert_draining.set()
+ return original_gather(*tasks, **kwargs)
+
+ def commit_after_release(documents):
+ events.append("insert:start")
+ worker_entered.set()
+ assert worker_release.wait(3), "test insert executor gate timed out"
+ store.calls.append(copy.deepcopy(documents))
+ store.persisted.extend(copy.deepcopy(documents))
+ events.append("insert:commit")
+ return ["new-row"]
+
+ async def blocked_insert(documents, ids=None, executor=None):
+ insert_tasks.append(asyncio.current_task())
+ future = asyncio.get_running_loop().run_in_executor(
+ executor, commit_after_release, documents
+ )
+ insert_futures.append(future)
+ return await future
+
+ def remove_attempt(metadata):
+ events.append("rollback")
+ store.rollbacks.append(metadata)
+ store.persisted = [
+ doc
+ for doc in store.persisted
+ if not all(
+ doc.metadata.get(key) == value for key, value in metadata.items()
+ )
+ ]
+
+ async def rollback(metadata, executor=None):
+ await asyncio.get_running_loop().run_in_executor(
+ executor, remove_attempt, metadata
+ )
+
+ monkeypatch.setattr(asyncio, "Queue", ControlledQueue)
+ monkeypatch.setattr(asyncio, "gather", observe_gather)
+ monkeypatch.setattr(store, "aadd_documents", blocked_insert)
+ monkeypatch.setattr(store, "delete_by_metadata", rollback)
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ app.state.thread_pool = executor
+ request = disconnect_request(client, uploads, route, "pdf")
+ connection = ASGIConnection(app, request, await request.aread())
+ try:
+ assert await asyncio.to_thread(worker_entered.wait, 1)
+ connection.disconnect()
+ await asyncio.wait_for(wrapper_entered.wait(), 1)
+ second_batch.set()
+ await asyncio.wait_for(insert_draining.wait(), 1)
+ assert state["operation"].cancelling() == 0
+ wrapper_release.set()
+ await asyncio.sleep(0.05)
+ assert state["operation"].cancelling() == 0
+ if repeat_cancel:
+ # The real disconnect started cleanup; stress subsequent owner cancellation.
+ state["operation"].cancel()
+ await asyncio.sleep(0)
+ state["operation"].cancel()
+ await asyncio.sleep(0.05)
+ assert not connection.task.done()
+ assert not insert_futures[0].cancelled()
+ assert not insert_tasks[0].cancelled()
+ assert store.persisted == [old] and not store.rollbacks
+ assert list(uploads.rglob("*.pdf"))
+ worker_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert events == ["insert:start", "insert:commit", "rollback"]
+ assert len(insert_tasks) == 1 and len(store.calls) == 1
+ assert store.persisted == [old]
+ assert store.rollbacks[0]["_rag_ingestion_attempt_id"] != "old"
+ assert [
+ m["status"]
+ for m in connection.sent
+ if m["type"] == "http.response.start"
+ ] == [499]
+ assert len(list(uploads.rglob("*.pdf"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ assert connection.task.cancelling() == 0
+ finally:
+ second_batch.set()
+ wrapper_release.set()
+ worker_release.set()
+ await connection.close()
+ app.state.thread_pool = original_executor
+
+
+@pytest.mark.parametrize("route", ["/embed", "/embed-upload", "/local/embed"])
+async def test_repeated_cancellation_during_preparation_still_drains_without_writes(
+ api, monkeypatch, freeze_disconnect_wrapper, route
+):
+ import threading
+
+ client, store, uploads = api
+ wrapper_entered, wrapper_release, state = freeze_disconnect_wrapper
+ worker_entered = threading.Event()
+ worker_release = threading.Event()
+ worker_done = threading.Event()
+ futures = []
+ prepare = routes._prepare_documents_sync
+ drain = routes._drain_request_worker
+
+ def blocked_prepare(*args):
+ worker_entered.set()
+ assert worker_release.wait(3), "test preparation executor gate timed out"
+ result = prepare(*args)
+ worker_done.set()
+ return result
+
+ async def observe_drain(future):
+ if routes._active_extraction_request.get() is not None and not futures:
+ futures.append(future)
+ return await drain(future)
+
+ monkeypatch.setattr(routes, "_prepare_documents_sync", blocked_prepare)
+ monkeypatch.setattr(routes, "_drain_request_worker", observe_drain)
+ request = disconnect_request(client, uploads, route, "pdf")
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ assert await asyncio.to_thread(worker_entered.wait, 1)
+ connection.disconnect()
+ await asyncio.wait_for(wrapper_entered.wait(), 1)
+ wrapper_release.set()
+ await asyncio.sleep(0.05)
+ for _ in range(2):
+ state["operation"].cancel()
+ await asyncio.sleep(0)
+ await asyncio.sleep(0.05)
+ assert not connection.task.done() and not worker_done.is_set()
+ assert not futures[0].cancelled()
+ assert not store.calls and not store.persisted and not store.rollbacks
+ assert list(uploads.rglob("*.pdf"))
+ worker_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert worker_done.is_set()
+ assert not store.calls and not store.persisted and not store.rollbacks
+ assert [
+ m["status"] for m in connection.sent if m["type"] == "http.response.start"
+ ] == [499]
+ assert len(list(uploads.rglob("*.pdf"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ assert connection.task.cancelling() == 0
+ finally:
+ wrapper_release.set()
+ worker_release.set()
+ await connection.close()
+
+
+async def test_repeated_cancellation_cannot_abandon_upload_unlink(
+ api, monkeypatch, freeze_disconnect_wrapper
+):
+ client, _, uploads = api
+ wrapper_entered, wrapper_release, state = freeze_disconnect_wrapper
+ unlink_entered = asyncio.Event()
+ unlink_release = asyncio.Event()
+ unlink_tasks = []
+ remove = routes.aiofiles.os.remove
+
+ async def delayed_remove(*args, **kwargs):
+ unlink_tasks.append(asyncio.current_task())
+ unlink_entered.set()
+ await unlink_release.wait()
+ await remove(*args, **kwargs)
+
+ monkeypatch.setattr(routes.aiofiles.os, "remove", delayed_remove)
+ request = disconnect_request(client, uploads, "/text", "pdf")
+ connection = ASGIConnection(client._transport.app, request, await request.aread())
+ try:
+ await asyncio.wait_for(unlink_entered.wait(), 1)
+ connection.disconnect()
+ await asyncio.wait_for(wrapper_entered.wait(), 1)
+ wrapper_release.set()
+ await asyncio.sleep(0.05)
+ for _ in range(2):
+ state["operation"].cancel()
+ await asyncio.sleep(0)
+ await asyncio.sleep(0.05)
+ assert not connection.task.done() and not unlink_tasks[0].cancelled()
+ assert list(uploads.rglob("*.pdf"))
+ unlink_release.set()
+ await asyncio.wait_for(asyncio.shield(connection.task), 1)
+ assert not list(uploads.rglob("*.pdf"))
+ assert [
+ m["status"] for m in connection.sent if m["type"] == "http.response.start"
+ ] == [499]
+ finally:
+ wrapper_release.set()
+ unlink_release.set()
+ await connection.close()
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_default_service_unavailable_fails_closed_without_affecting_text(
+ api, processor, monkeypatch, route, format
+):
+ client, store, uploads = api
+ for key in (
+ "RAG_EXTRACTION_ENABLED",
+ "RAG_EXTRACTION_FORMATS",
+ "RAG_EXTRACTION_SOCKET",
+ "RAG_EXTRACTION_URL",
+ "RAG_EXTRACTION_TOKEN_FILE",
+ "RAG_EXTRACTION_DEADLINE_SECONDS",
+ ):
+ monkeypatch.delenv(key, raising=False)
+ monkeypatch.setenv("RAG_EXTRACTION_TOKEN_FILE", processor.token_file)
+ routes.extraction_client = ExtractionClient(ExtractionSettings.from_env())
+
+ async def unavailable(host, port, **kwargs):
+ assert (host, port) == ("extraction", 8000)
+ raise ConnectionRefusedError()
+
+ monkeypatch.setattr(asyncio, "open_connection", unavailable)
+ response = await post(client, uploads, route, format)
+ assert response.status_code == 503
+ assert response.json()["detail"]["code"] == "UNAVAILABLE"
+ assert not store.calls and not store.persisted and not processor.requests
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ plaintext = await post(client, uploads, "/text", "txt", b"plain text", "text/plain")
+ assert plaintext.status_code == 200 and plaintext.json()["text"] == "plain text"
+
+
+@pytest.mark.parametrize("route", ROUTES)
+@pytest.mark.parametrize("failure", ["missing", "wrong", "extract-unauthorized"])
+@pytest.mark.parametrize("format", ["pdf", "docx"])
+async def test_service_auth_failure_is_unavailable_not_user_unauthorized(
+ api, processor, route, failure, format, caplog
+):
+ client, store, uploads = api
+ if failure == "missing":
+ Path(processor.token_file).unlink()
+ elif failure == "wrong":
+ Path(processor.token_file).write_text("cd" * 32)
+ else:
+ processor.error = (401, "UNAUTHORIZED")
+ response = await post(client, uploads, route, format)
+ assert response.status_code == 503
+ assert response.json()["detail"] == {"code": "UNAVAILABLE"}
+ assert not store.calls and not store.persisted and not store.rollbacks
+ assert (
+ len(processor.requests)
+ == {"missing": 0, "wrong": 1, "extract-unauthorized": 2}[failure]
+ )
+ assert "Bearer" not in caplog.text and processor.token not in caplog.text
+ assert "cd" * 32 not in caplog.text
+ assert len(list(uploads.rglob(f"*.{format}"))) == (
+ 1 if route == "/local/embed" else 0
+ )
+ plaintext = await post(client, uploads, "/text", "txt", b"plain text", "text/plain")
+ assert plaintext.status_code == 200
diff --git a/tests/extraction/test_upload_transport.py b/tests/extraction/test_upload_transport.py
new file mode 100644
index 00000000..c5e574cf
--- /dev/null
+++ b/tests/extraction/test_upload_transport.py
@@ -0,0 +1,135 @@
+import asyncio
+import socket
+import ssl
+
+import pytest
+
+from app.services.extraction_client import (
+ CHUNK_BYTES,
+ INPUT_BYTES,
+ ExtractionSettings,
+ ExtractionClient,
+ ExtractionError,
+)
+
+
+@pytest.mark.parametrize("https", [False, True])
+@pytest.mark.parametrize("cancel", [False, True])
+async def test_backpressured_upload_abort_stops_input(
+ processor, input_file, tls_certificate, monkeypatch, https, cancel
+):
+ with open(input_file, "wb") as file:
+ file.truncate(INPUT_BYTES)
+ entered, release, closed = asyncio.Event(), asyncio.Event(), asyncio.Event()
+ received = 0
+ tasks = set()
+
+ async def handle(reader, writer):
+ nonlocal received
+ tasks.add(asyncio.current_task())
+ try:
+ await reader.readuntil(b"\r\n\r\n")
+ entered.set()
+ await release.wait()
+ while chunk := await reader.read(CHUNK_BYTES):
+ received += len(chunk)
+ except OSError:
+ pass
+ finally:
+ writer.close()
+ try:
+ await writer.wait_closed()
+ except OSError:
+ pass
+ if entered.is_set():
+ closed.set()
+ tasks.discard(asyncio.current_task())
+
+ context = None
+ if https:
+ monkeypatch.setenv("SSL_CERT_FILE", str(tls_certificate[0]))
+ context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+ context.load_cert_chain(*tls_certificate)
+ server = await asyncio.start_server(handle, "127.0.0.1", 0, ssl=context)
+ connect = asyncio.open_connection
+
+ async def bounded_socket(*args, **kwargs):
+ reader, writer = await connect(*args, **kwargs)
+ writer.get_extra_info("socket").setsockopt(
+ socket.SOL_SOCKET, socket.SO_SNDBUF, 8192
+ )
+ return reader, writer
+
+ monkeypatch.setattr(asyncio, "open_connection", bounded_socket)
+ async with server:
+ client = ExtractionClient(
+ ExtractionSettings(
+ url=f"{'https://localhost' if https else 'http://127.0.0.1'}:{server.sockets[0].getsockname()[1]}",
+ token_file=processor.token_file,
+ deadline=10 if cancel else 0.2,
+ )
+ )
+ request = client._request
+
+ async def extraction_only(method, *args, **kwargs):
+ if method == "GET":
+ return processor.capabilities, ""
+ return await request(method, *args, **kwargs)
+
+ monkeypatch.setattr(client, "_request", extraction_only)
+ task = asyncio.create_task(client.extract(input_file, "pdf"))
+ try:
+ await asyncio.wait_for(entered.wait(), 1)
+ if cancel:
+ task.cancel()
+ with pytest.raises(
+ asyncio.CancelledError if cancel else ExtractionError
+ ) as caught:
+ await asyncio.wait_for(asyncio.shield(task), 1)
+ if not cancel:
+ assert caught.value.code == "DEADLINE"
+ release.set()
+ await asyncio.wait_for(closed.wait(), 1)
+ assert received < INPUT_BYTES and client._admitted == 0
+ finally:
+ release.set()
+ task.cancel()
+ await asyncio.gather(task, return_exceptions=True)
+ for active in list(tasks):
+ active.cancel()
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+
+async def test_repeated_cancellation_drains_aborted_writer(
+ processor, input_file, monkeypatch
+):
+ processor.hold = True
+ entered, release = asyncio.Event(), asyncio.Event()
+ wait_closed = asyncio.StreamWriter.wait_closed
+
+ async def delayed_close(writer):
+ await wait_closed(writer)
+ if processor.entered.is_set():
+ entered.set()
+ await release.wait()
+
+ monkeypatch.setattr(asyncio.StreamWriter, "wait_closed", delayed_close)
+ client = processor.client()
+ task = asyncio.create_task(client.extract(input_file, "pdf"))
+ try:
+ await asyncio.wait_for(processor.entered.wait(), 1)
+ task.cancel()
+ await asyncio.wait_for(entered.wait(), 1)
+ for _ in range(2):
+ task.cancel()
+ await asyncio.sleep(0)
+ assert not task.done() and client._admitted == 1
+ release.set()
+ with pytest.raises(asyncio.CancelledError):
+ await asyncio.wait_for(task, 1)
+ await asyncio.wait_for(processor.closed.wait(), 1)
+ assert client._admitted == 0
+ finally:
+ release.set()
+ task.cancel()
+ await asyncio.gather(task, return_exceptions=True)
diff --git a/tests/utils/test_document_loader.py b/tests/utils/test_document_loader.py
index 8baafa78..051fe60b 100644
--- a/tests/utils/test_document_loader.py
+++ b/tests/utils/test_document_loader.py
@@ -1,7 +1,3 @@
-import os
-from collections.abc import Iterator
-from unittest.mock import MagicMock, patch
-
import pytest
from app.utils.document_loader import get_loader, clean_text, process_documents
@@ -46,17 +42,6 @@ def test_process_documents():
assert "# PAGE 2" in processed
-def test_safe_pdf_loader_class():
- """Test that SafePyPDFLoader class can be instantiated"""
- from app.utils.document_loader import SafePyPDFLoader
-
- # Test instantiation
- loader = SafePyPDFLoader("dummy.pdf", extract_images=True)
- assert loader.filepath == "dummy.pdf"
- assert loader.extract_images == True
- assert loader._temp_filepath is None
-
-
def test_get_loader_text_lazy_load(tmp_path):
"""Test that lazy_load returns an iterator yielding documents."""
file_path = tmp_path / "test.txt"
@@ -69,111 +54,6 @@ def test_get_loader_text_lazy_load(tmp_path):
assert hasattr(data[0], "page_content")
-def test_get_loader_pdf(tmp_path):
- """Test get_loader returns SafePyPDFLoader for PDF files"""
- # Create a dummy PDF file path (doesn't need to be real for this test)
- file_path = tmp_path / "test.pdf"
- file_path.write_text("dummy content") # Not a real PDF, but that's OK for this test
-
- loader, known_type, file_ext = get_loader(
- "test.pdf", "application/pdf", str(file_path)
- )
-
- # Check that we get our SafePyPDFLoader
- from app.utils.document_loader import SafePyPDFLoader
-
- assert isinstance(loader, SafePyPDFLoader)
- assert known_type is True
- assert file_ext == "pdf"
-
-
-def test_safe_pdf_loader_lazy_load():
- """Test that SafePyPDFLoader.lazy_load() returns an Iterator."""
- from app.utils.document_loader import SafePyPDFLoader
-
- loader = SafePyPDFLoader("dummy.pdf", extract_images=False)
- assert hasattr(loader, "lazy_load")
- result = loader.lazy_load()
- assert isinstance(result, Iterator)
-
-
-def test_safe_pdf_loader_fallback_no_duplicate_pages():
- """Fallback after mid-stream KeyError must not duplicate already-yielded pages."""
- from app.utils.document_loader import SafePyPDFLoader
-
- fallback_docs = [Document(page_content=f"fallback page {i}") for i in range(5)]
-
- def primary_gen():
- yield Document(page_content="partial page 0")
- yield Document(page_content="partial page 1")
- raise KeyError("/Filter")
-
- def fallback_gen():
- yield from fallback_docs
-
- loader = SafePyPDFLoader("dummy.pdf", extract_images=True)
-
- with patch("app.utils.document_loader.PyPDFLoader") as MockPDF:
- primary_instance = MagicMock()
- primary_instance.lazy_load.side_effect = primary_gen
- fallback_instance = MagicMock()
- fallback_instance.lazy_load.side_effect = fallback_gen
- MockPDF.side_effect = [primary_instance, fallback_instance]
-
- result = list(loader.lazy_load())
-
- # Must be exactly the 5 fallback pages, NOT 2 partial + 5 fallback = 7
- assert len(result) == 5
- assert result[0].page_content == "fallback page 0"
- assert result[-1].page_content == "fallback page 4"
-
-
-def test_safe_pdf_loader_fallback_via_load():
- """load() delegates to lazy_load(), so fallback must also be correct via load()."""
- from app.utils.document_loader import SafePyPDFLoader
-
- fallback_docs = [Document(page_content=f"fb {i}") for i in range(3)]
-
- def primary_gen():
- yield Document(page_content="partial 0")
- raise KeyError("/Filter")
-
- def fallback_gen():
- yield from fallback_docs
-
- loader = SafePyPDFLoader("dummy.pdf", extract_images=True)
-
- with patch("app.utils.document_loader.PyPDFLoader") as MockPDF:
- primary_instance = MagicMock()
- primary_instance.lazy_load.side_effect = primary_gen
- fallback_instance = MagicMock()
- fallback_instance.lazy_load.side_effect = fallback_gen
- MockPDF.side_effect = [primary_instance, fallback_instance]
-
- result = loader.load()
-
- assert len(result) == 3
- assert result[0].page_content == "fb 0"
-
-
-def test_safe_pdf_loader_non_filter_error_propagates():
- """KeyError that isn't /Filter should propagate, not silently fallback."""
- from app.utils.document_loader import SafePyPDFLoader
-
- def bad_gen():
- raise KeyError("SomeOtherKey")
-
- loader = SafePyPDFLoader("dummy.pdf", extract_images=True)
-
- with patch("app.utils.document_loader.PyPDFLoader") as MockPDF:
- instance = MagicMock()
- instance.lazy_load.side_effect = bad_gen
- MockPDF.return_value = instance
-
- with pytest.raises(KeyError, match="SomeOtherKey"):
- list(loader.lazy_load())
-
-
MARKDOWN_SAMPLE = (
"# Heading\n\n"
"**bold** and *italic* text with a [link](https://example.com).\n\n"
@@ -236,26 +116,9 @@ def test_get_loader_markdown_raw_text_by_extension_only(tmp_path):
assert isinstance(loader, TextLoader)
-def test_get_loader_raw_text_leaves_pdf_alone(tmp_path):
- """raw_text must not disturb binary formats — PDF still uses the PDF loader."""
- from app.utils.document_loader import SafePyPDFLoader
-
- file_path = tmp_path / "doc.pdf"
- file_path.write_text("not a real pdf")
-
- loader, _, file_ext = get_loader(
- "doc.pdf", "application/pdf", str(file_path), raw_text=True
- )
-
- assert isinstance(loader, SafePyPDFLoader)
- assert file_ext == "pdf"
-
-
@pytest.mark.parametrize(
"filename, expected_loader_name",
[
- ("doc.pdf", "SafePyPDFLoader"),
- ("report.docx", "Docx2txtLoader"),
("book.epub", "UnstructuredEPubLoader"),
("data.xlsx", "UnstructuredExcelLoader"),
("slides.pptx", "UnstructuredPowerPointLoader"),
@@ -273,8 +136,30 @@ def test_get_loader_raw_text_respects_binary_extensions_over_markdown_mime(
file_path = tmp_path / filename
file_path.write_text("placeholder binary content")
- loader, _, _ = get_loader(
- filename, "text/markdown", str(file_path), raw_text=True
- )
+ loader, _, _ = get_loader(filename, "text/markdown", str(file_path), raw_text=True)
assert type(loader).__name__ == expected_loader_name
+
+
+@pytest.mark.parametrize(
+ "filename,mime",
+ [
+ ("doc.pdf", "application/pdf"),
+ ("doc.pdf", "text/markdown"),
+ ("doc.docx", "text/xml"),
+ ("doc.docx", "text/markdown"),
+ ("opaque.bin", "application/pdf"),
+ (
+ "opaque.bin",
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+ ),
+ ],
+)
+@pytest.mark.parametrize("raw_text", [False, True])
+def test_migrated_formats_cannot_construct_local_loader(
+ tmp_path, filename, mime, raw_text
+):
+ path = tmp_path / filename
+ path.write_bytes(b"not parsed locally")
+ with pytest.raises(ValueError, match="asynchronous extraction service"):
+ get_loader(filename, mime, str(path), raw_text=raw_text)
diff --git a/tests/utils/test_lazy_load.py b/tests/utils/test_lazy_load.py
index 9506f95f..144dd9e0 100644
--- a/tests/utils/test_lazy_load.py
+++ b/tests/utils/test_lazy_load.py
@@ -4,7 +4,7 @@
Verifies that every loader returned by get_loader() supports lazy_load(),
returns a generator/iterator, and yields valid Document objects with content.
Also includes memory benchmarks comparing lazy_load() vs load() to document
-which loaders benefit from streaming (CSV, PDF) vs which don't (Unstructured*).
+which local loaders benefit from streaming (CSV) vs which don't (Unstructured*).
"""
import gc
@@ -12,12 +12,11 @@
import shutil
import tracemalloc
import zipfile
-from collections.abc import Iterator
import pytest
from langchain_core.documents import Document
-from app.utils.document_loader import get_loader, SafePyPDFLoader
+from app.utils.document_loader import get_loader
# ---------------------------------------------------------------------------
# Environment checks — these deps aren't guaranteed in every CI runner
@@ -51,85 +50,6 @@
# ---------------------------------------------------------------------------
-def _make_pdf(path, *, num_pages=1):
- """Create a PDF with extractable text using pypdf.
-
- When *num_pages* > 1 each page gets unique text so tracemalloc can
- observe the difference between holding 1 page vs N pages in memory.
- """
- from pypdf import PdfWriter
- from pypdf.generic import (
- DecodedStreamObject,
- DictionaryObject,
- NameObject,
- )
-
- writer = PdfWriter()
- for i in range(num_pages):
- writer.add_blank_page(width=612, height=792)
- page = writer.pages[i]
-
- font_dict = DictionaryObject()
- font_dict[NameObject("/Type")] = NameObject("/Font")
- font_dict[NameObject("/Subtype")] = NameObject("/Type1")
- font_dict[NameObject("/BaseFont")] = NameObject("/Helvetica")
-
- resources = DictionaryObject()
- font_resources = DictionaryObject()
- font_resources[NameObject("/F1")] = font_dict
- resources[NameObject("/Font")] = font_resources
- page[NameObject("/Resources")] = resources
-
- # Pad each page with unique filler so memory differences are measurable
- filler = f"PAGE {i} " + ("X" * 2000)
- content = f"BT /F1 12 Tf 100 700 Td ({filler}) Tj ET".encode()
- stream = DecodedStreamObject()
- stream.set_data(content)
- page[NameObject("/Contents")] = stream
-
- with open(path, "wb") as f:
- writer.write(f)
-
-
-def _make_docx(path):
- """Create a minimal DOCX (Office Open XML zip) with text content."""
- content_types = (
- ''
- ''
- ' '
- ' '
- ' '
- ""
- )
- rels = (
- ''
- ''
- ' '
- ""
- )
- document = (
- ''
- ''
- " "
- " Hello from lazy_load DOCX test"
- " "
- ""
- )
- word_rels = (
- ''
- ''
- ""
- )
- with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as z:
- z.writestr("[Content_Types].xml", content_types)
- z.writestr("_rels/.rels", rels)
- z.writestr("word/document.xml", document)
- z.writestr("word/_rels/document.xml.rels", word_rels)
-
-
def _make_xlsx(path):
"""Create a minimal XLSX with openpyxl."""
from openpyxl import Workbook
@@ -283,14 +203,6 @@ def _make_large_csv(path, num_rows=500):
"Hello from lazy_load PY test",
id="py-source",
),
- pytest.param("test.pdf", "application/pdf", _make_pdf, "PAGE 0", id="pdf"),
- pytest.param(
- "test.docx",
- "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
- _make_docx,
- "lazy_load DOCX test",
- id="docx",
- ),
pytest.param(
"test.xlsx",
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
@@ -383,41 +295,6 @@ def test_lazy_load_matches_load(
), f"{type(loader).__name__} doc[{i}]: content mismatch between load() and lazy_load()"
-# ---------------------------------------------------------------------------
-# SafePyPDFLoader-specific tests
-# ---------------------------------------------------------------------------
-
-
-def test_safe_pdf_loader_lazy_load_is_iterator(tmp_path):
- """SafePyPDFLoader.lazy_load() should return an Iterator."""
- pdf_path = tmp_path / "gen_test.pdf"
- _make_pdf(str(pdf_path))
-
- loader = SafePyPDFLoader(str(pdf_path), extract_images=False)
- result = loader.lazy_load()
- assert isinstance(result, Iterator)
-
- # Consuming the iterator should yield documents
- docs = list(result)
- assert len(docs) > 0
-
-
-def test_safe_pdf_loader_load_delegates_to_lazy_load(tmp_path):
- """SafePyPDFLoader.load() should produce the same results as list(lazy_load())."""
- pdf_path = tmp_path / "delegate_test.pdf"
- _make_pdf(str(pdf_path))
-
- loader1 = SafePyPDFLoader(str(pdf_path), extract_images=False)
- loader2 = SafePyPDFLoader(str(pdf_path), extract_images=False)
-
- load_docs = loader1.load()
- lazy_docs = list(loader2.lazy_load())
-
- assert len(load_docs) == len(lazy_docs)
- for ld, lz in zip(load_docs, lazy_docs):
- assert ld.page_content == lz.page_content
-
-
# ---------------------------------------------------------------------------
# CSV with non-UTF-8 encoding
# ---------------------------------------------------------------------------
@@ -486,37 +363,6 @@ def _measure_lazy_load_streaming(loader_factory):
return texts, peak
-class TestMemoryBenchmarkPDF:
- """Memory comparison for PyPDFLoader -- the loader most likely to benefit
- from lazy_load() since it yields page-by-page."""
-
- NUM_PAGES = 50
-
- @pytest.fixture()
- def pdf_path(self, tmp_path):
- path = tmp_path / "bench.pdf"
- _make_pdf(str(path), num_pages=self.NUM_PAGES)
- return str(path)
-
- def test_pdf_streaming_lazy_load_peak_memory(self, pdf_path):
- """Streaming lazy_load() should use <= peak memory vs load()."""
-
- def factory():
- return SafePyPDFLoader(pdf_path, extract_images=False)
-
- load_docs, peak_load = _measure_load(factory)
- texts, peak_lazy_stream = _measure_lazy_load_streaming(factory)
-
- assert len(load_docs) == self.NUM_PAGES
- assert len(texts) == self.NUM_PAGES
-
- # Streaming should not use MORE memory than eager (allow 10% noise)
- assert peak_lazy_stream <= peak_load * 1.10, (
- f"streaming peak ({peak_lazy_stream:,} B) exceeded "
- f"load() peak ({peak_load:,} B) by >10%"
- )
-
-
class TestMemoryBenchmarkCSV:
"""Memory comparison for CSVLoader -- yields one Document per row."""
@@ -571,11 +417,6 @@ def factory():
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
id="pptx",
),
- pytest.param(
- "test.docx",
- "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
- id="docx",
- ),
]
# Map extensions to their file creators
@@ -595,7 +436,6 @@ def factory():
".epub": lambda p: _make_epub(p),
".xlsx": lambda p: _make_xlsx(p),
".pptx": lambda p: _make_pptx(p),
- ".docx": lambda p: _make_docx(p),
}