From f4ecf3ca725d4e0283756c545b07c8bfa5825788 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 23 Sep 2025 19:43:32 -0700 Subject: [PATCH] [Feat] Fixes for LiteLLM Proxy CLI to Auth to Gateway (#14836) * fix: error msg from updating key * fix _create_new_cli_key * fix validate_key_team_change * fix interface for chat * ruff fix * fix auth for keys * get_litellm_gateway_api_key * fix chat * test fix * linting fix * fix mypy * test_validate_key_team_change_with_member_permissions --- litellm/proxy/client/chat.py | 93 +++- litellm/proxy/client/cli/commands/auth.py | 4 +- litellm/proxy/client/cli/commands/chat.py | 402 ++++++++++++++---- litellm/proxy/client/cli/interface.py | 2 +- litellm/proxy/client/client.py | 4 +- litellm/proxy/client/keys.py | 9 +- .../key_management_endpoints.py | 19 +- litellm/proxy/management_endpoints/ui_sso.py | 1 - .../proxy/client/cli/test_chat_commands.py | 248 ----------- .../test_key_management_endpoints.py | 58 +++ 10 files changed, 504 insertions(+), 336 deletions(-) delete mode 100644 tests/test_litellm/proxy/client/cli/test_chat_commands.py diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index bf5ddbe85d..91fc33002b 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -1,5 +1,8 @@ +import json +from typing import Any, Dict, Iterator, List, Optional, Union + import requests -from typing import List, Dict, Any, Optional, Union + from .exceptions import UnauthorizedError @@ -99,3 +102,91 @@ class ChatClient: if e.response.status_code == 401: raise UnauthorizedError(e) raise + + def completions_stream( + self, + model: str, + messages: List[Dict[str, str]], + temperature: Optional[float] = None, + top_p: Optional[float] = None, + n: Optional[int] = None, + max_tokens: Optional[int] = None, + presence_penalty: Optional[float] = None, + frequency_penalty: Optional[float] = None, + user: Optional[str] = None, + ) -> Iterator[Dict[str, Any]]: + """ + Create a streaming chat completion. + + Args: + model (str): The model to use for completion + messages (List[Dict[str, str]]): The messages to generate a completion for + temperature (Optional[float]): Sampling temperature between 0 and 2 + top_p (Optional[float]): Nucleus sampling parameter between 0 and 1 + n (Optional[int]): Number of completions to generate + max_tokens (Optional[int]): Maximum number of tokens to generate + presence_penalty (Optional[float]): Presence penalty between -2.0 and 2.0 + frequency_penalty (Optional[float]): Frequency penalty between -2.0 and 2.0 + user (Optional[str]): Unique identifier for the end user + + Yields: + Dict[str, Any]: Streaming response chunks from the server + + Raises: + UnauthorizedError: If the request fails with a 401 status code + requests.exceptions.RequestException: If the request fails with any other error + """ + url = f"{self._base_url}/chat/completions" + + # Build request data with required fields + data: Dict[str, Any] = { + "model": model, + "messages": messages, + "stream": True + } + + # Add optional parameters if provided + if temperature is not None: + data["temperature"] = temperature + if top_p is not None: + data["top_p"] = top_p + if n is not None: + data["n"] = n + if max_tokens is not None: + data["max_tokens"] = max_tokens + if presence_penalty is not None: + data["presence_penalty"] = presence_penalty + if frequency_penalty is not None: + data["frequency_penalty"] = frequency_penalty + if user is not None: + data["user"] = user + + # Make streaming request + session = requests.Session() + try: + response = session.post( + url, + headers=self._get_headers(), + json=data, + stream=True + ) + response.raise_for_status() + + # Parse SSE stream + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + data_str = line[6:] # Remove 'data: ' prefix + if data_str.strip() == '[DONE]': + break + try: + chunk = json.loads(data_str) + yield chunk + except json.JSONDecodeError: + continue + + except requests.exceptions.HTTPError as e: + if e.response.status_code == 401: + raise UnauthorizedError(e) + raise diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 54d51db8b4..78cbb5ecf8 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -281,12 +281,12 @@ def prompt_team_selection_fallback(teams: List[Dict[str, Any]]) -> Optional[Dict def update_key_with_team(base_url: str, api_key: str, team_id: str) -> bool: """Update the API key to be associated with the selected team""" - + from litellm.proxy._types import SpecialModelNames from litellm.proxy.client import Client client = Client(base_url=base_url, api_key=api_key) try: - client.keys.update(key=api_key, team_id=team_id) + client.keys.update(key=api_key, team_id=team_id, models=[SpecialModelNames.all_team_models.value]) click.echo(f"✅ Successfully assigned key to team: {team_id}") return True except requests.exceptions.HTTPError as e: diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index ea51f74652..41ded68ed0 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -1,42 +1,96 @@ import json -from typing import Optional +import sys +from typing import Any, Dict, List, Optional import click -import rich import requests +from rich.console import Console +from rich.panel import Panel +from rich.prompt import Prompt +from rich.table import Table +from ... import Client from ...chat import ChatClient -@click.group() -def chat(): - """Chat with models through the LiteLLM proxy server""" - pass +def _get_available_models(ctx: click.Context) -> List[Dict[str, Any]]: + """Get list of available models from the proxy server""" + try: + client = Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + models_list = client.models.list() + # Ensure we return a list of dictionaries + if isinstance(models_list, list): + # Filter to ensure all items are dictionaries + return [model for model in models_list if isinstance(model, dict)] + return [] + except Exception as e: + click.echo(f"Warning: Could not fetch models list: {e}", err=True) + return [] -@chat.command() -@click.argument("model") -@click.option( - "--message", - "-m", - multiple=True, - help="Messages in 'role:content' format (e.g. 'user:Hello'). Can be specified multiple times.", -) +def _select_model(console: Console, available_models: List[Dict[str, Any]]) -> Optional[str]: + """Interactive model selection""" + if not available_models: + console.print("[yellow]No models available or could not fetch models list.[/yellow]") + model_name = Prompt.ask("Please enter a model name") + return model_name if model_name.strip() else None + + # Display available models in a table + table = Table(title="Available Models") + table.add_column("Index", style="cyan", no_wrap=True) + table.add_column("Model ID", style="green") + table.add_column("Owned By", style="yellow") + MAX_MODELS_TO_DISPLAY = 200 + + models_to_display: List[Dict[str, Any]] = available_models[:MAX_MODELS_TO_DISPLAY] + for i, model in enumerate(models_to_display): # Limit to first 200 models + table.add_row( + str(i + 1), + str(model.get("id", "")), + str(model.get("owned_by", "")) + ) + + if len(available_models) > MAX_MODELS_TO_DISPLAY: + console.print(f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]") + + console.print(table) + + while True: + try: + choice = Prompt.ask( + "\nSelect a model by entering the index number (or type a model name directly)", + default="1" + ).strip() + + # Try to parse as index + try: + index = int(choice) - 1 + if 0 <= index < len(available_models): + return available_models[index]["id"] + else: + console.print(f"[red]Invalid index. Please enter a number between 1 and {len(available_models)}[/red]") + continue + except ValueError: + # Not a number, treat as model name + if choice: + return choice + else: + console.print("[red]Please enter a valid model name or index[/red]") + continue + + except KeyboardInterrupt: + console.print("\n[yellow]Model selection cancelled.[/yellow]") + return None + + +@click.command() +@click.argument("model", required=False) @click.option( "--temperature", "-t", type=float, - help="Sampling temperature between 0 and 2", -) -@click.option( - "--top-p", - type=float, - help="Nucleus sampling parameter between 0 and 1", -) -@click.option( - "--n", - type=int, - help="Number of completions to generate", + default=0.7, + help="Sampling temperature between 0 and 2 (default: 0.7)", ) @click.option( "--max-tokens", @@ -44,65 +98,271 @@ def chat(): help="Maximum number of tokens to generate", ) @click.option( - "--presence-penalty", - type=float, - help="Presence penalty between -2.0 and 2.0", -) -@click.option( - "--frequency-penalty", - type=float, - help="Frequency penalty between -2.0 and 2.0", -) -@click.option( - "--user", + "--system", + "-s", type=str, - help="Unique identifier for the end user", + help="System message to set the behavior of the assistant", ) @click.pass_context -def completions( +def chat( ctx: click.Context, - model: str, - message: tuple[str, ...], - temperature: Optional[float] = None, - top_p: Optional[float] = None, - n: Optional[int] = None, + model: Optional[str], + temperature: float, max_tokens: Optional[int] = None, - presence_penalty: Optional[float] = None, - frequency_penalty: Optional[float] = None, - user: Optional[str] = None, + system: Optional[str] = None, ): - """Create a chat completion""" - if not message: - raise click.UsageError("At least one message is required") - - # Parse messages from role:content format - messages = [] - for msg in message: - try: - role, content = msg.split(":", 1) - messages.append({"role": role.strip(), "content": content.strip()}) - except ValueError: - raise click.BadParameter(f"Invalid message format: {msg}. Expected format: 'role:content'") - + """Interactive chat with streaming responses + + Examples: + + # Chat with a specific model + litellm-proxy chat gpt-4 + + # Chat without specifying model (will show model selection) + litellm-proxy chat + + # Chat with custom settings + litellm-proxy chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" + """ + console = Console() + + # If no model specified, show model selection + if not model: + available_models = _get_available_models(ctx) + model = _select_model(console, available_models) + if not model: + console.print("[red]No model selected. Exiting.[/red]") + return + client = ChatClient(ctx.obj["base_url"], ctx.obj["api_key"]) + + # Initialize conversation history + messages: List[Dict[str, Any]] = [] + + # Add system message if provided + if system: + messages.append({"role": "system", "content": system}) + + # Display welcome message + console.print(Panel.fit( + f"[bold blue]LiteLLM Interactive Chat[/bold blue]\n" + f"Model: [green]{model}[/green]\n" + f"Temperature: [yellow]{temperature}[/yellow]\n" + f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n" + f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n" + f"Type '/help' for more commands.", + title="🤖 Chat Session" + )) + try: - response = client.completions( + while True: + # Get user input + try: + user_input = console.input("\n[bold cyan]You:[/bold cyan] ").strip() + except (EOFError, KeyboardInterrupt): + console.print("\n[yellow]Chat session ended.[/yellow]") + break + + # Handle special commands + should_exit, messages, new_model = _handle_special_commands( + console, user_input, messages, system, ctx + ) + + if should_exit: + break + if new_model: + model = new_model + + # Check if this was a special command that was handled (not a normal message) + if user_input.lower().startswith(('/quit', '/exit', '/q', '/help', '/clear', '/history', '/save', '/load', '/model')) or not user_input: + continue + + # Add user message to conversation + messages.append({"role": "user", "content": user_input}) + + # Display assistant label + console.print("\n[bold green]Assistant:[/bold green]") + + # Stream the response + assistant_content = _stream_response( + console=console, + client=client, + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + ) + + # Add assistant message to conversation history + if assistant_content: + messages.append({"role": "assistant", "content": assistant_content}) + else: + console.print("[red]Error: No content received from the model[/red]") + + except KeyboardInterrupt: + console.print("\n[yellow]Chat session interrupted.[/yellow]") + + +def _show_help(console: Console): + """Show help for interactive chat commands""" + help_text = """ +[bold]Interactive Chat Commands:[/bold] + +[cyan]/help[/cyan] - Show this help message +[cyan]/quit[/cyan] - Exit the chat session (also /exit, /q) +[cyan]/clear[/cyan] - Clear conversation history +[cyan]/history[/cyan] - Show conversation history +[cyan]/model[/cyan] - Switch to a different model +[cyan]/save [/cyan] - Save conversation to file +[cyan]/load [/cyan] - Load conversation from file + +[bold]Tips:[/bold] +- Your conversation history is maintained during the session +- Use Ctrl+C to interrupt at any time +- Responses are streamed in real-time +- You can switch models mid-conversation with /model + """ + console.print(Panel(help_text, title="Help")) + + +def _show_history(console: Console, messages: List[Dict[str, Any]]): + """Show conversation history""" + if not messages: + console.print("[yellow]No conversation history.[/yellow]") + return + + console.print(Panel.fit("[bold]Conversation History[/bold]", title="History")) + + for i, message in enumerate(messages, 1): + role = message["role"] + content = message["content"] + + if role == "system": + console.print(f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]") + elif role == "user": + console.print(f"{i}. [bold cyan]You:[/bold cyan] {content}") + elif role == "assistant": + console.print(f"{i}. [bold green]Assistant:[/bold green] {content[:100]}{'...' if len(content) > 100 else ''}") + + +def _save_conversation(console: Console, messages: List[Dict[str, Any]], command: str): + """Save conversation to a file""" + parts = command.split() + if len(parts) < 2: + console.print("[red]Usage: /save [/red]") + return + + filename = parts[1] + if not filename.endswith('.json'): + filename += '.json' + + try: + with open(filename, 'w') as f: + json.dump(messages, f, indent=2) + console.print(f"[green]Conversation saved to {filename}[/green]") + except Exception as e: + console.print(f"[red]Error saving conversation: {e}[/red]") + + +def _load_conversation(console: Console, command: str, system: Optional[str]) -> List[Dict[str, Any]]: + """Load conversation from a file""" + parts = command.split() + if len(parts) < 2: + console.print("[red]Usage: /load [/red]") + return [] + + filename = parts[1] + if not filename.endswith('.json'): + filename += '.json' + + try: + with open(filename, 'r') as f: + messages = json.load(f) + console.print(f"[green]Conversation loaded from {filename}[/green]") + return messages + except FileNotFoundError: + console.print(f"[red]File not found: {filename}[/red]") + except Exception as e: + console.print(f"[red]Error loading conversation: {e}[/red]") + + # Return empty list or just system message if load failed + if system: + return [{"role": "system", "content": system}] + return [] + + +def _handle_special_commands( + console: Console, + user_input: str, + messages: List[Dict[str, Any]], + system: Optional[str], + ctx: click.Context +) -> tuple[bool, List[Dict[str, Any]], Optional[str]]: + """Handle special chat commands. Returns (should_exit, updated_messages, updated_model)""" + if user_input.lower() in ['/quit', '/exit', '/q']: + console.print("[yellow]Chat session ended.[/yellow]") + return True, messages, None + elif user_input.lower() == '/help': + _show_help(console) + return False, messages, None + elif user_input.lower() == '/clear': + new_messages = [] + if system: + new_messages.append({"role": "system", "content": system}) + console.print("[green]Conversation history cleared.[/green]") + return False, new_messages, None + elif user_input.lower() == '/history': + _show_history(console, messages) + return False, messages, None + elif user_input.lower().startswith('/save'): + _save_conversation(console, messages, user_input) + return False, messages, None + elif user_input.lower().startswith('/load'): + new_messages = _load_conversation(console, user_input, system) + return False, new_messages, None + elif user_input.lower() == '/model': + available_models = _get_available_models(ctx) + new_model = _select_model(console, available_models) + if new_model: + console.print(f"[green]Switched to model: {new_model}[/green]") + return False, messages, new_model + return False, messages, None + elif not user_input: + return False, messages, None + + # Not a special command + return False, messages, None + + +def _stream_response(console: Console, client: ChatClient, model: str, messages: List[Dict[str, Any]], temperature: float, max_tokens: Optional[int]) -> Optional[str]: + """Stream the model response and return the complete content""" + try: + assistant_content = "" + for chunk in client.completions_stream( model=model, messages=messages, temperature=temperature, - top_p=top_p, - n=n, max_tokens=max_tokens, - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, - user=user, - ) - rich.print_json(data=response) + ): + if "choices" in chunk and len(chunk["choices"]) > 0: + delta = chunk["choices"][0].get("delta", {}) + content = delta.get("content", "") + if content: + assistant_content += content + console.print(content, end="") + sys.stdout.flush() + + console.print() # Add newline after streaming + return assistant_content if assistant_content else None + except requests.exceptions.HTTPError as e: - click.echo(f"Error: HTTP {e.response.status_code}", err=True) + console.print(f"\n[red]Error: HTTP {e.response.status_code}[/red]") try: error_body = e.response.json() - rich.print_json(data=error_body) + console.print(f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]") except json.JSONDecodeError: - click.echo(e.response.text, err=True) - raise click.Abort() + console.print(f"[red]{e.response.text}[/red]") + return None + except Exception as e: + console.print(f"\n[red]Error: {str(e)}[/red]") + return None \ No newline at end of file diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index 78aeb44235..2c6f5f10b4 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -84,7 +84,7 @@ def show_commands(): ("whoami", "Show current authentication status"), ("models", "Manage and view model configurations"), ("credentials", "Manage API credentials"), - ("chat", "Interactive chat with models"), + ("chat", "Interactive streaming chat with models"), ("http", "Make HTTP requests to the proxy"), ("keys", "Manage API keys"), ("teams", "Manage teams and team assignments"), diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index 8ed9a4ff89..c9066f70de 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -1,5 +1,7 @@ from typing import Optional +from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key + from .chat import ChatClient from .credentials import CredentialsManagementClient from .http_client import HTTPClient @@ -27,7 +29,7 @@ class Client: timeout: Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present - self._api_key = api_key + self._api_key = get_litellm_gateway_api_key() or api_key # Initialize resource clients diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index fc307648a2..50fd7b9d9c 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -274,14 +274,15 @@ class KeysManagementClient: data["aliases"] = aliases request = requests.Request("POST", url, headers=self._get_headers(), json=data) session = requests.Session() + response_text: Optional[str] = None try: response = session.send(request.prepare()) + response_text = response.text response.raise_for_status() return response.json() - except requests.exceptions.HTTPError as e: - if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + except Exception: + raise Exception(f"Error updating key: {response_text}") + def info(self, key: str, return_request: bool = False) -> Union[Dict[str, Any], requests.Request]: """ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 026ae9e405..eb389450f9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1232,13 +1232,11 @@ def validate_key_team_change( ) # Check if the key's user_id is a member of the team + member_object = _get_user_in_team( + team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id + ) if key.user_id is not None: - is_member = False - for member in team.members_with_roles: - if member.user_id == key.user_id: - is_member = True - break - if not is_member: + if not member_object: raise HTTPException( status_code=403, detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", @@ -1265,10 +1263,17 @@ def validate_key_team_change( team_obj=team, ): return + # this teams member permissions allow updating a + elif TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( + team_member_object=member_object, + team_table=cast(LiteLLM_TeamTableCachedObj, team), + route=KeyManagementRoutes.KEY_UPDATE.value, + ): + return else: raise HTTPException( status_code=403, - detail=f"User={change_initiated_by.user_id} is not a Proxy Admin or Team Admin for team={team.team_id}.", + detail=f"User={change_initiated_by.user_id} is not a Proxy Admin or Team Admin for team={team.team_id}. Please ask your Proxy Admin to allow this action under 'Member Permissions' for this team.", ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 637a852087..2af940ec65 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -731,7 +731,6 @@ async def _create_new_cli_key( config={}, spend=0, user_id=user_id, - team_id="litellm-cli", table_name="key", token=key, ) diff --git a/tests/test_litellm/proxy/client/cli/test_chat_commands.py b/tests/test_litellm/proxy/client/cli/test_chat_commands.py deleted file mode 100644 index 1226c3557f..0000000000 --- a/tests/test_litellm/proxy/client/cli/test_chat_commands.py +++ /dev/null @@ -1,248 +0,0 @@ -import json -import os -import sys -from unittest.mock import MagicMock, patch - -import pytest -import requests -from click.testing import CliRunner - -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - - -from litellm.proxy.client.cli.main import cli - - -@pytest.fixture -def mock_chat_client(): - with patch("litellm.proxy.client.cli.commands.chat.ChatClient") as mock: - yield mock - - -@pytest.fixture -def cli_runner(): - return CliRunner() - - -def test_chat_completions_success(cli_runner, mock_chat_client): - # Mock response data - mock_response = { - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677858242, - "model": "gpt-4", - "choices": [ - { - "message": { - "role": "assistant", - "content": "Hello! How can I help you today?", - }, - "finish_reason": "stop", - "index": 0, - } - ], - } - mock_instance = mock_chat_client.return_value - mock_instance.completions.return_value = mock_response - - # Run command - result = cli_runner.invoke( - cli, - [ - "chat", - "completions", - "gpt-4", - "-m", - "user:Hello!", - "--temperature", - "0.7", - "--max-tokens", - "100", - ], - ) - - # Verify - assert result.exit_code == 0 - output_data = json.loads(result.output) - assert output_data == mock_response - mock_instance.completions.assert_called_once_with( - model="gpt-4", - messages=[{"role": "user", "content": "Hello!"}], - temperature=0.7, - max_tokens=100, - top_p=None, - n=None, - presence_penalty=None, - frequency_penalty=None, - user=None, - ) - - -def test_chat_completions_multiple_messages(cli_runner, mock_chat_client): - # Mock response data - mock_response = { - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677858242, - "model": "gpt-4", - "choices": [ - { - "message": { - "role": "assistant", - "content": "Paris has a population of about 2.2 million.", - }, - "finish_reason": "stop", - "index": 0, - } - ], - } - mock_instance = mock_chat_client.return_value - mock_instance.completions.return_value = mock_response - - # Run command - result = cli_runner.invoke( - cli, - [ - "chat", - "completions", - "gpt-4", - "-m", - "system:You are a helpful assistant", - "-m", - "user:What's the population of Paris?", - ], - ) - - # Verify - assert result.exit_code == 0 - output_data = json.loads(result.output) - assert output_data == mock_response - mock_instance.completions.assert_called_once_with( - model="gpt-4", - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": "What's the population of Paris?"}, - ], - temperature=None, - max_tokens=None, - top_p=None, - n=None, - presence_penalty=None, - frequency_penalty=None, - user=None, - ) - - -def test_chat_completions_no_messages(cli_runner, mock_chat_client): - # Run command without any messages - result = cli_runner.invoke(cli, ["chat", "completions", "gpt-4"]) - - # Verify - assert result.exit_code == 2 - assert "At least one message is required" in result.output - mock_instance = mock_chat_client.return_value - mock_instance.completions.assert_not_called() - - -def test_chat_completions_invalid_message_format(cli_runner, mock_chat_client): - # Run command with invalid message format - result = cli_runner.invoke( - cli, ["chat", "completions", "gpt-4", "-m", "invalid-format"] - ) - - # Verify - assert result.exit_code == 2 - assert "Invalid message format" in result.output - mock_instance = mock_chat_client.return_value - mock_instance.completions.assert_not_called() - - -def test_chat_completions_http_error(cli_runner, mock_chat_client): - # Mock HTTP error - mock_instance = mock_chat_client.return_value - mock_error_response = MagicMock() - mock_error_response.status_code = 400 - mock_error_response.json.return_value = { - "error": "Invalid request", - "message": "Invalid model specified", - } - mock_instance.completions.side_effect = requests.exceptions.HTTPError( - response=mock_error_response - ) - - # Run command - result = cli_runner.invoke( - cli, ["chat", "completions", "invalid-model", "-m", "user:Hello"] - ) - - # Verify - assert result.exit_code == 1 - assert "Error: HTTP 400" in result.output - assert "Invalid request" in result.output - assert "Invalid model specified" in result.output - - -def test_chat_completions_all_parameters(cli_runner, mock_chat_client): - # Mock response data - mock_response = { - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677858242, - "model": "gpt-4", - "choices": [ - { - "message": { - "role": "assistant", - "content": "Response with all parameters set", - }, - "finish_reason": "stop", - "index": 0, - } - ], - } - mock_instance = mock_chat_client.return_value - mock_instance.completions.return_value = mock_response - - # Run command with all available parameters - result = cli_runner.invoke( - cli, - [ - "chat", - "completions", - "gpt-4", - "-m", - "user:Test message", - "--temperature", - "0.7", - "--top-p", - "0.9", - "--n", - "1", - "--max-tokens", - "100", - "--presence-penalty", - "0.5", - "--frequency-penalty", - "0.5", - "--user", - "test-user", - ], - ) - - # Verify - assert result.exit_code == 0 - output_data = json.loads(result.output) - assert output_data == mock_response - mock_instance.completions.assert_called_once_with( - model="gpt-4", - messages=[{"role": "user", "content": "Test message"}], - temperature=0.7, - top_p=0.9, - n=1, - max_tokens=100, - presence_penalty=0.5, - frequency_penalty=0.5, - user="test-user", - ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index cd768e47ae..da64a866d5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _common_key_generation_helper, _list_key_helper, prepare_key_update_data, + validate_key_team_change, ) from litellm.proxy.proxy_server import app @@ -1033,3 +1034,60 @@ async def test_unblock_key_invalid_key_format(monkeypatch): assert exc_info.value.code == "400" assert "Invalid key format" in str(exc_info.value.message) + + +def test_validate_key_team_change_with_member_permissions(): + """ + Test validate_key_team_change function with team member permissions. + + This test covers the new logic that allows team members with specific + permissions to update keys, not just team admins. + """ + from unittest.mock import MagicMock, patch + + from litellm.proxy._types import KeyManagementRoutes + + # Create mock objects + mock_key = MagicMock() + mock_key.user_id = "test-user-123" + mock_key.models = ["gpt-4"] + mock_key.tpm_limit = None + mock_key.rpm_limit = None + + mock_team = MagicMock() + mock_team.team_id = "test-team-456" + mock_team.members_with_roles = [] + mock_team.tpm_limit = None + mock_team.rpm_limit = None + + mock_change_initiator = MagicMock() + mock_change_initiator.user_id = "test-user-123" + + mock_router = MagicMock() + + # Mock the member object returned by _get_user_in_team + mock_member_object = MagicMock() + + with patch('litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model'): + with patch('litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team') as mock_get_user: + with patch('litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin') as mock_is_admin: + with patch('litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint') as mock_has_perms: + + mock_get_user.return_value = mock_member_object + mock_is_admin.return_value = False + mock_has_perms.return_value = True + + # This should not raise an exception due to member permissions + validate_key_team_change( + key=mock_key, + team=mock_team, + change_initiated_by=mock_change_initiator, + llm_router=mock_router + ) + + # Verify the permission check was called with correct parameters + mock_has_perms.assert_called_once_with( + team_member_object=mock_member_object, + team_table=mock_team, + route=KeyManagementRoutes.KEY_UPDATE.value + )