From a885901c9b0ebaf8c1c7eb794d8a03b6c5d2e905 Mon Sep 17 00:00:00 2001 From: Chad Retz Date: Tue, 5 May 2026 11:55:25 -0500 Subject: [PATCH] Add User-Agent header to outbound requests --- baseten/client/_inference.py | 54 +++++++++++++++++++++---------- baseten/client/_management.py | 56 ++++++++++++++++++++++----------- baseten/client/_user_agent.py | 25 +++++++++++++++ tests/client/test_client.py | 53 +++++++++++++++++++++++++++++-- tests/client/test_user_agent.py | 27 ++++++++++++++++ 5 files changed, 177 insertions(+), 38 deletions(-) create mode 100644 baseten/client/_user_agent.py create mode 100644 tests/client/test_user_agent.py diff --git a/baseten/client/_inference.py b/baseten/client/_inference.py index 76d0706..350338d 100644 --- a/baseten/client/_inference.py +++ b/baseten/client/_inference.py @@ -1,11 +1,13 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import Any import httpx import baseten.client.inferenceapi +from baseten.client._user_agent import with_user_agent @dataclass(frozen=True) @@ -19,6 +21,9 @@ class InferenceClientOptions: api_key: str """API key for authentication.""" + headers: Mapping[str, str] | None = None + """Additional headers to send on every request.""" + model_id: str | None = None """Model ID. Mutually exclusive with *chain_id*.""" @@ -84,40 +89,48 @@ def __init__( self, *, api_key: str, + headers: Mapping[str, str] | None = None, model_id: str | None = None, chain_id: str | None = None, environment: str | None = None, base_url_override: str | None = None, - http_client: httpx.Client | None = None, + http_client_override: httpx.Client | None = None, close_http_client_on_close: bool | None = None, ) -> None: """Create a new synchronous inference client. Args: api_key: API key for authentication. + headers: Additional headers to send on every request. model_id: Model ID. Mutually exclusive with *chain_id*. chain_id: Chain ID. Mutually exclusive with *model_id*. environment: Environment name for regional routing (e.g. ``"production"``). base_url_override: Override the computed base URL. When set, *model_id*, *chain_id*, and *environment* are ignored. - http_client: Pre-configured httpx client. When provided, the - caller is responsible for setting base URL and auth headers. + http_client_override: Pre-configured httpx client. When provided, + the caller is responsible for setting base URL and all + headers. close_http_client_on_close: Whether :meth:`close` should close the underlying HTTP client. Defaults to ``True`` when the - client is created internally, ``False`` when *http_client* - is provided. + client is created internally, ``False`` when + *http_client_override* is provided. """ self._options = InferenceClientOptions( api_key=api_key, + headers=headers, model_id=model_id, chain_id=chain_id, environment=environment, base_url_override=base_url_override, ) - if http_client is None: + if http_client_override is None: + request_headers: dict[str, str] = {**(headers or {})} + # Empty api_key is an advanced opt-out from sending Authorization. + if api_key != "": + request_headers["Authorization"] = f"Api-Key {api_key}" self._http_client = httpx.Client( base_url=self._options.base_url, - headers={"Authorization": f"Api-Key {api_key}"}, + headers=with_user_agent(request_headers), ) self.close_http_client_on_close = ( True @@ -125,7 +138,7 @@ def __init__( else close_http_client_on_close ) else: - self._http_client = http_client + self._http_client = http_client_override self.close_http_client_on_close = ( False if close_http_client_on_close is None @@ -191,41 +204,48 @@ def __init__( self, *, api_key: str, + headers: Mapping[str, str] | None = None, model_id: str | None = None, chain_id: str | None = None, environment: str | None = None, base_url_override: str | None = None, - http_client: httpx.AsyncClient | None = None, + http_client_override: httpx.AsyncClient | None = None, close_http_client_on_close: bool | None = None, ) -> None: """Create a new asynchronous inference client. Args: api_key: API key for authentication. + headers: Additional headers to send on every request. model_id: Model ID. Mutually exclusive with *chain_id*. chain_id: Chain ID. Mutually exclusive with *model_id*. environment: Environment name for regional routing (e.g. ``"production"``). base_url_override: Override the computed base URL. When set, *model_id*, *chain_id*, and *environment* are ignored. - http_client: Pre-configured httpx async client. When provided, - the caller is responsible for setting base URL and auth - headers. + http_client_override: Pre-configured httpx async client. When + provided, the caller is responsible for setting base URL + and all headers. close_http_client_on_close: Whether :meth:`close` should close the underlying HTTP client. Defaults to ``True`` when the - client is created internally, ``False`` when *http_client* - is provided. + client is created internally, ``False`` when + *http_client_override* is provided. """ self._options = InferenceClientOptions( api_key=api_key, + headers=headers, model_id=model_id, chain_id=chain_id, environment=environment, base_url_override=base_url_override, ) - if http_client is None: + if http_client_override is None: + request_headers: dict[str, str] = {**(headers or {})} + # Empty api_key is an advanced opt-out from sending Authorization. + if api_key != "": + request_headers["Authorization"] = f"Api-Key {api_key}" self._http_client = httpx.AsyncClient( base_url=self._options.base_url, - headers={"Authorization": f"Api-Key {api_key}"}, + headers=with_user_agent(request_headers), ) self.close_http_client_on_close = ( True @@ -233,7 +253,7 @@ def __init__( else close_http_client_on_close ) else: - self._http_client = http_client + self._http_client = http_client_override self.close_http_client_on_close = ( False if close_http_client_on_close is None diff --git a/baseten/client/_management.py b/baseten/client/_management.py index 5bfe315..e6a7013 100644 --- a/baseten/client/_management.py +++ b/baseten/client/_management.py @@ -1,11 +1,13 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import Any import httpx import baseten.client.managementapi +from baseten.client._user_agent import with_user_agent @dataclass(frozen=True) @@ -19,6 +21,9 @@ class ManagementClientOptions: api_key: str """API key for authentication.""" + headers: Mapping[str, str] | None = None + """Additional headers to send on every request.""" + base_url_override: str | None = None """Explicit base URL override, or ``None`` to use the default.""" @@ -46,30 +51,37 @@ def __init__( self, *, api_key: str, + headers: Mapping[str, str] | None = None, base_url_override: str | None = None, - http_client: httpx.Client | None = None, + http_client_override: httpx.Client | None = None, close_http_client_on_close: bool | None = None, ) -> None: """Create a new synchronous management client. Args: api_key: API key for authentication. + headers: Additional headers to send on every request. base_url_override: Override the default base URL. When ``None``, :meth:`default_base_url` is used. - http_client: Pre-configured httpx client. When provided, the - caller is responsible for setting base URL and auth headers. + http_client_override: Pre-configured httpx client. When provided, + the caller is responsible for setting base URL and all + headers. close_http_client_on_close: Whether :meth:`close` should close the underlying HTTP client. Defaults to ``True`` when the - client is created internally, ``False`` when *http_client* - is provided. + client is created internally, ``False`` when + *http_client_override* is provided. """ self._options = ManagementClientOptions( - api_key=api_key, base_url_override=base_url_override + api_key=api_key, headers=headers, base_url_override=base_url_override ) - if http_client is None: + if http_client_override is None: + request_headers: dict[str, str] = {**(headers or {})} + # Empty api_key is an advanced opt-out from sending Authorization. + if api_key != "": + request_headers["Authorization"] = f"Api-Key {api_key}" self._http_client = httpx.Client( base_url=self._options.base_url, - headers={"Authorization": f"Api-Key {api_key}"}, + headers=with_user_agent(request_headers), ) self.close_http_client_on_close = ( True @@ -77,7 +89,7 @@ def __init__( else close_http_client_on_close ) else: - self._http_client = http_client + self._http_client = http_client_override self.close_http_client_on_close = ( False if close_http_client_on_close is None @@ -132,31 +144,37 @@ def __init__( self, *, api_key: str, + headers: Mapping[str, str] | None = None, base_url_override: str | None = None, - http_client: httpx.AsyncClient | None = None, + http_client_override: httpx.AsyncClient | None = None, close_http_client_on_close: bool | None = None, ) -> None: """Create a new asynchronous management client. Args: api_key: API key for authentication. + headers: Additional headers to send on every request. base_url_override: Override the default base URL. When ``None``, :meth:`default_base_url` is used. - http_client: Pre-configured httpx async client. When provided, - the caller is responsible for setting base URL and auth - headers. + http_client_override: Pre-configured httpx async client. When + provided, the caller is responsible for setting base URL + and all headers. close_http_client_on_close: Whether :meth:`close` should close the underlying HTTP client. Defaults to ``True`` when the - client is created internally, ``False`` when *http_client* - is provided. + client is created internally, ``False`` when + *http_client_override* is provided. """ self._options = ManagementClientOptions( - api_key=api_key, base_url_override=base_url_override + api_key=api_key, headers=headers, base_url_override=base_url_override ) - if http_client is None: + if http_client_override is None: + request_headers: dict[str, str] = {**(headers or {})} + # Empty api_key is an advanced opt-out from sending Authorization. + if api_key != "": + request_headers["Authorization"] = f"Api-Key {api_key}" self._http_client = httpx.AsyncClient( base_url=self._options.base_url, - headers={"Authorization": f"Api-Key {api_key}"}, + headers=with_user_agent(request_headers), ) self.close_http_client_on_close = ( True @@ -164,7 +182,7 @@ def __init__( else close_http_client_on_close ) else: - self._http_client = http_client + self._http_client = http_client_override self.close_http_client_on_close = ( False if close_http_client_on_close is None diff --git a/baseten/client/_user_agent.py b/baseten/client/_user_agent.py new file mode 100644 index 0000000..32078ab --- /dev/null +++ b/baseten/client/_user_agent.py @@ -0,0 +1,25 @@ +import platform +from collections.abc import Mapping +from importlib.metadata import PackageNotFoundError, version + + +def _package_version() -> str: + try: + return version("baseten") + except PackageNotFoundError: + return "dev" + + +def user_agent_header() -> str: + """Build a User-Agent value like ``baseten-python/0.9.0 (Python/3.13.2; Linux)``.""" + return ( + f"baseten-python/{_package_version()} " + f"(Python/{platform.python_version()}; {platform.system()})" + ) + + +def with_user_agent(headers: Mapping[str, str]) -> dict[str, str]: + """Return a copy of ``headers`` with our User-Agent set, unless one is already present.""" + if any(key.lower() == "user-agent" for key in headers): + return dict(headers) + return {**headers, "User-Agent": user_agent_header()} diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 3a25613..bead587 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -8,6 +8,7 @@ AsyncInferenceClient, AsyncManagementClient, InferenceClient, + InferenceClientOptions, ManagementClient, ManagementClientOptions, ) @@ -45,7 +46,7 @@ def test_management_close_http_client_default_false_when_provided() -> None: import httpx http_client = httpx.Client() - client = ManagementClient(api_key="test-key", http_client=http_client) + client = ManagementClient(api_key="test-key", http_client_override=http_client) assert client.close_http_client_on_close is False client.close() http_client.close() @@ -58,11 +59,33 @@ def test_management_context_manager() -> None: def test_management_options_splat() -> None: opts = ManagementClientOptions( - api_key="test-key", base_url_override="https://custom.example.com" + api_key="test-key", + headers={"X-Custom": "v"}, + base_url_override="https://custom.example.com", ) client = ManagementClient(**dataclasses.asdict(opts)) assert client.options.api_key == "test-key" + assert client.options.headers == {"X-Custom": "v"} assert client.options.base_url == "https://custom.example.com" + assert client.http_client.headers["X-Custom"] == "v" + client.close() + + +def test_management_user_agent_default() -> None: + client = ManagementClient(api_key="test-key") + assert client.http_client.headers["User-Agent"].startswith("baseten-python/") + client.close() + + +def test_management_user_agent_user_override() -> None: + client = ManagementClient(api_key="test-key", headers={"User-Agent": "custom/1.0"}) + assert client.http_client.headers["User-Agent"] == "custom/1.0" + client.close() + + +def test_management_empty_api_key_skips_authorization() -> None: + client = ManagementClient(api_key="") + assert "authorization" not in client.http_client.headers client.close() @@ -125,6 +148,32 @@ def test_inference_model_and_chain_mutually_exclusive() -> None: InferenceClient(api_key="test-key", model_id="abc", chain_id="def") +def test_inference_options_splat() -> None: + opts = InferenceClientOptions( + api_key="test-key", + headers={"X-Custom": "v"}, + model_id="abc", + environment="prod-us", + ) + client = InferenceClient(**dataclasses.asdict(opts)) + assert client.options.headers == {"X-Custom": "v"} + assert client.options.base_url == "https://model-abc-prod-us.api.baseten.co" + assert client.http_client.headers["X-Custom"] == "v" + client.close() + + +def test_inference_user_agent_default() -> None: + client = InferenceClient(api_key="test-key", model_id="abc") + assert client.http_client.headers["User-Agent"].startswith("baseten-python/") + client.close() + + +def test_inference_empty_api_key_skips_authorization() -> None: + client = InferenceClient(api_key="", model_id="abc") + assert "authorization" not in client.http_client.headers + client.close() + + def test_inference_options_frozen() -> None: client = InferenceClient(api_key="test-key", model_id="abc") with pytest.raises(dataclasses.FrozenInstanceError): diff --git a/tests/client/test_user_agent.py b/tests/client/test_user_agent.py new file mode 100644 index 0000000..d2095c8 --- /dev/null +++ b/tests/client/test_user_agent.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import re + +from baseten.client._user_agent import user_agent_header, with_user_agent + + +def test_user_agent_header_format() -> None: + assert re.match(r"^baseten-python/\S+ \(Python/\S+; \S+\)$", user_agent_header()) + + +def test_with_user_agent_sets_when_absent() -> None: + out = with_user_agent({}) + assert out["User-Agent"] == user_agent_header() + + +def test_with_user_agent_does_not_overwrite_any_case() -> None: + out = with_user_agent({"user-agent": "custom/1.0"}) + assert out["user-agent"] == "custom/1.0" + assert "User-Agent" not in out + + +def test_with_user_agent_returns_copy() -> None: + src: dict[str, str] = {} + out = with_user_agent(src) + assert src == {} + assert "User-Agent" in out