1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
|
from __future__ import annotations
import asyncio
import copy
import io
import json
import logging
from typing import Optional, Tuple, Any
import httpx
import requests
from requests.structures import CaseInsensitiveDict
from requests_toolbelt.multipart.encoder import MultipartEncoder
from unstructured_client._hooks.custom.common import UNSTRUCTURED_CLIENT_LOGGER_NAME
from unstructured_client._hooks.custom.form_utils import (
PARTITION_FORM_FILES_KEY,
PARTITION_FORM_SPLIT_PDF_PAGE_KEY,
PARTITION_FORM_SPLIT_PDF_ALLOW_FAILED_KEY,
PARTITION_FORM_PAGE_RANGE_KEY,
PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
FormData,
)
logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
def create_request_body(
form_data: FormData, page_content: io.BytesIO, filename: str, page_number: int
) -> MultipartEncoder:
payload = prepare_request_payload(form_data)
payload_fields: list[tuple[str, Any]] = []
for key, value in payload.items():
if isinstance(value, list):
payload_fields.extend([(key, list_value) for list_value in value])
else:
payload_fields.append((key, value))
payload_fields.append((PARTITION_FORM_FILES_KEY, (
filename,
page_content,
"application/pdf",
)))
payload_fields.append((PARTITION_FORM_STARTING_PAGE_NUMBER_KEY, str(page_number)))
body = MultipartEncoder(
fields=payload_fields
)
return body
def create_httpx_request(
original_request: requests.Request, body: MultipartEncoder
) -> httpx.Request:
headers = prepare_request_headers(original_request.headers)
return httpx.Request(
method="POST",
url=original_request.url or "",
content=body.to_string(),
headers={**headers, "Content-Type": body.content_type},
)
def create_request(
request: requests.PreparedRequest,
body: MultipartEncoder,
) -> requests.Request:
headers = prepare_request_headers(request.headers)
return requests.Request(
method="POST",
url=request.url or "",
data=body,
headers={**headers, "Content-Type": body.content_type},
)
async def call_api_async(
client: httpx.AsyncClient,
page: Tuple[io.BytesIO, int],
original_request: requests.Request,
form_data: FormData,
filename: str,
limiter: asyncio.Semaphore,
) -> tuple[int, dict]:
page_content, page_number = page
body = create_request_body(form_data, page_content, filename, page_number)
new_request = create_httpx_request(original_request, body)
async with limiter:
try:
response = await client.send(new_request)
return response.status_code, response.json()
except Exception:
logger.error("Failed to send request for page %d", page_number)
return 500, {}
def call_api(
client: Optional[requests.Session],
page: Tuple[io.BytesIO, int],
request: requests.PreparedRequest,
form_data: FormData,
filename: str,
) -> requests.Response:
if client is None:
raise RuntimeError("HTTP client not accessible!")
page_content, page_number = page
body = create_request_body(form_data, page_content, filename, page_number)
new_request = create_request(request, body)
prepared_request = client.prepare_request(new_request)
try:
return client.send(prepared_request)
except Exception:
logger.error("Failed to send request for page %d", page_number)
return requests.Response()
def prepare_request_headers(
headers: CaseInsensitiveDict[str],
) -> CaseInsensitiveDict[str]:
"""Prepare the request headers by removing the 'Content-Type' and 'Content-Length' headers.
Args:
headers: The original request headers.
Returns:
The modified request headers.
"""
headers = copy.deepcopy(headers)
headers.pop("Content-Type", None)
headers.pop("Content-Length", None)
return headers
def prepare_request_payload(form_data: FormData) -> FormData:
"""Prepares the request payload by removing unnecessary keys and updating the file.
Args:
form_data: The original form data.
Returns:
The updated request payload.
"""
payload = copy.deepcopy(form_data)
payload.pop(PARTITION_FORM_SPLIT_PDF_PAGE_KEY, None)
payload.pop(PARTITION_FORM_SPLIT_PDF_ALLOW_FAILED_KEY, None)
payload.pop(PARTITION_FORM_FILES_KEY, None)
payload.pop(PARTITION_FORM_PAGE_RANGE_KEY, None)
payload.pop(PARTITION_FORM_STARTING_PAGE_NUMBER_KEY, None)
updated_parameters = {
PARTITION_FORM_SPLIT_PDF_PAGE_KEY: "false",
}
payload.update(updated_parameters)
return payload
def create_response(response: requests.Response, elements: list) -> requests.Response:
"""
Creates a modified response object with updated content.
Args:
response: The original response object.
elements: The list of elements to be serialized and added to
the response.
Returns:
The modified response object with updated content.
"""
response_copy = copy.deepcopy(response)
content = json.dumps(elements).encode()
content_length = str(len(content))
response_copy.headers.update({"Content-Length": content_length})
setattr(response_copy, "_content", content)
return response_copy
def log_after_split_response(status_code: int, split_number: int):
if status_code == 200:
logger.info(
"Successfully partitioned set #%d, elements added to the final result.",
split_number,
)
else:
logger.warning(
"Failed to partition set #%d, its elements will be omitted in the final result.",
split_number,
)
|
