diff --git a/_test_unstructured_client/unit/test_ndjson_combine.py b/_test_unstructured_client/unit/test_ndjson_combine.py new file mode 100644 index 00000000..3f2b62f0 --- /dev/null +++ b/_test_unstructured_client/unit/test_ndjson_combine.py @@ -0,0 +1,213 @@ +"""Unit tests for the on-disk NDJSON recombination used by elements-file mode. + +This helper is what removes the four in-memory copies the split-PDF recombination used to +make (per-chunk list, flattened list, json.dumps blob, SDK re-parse). It must: + - handle chunk files that are JSON arrays (server returned application/json) + - handle chunk files that are already NDJSON (server honored application/x-ndjson) + - preserve element order across chunks + - round-trip payload strings byte-for-byte +""" + +import json + +import pytest + +import httpx + +from unstructured_client._hooks.custom.request_utils import ( + ELEMENTS_FILE_HEADER, + combine_chunk_files_to_ndjson, + create_elements_file_response, + write_chunk_body_to_temp, +) + + +def _elements(prefix, count): + return [ + { + "type": "Table" if i % 2 == 0 else "NarrativeText", + "text": f"{prefix}-{i}", + "metadata": {"page_number": i + 1, "image_base64": f"PAYLOAD{prefix}{i}" * 4}, + } + for i in range(count) + ] + + +def _write_json_array(path, elements): + path.write_text(json.dumps(elements), encoding="utf-8") + + +def _write_ndjson(path, elements): + with path.open("w", encoding="utf-8") as f: + for element in elements: + f.write(json.dumps(element)) + f.write("\n") + + +def _read_ndjson(path): + with open(path, encoding="utf-8") as f: + return [json.loads(line) for line in f if line.strip()] + + +def test_combines_json_array_chunks_in_order(tmp_path): + chunk_a = tmp_path / "a.json" + chunk_b = tmp_path / "b.json" + elements_a = _elements("a", 3) + elements_b = _elements("b", 2) + _write_json_array(chunk_a, elements_a) + _write_json_array(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 5 + assert _read_ndjson(out) == elements_a + elements_b + + +def test_combines_ndjson_chunks_without_parsing(tmp_path): + """The zero-parse path: both ends speak NDJSON.""" + chunk_a = tmp_path / "a.ndjson" + chunk_b = tmp_path / "b.ndjson" + elements_a = _elements("a", 4) + elements_b = _elements("b", 1) + _write_ndjson(chunk_a, elements_a) + _write_ndjson(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 5 + assert _read_ndjson(out) == elements_a + elements_b + + +def test_mixed_chunk_formats(tmp_path): + """A server upgraded mid-flight, or a retry served by an older pod.""" + chunk_a = tmp_path / "a.json" + chunk_b = tmp_path / "b.ndjson" + elements_a = _elements("a", 2) + elements_b = _elements("b", 2) + _write_json_array(chunk_a, elements_a) + _write_ndjson(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 4 + assert _read_ndjson(out) == elements_a + elements_b + + +@pytest.mark.parametrize("body", ["", " ", "\n\n"]) +def test_empty_chunk_files_are_skipped(tmp_path, body): + chunk_a = tmp_path / "a.json" + chunk_empty = tmp_path / "empty.json" + elements_a = _elements("a", 2) + _write_json_array(chunk_a, elements_a) + chunk_empty.write_text(body, encoding="utf-8") + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_empty)], str(out)) + + assert written == 2 + assert _read_ndjson(out) == elements_a + + +def test_empty_array_chunk_contributes_nothing(tmp_path): + """A chunk that legitimately produced no elements (e.g. blank pages).""" + chunk_a = tmp_path / "a.json" + chunk_b = tmp_path / "b.json" + _write_json_array(chunk_a, []) + elements_b = _elements("b", 3) + _write_json_array(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 3 + assert _read_ndjson(out) == elements_b + + +def test_payload_round_trips_exactly(tmp_path): + """Base64 payloads are the reason this path exists; they must survive unchanged.""" + payload = "A" * 100_000 + element = {"type": "Table", "text": "t", "metadata": {"image_base64": payload}} + chunk = tmp_path / "a.json" + _write_json_array(chunk, [element]) + + out = tmp_path / "combined.ndjson" + combine_chunk_files_to_ndjson([str(chunk)], str(out)) + + result = _read_ndjson(out) + assert result[0]["metadata"]["image_base64"] == payload + + +def test_non_ascii_is_preserved(tmp_path): + element = {"type": "NarrativeText", "text": "日本語 café ✓", "metadata": {}} + chunk = tmp_path / "a.json" + _write_json_array(chunk, [element]) + + out = tmp_path / "combined.ndjson" + combine_chunk_files_to_ndjson([str(chunk)], str(out)) + + assert _read_ndjson(out)[0]["text"] == "日本語 café ✓" + + +def test_write_chunk_body_to_temp_roundtrips(tmp_path): + """The cache_tmp_data=OFF path: an in-memory NDJSON body must spill verbatim. + + Regression guard. `ndjson_mode` used to also require cache_tmp_data, so with caching off + the server returned NDJSON while the hook took the JSON path and `res.json()` raised on a + body this client had itself requested. + """ + elements = _elements("x", 3) + body = "".join(json.dumps(e) + "\n" for e in elements).encode() + response = httpx.Response(status_code=200, content=body) + + path = write_chunk_body_to_temp(response, str(tmp_path)) + assert _read_ndjson(path) == elements + + +def test_combine_accepts_bodies_spilled_without_caching(tmp_path): + """End-to-end of the uncached path: spill two bodies, then combine them.""" + a, b = _elements("a", 2), _elements("b", 3) + ra = httpx.Response(200, content="".join(json.dumps(e) + "\n" for e in a).encode()) + rb = httpx.Response(200, content="".join(json.dumps(e) + "\n" for e in b).encode()) + paths = [write_chunk_body_to_temp(r, str(tmp_path)) for r in (ra, rb)] + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson(paths, str(out)) + + assert written == 5 + assert _read_ndjson(out) == a + b + + +def test_spilled_body_is_released_from_the_response(tmp_path): + """After spilling, the response must no longer hold the body. + + Regression guard for the real failure mode: every chunk response is retained in + `api_successful_responses` for failure bookkeeping, so spilling to disk without + releasing `_content` still accumulates the whole document in memory (125 chunks x + 32 MB = 4 GB observed on a 2500-page split). + """ + elements = _elements("a", 3) + body = "".join(json.dumps(e) + "\n" for e in elements).encode() + response = httpx.Response(status_code=200, content=body) + assert len(response.content) == len(body) + + path = write_chunk_body_to_temp(response, str(tmp_path)) + response._content = path.encode() + + # The body is on disk, and the response now costs a path rather than a payload. + assert _read_ndjson(path) == elements + assert response.text == path + assert len(response.content) < 512 + + +def test_elements_file_response_carries_path_in_header_and_body(tmp_path): + path = str(tmp_path / "combined.ndjson") + response = create_elements_file_response(path) + + assert response.status_code == 200 + assert response.headers[ELEMENTS_FILE_HEADER] == path + assert response.headers["Content-Type"] == "application/x-ndjson" + # Body-as-path mirrors the existing cached-chunk convention in the split hook. + assert response.text == path diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index bfc9cb0f..1fb24e79 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -4,6 +4,8 @@ import io import json import logging +import os +import tempfile from typing import Tuple, Any, BinaryIO, Optional from urllib.parse import urlparse @@ -277,6 +279,95 @@ def create_response(elements: list) -> httpx.Response: 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. diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index 912da332..df600b63 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -524,6 +524,11 @@ def __init__(self) -> None: self.allow_failed: dict[str, bool] = {} self.cache_tmp_data_feature: dict[str, bool] = {} self.cache_tmp_data_dir: dict[str, str] = {} + # NDJSON elements-file mode: when the caller asked for application/x-ndjson, the + # per-chunk temp files are concatenated on disk instead of being parsed and + # re-serialized, and the combined path is handed back via the response header. + self.ndjson_mode: dict[str, bool] = {} + self.ndjson_output_path: dict[str, str] = {} @staticmethod def _get_operation_id_from_request(request: Optional[httpx.Request]) -> Optional[str]: @@ -791,6 +796,14 @@ def _before_request_unlocked( self.cache_tmp_data_feature[operation_id] = cache_tmp_data_feature self.cache_tmp_data_dir[operation_id] = cache_tmp_data_dir self.concurrency_level[operation_id] = concurrency_level + # Depends ONLY on what the caller asked for. It must NOT also depend on + # cache_tmp_data: the Accept header is chosen by the caller while cache_tmp_data is + # a separate setting, so gating on both lets them disagree -- and when they do, the + # server returns NDJSON while the hook takes the JSON path and `res.json()` blows up + # on a body this very client requested. Both caching modes are handled below. + self.ndjson_mode[operation_id] = ( + request_utils.NDJSON_MEDIA_TYPE in request.headers.get("Accept", "") + ) timeout_seconds = _get_request_timeout_seconds(request) if timeout_seconds is None and hook_ctx.config.timeout_ms is not None: @@ -1382,6 +1395,8 @@ def _elements_from_task_responses( failed_responses: list[tuple[int, httpx.Response]] = [] transport_failure_count = 0 elements = [] + ndjson_mode = self.ndjson_mode.get(operation_id, False) + chunk_paths: list[str] = [] for response_number, res in task_responses: if res.status_code == 200: logger.debug( @@ -1390,7 +1405,28 @@ def _elements_from_task_responses( response_number, ) successful_responses.append(res) - if self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): + if ndjson_mode: + # Nothing is parsed here either way -- that is what keeps peak memory at + # ~one chunk during recombination. + if self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): + # Cached: the body was already streamed to a temp file and `res.text` + # holds that path. + chunk_paths.append(res.text) + else: + # Not cached: the body is in memory. Spill it verbatim (no parsing) + # and then RELEASE it -- every response object is retained in + # `successful_responses` for failure bookkeeping, so without this the + # spilled bodies stay resident and the whole document accumulates + # anyway (125 chunks x 32 MB = 4 GB observed). Overwriting `_content` + # with the path mirrors what the cached branch already does at the + # point of caching, so downstream `res.text` means the same thing in + # both branches. + spilled = request_utils.write_chunk_body_to_temp( + res, self.cache_tmp_data_dir.get(operation_id) + ) + res._content = spilled.encode() # pylint: disable=protected-access + chunk_paths.append(spilled) + elif self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): elements.append(load_elements_from_response(res)) else: elements.append(res.json()) @@ -1433,6 +1469,27 @@ def _elements_from_task_responses( total_chunks=len(task_responses), response=response, ) + if ndjson_mode: + # Combine on disk and stash the path; `_build_after_success_response` turns it + # into the response. Deliberately written to the cache *parent* dir, not the + # operation's TemporaryDirectory: that directory is cleaned up by + # `_clear_operation` -> `_finalize_operation_resources` immediately after + # after_success returns, which would delete this file before the caller could + # read it. The combined file therefore outlives the operation and the CALLER + # owns deleting it (see `PartitionResponse.elements_file`). + temp_dir_path = self.cache_tmp_data_dir.get(operation_id) or tempfile.gettempdir() + out_path = f"{temp_dir_path}/{uuid.uuid4()}.ndjson" + written = request_utils.combine_chunk_files_to_ndjson(chunk_paths, out_path) + self.ndjson_output_path[operation_id] = out_path + logger.info( + "split_pdf event=ndjson_combined operation_id=%s chunk_count=%d element_count=%d out_file=%s", + operation_id, + len(chunk_paths), + written, + Path(out_path).name, + ) + return [] + flattened_elements = [element for sublist in elements for element in sublist] return flattened_elements @@ -1464,6 +1521,19 @@ def _build_after_success_response( ) return self.api_failed_responses[operation_id][0] + # Elements-file mode: hand back the combined NDJSON path instead of a body. Checked + # before the `elements is None` guard because `elements` is intentionally empty + # here -- nothing was parsed. + if self.ndjson_mode.get(operation_id, False): + ndjson_path = self.ndjson_output_path.get(operation_id) + if ndjson_path is None: + logger.warning( + "split_pdf event=ndjson_missing_output operation_id=%s falling_back=true", + operation_id, + ) + return response + return request_utils.create_elements_file_response(ndjson_path) + if elements is None: return response @@ -1586,6 +1656,8 @@ def _clear_operation(self, operation_id: str) -> None: self.allow_failed.pop(operation_id, None) self.cache_tmp_data_feature.pop(operation_id, None) self.cache_tmp_data_dir.pop(operation_id, None) + self.ndjson_mode.pop(operation_id, None) + self.ndjson_output_path.pop(operation_id, None) self.pending_operation_ids.pop(operation_id, None) future = self.operation_futures.pop(operation_id, None) loop_holder = self.operation_loops.pop(operation_id, None) diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index a6ec3099..d8cc4b1a 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -14,6 +14,53 @@ class PartitionAcceptEnum(str, Enum): APPLICATION_JSON = "application/json" TEXT_CSV = "text/csv" + APPLICATION_X_NDJSON = "application/x-ndjson" + r"""Elements as NDJSON, one per line. The response is written to a temp file and + returned as `PartitionResponse.elements_file` rather than parsed into `elements`, so + the caller can stream a large document instead of holding it in memory.""" + + +def _ndjson_elements_file(http_res) -> str: + """Resolve an NDJSON response to a path on disk, without parsing the elements. + + Two shapes arrive here. When the split-PDF hook ran, it has already combined the + per-chunk temp files into one NDJSON file and passes the path through + `ELEMENTS_FILE_HEADER` -- nothing to do but read the header, so the document is never + materialized. Otherwise this is a real body from the server, which is written to a + temp file in bounded chunks. + """ + import tempfile + + from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_HEADER + + existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + if existing_path: + return existing_path + + with tempfile.NamedTemporaryFile( + mode="wb", prefix="unst_elements_", suffix=".ndjson", delete=False + ) as out: + for byte_chunk in http_res.iter_bytes(): + out.write(byte_chunk) + return out.name + + +async def _ndjson_elements_file_async(http_res) -> str: + """Async counterpart of `_ndjson_elements_file`.""" + import tempfile + + from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_HEADER + + existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + if existing_path: + return existing_path + + with tempfile.NamedTemporaryFile( + mode="wb", prefix="unst_elements_", suffix=".ndjson", delete=False + ) as out: + async for byte_chunk in http_res.aiter_bytes(): + out.write(byte_chunk) + return out.name class General(BaseSDK): @@ -128,6 +175,13 @@ def partition( content_type=http_res.headers.get("Content-Type") or "", raw_response=http_res, ) + if utils.match_response(http_res, "200", "application/x-ndjson"): + return operations.PartitionResponse( + elements_file=_ndjson_elements_file(http_res), + status_code=http_res.status_code, + content_type=http_res.headers.get("Content-Type") or "", + raw_response=http_res, + ) if utils.match_response(http_res, "422", "application/json"): response_data = unmarshal_json_response( errors.HTTPValidationErrorData, http_res @@ -253,6 +307,13 @@ async def partition_async( content_type=http_res.headers.get("Content-Type") or "", raw_response=http_res, ) + if utils.match_response(http_res, "200", "application/x-ndjson"): + return operations.PartitionResponse( + elements_file=await _ndjson_elements_file_async(http_res), + status_code=http_res.status_code, + content_type=http_res.headers.get("Content-Type") or "", + raw_response=http_res, + ) if utils.match_response(http_res, "422", "application/json"): response_data = unmarshal_json_response( errors.HTTPValidationErrorData, http_res diff --git a/src/unstructured_client/models/operations/partition.py b/src/unstructured_client/models/operations/partition.py index de4664fb..e447f641 100644 --- a/src/unstructured_client/models/operations/partition.py +++ b/src/unstructured_client/models/operations/partition.py @@ -78,6 +78,10 @@ class PartitionResponseTypedDict(TypedDict): r"""Successful Response""" elements: NotRequired[List[Dict[str, Any]]] r"""Successful Response""" + elements_file: NotRequired[str] + r"""Path to an NDJSON file of elements, one per line. Set instead of `elements` when + `application/x-ndjson` was requested, so a large document never has to be held in + memory as a parsed list.""" class PartitionResponse(BaseModel): @@ -95,3 +99,8 @@ class PartitionResponse(BaseModel): elements: Optional[List[Dict[str, Any]]] = None r"""Successful Response""" + + elements_file: Optional[str] = None + r"""Path to an NDJSON file of elements, one per line. Set instead of `elements` when + `application/x-ndjson` was requested, so a large document never has to be held in + memory as a parsed list. The caller owns the file and should delete it when done."""