diff --git a/.genignore b/.genignore index ea1fba41..83e366df 100644 --- a/.genignore +++ b/.genignore @@ -23,3 +23,14 @@ 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 + +# 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/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/README.md b/README.md index 17bfc6f4..ef00b584 100644 --- a/README.md +++ b/README.md @@ -427,6 +427,37 @@ 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.** + +> [!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 +from pathlib import Path + +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: + # missing_ok so a failure to open the file isn't masked by the cleanup. + Path(res.elements_file).unlink(missing_ok=True) +``` + ## File uploads 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/_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..22b178a7 --- /dev/null +++ b/_test_unstructured_client/unit/test_ndjson_elements_file.py @@ -0,0 +1,515 @@ +"""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 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, + 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, + _ndjson_elements_file_async, +) + + +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_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) + 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 chunk response must no longer hold the body. + + 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. + """ + 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 + + hook._elements_from_task_responses(operation_id, [(0, response)], started_at=0.0) + + # 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): + path = str(tmp_path / "combined.ndjson") + response = create_elements_file_response(path) + + assert response.status_code == 200 + 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 ------------------------------------------------ + + +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 + # Marks the operation live; `_clear_operation` removing it is what signals teardown. + hook.pending_operation_ids[operation_id] = operation_id + 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.extensions[ELEMENTS_FILE_EXTENSION_KEY] + 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_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. + + 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) + 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..f465e1d6 100644 --- a/_test_unstructured_client/unit/test_regeneration_guards.py +++ b/_test_unstructured_client/unit/test_regeneration_guards.py @@ -81,6 +81,33 @@ 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", + "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( 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..215983a8 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,129 @@ def create_response(elements: list) -> httpx.Response: return response +# 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 + + +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()) + 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 + + +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 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. + + 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)), + }, + extensions={ELEMENTS_FILE_EXTENSION_KEY: 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..9e31b3b3 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,15 @@ 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 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]: @@ -791,6 +817,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 +1417,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 +1430,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 +1493,96 @@ 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`). + + 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. + """ + temp_dir_path = self.cache_tmp_data_dir.get(operation_id) or tempfile.gettempdir() + out_path = f"{temp_dir_path}/{uuid.uuid4()}.ndjson" + 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 + + 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, + 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. + """ + 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, @@ -1451,6 +1598,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 +1610,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._claim_ndjson_output(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,7 +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.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/_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: diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index a6ec3099..4f1be006 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -2,18 +2,89 @@ 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 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_EXTENSION_KEY 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 _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. + + When the split-PDF hook ran it has already combined the per-chunk temp files into one + 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.extensions.get(ELEMENTS_FILE_EXTENSION_KEY) + if existing_path: + return existing_path + + 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: + """Async counterpart of `_ndjson_elements_file`.""" + existing_path = http_res.extensions.get(ELEMENTS_FILE_EXTENSION_KEY) + if existing_path: + return existing_path + + 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): @@ -128,6 +199,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 +331,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."""