feat: saving working efficient memory saving

Author: pawel-kmiecik

src/unstructured_client/_hooks/custom/pdf_utils.py

@@ -2,7 +2,7 @@
 
 import io
 import logging
-from typing import cast
+from typing import cast, Optional
 
 from pypdf import PdfReader
 from pypdf.errors import PdfReadError
@@ -17,7 +17,21 @@
 pdf_logger = logging.getLogger("pypdf")
 pdf_logger.setLevel(logging.ERROR)
 
+def read_pdf(file: shared.Files) -> Optional[PdfReader]:
+    """Reads the given PDF file.
 
+    Args:
+        file: The PDF file to be read.
+
+    Returns:
+        The PdfReader object if the file is a PDF, None otherwise.
+    """
+
+    try:
+        content = cast(bytes, file.content)
+        return PdfReader(io.BytesIO(content), strict=False)
+    except (PdfReadError, UnicodeDecodeError):
+        return None
 
 def is_pdf(file: shared.Files) -> bool:
     """Checks if the given file is a PDF.
@@ -31,10 +45,5 @@ def is_pdf(file: shared.Files) -> bool:
         True if the file is a PDF, False otherwise.
     """
 
-    try:
-        content = cast(bytes, file.content)
-        PdfReader(io.BytesIO(content), strict=True)
-    except (PdfReadError, UnicodeDecodeError):
-        return False
+    return read_pdf(file) is not None
 
-    return True

src/unstructured_client/_hooks/custom/request_utils.py

@@ -7,6 +7,7 @@
 from typing import Tuple, Any, BinaryIO
 
 import httpx
+from httpx._multipart import MultipartStream
 from requests_toolbelt.multipart.encoder import MultipartEncoder  # type: ignore
 
 from unstructured_client._hooks.custom.common import UNSTRUCTURED_CLIENT_LOGGER_NAME
@@ -18,7 +19,8 @@
     PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
     FormData,
 )
-from unstructured_client.utils import BackoffStrategy, Retries, RetryConfig, retry_async
+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)
 
@@ -53,32 +55,153 @@ def create_pdf_request_body(
         "application/pdf",
     )))
 
-    payload_fields.append((PARTITION_FORM_STARTING_PAGE_NUMBER_KEY, str(page_number)))
+    return MultipartEncoder(fields=payload_fields)
 
-    body = MultipartEncoder(
-        fields=payload_fields
-    )
-    return body
 
+def create_pdf_request_data(
+        form_data: FormData,
+        pdf_chunk: BinaryIO,
+        filename: str,
+        page_number: int
+) -> dict[str, Any]:
+    """Creates the request body for the partition API."
 
-async def call_api_async(
-    client: httpx.AsyncClient,
+    Args:
+        form_data: The form data.
+        pdf_chunk: The pdf chunk - can be both io.BytesIO or a file object (created with open())
+        filename: The filename.
+        page_number: The page number.
+
+    """
+    payload = prepare_request_payload(form_data)
+    payload[PARTITION_FORM_STARTING_PAGE_NUMBER_KEY] = str(page_number)
+
+    # payload[PARTITION_FORM_FILES_KEY] = (
+    #     filename,
+    #     pdf_chunk,
+    #     "application/pdf",
+    # )
+    #
+    # return MultipartEncoder(fields=payload_fields)
+    return payload
+
+def prepare_pdf_chunk_request_payload(form_data: FormData) -> FormData:
+    """Prepares the request payload by removing unnecessary keys and updating the file.
+
+    Args:
+        form_data: The original form data.
+
+    Returns:
+        The updated request payload.
+    """
+    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_STARTING_PAGE_NUMBER_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"
+    return chunk_payload
+
+def create_pdf_chunk_request_data(
+        form_data: FormData,
+        page_number: int
+) -> dict[str, Any]:
+    """Creates the request body for the partition API."
+
+    Args:
+        form_data: The form data.
+        pdf_chunk: The pdf chunk - can be both io.BytesIO or a file object (created with open())
+        filename: The filename.
+        page_number: The page number.
+
+    """
+    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_STARTING_PAGE_NUMBER_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,
-    form_data: FormData,
     filename: str,
-    limiter: asyncio.Semaphore,
-) -> httpx.Response:
+) -> httpx.Request:
+    """Creates a new request object with the updated payload for the partition API.
+
+    Args:
+        form_data: The form data.
+        pdf_chunk: The pdf chunk - can be both io.BytesIO or a file object (created with open())
+        original_request: The original request.
+        filename: The filename.
+
+    Returns:
+        The updated request object.
+    """
     pdf_chunk_file, page_number = pdf_chunk
-    body = create_pdf_request_body(form_data, pdf_chunk_file, filename, page_number)
+    data = create_pdf_chunk_request_data(form_data, page_number)
     original_headers = prepare_request_headers(original_request.headers)
 
-    new_request = httpx.Request(
+    pdf_chunk_partition_params = shared.PartitionParameters(
+        files=shared.Files(
+            content=pdf_chunk_file,
+            file_name=filename,
+            content_type="application/pdf",
+        ),
+        **data,
+    )
+    serialized_body = serialize_request_body(
+        pdf_chunk_partition_params,
+        False,
+        False,
+        "multipart",
+        shared.PartitionParameters,
+    )
+    # chunk_request = client.build_request(
+    #     method="POST",
+    #     url=original_request.url or "",
+    #     headers={**original_headers, "Content-Type": "application/pdf"},
+    #     files={PARTITION_FORM_FILES_KEY: (filename, pdf_chunk_file, "application/pdf")},
+    #     data=data,
+    # )
+    pdf_chunk_request = httpx.Request(
         method="POST",
         url=original_request.url or "",
-        content=body.to_string(),
-        headers={**original_headers, "Content-Type": body.content_type},
+        # headers={**original_headers, "Content-Type": serialized_body.media_type},
+        headers={**original_headers},
+        content=serialized_body.content,
+        data=serialized_body.data,
+        files=serialized_body.files,
     )
+    return pdf_chunk_request
+
+
+
+async def call_api_async(
+    client: httpx.AsyncClient,
+    pdf_chunk_request: httpx.Request,
+    pdf_chunk_file: BinaryIO,
+    limiter: asyncio.Semaphore,
+) -> httpx.Response:
+    # pdf_chunk_file, page_number = pdf_chunk
+    # body = create_pdf_request_body(form_data, pdf_chunk_file, filename, page_number)
+    # original_headers = prepare_request_headers(original_request.headers)
+
+    # new_request = httpx.Request(
+    #     method="POST",
+    #     url=original_request.url or "",
+    #     content=body.to_string(),
+    #     headers={**original_headers, "Content-Type": body.content_type},
+    # )
 
     one_second = 1000
     one_minute = 1000 * 60
@@ -101,7 +224,7 @@ async def call_api_async(
     ]
 
     async def do_request():
-        return await client.send(new_request)
+        return await client.send(pdf_chunk_request)
 
     async with limiter:
         try:
@@ -110,6 +233,9 @@ async def do_request():
                 Retries(retry_config, retryable_codes)
             )
             return response
+        except Exception as e:
+            print(e)
+            raise e
         finally:
             if not pdf_chunk_file.closed:
                 pdf_chunk_file.close()