futurehouse-client 0.3.18.dev109__py3-none-any.whl → 0.3.18.dev184__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 +89 -103
- futurehouse_client/models/__init__.py +10 -0
- futurehouse_client/models/app.py +93 -0
- futurehouse_client/utils/__init__.py +0 -3
- futurehouse_client/utils/auth.py +107 -0
- {futurehouse_client-0.3.18.dev109.dist-info → futurehouse_client-0.3.18.dev184.dist-info}/METADATA +1 -1
- futurehouse_client-0.3.18.dev184.dist-info/RECORD +17 -0
- {futurehouse_client-0.3.18.dev109.dist-info → futurehouse_client-0.3.18.dev184.dist-info}/WHEEL +1 -1
- futurehouse_client/utils/context.py +0 -16
- futurehouse_client-0.3.18.dev109.dist-info/RECORD +0 -17
- {futurehouse_client-0.3.18.dev109.dist-info → futurehouse_client-0.3.18.dev184.dist-info}/top_level.txt +0 -0
@@ -12,8 +12,7 @@ import sys
|
|
12
12
|
import tempfile
|
13
13
|
import time
|
14
14
|
import uuid
|
15
|
-
from collections.abc import Collection
|
16
|
-
from datetime import datetime
|
15
|
+
from collections.abc import Collection
|
17
16
|
from pathlib import Path
|
18
17
|
from types import ModuleType
|
19
18
|
from typing import Any, ClassVar, assert_never, cast
|
@@ -34,7 +33,6 @@ from httpx import (
|
|
34
33
|
RemoteProtocolError,
|
35
34
|
)
|
36
35
|
from ldp.agent import AgentConfig
|
37
|
-
from pydantic import BaseModel, ConfigDict, model_validator
|
38
36
|
from requests.exceptions import RequestException, Timeout
|
39
37
|
from tenacity import (
|
40
38
|
retry,
|
@@ -50,10 +48,18 @@ from futurehouse_client.models.app import (
|
|
50
48
|
APIKeyPayload,
|
51
49
|
AuthType,
|
52
50
|
JobDeploymentConfig,
|
51
|
+
PQATaskResponse,
|
53
52
|
Stage,
|
54
53
|
TaskRequest,
|
54
|
+
TaskResponse,
|
55
|
+
TaskResponseVerbose,
|
55
56
|
)
|
56
57
|
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
63
|
from futurehouse_client.utils.general import gather_with_concurrency
|
58
64
|
from futurehouse_client.utils.module_utils import (
|
59
65
|
OrganizationSelector,
|
@@ -65,7 +71,7 @@ from futurehouse_client.utils.monitoring import (
|
|
65
71
|
|
66
72
|
logger = logging.getLogger(__name__)
|
67
73
|
logging.basicConfig(
|
68
|
-
level=logging.
|
74
|
+
level=logging.WARNING,
|
69
75
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
70
76
|
stream=sys.stdout,
|
71
77
|
)
|
@@ -122,103 +128,11 @@ retry_if_connection_error = retry_if_exception_type((
|
|
122
128
|
FileUploadError,
|
123
129
|
))
|
124
130
|
|
125
|
-
|
126
|
-
class SimpleOrganization(BaseModel):
|
127
|
-
id: int
|
128
|
-
name: str
|
129
|
-
display_name: str
|
130
|
-
|
131
|
-
|
132
131
|
# 5 minute default for JWTs
|
133
132
|
JWT_TOKEN_CACHE_EXPIRY: int = 300 # seconds
|
134
133
|
DEFAULT_AGENT_TIMEOUT: int = 2400 # seconds
|
135
134
|
|
136
135
|
|
137
|
-
class TaskResponse(BaseModel):
|
138
|
-
"""Base class for task responses. This holds attributes shared over all futurehouse jobs."""
|
139
|
-
|
140
|
-
model_config = ConfigDict(extra="ignore")
|
141
|
-
|
142
|
-
status: str
|
143
|
-
query: str
|
144
|
-
user: str | None = None
|
145
|
-
created_at: datetime
|
146
|
-
job_name: str
|
147
|
-
public: bool
|
148
|
-
shared_with: list[SimpleOrganization] | None = None
|
149
|
-
build_owner: str | None = None
|
150
|
-
environment_name: str | None = None
|
151
|
-
agent_name: str | None = None
|
152
|
-
task_id: UUID | None = None
|
153
|
-
|
154
|
-
@model_validator(mode="before")
|
155
|
-
@classmethod
|
156
|
-
def validate_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
|
157
|
-
# Extract fields from environment frame state
|
158
|
-
if not isinstance(data, dict):
|
159
|
-
return data
|
160
|
-
# TODO: We probably want to remove these two once we define the final names.
|
161
|
-
data["job_name"] = data.get("crow")
|
162
|
-
data["query"] = data.get("task")
|
163
|
-
data["task_id"] = cast(UUID, data.get("id")) if data.get("id") else None
|
164
|
-
if not (metadata := data.get("metadata", {})):
|
165
|
-
return data
|
166
|
-
data["environment_name"] = metadata.get("environment_name")
|
167
|
-
data["agent_name"] = metadata.get("agent_name")
|
168
|
-
return data
|
169
|
-
|
170
|
-
|
171
|
-
class PQATaskResponse(TaskResponse):
|
172
|
-
model_config = ConfigDict(extra="ignore")
|
173
|
-
|
174
|
-
answer: str | None = None
|
175
|
-
formatted_answer: str | None = None
|
176
|
-
answer_reasoning: str | None = None
|
177
|
-
has_successful_answer: bool | None = None
|
178
|
-
total_cost: float | None = None
|
179
|
-
total_queries: int | None = None
|
180
|
-
|
181
|
-
@model_validator(mode="before")
|
182
|
-
@classmethod
|
183
|
-
def validate_pqa_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
|
184
|
-
if not isinstance(data, dict):
|
185
|
-
return data
|
186
|
-
if not (env_frame := data.get("environment_frame", {})):
|
187
|
-
return data
|
188
|
-
state = env_frame.get("state", {}).get("state", {})
|
189
|
-
response = state.get("response", {})
|
190
|
-
answer = response.get("answer", {})
|
191
|
-
usage = state.get("info", {}).get("usage", {})
|
192
|
-
|
193
|
-
# Add additional PQA specific fields to data so that pydantic can validate the model
|
194
|
-
data["answer"] = answer.get("answer")
|
195
|
-
data["formatted_answer"] = answer.get("formatted_answer")
|
196
|
-
data["answer_reasoning"] = answer.get("answer_reasoning")
|
197
|
-
data["has_successful_answer"] = answer.get("has_successful_answer")
|
198
|
-
data["total_cost"] = cast(float, usage.get("total_cost"))
|
199
|
-
data["total_queries"] = cast(int, usage.get("total_queries"))
|
200
|
-
|
201
|
-
return data
|
202
|
-
|
203
|
-
def clean_verbose(self) -> "TaskResponse":
|
204
|
-
"""Clean the verbose response from the server."""
|
205
|
-
self.request = None
|
206
|
-
self.response = None
|
207
|
-
return self
|
208
|
-
|
209
|
-
|
210
|
-
class TaskResponseVerbose(TaskResponse):
|
211
|
-
"""Class for responses to include all the fields of a task response."""
|
212
|
-
|
213
|
-
model_config = ConfigDict(extra="allow")
|
214
|
-
|
215
|
-
public: bool
|
216
|
-
agent_state: list[dict[str, Any]] | None = None
|
217
|
-
environment_frame: dict[str, Any] | None = None
|
218
|
-
metadata: dict[str, Any] | None = None
|
219
|
-
shared_with: list[SimpleOrganization] | None = None
|
220
|
-
|
221
|
-
|
222
136
|
class RestClient:
|
223
137
|
REQUEST_TIMEOUT: ClassVar[float] = 30.0 # sec
|
224
138
|
MAX_RETRY_ATTEMPTS: ClassVar[int] = 3
|
@@ -236,7 +150,13 @@ class RestClient:
|
|
236
150
|
api_key: str | None = None,
|
237
151
|
jwt: str | None = None,
|
238
152
|
headers: dict[str, str] | None = None,
|
153
|
+
verbose_logging: bool = False,
|
239
154
|
):
|
155
|
+
if verbose_logging:
|
156
|
+
logger.setLevel(logging.INFO)
|
157
|
+
else:
|
158
|
+
logger.setLevel(logging.WARNING)
|
159
|
+
|
240
160
|
self.base_url = service_uri or stage.value
|
241
161
|
self.stage = stage
|
242
162
|
self.auth_type = auth_type
|
@@ -360,6 +280,7 @@ class RestClient:
|
|
360
280
|
except Exception as e:
|
361
281
|
raise RestClientError(f"Error authenticating: {e!s}") from e
|
362
282
|
|
283
|
+
@refresh_token_on_auth_error()
|
363
284
|
def _check_job(self, name: str, organization: str) -> dict[str, Any]:
|
364
285
|
try:
|
365
286
|
response = self.client.get(
|
@@ -367,9 +288,19 @@ class RestClient:
|
|
367
288
|
)
|
368
289
|
response.raise_for_status()
|
369
290
|
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
|
370
300
|
except Exception as e:
|
371
301
|
raise JobFetchError(f"Error checking job: {e!s}") from e
|
372
302
|
|
303
|
+
@refresh_token_on_auth_error()
|
373
304
|
def _fetch_my_orgs(self) -> list[str]:
|
374
305
|
response = self.client.get(f"/v0.1/organizations?filter={True}")
|
375
306
|
response.raise_for_status()
|
@@ -432,6 +363,7 @@ class RestClient:
|
|
432
363
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
433
364
|
retry=retry_if_connection_error,
|
434
365
|
)
|
366
|
+
@refresh_token_on_auth_error()
|
435
367
|
def get_task(
|
436
368
|
self, task_id: str | None = None, history: bool = False, verbose: bool = False
|
437
369
|
) -> "TaskResponse":
|
@@ -467,8 +399,15 @@ class RestClient:
|
|
467
399
|
):
|
468
400
|
return PQATaskResponse(**data)
|
469
401
|
return TaskResponse(**data)
|
470
|
-
except
|
471
|
-
|
402
|
+
except HTTPStatusError as e:
|
403
|
+
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
404
|
+
raise AuthError(
|
405
|
+
e.response.status_code,
|
406
|
+
f"Authentication failed: {e}",
|
407
|
+
request=e.request,
|
408
|
+
response=e.response,
|
409
|
+
) from e
|
410
|
+
raise
|
472
411
|
except Exception as e:
|
473
412
|
raise TaskFetchError(f"Error getting task: {e!s}") from e
|
474
413
|
|
@@ -477,6 +416,7 @@ class RestClient:
|
|
477
416
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
478
417
|
retry=retry_if_connection_error,
|
479
418
|
)
|
419
|
+
@refresh_token_on_auth_error()
|
480
420
|
async def aget_task(
|
481
421
|
self, task_id: str | None = None, history: bool = False, verbose: bool = False
|
482
422
|
) -> "TaskResponse":
|
@@ -515,11 +455,19 @@ class RestClient:
|
|
515
455
|
):
|
516
456
|
return PQATaskResponse(**data)
|
517
457
|
return TaskResponse(**data)
|
518
|
-
except
|
519
|
-
|
458
|
+
except HTTPStatusError as e:
|
459
|
+
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
460
|
+
raise AuthError(
|
461
|
+
e.response.status_code,
|
462
|
+
f"Authentication failed: {e}",
|
463
|
+
request=e.request,
|
464
|
+
response=e.response,
|
465
|
+
) from e
|
466
|
+
raise
|
520
467
|
except Exception as e:
|
521
468
|
raise TaskFetchError(f"Error getting task: {e!s}") from e
|
522
469
|
|
470
|
+
@refresh_token_on_auth_error()
|
523
471
|
@retry(
|
524
472
|
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
|
525
473
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
@@ -543,6 +491,15 @@ class RestClient:
|
|
543
491
|
response.raise_for_status()
|
544
492
|
trajectory_id = response.json()["trajectory_id"]
|
545
493
|
self.trajectory_id = trajectory_id
|
494
|
+
except HTTPStatusError as e:
|
495
|
+
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
496
|
+
raise AuthError(
|
497
|
+
e.response.status_code,
|
498
|
+
f"Authentication failed: {e}",
|
499
|
+
request=e.request,
|
500
|
+
response=e.response,
|
501
|
+
) from e
|
502
|
+
raise
|
546
503
|
except Exception as e:
|
547
504
|
raise TaskFetchError(f"Error creating task: {e!s}") from e
|
548
505
|
return trajectory_id
|
@@ -552,6 +509,7 @@ class RestClient:
|
|
552
509
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
553
510
|
retry=retry_if_connection_error,
|
554
511
|
)
|
512
|
+
@refresh_token_on_auth_error()
|
555
513
|
async def acreate_task(self, task_data: TaskRequest | dict[str, Any]):
|
556
514
|
"""Create a new futurehouse task."""
|
557
515
|
if isinstance(task_data, dict):
|
@@ -570,6 +528,15 @@ class RestClient:
|
|
570
528
|
response.raise_for_status()
|
571
529
|
trajectory_id = response.json()["trajectory_id"]
|
572
530
|
self.trajectory_id = trajectory_id
|
531
|
+
except HTTPStatusError as e:
|
532
|
+
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
533
|
+
raise AuthError(
|
534
|
+
e.response.status_code,
|
535
|
+
f"Authentication failed: {e}",
|
536
|
+
request=e.request,
|
537
|
+
response=e.response,
|
538
|
+
) from e
|
539
|
+
raise
|
573
540
|
except Exception as e:
|
574
541
|
raise TaskFetchError(f"Error creating task: {e!s}") from e
|
575
542
|
return trajectory_id
|
@@ -720,11 +687,22 @@ class RestClient:
|
|
720
687
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
721
688
|
retry=retry_if_connection_error,
|
722
689
|
)
|
690
|
+
@refresh_token_on_auth_error()
|
723
691
|
def get_build_status(self, build_id: UUID | None = None) -> dict[str, Any]:
|
724
692
|
"""Get the status of a build."""
|
725
|
-
|
726
|
-
|
727
|
-
|
693
|
+
try:
|
694
|
+
build_id = build_id or self.build_id
|
695
|
+
response = self.client.get(f"/v0.1/builds/{build_id}")
|
696
|
+
response.raise_for_status()
|
697
|
+
except HTTPStatusError as e:
|
698
|
+
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
699
|
+
raise AuthError(
|
700
|
+
e.response.status_code,
|
701
|
+
f"Authentication failed: {e}",
|
702
|
+
request=e.request,
|
703
|
+
response=e.response,
|
704
|
+
) from e
|
705
|
+
raise
|
728
706
|
return response.json()
|
729
707
|
|
730
708
|
# TODO: Refactor later so we don't have to ignore PLR0915
|
@@ -733,6 +711,7 @@ class RestClient:
|
|
733
711
|
wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
|
734
712
|
retry=retry_if_connection_error,
|
735
713
|
)
|
714
|
+
@refresh_token_on_auth_error()
|
736
715
|
def create_job(self, config: JobDeploymentConfig) -> dict[str, Any]: # noqa: PLR0915
|
737
716
|
"""Creates a futurehouse job deployment from the environment and environment files.
|
738
717
|
|
@@ -907,6 +886,13 @@ class RestClient:
|
|
907
886
|
build_context = response.json()
|
908
887
|
self.build_id = build_context["build_id"]
|
909
888
|
except HTTPStatusError as e:
|
889
|
+
if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
|
890
|
+
raise AuthError(
|
891
|
+
e.response.status_code,
|
892
|
+
f"Authentication failed: {e}",
|
893
|
+
request=e.request,
|
894
|
+
response=e.response,
|
895
|
+
) from e
|
910
896
|
error_detail = response.json()
|
911
897
|
error_message = error_detail.get("detail", str(e))
|
912
898
|
raise JobCreationError(
|
@@ -3,10 +3,15 @@ from .app import (
|
|
3
3
|
DockerContainerConfiguration,
|
4
4
|
FramePath,
|
5
5
|
JobDeploymentConfig,
|
6
|
+
PQATaskResponse,
|
6
7
|
RuntimeConfig,
|
7
8
|
Stage,
|
8
9
|
Step,
|
10
|
+
TaskQueue,
|
11
|
+
TaskQueuesConfig,
|
9
12
|
TaskRequest,
|
13
|
+
TaskResponse,
|
14
|
+
TaskResponseVerbose,
|
10
15
|
)
|
11
16
|
|
12
17
|
__all__ = [
|
@@ -14,8 +19,13 @@ __all__ = [
|
|
14
19
|
"DockerContainerConfiguration",
|
15
20
|
"FramePath",
|
16
21
|
"JobDeploymentConfig",
|
22
|
+
"PQATaskResponse",
|
17
23
|
"RuntimeConfig",
|
18
24
|
"Stage",
|
19
25
|
"Step",
|
26
|
+
"TaskQueue",
|
27
|
+
"TaskQueuesConfig",
|
20
28
|
"TaskRequest",
|
29
|
+
"TaskResponse",
|
30
|
+
"TaskResponseVerbose",
|
21
31
|
]
|
futurehouse_client/models/app.py
CHANGED
@@ -1,6 +1,8 @@
|
|
1
1
|
import json
|
2
2
|
import os
|
3
3
|
import re
|
4
|
+
from collections.abc import Mapping
|
5
|
+
from datetime import datetime
|
4
6
|
from enum import StrEnum, auto
|
5
7
|
from pathlib import Path
|
6
8
|
from typing import TYPE_CHECKING, Any, ClassVar, Self, cast
|
@@ -646,3 +648,94 @@ class TaskRequest(BaseModel):
|
|
646
648
|
runtime_config: RuntimeConfig | None = Field(
|
647
649
|
default=None, description="All optional runtime parameters for the job"
|
648
650
|
)
|
651
|
+
|
652
|
+
|
653
|
+
class SimpleOrganization(BaseModel):
|
654
|
+
id: int
|
655
|
+
name: str
|
656
|
+
display_name: str
|
657
|
+
|
658
|
+
|
659
|
+
class TaskResponse(BaseModel):
|
660
|
+
"""Base class for task responses. This holds attributes shared over all futurehouse jobs."""
|
661
|
+
|
662
|
+
model_config = ConfigDict(extra="ignore")
|
663
|
+
|
664
|
+
status: str
|
665
|
+
query: str
|
666
|
+
user: str | None = None
|
667
|
+
created_at: datetime
|
668
|
+
job_name: str
|
669
|
+
public: bool
|
670
|
+
shared_with: list[SimpleOrganization] | None = None
|
671
|
+
build_owner: str | None = None
|
672
|
+
environment_name: str | None = None
|
673
|
+
agent_name: str | None = None
|
674
|
+
task_id: UUID | None = None
|
675
|
+
|
676
|
+
@model_validator(mode="before")
|
677
|
+
@classmethod
|
678
|
+
def validate_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
|
679
|
+
# Extract fields from environment frame state
|
680
|
+
if not isinstance(data, dict):
|
681
|
+
return data
|
682
|
+
# TODO: We probably want to remove these two once we define the final names.
|
683
|
+
data["job_name"] = data.get("crow")
|
684
|
+
data["query"] = data.get("task")
|
685
|
+
data["task_id"] = cast(UUID, data.get("id")) if data.get("id") else None
|
686
|
+
if not (metadata := data.get("metadata", {})):
|
687
|
+
return data
|
688
|
+
data["environment_name"] = metadata.get("environment_name")
|
689
|
+
data["agent_name"] = metadata.get("agent_name")
|
690
|
+
return data
|
691
|
+
|
692
|
+
|
693
|
+
class PQATaskResponse(TaskResponse):
|
694
|
+
model_config = ConfigDict(extra="ignore")
|
695
|
+
|
696
|
+
answer: str | None = None
|
697
|
+
formatted_answer: str | None = None
|
698
|
+
answer_reasoning: str | None = None
|
699
|
+
has_successful_answer: bool | None = None
|
700
|
+
total_cost: float | None = None
|
701
|
+
total_queries: int | None = None
|
702
|
+
|
703
|
+
@model_validator(mode="before")
|
704
|
+
@classmethod
|
705
|
+
def validate_pqa_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
|
706
|
+
if not isinstance(data, dict):
|
707
|
+
return data
|
708
|
+
if not (env_frame := data.get("environment_frame", {})):
|
709
|
+
return data
|
710
|
+
state = env_frame.get("state", {}).get("state", {})
|
711
|
+
response = state.get("response", {})
|
712
|
+
answer = response.get("answer", {})
|
713
|
+
usage = state.get("info", {}).get("usage", {})
|
714
|
+
|
715
|
+
# Add additional PQA specific fields to data so that pydantic can validate the model
|
716
|
+
data["answer"] = answer.get("answer")
|
717
|
+
data["formatted_answer"] = answer.get("formatted_answer")
|
718
|
+
data["answer_reasoning"] = answer.get("answer_reasoning")
|
719
|
+
data["has_successful_answer"] = answer.get("has_successful_answer")
|
720
|
+
data["total_cost"] = cast(float, usage.get("total_cost"))
|
721
|
+
data["total_queries"] = cast(int, usage.get("total_queries"))
|
722
|
+
|
723
|
+
return data
|
724
|
+
|
725
|
+
def clean_verbose(self) -> "TaskResponse":
|
726
|
+
"""Clean the verbose response from the server."""
|
727
|
+
self.request = None
|
728
|
+
self.response = None
|
729
|
+
return self
|
730
|
+
|
731
|
+
|
732
|
+
class TaskResponseVerbose(TaskResponse):
|
733
|
+
"""Class for responses to include all the fields of a task response."""
|
734
|
+
|
735
|
+
model_config = ConfigDict(extra="allow")
|
736
|
+
|
737
|
+
public: bool
|
738
|
+
agent_state: list[dict[str, Any]] | None = None
|
739
|
+
environment_frame: dict[str, Any] | None = None
|
740
|
+
metadata: dict[str, Any] | None = None
|
741
|
+
shared_with: list[SimpleOrganization] | None = None
|
@@ -0,0 +1,107 @@
|
|
1
|
+
import asyncio
|
2
|
+
import logging
|
3
|
+
from collections.abc import Callable, Coroutine
|
4
|
+
from functools import wraps
|
5
|
+
from typing import Any, Final, Optional, ParamSpec, TypeVar, overload
|
6
|
+
|
7
|
+
import httpx
|
8
|
+
from httpx import HTTPStatusError
|
9
|
+
|
10
|
+
logger = logging.getLogger(__name__)
|
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
|
45
|
+
|
46
|
+
|
47
|
+
@overload
|
48
|
+
def refresh_token_on_auth_error(
|
49
|
+
func: Callable[P, Coroutine[Any, Any, T]],
|
50
|
+
) -> Callable[P, Coroutine[Any, Any, T]]: ...
|
51
|
+
|
52
|
+
|
53
|
+
@overload
|
54
|
+
def refresh_token_on_auth_error(
|
55
|
+
func: None = None, *, max_retries: int = ...
|
56
|
+
) -> Callable[[Callable[P, T]], Callable[P, T]]: ...
|
57
|
+
|
58
|
+
|
59
|
+
def refresh_token_on_auth_error(func=None, max_retries=1):
|
60
|
+
"""Decorator that refreshes JWT token on 401/403 auth errors."""
|
61
|
+
|
62
|
+
def decorator(fn):
|
63
|
+
@wraps(fn)
|
64
|
+
def sync_wrapper(self, *args, **kwargs):
|
65
|
+
retries = 0
|
66
|
+
while True:
|
67
|
+
try:
|
68
|
+
return fn(self, *args, **kwargs)
|
69
|
+
except Exception as e:
|
70
|
+
if is_auth_error(e) and retries < max_retries:
|
71
|
+
retries += 1
|
72
|
+
status = get_status_code(e) or "Unknown"
|
73
|
+
logger.info(
|
74
|
+
f"Received auth error {status}, "
|
75
|
+
f"refreshing token and retrying (attempt {retries}/{max_retries})..."
|
76
|
+
)
|
77
|
+
self.auth_jwt = self._run_auth()
|
78
|
+
self._clients = {}
|
79
|
+
continue
|
80
|
+
raise
|
81
|
+
|
82
|
+
@wraps(fn)
|
83
|
+
async def async_wrapper(self, *args, **kwargs):
|
84
|
+
retries = 0
|
85
|
+
while True:
|
86
|
+
try:
|
87
|
+
return await fn(self, *args, **kwargs)
|
88
|
+
except Exception as e:
|
89
|
+
if is_auth_error(e) and retries < max_retries:
|
90
|
+
retries += 1
|
91
|
+
status = get_status_code(e) or "Unknown"
|
92
|
+
logger.info(
|
93
|
+
f"Received auth error {status}, "
|
94
|
+
f"refreshing token and retrying (attempt {retries}/{max_retries})..."
|
95
|
+
)
|
96
|
+
self.auth_jwt = self._run_auth()
|
97
|
+
self._clients = {}
|
98
|
+
continue
|
99
|
+
raise
|
100
|
+
|
101
|
+
if asyncio.iscoroutinefunction(fn):
|
102
|
+
return async_wrapper
|
103
|
+
return sync_wrapper
|
104
|
+
|
105
|
+
if callable(func):
|
106
|
+
return decorator(func)
|
107
|
+
return decorator
|
{futurehouse_client-0.3.18.dev109.dist-info → futurehouse_client-0.3.18.dev184.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.dev184
|
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
|
@@ -0,0 +1,17 @@
|
|
1
|
+
futurehouse_client/__init__.py,sha256=ddxO7JE97c6bt7LjNglZZ2Ql8bYCGI9laSFeh9MP6VU,344
|
2
|
+
futurehouse_client/clients/__init__.py,sha256=tFWqwIAY5PvwfOVsCje4imjTpf6xXNRMh_UHIKVI1_0,320
|
3
|
+
futurehouse_client/clients/job_client.py,sha256=uNkqQbeZw7wbA0qDWcIOwOykrosza-jev58paJZ_mbA,11150
|
4
|
+
futurehouse_client/clients/rest_client.py,sha256=0SN9cNy2QFVUS9U1L1KDf_ncSx2IV7phbGatdbNlY9w,46814
|
5
|
+
futurehouse_client/models/__init__.py,sha256=5x-f9AoM1hGzJBEHcHAXSt7tPeImST5oZLuMdwp0mXc,554
|
6
|
+
futurehouse_client/models/app.py,sha256=w_1e4F0IiC-BKeOLqYkABYo4U-Nka1S-F64S_eHB2KM,26421
|
7
|
+
futurehouse_client/models/client.py,sha256=n4HD0KStKLm6Ek9nL9ylP-bkK10yzAaD1uIDF83Qp_A,1828
|
8
|
+
futurehouse_client/models/rest.py,sha256=lgwkMIXz0af-49BYSkKeS7SRqvN3motqnAikDN4YGTc,789
|
9
|
+
futurehouse_client/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
10
|
+
futurehouse_client/utils/auth.py,sha256=Lq9mjSGc7iuRP6fmLICCS6KjzLHN6-tJUuhYp0XXrkE,3342
|
11
|
+
futurehouse_client/utils/general.py,sha256=A_rtTiYW30ELGEZlWCIArO7q1nEmqi8hUlmBRYkMQ_c,767
|
12
|
+
futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
|
13
|
+
futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
|
14
|
+
futurehouse_client-0.3.18.dev184.dist-info/METADATA,sha256=ckn5Ucj4fuMFejVZTtAuw3H7gSqj4NUiIpdLzp8W16I,12767
|
15
|
+
futurehouse_client-0.3.18.dev184.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
|
16
|
+
futurehouse_client-0.3.18.dev184.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
|
17
|
+
futurehouse_client-0.3.18.dev184.dist-info/RECORD,,
|
@@ -1,16 +0,0 @@
|
|
1
|
-
class UserContext:
|
2
|
-
"""A context manager for storing user information from the initial request."""
|
3
|
-
|
4
|
-
_user_jwt = None
|
5
|
-
|
6
|
-
@classmethod
|
7
|
-
def set_user_jwt(cls, jwt: str) -> None:
|
8
|
-
cls._user_jwt = jwt
|
9
|
-
|
10
|
-
@classmethod
|
11
|
-
def get_user_jwt(cls) -> str | None:
|
12
|
-
return cls._user_jwt
|
13
|
-
|
14
|
-
@classmethod
|
15
|
-
def clear_user_jwt(cls) -> None:
|
16
|
-
cls._user_jwt = None
|
@@ -1,17 +0,0 @@
|
|
1
|
-
futurehouse_client/__init__.py,sha256=ddxO7JE97c6bt7LjNglZZ2Ql8bYCGI9laSFeh9MP6VU,344
|
2
|
-
futurehouse_client/clients/__init__.py,sha256=tFWqwIAY5PvwfOVsCje4imjTpf6xXNRMh_UHIKVI1_0,320
|
3
|
-
futurehouse_client/clients/job_client.py,sha256=uNkqQbeZw7wbA0qDWcIOwOykrosza-jev58paJZ_mbA,11150
|
4
|
-
futurehouse_client/clients/rest_client.py,sha256=Qv54VFcGnCDbOoGFmfx8AZsXyAyZyZ4weK9RGKVePOE,47214
|
5
|
-
futurehouse_client/models/__init__.py,sha256=ta3jFLM_LsDz1rKDmx8rja8sT7WtSKoFvMgLF0yFpvA,342
|
6
|
-
futurehouse_client/models/app.py,sha256=yfZ9tyw4VATVAfYrU7aTdCNPSljLEho09_nIbh8oZDY,23174
|
7
|
-
futurehouse_client/models/client.py,sha256=n4HD0KStKLm6Ek9nL9ylP-bkK10yzAaD1uIDF83Qp_A,1828
|
8
|
-
futurehouse_client/models/rest.py,sha256=lgwkMIXz0af-49BYSkKeS7SRqvN3motqnAikDN4YGTc,789
|
9
|
-
futurehouse_client/utils/__init__.py,sha256=mCp1UP3eyrWyHpFAEHgKcazKCJ5g87MxlDIZVzi-5Tk,60
|
10
|
-
futurehouse_client/utils/context.py,sha256=MdYYldQfsqxnGJUDcMLkbRtKHb2srHfcMf6aR-oZh6Y,387
|
11
|
-
futurehouse_client/utils/general.py,sha256=A_rtTiYW30ELGEZlWCIArO7q1nEmqi8hUlmBRYkMQ_c,767
|
12
|
-
futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
|
13
|
-
futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
|
14
|
-
futurehouse_client-0.3.18.dev109.dist-info/METADATA,sha256=yeDUNer_jX0nDUAgV26q5ssYrcWmWlCpuqy6iRuSz_c,12767
|
15
|
-
futurehouse_client-0.3.18.dev109.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
|
16
|
-
futurehouse_client-0.3.18.dev109.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
|
17
|
-
futurehouse_client-0.3.18.dev109.dist-info/RECORD,,
|
File without changes
|