feat: NDJSON elements-file mode for partition (0.46.0)

Merged
BA
3mo ago
+1397-6

What

Adds an opt-in NDJSON response mode to partition() that returns elements as a path to a file on disk instead of a parsed list, and ships it as 0.46.0.

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)

PartitionResponse.elements_file is set instead of PartitionResponse.elements. The caller owns the file and must delete it. Requesting application/json remains the default and is entirely unchanged.

Why

On the split-PDF path the SDK rebuilt the whole document in memory in order to return it: a list per chunk, a flattened list, a json.dumps blob in create_response, and then the SDK's re-parse of that blob — four copies live at once, with the serialization step dominating peak usage. For documents with large metadata.image_base64 payloads this is the difference between a job completing and being OOM-killed.

In the new mode the per-chunk temp files are concatenated on disk and never parsed, so peak memory is roughly one chunk rather than the whole document.

How

  • combine_chunk_files_to_ndjson concatenates chunk files on disk. Each chunk is sniffed for its first non-whitespace character, so a server returning application/json still works; chunks that are already NDJSON are copied through without parsing.
  • 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.
  • The combined file is deliberately written outside the operation's TemporaryDirectory, which _clear_operation removes as soon as after_success returns.

Temp-file ownership

Everything this path creates is accounted for:

  • Spilled chunk bodies are written inside the operation's temp directory and unlinked once combined.
  • The combined file is deleted when a chunk failure means it is never handed back to the caller.
  • Recombination writes to a staging file that is atomically renamed into place only on success, so a malformed chunk cannot orphan a partial file.
  • No combined file is created at all when every chunk failed.

Security

The elements-file marker is an httpx response extension, not a response header. Extensions are populated by the transport, so a remote server cannot set the key. A header would be wire-controlled, and since callers are documented to open elements_file and then delete it, that would hand a hostile server an arbitrary local file to destroy. A real server body is always copied to a file this client creates.

Regeneration

elements_file is client-side only and can never come from the OpenAPI spec, so a regeneration would silently drop it. Both general.py and models/operations/partition.py are now in .genignore, and test_regeneration_guards.py fails if either entry is lost.

Known limitation

elements_file is set for every input, so callers need one code path. The memory saving, however, applies only to split PDFs.

An input is sent whole when it is not a PDF, when split_pdf_page=False, or when it has two pages or fewer — _before_request_unlocked short-circuits on split_size >= page_count and get_optimal_split_size floors at MIN_PAGES_PER_SPLIT = 2. For those, the body is read fully into memory before being written to disk, so peak is roughly 2x the body rather than bounded.

Bounding it means stream=True for NDJSON requests, which makes raw_response.content raise on the returned closed response — a user-visible change worth its own review. Tracked separately.

Note also that the deployed API does not currently emit application/x-ndjson, so the unsplit path reaches the JSON-to-NDJSON conversion rather than the streamed-body branch. That is not merely a spec omission: the service does not negotiate the response format on Accept at all. It selects the format from the output_format form field, and consults Accept only to choose multipart/mixed and to reject conflicting media types on multi-file uploads. NDJSON was therefore never going to arrive via Accept.

The service's 406 NOT_ACCEPTABLE on an unrecognized Accept is gated on multi-file uploads. This SDK sends a single file per request — PartitionParameters.files is one Files, and the split-PDF hook sends one chunk per request — so that branch is unreachable from here and the unsplit path cannot raise SDKError because of it. Server-side NDJSON support is tracked separately.

Testing

  • New _test_unstructured_client/unit/test_ndjson_elements_file.py — recombination across JSON-array / NDJSON / mixed chunk formats, order preservation, byte-exact payload round-trip, non-ASCII, temp-file lifecycle on success and failure, and regression guards for the header-spoofing and partial-output defects.
  • 235 unit tests and 64 contract tests pass; pylint 10.00/10; mypy clean.

🤖 Generated with Claude Code