Add user management functionality to Python client library & CLI (#10627)

* Add user mgmt functionality to client & CLI

* Add user mgmt client & CLI tests

* Add user mgmt client & CLI info to README.md's

* lint fixes: litellm/proxy/client/users.py

* Fix mypy errors
This commit is contained in:
Marc Abramowitz
2025-05-07 23:26:26 -07:00
committed by GitHub
parent d680feb5dc
commit fcaa4a9f30
9 changed files with 340 additions and 3 deletions
+29
View File
@@ -38,6 +38,7 @@ The client is organized into several resource clients for different functionalit
- `model_groups`: Model group management
- `keys`: API key management
- `credentials`: Credential management
- `users`: User management
## Chat Completions
@@ -157,6 +158,34 @@ groups = client.model_groups.list()
client.model_groups.delete(name="gpt4-group")
```
## Users Management
Manage users on your proxy:
```python
from litellm.proxy.client import UsersManagementClient
users = UsersManagementClient(base_url="http://localhost:4000", api_key="sk-test")
# List users
user_list = users.list_users()
# Get user info
user_info = users.get_user(user_id="u1")
# Create a new user
created = users.create_user({
"user_email": "a@b.com",
"user_role": "internal_user",
"user_alias": "Alice",
"teams": ["team1"],
"max_budget": 100.0
})
# Delete users
users.delete_user(["u1", "u2"])
```
## Low-Level HTTP Client
The client provides access to a low-level HTTP client for making direct requests
+2 -1
View File
@@ -3,5 +3,6 @@ from .chat import ChatClient
from .models import ModelsManagementClient
from .model_groups import ModelGroupsManagementClient
from .exceptions import UnauthorizedError
from .users import UsersManagementClient
__all__ = ["Client", "ChatClient", "ModelsManagementClient", "ModelGroupsManagementClient", "UnauthorizedError"]
__all__ = ["Client", "ChatClient", "ModelsManagementClient", "ModelGroupsManagementClient", "UsersManagementClient", "UnauthorizedError"]
+52
View File
@@ -254,6 +254,42 @@ Example:
litellm-proxy keys info --key sk-key1
```
### User Management
The CLI provides commands for managing users on your LiteLLM proxy server:
#### List Users
View all users:
```bash
litellm-proxy users list
```
#### Get User Info
Get information about a specific user:
```bash
litellm-proxy users get --id <user-id>
```
#### Create User
Create a new user:
```bash
litellm-proxy users create --email user@example.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0
```
#### Delete User
Delete one or more users by user_id:
```bash
litellm-proxy users delete <user-id-1> <user-id-2>
```
### Chat Commands
The CLI provides commands for interacting with chat models through your LiteLLM proxy server:
@@ -397,6 +433,22 @@ litellm-proxy http request POST /chat/completions \
-H "X-Custom-Header:value"
```
8. User management:
```bash
# List users
litellm-proxy users list
# Get user info
litellm-proxy users get --id u1
# Create a user
litellm-proxy users create --email a@b.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0
# Delete users
litellm-proxy users delete u1 u2
```
## Error Handling
The CLI will display appropriate error messages when:
@@ -0,0 +1,77 @@
import click
import rich
from ... import UsersManagementClient
@click.group()
def users():
"""Manage users on your LiteLLM proxy server"""
pass
@users.command("list")
@click.pass_context
def list_users(ctx: click.Context):
"""List all users"""
client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
users = client.list_users()
if isinstance(users, dict) and "users" in users:
users = users["users"]
if not users:
click.echo("No users found.")
return
from rich.table import Table
from rich.console import Console
table = Table(title="Users")
table.add_column("User ID", style="cyan")
table.add_column("Email", style="green")
table.add_column("Role", style="magenta")
table.add_column("Teams", style="yellow")
for user in users:
table.add_row(
str(user.get("user_id", "")),
str(user.get("user_email", "")),
str(user.get("user_role", "")),
", ".join(user.get("teams", []) or [])
)
console = Console()
console.print(table)
@users.command("get")
@click.option("--id", "user_id", help="ID of the user to retrieve")
@click.pass_context
def get_user(ctx: click.Context, user_id: str):
"""Get information about a specific user"""
client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
result = client.get_user(user_id=user_id)
rich.print_json(data=result)
@users.command("create")
@click.option("--email", required=True, help="User email")
@click.option("--role", default="internal_user", help="User role")
@click.option("--alias", default=None, help="User alias")
@click.option("--team", multiple=True, help="Team IDs (can specify multiple)")
@click.option("--max-budget", type=float, default=None, help="Max budget for user")
@click.pass_context
def create_user(ctx: click.Context, email, role, alias, team, max_budget):
"""Create a new user"""
client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
user_data = {
"user_email": email,
"user_role": role,
}
if alias:
user_data["user_alias"] = alias
if team:
user_data["teams"] = list(team)
if max_budget is not None:
user_data["max_budget"] = max_budget
result = client.create_user(user_data)
rich.print_json(data=result)
@users.command("delete")
@click.argument("user_ids", nargs=-1)
@click.pass_context
def delete_user(ctx: click.Context, user_ids):
"""Delete one or more users by user_id"""
client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
result = client.delete_user(list(user_ids))
rich.print_json(data=result)
+3
View File
@@ -10,6 +10,7 @@ from .commands.credentials import credentials
from .commands.chat import chat
from .commands.http import http
from .commands.keys import keys
from .commands.users import users
@click.group()
@@ -44,6 +45,8 @@ cli.add_command(chat)
cli.add_command(http)
# Add the keys command group
cli.add_command(keys)
# Add the users command group
cli.add_command(users)
if __name__ == "__main__":
+2 -2
View File
@@ -4,7 +4,7 @@ import requests
class UnauthorizedError(Exception):
"""Exception raised when the API returns a 401 Unauthorized response."""
def __init__(self, orig_exception: requests.exceptions.HTTPError):
def __init__(self, orig_exception: requests.exceptions.HTTPError | str):
self.orig_exception = orig_exception
super().__init__(str(orig_exception))
@@ -12,6 +12,6 @@ class UnauthorizedError(Exception):
class NotFoundError(Exception):
"""Exception raised when the API returns a 404 Not Found response or indicates a resource was not found."""
def __init__(self, orig_exception: requests.exceptions.HTTPError):
def __init__(self, orig_exception: requests.exceptions.HTTPError | str):
self.orig_exception = orig_exception
super().__init__(str(orig_exception))
+54
View File
@@ -0,0 +1,54 @@
import requests
from typing import List, Dict, Any, Optional
from .exceptions import UnauthorizedError, NotFoundError
class UsersManagementClient:
def __init__(self, base_url: str, api_key: Optional[str] = None):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
def _get_headers(self) -> Dict[str, str]:
headers = {"Content-Type": "application/json"}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
return headers
def list_users(self, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
"""List users (GET /user/list)"""
url = f"{self.base_url}/user/list"
response = requests.get(url, headers=self._get_headers(), params=params)
if response.status_code == 401:
raise UnauthorizedError(response.text)
response.raise_for_status()
return response.json().get("users", response.json())
def get_user(self, user_id: Optional[str] = None) -> Dict[str, Any]:
"""Get user info (GET /user/info)"""
url = f"{self.base_url}/user/info"
params = {"user_id": user_id} if user_id else {}
response = requests.get(url, headers=self._get_headers(), params=params)
if response.status_code == 401:
raise UnauthorizedError(response.text)
if response.status_code == 404:
raise NotFoundError(response.text)
response.raise_for_status()
return response.json()
def create_user(self, user_data: Dict[str, Any]) -> Dict[str, Any]:
"""Create a new user (POST /user/new)"""
url = f"{self.base_url}/user/new"
response = requests.post(url, headers=self._get_headers(), json=user_data)
if response.status_code == 401:
raise UnauthorizedError(response.text)
response.raise_for_status()
return response.json()
def delete_user(self, user_ids: List[str]) -> Dict[str, Any]:
"""Delete users (POST /user/delete)"""
url = f"{self.base_url}/user/delete"
response = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids})
if response.status_code == 401:
raise UnauthorizedError(response.text)
response.raise_for_status()
return response.json()
@@ -0,0 +1,54 @@
import pytest
from click.testing import CliRunner
from unittest.mock import patch
from litellm.proxy.client.cli import cli
@pytest.fixture
def cli_runner():
return CliRunner()
@pytest.fixture(autouse=True)
def mock_env():
with patch.dict("os.environ", {"LITELLM_PROXY_URL": "http://localhost:4000", "LITELLM_PROXY_API_KEY": "sk-test"}):
yield
@pytest.fixture
def mock_users_client():
with patch("litellm.proxy.client.cli.commands.users.UsersManagementClient") as MockClient:
yield MockClient
def test_users_list(cli_runner, mock_users_client):
mock_users_client.return_value.list_users.return_value = [
{"user_id": "u1", "user_email": "a@b.com", "user_role": "internal_user", "teams": ["t1"]},
{"user_id": "u2", "user_email": "b@b.com", "user_role": "proxy_admin", "teams": ["t2", "t3"]},
]
result = cli_runner.invoke(cli, ["users", "list"])
assert result.exit_code == 0
assert "u1" in result.output
assert "a@b.com" in result.output
assert "proxy_admin" in result.output
assert "t3" in result.output
mock_users_client.return_value.list_users.assert_called_once()
def test_users_get(cli_runner, mock_users_client):
mock_users_client.return_value.get_user.return_value = {"user_id": "u1", "user_email": "a@b.com"}
result = cli_runner.invoke(cli, ["users", "get", "--id", "u1"])
assert result.exit_code == 0
assert '"user_id": "u1"' in result.output
assert '"user_email": "a@b.com"' in result.output
mock_users_client.return_value.get_user.assert_called_once_with(user_id="u1")
def test_users_create(cli_runner, mock_users_client):
mock_users_client.return_value.create_user.return_value = {"user_id": "u1", "user_email": "a@b.com"}
result = cli_runner.invoke(cli, ["users", "create", "--email", "a@b.com", "--role", "internal_user"])
assert result.exit_code == 0
assert '"user_id": "u1"' in result.output
assert '"user_email": "a@b.com"' in result.output
mock_users_client.return_value.create_user.assert_called_once()
def test_users_delete(cli_runner, mock_users_client):
mock_users_client.return_value.delete_user.return_value = {"deleted": 1}
result = cli_runner.invoke(cli, ["users", "delete", "u1", "u2"])
assert result.exit_code == 0
assert '"deleted": 1' in result.output
mock_users_client.return_value.delete_user.assert_called_once_with(["u1", "u2"])
+67
View File
@@ -0,0 +1,67 @@
import pytest
from unittest.mock import patch, MagicMock
from litellm.proxy.client.users import UsersManagementClient, UnauthorizedError, NotFoundError
@pytest.fixture
def client():
return UsersManagementClient(base_url="http://localhost:4000", api_key="sk-test")
@patch("requests.get")
def test_list_users_success(mock_get, client):
mock_get.return_value.status_code = 200
mock_get.return_value.json.return_value = {"users": [{"user_id": "u1"}]}
users = client.list_users()
assert users == [{"user_id": "u1"}]
mock_get.assert_called_once()
@patch("requests.get")
def test_list_users_unauthorized(mock_get, client):
mock_get.return_value.status_code = 401
mock_get.return_value.text = "unauthorized"
with pytest.raises(UnauthorizedError):
client.list_users()
@patch("requests.get")
def test_get_user_success(mock_get, client):
mock_get.return_value.status_code = 200
mock_get.return_value.json.return_value = {"user_id": "u1"}
user = client.get_user(user_id="u1")
assert user["user_id"] == "u1"
mock_get.assert_called_once()
@patch("requests.get")
def test_get_user_404(mock_get, client):
mock_get.return_value.status_code = 404
mock_get.return_value.text = "not found"
with pytest.raises(NotFoundError):
client.get_user(user_id="u1")
@patch("requests.post")
def test_create_user_success(mock_post, client):
mock_post.return_value.status_code = 200
mock_post.return_value.json.return_value = {"user_id": "u1"}
user = client.create_user({"user_email": "a@b.com"})
assert user["user_id"] == "u1"
mock_post.assert_called_once()
@patch("requests.post")
def test_create_user_unauthorized(mock_post, client):
mock_post.return_value.status_code = 401
mock_post.return_value.text = "unauthorized"
with pytest.raises(UnauthorizedError):
client.create_user({"user_email": "a@b.com"})
@patch("requests.post")
def test_delete_user_success(mock_post, client):
mock_post.return_value.status_code = 200
mock_post.return_value.json.return_value = {"deleted": 1}
result = client.delete_user(["u1"])
assert result["deleted"] == 1
mock_post.assert_called_once()
@patch("requests.post")
def test_delete_user_unauthorized(mock_post, client):
mock_post.return_value.status_code = 401
mock_post.return_value.text = "unauthorized"
with pytest.raises(UnauthorizedError):
client.delete_user(["u1"])