feat: Adds a new environment variable HF_HUB_USER_AGENT_ORIGIN to set…

… origin of calls in user-agent (#2869)

* feat: Adds a new environment variable HF_HUB_USER_AGENT_ORIGIN to set origin of calls in user-agent. This can be used to track calls for partners for collaborations.

---------

Co-authored-by: Célina <hanouticelina@gmail.com>

src/huggingface_hub/constants.py

@@ -194,6 +194,9 @@ def _as_int(value: Optional[str]) -> Optional[int]:
 # Used to override the get request timeout on a system level
 HF_HUB_DOWNLOAD_TIMEOUT: int = _as_int(os.environ.get("HF_HUB_DOWNLOAD_TIMEOUT")) or DEFAULT_DOWNLOAD_TIMEOUT
 
+# Allows to add information about the requester in the user-agent (eg. partner name)
+HF_HUB_USER_AGENT_ORIGIN: Optional[str] = os.environ.get("HF_HUB_USER_AGENT_ORIGIN")
+
 # List frameworks that are handled by the InferenceAPI service. Useful to scan endpoints and check which models are
 # deployed and running. Since 95% of the models are using the top 4 frameworks listed below, we scan only those by
 # default. We still keep the full list of supported frameworks in case we want to scan all of them.

src/huggingface_hub/utils/_headers.py

@@ -213,6 +213,11 @@ def _http_user_agent(
     elif isinstance(user_agent, str):
         ua += "; " + user_agent
 
+    # Retrieve user-agent origin headers from environment variable
+    origin = constants.HF_HUB_USER_AGENT_ORIGIN
+    if origin is not None:
+        ua += "; origin/" + origin
+
     return _deduplicate_user_agent(ua)
 
 

tests/test_utils_headers.py

@@ -141,3 +141,17 @@ def test_user_agent_deduplicate(self) -> None:
             # 4. order is preserved
             "python/3.7; python/3.8; hf_hub/0.12; transformers/None; diffusers/0.12.1",
         )
+
+    @patch("huggingface_hub.utils._telemetry.constants.HF_HUB_USER_AGENT_ORIGIN", "custom-origin")
+    def test_user_agent_with_origin(self) -> None:
+        self.assertTrue(self._get_user_agent().endswith("origin/custom-origin"))
+
+    @patch("huggingface_hub.utils._telemetry.constants.HF_HUB_USER_AGENT_ORIGIN", "custom-origin")
+    def test_user_agent_with_origin_and_user_agent(self) -> None:
+        self.assertTrue(
+            self._get_user_agent(user_agent={"a": "b", "c": "d"}).endswith("a/b; c/d; origin/custom-origin")
+        )
+
+    @patch("huggingface_hub.utils._telemetry.constants.HF_HUB_USER_AGENT_ORIGIN", "custom-origin")
+    def test_user_agent_with_origin_and_user_agent_str(self) -> None:
+        self.assertTrue(self._get_user_agent(user_agent="a/b;c/d").endswith("a/b; c/d; origin/custom-origin"))