From fce575a512991e0f8620a2ef8d33e20e8046ca66 Mon Sep 17 00:00:00 2001 From: Yao You Date: Sat, 1 Aug 2026 13:04:24 -0500 Subject: [PATCH 01/11] 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 02/11] 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 03/11] 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 04/11] 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 05/11] 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) From 1ec2e25241bc4a03a31cdd7244eeb2b752366582 Mon Sep 17 00:00:00 2001 From: Yao You Date: Mon, 3 Aug 2026 10:26:39 -0500 Subject: [PATCH 06/11] test: assert orphan cleanup against the path production created The orphan checks asserted that a directory held no matching files, which only proves cleanup if the file landed in that directory to begin with. That held solely because the helpers route through the global tempfile.tempdir the tests patch -- an assumption the tests never checked. A change that passed an explicit dir would have made them pass while a partial file leaked into the real temp dir. A shared `_record_created_paths` spy now wraps the temp-file factory, delegating to the real one and recording each path it hands out. The tests assert on that path: one file was created, it was where the test expected, and it is gone. Verified by breaking it two ways. With the unlink removed the tests fail, as before. With the unlink removed AND the destination pinned to an explicit dir that ignores the patched tempdir -- the vacuous case, which the old directory-glob assertion passed -- they now fail too. Applied to the spill test as well. Review only raised the two copy-failure tests, but that one asserted on an empty directory for the same reason and could go vacuous the same way if `dir_` stopped being honored. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_elements_file.py | 42 +++++++++++++++++-- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index 22b178a7..75322d6d 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -20,6 +20,7 @@ import httpx +from unstructured_client import general from unstructured_client._hooks.custom import request_utils from unstructured_client._hooks.custom.request_utils import ( ELEMENTS_FILE_EXTENSION_KEY, @@ -178,7 +179,27 @@ 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): +def _record_created_paths(monkeypatch, module, attr): + """Spy on a temp-file factory, recording every path it hands out. + + Asserting on an empty directory only proves cleanup if the file landed in that + directory in the first place. Recording what production actually created keeps the + assertion honest even if the destination moves. + """ + created: list[str] = [] + real = getattr(module, attr) + + def _spy(*args, **kwargs): + result = real(*args, **kwargs) + # mkstemp returns (fd, path); NamedTemporaryFile returns a handle with .name. + created.append(result[1] if isinstance(result, tuple) else result.name) + return result + + monkeypatch.setattr(module, attr, _spy) + return created + + +def test_spill_failure_leaves_no_orphan_file(tmp_path, monkeypatch): """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, @@ -207,13 +228,18 @@ def _failing_fdopen(fd, mode): return _FailingWriter(real_fdopen(fd, mode)) response = httpx.Response(status_code=200, content=b'{"type": "Table"}\n') + created = _record_created_paths(monkeypatch, request_utils.tempfile, "mkstemp") 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()) == [] + # Assert against the path production actually created, so the check cannot pass + # vacuously if the file ever stops landing where the test expects. + assert len(created) == 1 + assert Path(created[0]).parent == tmp_path + assert not os.path.exists(created[0]) def test_elements_file_copy_failure_leaves_no_orphan(tmp_path, monkeypatch): @@ -223,6 +249,7 @@ def test_elements_file_copy_failure_leaves_no_orphan(tmp_path, monkeypatch): exactly what makes an interrupted copy leak. """ monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + created = _record_created_paths(monkeypatch, general, "_new_elements_file") class _FailingResponse: extensions: dict = {} @@ -234,13 +261,18 @@ def iter_bytes(self): with pytest.raises(httpx.ReadError): _ndjson_elements_file(_FailingResponse()) - assert list(tmp_path.glob("unst_elements_*")) == [] + # The recorded path is the one production created, so this cannot pass vacuously if + # the helper ever stops routing through the global tempdir. + assert len(created) == 1 + assert Path(created[0]).parent == tmp_path + assert not os.path.exists(created[0]) @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)) + created = _record_created_paths(monkeypatch, general, "_new_elements_file") class _FailingResponse: extensions: dict = {} @@ -252,7 +284,9 @@ async def aiter_bytes(self): with pytest.raises(httpx.ReadError): await _ndjson_elements_file_async(_FailingResponse()) - assert list(tmp_path.glob("unst_elements_*")) == [] + assert len(created) == 1 + assert Path(created[0]).parent == tmp_path + assert not os.path.exists(created[0]) def test_combine_accepts_bodies_spilled_without_caching(tmp_path): From 358e21ad3a7687ab85628f87c3b988d33355d71b Mon Sep 17 00:00:00 2001 From: Yao You Date: Mon, 3 Aug 2026 13:59:03 -0500 Subject: [PATCH 07/11] fix: set elements_file even when the server returns JSON The deployed API does not offer application/x-ndjson -- its spec (v1.5.69) declares only application/json and text/csv, and the string "ndjson" appears nowhere in the document. So any request that bypasses the split-PDF hook came back as JSON, the application/json branch matched first and won, and elements_file stayed None while elements was populated. The opt-in was silently ignored. That path is much wider than "non-PDF". _before_request_unlocked short-circuits when split_size >= page_count, and get_optimal_split_size floors at MIN_PAGES_PER_SPLIT=2, so one- and two-page PDFs are sent whole -- it would have read as intermittent to anyone hitting it on small inputs. The x-ndjson branch is now checked first, and a JSON body is written out as NDJSON when NDJSON was requested. elements_file is therefore set for every input and callers need one code path rather than two. This does not bound memory on the unsplit path -- the body is already read and parsing adds the element list on top -- it makes the contract uniform. The memory win remains the split-PDF path, where the hook concatenates chunk files on disk and this conversion is never reached. Also fixes the README example, which raised TypeError from open(None) and then raised the same TypeError again out of Path(None) in the finally, masking the first. Adds the end-to-end coverage whose absence let this through: all 24 existing tests were helper- or hook-level, so nothing exercised the media-type dispatch in general.py. Three tests now drive partition() over a mock transport -- server returns JSON, server returns NDJSON, and no override at all. The first fails against the old dispatch. They hold the client for the call rather than chaining off a temporary, since sdk.py registers a weakref.finalize that closes the transport. Smaller review notes: log a warning per skipped empty chunk so element_count can be reconciled against chunk_count; document that the non-"[" sniff assumes one JSON value per line rather than checking; trim the _combine_chunks_to_ndjson docstring. The regeneration comments no longer claim a daily workflow is about to drop the field -- generation is blocked at the Speakeasy account level and gen.lock has not moved since 2026-01. Whether that scaffolding is worth keeping is left open. Co-Authored-By: Claude Opus 5 (1M context) --- .genignore | 9 +- CHANGELOG.md | 2 +- README.md | 2 +- .../unit/test_ndjson_elements_file.py | 99 +++++++++++++++++++ .../unit/test_regeneration_guards.py | 8 +- .../_hooks/custom/request_utils.py | 13 ++- .../_hooks/custom/split_pdf_hook.py | 24 ++--- src/unstructured_client/general.py | 73 +++++++++++--- 8 files changed, 193 insertions(+), 37 deletions(-) diff --git a/.genignore b/.genignore index 83e366df..cb92e874 100644 --- a/.genignore +++ b/.genignore @@ -31,6 +31,11 @@ src/unstructured_client/utils/retries.py # 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 for that same custom elements_file field. Generated from the spec, so a +# regeneration would drop the row. +# +# Note: SDK generation is currently blocked at the Speakeasy account level +# ("generation access blocked"), so nothing can regenerate today and gen.lock has not +# moved since 2026-01. These entries are insurance for when that is restored, not a +# defence against an imminent run. docs/models/operations/partitionresponse.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 344fb4fd..eb910939 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ ## 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. +* 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. `elements_file` is set for every input, including ones that are not split and responses from a server that ignores the `Accept` header, so callers need only one code path — but the memory saving itself applies only to split PDFs (a PDF of more than two pages, with `split_pdf_page=True`). **The caller owns the returned file and is responsible for deleting it.** Requesting `application/json` (the default) is unchanged. ## 0.45.0 diff --git a/README.md b/README.md index ef00b584..39f9dc3a 100644 --- a/README.md +++ b/README.md @@ -434,7 +434,7 @@ 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. +> `elements_file` is always set when you pass this header, but the **memory saving** only applies to the split-PDF path. An input is only split when it is a PDF, `split_pdf_page=True` (the default), and it has more than two pages — `split_size` is floored at 2, so one- and two-page PDFs are sent whole. For those, and for non-PDFs, the response body is read fully into memory before being written to disk, so peak memory can reach roughly twice the body size. Example: ```python diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index 75322d6d..30235936 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -29,10 +29,13 @@ write_chunk_body_to_temp, ) from unstructured_client._hooks.custom.split_pdf_hook import SplitPdfHook +from unstructured_client import UnstructuredClient from unstructured_client.general import ( + PartitionAcceptEnum, _ndjson_elements_file, _ndjson_elements_file_async, ) +from unstructured_client.models import operations, shared def _elements(prefix, count): @@ -547,3 +550,99 @@ def test_no_combined_file_is_created_when_every_chunk_failed(tmp_path): assert operation_id not in hook.ndjson_output_path assert _ndjson_files_in(tmp_path) == [] + + +# --- end to end through partition() ----------------------------------------------- +# +# The gap these close: every other test here is helper- or hook-level, so nothing +# exercised the response dispatch in `general.py`. That is where the media-type branches +# are chosen, and it is where NDJSON mode silently fell back to `elements`. + + +def _mock_client(handler): + """Build a client over a mock transport. + + Callers must hold the returned client for the duration of the call: `sdk.py` registers + a `weakref.finalize` that closes the underlying httpx client, so chaining off a + temporary tears the transport down mid-request. + """ + return UnstructuredClient( + api_key_auth="x", + server_url="http://localhost:8000", + client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + +def _text_request(): + """A non-PDF input, so the split-PDF hook does not engage.""" + return operations.PartitionRequest( + partition_parameters=shared.PartitionParameters( + files=shared.Files(content=b"hello", file_name="doc.txt"), + ) + ) + + +def test_partition_sets_elements_file_when_server_ignores_the_accept_header(): + """The deployed API only offers application/json, so this is the live unsplit path. + + Regression guard: the JSON branch used to be matched first and win, populating + `elements` and leaving `elements_file` None even though NDJSON was requested. + """ + elements = [{"type": "Table", "text": "t0"}, {"type": "NarrativeText", "text": "t1"}] + + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/json"}, json=elements + ) + + client = _mock_client(handler) + res = client.general.partition( + request=_text_request(), + accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON, + ) + + assert res.elements_file is not None + assert res.elements is None + try: + assert _read_ndjson(res.elements_file) == elements + finally: + os.unlink(res.elements_file) + + +def test_partition_sets_elements_file_when_server_returns_ndjson(): + """The path that becomes live if the API ever honors the Accept header.""" + elements = [{"type": "Table", "text": "t0"}] + body = "".join(json.dumps(e) + "\n" for e in elements).encode() + + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/x-ndjson"}, content=body + ) + + client = _mock_client(handler) + res = client.general.partition( + request=_text_request(), + accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON, + ) + + assert res.elements is None + try: + assert _read_ndjson(res.elements_file) == elements + finally: + os.unlink(res.elements_file) + + +def test_partition_without_the_override_still_returns_elements(): + """The default must be untouched: no elements_file, elements populated as before.""" + elements = [{"type": "Table", "text": "t0"}] + + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/json"}, json=elements + ) + + client = _mock_client(handler) + res = client.general.partition(request=_text_request()) + + assert res.elements == elements + assert res.elements_file is None diff --git a/_test_unstructured_client/unit/test_regeneration_guards.py b/_test_unstructured_client/unit/test_regeneration_guards.py index f465e1d6..62dc76a1 100644 --- a/_test_unstructured_client/unit/test_regeneration_guards.py +++ b/_test_unstructured_client/unit/test_regeneration_guards.py @@ -86,6 +86,10 @@ def test_partition_response_keeps_elements_file(): Both the model and the enum value that selects it live in generated files; the .genignore entries are the only thing keeping them. + + Note: SDK generation is currently blocked at the Speakeasy account level, so this + guards a path that cannot execute today. Whether it is worth keeping is a live + question -- see the discussion on PR #347. """ from unstructured_client.general import PartitionAcceptEnum from unstructured_client.models import operations @@ -102,8 +106,8 @@ def test_partition_response_keeps_elements_file(): ): 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. + # The docs row is generated from the spec too, so a regeneration 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 diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index 215983a8..66bdaacd 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -310,7 +310,11 @@ def combine_chunk_files_to_ndjson(chunk_paths: list[str], out_path: str) -> int: 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. + says which. A leading `[` is treated as an array, and anything else is assumed to be + one JSON value per line -- an assumption, not a check, so a pretty-printed multi-line + object would be split into invalid lines. No endpoint returns that today. + + NDJSON chunks are copied through without parsing. Args: chunk_paths: Per-chunk response files, in element order. @@ -325,6 +329,13 @@ def combine_chunk_files_to_ndjson(chunk_paths: list[str], out_path: str) -> int: with open(chunk_path, "r", encoding="utf-8") as chunk: first_char = _first_non_space_char(chunk) if not first_char: + # A 200 with an empty body. Log it, because otherwise the document is + # silently short and the element count cannot be reconciled against + # the chunk count. + logger.warning( + "split_pdf event=ndjson_empty_chunk file=%s", + os.path.basename(chunk_path), + ) continue chunk.seek(0) diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index 9e31b3b3..bcf089fe 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -1509,22 +1509,14 @@ def _combine_chunks_to_ndjson(self, operation_id: str, chunk_paths: list[str]) - """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`). - - 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. - - 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. + The combined file goes in the cache *parent* directory rather than the operation's + TemporaryDirectory, which `_clear_operation` removes as soon as after_success + returns. It therefore outlives the operation and the caller owns deleting it (see + `PartitionResponse.elements_file`). + + Writing to a staging file and renaming keeps a partial file from surviving under + the final name if a malformed chunk makes the parse raise. Publishing is gated on + the operation still being live; see `_ndjson_lock`. """ temp_dir_path = self.cache_tmp_data_dir.get(operation_id) or tempfile.gettempdir() out_path = f"{temp_dir_path}/{uuid.uuid4()}.ndjson" diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index 4f1be006..7e14e0f2 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 json import os import tempfile from typing import Any, Dict, List, Mapping, Optional, Union, cast @@ -32,6 +33,36 @@ def _new_elements_file(): ) +def _ndjson_requested(accept_header_override: Optional[PartitionAcceptEnum]) -> bool: + return accept_header_override == PartitionAcceptEnum.APPLICATION_X_NDJSON + + +def _json_body_to_elements_file(http_res: httpx.Response) -> str: + """Write a JSON-array body out as NDJSON, so `elements_file` is set either way. + + The deployed API does not offer `application/x-ndjson`: its spec declares only + `application/json` and `text/csv`, so a request that skips the split-PDF hook comes + back as JSON no matter what was asked for. Without this, opting into NDJSON would + silently populate `elements` instead and leave `elements_file` as None, forcing every + caller to handle both shapes for one flag. + + This does not bound memory -- the body is already fully read by the time we get here, + and parsing it adds the element list on top. It exists to keep the contract uniform. + The memory win comes from the split-PDF path, where the hook concatenates chunk files + on disk and this function is never reached. + """ + out = _new_elements_file() + try: + with out: + for element in http_res.json(): + out.write(json.dumps(element, ensure_ascii=False).encode()) + out.write(b"\n") + except BaseException: + _discard_elements_file(out.name) + raise + return out.name + + def _discard_elements_file(path: str) -> None: """Remove a partially written elements file, ignoring a failure to do so. @@ -183,7 +214,21 @@ def partition( ) response_data: Any = None + 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, "200", "application/json"): + if _ndjson_requested(accept_header_override): + return operations.PartitionResponse( + elements_file=_json_body_to_elements_file(http_res), + status_code=http_res.status_code, + content_type=http_res.headers.get("Content-Type") or "", + raw_response=http_res, + ) return operations.PartitionResponse( elements=unmarshal_json_response( Optional[List[Dict[str, Any]]], http_res @@ -199,13 +244,6 @@ 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 @@ -315,7 +353,21 @@ async def partition_async( ) response_data: Any = None + 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, "200", "application/json"): + if _ndjson_requested(accept_header_override): + return operations.PartitionResponse( + elements_file=_json_body_to_elements_file(http_res), + status_code=http_res.status_code, + content_type=http_res.headers.get("Content-Type") or "", + raw_response=http_res, + ) return operations.PartitionResponse( elements=unmarshal_json_response( Optional[List[Dict[str, Any]]], http_res @@ -331,13 +383,6 @@ 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 From ba3e7f47922c988c30719d4d74848ca20a3d6abf Mon Sep 17 00:00:00 2001 From: Yao You Date: Mon, 3 Aug 2026 14:19:25 -0500 Subject: [PATCH 08/11] fix: validate the JSON fallback, key NDJSON off the sent Accept header Four review findings on the elements-file path, all reproduced first. The JSON-to-NDJSON conversion iterated `http_res.json()` directly, bypassing the response schema. A `null` body raised TypeError from iterating None, and a non-array body was worse than that: iterating a dict yields its keys, so `{"detail": "oops"}` was written out as a one-element document reading `"detail"`. It now goes through `unmarshal_json_response`, so a malformed 200 fails exactly as it does on the `elements` path, and `null` yields an empty file. NDJSON detection keyed off `accept_header_override` alone, but `http_headers` can replace Accept too -- and the split-PDF hook already keys off the request header. So one caller passing `http_headers={"Accept": "application/x-ndjson"}` got `elements_file` for a multi-page PDF and `elements` for a two-page one. Detection now reads the header off the built request, which is the same source the hook uses. The async path ran the parse-and-write inline, blocking the event loop for exactly the large bodies this mode exists for. Offloaded with `asyncio.to_thread`. Test `finally` blocks used a bare `os.unlink`, which masks an assertion failure above it with FileNotFoundError -- the same defect flagged in the README two rounds ago, reintroduced in the tests. Now `Path(...).unlink(missing_ok=True)`. Also removes a duplicated "The" left in `_combine_chunks_to_ndjson`'s docstring by the earlier trim. Four tests added, each confirmed to fail against the pre-fix source. The async one originally asserted `asyncio.to_thread` was called, which passed even with the offload removed because other SDK internals use it during the same call; it now records the thread the conversion ran on and compares it to the event-loop thread. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_elements_file.py | 112 +++++++++++++++++- .../_hooks/custom/split_pdf_hook.py | 2 +- src/unstructured_client/general.py | 34 ++++-- 3 files changed, 137 insertions(+), 11 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index 30235936..43d53dfa 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -13,6 +13,7 @@ import json import os import tempfile +import threading from pathlib import Path from unittest import mock @@ -35,7 +36,7 @@ _ndjson_elements_file, _ndjson_elements_file_async, ) -from unstructured_client.models import operations, shared +from unstructured_client.models import errors, operations, shared def _elements(prefix, count): @@ -606,7 +607,8 @@ def handler(_request): try: assert _read_ndjson(res.elements_file) == elements finally: - os.unlink(res.elements_file) + # missing_ok so a failed assertion above is not masked by the cleanup. + Path(res.elements_file).unlink(missing_ok=True) def test_partition_sets_elements_file_when_server_returns_ndjson(): @@ -629,7 +631,8 @@ def handler(_request): try: assert _read_ndjson(res.elements_file) == elements finally: - os.unlink(res.elements_file) + # missing_ok so a failed assertion above is not masked by the cleanup. + Path(res.elements_file).unlink(missing_ok=True) def test_partition_without_the_override_still_returns_elements(): @@ -646,3 +649,106 @@ def handler(_request): assert res.elements == elements assert res.elements_file is None + + +def test_partition_sets_elements_file_when_accept_set_via_http_headers(): + """`http_headers` can replace Accept too, and must behave the same as the override. + + Regression guard: this used to key off `accept_header_override` alone, so the split + hook (which reads the request header) and the unsplit path disagreed for one caller. + """ + elements = [{"type": "Table", "text": "t0"}] + + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/json"}, json=elements + ) + + client = _mock_client(handler) + res = client.general.partition( + request=_text_request(), + http_headers={"Accept": "application/x-ndjson"}, + ) + + assert res.elements is None + try: + assert _read_ndjson(res.elements_file) == elements + finally: + Path(res.elements_file).unlink(missing_ok=True) + + +def test_partition_ndjson_rejects_a_non_list_json_body(): + """A malformed 200 must fail the same way it would on the `elements` path. + + Iterating the raw JSON would write a dict's *keys* out as elements, so a + `{"detail": ...}` body became a one-element document reading `"detail"`. + """ + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/json"}, json={"detail": "oops"} + ) + + client = _mock_client(handler) + with pytest.raises(errors.ResponseValidationError): + client.general.partition( + request=_text_request(), + accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON, + ) + + +def test_partition_ndjson_handles_a_null_json_body(): + """`null` used to raise TypeError from iterating None; it now yields an empty file.""" + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/json"}, content=b"null" + ) + + client = _mock_client(handler) + res = client.general.partition( + request=_text_request(), + accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON, + ) + + try: + assert _read_ndjson(res.elements_file) == [] + finally: + Path(res.elements_file).unlink(missing_ok=True) + + +@pytest.mark.asyncio +async def test_partition_async_ndjson_does_not_block_on_conversion(): + """The async path must offload the parse-and-write, not run it on the event loop.""" + elements = [{"type": "Table", "text": "t0"}] + + def handler(_request): + return httpx.Response( + 200, headers={"Content-Type": "application/json"}, json=elements + ) + + client = UnstructuredClient( + api_key_auth="x", + server_url="http://localhost:8000", + async_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + # Assert on the thread the conversion actually ran on. Patching `asyncio.to_thread` + # and checking `.called` proves nothing -- other SDK internals use it during the same + # call, so that assertion passes even with the offload removed. + event_loop_thread = threading.get_ident() + ran_on = {} + real_convert = general._json_body_to_elements_file + + def _spy(http_res): + ran_on["thread"] = threading.get_ident() + return real_convert(http_res) + + with mock.patch.object(general, "_json_body_to_elements_file", _spy): + res = await client.general.partition_async( + request=_text_request(), + accept_header_override=PartitionAcceptEnum.APPLICATION_X_NDJSON, + ) + + assert ran_on["thread"] != event_loop_thread, "conversion ran on the event loop" + try: + assert _read_ndjson(res.elements_file) == elements + finally: + Path(res.elements_file).unlink(missing_ok=True) diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index bcf089fe..fed214d3 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -1509,7 +1509,7 @@ def _combine_chunks_to_ndjson(self, operation_id: str, chunk_paths: list[str]) - """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 - The combined file goes in the cache *parent* directory rather than the operation's + combined file goes in the cache *parent* directory rather than the operation's TemporaryDirectory, which `_clear_operation` removes as soon as after_success returns. It therefore outlives the operation and the caller owns deleting it (see `PartitionResponse.elements_file`). diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index 7e14e0f2..4694c7d7 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -1,6 +1,7 @@ """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.""" from .basesdk import BaseSDK +import asyncio from enum import Enum import httpx import json @@ -12,7 +13,10 @@ 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_EXTENSION_KEY +from unstructured_client._hooks.custom.request_utils import ( + ELEMENTS_FILE_EXTENSION_KEY, + NDJSON_MEDIA_TYPE, +) from unstructured_client.utils.unmarshal_json_response import unmarshal_json_response @@ -33,8 +37,16 @@ def _new_elements_file(): ) -def _ndjson_requested(accept_header_override: Optional[PartitionAcceptEnum]) -> bool: - return accept_header_override == PartitionAcceptEnum.APPLICATION_X_NDJSON +def _ndjson_requested(request: httpx.Request) -> bool: + """Whether NDJSON was asked for, read off the Accept header actually sent. + + Keyed off the request rather than `accept_header_override` because `http_headers` can + replace Accept as well, and the split-PDF hook already keys off this same header. + Reading only the override would let the split and unsplit paths disagree for the same + caller: a multi-page PDF would come back as `elements_file` while a two-page one came + back as `elements`. + """ + return NDJSON_MEDIA_TYPE in request.headers.get("Accept", "") def _json_body_to_elements_file(http_res: httpx.Response) -> str: @@ -50,11 +62,17 @@ def _json_body_to_elements_file(http_res: httpx.Response) -> str: and parsing it adds the element list on top. It exists to keep the contract uniform. The memory win comes from the split-PDF path, where the hook concatenates chunk files on disk and this function is never reached. + + The body goes through `unmarshal_json_response` rather than a bare `json()` so that a + payload which is not a list of objects fails the same way it would on the `elements` + path. Iterating raw JSON would crash on `null` and, worse, silently write a dict's + keys out as elements. """ + elements = unmarshal_json_response(Optional[List[Dict[str, Any]]], http_res) out = _new_elements_file() try: with out: - for element in http_res.json(): + for element in elements or []: out.write(json.dumps(element, ensure_ascii=False).encode()) out.write(b"\n") except BaseException: @@ -222,7 +240,7 @@ def partition( raw_response=http_res, ) if utils.match_response(http_res, "200", "application/json"): - if _ndjson_requested(accept_header_override): + if _ndjson_requested(req): return operations.PartitionResponse( elements_file=_json_body_to_elements_file(http_res), status_code=http_res.status_code, @@ -361,9 +379,11 @@ async def partition_async( raw_response=http_res, ) if utils.match_response(http_res, "200", "application/json"): - if _ndjson_requested(accept_header_override): + if _ndjson_requested(req): return operations.PartitionResponse( - elements_file=_json_body_to_elements_file(http_res), + elements_file=await asyncio.to_thread( + _json_body_to_elements_file, http_res + ), status_code=http_res.status_code, content_type=http_res.headers.get("Content-Type") or "", raw_response=http_res, From 140ced000c524858257866ea5098b104aabda6ef Mon Sep 17 00:00:00 2001 From: Yao You Date: Mon, 3 Aug 2026 14:29:57 -0500 Subject: [PATCH 09/11] docs: correct why the unsplit path receives JSON The API does not negotiate response format on Accept at all: it reads the output_format form field, and only consults Accept to select multipart/mixed and to reject conflicting media types on multi-file uploads. Saying the spec merely omits application/x-ndjson understated it -- NDJSON was never going to arrive via Accept regardless of what the spec declared. Comment only; no behavior change. Co-Authored-By: Claude Opus 5 (1M context) --- src/unstructured_client/general.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index 4694c7d7..19aecbda 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -52,11 +52,13 @@ def _ndjson_requested(request: httpx.Request) -> bool: def _json_body_to_elements_file(http_res: httpx.Response) -> str: """Write a JSON-array body out as NDJSON, so `elements_file` is set either way. - The deployed API does not offer `application/x-ndjson`: its spec declares only - `application/json` and `text/csv`, so a request that skips the split-PDF hook comes - back as JSON no matter what was asked for. Without this, opting into NDJSON would - silently populate `elements` instead and leave `elements_file` as None, forcing every - caller to handle both shapes for one flag. + The deployed API does not offer `application/x-ndjson`, and does not in fact negotiate + the response format on `Accept` at all -- it reads the `output_format` form field, and + only consults `Accept` to select `multipart/mixed` and to reject conflicting media + types on multi-file uploads. So a request that skips the split-PDF hook comes back as + JSON no matter what was asked for. Without this, opting into NDJSON would silently + populate `elements` instead and leave `elements_file` as None, forcing every caller to + handle both shapes for one flag. This does not bound memory -- the body is already fully read by the time we get here, and parsing it adds the element list on top. It exists to keep the contract uniform. From 19ca812f4e3817412ba4ec814a2425008becc8a4 Mon Sep 17 00:00:00 2001 From: Yao You Date: Mon, 3 Aug 2026 14:36:12 -0500 Subject: [PATCH 10/11] fix: clean up the async conversion's output file when cancelled Offloading the JSON-to-NDJSON conversion to a thread last commit removed the event loop block but introduced a leak: a thread cannot be cancelled, so on cancellation the conversion still runs to completion and creates its file, while the awaiting coroutine has already raised CancelledError and discarded the path. Nothing was left that could delete it. Reproduced -- cancel mid-conversion and an unst_elements_*.ndjson survives. The conversion is now awaited through asyncio.shield, which keeps a handle on the thread's result after the caller stops waiting, and a done callback unlinks the finished file. A conversion that raised needs no callback, since it already removes its own partial file. This is the same shape as the split hook's cancellation race: work that outlives the operation that requested it, publishing a path nobody will read. It did not exist on the inline version, which had no await point to cancel at. Regression test asserts against the path production actually created, so it cannot pass by observing that no file was ever made. Confirmed to fail against the plain awaited to_thread. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_elements_file.py | 40 +++++++++++++++++++ src/unstructured_client/general.py | 31 ++++++++++++-- 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index 43d53dfa..9358fe81 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -9,11 +9,13 @@ - leave no temp files behind other than the combined file the caller owns """ +import asyncio import errno import json import os import tempfile import threading +import time from pathlib import Path from unittest import mock @@ -752,3 +754,41 @@ def _spy(http_res): assert _read_ndjson(res.elements_file) == elements finally: Path(res.elements_file).unlink(missing_ok=True) + + +@pytest.mark.asyncio +async def test_async_conversion_cleans_up_when_cancelled(tmp_path, monkeypatch): + """A cancelled conversion must not orphan the file its thread went on to create. + + Regression guard for the cost of the offload: a thread cannot be cancelled, so the + conversion runs to completion regardless, and a plain `await asyncio.to_thread(...)` + discards the path it returned -- leaving nothing that could delete the file. + """ + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + created = _record_created_paths(monkeypatch, general, "_new_elements_file") + real_convert = general._json_body_to_elements_file + + def _slow_convert(http_res): + time.sleep(0.3) + return real_convert(http_res) + + monkeypatch.setattr(general, "_json_body_to_elements_file", _slow_convert) + response = httpx.Response( + 200, headers={"Content-Type": "application/json"}, json=[{"type": "Table"}] + ) + + task = asyncio.ensure_future(general._json_body_to_elements_file_async(response)) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # Let the shielded thread finish and its cleanup callback run. + for _ in range(200): + await asyncio.sleep(0.02) + if created and not os.path.exists(created[0]): + break + + # Non-vacuous: a file really was created, and it is now gone. + assert len(created) == 1 + assert not os.path.exists(created[0]) diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index 19aecbda..c2372c4c 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -83,6 +83,33 @@ def _json_body_to_elements_file(http_res: httpx.Response) -> str: return out.name +async def _json_body_to_elements_file_async(http_res: httpx.Response) -> str: + """Run the conversion off the event loop, cleaning up if the caller gives up. + + The work is offloaded so it does not block the loop, but a thread cannot be cancelled: + on cancellation the conversion still runs to completion and creates its file, and a + plain `await asyncio.to_thread(...)` discards the returned path, so nothing is left + that could delete it. Shielding keeps a handle on the thread's result so the finished + file can be removed once the caller has stopped waiting for it. + """ + task = asyncio.ensure_future(asyncio.to_thread(_json_body_to_elements_file, http_res)) + try: + return await asyncio.shield(task) + except BaseException: + task.add_done_callback(_discard_abandoned_elements_file) + raise + + +def _discard_abandoned_elements_file(task: "asyncio.Future[str]") -> None: + """Delete the file a conversion produced after its caller stopped waiting.""" + if task.cancelled(): + return + if task.exception() is not None: + # The conversion raised, and it already removed its own partial file. + return + _discard_elements_file(task.result()) + + def _discard_elements_file(path: str) -> None: """Remove a partially written elements file, ignoring a failure to do so. @@ -383,9 +410,7 @@ async def partition_async( if utils.match_response(http_res, "200", "application/json"): if _ndjson_requested(req): return operations.PartitionResponse( - elements_file=await asyncio.to_thread( - _json_body_to_elements_file, http_res - ), + elements_file=await _json_body_to_elements_file_async(http_res), status_code=http_res.status_code, content_type=http_res.headers.get("Content-Type") or "", raw_response=http_res, From 735b84085053ce86626feea873768dfac4be17ed Mon Sep 17 00:00:00 2001 From: Yao You Date: Mon, 3 Aug 2026 14:50:39 -0500 Subject: [PATCH 11/11] test: sequence the cancellation test with events, not sleeps The test raced. It relied on a fixed `time.sleep(0.3)` in the worker outlasting an `await asyncio.sleep(0.05)` before `task.cancel()`. On a loaded runner the conversion could finish first, leaving nothing to cancel, and the test would fail with "DID NOT RAISE" -- a spurious failure rather than a real one. Now sequenced with three `threading.Event` handshakes: the worker signals that it has entered the conversion, the test cancels only after that, and only then releases the worker, so completion is necessarily post-cancellation. A spy on `_discard_elements_file` signals that cleanup ran, so the assertion waits on the actual event rather than polling the filesystem on a timer. The `wait` timeouts are deadlock guards, never a duration anything waits out. Also drops the 0.02s x 200 polling loop, so the test now takes ~0.15s instead of ~0.5s, and reports "cleanup never ran" instead of a bare assertion when it fails. Still confirmed to fail against the plain awaited to_thread. Ran 20x clean, and 10x clean under eight busy cores. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_elements_file.py | 38 ++++++++++++++----- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_elements_file.py b/_test_unstructured_client/unit/test_ndjson_elements_file.py index 9358fe81..5aac271d 100644 --- a/_test_unstructured_client/unit/test_ndjson_elements_file.py +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -15,7 +15,6 @@ import os import tempfile import threading -import time from pathlib import Path from unittest import mock @@ -40,6 +39,9 @@ ) from unstructured_client.models import errors, operations, shared +# Deadlock guard for the event handshakes below, never a value the tests wait out. +TIMEOUT = 10 + def _elements(prefix, count): return [ @@ -763,31 +765,47 @@ async def test_async_conversion_cleans_up_when_cancelled(tmp_path, monkeypatch): Regression guard for the cost of the offload: a thread cannot be cancelled, so the conversion runs to completion regardless, and a plain `await asyncio.to_thread(...)` discards the path it returned -- leaving nothing that could delete the file. + + Sequenced with events rather than sleeps. Timing the cancellation against a fixed + sleep would race on a loaded runner: if the conversion finished first there would be + nothing to cancel and the test would fail spuriously. The timeouts here are only + deadlocks guards, never the thing being waited on. """ monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) created = _record_created_paths(monkeypatch, general, "_new_elements_file") + + conversion_started = threading.Event() + allow_conversion = threading.Event() + cleanup_done = threading.Event() real_convert = general._json_body_to_elements_file + real_discard = general._discard_elements_file - def _slow_convert(http_res): - time.sleep(0.3) + def _blocked_convert(http_res): + conversion_started.set() + assert allow_conversion.wait(TIMEOUT), "test never released the conversion" return real_convert(http_res) - monkeypatch.setattr(general, "_json_body_to_elements_file", _slow_convert) + def _observed_discard(path): + real_discard(path) + cleanup_done.set() + + monkeypatch.setattr(general, "_json_body_to_elements_file", _blocked_convert) + monkeypatch.setattr(general, "_discard_elements_file", _observed_discard) response = httpx.Response( 200, headers={"Content-Type": "application/json"}, json=[{"type": "Table"}] ) task = asyncio.ensure_future(general._json_body_to_elements_file_async(response)) - await asyncio.sleep(0.05) + + # Waiting in a worker thread keeps the event loop free to run the task. + assert await asyncio.to_thread(conversion_started.wait, TIMEOUT), "conversion never ran" task.cancel() with pytest.raises(asyncio.CancelledError): await task - # Let the shielded thread finish and its cleanup callback run. - for _ in range(200): - await asyncio.sleep(0.02) - if created and not os.path.exists(created[0]): - break + # Only now let the conversion complete, so it necessarily finishes post-cancellation. + allow_conversion.set() + assert await asyncio.to_thread(cleanup_done.wait, TIMEOUT), "cleanup never ran" # Non-vacuous: a file really was created, and it is now gone. assert len(created) == 1