request_utils.py

from __future__ import annotations

import asyncio
import io
import json
import logging
import os
import tempfile
from typing import Tuple, Any, BinaryIO, Optional
from urllib.parse import urlparse

import httpx
from httpx import URL
from httpx._multipart import DataField, FileField

from unstructured_client._hooks.custom.common import UNSTRUCTURED_CLIENT_LOGGER_NAME
from unstructured_client._hooks.custom.form_utils import (
    PARTITION_FORM_FILES_KEY,
    PARTITION_FORM_SPLIT_PDF_PAGE_KEY,
    PARTITION_FORM_SPLIT_PDF_ALLOW_FAILED_KEY,
    PARTITION_FORM_SPLIT_CACHE_TMP_DATA_KEY,
    PARTITION_FORM_SPLIT_CACHE_TMP_DATA_DIR_KEY,
    PARTITION_FORM_PAGE_RANGE_KEY,
    PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
    FormData,
)
from unstructured_client.models import shared
from unstructured_client.utils import (
    BackoffStrategy,
    Retries,
    RetryConfig,
    retry_async,
    serialize_request_body,
)

logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)


def create_default_retry_config() -> RetryConfig:
    one_second = 1000
    one_minute = 1000 * 60
    return RetryConfig(
        "backoff",
        BackoffStrategy(
            initial_interval=one_second * 3,
            max_interval=one_minute * 12,
            max_elapsed_time=one_minute * 30,
            exponent=1.88,
        ),
        retry_connection_errors=True,
    )


def create_split_retry_config(retry_config: Optional[RetryConfig]) -> RetryConfig:
    if retry_config is None:
        return create_default_retry_config()

    # Split chunk requests run after the top-level dummy request has already
    # succeeded, so they must preserve their own transport-level retry budget.
    return RetryConfig(
        retry_config.strategy,
        retry_config.backoff,
        True,
    )


def get_multipart_stream_fields(request: httpx.Request) -> dict[str, Any]:
    """Extracts the multipart fields from the request.

    Args:
        request: The request object.

    Returns:
        The multipart fields.

    Raises:
        Exception: If the filename is not set
    """
    content_type = request.headers.get("Content-Type", "")
    if "multipart" not in content_type:
        return {}
    if request.stream is None or not hasattr(request.stream, "fields"):
        return {}
    fields = request.stream.fields

    mapped_fields: dict[str, Any] = {}
    for field in fields:
        if isinstance(field, DataField):
            if "[]" in field.name:
                name = field.name.replace("[]", "")
                if name not in mapped_fields:
                    mapped_fields[name] = []
                mapped_fields[name].append(field.value)
            mapped_fields[field.name] = field.value
        elif isinstance(field, FileField):
            if field.filename is None or not field.filename.strip():
                raise ValueError("Filename can't be an empty string.")
            mapped_fields[field.name] = {
                "filename": field.filename,
                "content_type": field.headers.get("Content-Type", ""),
                "file": field.file,
            }
    return mapped_fields


def create_pdf_chunk_request_params(
    form_data: FormData, page_number: int
) -> dict[str, Any]:
    """Creates the request body for the partition API."

    Args:
        form_data: The form data.
        page_number: The page number.

    Returns:
        The updated request payload for the chunk.
    """
    fields_to_drop = [
        PARTITION_FORM_SPLIT_PDF_PAGE_KEY,
        PARTITION_FORM_SPLIT_PDF_ALLOW_FAILED_KEY,
        PARTITION_FORM_FILES_KEY,
        PARTITION_FORM_PAGE_RANGE_KEY,
        PARTITION_FORM_PAGE_RANGE_KEY.replace("[]", ""),
        PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
        PARTITION_FORM_SPLIT_CACHE_TMP_DATA_KEY,
        PARTITION_FORM_SPLIT_CACHE_TMP_DATA_DIR_KEY,
    ]
    chunk_payload = {
        key: form_data[key] for key in form_data if key not in fields_to_drop
    }
    chunk_payload[PARTITION_FORM_SPLIT_PDF_PAGE_KEY] = "false"
    chunk_payload[PARTITION_FORM_STARTING_PAGE_NUMBER_KEY] = str(page_number)
    return chunk_payload


def create_pdf_chunk_request(
    form_data: FormData,
    pdf_chunk: Tuple[BinaryIO, int],
    original_request: httpx.Request,
    filename: str,
) -> httpx.Request:
    """Creates a new request object with the updated payload for the partition API.

    Args:
        form_data: The form data.
        pdf_chunk: Tuple of pdf chunk contents (can be both io.BytesIO or
            a file object created with e.g. open()) and the page number.
        original_request: The original request.
        filename: The filename.

    Returns:
        The updated request object.
    """
    pdf_chunk_file, page_number = pdf_chunk
    data = create_pdf_chunk_request_params(form_data, page_number)
    original_headers = prepare_request_headers(original_request.headers)

    pdf_chunk_content: BinaryIO | bytes = (
        pdf_chunk_file.getvalue()
        if isinstance(pdf_chunk_file, io.BytesIO)
        else pdf_chunk_file
    )

    pdf_chunk_partition_params = shared.PartitionParameters(
        files=shared.Files(
            content=pdf_chunk_content,
            file_name=filename,
            content_type="application/pdf",
        ),
        **data,
    )
    serialized_body = serialize_request_body(
        pdf_chunk_partition_params,
        False,
        False,
        "multipart",
        shared.PartitionParameters,
    )
    if serialized_body is None:
        raise ValueError("Failed to serialize the request body.")
    return httpx.Request(
        method="POST",
        url=original_request.url or "",
        headers={**original_headers},
        content=serialized_body.content,
        data=serialized_body.data,
        files=serialized_body.files,
    )


async def call_api_async(
    client: httpx.AsyncClient,
    pdf_chunk_request: httpx.Request,
    pdf_chunk_file: BinaryIO,
    limiter: asyncio.Semaphore,
    retry_config: Optional[RetryConfig] = None,
    operation_id: Optional[str] = None,
    chunk_index: Optional[int] = None,
    page_number: Optional[int] = None,
) -> httpx.Response:
    retryable_codes = ["5xx"]
    effective_retry_config = create_split_retry_config(retry_config)

    async def do_request():
        return await client.send(pdf_chunk_request)

    async with limiter:
        try:
            logger.debug(
                "split_pdf event=chunk_request_send operation_id=%s chunk_index=%s page_number=%s retry_config_mode=%s retry_connection_errors=%s",
                operation_id,
                chunk_index,
                page_number,
                "sdk_custom" if retry_config is not None else "sdk_default_or_unset",
                effective_retry_config.retry_connection_errors,
            )
            response = await retry_async(
                do_request, Retries(effective_retry_config, retryable_codes)
            )
            logger.debug(
                "split_pdf event=chunk_request_response operation_id=%s chunk_index=%s page_number=%s status_code=%d",
                operation_id,
                chunk_index,
                page_number,
                response.status_code,
            )
            return response
        except Exception as e:
            logger.error(
                "split_pdf event=chunk_request_error operation_id=%s chunk_index=%s page_number=%s error_type=%s error=%s",
                operation_id,
                chunk_index,
                page_number,
                type(e).__name__,
                e,
                exc_info=e,
            )
            raise e
        finally:
            if not isinstance(pdf_chunk_file, io.BytesIO) and not pdf_chunk_file.closed:
                pdf_chunk_file.close()


def prepare_request_headers(
    headers: httpx.Headers,
) -> httpx.Headers:
    """Prepare the request headers by removing the 'Content-Type' and 'Content-Length' headers.

    Args:
        headers: The original request headers.

    Returns:
        The modified request headers.
    """
    new_headers = headers.copy()
    new_headers.pop("Content-Type", None)
    new_headers.pop("Content-Length", None)
    return new_headers


def create_response(elements: list) -> httpx.Response:
    """
    Creates a modified response object with updated content.

    Args:
        elements: The list of elements to be serialized and added to
        the response.

    Returns:
        The modified response object with updated content.
    """
    response = httpx.Response(
        status_code=200, headers={"Content-Type": "application/json"}
    )
    content = json.dumps(elements).encode()
    content_length = str(len(content))
    response.headers.update({"Content-Length": content_length})
    setattr(response, "_content", content)
    return response


ELEMENTS_FILE_HEADER = "x-unstructured-elements-file"
NDJSON_MEDIA_TYPE = "application/x-ndjson"


def combine_chunk_files_to_ndjson(chunk_paths: list[str], out_path: str) -> int:
    """Combine per-chunk split-PDF response files into one NDJSON file on disk.

    This is the whole point of the elements-file path: the split-PDF hook already streams
    each chunk response to a temp file without reading it into memory
    (`call_api_async`), but recombining them used to build a list per chunk, a flattened
    list, a `json.dumps` blob, and then the SDK re-parsed that blob -- four full copies of
    the document. Concatenating on disk keeps peak memory at ~one chunk.

    Each chunk file is either a JSON array (a server that returns
    `application/json`) or already NDJSON (a server that honors
    `application/x-ndjson`). The first non-whitespace byte tells us which, so the
    fast path -- byte-level concatenation with no parsing at all -- is taken whenever both
    ends speak NDJSON, and the JSON-array path still works against an older server.

    Returns the number of elements written.
    """
    total = 0
    with open(out_path, "w", encoding="utf-8") as out:
        for chunk_path in chunk_paths:
            with open(chunk_path, "r", encoding="utf-8") as chunk:
                first_char = ""
                while True:
                    char = chunk.read(1)
                    if char == "":
                        break
                    if not char.isspace():
                        first_char = char
                        break
                if first_char == "":
                    # Empty chunk file; nothing to append.
                    continue
                chunk.seek(0)

                if first_char == "[":
                    # JSON array: bounded by one chunk (20 pages by default), so a plain
                    # load is fine and avoids taking on a streaming-parser dependency.
                    for element in json.load(chunk):
                        out.write(json.dumps(element, ensure_ascii=False))
                        out.write("\n")
                        total += 1
                else:
                    # Already NDJSON: copy through, no parsing.
                    for line in chunk:
                        line = line.strip()
                        if line:
                            out.write(line)
                            out.write("\n")
                            total += 1
    return total


def write_chunk_body_to_temp(response: httpx.Response, dir_: Optional[str] = None) -> str:
    """Spill an in-memory chunk response body to a temp file, returning its path.

    Needed for elements-file mode when `cache_tmp_data` is OFF: there is no cached file to
    reference, so the body is written out verbatim (no parsing) to give
    `combine_chunk_files_to_ndjson` the same uniform input it gets in the cached case.
    """
    fd, path = tempfile.mkstemp(suffix=".ndjson", dir=dir_ or tempfile.gettempdir())
    with os.fdopen(fd, "wb") as f:
        f.write(response.content)
    return path


def create_elements_file_response(elements_file: str) -> httpx.Response:
    """Synthetic 200 whose payload is a path to an NDJSON file of elements.

    Mirrors the existing convention in the split hook, where a cached chunk response
    carries its temp-file path as the body. The path is also exposed as a header so the
    SDK can distinguish this from a real NDJSON body streamed from the server.
    """
    content = elements_file.encode()
    response = httpx.Response(
        status_code=200,
        headers={
            "Content-Type": NDJSON_MEDIA_TYPE,
            "Content-Length": str(len(content)),
            ELEMENTS_FILE_HEADER: elements_file,
        },
    )
    setattr(response, "_content", content)
    return response


def get_base_url(url: str | URL) -> str:
    """Extracts the base URL from the given URL.

    Args:
        url: The URL.

    Returns:
        The base URL.
    """
    parsed_url = urlparse(str(url))
    return f"{parsed_url.scheme}://{parsed_url.netloc}"