futurehouse-client 0.3.18.dev186__py3-none-any.whl → 0.3.18.dev195__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- futurehouse_client/clients/rest_client.py +39 -148
- futurehouse_client/utils/auth.py +83 -101
- {futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/METADATA +1 -1
- {futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/RECORD +6 -6
- {futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/WHEEL +0 -0
- {futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/top_level.txt +0 -0
@@ -15,7 +15,7 @@ import uuid
|
|
15
15
|
from collections.abc import Collection
|
16
16
|
from pathlib import Path
|
17
17
|
from types import ModuleType
|
18
|
-
from typing import Any, ClassVar,
|
18
|
+
from typing import Any, ClassVar, cast
|
19
19
|
from uuid import UUID
|
20
20
|
|
21
21
|
import cloudpickle
|
@@ -45,7 +45,6 @@ from tqdm.asyncio import tqdm
|
|
45
45
|
|
46
46
|
from futurehouse_client.clients import JobNames
|
47
47
|
from futurehouse_client.models.app import (
|
48
|
-
APIKeyPayload,
|
49
48
|
AuthType,
|
50
49
|
JobDeploymentConfig,
|
51
50
|
PQATaskResponse,
|
@@ -55,11 +54,7 @@ from futurehouse_client.models.app import (
|
|
55
54
|
TaskResponseVerbose,
|
56
55
|
)
|
57
56
|
from futurehouse_client.models.rest import ExecutionStatus
|
58
|
-
from futurehouse_client.utils.auth import
|
59
|
-
AUTH_ERRORS_TO_RETRY_ON,
|
60
|
-
AuthError,
|
61
|
-
refresh_token_on_auth_error,
|
62
|
-
)
|
57
|
+
from futurehouse_client.utils.auth import RefreshingJWT
|
63
58
|
from futurehouse_client.utils.general import gather_with_concurrency
|
64
59
|
from futurehouse_client.utils.module_utils import (
|
65
60
|
OrganizationSelector,
|
@@ -128,8 +123,6 @@ retry_if_connection_error = retry_if_exception_type((
|
|
128
123
|
FileUploadError,
|
129
124
|
))
|
130
125
|
|
131
|
-
# 5 minute default for JWTs
|
132
|
-
JWT_TOKEN_CACHE_EXPIRY: int = 300 # seconds
|
133
126
|
DEFAULT_AGENT_TIMEOUT: int = 2400 # seconds
|
134
127
|
|
135
128
|
|
@@ -163,69 +156,85 @@ class RestClient:
|
|
163
156
|
self.api_key = api_key
|
164
157
|
self._clients: dict[str, Client | AsyncClient] = {}
|
165
158
|
self.headers = headers or {}
|
166
|
-
self.
|
159
|
+
self.jwt = jwt
|
167
160
|
self.organizations: list[str] = self._filter_orgs(organization)
|
168
161
|
|
169
162
|
@property
|
170
163
|
def client(self) -> Client:
|
171
|
-
"""
|
172
|
-
return cast(Client, self.get_client("application/json",
|
164
|
+
"""Authenticated HTTP client for regular API calls."""
|
165
|
+
return cast(Client, self.get_client("application/json", authenticated=True))
|
173
166
|
|
174
167
|
@property
|
175
168
|
def async_client(self) -> AsyncClient:
|
176
|
-
"""
|
169
|
+
"""Authenticated async HTTP client for regular API calls."""
|
177
170
|
return cast(
|
178
171
|
AsyncClient,
|
179
|
-
self.get_client("application/json",
|
172
|
+
self.get_client("application/json", authenticated=True, async_client=True),
|
180
173
|
)
|
181
174
|
|
182
175
|
@property
|
183
|
-
def
|
184
|
-
"""
|
185
|
-
return cast(Client, self.get_client("application/json",
|
176
|
+
def unauthenticated_client(self) -> Client:
|
177
|
+
"""Unauthenticated HTTP client for auth operations to avoid recursion."""
|
178
|
+
return cast(Client, self.get_client("application/json", authenticated=False))
|
186
179
|
|
187
180
|
@property
|
188
181
|
def multipart_client(self) -> Client:
|
189
|
-
"""
|
190
|
-
return cast(Client, self.get_client(None,
|
182
|
+
"""Authenticated HTTP client for multipart uploads."""
|
183
|
+
return cast(Client, self.get_client(None, authenticated=True))
|
191
184
|
|
192
185
|
def get_client(
|
193
186
|
self,
|
194
187
|
content_type: str | None = "application/json",
|
195
|
-
|
196
|
-
|
188
|
+
authenticated: bool = True,
|
189
|
+
async_client: bool = False,
|
197
190
|
) -> Client | AsyncClient:
|
198
191
|
"""Return a cached HTTP client or create one if needed.
|
199
192
|
|
200
193
|
Args:
|
201
194
|
content_type: The desired content type header. Use None for multipart uploads.
|
202
|
-
|
203
|
-
|
195
|
+
authenticated: Whether the client should include authentication.
|
196
|
+
async_client: Whether to use an async client.
|
204
197
|
|
205
198
|
Returns:
|
206
199
|
An HTTP client configured with the appropriate headers.
|
207
200
|
"""
|
208
|
-
# Create a composite key based on content type and auth flag
|
209
|
-
key = f"{content_type or 'multipart'}_{
|
201
|
+
# Create a composite key based on content type and auth flag
|
202
|
+
key = f"{content_type or 'multipart'}_{authenticated}_{async_client}"
|
203
|
+
|
210
204
|
if key not in self._clients:
|
211
205
|
headers = copy.deepcopy(self.headers)
|
212
|
-
|
213
|
-
|
206
|
+
auth = None
|
207
|
+
|
208
|
+
if authenticated:
|
209
|
+
auth = RefreshingJWT(
|
210
|
+
# authenticated=False will always return a synchronous client
|
211
|
+
auth_client=cast(
|
212
|
+
Client, self.get_client("application/json", authenticated=False)
|
213
|
+
),
|
214
|
+
auth_type=self.auth_type,
|
215
|
+
api_key=self.api_key,
|
216
|
+
jwt=self.jwt,
|
217
|
+
)
|
218
|
+
|
214
219
|
if content_type:
|
215
220
|
headers["Content-Type"] = content_type
|
221
|
+
|
216
222
|
self._clients[key] = (
|
217
223
|
AsyncClient(
|
218
224
|
base_url=self.base_url,
|
219
225
|
headers=headers,
|
220
226
|
timeout=self.REQUEST_TIMEOUT,
|
227
|
+
auth=auth,
|
221
228
|
)
|
222
|
-
if
|
229
|
+
if async_client
|
223
230
|
else Client(
|
224
231
|
base_url=self.base_url,
|
225
232
|
headers=headers,
|
226
233
|
timeout=self.REQUEST_TIMEOUT,
|
234
|
+
auth=auth,
|
227
235
|
)
|
228
236
|
)
|
237
|
+
|
229
238
|
return self._clients[key]
|
230
239
|
|
231
240
|
def close(self):
|
@@ -255,32 +264,6 @@ class RestClient:
|
|
255
264
|
raise ValueError(f"Organization '{organization}' not found.")
|
256
265
|
return filtered_orgs
|
257
266
|
|
258
|
-
def _run_auth(self, jwt: str | None = None) -> str:
|
259
|
-
auth_payload: APIKeyPayload | None
|
260
|
-
if self.auth_type == AuthType.API_KEY:
|
261
|
-
auth_payload = APIKeyPayload(api_key=self.api_key)
|
262
|
-
elif self.auth_type == AuthType.JWT:
|
263
|
-
auth_payload = None
|
264
|
-
else:
|
265
|
-
assert_never(self.auth_type)
|
266
|
-
try:
|
267
|
-
# Use the unauthenticated client for login
|
268
|
-
if auth_payload:
|
269
|
-
response = self.auth_client.post(
|
270
|
-
"/auth/login", json=auth_payload.model_dump()
|
271
|
-
)
|
272
|
-
response.raise_for_status()
|
273
|
-
token_data = response.json()
|
274
|
-
elif jwt:
|
275
|
-
token_data = {"access_token": jwt, "expires_in": JWT_TOKEN_CACHE_EXPIRY}
|
276
|
-
else:
|
277
|
-
raise ValueError("JWT token required for JWT authentication.")
|
278
|
-
|
279
|
-
return token_data["access_token"]
|
280
|
-
except Exception as e:
|
281
|
-
raise RestClientError(f"Error authenticating: {e!s}") from e
|
282
|
-
|
283
|
-
@refresh_token_on_auth_error()
|
284
267
|
def _check_job(self, name: str, organization: str) -> dict[str, Any]:
|
285
268
|
try:
|
286
269
|
response = self.client.get(
|
@@ -288,19 +271,9 @@ class RestClient:
|
|
288
271
|
)
|
289
272
|
response.raise_for_status()
|
290
273
|
return response.json()
|
291
|
-
except HTTPStatusError as e:
|
292
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
293
|
-
raise AuthError(
|
294
|
-
e.response.status_code,
|
295
|
-
f"Authentication failed: {e}",
|
296
|
-
request=e.request,
|
297
|
-
response=e.response,
|
298
|
-
) from e
|
299
|
-
raise
|
300
274
|
except Exception as e:
|
301
275
|
raise JobFetchError(f"Error checking job: {e!s}") from e
|
302
276
|
|
303
|
-
@refresh_token_on_auth_error()
|
304
277
|
def _fetch_my_orgs(self) -> list[str]:
|
305
278
|
response = self.client.get(f"/v0.1/organizations?filter={True}")
|
306
279
|
response.raise_for_status()
|
@@ -358,7 +331,6 @@ class RestClient:
|
|
358
331
|
if not files:
|
359
332
|
raise TaskFetchError(f"No files found in {path}")
|
360
333
|
|
361
|
-
@refresh_token_on_auth_error()
|
362
334
|
@retry(
|
363
335
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
364
336
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -400,19 +372,9 @@ class RestClient:
|
|
400
372
|
):
|
401
373
|
return PQATaskResponse(**data)
|
402
374
|
return TaskResponse(**data)
|
403
|
-
except HTTPStatusError as e:
|
404
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
405
|
-
raise AuthError(
|
406
|
-
e.response.status_code,
|
407
|
-
f"Authentication failed: {e}",
|
408
|
-
request=e.request,
|
409
|
-
response=e.response,
|
410
|
-
) from e
|
411
|
-
raise
|
412
375
|
except Exception as e:
|
413
376
|
raise TaskFetchError(f"Error getting task: {e!s}") from e
|
414
377
|
|
415
|
-
@refresh_token_on_auth_error()
|
416
378
|
@retry(
|
417
379
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
418
380
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -456,19 +418,9 @@ class RestClient:
|
|
456
418
|
):
|
457
419
|
return PQATaskResponse(**data)
|
458
420
|
return TaskResponse(**data)
|
459
|
-
except HTTPStatusError as e:
|
460
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
461
|
-
raise AuthError(
|
462
|
-
e.response.status_code,
|
463
|
-
f"Authentication failed: {e}",
|
464
|
-
request=e.request,
|
465
|
-
response=e.response,
|
466
|
-
) from e
|
467
|
-
raise
|
468
421
|
except Exception as e:
|
469
422
|
raise TaskFetchError(f"Error getting task: {e!s}") from e
|
470
423
|
|
471
|
-
@refresh_token_on_auth_error()
|
472
424
|
@retry(
|
473
425
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
474
426
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -492,20 +444,10 @@ class RestClient:
|
|
492
444
|
response.raise_for_status()
|
493
445
|
trajectory_id = response.json()["trajectory_id"]
|
494
446
|
self.trajectory_id = trajectory_id
|
495
|
-
except HTTPStatusError as e:
|
496
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
497
|
-
raise AuthError(
|
498
|
-
e.response.status_code,
|
499
|
-
f"Authentication failed: {e}",
|
500
|
-
request=e.request,
|
501
|
-
response=e.response,
|
502
|
-
) from e
|
503
|
-
raise
|
504
447
|
except Exception as e:
|
505
448
|
raise TaskFetchError(f"Error creating task: {e!s}") from e
|
506
449
|
return trajectory_id
|
507
450
|
|
508
|
-
@refresh_token_on_auth_error()
|
509
451
|
@retry(
|
510
452
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
511
453
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -529,15 +471,6 @@ class RestClient:
|
|
529
471
|
response.raise_for_status()
|
530
472
|
trajectory_id = response.json()["trajectory_id"]
|
531
473
|
self.trajectory_id = trajectory_id
|
532
|
-
except HTTPStatusError as e:
|
533
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
534
|
-
raise AuthError(
|
535
|
-
e.response.status_code,
|
536
|
-
f"Authentication failed: {e}",
|
537
|
-
request=e.request,
|
538
|
-
response=e.response,
|
539
|
-
) from e
|
540
|
-
raise
|
541
474
|
except Exception as e:
|
542
475
|
raise TaskFetchError(f"Error creating task: {e!s}") from e
|
543
476
|
return trajectory_id
|
@@ -683,7 +616,6 @@ class RestClient:
|
|
683
616
|
for task_id in trajectory_ids
|
684
617
|
]
|
685
618
|
|
686
|
-
@refresh_token_on_auth_error()
|
687
619
|
@retry(
|
688
620
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
689
621
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -695,19 +627,11 @@ class RestClient:
|
|
695
627
|
build_id = build_id or self.build_id
|
696
628
|
response = self.client.get(f"/v0.1/builds/{build_id}")
|
697
629
|
response.raise_for_status()
|
698
|
-
except
|
699
|
-
|
700
|
-
raise AuthError(
|
701
|
-
e.response.status_code,
|
702
|
-
f"Authentication failed: {e}",
|
703
|
-
request=e.request,
|
704
|
-
response=e.response,
|
705
|
-
) from e
|
706
|
-
raise
|
630
|
+
except Exception as e:
|
631
|
+
raise JobFetchError(f"Error getting build status: {e!s}") from e
|
707
632
|
return response.json()
|
708
633
|
|
709
634
|
# TODO: Refactor later so we don't have to ignore PLR0915
|
710
|
-
@refresh_token_on_auth_error()
|
711
635
|
@retry(
|
712
636
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
713
637
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -887,13 +811,6 @@ class RestClient:
|
|
887
811
|
build_context = response.json()
|
888
812
|
self.build_id = build_context["build_id"]
|
889
813
|
except HTTPStatusError as e:
|
890
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
891
|
-
raise AuthError(
|
892
|
-
e.response.status_code,
|
893
|
-
f"Authentication failed: {e}",
|
894
|
-
request=e.request,
|
895
|
-
response=e.response,
|
896
|
-
) from e
|
897
814
|
error_detail = response.json()
|
898
815
|
error_message = error_detail.get("detail", str(e))
|
899
816
|
raise JobCreationError(
|
@@ -974,7 +891,6 @@ class RestClient:
|
|
974
891
|
except Exception as e:
|
975
892
|
raise FileUploadError(f"Error uploading directory {dir_path}: {e}") from e
|
976
893
|
|
977
|
-
@refresh_token_on_auth_error()
|
978
894
|
def _upload_single_file(
|
979
895
|
self,
|
980
896
|
job_name: str,
|
@@ -1048,20 +964,10 @@ class RestClient:
|
|
1048
964
|
)
|
1049
965
|
|
1050
966
|
logger.info(f"Successfully uploaded {file_name}")
|
1051
|
-
except HTTPStatusError as e:
|
1052
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
1053
|
-
raise AuthError(
|
1054
|
-
e.response.status_code,
|
1055
|
-
f"Authentication failed: {e}",
|
1056
|
-
request=e.request,
|
1057
|
-
response=e.response,
|
1058
|
-
) from e
|
1059
|
-
raise
|
1060
967
|
except Exception as e:
|
1061
968
|
logger.exception(f"Error uploading file {file_path}")
|
1062
969
|
raise FileUploadError(f"Error uploading file {file_path}: {e}") from e
|
1063
970
|
|
1064
|
-
@refresh_token_on_auth_error()
|
1065
971
|
@retry(
|
1066
972
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
1067
973
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -1098,13 +1004,6 @@ class RestClient:
|
|
1098
1004
|
response.raise_for_status()
|
1099
1005
|
return response.json()
|
1100
1006
|
except HTTPStatusError as e:
|
1101
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
1102
|
-
raise AuthError(
|
1103
|
-
e.response.status_code,
|
1104
|
-
f"Authentication failed: {e}",
|
1105
|
-
request=e.request,
|
1106
|
-
response=e.response,
|
1107
|
-
) from e
|
1108
1007
|
logger.exception(
|
1109
1008
|
f"Error listing files for job {job_name}, trajectory {trajectory_id}, upload_id {upload_id}: {e.response.text}"
|
1110
1009
|
)
|
@@ -1117,7 +1016,6 @@ class RestClient:
|
|
1117
1016
|
)
|
1118
1017
|
raise RestClientError(f"Error listing files: {e!s}") from e
|
1119
1018
|
|
1120
|
-
@refresh_token_on_auth_error()
|
1121
1019
|
@retry(
|
1122
1020
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
1123
1021
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -1165,13 +1063,6 @@ class RestClient:
|
|
1165
1063
|
|
1166
1064
|
logger.info(f"File {file_path} downloaded to {destination_path}")
|
1167
1065
|
except HTTPStatusError as e:
|
1168
|
-
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
1169
|
-
raise AuthError(
|
1170
|
-
e.response.status_code,
|
1171
|
-
f"Authentication failed: {e}",
|
1172
|
-
request=e.request,
|
1173
|
-
response=e.response,
|
1174
|
-
) from e
|
1175
1066
|
logger.exception(
|
1176
1067
|
f"Error downloading file {file_path} for job {job_name}, trajectory_id {trajectory_id}: {e.response.text}"
|
1177
1068
|
)
|
futurehouse_client/utils/auth.py
CHANGED
@@ -1,107 +1,89 @@
|
|
1
|
-
import asyncio
|
2
1
|
import logging
|
3
|
-
from
|
4
|
-
from functools import wraps
|
5
|
-
from typing import Any, Final, Optional, ParamSpec, TypeVar, overload
|
2
|
+
from typing import ClassVar, Final
|
6
3
|
|
7
4
|
import httpx
|
8
|
-
from httpx import HTTPStatusError
|
9
5
|
|
10
|
-
|
11
|
-
|
12
|
-
T = TypeVar("T")
|
13
|
-
P = ParamSpec("P")
|
14
|
-
|
15
|
-
AUTH_ERRORS_TO_RETRY_ON: Final[set[int]] = {
|
16
|
-
httpx.codes.UNAUTHORIZED,
|
17
|
-
httpx.codes.FORBIDDEN,
|
18
|
-
}
|
19
|
-
|
20
|
-
|
21
|
-
class AuthError(Exception):
|
22
|
-
"""Raised when authentication fails with 401/403 status."""
|
23
|
-
|
24
|
-
def __init__(self, status_code: int, message: str, request=None, response=None):
|
25
|
-
self.status_code = status_code
|
26
|
-
self.request = request
|
27
|
-
self.response = response
|
28
|
-
super().__init__(message)
|
29
|
-
|
30
|
-
|
31
|
-
def is_auth_error(e: Exception) -> bool:
|
32
|
-
if isinstance(e, AuthError):
|
33
|
-
return True
|
34
|
-
if isinstance(e, HTTPStatusError):
|
35
|
-
return e.response.status_code in AUTH_ERRORS_TO_RETRY_ON
|
36
|
-
return False
|
37
|
-
|
38
|
-
|
39
|
-
def get_status_code(e: Exception) -> Optional[int]:
|
40
|
-
if isinstance(e, AuthError):
|
41
|
-
return e.status_code
|
42
|
-
if isinstance(e, HTTPStatusError):
|
43
|
-
return e.response.status_code
|
44
|
-
return None
|
6
|
+
from futurehouse_client.models.app import APIKeyPayload, AuthType
|
45
7
|
|
8
|
+
logger = logging.getLogger(__name__)
|
46
9
|
|
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
|
-
|
10
|
+
INVALID_REFRESH_TYPE_MSG: Final[str] = (
|
11
|
+
"API key auth is required to refresh auth tokens."
|
12
|
+
)
|
13
|
+
JWT_TOKEN_CACHE_EXPIRY: int = 300 # seconds
|
14
|
+
|
15
|
+
|
16
|
+
def _run_auth(
|
17
|
+
client: httpx.Client,
|
18
|
+
auth_type: AuthType = AuthType.API_KEY,
|
19
|
+
api_key: str | None = None,
|
20
|
+
jwt: str | None = None,
|
21
|
+
) -> str:
|
22
|
+
auth_payload: APIKeyPayload | None
|
23
|
+
if auth_type == AuthType.API_KEY:
|
24
|
+
auth_payload = APIKeyPayload(api_key=api_key)
|
25
|
+
elif auth_type == AuthType.JWT:
|
26
|
+
auth_payload = None
|
27
|
+
try:
|
28
|
+
if auth_payload:
|
29
|
+
response = client.post("/auth/login", json=auth_payload.model_dump())
|
30
|
+
response.raise_for_status()
|
31
|
+
token_data = response.json()
|
32
|
+
elif jwt:
|
33
|
+
token_data = {"access_token": jwt, "expires_in": JWT_TOKEN_CACHE_EXPIRY}
|
34
|
+
else:
|
35
|
+
raise ValueError("JWT token required for JWT authentication.")
|
36
|
+
|
37
|
+
return token_data["access_token"]
|
38
|
+
except Exception as e:
|
39
|
+
raise Exception("Failed to authenticate") from e # noqa: TRY002
|
40
|
+
|
41
|
+
|
42
|
+
class RefreshingJWT(httpx.Auth):
|
43
|
+
"""Automatically (re-)inject a JWT and transparently retry exactly once when we hit a 401/403."""
|
44
|
+
|
45
|
+
RETRY_STATUSES: ClassVar[set[int]] = {
|
46
|
+
httpx.codes.UNAUTHORIZED,
|
47
|
+
httpx.codes.FORBIDDEN,
|
48
|
+
}
|
49
|
+
|
50
|
+
def __init__(
|
51
|
+
self,
|
52
|
+
auth_client: httpx.Client,
|
53
|
+
auth_type: AuthType = AuthType.API_KEY,
|
54
|
+
api_key: str | None = None,
|
55
|
+
jwt: str | None = None,
|
56
|
+
):
|
57
|
+
self.auth_type = auth_type
|
58
|
+
self.auth_client = auth_client
|
59
|
+
self.api_key = api_key
|
60
|
+
self._jwt = _run_auth(
|
61
|
+
client=auth_client,
|
62
|
+
jwt=jwt,
|
63
|
+
auth_type=auth_type,
|
64
|
+
api_key=api_key,
|
65
|
+
)
|
66
|
+
|
67
|
+
def refresh_token(self):
|
68
|
+
if self.auth_type == AuthType.JWT:
|
69
|
+
logger.error(INVALID_REFRESH_TYPE_MSG)
|
70
|
+
raise ValueError(INVALID_REFRESH_TYPE_MSG)
|
71
|
+
self._jwt = _run_auth(
|
72
|
+
client=self.auth_client,
|
73
|
+
auth_type=self.auth_type,
|
74
|
+
api_key=self.api_key,
|
75
|
+
)
|
76
|
+
|
77
|
+
def auth_flow(self, request):
|
78
|
+
request.headers["Authorization"] = f"Bearer {self._jwt}"
|
79
|
+
response = yield request
|
80
|
+
|
81
|
+
# If it failed, refresh once and replay the request
|
82
|
+
if response.status_code in self.RETRY_STATUSES:
|
83
|
+
logger.info(
|
84
|
+
"Received %s, refreshing token and retrying …",
|
85
|
+
response.status_code,
|
86
|
+
)
|
87
|
+
self.refresh_token()
|
88
|
+
request.headers["Authorization"] = f"Bearer {self._jwt}"
|
89
|
+
yield request # second (and final) attempt, again or use a while loop
|
{futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/METADATA
RENAMED
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: futurehouse-client
|
3
|
-
Version: 0.3.18.
|
3
|
+
Version: 0.3.18.dev195
|
4
4
|
Summary: A client for interacting with endpoints of the FutureHouse service.
|
5
5
|
Author-email: FutureHouse technical staff <hello@futurehouse.org>
|
6
6
|
Classifier: Operating System :: OS Independent
|
{futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/RECORD
RENAMED
@@ -1,17 +1,17 @@
|
|
1
1
|
futurehouse_client/__init__.py,sha256=ddxO7JE97c6bt7LjNglZZ2Ql8bYCGI9laSFeh9MP6VU,344
|
2
2
|
futurehouse_client/clients/__init__.py,sha256=tFWqwIAY5PvwfOVsCje4imjTpf6xXNRMh_UHIKVI1_0,320
|
3
3
|
futurehouse_client/clients/job_client.py,sha256=uNkqQbeZw7wbA0qDWcIOwOykrosza-jev58paJZ_mbA,11150
|
4
|
-
futurehouse_client/clients/rest_client.py,sha256=
|
4
|
+
futurehouse_client/clients/rest_client.py,sha256=6HQF3YXDnSdGxAoXpB_wU6Vhcqhp5OB5SNuGQJ6Hseo,43454
|
5
5
|
futurehouse_client/models/__init__.py,sha256=5x-f9AoM1hGzJBEHcHAXSt7tPeImST5oZLuMdwp0mXc,554
|
6
6
|
futurehouse_client/models/app.py,sha256=w_1e4F0IiC-BKeOLqYkABYo4U-Nka1S-F64S_eHB2KM,26421
|
7
7
|
futurehouse_client/models/client.py,sha256=n4HD0KStKLm6Ek9nL9ylP-bkK10yzAaD1uIDF83Qp_A,1828
|
8
8
|
futurehouse_client/models/rest.py,sha256=lgwkMIXz0af-49BYSkKeS7SRqvN3motqnAikDN4YGTc,789
|
9
9
|
futurehouse_client/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
10
|
-
futurehouse_client/utils/auth.py,sha256=
|
10
|
+
futurehouse_client/utils/auth.py,sha256=0V161S9jW4vbTCoJJrOtNzWXQkAVyzdGM3yefGgJ578,2808
|
11
11
|
futurehouse_client/utils/general.py,sha256=A_rtTiYW30ELGEZlWCIArO7q1nEmqi8hUlmBRYkMQ_c,767
|
12
12
|
futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
|
13
13
|
futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
|
14
|
-
futurehouse_client-0.3.18.
|
15
|
-
futurehouse_client-0.3.18.
|
16
|
-
futurehouse_client-0.3.18.
|
17
|
-
futurehouse_client-0.3.18.
|
14
|
+
futurehouse_client-0.3.18.dev195.dist-info/METADATA,sha256=yM1NbN2au3MmkfIkkuT85eYahKYTmnBuaWCQ1OvQ97A,12767
|
15
|
+
futurehouse_client-0.3.18.dev195.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
|
16
|
+
futurehouse_client-0.3.18.dev195.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
|
17
|
+
futurehouse_client-0.3.18.dev195.dist-info/RECORD,,
|
{futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.18.dev195.dist-info}/WHEEL
RENAMED
File without changes
|
File without changes
|