Skip to content

Commit d10d877

Browse files
committed
add text/csv handling to the python sdk
1 parent 529fe87 commit d10d877

File tree

2 files changed

+17
-5
lines changed

2 files changed

+17
-5
lines changed

src/unstructured_client/_hooks/custom/request_utils.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,12 @@ def create_response(elements: list) -> httpx.Response:
150150
Returns:
151151
The modified response object with updated content.
152152
"""
153-
response = httpx.Response(status_code=200, headers={"Content-Type": "application/json"})
154-
content = json.dumps(elements).encode()
153+
if not isinstance(elements[0], dict):
154+
response = httpx.Response(status_code=200, headers={"Content-Type": "text/csv"})
155+
content = b''.join(elements)
156+
else:
157+
response = httpx.Response(status_code=200, headers={"Content-Type": "application/json"})
158+
content = json.dumps(elements).encode()
155159
content_length = str(len(content))
156160
response.headers.update({"Content-Length": content_length})
157161
setattr(response, "_content", content)

src/unstructured_client/_hooks/custom/split_pdf_hook.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,10 @@ def _await_elements(
389389
response_number,
390390
)
391391
successful_responses.append(res)
392-
elements.append(res.json())
392+
if res.headers["Content-Type"] == "application/json":
393+
elements.append(res.json())
394+
else:
395+
elements.append(res.content)
393396
else:
394397
error_message = f"Failed to partition set {response_number}."
395398

@@ -401,7 +404,12 @@ def _await_elements(
401404

402405
self.api_successful_responses[operation_id] = successful_responses
403406
self.api_failed_responses[operation_id] = failed_responses
404-
flattened_elements = [element for sublist in elements for element in sublist]
407+
flattened_elements = []
408+
for sublist in elements:
409+
if isinstance(sublist, list):
410+
flattened_elements.extend(sublist)
411+
else:
412+
flattened_elements.append(sublist)
405413
return flattened_elements
406414

407415
def after_success(
@@ -423,7 +431,7 @@ def after_success(
423431
"""
424432
# Grab the correct id out of the dummy request
425433
operation_id = response.request.headers.get("operation_id")
426-
434+
# breakpoint();
427435
elements = self._await_elements(operation_id)
428436

429437
# if fails are disallowed, return the first failed response

0 commit comments

Comments
 (0)