mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-06 10:21:32 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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)
|
||||
@@ -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__":
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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"])
|
||||
@@ -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"])
|
||||
Reference in New Issue
Block a user