import typing
import time
import requests
if typing.TYPE_CHECKING:
from azure.core.credentials import TokenCredential
class TokenRefresh(requests.auth.AuthBase):
def __init__(self, credential: "TokenCredential", scopes: typing.List[str]) -> None:
self.credential = credential
self.scopes = scopes
self.cached_token: typing.Optional[str] = None
def __call__(self, req):
if not self.cached_token or self.cached_token.expires_on - time.time() < 300:
self.cached_token = self.credential.get_token(*self.scopes)
req.headers["Authorization"] = f"Bearer {self.cached_token.token}"
return req
session = requests.Session()
session.auth = TokenRefresh(default_credential, ["https://cognitiveservices.azure.com/.default"])
openai.requestssession = session