diff --git a/src/replit/ai/modelfarm/replit_identity_token_manager.py b/src/replit/ai/modelfarm/replit_identity_token_manager.py index 57cf640..c275bf2 100644 --- a/src/replit/ai/modelfarm/replit_identity_token_manager.py +++ b/src/replit/ai/modelfarm/replit_identity_token_manager.py @@ -24,15 +24,15 @@ def __init__(self, token_timeout: int = 300): self.last_update: Optional[float] = None self.token: Optional[str] = None self.__set_token_type() - if self.token_type != "L402": - self.__update_token() + self.__update_token() def __set_token_type(self): """Sets the type of token to be used. Bearer if in replit, L402 if not.""" - if self.__in_deployment(): - self.token_type = "Bearer" - else: + if self.__using_L402(): self.token_type = "L402" + else : + self.token_type = "Bearer" + def get_token(self) -> Optional[str]: """Returns the token, updates if the current token has expired. @@ -58,14 +58,16 @@ def get_new_token(self) -> str: Returns: str: The most recent token. """ - try: + try: if self.__in_deployment(): return self.get_deployment_token() - return self.get_interactive_token() - # If the environment variables are not set, they're not in Replit, so use L402 - except MissingEnvironmentVariable: - print("Using L402 token") - return self.get_l402_token() + elif self.__using_L402(): + return self.get_l402_token() + else: + return self.get_interactive_token() + # Fallback, try to let them use L402 + except: + return self.get_402_token() def get_deployment_token(self) -> str: """Fetches deployment token from hostindpid1. @@ -108,7 +110,7 @@ def get_l402_token(self) -> str: try: l402_token = self.get_env_var("REPLIT_L402_TOKEN") l402_preimage = self.get_env_var("REPLIT_L402_PREIMAGE") - return "L402 " + l402_token + ":" + l402_preimage + return l402_token + ":" + l402_preimage except MissingEnvironmentVariable: print("Generating new L402 token") res = requests.get("https://matador-replit.kody.repl.co") @@ -142,6 +144,14 @@ def __in_deployment(self) -> bool: """ return "REPLIT_DEPLOYMENT" in os.environ + def __using_L402(self) -> bool: + """Determines if using L402. + + Returns: + bool: True if using L402, False otherwise. + """ + return "REPLIT_L402_TOKEN" in os.environ and "REPLIT_L402_PREIMAGE" in os.environ + def split_L402(l402: str) -> tuple[str, str]: """Splits L402 token into token and invoice.