mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-14 01:07:15 +00:00
[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
This commit is contained in:
Ishaan Jaff
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 <name>[/cyan] - Save conversation to file
|
||||
[cyan]/load <name>[/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 <filename>[/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 <filename>[/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
|
||||
@@ -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"),
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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]:
|
||||
"""
|
||||
|
||||
@@ -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.",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user