diff --git a/.github/workflows/docker-build-workflow.yaml b/.github/workflows/docker-build-workflow.yaml index 116bc990..5338994e 100644 --- a/.github/workflows/docker-build-workflow.yaml +++ b/.github/workflows/docker-build-workflow.yaml @@ -45,6 +45,12 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} + # Registry names must be lowercase, and the repository owner is not guaranteed to be. + - name: Resolve image name + env: + IMAGE_NAME: ${{ matrix.image_name }} + run: printf 'GHCR_IMAGE=ghcr.io/%s/%s\n' "${GITHUB_REPOSITORY_OWNER,,}" "$IMAGE_NAME" >> "$GITHUB_ENV" + # Build and push Docker images for each target - name: Build and push Docker images uses: docker/build-push-action@v5 @@ -53,7 +59,63 @@ jobs: file: ${{ matrix.file }} push: true tags: | - ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}:${{ github.sha }} - ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}:${{ inputs.tag }} + ${{ env.GHCR_IMAGE }}:${{ github.sha }} + ${{ env.GHCR_IMAGE }}:${{ inputs.tag }} platforms: linux/amd64,linux/arm64 - target: ${{ matrix.target }} \ No newline at end of file + target: ${{ matrix.target }} + + # After the repository moved to another owner, deployments still pull the old GHCR + # namespace. Copying each published tag there keeps them updating until the deprecation + # window closes. Copies are cross-repo blob mounts within GHCR, so no image data is + # rebuilt or re-uploaded. + # + # It runs in its own job so the credential that can write to the legacy namespace lives + # in an environment, out of reach of workflows on other branches, and so a mirror failure + # cannot fail a publish. + # + # Set the LEGACY_GHCR_OWNER variable and the LEGACY_GHCR_TOKEN secret (a token with + # write:packages for that owner) to turn this on. + mirror: + name: Mirror ${{ matrix.image_name }} to the legacy namespace + if: ${{ vars.LEGACY_GHCR_OWNER != '' }} + needs: build + runs-on: ubuntu-latest + timeout-minutes: 15 + continue-on-error: true + environment: legacy-mirror + strategy: + fail-fast: false + matrix: + image_name: [librechat-rag-api-dev, librechat-rag-api-dev-lite] + + steps: + - name: Copy the published tags + env: + LEGACY_OWNER: ${{ vars.LEGACY_GHCR_OWNER }} + LEGACY_TOKEN: ${{ secrets.LEGACY_GHCR_TOKEN }} + IMAGE_NAME: ${{ matrix.image_name }} + TAGS: "${{ github.sha }} ${{ inputs.tag }}" + run: | + set -euo pipefail + + legacy_owner="${LEGACY_OWNER,,}" + if [ "$legacy_owner" = "${GITHUB_REPOSITORY_OWNER,,}" ]; then + echo "Legacy owner matches the current owner; nothing to mirror." + exit 0 + fi + if [ -z "$LEGACY_TOKEN" ]; then + echo "::warning::LEGACY_GHCR_OWNER is set but LEGACY_GHCR_TOKEN is not; skipping the mirror." + exit 0 + fi + + # Registry names must be lowercase, and the repository owner is not guaranteed to be. + source_image="ghcr.io/${GITHUB_REPOSITORY_OWNER,,}/${IMAGE_NAME}" + legacy_image="ghcr.io/${legacy_owner}/${IMAGE_NAME}" + + # This login reads both namespaces: it replaces the default ghcr.io credential. + echo "$LEGACY_TOKEN" | docker login ghcr.io -u "$legacy_owner" --password-stdin + + for tag in $TAGS; do + echo "Mirroring ${source_image}:${tag} -> ${legacy_image}:${tag}" + docker buildx imagetools create -t "${legacy_image}:${tag}" "${source_image}:${tag}" + done diff --git a/.github/workflows/images.yaml b/.github/workflows/images.yaml index 12e8099c..8b04a503 100644 --- a/.github/workflows/images.yaml +++ b/.github/workflows/images.yaml @@ -57,6 +57,12 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} + # Registry names must be lowercase, and the repository owner is not guaranteed to be. + - name: Resolve image name + env: + IMAGE_NAME: ${{ matrix.image_name }} + run: printf 'GHCR_IMAGE=ghcr.io/%s/%s\n' "${GITHUB_REPOSITORY_OWNER,,}" "$IMAGE_NAME" >> "$GITHUB_ENV" + # Build and push Docker images for each target - name: Build and push Docker images uses: docker/build-push-action@v5 @@ -65,9 +71,65 @@ jobs: file: ${{ matrix.file }} push: true tags: | - ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}:${{ github.sha }} - ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}:latest + ${{ env.GHCR_IMAGE }}:${{ github.sha }} + ${{ env.GHCR_IMAGE }}:latest platforms: linux/amd64,linux/arm64 target: ${{ matrix.target }} cache-from: type=gha cache-to: type=gha,mode=max + + # After the repository moved to another owner, deployments still pull the old GHCR + # namespace. Copying each published tag there keeps them updating until the deprecation + # window closes. Copies are cross-repo blob mounts within GHCR, so no image data is + # rebuilt or re-uploaded. + # + # It runs in its own job so the credential that can write to the legacy namespace lives + # in an environment, out of reach of workflows on other branches, and so a mirror failure + # cannot fail a publish. + # + # Set the LEGACY_GHCR_OWNER variable and the LEGACY_GHCR_TOKEN secret (a token with + # write:packages for that owner) to turn this on. + mirror: + name: Mirror ${{ matrix.image_name }} to the legacy namespace + if: ${{ vars.LEGACY_GHCR_OWNER != '' }} + needs: build + runs-on: ubuntu-latest + timeout-minutes: 15 + continue-on-error: true + environment: legacy-mirror + strategy: + fail-fast: false + matrix: + image_name: [librechat-rag-api-dev, librechat-rag-api-dev-lite] + + steps: + - name: Copy the published tags + env: + LEGACY_OWNER: ${{ vars.LEGACY_GHCR_OWNER }} + LEGACY_TOKEN: ${{ secrets.LEGACY_GHCR_TOKEN }} + IMAGE_NAME: ${{ matrix.image_name }} + TAGS: "${{ github.sha }} latest" + run: | + set -euo pipefail + + legacy_owner="${LEGACY_OWNER,,}" + if [ "$legacy_owner" = "${GITHUB_REPOSITORY_OWNER,,}" ]; then + echo "Legacy owner matches the current owner; nothing to mirror." + exit 0 + fi + if [ -z "$LEGACY_TOKEN" ]; then + echo "::warning::LEGACY_GHCR_OWNER is set but LEGACY_GHCR_TOKEN is not; skipping the mirror." + exit 0 + fi + + # Registry names must be lowercase, and the repository owner is not guaranteed to be. + source_image="ghcr.io/${GITHUB_REPOSITORY_OWNER,,}/${IMAGE_NAME}" + legacy_image="ghcr.io/${legacy_owner}/${IMAGE_NAME}" + + # This login reads both namespaces: it replaces the default ghcr.io credential. + echo "$LEGACY_TOKEN" | docker login ghcr.io -u "$legacy_owner" --password-stdin + + for tag in $TAGS; do + echo "Mirroring ${source_image}:${tag} -> ${legacy_image}:${tag}" + docker buildx imagetools create -t "${legacy_image}:${tag}" "${source_image}:${tag}" + done diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 3aaeaaff..207f37ee 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -13,5 +13,8 @@ permissions: jobs: call-docker-build: uses: ./.github/workflows/docker-build-workflow.yaml + # The mirror step reads LEGACY_GHCR_TOKEN, and a called workflow gets no + # secrets unless they are passed. + secrets: inherit with: tag: ${{ github.ref_name }} diff --git a/Dockerfile b/Dockerfile index 01ba8d50..8d7a836e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,8 +2,10 @@ FROM python:3.12 AS main WORKDIR /app -# Install pandoc and netcat +# Install pandoc and netcat. `upgrade` pulls Debian security fixes (e.g. +# perl-base) that the base image tag may not carry yet. RUN apt-get update \ + && apt-get upgrade -y --no-install-recommends \ && apt-get install -y --no-install-recommends \ pandoc \ netcat-openbsd \ @@ -13,14 +15,6 @@ RUN apt-get update \ COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -ENV NLTK_DATA=/app/nltk_data - -# Download standard NLTK data, to prevent unstructured from downloading packages at runtime. -# Uses the Python API with explicit package IDs (rather than `python -m nltk.downloader`), -# which never falls back to the interactive prompt that EOFs during a non-interactive build. -# A failed download now fails the build (upstream #315) instead of leaving the image to -# fetch at runtime. -RUN python -c "import nltk, sys; packages = ('punkt_tab', 'averaged_perceptron_tagger', 'averaged_perceptron_tagger_eng'); sys.exit(0 if all(nltk.download(package, download_dir='/app/nltk_data') for package in packages) else 1)" # Disable Unstructured analytics ENV SCARF_NO_ANALYTICS=true diff --git a/Dockerfile.lite b/Dockerfile.lite index 163c2d2b..ddc6d572 100644 --- a/Dockerfile.lite +++ b/Dockerfile.lite @@ -2,8 +2,11 @@ FROM python:3.12-slim AS lite WORKDIR /app -# Install pandoc (DOCX /text extraction) and netcat (DB wait-for in entrypoint) +# Install pandoc (DOCX /text extraction) and netcat (DB wait-for in entrypoint). +# `upgrade` pulls Debian security fixes (e.g. perl-base) that the base image +# tag may not carry yet. RUN apt-get update \ + && apt-get upgrade -y --no-install-recommends \ && apt-get install -y --no-install-recommends \ pandoc \ netcat-openbsd \ @@ -17,14 +20,6 @@ RUN pip install --no-cache-dir --no-compile -r requirements.lite.txt \ -type d -name tests -prune -exec rm -rf {} + \ && find /usr/local/lib/python3.12/site-packages \ -type d -name __pycache__ -prune -exec rm -rf {} + -ENV NLTK_DATA=/app/nltk_data - -# Download standard NLTK data, to prevent unstructured from downloading packages at runtime. -# Uses the Python API with explicit package IDs (rather than `python -m nltk.downloader`), -# which never falls back to the interactive prompt that EOFs during a non-interactive build. -# A failed download now fails the build (upstream #315) instead of leaving the image to -# fetch at runtime. -RUN python -c "import nltk, sys; packages = ('punkt_tab', 'averaged_perceptron_tagger', 'averaged_perceptron_tagger_eng'); sys.exit(0 if all(nltk.download(package, download_dir='/app/nltk_data') for package in packages) else 1)" # Disable Unstructured analytics ENV SCARF_NO_ANALYTICS=true diff --git a/README.md b/README.md index 2be1b14a..0d81a103 100644 --- a/README.md +++ b/README.md @@ -209,6 +209,11 @@ The following environment variables are required to run the application: > **Upgrading:** these defaults change what `/text` returns for `.docx` files — tables come back as pipe tables (or an HTML `` for cells holding block content) instead of padded grid tables, and the cleanup pass runs. Anything downstream that depends on the previous Markdown shape byte for byte — a prompt template keyed on grid tables, or a stored extraction being compared against a fresh one — sees different text. Set `DOCX_TEXT_TABLE_STYLE=grid`, `DOCX_TEXT_CLEANUP=False` and `DOCX_TEXT_STRIP_HEADING_ANCHORS=False` to keep the old output. - `EMAIL_INCLUDE_HEADERS`: (Optional) Boolean. When "True" (default), the From/To/Cc/Subject/Date headers are prepended to extracted email text for `.eml` and `.msg` files. Set "False" to extract the body only. +- `MISTRAL_API_KEY`: (Optional) API key for Mistral OCR, used for PDFs and standalone images. OCR is unavailable without it. +- `MISTRAL_OCR_MODEL`: (Optional) Mistral OCR model. Default value is `mistral-ocr-latest`. +- `MISTRAL_SERVER_URL`: (Optional) Base URL of the Mistral API. Leave unset for the SDK default (`https://api.mistral.ai`); set `https://api.eu.mistral.ai` to keep OCR traffic on Mistral's [EU regional endpoint](https://docs.mistral.ai/inference/regional-inference). Regional endpoints don't offer the Files API, which is fine here: documents and images are sent inline as base64. Confirm with Mistral that your OCR model is served on the regional endpoint before relying on it. + - Note: LibreChat's own `mistral_ocr` strategy uploads through the Files API, so on an EU-only deployment route PDFs and images to this service's `/text` endpoint instead. + - For EU embeddings as well, point `RAG_OPENAI_BASEURL` at an EU OpenAI-compatible endpoint (e.g. `https://eu.openrouter.ai/api/v1`) with a non-OpenAI model such as `mistralai/mistral-embed-2312` (1024 dimensions), and set `RAG_CHECK_EMBEDDING_CTX_LENGTH=False`: with the check on, inputs are sent as OpenAI tiktoken token IDs, which non-OpenAI models don't accept. Switching embedding models requires re-indexing existing files (see `EMBEDDINGS_DIMENSIONS`). - Standalone image uploads (`.png`, `.jpg/.jpeg`, `.gif`, `.bmp`, `.tif/.tiff`, `.webp`) — including multi-page TIFF — are run through **Mistral OCR**, yielding one result per page. Requires `MISTRAL_API_KEY` (same key used for PDF OCR). This is separate from PDF handling; images embedded inside PDFs are not extracted. - `IMAGE_OCR_MAX_PAGES`: (Optional) Max number of frames/pages of a single image upload (e.g. animated GIF or large multi-page TIFF) sent to OCR; frames beyond the cap are skipped with a warning. Default value is `100`. - `DEBUG_RAG_API`: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes diff --git a/app/config.py b/app/config.py index d37e12b4..2a34257c 100644 --- a/app/config.py +++ b/app/config.py @@ -221,6 +221,9 @@ async def dispatch(self, request, call_next): # Mistral OCR settings MISTRAL_API_KEY = get_env_variable("MISTRAL_API_KEY", "") MISTRAL_OCR_MODEL = get_env_variable("MISTRAL_OCR_MODEL", "mistral-ocr-latest") +# Override the Mistral API base URL, e.g. https://api.eu.mistral.ai for EU-only +# deployments. Empty keeps the SDK default (https://api.mistral.ai). +MISTRAL_SERVER_URL = get_env_variable("MISTRAL_SERVER_URL", "").rstrip("/") # DOCX handling for the /text endpoint: when enabled, .docx files are converted # via pandoc so tracked changes and comments are preserved (the embedding path diff --git a/app/routes/document_routes.py b/app/routes/document_routes.py index 2360174a..c3cdb2ee 100644 --- a/app/routes/document_routes.py +++ b/app/routes/document_routes.py @@ -91,6 +91,7 @@ def _apply_distance_threshold(documents): QueryMultipleBody, ) from app.services.vector_store.async_pg_vector import AsyncPgVector +from app.utils.async_utils import run_in_executor from app.utils.document_loader import ( get_loader, clean_text, @@ -534,7 +535,8 @@ async def delete_documents( raise HTTPException(status_code=500, detail=str(e)) -# Cache the embedding function with LRU cache +# Cache query vectors only, never scoped search results. Both cache hits and +# misses run in the request worker pool so a miss cannot block the event loop. @lru_cache(maxsize=128) def get_cached_query_embedding(query: str): return vector_store.embedding_function.embed_query(query) @@ -548,7 +550,9 @@ async def query_embeddings_by_file_id( scope = resolve_scope(request, body.entity_id) try: - embedding = get_cached_query_embedding(body.query) + embedding = await run_in_executor( + request.app.state.thread_pool, get_cached_query_embedding, body.query + ) query_filter = scope.predicate(file_clause(body.file_id)) if isinstance(vector_store, AsyncPgVector): @@ -559,8 +563,12 @@ async def query_embeddings_by_file_id( executor=request.app.state.thread_pool, ) else: - documents = vector_store.similarity_search_with_score_by_vector( - embedding, k=body.k, filter=query_filter + documents = await run_in_executor( + request.app.state.thread_pool, + vector_store.similarity_search_with_score_by_vector, + embedding, + k=body.k, + filter=query_filter, ) return _apply_distance_threshold(documents) @@ -1460,7 +1468,9 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody scope = resolve_scope(request, body.entity_id) try: # Get the embedding of the query text - embedding = get_cached_query_embedding(body.query) + embedding = await run_in_executor( + request.app.state.thread_pool, get_cached_query_embedding, body.query + ) query_filter = scope.predicate(files_clause(body.file_ids)) # Perform similarity search with the query embedding and filter by the file_ids in metadata @@ -1472,8 +1482,12 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody executor=request.app.state.thread_pool, ) else: - documents = vector_store.similarity_search_with_score_by_vector( - embedding, k=body.k, filter=query_filter + documents = await run_in_executor( + request.app.state.thread_pool, + vector_store.similarity_search_with_score_by_vector, + embedding, + k=body.k, + filter=query_filter, ) documents = _apply_distance_threshold(documents) diff --git a/app/services/vector_store/async_pg_vector.py b/app/services/vector_store/async_pg_vector.py index 22e093a8..7b93d010 100644 --- a/app/services/vector_store/async_pg_vector.py +++ b/app/services/vector_store/async_pg_vector.py @@ -2,6 +2,7 @@ import asyncio from concurrent.futures import Executor from langchain_core.documents import Document +from app.utils.async_utils import run_in_executor from .extended_pg_vector import ExtendedPgVector T = TypeVar("T") @@ -28,21 +29,8 @@ async def _run_in_executor( *args: Any, **kwargs: Any, ) -> T: - """Run a sync callable in a thread pool executor. - - Wraps the call to convert StopIteration into RuntimeError. - StopIteration cannot be set on an asyncio.Future — it raises - TypeError and leaves the Future pending forever. - """ - - def wrapper() -> T: - try: - return func(*args, **kwargs) - except StopIteration as exc: - raise RuntimeError from exc - - loop = asyncio.get_running_loop() - return await loop.run_in_executor(executor, wrapper) + """Run a sync callable with the shared exception-safe executor bridge.""" + return await run_in_executor(executor, func, *args, **kwargs) async def get_all_ids(self, owners: Sequence[str], executor=None) -> list[str]: executor = executor or self._get_thread_pool() diff --git a/app/utils/async_utils.py b/app/utils/async_utils.py new file mode 100644 index 00000000..2adc69b8 --- /dev/null +++ b/app/utils/async_utils.py @@ -0,0 +1,29 @@ +"""Bridges for blocking operations used by async request paths.""" + +import asyncio +from concurrent.futures import Executor +from typing import Any, Callable, TypeVar + +T = TypeVar("T") + + +async def run_in_executor( + executor: Executor | None, + func: Callable[..., T], + *args: Any, + **kwargs: Any, +) -> T: + """Run a sync callable in the supplied pool without blocking the event loop. + + StopIteration cannot be set on an asyncio.Future: convert it into an error + that the awaiting request can handle instead of leaving it pending forever. + Cancelling the await does not stop a callable already running in a worker. + """ + + def wrapper() -> T: + try: + return func(*args, **kwargs) + except StopIteration as exc: + raise RuntimeError from exc + + return await asyncio.get_running_loop().run_in_executor(executor, wrapper) diff --git a/app/utils/document_loader.py b/app/utils/document_loader.py index 6c51837e..44ab09db 100644 --- a/app/utils/document_loader.py +++ b/app/utils/document_loader.py @@ -16,6 +16,7 @@ logger, MISTRAL_API_KEY, MISTRAL_OCR_MODEL, + MISTRAL_SERVER_URL, DOCX_TEXT_USE_PANDOC, DOCX_TEXT_TRACK_CHANGES, DOCX_TEXT_INCLUDE_HEADERS_FOOTERS, @@ -319,6 +320,11 @@ def _mistral_ocr_client(): "MISTRAL_API_KEY is not set. Please configure it in environment variables." ) + # Only pass server_url when configured so the SDK keeps its own default. + # Regional endpoints (e.g. api.eu.mistral.ai) lack the Files API; OCR here + # sends documents inline as base64, so it works against them unchanged. + if MISTRAL_SERVER_URL: + return Mistral(api_key=MISTRAL_API_KEY, server_url=MISTRAL_SERVER_URL) return Mistral(api_key=MISTRAL_API_KEY) diff --git a/requirements.lite.txt b/requirements.lite.txt index 69096ef3..6a0c7f02 100644 --- a/requirements.lite.txt +++ b/requirements.lite.txt @@ -9,8 +9,10 @@ fastapi==0.115.12 psycopg2-binary==2.9.9 pgvector==0.2.5 uvicorn==0.28.0 -pypdf==6.14.2 -unstructured==0.18.32 +pypdf==6.19.0 +unstructured==0.27.8 +# unstructured >=0.24 lazily downloads this spaCy model from GitHub on first use; install it at build time instead (same pin/hash unstructured 0.27.8 expects). +en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85 markdown==3.8.2 networkx==3.2.1 pandas==2.2.1 @@ -20,13 +22,13 @@ python-docx==1.1.2 # .docx header/footer extraction (pandoc drops them) pypandoc==1.15 python-oxmsg==0.0.2 # .msg (Outlook) extraction; also a transitive of unstructured Pillow==12.3.0 # normalize standalone images (incl. multi-page TIFF) for OCR -PyJWT==2.13.0 +PyJWT==2.15.1 asyncpg==0.29.0 python-multipart==0.0.32 aiofiles==24.1.0 -pymongo==4.17.0 +pymongo==4.18.2 langchain-mongodb==0.11.0 -cryptography==48.0.1 +cryptography==50.0.1 python-magic==0.4.27 python-pptx==1.0.2 xlrd==2.0.2 diff --git a/requirements.txt b/requirements.txt index 6dbc8abc..e77f4368 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,8 +11,10 @@ fastapi==0.115.12 psycopg2-binary==2.9.9 pgvector==0.2.5 uvicorn==0.28.0 -pypdf==6.14.2 -unstructured==0.18.32 +pypdf==6.19.0 +unstructured==0.27.8 +# unstructured >=0.24 lazily downloads this spaCy model from GitHub on first use; install it at build time instead (same pin/hash unstructured 0.27.8 expects). +en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85 markdown==3.8.2 networkx==3.2.1 pandas==2.2.1 @@ -22,19 +24,19 @@ python-docx==1.1.2 # .docx header/footer extraction (pandoc drops them) pypandoc==1.15 python-oxmsg==0.0.2 # .msg (Outlook) extraction; also a transitive of unstructured Pillow==12.3.0 # normalize standalone images (incl. multi-page TIFF) for OCR -PyJWT==2.13.0 +PyJWT==2.15.1 asyncpg==0.29.0 python-multipart==0.0.32 sentence_transformers==3.1.1 aiofiles==24.1.0 rapidocr-onnxruntime==1.4.4 opencv-python-headless==4.9.0.80 -pymongo==4.17.0 +pymongo==4.18.2 langchain-mongodb==0.11.0 langchain-ollama==1.1.0 langchain-huggingface==1.2.2 langchain-google-genai==4.2.5 -cryptography==48.0.1 +cryptography==50.0.1 python-magic==0.4.27 python-pptx==1.0.2 xlrd==2.0.2 diff --git a/tests/test_query_execution.py b/tests/test_query_execution.py new file mode 100644 index 00000000..4e4d7518 --- /dev/null +++ b/tests/test_query_execution.py @@ -0,0 +1,255 @@ +"""Exercise query dispatch with real worker threads and no external services.""" + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from fastapi import HTTPException +from langchain_core.documents import Document + +from app.config import VectorDBType +from app.models import QueryMultipleBody, QueryRequestBody +from app.routes import document_routes +from app.services.vector_store.async_pg_vector import AsyncPgVector +from app.services.vector_store.extended_pg_vector import ExtendedPgVector + + +class DummyAsyncPgVector(AsyncPgVector): + def __init__(self, embedding_function): + self.embedding_function = embedding_function + self._thread_pool = None + self._bind = None + + +@pytest.fixture(params=["single", "multiple"]) +def query_route(request): + return request.param + + +@pytest.fixture(params=["async", "sync"]) +def query_context(request, monkeypatch, query_route): + """Use the real async adapter or the route's synchronous fallback.""" + embedding = Mock(return_value=[0.1, 0.2, 0.3]) + provider = SimpleNamespace(embed_query=embedding) + hits = [ + (Document(page_content="near", metadata={"file_id": "file-1"}), 0.2), + (Document(page_content="far", metadata={"file_id": "file-1"}), 0.8), + ] + search = Mock(return_value=hits) + if request.param == "async": + store = DummyAsyncPgVector(provider) + monkeypatch.setattr( + ExtendedPgVector, + "similarity_search_with_score_by_vector", + lambda self, *args, **kwargs: search(*args, **kwargs), + ) + else: + store = SimpleNamespace( + embedding_function=provider, + similarity_search_with_score_by_vector=search, + ) + monkeypatch.setattr(document_routes, "vector_store", store) + monkeypatch.setattr(document_routes, "RAG_DISTANCE_THRESHOLD", None) + monkeypatch.setattr(document_routes, "VECTOR_DB_TYPE", VectorDBType.PGVECTOR) + document_routes.get_cached_query_embedding.cache_clear() + + with ThreadPoolExecutor(max_workers=2, thread_name_prefix="query-test") as pool: + req = SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(thread_pool=pool)), + state=SimpleNamespace(user={"id": "owner-1"}), + ) + + async def query(text="question", entity_id=None): + body_args = dict(query=text, k=7, entity_id=entity_id) + if query_route == "single": + return await document_routes.query_embeddings_by_file_id( + QueryRequestBody(file_id="file-1", **body_args), req + ) + return await document_routes.query_embeddings_by_file_ids( + req, QueryMultipleBody(file_ids=["file-1", "file-2"], **body_args) + ) + + try: + yield SimpleNamespace( + query=query, + embedding=embedding, + search=search, + hits=hits, + request=req, + ) + finally: + # Drain cancelled read-only work before clearing its shared cache. + pool.shutdown(wait=True) + document_routes.get_cached_query_embedding.cache_clear() + + +def _block_stage(context, stage, loop, started, release, finished): + operation = getattr(context, stage) + result = operation.return_value + + def blocking(*args, **kwargs): + assert threading.current_thread().name.startswith("query-test") + loop.call_soon_threadsafe(started.set) + try: + if not release.wait(timeout=5): + raise TimeoutError("event loop did not release the worker") + return result + finally: + loop.call_soon_threadsafe(finished.set) + + operation.side_effect = blocking + + +@pytest.mark.parametrize("stage", ["embedding", "search"]) +async def test_blocking_stages_leave_event_loop_responsive(query_context, stage): + started, finished = asyncio.Event(), asyncio.Event() + release = threading.Event() + _block_stage( + query_context, stage, asyncio.get_running_loop(), started, release, finished + ) + task = asyncio.create_task(query_context.query()) + try: + await asyncio.wait_for(started.wait(), timeout=2) + # This coroutine can resume while the provider or database is still blocked. + assert not task.done() + assert not finished.is_set() + finally: + release.set() + result = await asyncio.wait_for(task, timeout=2) + assert result == query_context.hits + + +async def test_concurrent_queries_share_the_configured_worker_limit(query_context): + loop = asyncio.get_running_loop() + saturated = asyncio.Event() + release = threading.Event() + lock = threading.Lock() + started = 0 + + def embed(query): + nonlocal started + with lock: + started += 1 + if started == 2: + loop.call_soon_threadsafe(saturated.set) + if not release.wait(timeout=5): + raise TimeoutError("event loop did not release the workers") + return [0.1, 0.2, 0.3] + + query_context.embedding.side_effect = embed + tasks = [asyncio.create_task(query_context.query(str(i))) for i in range(5)] + try: + await asyncio.wait_for(saturated.wait(), timeout=2) + # Two workers are occupied. The remaining requests must wait in this + # same pool, not start in a separate or unbounded set of threads. + with lock: + assert started == 2 + assert all(not task.done() for task in tasks) + finally: + release.set() + results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=2) + assert results == [query_context.hits] * 5 + assert started == 5 + + +async def test_cache_reuses_vectors_but_not_scoped_results(query_context, query_route): + assert await query_context.query() == query_context.hits + query_context.request.state.user = {"id": "owner-2"} + assert await query_context.query(entity_id="agent-1") == query_context.hits + query_context.embedding.assert_called_once_with("question") + assert query_context.search.call_count == 2 + expected_file_clause = ( + {"file_id": {"$eq": "file-1"}} + if query_route == "single" + else {"file_id": {"$in": ["file-1", "file-2"]}} + ) + for call, owners in zip( + query_context.search.call_args_list, + [["owner-1"], ["agent-1", "owner-2"]], + ): + # The real async adapter forwards positional arguments; the sync path + # forwards keyword arguments. Both must carry the same pre-ranking scope. + embedding = call.args[0] + k = call.args[1] if len(call.args) > 1 else call.kwargs["k"] + predicate = call.args[2] if len(call.args) > 2 else call.kwargs["filter"] + assert embedding == [0.1, 0.2, 0.3] + assert k == 7 + assert predicate == { + "$and": [expected_file_clause, {"user_id": {"$in": owners}}] + } + + +@pytest.mark.parametrize("stage", ["embedding", "search"]) +@pytest.mark.parametrize( + "error,status", + [ + (ValueError("unavailable"), 500), + (HTTPException(429, "retry"), 429), + (StopIteration("exhausted"), 500), + ], +) +async def test_worker_errors_propagate_and_allow_retry( + query_context, stage, error, status +): + operation = getattr(query_context, stage) + operation.side_effect = error + with pytest.raises(HTTPException) as caught: + await asyncio.wait_for(query_context.query(), timeout=2) + assert caught.value.status_code == status + if stage == "embedding": + query_context.search.assert_not_called() + operation.side_effect = None + assert await query_context.query() == query_context.hits + assert operation.call_count == 2 + + +@pytest.mark.parametrize("stage", ["embedding", "search"]) +async def test_cancellation_propagates_while_read_only_worker_finishes( + query_context, stage +): + started, finished = asyncio.Event(), asyncio.Event() + release = threading.Event() + _block_stage( + query_context, stage, asyncio.get_running_loop(), started, release, finished + ) + task = asyncio.create_task(query_context.query()) + try: + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=2) + assert not finished.is_set() + finally: + release.set() + await asyncio.wait_for(finished.wait(), timeout=2) + await asyncio.gather(task, return_exceptions=True) + if stage == "embedding": + query_context.search.assert_not_called() + + +@pytest.mark.parametrize("db_type", [VectorDBType.PGVECTOR, VectorDBType.ATLAS_MONGO]) +async def test_distance_threshold_semantics_unchanged( + query_context, monkeypatch, db_type +): + monkeypatch.setattr(document_routes, "VECTOR_DB_TYPE", db_type) + monkeypatch.setattr(document_routes, "RAG_DISTANCE_THRESHOLD", 0.3) + expected = ( + query_context.hits[:1] + if db_type == VectorDBType.PGVECTOR + else query_context.hits + ) + assert await query_context.query() == expected + + +async def test_empty_result_contract_unchanged(query_context, query_route): + query_context.search.return_value = [] + if query_route == "single": + assert await query_context.query() == [] + else: + with pytest.raises(HTTPException) as caught: + await query_context.query() + assert caught.value.status_code == 404 + assert caught.value.detail == "No documents found for the given query" diff --git a/tests/utils/test_document_loader.py b/tests/utils/test_document_loader.py index 1de02e5f..b1b81a21 100644 --- a/tests/utils/test_document_loader.py +++ b/tests/utils/test_document_loader.py @@ -203,6 +203,45 @@ def test_safe_pdf_loader_requires_api_key(): loader.load() +def test_mistral_client_uses_sdk_default_server_without_override(): + """Without MISTRAL_SERVER_URL the client is built with the SDK's default server.""" + from app.utils import document_loader + + fake_module, _ = _make_mistral_module(pages=[]) + + with patch.dict( + "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object( + document_loader, "MISTRAL_API_KEY", "test-key" + ), patch.object(document_loader, "MISTRAL_SERVER_URL", ""): + document_loader._mistral_ocr_client() + + fake_module.Mistral.assert_called_once_with(api_key="test-key") + + +def test_mistral_client_uses_configured_server_url(): + """MISTRAL_SERVER_URL (e.g. the EU regional endpoint) is passed to the client.""" + from app.utils import document_loader + from app.utils.document_loader import SafePyPDFLoader + + fake_module, client = _make_mistral_module(pages=[]) + + loader = SafePyPDFLoader("dummy.pdf") + with patch.dict( + "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object( + document_loader, "MISTRAL_API_KEY", "test-key" + ), patch.object( + document_loader, "MISTRAL_SERVER_URL", "https://api.eu.mistral.ai" + ), patch.object(SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data"): + loader.load() + + fake_module.Mistral.assert_called_once_with( + api_key="test-key", server_url="https://api.eu.mistral.ai" + ) + client.ocr.process.assert_called_once() + + MARKDOWN_SAMPLE = ( "# Heading\n\n" "**bold** and *italic* text with a [link](https://example.com).\n\n"