diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 82708d412c9..a444e09c6b3 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -413,6 +413,14 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, ) + elif original_exception.status_code == 403 and "spending-limit" in error_str: + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) elif original_exception.status_code == 404: raise NotFoundError( message=f"NotFoundError: {exception_provider} - {message}", diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 1fd79db15a6..22aa7d79c21 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -55,6 +55,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", + "xai_oauth_token_file", } ) | AWS_CREDENTIAL_KWARGS_KEYS diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 747ee0b6c49..6e6353d015e 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -68,8 +68,10 @@ class XAIChatConfig(OpenAIGPTConfig): dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: + raw_token_file: Final = (litellm_params or {}).get("xai_oauth_token_file") + token_file: Final = raw_token_file if isinstance(raw_token_file, str) else None try: - headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}" + headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator(auth_file=token_file).get_access_token()}" except XAIOAuthError as exc: raise AuthenticationError( model=model, @@ -103,7 +105,9 @@ class XAIChatConfig(OpenAIGPTConfig): dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: - api_base = XAIOAuthAuthenticator().get_api_base() + raw_token_file: Final = (litellm_params or {}).get("xai_oauth_token_file") + token_file: Final = raw_token_file if isinstance(raw_token_file, str) else None + api_base = XAIOAuthAuthenticator(auth_file=token_file).get_api_base() return super().get_complete_url( api_base=api_base, diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py index e8196ec6cb9..ae8d47c2042 100644 --- a/litellm/llms/xai/oauth.py +++ b/litellm/llms/xai/oauth.py @@ -121,9 +121,22 @@ class _CallbackServer(HTTPServer): class XAIOAuthAuthenticator: - def __init__(self, http_client: httpx.Client | HTTPHandler | None = None) -> None: - self.token_dir = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") - self.auth_file = os.path.join(self.token_dir, get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json") + def __init__( + self, + http_client: httpx.Client | HTTPHandler | None = None, + auth_file: str | None = None, + token_dir: str | None = None, + ) -> None: + self.token_dir = ( + token_dir or get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") + ) + if auth_file: + self.auth_file = auth_file if os.path.isabs(auth_file) else os.path.join(self.token_dir, auth_file) + else: + self.auth_file = os.path.join( + self.token_dir, + get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json", + ) self.http_client = http_client def get_api_base(self) -> str: diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 646d6798783..a1e53da1ab3 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -231,8 +231,9 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ) if should_use_xai_oauth(litellm_params.model_dump()): + token_file: Final = litellm_params.xai_oauth_token_file try: - api_key = XAIOAuthAuthenticator().get_access_token() + api_key = XAIOAuthAuthenticator(auth_file=token_file).get_access_token() except XAIOAuthError as exc: raise AuthenticationError( model=model, @@ -268,7 +269,9 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): api_key: Final = XAIModelInfo.get_api_key(litellm_params.get("api_key"), legacy_generic_before_env=True) if should_use_xai_oauth(litellm_params) and not api_key: - api_base = XAIOAuthAuthenticator().get_api_base() + raw_token_file: Final = litellm_params.get("xai_oauth_token_file") + token_file: Final = raw_token_file if isinstance(raw_token_file, str) else None + api_base = XAIOAuthAuthenticator(auth_file=token_file).get_api_base() else: api_base = api_base or litellm.api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a5..35b78aa729c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5515,6 +5515,7 @@ def completion( tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), + xai_oauth_token_file=kwargs.get("xai_oauth_token_file"), gigachat_scope=kwargs.get("gigachat_scope"), gigachat_auth_url=kwargs.get("gigachat_auth_url"), gigachat_access_token=kwargs.get("gigachat_access_token"), diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e245367b1b4..f5ba98e1dde 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -997,10 +997,19 @@ def run_server( prometheus_metrics_port: int | None, ): if cli_args: - if cli_args == ("xai-oauth", "login"): + if len(cli_args) >= 2 and cli_args[0] == "xai-oauth" and cli_args[1] == "login": from litellm.llms.xai.oauth import XAIOAuthAuthenticator + from litellm.secret_managers.main import get_secret_str - authenticator: Final = XAIOAuthAuthenticator() + account: Final = cli_args[2] if len(cli_args) >= 3 else None + token_dir: Final = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser( + "~/.config/litellm/xai_oauth" + ) + authenticator: Final = ( + XAIOAuthAuthenticator(auth_file=os.path.join(token_dir, f"auth-{account}.json")) + if account + else XAIOAuthAuthenticator() + ) auth_data: Final = authenticator.login() click.echo(f"xAI OAuth login successful. Credentials saved to {authenticator.auth_file}.") if auth_data.get("expires_at"): diff --git a/litellm/types/router.py b/litellm/types/router.py index 0db482d8a58..c294a270408 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -344,6 +344,14 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): default=False, description="Use stored xAI OAuth credentials when no xAI API key is configured.", ) + xai_oauth_token_file: str | None = Field( + default=None, + description=( + "Per-deployment xAI OAuth token file path. When set with use_xai_oauth=True, " + "the request reads that account's OAuth credentials instead of the global " + "token file, enabling multi-account SuperGrok routing." + ), + ) model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: bool | None = False model_info: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9b5fb08a45f..3b2d6a636a2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3790,6 +3790,7 @@ all_litellm_params = ( "enable_tag_filtering", "enable_json_schema_validation", "use_xai_oauth", + "xai_oauth_token_file", "auto_router_config_path", "auto_router_config", "auto_router_default_model", diff --git a/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py b/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py new file mode 100644 index 00000000000..5f3bcde819f --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py @@ -0,0 +1,224 @@ +import json +import os +import time + +import pytest +from click.testing import CliRunner + +import litellm +from litellm.litellm_core_utils.exception_mapping_utils import _map_openai_exception +from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.oauth import XAIOAuthAuthenticator +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams + + +def test_authenticator_accepts_explicit_absolute_auth_file(tmp_path): + custom = tmp_path / "custom-dir" / "alice.json" + custom.parent.mkdir(parents=True) + custom.write_text( + json.dumps( + { + "access_token": "alice-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + ) + + auth = XAIOAuthAuthenticator(auth_file=str(custom)) + + assert auth.auth_file == str(custom) + assert auth.get_access_token() == "alice-token" + + +def test_authenticator_relative_auth_file_joins_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + (token_dir / "auth-bob.json").write_text( + json.dumps( + { + "access_token": "bob-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + auth = XAIOAuthAuthenticator(auth_file="auth-bob.json") + + assert auth.auth_file == str(token_dir / "auth-bob.json") + assert auth.get_access_token() == "bob-token" + + +def test_authenticator_explicit_auth_file_overrides_env(tmp_path, monkeypatch): + env_file = tmp_path / "env.json" + env_file.write_text( + json.dumps( + { + "access_token": "env-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + explicit_file = tmp_path / "explicit.json" + explicit_file.write_text( + json.dumps( + { + "access_token": "explicit-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + monkeypatch.setenv("XAI_OAUTH_AUTH_FILE", str(env_file)) + + auth = XAIOAuthAuthenticator(auth_file=str(explicit_file)) + + assert auth.get_access_token() == "explicit-token" + + +def test_chat_config_uses_per_deployment_token_file(tmp_path): + alice_file = tmp_path / "auth-alice.json" + alice_file.write_text( + json.dumps( + { + "access_token": "alice-chat-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + ) + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={ + "use_xai_oauth": True, + "xai_oauth_token_file": str(alice_file), + }, + api_key=None, + ) + + assert headers["Authorization"] == "Bearer alice-chat-token" + + +def test_chat_config_multi_deployment_token_isolation(tmp_path): + alice = tmp_path / "auth-alice.json" + alice.write_text( + json.dumps( + { + "access_token": "alice-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + bob = tmp_path / "auth-bob.json" + bob.write_text( + json.dumps( + { + "access_token": "bob-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + + headers_alice = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={ + "use_xai_oauth": True, + "xai_oauth_token_file": str(alice), + }, + api_key=None, + ) + headers_bob = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={ + "use_xai_oauth": True, + "xai_oauth_token_file": str(bob), + }, + api_key=None, + ) + + assert headers_alice["Authorization"] == "Bearer alice-token" + assert headers_bob["Authorization"] == "Bearer bob-token" + + +def test_responses_config_uses_per_deployment_token_file(tmp_path): + bob = tmp_path / "auth-bob.json" + bob.write_text( + json.dumps( + { + "access_token": "bob-responses-token", + "refresh_token": "r", + "expires_at": time.time() + 3600, + } + ) + ) + + headers = XAIResponsesAPIConfig().validate_environment( + headers={}, + model="grok-4", + litellm_params=GenericLiteLLMParams( + use_xai_oauth=True, + xai_oauth_token_file=str(bob), + ), + ) + + assert headers["Authorization"] == "Bearer bob-responses-token" + + +def test_proxy_cli_xai_oauth_login_with_account(monkeypatch, tmp_path): + from litellm.proxy.proxy_cli import run_server + + captured: dict[str, str | None] = {} + + class FakeAuthenticator: + def __init__(self, http_client=None, auth_file=None, token_dir=None): + captured["auth_file"] = auth_file + self.auth_file = auth_file or "/tmp/xai-oauth-auth.json" + + def login(self): + return {"expires_at": 1234567890} + + monkeypatch.setattr("litellm.llms.xai.oauth.XAIOAuthAuthenticator", FakeAuthenticator) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) + + result = CliRunner().invoke(run_server, ["xai-oauth", "login", "alice"]) + + assert result.exit_code == 0 + expected = os.path.join(str(tmp_path), "auth-alice.json") + assert captured["auth_file"] == expected + assert expected in result.output + + +def test_spending_limit_403_maps_to_rate_limit_error(): + class FakeHTTPError(Exception): + def __init__(self): + self.status_code = 403 + self.message = "personal-team-blocked:spending-limit" + self.response = None + super().__init__(self.message) + + with pytest.raises(litellm.RateLimitError): + _map_openai_exception( + model="grok-4", + original_exception=FakeHTTPError(), + custom_llm_provider="xai", + error_str="personal-team-blocked:spending-limit", + exception_type="APIStatusError", + exception_provider="XaiException", + extra_information="", + )