diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 042f1556461..ac9dd5998e2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -53,23 +53,17 @@ except Exception: def get_default_headers() -> dict: """ Get default headers for HTTP requests. - - Respects litellm.disable_default_user_agent flag to allow users to disable - the automatic User-Agent header injection or override it completely. - - Returns: - dict: Default headers (may be empty if user disabled defaults) - """ - import litellm - - if getattr(litellm, "disable_default_user_agent", False): - return {} - - return { - "User-Agent": f"litellm/{version}", - } -# Initialize headers - will be empty if disable_default_user_agent is True + - Default: `User-Agent: litellm/{version}` + - Override: set `LITELLM_USER_AGENT` to fully override the header value. + """ + user_agent = os.environ.get("LITELLM_USER_AGENT") + if user_agent is not None: + return {"User-Agent": user_agent} + + return {"User-Agent": f"litellm/{version}"} + +# Initialize headers (User-Agent) headers = get_default_headers() # https://www.python-httpx.org/advanced/timeouts @@ -389,7 +383,7 @@ class AsyncHTTPHandler: shared_session=shared_session, ) - # Get default headers - will be empty if disable_default_user_agent is True + # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) default_headers = get_default_headers() return httpx.AsyncClient( @@ -920,7 +914,7 @@ class HTTPHandler: # /path/to/client.pem cert = os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate) - # Get default headers - will be empty if disable_default_user_agent is True + # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) default_headers = get_default_headers() if not disable_default_headers else None if client is None: diff --git a/litellm/llms/custom_httpx/httpx_handler.py b/litellm/llms/custom_httpx/httpx_handler.py index 1b61a312318..491cd97f7db 100644 --- a/litellm/llms/custom_httpx/httpx_handler.py +++ b/litellm/llms/custom_httpx/httpx_handler.py @@ -1,3 +1,4 @@ +import os from typing import Optional, Union import httpx @@ -10,28 +11,19 @@ except Exception: def get_default_headers() -> dict: """ Get default headers for HTTP requests. - - Respects litellm.disable_default_user_agent flag to allow users to disable - the automatic User-Agent header injection or override it completely. - - Returns: - dict: Default headers (may be empty if user disabled defaults) + + - Default: `User-Agent: litellm/{version}` + - Override: set `LITELLM_USER_AGENT` to fully override the header value. """ - import litellm - - if getattr(litellm, "disable_default_user_agent", False): - return {} - - return { - "User-Agent": f"litellm/{version}", - } - -# Initialize headers - will be empty if disable_default_user_agent is True -headers = get_default_headers() + user_agent = os.environ.get("LITELLM_USER_AGENT") + if user_agent is not None: + return {"User-Agent": user_agent} + return {"User-Agent": f"litellm/{version}"} class HTTPHandler: def __init__(self, concurrent_limit=1000): + headers = get_default_headers() # Create a client with a connection pool self.client = httpx.AsyncClient( limits=httpx.Limits( diff --git a/tests/test_litellm/test_user_agent_customization.py b/tests/test_litellm/test_user_agent_customization.py deleted file mode 100644 index af81b7e81e5..00000000000 --- a/tests/test_litellm/test_user_agent_customization.py +++ /dev/null @@ -1,209 +0,0 @@ -""" -Test User-Agent header customization -Tests for Issue #19017: Option to disable or customize default User-Agent header -""" - -import os -import sys -from unittest.mock import MagicMock, Mock, patch - -import pytest - -sys.path.insert(0, os.path.abspath("../..")) - -import litellm -from litellm import completion - - -def test_default_user_agent_is_set(): - """ - Test that by default, litellm sets the User-Agent header. - """ - from litellm.llms.custom_httpx.http_handler import get_default_headers - from litellm._version import version - - # Reset to default - litellm.disable_default_user_agent = False - - headers = get_default_headers() - assert "User-Agent" in headers - assert headers["User-Agent"] == f"litellm/{version}" - - -def test_disable_default_user_agent(): - """ - Test that setting litellm.disable_default_user_agent = True prevents - the default User-Agent header from being set. - """ - from litellm.llms.custom_httpx.http_handler import get_default_headers - - # Disable default User-Agent - litellm.disable_default_user_agent = True - - headers = get_default_headers() - assert headers == {} - - # Reset to default - litellm.disable_default_user_agent = False - - -def test_custom_user_agent_via_extra_headers(): - """ - Test that users can provide their own User-Agent via extra_headers. - This is critical for Claude Code credentials that require specific User-Agent. - """ - import httpx - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - # Disable default User-Agent - litellm.disable_default_user_agent = True - - # Create HTTP handler - handler = HTTPHandler() - - # Custom User-Agent for Claude Code - custom_headers = {"User-Agent": "Claude Code/1.0.0"} - - # Build request with custom headers - req = handler.client.build_request( - "POST", - "https://api.anthropic.com/v1/messages", - headers=custom_headers, - json={"test": "data"} - ) - - # Verify custom User-Agent is used - assert "User-Agent" in req.headers - assert req.headers["User-Agent"] == "Claude Code/1.0.0" - - # Reset to default - litellm.disable_default_user_agent = False - - -def test_env_var_disable_default_user_agent(): - """ - Test that LITELLM_DISABLE_DEFAULT_USER_AGENT environment variable works. - """ - from litellm.llms.custom_httpx.http_handler import get_default_headers - - # Test with env var - with patch.dict(os.environ, {"LITELLM_DISABLE_DEFAULT_USER_AGENT": "True"}): - # Manually set the flag (in real usage, this would be done at import time) - litellm.disable_default_user_agent = True - - headers = get_default_headers() - assert headers == {} - - # Reset to default - litellm.disable_default_user_agent = False - - -@pytest.mark.asyncio -async def test_async_http_handler_respects_disable_flag(): - """ - Test that AsyncHTTPHandler also respects the disable_default_user_agent flag. - """ - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_default_headers - - # Disable default User-Agent - litellm.disable_default_user_agent = True - - # Create async handler - handler = AsyncHTTPHandler() - - # Check that headers are empty - headers = get_default_headers() - assert headers == {} - - await handler.close() - - # Reset to default - litellm.disable_default_user_agent = False - - -def test_override_user_agent_without_disabling(): - """ - Test that users can override User-Agent by passing it in extra_headers, - even without disabling the default. - - Note: httpx will use the last header value when building the request. - """ - import httpx - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - # Default User-Agent is enabled - litellm.disable_default_user_agent = False - - # Create HTTP handler (will have default User-Agent) - handler = HTTPHandler() - - # Custom User-Agent provided in request - custom_headers = {"User-Agent": "MyCustomAgent/2.0.0"} - - # Build request with custom headers - httpx merges headers - req = handler.client.build_request( - "POST", - "https://api.anthropic.com/v1/messages", - headers=custom_headers, - json={"test": "data"} - ) - - # The custom User-Agent should override the default - assert "User-Agent" in req.headers - # httpx uses the request-level header over the client-level header - assert req.headers["User-Agent"] == "MyCustomAgent/2.0.0" - - -def test_claude_code_use_case(): - """ - Test the specific use case from Issue #19017: - Claude Code credentials that require specific User-Agent. - """ - # Disable default User-Agent globally - litellm.disable_default_user_agent = True - - # This is what the user would do in their code - custom_headers = {"User-Agent": "Claude Code"} - - # Verify the headers can be passed through - from litellm.llms.custom_httpx.http_handler import get_default_headers - default_headers = get_default_headers() - - # Default headers should be empty - assert default_headers == {} - - # Custom headers would be used in the actual request - assert custom_headers["User-Agent"] == "Claude Code" - - # Reset - litellm.disable_default_user_agent = False - - -def test_backwards_compatibility(): - """ - Test that existing code continues to work without any changes. - By default, the User-Agent header should still be set. - """ - from litellm.llms.custom_httpx.http_handler import get_default_headers - from litellm._version import version - - # Ensure default behavior is maintained - litellm.disable_default_user_agent = False - - headers = get_default_headers() - assert "User-Agent" in headers - assert headers["User-Agent"] == f"litellm/{version}" - - # Create HTTP handler - from litellm.llms.custom_httpx.http_handler import HTTPHandler - handler = HTTPHandler() - - # Build a request - req = handler.client.build_request( - "GET", - "https://api.openai.com/v1/models" - ) - - # Default User-Agent should be present - assert "User-Agent" in req.headers - assert "litellm" in req.headers["User-Agent"]