classAuth:id_token=api_key=model_key=Falsedef__init__(self,api_key='',verbose=False):""" Initialize the Auth class with an optional API key. Args: api_key (str, optional): May be an API key or a combination API key and model ID, i.e. key_id """# Split the input API key in case it contains a combined key_model and keep only the API key partapi_key=api_key.split('_')[0]# Set API key attribute as value passed or SETTINGS API key if none passedself.api_key=api_keyorSETTINGS.get('api_key','')# If an API key is providedifself.api_key:# If the provided API key matches the API key in the SETTINGSifself.api_key==SETTINGS.get('api_key'):# Log that the user is already logged inifverbose:LOGGER.info(f'{PREFIX}Authenticated ✅')returnelse:# Attempt to authenticate with the provided API keysuccess=self.authenticate()# If the API key is not provided and the environment is a Google Colab notebookelifis_colab():# Attempt to authenticate using browser cookiessuccess=self.auth_with_cookies()else:# Request an API keysuccess=self.request_api_key()# Update SETTINGS with the new API key after successful authenticationifsuccess:set_settings({'api_key':self.api_key})# Log that the new login was successfulifverbose:LOGGER.info(f'{PREFIX}New authentication successful ✅')elifverbose:LOGGER.info(f'{PREFIX}Retrieve API key from {API_KEY_URL}')defrequest_api_key(self,max_attempts=3):""" Prompt the user to input their API key. Returns the model ID. """importgetpassforattemptsinrange(max_attempts):LOGGER.info(f'{PREFIX}Login. Attempt {attempts+1} of {max_attempts}')input_key=getpass.getpass(f'Enter API key from {API_KEY_URL} ')self.api_key=input_key.split('_')[0]# remove model id if presentifself.authenticate():returnTrueraiseConnectionError(emojis(f'{PREFIX}Failed to authenticate ❌'))defauthenticate(self)->bool:""" Attempt to authenticate with the server using either id_token or API key. Returns: bool: True if authentication is successful, False otherwise. """try:header=self.get_auth_header()ifheader:r=requests.post(f'{HUB_API_ROOT}/v1/auth',headers=header)ifnotr.json().get('success',False):raiseConnectionError('Unable to authenticate.')returnTrueraiseConnectionError('User has not authenticated locally.')exceptConnectionError:self.id_token=self.api_key=False# reset invalidLOGGER.warning(f'{PREFIX}Invalid API key ⚠️')returnFalsedefauth_with_cookies(self)->bool:""" Attempt to fetch authentication via cookies and set id_token. User must be logged in to HUB and running in a supported browser. Returns: bool: True if authentication is successful, False otherwise. """ifnotis_colab():returnFalse# Currently only works with Colabtry:authn=request_with_credentials(f'{HUB_API_ROOT}/v1/auth/auto')ifauthn.get('success',False):self.id_token=authn.get('data',{}).get('idToken',None)self.authenticate()returnTrueraiseConnectionError('Unable to fetch browser authentication details.')exceptConnectionError:self.id_token=False# reset invalidreturnFalsedefget_auth_header(self):""" Get the authentication header for making API requests. Returns: (dict): The authentication header if id_token or API key is set, None otherwise. """ifself.id_token:return{'authorization':f'Bearer {self.id_token}'}elifself.api_key:return{'x-api-key':self.api_key}else:returnNonedefget_state(self)->bool:""" Get the authentication state. Returns: bool: True if either id_token or API key is set, False otherwise. """returnself.id_tokenorself.api_keydefset_api_key(self,key:str):""" Set the API key for authentication. Args: key (str): The API key string. """self.api_key=key
__init__(api_key='',verbose=False)
Initialize the Auth class with an optional API key.
Parameters:
Name
Type
Description
Default
api_key
str
May be an API key or a combination API key and model ID, i.e. key_id
def__init__(self,api_key='',verbose=False):""" Initialize the Auth class with an optional API key. Args: api_key (str, optional): May be an API key or a combination API key and model ID, i.e. key_id """# Split the input API key in case it contains a combined key_model and keep only the API key partapi_key=api_key.split('_')[0]# Set API key attribute as value passed or SETTINGS API key if none passedself.api_key=api_keyorSETTINGS.get('api_key','')# If an API key is providedifself.api_key:# If the provided API key matches the API key in the SETTINGSifself.api_key==SETTINGS.get('api_key'):# Log that the user is already logged inifverbose:LOGGER.info(f'{PREFIX}Authenticated ✅')returnelse:# Attempt to authenticate with the provided API keysuccess=self.authenticate()# If the API key is not provided and the environment is a Google Colab notebookelifis_colab():# Attempt to authenticate using browser cookiessuccess=self.auth_with_cookies()else:# Request an API keysuccess=self.request_api_key()# Update SETTINGS with the new API key after successful authenticationifsuccess:set_settings({'api_key':self.api_key})# Log that the new login was successfulifverbose:LOGGER.info(f'{PREFIX}New authentication successful ✅')elifverbose:LOGGER.info(f'{PREFIX}Retrieve API key from {API_KEY_URL}')
auth_with_cookies()
Attempt to fetch authentication via cookies and set id_token.
User must be logged in to HUB and running in a supported browser.
Returns:
Name
Type
Description
bool
bool
True if authentication is successful, False otherwise.
defauth_with_cookies(self)->bool:""" Attempt to fetch authentication via cookies and set id_token. User must be logged in to HUB and running in a supported browser. Returns: bool: True if authentication is successful, False otherwise. """ifnotis_colab():returnFalse# Currently only works with Colabtry:authn=request_with_credentials(f'{HUB_API_ROOT}/v1/auth/auto')ifauthn.get('success',False):self.id_token=authn.get('data',{}).get('idToken',None)self.authenticate()returnTrueraiseConnectionError('Unable to fetch browser authentication details.')exceptConnectionError:self.id_token=False# reset invalidreturnFalse
authenticate()
Attempt to authenticate with the server using either id_token or API key.
Returns:
Name
Type
Description
bool
bool
True if authentication is successful, False otherwise.
defauthenticate(self)->bool:""" Attempt to authenticate with the server using either id_token or API key. Returns: bool: True if authentication is successful, False otherwise. """try:header=self.get_auth_header()ifheader:r=requests.post(f'{HUB_API_ROOT}/v1/auth',headers=header)ifnotr.json().get('success',False):raiseConnectionError('Unable to authenticate.')returnTrueraiseConnectionError('User has not authenticated locally.')exceptConnectionError:self.id_token=self.api_key=False# reset invalidLOGGER.warning(f'{PREFIX}Invalid API key ⚠️')returnFalse
get_auth_header()
Get the authentication header for making API requests.
Returns:
Type
Description
dict
The authentication header if id_token or API key is set, None otherwise.
defget_auth_header(self):""" Get the authentication header for making API requests. Returns: (dict): The authentication header if id_token or API key is set, None otherwise. """ifself.id_token:return{'authorization':f'Bearer {self.id_token}'}elifself.api_key:return{'x-api-key':self.api_key}else:returnNone
get_state()
Get the authentication state.
Returns:
Name
Type
Description
bool
bool
True if either id_token or API key is set, False otherwise.
defget_state(self)->bool:""" Get the authentication state. Returns: bool: True if either id_token or API key is set, False otherwise. """returnself.id_tokenorself.api_key
request_api_key(max_attempts=3)
Prompt the user to input their API key. Returns the model ID.
defrequest_api_key(self,max_attempts=3):""" Prompt the user to input their API key. Returns the model ID. """importgetpassforattemptsinrange(max_attempts):LOGGER.info(f'{PREFIX}Login. Attempt {attempts+1} of {max_attempts}')input_key=getpass.getpass(f'Enter API key from {API_KEY_URL} ')self.api_key=input_key.split('_')[0]# remove model id if presentifself.authenticate():returnTrueraiseConnectionError(emojis(f'{PREFIX}Failed to authenticate ❌'))