scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(xai): support multi-account SuperGrok OAuth with automatic failover

Browse source 
Add per-deployment xai_oauth_token_file so multiple SuperGrok accounts
can share one proxy. Map xAI 403 spending-limit to RateLimitError so the
router can fail over to another deployment

Co-authored-by: HX <HXhlx@users.noreply.github.com>
This commit is contained in:
hx 2026-09-08 11:16:18 +00:00
parent 1af7a403c6
commit dfab8e2969
10 changed files with 281 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -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}",

1
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -55,6 +55,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
"itpm",
"otpm",
"use_xai_oauth",
"xai_oauth_token_file",
}
)
| AWS_CREDENTIAL_KWARGS_KEYS

8
litellm/llms/xai/chat/transformation.py
View file

@ -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,

19
litellm/llms/xai/oauth.py
View file

@ -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:

7
litellm/llms/xai/responses/transformation.py
View file

@ -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

1
litellm/main.py
View file

@ -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"),

13
litellm/proxy/proxy_cli.py
View file

@ -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"):

8
litellm/types/router.py
View file

@ -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

1
litellm/types/utils.py
View file

@ -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",

224
tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py Normal file
View file

@ -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="",
)