Add optional cache to whoami (#3568)

* Add optional cache to whoami

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Abubakar Abid <[email protected]>

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Abubakar Abid <[email protected]>

* Fix token=False not working

* add test for token=False

* Update src/huggingface_hub/hf_api.py

Co-authored-by: célina <[email protected]>

* Update src/huggingface_hub/hf_api.py

Co-authored-by: célina <[email protected]>

* code quality

---------

Co-authored-by: Abubakar Abid <[email protected]>
Co-authored-by: célina <[email protected]>

Author: 3 people

src/huggingface_hub/hf_api.py

@@ -75,6 +75,7 @@
     BadRequestError,
     GatedRepoError,
     HfHubHTTPError,
+    LocalTokenNotFoundError,
     RemoteEntryNotFoundError,
     RepositoryNotFoundError,
     RevisionNotFoundError,
@@ -1694,6 +1695,9 @@ def __init__(
         self.headers = headers
         self._thread_pool: Optional[ThreadPoolExecutor] = None
 
+        # /whoami-v2 is the only endpoint for which we may want to cache results
+        self._whoami_cache: dict[str, dict] = {}
+
     def run_as_future(self, fn: Callable[..., R], *args, **kwargs) -> Future[R]:
         """
         Run a method in the background and return a Future instance.
@@ -1735,39 +1739,74 @@ def run_as_future(self, fn: Callable[..., R], *args, **kwargs) -> Future[R]:
         return self._thread_pool.submit(fn, *args, **kwargs)
 
     @validate_hf_hub_args
-    def whoami(self, token: Union[bool, str, None] = None) -> dict:
+    def whoami(self, token: Union[bool, str, None] = None, *, cache: bool = False) -> dict:
         """
         Call HF API to know "whoami".
 
+        If passing `cache=True`, the result will be cached for subsequent calls for the duration of the Python process. This is useful if you plan to call
+        `whoami` multiple times as this endpoint is heavily rate-limited for security reasons.
+
         Args:
             token (`bool` or `str`, *optional*):
                 A valid user access token (string). Defaults to the locally saved
                 token, which is the recommended method for authentication (see
                 https://huggingface.co/docs/huggingface_hub/quick-start#authentication).
                 To disable authentication, pass `False`.
+            cache (`bool`, *optional*):
+                Whether to cache the result of the `whoami` call for subsequent calls.
+                If an error occurs during the first call, it won't be cached.
+                Defaults to `False`.
         """
         # Get the effective token using the helper function get_token
-        effective_token = token or self.token or get_token() or True
+        token = self.token if token is None else token
+        if token is False:
+            raise ValueError("Cannot use `token=False` with `whoami` method as it requires authentication.")
+        if token is True or token is None:
+            token = get_token()
+        if token is None:
+            raise LocalTokenNotFoundError(
+                "Token is required to call the /whoami-v2 endpoint, but no token found. You must provide a token or be logged in to "
+                "Hugging Face with `hf auth login` or `huggingface_hub.login`. See https://huggingface.co/settings/tokens."
+            )
+
+        if cache and (cached_token := self._whoami_cache.get(token)):
+            return cached_token
+
+        # Call Hub
+        output = self._inner_whoami(token=token)
+
+        # Cache result and return
+        if cache:
+            self._whoami_cache[token] = output
+        return output
+
+    def _inner_whoami(self, token: str) -> dict:
         r = get_session().get(
             f"{self.endpoint}/api/whoami-v2",
-            headers=self._build_hf_headers(token=effective_token),
+            headers=self._build_hf_headers(token=token),
         )
         try:
             hf_raise_for_status(r)
         except HfHubHTTPError as e:
             if e.response.status_code == 401:
                 error_message = "Invalid user token."
                 # Check which token is the effective one and generate the error message accordingly
-                if effective_token == _get_token_from_google_colab():
+                if token == _get_token_from_google_colab():
                     error_message += " The token from Google Colab vault is invalid. Please update it from the UI."
-                elif effective_token == _get_token_from_environment():
+                elif token == _get_token_from_environment():
                     error_message += (
                         " The token from HF_TOKEN environment variable is invalid. "
                         "Note that HF_TOKEN takes precedence over `hf auth login`."
                     )
-                elif effective_token == _get_token_from_file():
+                elif token == _get_token_from_file():
                     error_message += " The token stored is invalid. Please run `hf auth login` to update it."
                 raise HfHubHTTPError(error_message, response=e.response) from e
+            if e.response.status_code == 429:
+                error_message = (
+                    "You've hit the rate limit for the /whoami-v2 endpoint, which is intentionally strict for security reasons."
+                    " If you're calling it often, consider caching the response with `whoami(..., cache=True)`."
+                )
+                raise HfHubHTTPError(error_message, response=e.response) from e
             raise
         return r.json()
 

tests/test_hf_api.py

@@ -170,26 +170,68 @@ def test_file_exists(self):
 class HfApiEndpointsTest(HfApiCommonTest):
     def test_whoami_with_passing_token(self):
         info = self._api.whoami(token=self._token)
-        self.assertEqual(info["name"], USER)
-        self.assertEqual(info["fullname"], FULL_NAME)
-        self.assertIsInstance(info["orgs"], list)
+        assert info["name"] == USER
+        assert info["fullname"] == FULL_NAME
+        assert isinstance(info["orgs"], list)
         valid_org = [org for org in info["orgs"] if org["name"] == "valid_org"][0]
-        self.assertEqual(valid_org["fullname"], "Dummy Org")
+        assert valid_org["fullname"] == "Dummy Org"
 
-    @patch("huggingface_hub.utils._headers.get_token", return_value=TOKEN)
+    @patch("huggingface_hub.hf_api.get_token", return_value=TOKEN)
     def test_whoami_with_implicit_token_from_login(self, mock_get_token: Mock) -> None:
         """Test using `whoami` after a `hf auth login`."""
         with patch.object(self._api, "token", None):  # no default token
             info = self._api.whoami()
-        self.assertEqual(info["name"], USER)
+        assert info["name"] == USER
 
     @patch("huggingface_hub.utils._headers.get_token")
     def test_whoami_with_implicit_token_from_hf_api(self, mock_get_token: Mock) -> None:
         """Test using `whoami` with token from the HfApi client."""
         info = self._api.whoami()
-        self.assertEqual(info["name"], USER)
+        assert info["name"] == USER
         mock_get_token.assert_not_called()
 
+    def test_whoami_with_caching(self) -> None:
+        # Don't use class instance to avoid cache sharing
+        api = HfApi(endpoint=ENDPOINT_STAGING, token=TOKEN)
+        assert api._whoami_cache == {}
+
+        assert api.whoami(cache=True)["name"] == USER
+
+        # Value in cache
+        assert len(api._whoami_cache) == 1
+        assert TOKEN in api._whoami_cache
+        mocked_value = Mock()
+        api._whoami_cache[TOKEN] = mocked_value
+
+        # Call again => use cache
+        assert api.whoami(cache=True) == mocked_value
+
+        # Cache not shared between HfApi instances
+        api_bis = HfApi(endpoint=ENDPOINT_STAGING, token=TOKEN)
+        assert api_bis._whoami_cache == {}
+        assert api_bis.whoami(cache=True)["name"] == USER
+
+    def test_whoami_rate_limit_suggest_caching(self) -> None:
+        with patch("huggingface_hub.hf_api.hf_raise_for_status") as mock:
+            mock.side_effect = HfHubHTTPError(message="Fake error.", response=Mock(status_code=429))
+            with pytest.raises(
+                HfHubHTTPError, match=r".*consider caching the response with `whoami\(..., cache=True\)`.*"
+            ):
+                self._api.whoami()
+
+    def test_whoami_with_token_false(self):
+        """Test that using `token=False` raises an error.
+
+        Regression test for https://github.com/huggingface/huggingface_hub/pull/3568#discussion_r2557248898.
+
+        Before the fix, local token was used even when `token=False` was passed (which is not intended).
+        """
+        with self.assertRaises(ValueError):
+            self._api.whoami(token=False)
+
+        with self.assertRaises(ValueError):
+            HfApi(token=False).whoami()
+
     def test_delete_repo_error_message(self):
         # test for #751
         # See https://github.com/huggingface/huggingface_hub/issues/751