feat: support HF_ENDPOINT environment variable for model downloads

fastembed/common/model_management.py

@@ -1,14 +1,14 @@
-import os
-import time
 import json
+import os
 import shutil
 import tarfile
+import time
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, TypeVar, Generic
+from typing import Any, Generic, TypeVar
 
 import requests
-from huggingface_hub import snapshot_download, model_info, list_repo_tree
+from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.hf_api import RepoFile
 from huggingface_hub.utils import (
     RepositoryNotFoundError,
@@ -17,6 +17,7 @@
 )
 from loguru import logger
 from tqdm import tqdm
+
 from fastembed.common.model_description import BaseModelDescription
 
 T = TypeVar("T", bound=BaseModelDescription)
@@ -214,6 +215,9 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]])
         snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
         metadata_file = snapshot_dir / cls.METADATA_FILE
 
+        hf_endpoint = os.environ.get("HF_ENDPOINT")
+        hf_api = HfApi(endpoint=hf_endpoint)
+
         if local_files_only:
             disable_progress_bars()
             if metadata_file.exists():
@@ -228,12 +232,15 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]])
                 allow_patterns=allow_patterns,
                 cache_dir=cache_dir,
                 local_files_only=local_files_only,
+                endpoint=hf_endpoint,
                 **kwargs,
             )
             return result
 
-        repo_revision = model_info(hf_source_repo).sha
-        repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model"))
+        repo_revision = hf_api.model_info(hf_source_repo).sha
+        repo_tree = list(
+            hf_api.list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model")
+        )
 
         allowed_extensions = {".json", ".onnx", ".txt"}
         repo_files = (
@@ -260,6 +267,7 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]])
             allow_patterns=allow_patterns,
             cache_dir=cache_dir,
             local_files_only=local_files_only,
+            endpoint=hf_endpoint,
             **kwargs,
         )
 

tests/test_common.py

@@ -1,10 +1,14 @@
+import os
+from unittest.mock import MagicMock, patch
+
 from fastembed import (
-    TextEmbedding,
-    SparseTextEmbedding,
     ImageEmbedding,
     LateInteractionMultimodalEmbedding,
     LateInteractionTextEmbedding,
+    SparseTextEmbedding,
+    TextEmbedding,
 )
+from fastembed.common.model_management import ModelManagement
 
 
 def test_text_list_supported_models():
@@ -28,3 +32,75 @@ def test_text_list_supported_models():
         assert "model_file" in description and description["model_file"]
         assert "sources" in description and description["sources"]
         assert "hf" in description["sources"] or "url" in description["sources"]
+
+
+def _run_download_with_mocks(tmp_path, extra_env):
+    """Run download_files_from_huggingface with all network calls mocked out.
+
+    Uses autospec=True so mocks are checked against the real huggingface_hub
+    signatures - a kwarg that the real API doesn't accept (e.g. passing
+    `endpoint` to a bound HfApi method instead of its constructor) raises a
+    TypeError here, the same as it would against the real library.
+    """
+    mock_hf_api_instance = MagicMock()
+    mock_hf_api_instance.model_info.return_value = MagicMock(sha="abc123")
+    mock_hf_api_instance.list_repo_tree.return_value = []
+
+    with (
+        patch.dict(os.environ, extra_env),
+        patch(
+            "fastembed.common.model_management.HfApi",
+            autospec=True,
+            return_value=mock_hf_api_instance,
+        ) as mock_hf_api_cls,
+        patch(
+            "fastembed.common.model_management.snapshot_download",
+            autospec=True,
+            return_value=str(tmp_path),
+        ) as mock_sd,
+        # skip post-download metadata verification so the function completes cleanly
+        patch.object(ModelManagement, "METADATA_FILE", "__nonexistent__"),
+    ):
+        ModelManagement.download_files_from_huggingface(
+            hf_source_repo="test-org/test-model",
+            cache_dir=str(tmp_path),
+            extra_patterns=["*.onnx"],
+        )
+        return mock_hf_api_cls, mock_hf_api_instance, mock_sd
+
+
+def test_hf_endpoint_forwarded_to_hub_calls(tmp_path):
+    """HF_ENDPOINT env var must be forwarded to HfApi and snapshot_download."""
+    custom_endpoint = "https://hf-mirror.example.com"
+    mock_hf_api_cls, mock_hf_api_instance, mock_sd = _run_download_with_mocks(
+        tmp_path, {"HF_ENDPOINT": custom_endpoint}
+    )
+
+    _, api_kwargs = mock_hf_api_cls.call_args
+    assert api_kwargs.get("endpoint") == custom_endpoint, (
+        f"HfApi should be constructed with endpoint={custom_endpoint!r}, got {api_kwargs}"
+    )
+    mock_hf_api_instance.model_info.assert_called_once()
+    mock_hf_api_instance.list_repo_tree.assert_called_once()
+
+    _, sd_kwargs = mock_sd.call_args
+    assert sd_kwargs.get("endpoint") == custom_endpoint, (
+        f"snapshot_download should receive endpoint={custom_endpoint!r}, got {sd_kwargs}"
+    )
+
+
+def test_no_hf_endpoint_no_extra_kwarg(tmp_path):
+    """When HF_ENDPOINT is not set, endpoint must be None for HfApi and snapshot_download."""
+    env_without_hf_endpoint = {k: v for k, v in os.environ.items() if k != "HF_ENDPOINT"}
+    with patch.dict(os.environ, env_without_hf_endpoint, clear=True):
+        mock_hf_api_cls, _, mock_sd = _run_download_with_mocks(tmp_path, {})
+
+    _, api_kwargs = mock_hf_api_cls.call_args
+    assert api_kwargs.get("endpoint") is None, (
+        f"HfApi endpoint should be None when HF_ENDPOINT is unset, got {api_kwargs}"
+    )
+
+    _, sd_kwargs = mock_sd.call_args
+    assert sd_kwargs.get("endpoint") is None, (
+        f"snapshot_download endpoint should be None when HF_ENDPOINT is unset, got {sd_kwargs}"
+    )