scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 156

fix(databricks): respect custom User-Agent set via extra_headers

Browse source 
Currently, `databricks_validate_environment()` unconditionally
overwrites the User-Agent header, discarding any value the caller
explicitly set via `extra_headers`. This makes it impossible to
pass a custom User-Agent string to Databricks endpoints.

This makes the assignment conditional: if a User-Agent is already
present in the headers dict (e.g. set via extra_headers), it is
preserved. Otherwise, the default user agent is built as before.

Databricks is the only provider whose `validate_environment`
actively sets User-Agent. Other providers only add Authorization
and Content-Type, leaving User-Agent untouched when set by the
caller. This fix aligns the Databricks behavior with that pattern.
This commit is contained in:
Jose Maria Vilaplana 2026-04-17 17:16:29 +02:00
parent b8f7d61400
commit 2987419117
2 changed files with 38 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/llms/databricks/common_utils.py
View file

@ -387,8 +387,10 @@ class DatabricksBase:
if api_key is not None:
headers["Authorization"] = f"Bearer {api_key}"
# Set User-Agent with optional custom prefix
headers["User-Agent"] = self._build_user_agent(custom_user_agent)
# Set User-Agent with optional custom prefix, but don't overwrite
# if the caller explicitly set one (e.g. via extra_headers)
if "User-Agent" not in headers:
headers["User-Agent"] = self._build_user_agent(custom_user_agent)
# Debug logging with redaction (never log actual tokens)
verbose_logger.debug(

34
tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
View file

@ -334,6 +334,40 @@ class TestValidateEnvironmentUserAgent:
assert headers["User-Agent"].startswith("mycompany_litellm/")
def test_extra_headers_user_agent_preserved(self, monkeypatch):
"""User-Agent set via headers dict (extra_headers) is not overwritten."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers={"User-Agent": "My-Custom-Agent/1.0"},
)
assert headers["User-Agent"] == "My-Custom-Agent/1.0"
def test_default_user_agent_when_not_in_headers(self, monkeypatch):
"""Default User-Agent is set when headers dict has no User-Agent."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers={},
)
assert headers["User-Agent"].startswith("litellm/")
class TestSDKPartnerTelemetry:
"""Test that SDK partner telemetry is registered."""