From fce575a512991e0f8620a2ef8d33e20e8046ca66 Mon Sep 17 00:00:00 2001 From: Yao You Date: Sat, 1 Aug 2026 13:04:24 -0500 Subject: [PATCH 1/5] feat: NDJSON elements-file mode for partition On the split-PDF path the SDK rebuilt the whole document in memory to return it: a list per chunk, a flattened list, a json.dumps blob in create_response, and the SDK's re-parse of that blob -- four copies live at once, with the serialization step dominating peak usage on large documents. Passing accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON now returns PartitionResponse.elements_file -- a path to an NDJSON file, one element per line -- instead of PartitionResponse.elements. The per-chunk temp files are concatenated on disk and never parsed, so peak memory is roughly one chunk rather than the whole document. Chunk files are sniffed for their first non-whitespace character, so a server returning application/json still works; NDJSON chunks are copied through untouched. Requesting application/json remains the default and is unchanged. The caller owns the returned file and must delete it. It is deliberately written outside the operation's TemporaryDirectory, which _clear_operation removes as soon as after_success returns. ndjson_mode depends only on the Accept header, never on split_pdf_cache_tmp_data. Those are set by different parties, so gating on both let them disagree: the server would return NDJSON while the hook took the JSON path and res.json() raised on a body this client had itself requested. Both caching modes are handled -- a cached chunk contributes its existing temp-file path, an uncached one spills its body verbatim and then releases it, since every response is retained in api_successful_responses and leaving _content set would keep the document resident regardless. general.py and models/operations/partition.py are both .genignore'd: elements_file is client-side only and can never come from the OpenAPI spec, so a regeneration would silently drop it. test_regeneration_guards.py fails if either entry is lost. Co-Authored-By: Claude Opus 5 (1M context) --- .genignore | 7 + README.md | 27 ++ .../unit/test_ndjson_elements_file.py | 322 ++++++++++++++++++ .../unit/test_regeneration_guards.py | 21 ++ docs/models/operations/partitionresponse.md | 3 +- .../_hooks/custom/request_utils.py | 111 +++++- .../_hooks/custom/split_pdf_hook.py | 123 ++++++- src/unstructured_client/general.py | 59 ++++ .../models/operations/partition.py | 9 + 9 files changed, 679 insertions(+), 3 deletions(-) create mode 100644 _test_unstructured_client/unit/test_ndjson_elements_file.py diff --git a/.genignore b/.genignore index ea1fba41..b68a3c44 100644 --- a/.genignore +++ b/.genignore @@ -23,3 +23,10 @@ src/unstructured_client/general.py # Custom min_attempts / absolute_max_elapsed_time_ms fields on BackoffStrategy. # Push upstream to Speakeasy templates to remove this entry. src/unstructured_client/utils/retries.py + +# Custom elements_file field on PartitionResponse, for the NDJSON elements-file mode. +# The field is client-side only - the server never returns it - so it cannot come from +# the OpenAPI spec, and regenerating would drop it. If /general/v0/general gains a new +# response field, follow the same procedure as general.py above. +# See test_regeneration_guards.py::test_partition_response_keeps_elements_file. +src/unstructured_client/models/operations/partition.py diff --git a/README.md b/README.md index 17bfc6f4..b329ca40 100644 --- a/README.md +++ b/README.md @@ -427,6 +427,33 @@ req = operations.PartitionRequest( ) ``` +### Streaming elements to a file instead of memory + +For very large documents, the parsed element list can dominate the client's memory: the split-PDF path holds a list per chunk, a flattened list, a serialized blob and the SDK's re-parse of that blob. Request `application/x-ndjson` to skip all of it. The chunk responses are concatenated on disk and you get back a path in `elements_file` instead of a list in `elements`, which keeps peak memory at roughly one chunk. + +**You own the returned file and are responsible for deleting it.** + +Example: +```python +import json +import os + +from unstructured_client.general import PartitionAcceptEnum + +res = client.general.partition( + request=req, + accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON, +) + +try: + with open(res.elements_file, encoding="utf-8") as f: + for line in f: + element = json.loads(line) + ... +finally: + os.unlink(res.elements_file) +``` + ## File uploads diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py new file mode 100644 index 00000000..075b6980 --- /dev/null +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -0,0 +1,322 @@ +"""Unit tests for NDJSON elements-file mode. + +The on-disk recombination replaces the four in-memory copies the split-PDF path 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 + - leave no temp files behind other than the combined file the caller owns +""" + +import json +import os +from pathlib import Path + +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, +) +from unstructured_client._hooks.custom.split_pdf_hook import SplitPdfHook + + +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: 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, defeating the point of spilling. + """ + 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 + + +# --- hook-level temp-file lifecycle ------------------------------------------------ + + +def _ndjson_response(elements): + body = "".join(json.dumps(e) + "\n" for e in elements).encode() + return httpx.Response(status_code=200, content=body) + + +def _hook_in_ndjson_mode(operation_id, tmp_path): + """A hook set up as `before_request` would leave it for an uncached NDJSON run.""" + hook = SplitPdfHook() + hook.ndjson_mode[operation_id] = True + hook.cache_tmp_data_feature[operation_id] = False + hook.cache_tmp_data_dir[operation_id] = str(tmp_path) + hook.allow_failed[operation_id] = False + return hook + + +def _ndjson_files_in(directory): + return sorted(p.name for p in Path(directory).glob("*.ndjson")) + + +def test_spilled_chunk_files_are_deleted_after_combining(tmp_path): + """Regression guard: spilled chunks used to be left in the temp dir forever. + + `cache_tmp_data` defaults to off, so this is the default path. One file per chunk + accumulating for the lifetime of the host is a disk leak, not a memory one. + """ + operation_id = "op-cleanup" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + responses = [(0, _ndjson_response(_elements("a", 2))), (1, _ndjson_response(_elements("b", 3)))] + + hook._elements_from_task_responses(operation_id, responses, started_at=0.0) + + combined = hook.ndjson_output_path[operation_id] + # Exactly one file survives: the combined output the caller owns. + assert _ndjson_files_in(tmp_path) == [Path(combined).name] + assert len(_read_ndjson(combined)) == 5 + + +def test_spilled_chunks_land_in_the_operation_tempdir_when_one_exists(tmp_path): + """Spilling into the operation's tempdir means cleanup happens even if we miss it.""" + operation_id = "op-tempdir" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + operation_dir = tmp_path / "unstructured_client_op" + operation_dir.mkdir() + + class _FakeTempDir: + name = str(operation_dir) + + hook.tempdirs[operation_id] = _FakeTempDir() # type: ignore[assignment] + spill_dir = hook._operation_tempdir_path(operation_id) + + assert spill_dir == str(operation_dir) + + +def test_combined_file_is_returned_on_success(tmp_path): + operation_id = "op-success" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + elements = _elements("a", 4) + + hook._elements_from_task_responses( + operation_id, [(0, _ndjson_response(elements))], started_at=0.0 + ) + response = hook._build_after_success_response(operation_id, httpx.Response(200), []) + + combined = response.headers[ELEMENTS_FILE_HEADER] + assert os.path.exists(combined) + assert _read_ndjson(combined) == elements + + +def test_combined_file_is_discarded_when_a_failure_response_is_returned(tmp_path): + """Regression guard: on the failure path nothing downstream learns the path. + + `_build_after_success_response` returns the failed chunk response instead, so the + combined file would be leaked for the lifetime of the host. + """ + operation_id = "op-strict-failure" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + responses = [ + (0, _ndjson_response(_elements("a", 2))), + (1, httpx.Response(status_code=500, content=b"boom")), + ] + + hook._elements_from_task_responses(operation_id, responses, started_at=0.0) + combined = hook.ndjson_output_path[operation_id] + assert os.path.exists(combined) + + response = hook._build_after_success_response(operation_id, httpx.Response(200), []) + + assert response.status_code == 500 + assert not os.path.exists(combined) + assert _ndjson_files_in(tmp_path) == [] + + +def test_no_combined_file_is_created_when_every_chunk_failed(tmp_path): + operation_id = "op-all-failed" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + responses = [(0, httpx.Response(status_code=500, content=b"boom"))] + + hook._elements_from_task_responses(operation_id, responses, started_at=0.0) + + assert operation_id not in hook.ndjson_output_path + assert _ndjson_files_in(tmp_path) == [] diff --git a/_test_unstructured_client/unit/test_regeneration_guards.py b/_test_unstructured_client/unit/test_regeneration_guards.py index 1286057b..48f04a57 100644 --- a/_test_unstructured_client/unit/test_regeneration_guards.py +++ b/_test_unstructured_client/unit/test_regeneration_guards.py @@ -81,6 +81,27 @@ def test_ci_installs_with_locked_uv_sync(): assert "run: make install" in workflow +def test_partition_response_keeps_elements_file(): + """`elements_file` is client-side only, so no spec change can restore it after a regen. + + Both the model and the enum value that selects it live in generated files; the + .genignore entries are the only thing keeping them. + """ + from unstructured_client.general import PartitionAcceptEnum + from unstructured_client.models import operations + + assert "elements_file" in operations.PartitionResponse.model_fields + assert "elements_file" in operations.PartitionResponseTypedDict.__annotations__ + assert PartitionAcceptEnum.APPLICATION_X_NDJSON.value == "application/x-ndjson" + + genignore = (REPO_ROOT / ".genignore").read_text() + for path in ( + "src/unstructured_client/general.py", + "src/unstructured_client/models/operations/partition.py", + ): + assert path in genignore, f"{path} carries custom code and must stay in .genignore" + + def test_body_create_job_input_files_are_serialized_as_multipart_files(): request = shared.BodyCreateJob( request_data="{}", diff --git a/docs/models/operations/partitionresponse.md b/docs/models/operations/partitionresponse.md index d19430ae..b5b19f50 100644 --- a/docs/models/operations/partitionresponse.md +++ b/docs/models/operations/partitionresponse.md @@ -9,4 +9,5 @@ | `status_code` | *int* | :heavy_check_mark: | HTTP response status code for this operation | | `raw_response` | [httpx.Response](https://www.python-httpx.org/api/#response) | :heavy_check_mark: | Raw HTTP response; suitable for custom response parsing | | `csv_elements` | *Optional[str]* | :heavy_minus_sign: | Successful Response | -| `elements` | List[Dict[str, *Any*]] | :heavy_minus_sign: | Successful Response | \ No newline at end of file +| `elements` | List[Dict[str, *Any*]] | :heavy_minus_sign: | Successful Response | +| `elements_file` | *Optional[str]* | :heavy_minus_sign: | Path to an NDJSON file of elements, one per line. Set instead of `elements` when `application/x-ndjson` was requested. The caller owns the file and should delete it when done. | \ No newline at end of file diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index bfc9cb0f..be8fe31f 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -4,7 +4,9 @@ import io import json import logging -from typing import Tuple, Any, BinaryIO, Optional +import os +import tempfile +from typing import Tuple, Any, BinaryIO, Optional, TextIO from urllib.parse import urlparse import httpx @@ -277,6 +279,113 @@ def create_response(elements: list) -> httpx.Response: return response +ELEMENTS_FILE_HEADER = "x-unstructured-elements-file" +NDJSON_MEDIA_TYPE = "application/x-ndjson" + +_SNIFF_BLOCK_SIZE = 64 + + +def _first_non_space_char(stream: TextIO) -> str: + """Return the first non-whitespace character in `stream`, or "" if there is none.""" + while True: + block = stream.read(_SNIFF_BLOCK_SIZE) + if not block: + return "" + stripped = block.lstrip() + if stripped: + return stripped[0] + + +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. + + Recombining chunks by parsing them builds four full copies of the document (a list + per chunk, a flattened list, a `json.dumps` blob, and the SDK's re-parse of that + blob). Concatenating on disk keeps peak memory at roughly one chunk instead. + + A chunk file is either a JSON array (server returned `application/json`) or already + NDJSON (server honored `application/x-ndjson`); the first non-whitespace character + says which. NDJSON chunks are copied through without parsing. + + Args: + chunk_paths: Per-chunk response files, in element order. + out_path: File to write the combined NDJSON to. + + 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 = _first_non_space_char(chunk) + if not first_char: + continue + chunk.seek(0) + + if first_char == "[": + # A chunk is bounded (20 pages by default), so a plain load 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: + 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. + + Used by elements-file mode when `cache_tmp_data` is off and there is no cached file + to reference. The body is written verbatim so that `combine_chunk_files_to_ndjson` + gets the same input it does in the cached case. + + Args: + response: The chunk response whose body should be spilled. + dir_: Directory to create the file in. Defaults to the system temp directory. + + Returns: + The path to the spilled file. The caller owns deleting it. + """ + 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: + """Create a synthetic 200 response whose payload is a path to an NDJSON file. + + Mirrors the split hook's existing convention of a cached chunk response carrying its + temp-file path as the body. The path is also set as a header so the SDK can tell this + apart from a real NDJSON body streamed from the server. + + Args: + elements_file: Path to the combined NDJSON file of elements. + + Returns: + The synthetic response. + """ + 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..5d64aa49 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -487,6 +487,23 @@ def load_elements_from_response(response: httpx.Response) -> list[dict]: return json.load(file) +def _unlink_quietly(paths: Iterable[str]) -> None: + """Delete temp files, logging rather than raising if one cannot be removed. + + Cleanup runs on the success path, so a failure to unlink must never turn a completed + partition into an error. + """ + for path in paths: + try: + os.unlink(path) + except OSError as exc: + logger.warning( + "split_pdf event=temp_file_cleanup_failed file=%s error=%s", + Path(path).name, + exc, + ) + + class SplitPdfHook(SDKInitHook, BeforeRequestHook, AfterSuccessHook, AfterErrorHook): """ A hook class that splits a PDF file into multiple pages and sends each page as @@ -524,6 +541,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 asks 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 +813,15 @@ 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, never on cache_tmp_data. The Accept + # header is chosen by the caller while cache_tmp_data is a separate split-PDF + # setting, so gating on both lets them disagree: the server would return NDJSON + # while the hook took the JSON path and `res.json()` raised on a body this client + # had itself requested. Both caching modes are handled in + # `_elements_from_task_responses`. + 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 +1413,11 @@ 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] = [] + # Subset of `chunk_paths` this method created itself, and so must clean up. The + # rest belong to the operation's tempdir and are removed with it. + spilled_chunk_paths: list[str] = [] for response_number, res in task_responses: if res.status_code == 200: logger.debug( @@ -1390,7 +1426,27 @@ 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: + # Neither branch parses the body; that is what keeps peak memory at + # roughly one chunk during recombination. + if self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): + # The body was already streamed to a temp file owned by the + # operation's tempdir, and `res.text` holds that path. + chunk_paths.append(res.text) + else: + # The body is in memory, so spill it verbatim and then release it: + # every response stays in `successful_responses` for failure + # bookkeeping, so leaving `_content` set would keep the whole + # document resident anyway. Overwriting it with the path matches + # what the cached branch does, so `res.text` means the same thing + # in both. + spilled = request_utils.write_chunk_body_to_temp( + res, self._operation_tempdir_path(operation_id) + ) + res._content = spilled.encode() # pylint: disable=protected-access + spilled_chunk_paths.append(spilled) + 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,9 +1489,57 @@ def _elements_from_task_responses( total_chunks=len(task_responses), response=response, ) + if ndjson_mode: + try: + if chunk_paths: + self._combine_chunks_to_ndjson(operation_id, chunk_paths) + finally: + # These are ours; the cached chunk files belong to the operation tempdir. + _unlink_quietly(spilled_chunk_paths) + return [] + flattened_elements = [element for sublist in elements for element in sublist] return flattened_elements + def _combine_chunks_to_ndjson(self, operation_id: str, chunk_paths: list[str]) -> None: + """Concatenate the per-chunk files into one NDJSON file and record its path. + + `_build_after_success_response` turns the recorded path into the response. The + combined file goes in the cache *parent* directory rather than the operation's + TemporaryDirectory, because `_clear_operation` cleans that directory up as soon as + after_success returns, which would delete the 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, + ) + + def _operation_tempdir_path(self, operation_id: str) -> Optional[str]: + """Directory to spill chunk bodies into, preferring the operation's own tempdir.""" + tempdir = self.tempdirs.get(operation_id) + if tempdir is not None: + return tempdir.name + return self.cache_tmp_data_dir.get(operation_id) + + def _discard_ndjson_output(self, operation_id: str) -> None: + """Delete the combined NDJSON file when it will not be handed to the caller. + + Once we return a failure response instead, nothing downstream learns the path, so + without this the combined file is leaked for the lifetime of the host. + """ + out_path = self.ndjson_output_path.pop(operation_id, None) + if out_path is not None: + _unlink_quietly([out_path]) + def _build_after_success_response( self, operation_id: str, @@ -1451,6 +1555,7 @@ def _build_after_success_response( "split_pdf event=top_level_failure operation_id=%s mode=strict failed_response_selected=true", operation_id, ) + self._discard_ndjson_output(operation_id) return self.api_failed_responses[operation_id][0] if ( @@ -1462,8 +1567,22 @@ def _build_after_success_response( "split_pdf event=top_level_failure operation_id=%s mode=allow_failed reason=no_successful_chunks", operation_id, ) + self._discard_ndjson_output(operation_id) 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 +1705,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..3ecf6f67 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -2,18 +2,63 @@ from .basesdk import BaseSDK from enum import Enum +import httpx +import tempfile from typing import Any, Dict, List, Mapping, Optional, Union, cast from unstructured_client import utils from unstructured_client._hooks import HookContext from unstructured_client.models import errors, operations, shared from unstructured_client.types import BaseModel, OptionalNullable, UNSET from unstructured_client._hooks.custom.clean_server_url_hook import clean_server_url +from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_HEADER from unstructured_client.utils.unmarshal_json_response import unmarshal_json_response 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 + a large document never has to be held in memory.""" + + +# NDJSON elements-file support. `partition.py` and this module are both .genignore'd so +# these edits survive regeneration; see the notes in .genignore. +def _new_elements_file(): + return tempfile.NamedTemporaryFile( # pylint: disable=consider-using-with + mode="wb", prefix="unst_elements_", suffix=".ndjson", delete=False + ) + + +def _ndjson_elements_file(http_res: httpx.Response) -> str: + """Resolve an NDJSON response to a path on disk, without parsing the elements. + + 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`, so there is nothing to + do but read the header. Otherwise this is a real body from the server, which is + streamed to a temp file. + """ + existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + if existing_path: + return existing_path + + with _new_elements_file() 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: httpx.Response) -> str: + """Async counterpart of `_ndjson_elements_file`.""" + existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + if existing_path: + return existing_path + + with _new_elements_file() as out: + async for byte_chunk in http_res.aiter_bytes(): + out.write(byte_chunk) + return out.name class General(BaseSDK): @@ -128,6 +173,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 +305,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.""" From 7210084c52a5447bd760e1befaf7a43e758d1c70 Mon Sep 17 00:00:00 2001 From: Yao You Date: Sat, 1 Aug 2026 13:59:34 -0500 Subject: [PATCH 2/5] fix: harden the NDJSON elements-file path Two defects found in review of the elements-file mode. The elements-file marker was an `x-unstructured-elements-file` response header, and any response carrying it was trusted as an SDK-created path. Headers come off the wire, so a server could name an arbitrary local file -- and callers are documented to open `elements_file` and then delete it, making this an arbitrary-file delete rather than just a disclosure. The marker is now an httpx response extension, which is populated by the transport and cannot be set remotely; a real server body is always copied to a file this client creates. Recombination also wrote straight to its final UUID path while recording that path only on success, so a malformed chunk left a partial file behind under a name nothing owned -- the combined file is deliberately outside the operation's TemporaryDirectory, so nothing else cleaned it up. It now writes to a staging file renamed into place atomically, and unlinks it on any exception. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_elements_file.py | 54 +++++++++++++++++-- .../_hooks/custom/request_utils.py | 14 +++-- .../_hooks/custom/split_pdf_hook.py | 15 +++++- src/unstructured_client/general.py | 15 +++--- 4 files changed, 83 insertions(+), 15 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index 075b6980..eac9fb95 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -18,12 +18,13 @@ import httpx from unstructured_client._hooks.custom.request_utils import ( - ELEMENTS_FILE_HEADER, + ELEMENTS_FILE_EXTENSION_KEY, combine_chunk_files_to_ndjson, create_elements_file_response, write_chunk_body_to_temp, ) from unstructured_client._hooks.custom.split_pdf_hook import SplitPdfHook +from unstructured_client.general import _ndjson_elements_file def _elements(prefix, count): @@ -205,17 +206,43 @@ def test_spilled_body_is_released_from_the_response(tmp_path): assert len(response.content) < 512 -def test_elements_file_response_carries_path_in_header_and_body(tmp_path): +def test_elements_file_response_carries_path_in_extension_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.extensions[ELEMENTS_FILE_EXTENSION_KEY] == 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 +def test_server_cannot_name_a_local_file_via_a_response_header(tmp_path): + """The elements-file marker must not be reachable from the wire. + + Callers are documented to open `elements_file` and then delete it, so trusting a + server-supplied path would hand a hostile server an arbitrary local file to destroy. + A header must be ignored and the body copied to a file this client created. + """ + victim = tmp_path / "victim" + victim.write_text("do not touch", encoding="utf-8") + response = httpx.Response( + 200, + headers={ + "content-type": "application/x-ndjson", + "x-unstructured-elements-file": str(victim), + }, + content=b'{"safe": true}\n', + ) + + resolved = _ndjson_elements_file(response) + + assert resolved != str(victim) + assert victim.read_text(encoding="utf-8") == "do not touch" + assert _read_ndjson(resolved) == [{"safe": True}] + os.unlink(resolved) + + # --- hook-level temp-file lifecycle ------------------------------------------------ @@ -282,7 +309,7 @@ def test_combined_file_is_returned_on_success(tmp_path): ) response = hook._build_after_success_response(operation_id, httpx.Response(200), []) - combined = response.headers[ELEMENTS_FILE_HEADER] + combined = response.extensions[ELEMENTS_FILE_EXTENSION_KEY] assert os.path.exists(combined) assert _read_ndjson(combined) == elements @@ -311,6 +338,25 @@ def test_combined_file_is_discarded_when_a_failure_response_is_returned(tmp_path assert _ndjson_files_in(tmp_path) == [] +def test_malformed_chunk_leaves_no_partial_output_behind(tmp_path): + """Regression guard: recombination that raises must not orphan a partial file. + + The combined file is the one artifact here that no temp directory owns, so a partial + one would outlive the failed operation. + """ + operation_id = "op-malformed" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + + with pytest.raises(json.JSONDecodeError): + hook._elements_from_task_responses( + operation_id, [(0, httpx.Response(200, content=b"[not-json"))], started_at=0.0 + ) + + assert operation_id not in hook.ndjson_output_path + assert _ndjson_files_in(tmp_path) == [] + assert list(Path(tmp_path).glob("*.partial")) == [] + + def test_no_combined_file_is_created_when_every_chunk_failed(tmp_path): operation_id = "op-all-failed" hook = _hook_in_ndjson_mode(operation_id, tmp_path) diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index be8fe31f..f48d04bf 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -279,7 +279,12 @@ def create_response(elements: list) -> httpx.Response: return response -ELEMENTS_FILE_HEADER = "x-unstructured-elements-file" +# Marks a response this client synthesized, whose body is a path to a local file rather +# than elements. It is deliberately an httpx extension and NOT a response header: a +# header is wire-controlled, so a server could name any local path and have the SDK hand +# it to the caller, who then opens and (per the documented usage) deletes it. Extensions +# are populated by the transport, so a remote server cannot set this key. +ELEMENTS_FILE_EXTENSION_KEY = "unstructured_elements_file" NDJSON_MEDIA_TYPE = "application/x-ndjson" _SNIFF_BLOCK_SIZE = 64 @@ -364,8 +369,9 @@ def create_elements_file_response(elements_file: str) -> httpx.Response: """Create a synthetic 200 response whose payload is a path to an NDJSON file. Mirrors the split hook's existing convention of a cached chunk response carrying its - temp-file path as the body. The path is also set as a header so the SDK can tell this - apart from a real NDJSON body streamed from the server. + temp-file path as the body. The path is also recorded in `ELEMENTS_FILE_EXTENSION_KEY` + so the SDK can tell this apart from a real NDJSON body streamed from the server + without trusting anything that came off the wire. Args: elements_file: Path to the combined NDJSON file of elements. @@ -379,8 +385,8 @@ def create_elements_file_response(elements_file: str) -> httpx.Response: headers={ "Content-Type": NDJSON_MEDIA_TYPE, "Content-Length": str(len(content)), - ELEMENTS_FILE_HEADER: elements_file, }, + extensions={ELEMENTS_FILE_EXTENSION_KEY: elements_file}, ) setattr(response, "_content", content) return response diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index 5d64aa49..5c6ed538 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -1510,10 +1510,23 @@ def _combine_chunks_to_ndjson(self, operation_id: str, chunk_paths: list[str]) - after_success returns, which would delete the file before the caller could read it. The combined file therefore outlives the operation and the caller owns deleting it (see `PartitionResponse.elements_file`). + + Recombination writes to a staging file that is renamed into place only once it + completes. A malformed chunk makes the parse raise partway through, and since the + combined file is the one thing here nothing else owns, a partial one would survive + the failed operation as an orphan under the final name. """ 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) + fd, staging_path = tempfile.mkstemp(suffix=".ndjson.partial", dir=temp_dir_path) + os.close(fd) + try: + written = request_utils.combine_chunk_files_to_ndjson(chunk_paths, staging_path) + # Same directory, so this is atomic; the caller never sees a partial file. + os.replace(staging_path, out_path) + except BaseException: + _unlink_quietly([staging_path]) + raise 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", diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index 3ecf6f67..0ce9e133 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -10,7 +10,7 @@ from unstructured_client.models import errors, operations, shared from unstructured_client.types import BaseModel, OptionalNullable, UNSET from unstructured_client._hooks.custom.clean_server_url_hook import clean_server_url -from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_HEADER +from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_EXTENSION_KEY from unstructured_client.utils.unmarshal_json_response import unmarshal_json_response @@ -35,11 +35,14 @@ def _ndjson_elements_file(http_res: httpx.Response) -> str: """Resolve an NDJSON response to a path on disk, without parsing the elements. 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`, so there is nothing to - do but read the header. Otherwise this is a real body from the server, which is - streamed to a temp file. + NDJSON file and records the path in `ELEMENTS_FILE_EXTENSION_KEY`. Otherwise this is a + real body from the server, which is written to a temp file we create. + + Only the extension is trusted. A path taken from a response header would be + server-controlled, letting a hostile server name any local file for the caller to open + and then delete. """ - existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + existing_path = http_res.extensions.get(ELEMENTS_FILE_EXTENSION_KEY) if existing_path: return existing_path @@ -51,7 +54,7 @@ def _ndjson_elements_file(http_res: httpx.Response) -> str: async def _ndjson_elements_file_async(http_res: httpx.Response) -> str: """Async counterpart of `_ndjson_elements_file`.""" - existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + existing_path = http_res.extensions.get(ELEMENTS_FILE_EXTENSION_KEY) if existing_path: return existing_path From c01a6d1de5ae8fc2c66772c064c7bb4eb4406d79 Mon Sep 17 00:00:00 2001 From: Yao You Date: Sat, 1 Aug 2026 13:04:24 -0500 Subject: [PATCH 3/5] chore: release 0.46.0 Ships the NDJSON elements-file mode. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 5 +++++ RELEASES.md | 10 ++++++++++ src/unstructured_client/_version.py | 4 ++-- 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0661441e..344fb4fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,8 @@ +## 0.46.0 + +### Features +* Add an NDJSON elements-file mode to `partition()`. Pass `accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON` to get `PartitionResponse.elements_file` — a path to an NDJSON file with one element per line — instead of `PartitionResponse.elements`. On the split-PDF path the per-chunk temp files are concatenated on disk rather than parsed, flattened, re-serialized with `json.dumps` and re-parsed by the SDK, which held four copies of the document in memory at once. Peak memory becomes roughly one chunk instead of the whole document. **The caller owns the returned file and is responsible for deleting it.** Requesting `application/json` (the default) is unchanged. + ## 0.45.0 ### Features diff --git a/RELEASES.md b/RELEASES.md index 51643612..d07c7d53 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -1241,3 +1241,13 @@ Based on: - [python v0.45.0] . ### Releases - [PyPI v0.45.0] https://pypi.org/project/unstructured-client/0.45.0 - . + +## 2026-08-01 00:00:00 +### Changes +Based on: +- OpenAPI Doc +- Speakeasy CLI 1.601.0 (2.680.0) https://github.com/speakeasy-api/speakeasy +### Generated +- [python v0.46.0] . +### Releases +- [PyPI v0.46.0] https://pypi.org/project/unstructured-client/0.46.0 - . diff --git a/src/unstructured_client/_version.py b/src/unstructured_client/_version.py index 8a14e983..6fa28865 100644 --- a/src/unstructured_client/_version.py +++ b/src/unstructured_client/_version.py @@ -3,10 +3,10 @@ import importlib.metadata __title__: str = "unstructured-client" -__version__: str = "0.45.0" +__version__: str = "0.46.0" __openapi_doc_version__: str = "1.2.31" __gen_version__: str = "2.680.0" -__user_agent__: str = "speakeasy-sdk/python 0.45.0 2.680.0 1.2.31 unstructured-client" +__user_agent__: str = "speakeasy-sdk/python 0.46.0 2.680.0 1.2.31 unstructured-client" try: if __package__ is not None: From 5fc5a14247fa02f404015b80b46f311afacd3dd1 Mon Sep 17 00:00:00 2001 From: Yao You Date: Sun, 2 Aug 2026 10:37:19 -0500 Subject: [PATCH 4/5] fix: close remaining temp-file and cancellation gaps in NDJSON mode Addresses review findings on the elements-file path. Orphaned files on failure. Both `_ndjson_elements_file` helpers create their destination with delete=False, so a body that raises or is cancelled partway through left the partial copy behind; they now unlink it and re-raise. `write_chunk_body_to_temp` had the same shape -- the caller only registers the path for cleanup once the function returns, so a failed write orphaned the file, and a full disk is exactly the failure that repeats. Cancellation race. Recombination runs in a worker thread that cancellation cannot interrupt, so `_clear_operation` could tear the operation down while it was still running; publishing the path afterwards resurrected a cleared dict entry and orphaned the file. Publishing is now gated on the operation still being live, under a lock that `_clear_operation` also takes when dropping `pending_operation_ids`. Ownership is explicit either way: the success path claims the path out of `ndjson_output_path`, so anything still recorded at teardown was never delivered and is safe to delete. Docs regeneration. docs/models/operations/partitionresponse.md is generated and tracked in gen.lock, and the generation workflow runs on a daily cron, so the elements_file row would have been dropped within a day of merging. Added to .genignore alongside the model, and the regeneration guard now asserts the row. README. Noted that the memory saving applies to the split-PDF path -- unsplit inputs still buffer the body -- so the caveat is visible where the feature is advertised rather than only in the PR. The example's cleanup used a bare unlink in a finally, which would mask a failure to open the file with FileNotFoundError; it now uses Path.unlink(missing_ok=True). The spilled-body regression guard asserted against `_content` it had assigned itself, so it could not fail if the hook stopped releasing the body. It now drives `_elements_from_task_responses`, and was confirmed to fail with the release removed. Co-Authored-By: Claude Opus 5 (1M context) --- .genignore | 4 + README.md | 8 +- .../unit/test_ndjson_elements_file.py | 89 ++++++++++++++++--- .../unit/test_regeneration_guards.py | 6 ++ .../_hooks/custom/request_utils.py | 14 ++- .../_hooks/custom/split_pdf_hook.py | 51 +++++++++-- src/unstructured_client/general.py | 39 ++++++-- 7 files changed, 179 insertions(+), 32 deletions(-) diff --git a/.genignore b/.genignore index b68a3c44..83e366df 100644 --- a/.genignore +++ b/.genignore @@ -30,3 +30,7 @@ src/unstructured_client/utils/retries.py # response field, follow the same procedure as general.py above. # See test_regeneration_guards.py::test_partition_response_keeps_elements_file. src/unstructured_client/models/operations/partition.py + +# Docs for that same custom elements_file field. This file is generated from the spec +# and the daily generation workflow would otherwise drop the row on its next run. +docs/models/operations/partitionresponse.md diff --git a/README.md b/README.md index b329ca40..ef00b584 100644 --- a/README.md +++ b/README.md @@ -433,10 +433,13 @@ For very large documents, the parsed element list can dominate the client's memo **You own the returned file and are responsible for deleting it.** +> [!NOTE] +> The memory saving applies to the split-PDF path, i.e. a PDF with `split_pdf_page=True` (the default). For unsplit inputs — a non-PDF file, or `split_pdf_page=False` — the response body is still read fully into memory before being written to disk, so peak memory can reach roughly twice the body size. You still get `elements_file` either way. + Example: ```python import json -import os +from pathlib import Path from unstructured_client.general import PartitionAcceptEnum @@ -451,7 +454,8 @@ try: element = json.loads(line) ... finally: - os.unlink(res.elements_file) + # missing_ok so a failure to open the file isn't masked by the cleanup. + Path(res.elements_file).unlink(missing_ok=True) ``` diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index eac9fb95..e1d9b9c0 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -186,24 +186,32 @@ def test_combine_accepts_bodies_spilled_without_caching(tmp_path): def test_spilled_body_is_released_from_the_response(tmp_path): - """After spilling, the response must no longer hold the body. + """After spilling, the chunk response must no longer hold the body. - Regression guard: 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, defeating the point of spilling. + Regression guard for the release step in `_elements_from_task_responses`: every chunk + response stays in `api_successful_responses` for failure bookkeeping, so spilling to + disk without clearing `_content` still accumulates the whole document in memory, + defeating the point of spilling. + + Driven through the hook on purpose. The release happens there, not in + `write_chunk_body_to_temp`, so a test that clears `_content` itself would still pass + if the hook ever stopped doing it. """ - 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) + operation_id = "op-release" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + elements = [ + {"type": "Table", "text": f"t{i}", "metadata": {"image_base64": "A" * 2000}} + for i in range(3) + ] + response = _ndjson_response(elements) + assert len(response.content) > 512 - path = write_chunk_body_to_temp(response, str(tmp_path)) - response._content = path.encode() + hook._elements_from_task_responses(operation_id, [(0, response)], started_at=0.0) - # 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 + # The response now costs a path rather than a payload... assert len(response.content) < 512 + # ...and the elements still made it into the combined output. + assert _read_ndjson(hook.ndjson_output_path[operation_id]) == elements def test_elements_file_response_carries_path_in_extension_and_body(tmp_path): @@ -258,6 +266,8 @@ def _hook_in_ndjson_mode(operation_id, tmp_path): hook.cache_tmp_data_feature[operation_id] = False hook.cache_tmp_data_dir[operation_id] = str(tmp_path) hook.allow_failed[operation_id] = False + # Marks the operation live; `_clear_operation` removing it is what signals teardown. + hook.pending_operation_ids[operation_id] = operation_id return hook @@ -338,6 +348,59 @@ def test_combined_file_is_discarded_when_a_failure_response_is_returned(tmp_path assert _ndjson_files_in(tmp_path) == [] +def test_output_is_discarded_when_the_operation_was_cleared_mid_recombination(tmp_path): + """Recombination runs in a thread that cancellation cannot interrupt. + + If `_clear_operation` tears the operation down first, publishing the path would both + resurrect a cleared dict entry and orphan the file, since nothing will ever read it. + """ + operation_id = "op-cancelled" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + # `before_request` registers this; `_clear_operation` removing it is what marks the + # operation dead. Simulate the teardown having already happened. + hook.pending_operation_ids.pop(operation_id, None) + + hook._elements_from_task_responses( + operation_id, [(0, _ndjson_response(_elements("a", 2)))], started_at=0.0 + ) + + assert operation_id not in hook.ndjson_output_path + assert _ndjson_files_in(tmp_path) == [] + + +def test_clear_operation_deletes_an_unclaimed_output_file(tmp_path): + """A path still recorded at teardown was never handed to the caller, so it is ours.""" + operation_id = "op-unclaimed" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + + hook._elements_from_task_responses( + operation_id, [(0, _ndjson_response(_elements("a", 2)))], started_at=0.0 + ) + combined = hook.ndjson_output_path[operation_id] + assert os.path.exists(combined) + + hook._clear_operation(operation_id) + + assert not os.path.exists(combined) + + +def test_clear_operation_keeps_an_output_file_the_caller_claimed(tmp_path): + """The success path hands the path over, so teardown must not delete it.""" + operation_id = "op-claimed" + hook = _hook_in_ndjson_mode(operation_id, tmp_path) + + hook._elements_from_task_responses( + operation_id, [(0, _ndjson_response(_elements("a", 2)))], started_at=0.0 + ) + response = hook._build_after_success_response(operation_id, httpx.Response(200), []) + combined = response.extensions[ELEMENTS_FILE_EXTENSION_KEY] + + hook._clear_operation(operation_id) + + assert os.path.exists(combined) + assert len(_read_ndjson(combined)) == 2 + + def test_malformed_chunk_leaves_no_partial_output_behind(tmp_path): """Regression guard: recombination that raises must not orphan a partial file. diff --git a/_test_unstructured_client/unit/test_regeneration_guards.py b/_test_unstructured_client/unit/test_regeneration_guards.py index 48f04a57..f465e1d6 100644 --- a/_test_unstructured_client/unit/test_regeneration_guards.py +++ b/_test_unstructured_client/unit/test_regeneration_guards.py @@ -98,9 +98,15 @@ def test_partition_response_keeps_elements_file(): for path in ( "src/unstructured_client/general.py", "src/unstructured_client/models/operations/partition.py", + "docs/models/operations/partitionresponse.md", ): assert path in genignore, f"{path} carries custom code and must stay in .genignore" + # The docs row is generated from the spec too, so the daily generation workflow would + # drop it without the .genignore entry above. + response_docs = (REPO_ROOT / "docs/models/operations/partitionresponse.md").read_text() + assert "elements_file" in response_docs + def test_body_create_job_input_files_are_serialized_as_multipart_files(): request = shared.BodyCreateJob( diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index f48d04bf..215983a8 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -360,8 +360,18 @@ def write_chunk_body_to_temp(response: httpx.Response, dir_: Optional[str] = Non The path to the spilled file. The caller owns deleting it. """ fd, path = tempfile.mkstemp(suffix=".ndjson", dir=dir_ or tempfile.gettempdir()) - with os.fdopen(fd, "wb") as f: - f.write(response.content) + try: + with os.fdopen(fd, "wb") as f: + f.write(response.content) + except BaseException: + # The caller only registers this path for cleanup once we return, so a failed + # write (a full disk, most likely) has to clean up after itself or the file is + # orphaned -- and a full disk is exactly the case that repeats. + try: + os.unlink(path) + except OSError: + pass + raise return path diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index 5c6ed538..9e31b3b3 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -543,9 +543,13 @@ def __init__(self) -> None: self.cache_tmp_data_dir: dict[str, str] = {} # NDJSON elements-file mode: when the caller asks 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. + # re-serialized, and the combined path is handed back on the response. self.ndjson_mode: dict[str, bool] = {} self.ndjson_output_path: dict[str, str] = {} + # Guards publishing the combined file against concurrent operation teardown. + # Recombination runs in a worker thread that a cancelled operation cannot stop, so + # it can still be running when `_clear_operation` fires on the event-loop thread. + self._ndjson_lock = threading.Lock() @staticmethod def _get_operation_id_from_request(request: Optional[httpx.Request]) -> Optional[str]: @@ -1515,6 +1519,12 @@ def _combine_chunks_to_ndjson(self, operation_id: str, chunk_paths: list[str]) - completes. A malformed chunk makes the parse raise partway through, and since the combined file is the one thing here nothing else owns, a partial one would survive the failed operation as an orphan under the final name. + + Publishing is also gated on the operation still being live. This runs in a worker + thread that cancellation cannot interrupt, so `_clear_operation` may already have + torn the operation down by the time we finish; recording the path then would both + resurrect a cleared dict entry and orphan the file, since nothing downstream will + ever read it. """ temp_dir_path = self.cache_tmp_data_dir.get(operation_id) or tempfile.gettempdir() out_path = f"{temp_dir_path}/{uuid.uuid4()}.ndjson" @@ -1527,7 +1537,16 @@ def _combine_chunks_to_ndjson(self, operation_id: str, chunk_paths: list[str]) - except BaseException: _unlink_quietly([staging_path]) raise - self.ndjson_output_path[operation_id] = out_path + + with self._ndjson_lock: + if operation_id not in self.pending_operation_ids: + _unlink_quietly([out_path]) + logger.warning( + "split_pdf event=ndjson_output_discarded operation_id=%s reason=operation_cleared", + operation_id, + ) + return + 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, @@ -1549,10 +1568,21 @@ def _discard_ndjson_output(self, operation_id: str) -> None: Once we return a failure response instead, nothing downstream learns the path, so without this the combined file is leaked for the lifetime of the host. """ - out_path = self.ndjson_output_path.pop(operation_id, None) + with self._ndjson_lock: + out_path = self.ndjson_output_path.pop(operation_id, None) if out_path is not None: _unlink_quietly([out_path]) + def _claim_ndjson_output(self, operation_id: str) -> Optional[str]: + """Take the combined file's path, transferring ownership of it to the caller. + + Removing it from `ndjson_output_path` is what makes cleanup unambiguous: a path + still recorded when the operation is torn down is one the caller never received, + so `_clear_operation` can delete it without risking the file it just handed over. + """ + with self._ndjson_lock: + return self.ndjson_output_path.pop(operation_id, None) + def _build_after_success_response( self, operation_id: str, @@ -1587,7 +1617,7 @@ def _build_after_success_response( # 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) + ndjson_path = self._claim_ndjson_output(operation_id) if ndjson_path is None: logger.warning( "split_pdf event=ndjson_missing_output operation_id=%s falling_back=true", @@ -1718,9 +1748,16 @@ 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) + with self._ndjson_lock: + self.ndjson_mode.pop(operation_id, None) + # Anything still recorded here was never claimed by the caller -- the + # operation was torn down first -- so this is ours to delete. Dropping + # `pending_operation_ids` under the same lock is what stops a recombination + # worker still running in a thread from publishing after this point. + undelivered_ndjson = self.ndjson_output_path.pop(operation_id, None) + self.pending_operation_ids.pop(operation_id, None) + if undelivered_ndjson is not None: + _unlink_quietly([undelivered_ndjson]) future = self.operation_futures.pop(operation_id, None) loop_holder = self.operation_loops.pop(operation_id, None) executor = self.executors.pop(operation_id, None) diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index 0ce9e133..4f1be006 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -3,6 +3,7 @@ from .basesdk import BaseSDK from enum import Enum import httpx +import os import tempfile from typing import Any, Dict, List, Mapping, Optional, Union, cast from unstructured_client import utils @@ -31,6 +32,18 @@ def _new_elements_file(): ) +def _discard_elements_file(path: str) -> None: + """Remove a partially written elements file, ignoring a failure to do so. + + The copy is created with delete=False so it can outlive this function, which means a + body that fails or is cancelled partway through would otherwise leave the file behind. + """ + try: + os.unlink(path) + except OSError: + pass + + def _ndjson_elements_file(http_res: httpx.Response) -> str: """Resolve an NDJSON response to a path on disk, without parsing the elements. @@ -46,10 +59,15 @@ def _ndjson_elements_file(http_res: httpx.Response) -> str: if existing_path: return existing_path - with _new_elements_file() as out: - for byte_chunk in http_res.iter_bytes(): - out.write(byte_chunk) - return out.name + out = _new_elements_file() + try: + with out: + for byte_chunk in http_res.iter_bytes(): + out.write(byte_chunk) + except BaseException: + _discard_elements_file(out.name) + raise + return out.name async def _ndjson_elements_file_async(http_res: httpx.Response) -> str: @@ -58,10 +76,15 @@ async def _ndjson_elements_file_async(http_res: httpx.Response) -> str: if existing_path: return existing_path - with _new_elements_file() as out: - async for byte_chunk in http_res.aiter_bytes(): - out.write(byte_chunk) - return out.name + out = _new_elements_file() + try: + with out: + async for byte_chunk in http_res.aiter_bytes(): + out.write(byte_chunk) + except BaseException: + _discard_elements_file(out.name) + raise + return out.name class General(BaseSDK): From 3080cf6077f1d77e084f8cd45ef36fd39e905d07 Mon Sep 17 00:00:00 2001 From: Yao You Date: Sun, 2 Aug 2026 12:44:29 -0500 Subject: [PATCH 5/5] test: pin the orphan-cleanup behavior on NDJSON write failures The unlink-on-failure paths added in the previous commit had no coverage -- the existing tests only walked the success roundtrip, so the cleanup could have been removed without anything going red. Three tests, each confirmed to fail with its corresponding unlink removed: - write_chunk_body_to_temp: os.fdopen is patched so the write raises ENOSPC. The wrapper still closes the real handle, so the fd is not leaked by the test itself. - _ndjson_elements_file and its async counterpart: a response whose byte iterator raises partway through. tempfile.tempdir is redirected at the test's tmp_path so the assertion can see whether anything was left behind. Each asserts both halves of the contract: no file survives, and the original exception still propagates rather than being swallowed by the cleanup. Coverage for the general.py pair was not requested in review, but those helpers grew the same delete=False cleanup in the same commit and had the same gap. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_elements_file.py | 86 ++++++++++++++++++- 1 file changed, 85 insertions(+), 1 deletion(-) diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index e1d9b9c0..22b178a7 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -9,14 +9,18 @@ - leave no temp files behind other than the combined file the caller owns """ +import errno import json import os +import tempfile from pathlib import Path +from unittest import mock import pytest import httpx +from unstructured_client._hooks.custom import request_utils from unstructured_client._hooks.custom.request_utils import ( ELEMENTS_FILE_EXTENSION_KEY, combine_chunk_files_to_ndjson, @@ -24,7 +28,10 @@ write_chunk_body_to_temp, ) from unstructured_client._hooks.custom.split_pdf_hook import SplitPdfHook -from unstructured_client.general import _ndjson_elements_file +from unstructured_client.general import ( + _ndjson_elements_file, + _ndjson_elements_file_async, +) def _elements(prefix, count): @@ -171,6 +178,83 @@ def test_write_chunk_body_to_temp_roundtrips(tmp_path): assert _read_ndjson(path) == elements +def test_spill_failure_leaves_no_orphan_file(tmp_path): + """A write that fails partway must take its own temp file with it. + + The caller only registers the returned path for cleanup once this function returns, + so an orphan here is permanent -- and a full disk, the likeliest cause, is exactly + the failure that repeats on every retry. + """ + real_fdopen = os.fdopen + + class _FailingWriter: + """Wraps the real handle so the fd is still closed, but the write blows up.""" + + def __init__(self, handle): + self._handle = handle + + def write(self, _data): + raise OSError(errno.ENOSPC, "No space left on device") + + def __enter__(self): + return self + + def __exit__(self, *_exc): + self._handle.close() + return False + + def _failing_fdopen(fd, mode): + return _FailingWriter(real_fdopen(fd, mode)) + + response = httpx.Response(status_code=200, content=b'{"type": "Table"}\n') + + with mock.patch.object(request_utils.os, "fdopen", _failing_fdopen): + with pytest.raises(OSError) as excinfo: + write_chunk_body_to_temp(response, str(tmp_path)) + + assert excinfo.value.errno == errno.ENOSPC + assert list(tmp_path.iterdir()) == [] + + +def test_elements_file_copy_failure_leaves_no_orphan(tmp_path, monkeypatch): + """A body that dies mid-copy must not leave the partial file behind. + + The destination is created with delete=False so it can outlive the helper, which is + exactly what makes an interrupted copy leak. + """ + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + + class _FailingResponse: + extensions: dict = {} + + def iter_bytes(self): + yield b'{"type": "Table"}\n' + raise httpx.ReadError("connection dropped") + + with pytest.raises(httpx.ReadError): + _ndjson_elements_file(_FailingResponse()) + + assert list(tmp_path.glob("unst_elements_*")) == [] + + +@pytest.mark.asyncio +async def test_elements_file_copy_failure_leaves_no_orphan_async(tmp_path, monkeypatch): + """Async counterpart of `test_elements_file_copy_failure_leaves_no_orphan`.""" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + + class _FailingResponse: + extensions: dict = {} + + async def aiter_bytes(self): + yield b'{"type": "Table"}\n' + raise httpx.ReadError("connection dropped") + + with pytest.raises(httpx.ReadError): + await _ndjson_elements_file_async(_FailingResponse()) + + assert list(tmp_path.glob("unst_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)