From fcaa4a9f30db187dec6562e2cba324622eb9be0a Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Wed, 7 May 2025 23:26:26 -0700 Subject: [PATCH] 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 --- litellm/proxy/client/README.md | 29 +++++++ litellm/proxy/client/__init__.py | 3 +- litellm/proxy/client/cli/README.md | 52 +++++++++++++ litellm/proxy/client/cli/commands/users.py | 77 +++++++++++++++++++ litellm/proxy/client/cli/main.py | 3 + litellm/proxy/client/exceptions.py | 4 +- litellm/proxy/client/users.py | 54 +++++++++++++ .../proxy/client/cli/test_users_commands.py | 54 +++++++++++++ tests/litellm/proxy/client/test_users.py | 67 ++++++++++++++++ 9 files changed, 340 insertions(+), 3 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/users.py create mode 100644 litellm/proxy/client/users.py create mode 100644 tests/litellm/proxy/client/cli/test_users_commands.py create mode 100644 tests/litellm/proxy/client/test_users.py diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index c83944288f..06340e6dc1 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -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 diff --git a/litellm/proxy/client/__init__.py b/litellm/proxy/client/__init__.py index 6eb8cd1d42..2d71458ed6 100644 --- a/litellm/proxy/client/__init__.py +++ b/litellm/proxy/client/__init__.py @@ -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"] diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 12e8428caf..74d1e654f0 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -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 +``` + +#### 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 +``` + ### 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: diff --git a/litellm/proxy/client/cli/commands/users.py b/litellm/proxy/client/cli/commands/users.py new file mode 100644 index 0000000000..9887a8d0df --- /dev/null +++ b/litellm/proxy/client/cli/commands/users.py @@ -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) \ No newline at end of file diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index b4ddc900af..f2ee28c44b 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -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__": diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index 4975fb8452..ee6c418fbe 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -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)) diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py new file mode 100644 index 0000000000..66e2d76bee --- /dev/null +++ b/litellm/proxy/client/users.py @@ -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() diff --git a/tests/litellm/proxy/client/cli/test_users_commands.py b/tests/litellm/proxy/client/cli/test_users_commands.py new file mode 100644 index 0000000000..3489855e81 --- /dev/null +++ b/tests/litellm/proxy/client/cli/test_users_commands.py @@ -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"]) \ No newline at end of file diff --git a/tests/litellm/proxy/client/test_users.py b/tests/litellm/proxy/client/test_users.py new file mode 100644 index 0000000000..8cc5771343 --- /dev/null +++ b/tests/litellm/proxy/client/test_users.py @@ -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"]) \ No newline at end of file