forked from runpod/runpod-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathhttp_client.py
More file actions
53 lines (41 loc) · 1.24 KB
/
http_client.py
File metadata and controls
53 lines (41 loc) · 1.24 KB
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
"""
HTTP Client abstractions
"""
import os
import requests
from aiohttp import ClientSession, ClientTimeout, TCPConnector, ClientResponseError
from .cli.groups.config.functions import get_credentials
from .user_agent import USER_AGENT
class TooManyRequests(ClientResponseError):
pass
def get_auth_header():
"""
Produce a header dict with the `Authorization` key derived from
credentials.get("api_key") OR os.getenv('RUNPOD_AI_API_KEY')
"""
if credentials := get_credentials():
auth = credentials.get("api_key", "")
else:
auth = os.getenv("RUNPOD_AI_API_KEY", "")
return {
"Content-Type": "application/json",
"Authorization": auth,
"User-Agent": USER_AGENT,
}
def AsyncClientSession(*args, **kwargs): # pylint: disable=invalid-name
"""
Deprecation from aiohttp.ClientSession forbids inheritance.
This is now a factory method
"""
return ClientSession(
connector=TCPConnector(limit=0),
headers=get_auth_header(),
timeout=ClientTimeout(600, ceil_threshold=400),
*args,
**kwargs,
)
class SyncClientSession(requests.Session):
"""
Inherits requests.Session to override `request()` method for tracing
"""
pass