backend/open_webui/retrieval/loaders/mistral.py

import requests
import aiohttp
import asyncio
import logging
import os
import sys
import time
from typing import List, Dict, Any
from contextlib import asynccontextmanager

from langchain_core.documents import Document
from open_webui.env import SRC_LOG_LEVELS, GLOBAL_LOG_LEVEL

logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])


class MistralLoader:
    """
    Enhanced Mistral OCR loader with both sync and async support.
    Loads documents by processing them through the Mistral OCR API.

    Performance Optimizations:
    - Differentiated timeouts for different operations
    - Intelligent retry logic with exponential backoff
    - Memory-efficient file streaming for large files
    - Connection pooling and keepalive optimization
    - Semaphore-based concurrency control for batch processing
    - Enhanced error handling with retryable error classification
    """

    BASE_API_URL = "https://api.mistral.ai/v1"

    def __init__(
        self,
        api_key: str,
        file_path: str,
        timeout: int = 300,  # 5 minutes default
        max_retries: int = 3,
        enable_debug_logging: bool = False,
    ):
        """
        Initializes the loader with enhanced features.

        Args:
            api_key: Your Mistral API key.
            file_path: The local path to the PDF file to process.
            timeout: Request timeout in seconds.
            max_retries: Maximum number of retry attempts.
            enable_debug_logging: Enable detailed debug logs.
        """
        if not api_key:
            raise ValueError("API key cannot be empty.")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found at {file_path}")

        self.api_key = api_key
        self.file_path = file_path
        self.timeout = timeout
        self.max_retries = max_retries
        self.debug = enable_debug_logging

        # PERFORMANCE OPTIMIZATION: Differentiated timeouts for different operations
        # This prevents long-running OCR operations from affecting quick operations
        # and improves user experience by failing fast on operations that should be quick
        self.upload_timeout = min(
            timeout, 120
        )  # Cap upload at 2 minutes - prevents hanging on large files
        self.url_timeout = (
            30  # URL requests should be fast - fail quickly if API is slow
        )
        self.ocr_timeout = (
            timeout  # OCR can take the full timeout - this is the heavy operation
        )
        self.cleanup_timeout = (
            30  # Cleanup should be quick - don't hang on file deletion
        )

        # PERFORMANCE OPTIMIZATION: Pre-compute file info to avoid repeated filesystem calls
        # This avoids multiple os.path.basename() and os.path.getsize() calls during processing
        self.file_name = os.path.basename(file_path)
        self.file_size = os.path.getsize(file_path)

        # ENHANCEMENT: Added User-Agent for better API tracking and debugging
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "User-Agent": "OpenWebUI-MistralLoader/2.0",  # Helps API provider track usage
        }

    def _debug_log(self, message: str, *args) -> None:
        """
        PERFORMANCE OPTIMIZATION: Conditional debug logging for performance.

        Only processes debug messages when debug mode is enabled, avoiding
        string formatting overhead in production environments.
        """
        if self.debug:
            log.debug(message, *args)

    def _handle_response(self, response: requests.Response) -> Dict[str, Any]:
        """Checks response status and returns JSON content."""
        try:
            response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
            # Handle potential empty responses for certain successful requests (e.g., DELETE)
            if response.status_code == 204 or not response.content:
                return {}  # Return empty dict if no content
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            log.error(f"HTTP error occurred: {http_err} - Response: {response.text}")
            raise
        except requests.exceptions.RequestException as req_err:
            log.error(f"Request exception occurred: {req_err}")
            raise
        except ValueError as json_err:  # Includes JSONDecodeError
            log.error(f"JSON decode error: {json_err} - Response: {response.text}")
            raise  # Re-raise after logging

    async def _handle_response_async(
        self, response: aiohttp.ClientResponse
    ) -> Dict[str, Any]:
        """Async version of response handling with better error info."""
        try:
            response.raise_for_status()

            # Check content type
            content_type = response.headers.get("content-type", "")
            if "application/json" not in content_type:
                if response.status == 204:
                    return {}
                text = await response.text()
                raise ValueError(
                    f"Unexpected content type: {content_type}, body: {text[:200]}..."
                )

            return await response.json()

        except aiohttp.ClientResponseError as e:
            error_text = await response.text() if response else "No response"
            log.error(f"HTTP {e.status}: {e.message} - Response: {error_text[:500]}")
            raise
        except aiohttp.ClientError as e:
            log.error(f"Client error: {e}")
            raise
        except Exception as e:
            log.error(f"Unexpected error processing response: {e}")
            raise

    def _is_retryable_error(self, error: Exception) -> bool:
        """
        ENHANCEMENT: Intelligent error classification for retry logic.

        Determines if an error is retryable based on its type and status code.
        This prevents wasting time retrying errors that will never succeed
        (like authentication errors) while ensuring transient errors are retried.

        Retryable errors:
        - Network connection errors (temporary network issues)
        - Timeouts (server might be temporarily overloaded)
        - Server errors (5xx status codes - server-side issues)
        - Rate limiting (429 status - temporary throttling)

        Non-retryable errors:
        - Authentication errors (401, 403 - won't fix with retry)
        - Bad request errors (400 - malformed request)
        - Not found errors (404 - resource doesn't exist)
        """
        if isinstance(error, requests.exceptions.ConnectionError):
            return True  # Network issues are usually temporary
        if isinstance(error, requests.exceptions.Timeout):
            return True  # Timeouts might resolve on retry
        if isinstance(error, requests.exceptions.HTTPError):
            # Only retry on server errors (5xx) or rate limits (429)
            if hasattr(error, "response") and error.response is not None:
                status_code = error.response.status_code
                return status_code >= 500 or status_code == 429
            return False
        if isinstance(
            error, (aiohttp.ClientConnectionError, aiohttp.ServerTimeoutError)
        ):
            return True  # Async network/timeout errors are retryable
        if isinstance(error, aiohttp.ClientResponseError):
            return error.status >= 500 or error.status == 429
        return False  # All other errors are non-retryable

    def _retry_request_sync(self, request_func, *args, **kwargs):
        """
        ENHANCEMENT: Synchronous retry logic with intelligent error classification.

        Uses exponential backoff with jitter to avoid thundering herd problems.
        The wait time increases exponentially but is capped at 30 seconds to
        prevent excessive delays. Only retries errors that are likely to succeed
        on subsequent attempts.
        """
        for attempt in range(self.max_retries):
            try:
                return request_func(*args, **kwargs)
            except Exception as e:
                if attempt == self.max_retries - 1 or not self._is_retryable_error(e):
                    raise

                # PERFORMANCE OPTIMIZATION: Exponential backoff with cap
                # Prevents overwhelming the server while ensuring reasonable retry delays
                wait_time = min((2**attempt) + 0.5, 30)  # Cap at 30 seconds
                log.warning(
                    f"Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. "
                    f"Retrying in {wait_time}s..."
                )
                time.sleep(wait_time)

    async def _retry_request_async(self, request_func, *args, **kwargs):
        """
        ENHANCEMENT: Async retry logic with intelligent error classification.

        Async version of retry logic that doesn't block the event loop during
        wait periods. Uses the same exponential backoff strategy as sync version.
        """
        for attempt in range(self.max_retries):
            try:
                return await request_func(*args, **kwargs)
            except Exception as e:
                if attempt == self.max_retries - 1 or not self._is_retryable_error(e):
                    raise

                # PERFORMANCE OPTIMIZATION: Non-blocking exponential backoff
                wait_time = min((2**attempt) + 0.5, 30)  # Cap at 30 seconds
                log.warning(
                    f"Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. "
                    f"Retrying in {wait_time}s..."
                )
                await asyncio.sleep(wait_time)  # Non-blocking wait

    def _upload_file(self) -> str:
        """
        PERFORMANCE OPTIMIZATION: Enhanced file upload with streaming consideration.

        Uploads the file to Mistral for OCR processing (sync version).
        Uses context manager for file handling to ensure proper resource cleanup.
        Although streaming is not enabled for this endpoint, the file is opened
        in a context manager to minimize memory usage duration.
        """
        log.info("Uploading file to Mistral API")
        url = f"{self.BASE_API_URL}/files"

        def upload_request():
            # MEMORY OPTIMIZATION: Use context manager to minimize file handle lifetime
            # This ensures the file is closed immediately after reading, reducing memory usage
            with open(self.file_path, "rb") as f:
                files = {"file": (self.file_name, f, "application/pdf")}
                data = {"purpose": "ocr"}

                # NOTE: stream=False is required for this endpoint
                # The Mistral API doesn't support chunked uploads for this endpoint
                response = requests.post(
                    url,
                    headers=self.headers,
                    files=files,
                    data=data,
                    timeout=self.upload_timeout,  # Use specialized upload timeout
                    stream=False,  # Keep as False for this endpoint
                )

            return self._handle_response(response)

        try:
            response_data = self._retry_request_sync(upload_request)
            file_id = response_data.get("id")
            if not file_id:
                raise ValueError("File ID not found in upload response.")
            log.info(f"File uploaded successfully. File ID: {file_id}")
            return file_id
        except Exception as e:
            log.error(f"Failed to upload file: {e}")
            raise

    async def _upload_file_async(self, session: aiohttp.ClientSession) -> str:
        """Async file upload with streaming for better memory efficiency."""
        url = f"{self.BASE_API_URL}/files"

        async def upload_request():
            # Create multipart writer for streaming upload
            writer = aiohttp.MultipartWriter("form-data")

            # Add purpose field
            purpose_part = writer.append("ocr")
            purpose_part.set_content_disposition("form-data", name="purpose")

            # Add file part with streaming
            file_part = writer.append_payload(
                aiohttp.streams.FilePayload(
                    self.file_path,
                    filename=self.file_name,
                    content_type="application/pdf",
                )
            )
            file_part.set_content_disposition(
                "form-data", name="file", filename=self.file_name
            )

            self._debug_log(
                f"Uploading file: {self.file_name} ({self.file_size:,} bytes)"
            )

            async with session.post(
                url,
                data=writer,
                headers=self.headers,
                timeout=aiohttp.ClientTimeout(total=self.upload_timeout),
            ) as response:
                return await self._handle_response_async(response)

        response_data = await self._retry_request_async(upload_request)

        file_id = response_data.get("id")
        if not file_id:
            raise ValueError("File ID not found in upload response.")

        log.info(f"File uploaded successfully. File ID: {file_id}")
        return file_id

    def _get_signed_url(self, file_id: str) -> str:
        """Retrieves a temporary signed URL for the uploaded file (sync version)."""
        log.info(f"Getting signed URL for file ID: {file_id}")
        url = f"{self.BASE_API_URL}/files/{file_id}/url"
        params = {"expiry": 1}
        signed_url_headers = {**self.headers, "Accept": "application/json"}

        def url_request():
            response = requests.get(
                url, headers=signed_url_headers, params=params, timeout=self.url_timeout
            )
            return self._handle_response(response)

        try:
            response_data = self._retry_request_sync(url_request)
            signed_url = response_data.get("url")
            if not signed_url:
                raise ValueError("Signed URL not found in response.")
            log.info("Signed URL received.")
            return signed_url
        except Exception as e:
            log.error(f"Failed to get signed URL: {e}")
            raise

    async def _get_signed_url_async(
        self, session: aiohttp.ClientSession, file_id: str
    ) -> str:
        """Async signed URL retrieval."""
        url = f"{self.BASE_API_URL}/files/{file_id}/url"
        params = {"expiry": 1}

        headers = {**self.headers, "Accept": "application/json"}

        async def url_request():
            self._debug_log(f"Getting signed URL for file ID: {file_id}")
            async with session.get(
                url,
                headers=headers,
                params=params,
                timeout=aiohttp.ClientTimeout(total=self.url_timeout),
            ) as response:
                return await self._handle_response_async(response)

        response_data = await self._retry_request_async(url_request)

        signed_url = response_data.get("url")
        if not signed_url:
            raise ValueError("Signed URL not found in response.")

        self._debug_log("Signed URL received successfully")
        return signed_url

    def _process_ocr(self, signed_url: str) -> Dict[str, Any]:
        """Sends the signed URL to the OCR endpoint for processing (sync version)."""
        log.info("Processing OCR via Mistral API")
        url = f"{self.BASE_API_URL}/ocr"
        ocr_headers = {
            **self.headers,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        payload = {
            "model": "mistral-ocr-latest",
            "document": {
                "type": "document_url",
                "document_url": signed_url,
            },
            "include_image_base64": False,
        }

        def ocr_request():
            response = requests.post(
                url, headers=ocr_headers, json=payload, timeout=self.ocr_timeout
            )
            return self._handle_response(response)

        try:
            ocr_response = self._retry_request_sync(ocr_request)
            log.info("OCR processing done.")
            self._debug_log("OCR response: %s", ocr_response)
            return ocr_response
        except Exception as e:
            log.error(f"Failed during OCR processing: {e}")
            raise

    async def _process_ocr_async(
        self, session: aiohttp.ClientSession, signed_url: str
    ) -> Dict[str, Any]:
        """Async OCR processing with timing metrics."""
        url = f"{self.BASE_API_URL}/ocr"

        headers = {
            **self.headers,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

        payload = {
            "model": "mistral-ocr-latest",
            "document": {
                "type": "document_url",
                "document_url": signed_url,
            },
            "include_image_base64": False,
        }

        async def ocr_request():
            log.info("Starting OCR processing via Mistral API")
            start_time = time.time()

            async with session.post(
                url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=self.ocr_timeout),
            ) as response:
                ocr_response = await self._handle_response_async(response)

            processing_time = time.time() - start_time
            log.info(f"OCR processing completed in {processing_time:.2f}s")

            return ocr_response

        return await self._retry_request_async(ocr_request)

    def _delete_file(self, file_id: str) -> None:
        """Deletes the file from Mistral storage (sync version)."""
        log.info(f"Deleting uploaded file ID: {file_id}")
        url = f"{self.BASE_API_URL}/files/{file_id}"

        try:
            response = requests.delete(
                url, headers=self.headers, timeout=self.cleanup_timeout
            )
            delete_response = self._handle_response(response)
            log.info(f"File deleted successfully: {delete_response}")
        except Exception as e:
            # Log error but don't necessarily halt execution if deletion fails
            log.error(f"Failed to delete file ID {file_id}: {e}")

    async def _delete_file_async(
        self, session: aiohttp.ClientSession, file_id: str
    ) -> None:
        """Async file deletion with error tolerance."""
        try:

            async def delete_request():
                self._debug_log(f"Deleting file ID: {file_id}")
                async with session.delete(
                    url=f"{self.BASE_API_URL}/files/{file_id}",
                    headers=self.headers,
                    timeout=aiohttp.ClientTimeout(
                        total=self.cleanup_timeout
                    ),  # Shorter timeout for cleanup
                ) as response:
                    return await self._handle_response_async(response)

            await self._retry_request_async(delete_request)
            self._debug_log(f"File {file_id} deleted successfully")

        except Exception as e:
            # Don't fail the entire process if cleanup fails
            log.warning(f"Failed to delete file ID {file_id}: {e}")

    @asynccontextmanager
    async def _get_session(self):
        """Context manager for HTTP session with optimized settings."""
        connector = aiohttp.TCPConnector(
            limit=20,  # Increased total connection limit for better throughput
            limit_per_host=10,  # Increased per-host limit for API endpoints
            ttl_dns_cache=600,  # Longer DNS cache TTL (10 minutes)
            use_dns_cache=True,
            keepalive_timeout=60,  # Increased keepalive for connection reuse
            enable_cleanup_closed=True,
            force_close=False,  # Allow connection reuse
            resolver=aiohttp.AsyncResolver(),  # Use async DNS resolver
        )

        timeout = aiohttp.ClientTimeout(
            total=self.timeout,
            connect=30,  # Connection timeout
            sock_read=60,  # Socket read timeout
        )

        async with aiohttp.ClientSession(
            connector=connector,
            timeout=timeout,
            headers={"User-Agent": "OpenWebUI-MistralLoader/2.0"},
            raise_for_status=False,  # We handle status codes manually
        ) as session:
            yield session

    def _process_results(self, ocr_response: Dict[str, Any]) -> List[Document]:
        """Process OCR results into Document objects with enhanced metadata and memory efficiency."""
        pages_data = ocr_response.get("pages")
        if not pages_data:
            log.warning("No pages found in OCR response.")
            return [
                Document(
                    page_content="No text content found",
                    metadata={"error": "no_pages", "file_name": self.file_name},
                )
            ]

        documents = []
        total_pages = len(pages_data)
        skipped_pages = 0

        # Process pages in a memory-efficient way
        for page_data in pages_data:
            page_content = page_data.get("markdown")
            page_index = page_data.get("index")  # API uses 0-based index

            if page_content is None or page_index is None:
                skipped_pages += 1
                self._debug_log(
                    f"Skipping page due to missing 'markdown' or 'index'. Data keys: {list(page_data.keys())}"
                )
                continue

            # Clean up content efficiently with early exit for empty content
            if isinstance(page_content, str):
                cleaned_content = page_content.strip()
            else:
                cleaned_content = str(page_content).strip()

            if not cleaned_content:
                skipped_pages += 1
                self._debug_log(f"Skipping empty page {page_index}")
                continue

            # Create document with optimized metadata
            documents.append(
                Document(
                    page_content=cleaned_content,
                    metadata={
                        "page": page_index,  # 0-based index from API
                        "page_label": page_index + 1,  # 1-based label for convenience
                        "total_pages": total_pages,
                        "file_name": self.file_name,
                        "file_size": self.file_size,
                        "processing_engine": "mistral-ocr",
                        "content_length": len(cleaned_content),
                    },
                )
            )

        if skipped_pages > 0:
            log.info(
                f"Processed {len(documents)} pages, skipped {skipped_pages} empty/invalid pages"
            )

        if not documents:
            # Case where pages existed but none had valid markdown/index
            log.warning(
                "OCR response contained pages, but none had valid content/index."
            )
            return [
                Document(
                    page_content="No valid text content found in document",
                    metadata={
                        "error": "no_valid_pages",
                        "total_pages": total_pages,
                        "file_name": self.file_name,
                    },
                )
            ]

        return documents

    def load(self) -> List[Document]:
        """
        Executes the full OCR workflow: upload, get URL, process OCR, delete file.
        Synchronous version for backward compatibility.

        Returns:
            A list of Document objects, one for each page processed.
        """
        file_id = None
        start_time = time.time()

        try:
            # 1. Upload file
            file_id = self._upload_file()

            # 2. Get Signed URL
            signed_url = self._get_signed_url(file_id)

            # 3. Process OCR
            ocr_response = self._process_ocr(signed_url)

            # 4. Process results
            documents = self._process_results(ocr_response)

            total_time = time.time() - start_time
            log.info(
                f"Sync OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents"
            )

            return documents

        except Exception as e:
            total_time = time.time() - start_time
            log.error(
                f"An error occurred during the loading process after {total_time:.2f}s: {e}"
            )
            # Return an error document on failure
            return [
                Document(
                    page_content=f"Error during processing: {e}",
                    metadata={
                        "error": "processing_failed",
                        "file_name": self.file_name,
                    },
                )
            ]
        finally:
            # 5. Delete file (attempt even if prior steps failed after upload)
            if file_id:
                try:
                    self._delete_file(file_id)
                except Exception as del_e:
                    # Log deletion error, but don't overwrite original error if one occurred
                    log.error(
                        f"Cleanup error: Could not delete file ID {file_id}. Reason: {del_e}"
                    )

    async def load_async(self) -> List[Document]:
        """
        Asynchronous OCR workflow execution with optimized performance.

        Returns:
            A list of Document objects, one for each page processed.
        """
        file_id = None
        start_time = time.time()

        try:
            async with self._get_session() as session:
                # 1. Upload file with streaming
                file_id = await self._upload_file_async(session)

                # 2. Get signed URL
                signed_url = await self._get_signed_url_async(session, file_id)

                # 3. Process OCR
                ocr_response = await self._process_ocr_async(session, signed_url)

                # 4. Process results
                documents = self._process_results(ocr_response)

                total_time = time.time() - start_time
                log.info(
                    f"Async OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents"
                )

                return documents

        except Exception as e:
            total_time = time.time() - start_time
            log.error(f"Async OCR workflow failed after {total_time:.2f}s: {e}")
            return [
                Document(
                    page_content=f"Error during OCR processing: {e}",
                    metadata={
                        "error": "processing_failed",
                        "file_name": self.file_name,
                    },
                )
            ]
        finally:
            # 5. Cleanup - always attempt file deletion
            if file_id:
                try:
                    async with self._get_session() as session:
                        await self._delete_file_async(session, file_id)
                except Exception as cleanup_error:
                    log.error(f"Cleanup failed for file ID {file_id}: {cleanup_error}")

    @staticmethod
    async def load_multiple_async(
        loaders: List["MistralLoader"],
        max_concurrent: int = 5,  # Limit concurrent requests
    ) -> List[List[Document]]:
        """
        Process multiple files concurrently with controlled concurrency.

        Args:
            loaders: List of MistralLoader instances
            max_concurrent: Maximum number of concurrent requests

        Returns:
            List of document lists, one for each loader
        """
        if not loaders:
            return []

        log.info(
            f"Starting concurrent processing of {len(loaders)} files with max {max_concurrent} concurrent"
        )
        start_time = time.time()

        # Use semaphore to control concurrency
        semaphore = asyncio.Semaphore(max_concurrent)

        async def process_with_semaphore(loader: "MistralLoader") -> List[Document]:
            async with semaphore:
                return await loader.load_async()

        # Process all files with controlled concurrency
        tasks = [process_with_semaphore(loader) for loader in loaders]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Handle any exceptions in results
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                log.error(f"File {i} failed: {result}")
                processed_results.append(
                    [
                        Document(
                            page_content=f"Error processing file: {result}",
                            metadata={
                                "error": "batch_processing_failed",
                                "file_index": i,
                            },
                        )
                    ]
                )
            else:
                processed_results.append(result)

        # MONITORING: Log comprehensive batch processing statistics
        total_time = time.time() - start_time
        total_docs = sum(len(docs) for docs in processed_results)
        success_count = sum(
            1 for result in results if not isinstance(result, Exception)
        )
        failure_count = len(results) - success_count

        log.info(
            f"Batch processing completed in {total_time:.2f}s: "
            f"{success_count} files succeeded, {failure_count} files failed, "
            f"produced {total_docs} total documents"
        )

        return processed_results