diff --git a/docker/docker-compose.quickstart.yml b/docker/docker-compose.quickstart.yml index a1d47e323ff..98366a73e58 100644 --- a/docker/docker-compose.quickstart.yml +++ b/docker/docker-compose.quickstart.yml @@ -11,7 +11,7 @@ # beyond local evaluation, pin the image to a specific release tag. services: litellm: - image: docker.litellm.ai/berriai/litellm:main-stable + image: docker.litellm.ai/berriai/litellm:latest ports: # LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart # script sets it to "127.0.0.1:" so new installs listen on this machine only. diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a469c378db6..a1abed202a2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2271,7 +2271,13 @@ async def user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, ) if isinstance(resolve_result, UserAPIKeyAuth): - valid_token = resolve_result + mapped_end_user_id: Final = ( + jwt_handler.get_end_user_id( # pyright: ignore[reportUnknownMemberType] # untyped dict token + token=jwt_claims or {}, default_value=None + ) + or None + ) + valid_token = resolve_result.model_copy(update={"end_user_id": mapped_end_user_id}) api_key = valid_token.token or "" valid_token.jwt_claims = jwt_claims do_standard_jwt_auth = False @@ -2503,7 +2509,9 @@ async def user_api_key_auth_builder( _end_user_object = None end_user_params: Final = {} - raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, safe_get_request_headers(request)) + raw_end_user_id: Final = ( + valid_token.end_user_id if valid_token is not None else None + ) or get_end_user_id_from_request_body(request_data, safe_get_request_headers(request)) end_user_id = await resolve_and_validate_end_user_id( raw_end_user_id=raw_end_user_id, prisma_client=prisma_client, diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 78d03ecc6e1..ff4f2c4fd1e 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -381,7 +381,7 @@ class ContentBlockDeltaEvent(TypedDict, total=False): class PerformanceConfigBlock(TypedDict): - latency: Literal["optimized", "throughput"] + latency: Literal["standard", "optimized"] class JsonSchemaDefinition(TypedDict, total=False): diff --git a/scripts/quickstart.ps1 b/scripts/quickstart.ps1 index aa6de6e247d..2bc89f3513b 100644 --- a/scripts/quickstart.ps1 +++ b/scripts/quickstart.ps1 @@ -9,10 +9,16 @@ # # Asks where to keep the files, then offers to copy the admin password and # open the admin UI; every question has a default you accept by pressing Enter. +# Run it again for a menu to open, stop, start, update, or remove the install. # It asks nothing when input is not a console, under CI or Claude Code, or # with -Yes. # -# -Yes, or LITELLM_YES=1 no questions: install to ~\litellm-gateway, don't open a browser +# It also takes an action, which the litellm-gateway command it installs takes too: +# & ([scriptblock]::Create((irm https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.ps1))) stop +# status, start (up), stop (down), restart, logs, open, update, uninstall +# +# -Yes, or LITELLM_YES=1 no questions: install to ~\litellm-gateway, don't open a browser, +# and let uninstall delete the database without asking # LITELLM_DIR folder to install into (skips the folder question) # LITELLM_PORT port for the gateway (default 4000, or the next free one) # NO_COLOR plain output, which agents, CI, and log files always get @@ -29,7 +35,7 @@ param([switch]$Yes) function Invoke-LiteLLMQuickstart { - param([switch]$Yes) + param([switch]$Yes, [object[]]$Rest = @()) $ErrorActionPreference = 'Stop' # Windows PowerShell draws a progress bar for every web request, which makes them crawl. $ProgressPreference = 'SilentlyContinue' @@ -37,8 +43,64 @@ function Invoke-LiteLLMQuickstart { $composeUrl = 'https://raw.githubusercontent.com/BerriAI/litellm/main/docker/docker-compose.quickstart.yml' if ($env:LITELLM_COMPOSE_URL) { $composeUrl = $env:LITELLM_COMPOSE_URL } + $scriptUrl = 'https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.ps1' + if ($env:LITELLM_SCRIPT_URL) { $scriptUrl = $env:LITELLM_SCRIPT_URL } $onWindows = ($PSVersionTable.PSEdition -eq 'Desktop') -or ((Test-Path variable:IsWindows) -and $IsWindows) + # A step that cannot go on stops the whole run with exit code 1, from any depth. + $exitSignal = 'LITELLM_QUICKSTART_EXIT' + function Stop-Quickstart { throw $exitSignal } + + function Show-Usage([IO.TextWriter]$To) { + @( + 'Usage: quickstart.ps1 [-Yes] [action]' + '' + 'With no action it installs the LiteLLM Gateway, or shows a menu to manage the' + 'install it finds. Actions:' + ' status whether it is running, and where' + ' start start it (also: up)' + ' stop stop it; models, keys, and logs are kept (also: down)' + ' restart stop it and start it again' + " logs show the gateway's logs" + ' open open the admin UI in your browser' + ' update download the latest image and restart on it' + ' uninstall remove its containers, its database, and its files' + ' help show this help (also: -h, --help)' + '' + ' -Yes, -y ask nothing; without a console, uninstall needs it' + '' + 'From the install command:' + ' & ([scriptblock]::Create((irm https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.ps1))) status' + ) | ForEach-Object { $To.WriteLine($_) } + } + + # Windows PowerShell passes -y and --yes through as plain arguments. + $noQuestions = [bool]$Yes + # A port set by the caller is never moved; later steps set LITELLM_PORT too. + $portPinned = [bool]$env:LITELLM_PORT + $action = '' + foreach ($arg in $Rest) { + $a = "$arg" + if ($a -in '-y', '--yes', '-yes') { + $noQuestions = $true + } elseif ($a -in '-h', '--help', '-help', 'help') { + Show-Usage ([Console]::Out) + return 0 + } elseif ($a -in 'status', 'start', 'up', 'stop', 'down', 'restart', 'logs', 'open', 'update', 'uninstall') { + if ($action) { + [Console]::Error.WriteLine("Choose one action, not both $action and $a") + return 1 + } + $action = $a.ToLowerInvariant() + } else { + [Console]::Error.WriteLine("Unknown option: $a") + Show-Usage ([Console]::Error) + return 1 + } + } + if ($action -eq 'up') { $action = 'start' } elseif ($action -eq 'down') { $action = 'stop' } + if ($env:LITELLM_YES) { $noQuestions = $true } + # ------------------------------------------------------------ output # A person at a console gets colors, step marks, and a spinner. Agents, CI, @@ -127,6 +189,16 @@ function Invoke-LiteLLMQuickstart { return [pscustomobject]@{ Code = $code; Out = @($out) } } + # Run a native program with its output on the console, never in this + # function's results; returns its exit code. + function Invoke-Shown { + param([string]$File, [string[]]$Arguments) + $prev = $ErrorActionPreference + $ErrorActionPreference = 'Continue' + try { & $File @Arguments 2>&1 | ForEach-Object { [Console]::Out.WriteLine("$_") }; return $LASTEXITCODE } + finally { $ErrorActionPreference = $prev } + } + # Write a text file as UTF-8 without a byte order mark, with LF line endings, # which is what Docker Compose and git expect. function Write-PlainText([string]$Path, [string]$Text, [switch]$Append) { @@ -154,7 +226,7 @@ function Invoke-LiteLLMQuickstart { # ------------------------------------------------------------ questions - $interactive = (-not $Yes) -and (-not $env:LITELLM_YES) -and (-not $env:CI) -and (-not $env:CLAUDECODE) -and + $interactive = (-not $noQuestions) -and (-not $env:CI) -and (-not $env:CLAUDECODE) -and (-not [Console]::IsInputRedirected) -and [Environment]::UserInteractive # Menu "Question" Default @(@('label', 'note'), ...) -> the 1-based pick. @@ -224,19 +296,29 @@ function Invoke-LiteLLMQuickstart { # ------------------------------------------------------------ steps function Test-PortFree([int]$Port) { + # Windows lets a second program bind a port that another one holds on a + # different address, so ask it what is listening first. + if ($onWindows -and (Get-Command Get-NetTCPConnection -ErrorAction SilentlyContinue)) { + try { + if (Get-NetTCPConnection -State Listen -LocalPort $Port -ErrorAction SilentlyContinue) { return $false } + } catch { <# no answer: binding below decides #> } + } # Binding answers right away and is accurate. Probing with a connect is # not: the WSL2 localhost relay (Podman, Rancher Desktop) swallows the # refusal on closed ports, so every connect waits out its timeout and - # closed ports look taken. - $listener = New-Object System.Net.Sockets.TcpListener([Net.IPAddress]::Loopback, $Port) - try { - $listener.Start() - return $true - } catch { - return $false - } finally { - $listener.Stop() + # closed ports look taken. A listener on all addresses only blocks the + # second bind, so try both. + foreach ($address in [Net.IPAddress]::Loopback, [Net.IPAddress]::Any) { + $listener = New-Object System.Net.Sockets.TcpListener($address, $Port) + try { + $listener.Start() + } catch { + return $false + } finally { + $listener.Stop() + } } + return $true } function Test-Ready([int]$Port) { @@ -250,13 +332,7 @@ function Invoke-LiteLLMQuickstart { # show if it fails. Without styling it runs in the open. function Invoke-WithSpinner { param([string]$Label, [string]$Detail, [string]$File, [string[]]$Arguments) - if (-not $style) { - $prev = $ErrorActionPreference - $ErrorActionPreference = 'Continue' - try { & $File @Arguments 2>&1 | ForEach-Object { [Console]::Out.WriteLine("$_") }; $code = $LASTEXITCODE } - finally { $ErrorActionPreference = $prev } - return $code - } + if (-not $style) { return Invoke-Shown $File $Arguments } $out = [IO.Path]::GetTempFileName(); $err = [IO.Path]::GetTempFileName() try { $p = Start-Process -FilePath $File -ArgumentList $Arguments -NoNewWindow -PassThru -RedirectStandardOutput $out -RedirectStandardError $err @@ -363,9 +439,613 @@ function Invoke-LiteLLMQuickstart { Say "{a:$($sym.BL)$rule$($sym.BR)}" } + # ------------------------------------------------------------ the install + + # The install this run manages, filled in by Use-Install. + $inst = @{ Dir = ''; Project = ''; Port = 0; Url = ''; Command = '' } + $cf = @('compose', '-f', 'docker-compose.quickstart.yml') + + function Read-EnvValue([string]$Name) { + $p = Join-Path (Get-Location).ProviderPath '.env' + if (-not (Test-Path -LiteralPath $p)) { return '' } + $m = Get-Content -LiteralPath $p | + Where-Object { $_.StartsWith("$Name=") } | Select-Object -Last 1 + if ($m) { return ($m -split '=', 2)[1] } + return '' + } + + # An existing install: LITELLM_DIR, this folder, .\litellm-gateway, or + # ~\litellm-gateway, whichever has a .env first. Empty when there is none. + function Find-Install { + $here = (Get-Location).ProviderPath + if ($env:LITELLM_DIR) { + if (Test-Path -LiteralPath (Join-Path $env:LITELLM_DIR '.env')) { return $env:LITELLM_DIR } + } elseif ((Test-Path -LiteralPath (Join-Path $here '.env')) -and (Test-Path -LiteralPath (Join-Path $here 'docker-compose.quickstart.yml'))) { + return $here + } elseif (Test-Path -LiteralPath (Join-Path $here 'litellm-gateway/.env')) { + return (Join-Path $here 'litellm-gateway') + } elseif (Test-Path -LiteralPath (Join-Path $HOME 'litellm-gateway/.env')) { + return (Join-Path $HOME 'litellm-gateway') + } + return '' + } + + # What holds a port: Owner describes it, FreeIt is the command that frees it, + # and Project is set when it is another quickstart install. All stay empty + # when the holder cannot be named, such as another user's process. + function Get-PortOwner([int]$Port) { + $o = @{ Owner = ''; FreeIt = ''; Project = '' } + $name = '' + foreach ($line in (Invoke-Native $docker @('ps', '--format', '{{.Names}} {{.Ports}}')).Out) { + if ("$line" -match "[:]$Port->") { $name = ("$line" -split ' ')[0]; break } + } + if ($name) { + # Labels come as JSON: Windows PowerShell mangles the quotes an index template needs. + $labels = Invoke-Native $docker @('inspect', '--format', '{{json .Config.Labels}}', $name) + if ($labels.Code -eq 0) { + try { + $json = ($labels.Out -join "`n") | ConvertFrom-Json + if ($json -and $json.PSObject.Properties['com.docker.compose.project']) { $o.Project = "$($json.'com.docker.compose.project')" } + } catch { <# unreadable labels: treat it as a plain container #> } + } + # Only another quickstart install keeps its project, since that is the one + # this script offers to stop. Whoever started the container set the label, + # so it must also be a name Compose would accept. + if ($o.Project -cmatch '^litellm-gateway[a-z0-9_-]*$') { + $o.Owner = "another LiteLLM quickstart install ($engineName project $($o.Project))" + $o.FreeIt = "$engine compose -p $($o.Project) stop" + } else { + $o.Project = '' + $o.Owner = "the $engineName container $name"; $o.FreeIt = "$engine stop $name" + } + return $o + } + $procId = 0; $name = ''; $args0 = '' + if ($onWindows) { + try { + $conn = Get-NetTCPConnection -State Listen -LocalPort $Port -ErrorAction SilentlyContinue | Select-Object -First 1 + if ($conn) { + $procId = [int]$conn.OwningProcess + $name = (Get-Process -Id $procId -ErrorAction SilentlyContinue).ProcessName + $args0 = "$((Get-CimInstance Win32_Process -Filter "ProcessId=$procId" -ErrorAction SilentlyContinue).CommandLine)" + } + } catch { <# not allowed to look: leave it unnamed #> } + if ($procId) { $o.FreeIt = "Stop-Process -Id $procId" } + } elseif (Get-Command lsof -CommandType Application -ErrorAction SilentlyContinue) { + $found = (Invoke-Native lsof @('-nP', "-iTCP:$Port", '-sTCP:LISTEN', '-t')).Out | Select-Object -First 1 + if ($found) { + $procId = [int]"$found" + $name = "$((Invoke-Native ps @('-o', 'comm=', '-p', "$procId")).Out | Select-Object -First 1)" + if ($name) { $name = Split-Path -Leaf $name.Trim() } + $args0 = "$((Invoke-Native ps @('-o', 'args=', '-p', "$procId")).Out | Select-Object -First 1)" + $o.FreeIt = "kill $procId" + } + } + if (-not $procId) { return $o } + if ($args0 -match 'litellm') { + $o.Owner = "a litellm server (pid $procId), maybe one started after pip install litellm" + } else { + if (-not $name) { $name = 'a program' } + $o.Owner = "$name (pid $procId)" + } + return $o + } + + function Select-Port { + $saved = Read-EnvValue 'LITELLM_PORT' + if ($env:LITELLM_PORT) { + $port = [int]$env:LITELLM_PORT + } elseif ($saved) { + $port = [int]$saved + } elseif (Test-Path -LiteralPath '.env') { + $port = 4000 # an existing install without a saved port runs on the compose default + } else { + $port = 4000 + while (-not (Test-PortFree $port)) { + $port++ + if ($port -gt 4099) { + Fail 'Ports 4000 to 4099 are all in use.' 'Set LITELLM_PORT to a free port and run this again.' + Stop-Quickstart + } + } + if ($port -eq 4000) { + Step 'Port {b:4000} is free' + } else { + $o = Get-PortOwner 4000 + $by = if ($o.Owner) { ' by ' + (Lit $o.Owner) } else { '' } + Warn "Port 4000 is in use$by, so LiteLLM will use {b:$port}" + } + } + $env:LITELLM_PORT = "$port" + return $port + } + + # Compose prefers values already set in the environment over .env, so drop + # inherited ones: .env stays the only source for keys and the project name. + function Clear-ComposeEnv { + foreach ($name in 'LITELLM_MASTER_KEY', 'LITELLM_SALT_KEY', 'POSTGRES_PASSWORD', 'COMPOSE_PROJECT_NAME') { + Remove-Item "Env:$name" -ErrorAction SilentlyContinue + } + # The bind address follows .env when .env sets it. For an older .env without + # it, a value already in the environment is kept. + if (Get-Content -LiteralPath '.env' | Where-Object { $_ -match '^LITELLM_BIND=' }) { Remove-Item Env:LITELLM_BIND -ErrorAction SilentlyContinue } + } + + function Use-Install([string]$Dir) { + Set-Location -LiteralPath $Dir + $inst.Dir = (Get-Location).ProviderPath + $inst.Project = Read-EnvValue 'COMPOSE_PROJECT_NAME' + # Installs from before COMPOSE_PROJECT_NAME have Compose's default, the folder name. + if (-not $inst.Project) { $inst.Project = (Split-Path -Leaf $inst.Dir).ToLowerInvariant() -replace '[^a-z0-9_-]', '' } + $inst.Port = Select-Port + Clear-ComposeEnv + $inst.Url = "http://localhost:$($inst.Port)/ui" + } + + function Save-Compose { + Invoke-WebRequest -UseBasicParsing -Uri $composeUrl -OutFile (Join-Path (Get-Location).ProviderPath 'docker-compose.quickstart.yml') + } + + # Where a symbolic link points, or nothing when the path is not one. + function Get-LinkTarget([string]$Path) { + $item = Get-Item -LiteralPath $Path -Force -ErrorAction SilentlyContinue + if ($item -and $item.LinkType -eq 'SymbolicLink') { return "$(@($item.Target)[0])" } + return '' + } + + function Test-BinOnPath { + $bin = Join-Path $HOME '.local/bin' + return (-not $onWindows) -and (@($env:PATH -split ':') -contains $bin) -and (Test-Path -LiteralPath $bin -PathType Container) + } + + # The litellm-gateway command: the copy of this script kept beside .env, run + # against this folder, so up, down, status, and the rest work from anywhere. + function Install-Command { + $dir = $inst.Dir + $tmp = Join-Path $dir ('quickstart.ps1.' + [IO.Path]::GetRandomFileName()) + try { + Invoke-WebRequest -UseBasicParsing -Uri $scriptUrl -OutFile $tmp + if ([IO.File]::ReadAllText($tmp) -match '(?m)^\$code = Invoke-LiteLLMQuickstart') { + Move-Item -LiteralPath $tmp -Destination (Join-Path $dir 'quickstart.ps1') -Force + } + } catch { + # Offline, or a partial download: keep the copy already here, if any. + } finally { + Remove-Item -LiteralPath $tmp -Force -ErrorAction SilentlyContinue + } + if (-not (Test-Path -LiteralPath (Join-Path $dir 'quickstart.ps1'))) { return } + if ($onWindows) { + # cmd reads a batch file line by line as it runs, so the last line exits + # before uninstall has deleted the file under it. + $path = Join-Path $dir 'litellm-gateway.cmd' + $text = (@( + '@echo off' + 'setlocal' + 'rem Manages the LiteLLM Gateway in this folder. Run it for a menu, or with one of' + 'rem status, up, down, restart, logs, open, update, uninstall' + 'set "LITELLM_DIR=%~dp0."' + 'powershell -NoProfile -ExecutionPolicy Bypass -File "%~dp0quickstart.ps1" %* & exit /b' + ) -join "`r`n") + "`r`n" + } else { + $path = Join-Path $dir 'litellm-gateway' + $shell = if (Get-Command pwsh -CommandType Application -ErrorAction SilentlyContinue) { 'pwsh' } else { (Get-Process -Id $PID).Path } + $text = (@( + '#!/bin/sh' + '# Manages the LiteLLM Gateway in this folder. Run it for a menu, or with one of' + '# status, up, down, restart, logs, open, update, uninstall' + ("export LITELLM_DIR='" + $dir.Replace("'", "'\''") + "'") + ("exec '" + $shell.Replace("'", "'\''") + "' -NoProfile -File `"`$LITELLM_DIR/quickstart.ps1`" `"`$@`"") + ) -join "`n") + "`n" + } + $current = if (Test-Path -LiteralPath $path) { [IO.File]::ReadAllText($path) } else { '' } + if ($current -ne $text) { Write-PlainText $path $text } + if ($onWindows) { return } + $null = Invoke-Native chmod @('755', $path) + # Linked into ~/.local/bin when that is on PATH. A link another install made + # stays, unless that install is gone. + if (-not (Test-BinOnPath)) { return } + $bin = Join-Path $HOME '.local/bin' + $link = Join-Path $bin 'litellm-gateway' + $target = Get-LinkTarget $link + if ($target) { + if ($target -eq $path -or (Test-Path -LiteralPath ([IO.Path]::Combine($bin, $target)))) { return } + Remove-Item -LiteralPath $link -Force + } elseif (Test-Path -LiteralPath $link) { + return + } + try { $null = New-Item -ItemType SymbolicLink -Path $link -Target $path } catch { <# not allowed: the full path still works #> } + } + + # How this install's litellm-gateway command runs from anywhere, or nothing. + function Find-Command { + $name = if ($onWindows) { 'litellm-gateway.cmd' } else { 'litellm-gateway' } + $path = Join-Path $inst.Dir $name + if (-not (Test-Path -LiteralPath $path)) { return '' } + if ((Test-BinOnPath) -and (Get-LinkTarget (Join-Path $HOME '.local/bin/litellm-gateway')) -eq $path) { return 'litellm-gateway' } + $command = Tildify $path + if ($command -match ' ') { $command = "& `"$command`"" } + return $command + } + + # The command for an action: the litellm-gateway command, or Compose itself. + function Hint([string]$What) { + if ($inst.Command) { return "$($inst.Command) $What" } + $verb = switch ($What) { 'up' { 'up -d' } 'down' { 'stop' } 'uninstall' { 'down --volumes' } 'logs' { 'logs litellm' } default { $What } } + $where = Tildify $inst.Dir + if ($where -match ' ') { $where = "`"$where`"" } + return "cd $where; $engine compose -f docker-compose.quickstart.yml $verb" + } + + # The gateway container's status, such as "Up 2 hours", or nothing when there is none. + function Get-GatewayStatus { + $r = Invoke-Native $docker @('ps', '-a', '--filter', "label=com.docker.compose.project=$($inst.Project)", + '--filter', 'label=com.docker.compose.service=litellm', '--format', '{{.Status}}') + return "$($r.Out | Select-Object -First 1)" + } + + function Test-Running { return (Get-GatewayStatus) -like 'Up*' } + + function Test-HasContainers([switch]$All) { + $a = @('ps', '-q') + if ($All) { $a += '-a' } + $a += @('--filter', "label=com.docker.compose.project=$($inst.Project)") + return [bool](@((Invoke-Native $docker $a).Out | Where-Object { $_ }).Count) + } + + # The running version, from the API description: the first "version" in it is + # the gateway's, and it needs no key. + function Get-GatewayVersion { + try { + $r = Invoke-WebRequest -UseBasicParsing -TimeoutSec 5 -Uri "http://127.0.0.1:$($inst.Port)/openapi.json" + if ("$($r.Content)" -match '"version": *"([^"]*)"') { return $Matches[1] } + } catch { <# not up yet #> } + return '' + } + + # How many requests the gateway has handled, failed ones included, counted + # from its request logs. It needs the admin key. Empty when the gateway cannot say. + function Get-RequestCount { + $key = Read-EnvValue 'LITELLM_MASTER_KEY' + if (-not $key) { return '' } + $today = [DateTime]::UtcNow.ToString('yyyy-MM-dd', [Globalization.CultureInfo]::InvariantCulture) + try { + $r = Invoke-WebRequest -UseBasicParsing -TimeoutSec 5 -Headers @{ Authorization = "Bearer $key" } ` + -Uri "http://127.0.0.1:$($inst.Port)/global/activity?start_date=2000-01-01&end_date=$today" + if ("$($r.Content)" -match '"sum_api_requests": *([0-9]+)') { return $Matches[1] } + } catch { <# an older gateway, or no database yet #> } + return '' + } + + function Get-Images { + return @((Invoke-Native $docker ($cf + @('config', '--images'))).Out | Where-Object { $_ } | ForEach-Object { "$_" }) + } + + # Before starting: if something else holds the port, say what and how to free it. + # Moves this install to another port and saves it in .env. Writing the file in + # place keeps its permissions, so only you can still read the keys. + function Move-Port([int]$Port) { + $p = Join-Path $inst.Dir '.env' + $lines = @(Get-Content -LiteralPath $p -ErrorAction Stop | Where-Object { -not $_.StartsWith('LITELLM_PORT=') }) + Write-PlainText $p ((($lines + "LITELLM_PORT=$Port") -join "`n") + "`n") + $env:LITELLM_PORT = "$Port" + $inst.Port = $Port + $inst.Url = "http://localhost:$Port/ui" + } + + # Before starting: if something else holds the port, stop it when it is another + # quickstart install, or move to a free port when it is anything else. + function Confirm-PortFree { + $port = $inst.Port + if ((Test-Running) -or (Test-PortFree $port)) { return } + $o = Get-PortOwner $port + $who = if ($o.Owner) { $o.Owner } else { 'another program' } + $next = 0 + if ($o.Project) { + if ($interactive) { + $choice = Menu "Port $port is in use by $who. Stop it and start this one?" 1 @(@('Yes, stop it', $o.FreeIt), @('No')) + if ($choice -eq 1 -and (Invoke-WithSpinner 'Stopping it' '' $docker @('compose', '-p', $o.Project, 'stop')) -eq 0) { + Step ('Stopped ' + (Lit $o.Owner)) + return + } + } + } elseif (-not $portPinned) { + foreach ($candidate in ($port + 1)..($port + 100)) { + if (Test-PortFree $candidate) { $next = $candidate; break } + } + if ($next -and $interactive) { + $where = 'saved in ' + (Join-Path (Tildify $inst.Dir) '.env') + $choice = Menu "Port $port is in use by $who. Use port $next instead?" 1 @(@("Yes, use $next", $where), @('No')) + if ($choice -eq 1) { + Move-Port $next + Step "LiteLLM moved from port $port to {b:$next}" + return + } + } + } + $instead = if ($next) { ", such as $next" } else { '' } + Fail (Lit "Port $port is in use by $who.") ('Free it, or set LITELLM_PORT in ' + (Lit (Join-Path (Tildify $inst.Dir) '.env')) + " to another port$instead.") + if ($o.FreeIt) { Say (' To free it: ' + (Lit $o.FreeIt)) -Err } + Stop-Quickstart + } + + function Start-Gateway { + Confirm-PortFree + if (-not (Test-Path -LiteralPath 'docker-compose.quickstart.yml')) { Save-Compose } + $started = Get-Date + $detail = '' + foreach ($image in Get-Images) { + if ((Invoke-Native $docker @('image', 'inspect', $image)).Code -ne 0) { $detail = 'downloading images, first run only' } + } + if (-not $style) { if ($detail) { Say "Starting LiteLLM and Postgres ($detail)..." } else { Say 'Starting LiteLLM and Postgres...' } } + if ((Invoke-WithSpinner 'Starting LiteLLM and Postgres' $detail $docker ($cf + @('up', '-d'))) -ne 0) { + Fail "$engineName could not start LiteLLM and Postgres." 'Its output is above.' + Stop-Quickstart + } + if (-not (Wait-Ready $inst.Port)) { + Fail 'The gateway did not become ready in 3 minutes.' 'See what it logged:' + Say (' ' + (Lit (Hint 'logs'))) -Err + Stop-Quickstart + } + Step "LiteLLM is up at {a:$($inst.Url)} {d:$($sym.Dot) $(Elapsed $started)}" + } + + function Stop-Gateway { + if (-not (Test-HasContainers)) { Step 'LiteLLM is already stopped'; return } + if ((Invoke-WithSpinner 'Stopping LiteLLM and Postgres' '' $docker ($cf + @('stop'))) -ne 0) { + Fail "$engineName could not stop LiteLLM and Postgres." 'Its output is above.' + Stop-Quickstart + } + Step ("Stopped LiteLLM. Your models, keys, and logs are kept {d:$($sym.Dot) start it again: " + (Lit (Hint 'up')) + '}') + } + + function Restart-Gateway { + if ((Test-HasContainers) -and (Invoke-WithSpinner 'Stopping LiteLLM and Postgres' '' $docker ($cf + @('stop'))) -ne 0) { + Fail "$engineName could not stop LiteLLM and Postgres." 'Its output is above.' + Stop-Quickstart + } + Start-Gateway + } + + function Update-Gateway { + $before = '' + if (Test-Running) { $before = Get-GatewayVersion } + Save-Compose + Install-Command + $inst.Command = Find-Command + $images = Get-Images + $old = @(foreach ($image in $images) { + $r = Invoke-Native $docker @('image', 'inspect', '--format', '{{.Id}}', $image) + if ($r.Code -eq 0) { "$($r.Out | Select-Object -First 1)" } + }) + if (-not $style) { Say 'Downloading the latest LiteLLM and Postgres images...' } + if ((Invoke-WithSpinner 'Downloading the latest images' '' $docker ($cf + @('pull'))) -ne 0) { + Fail "$engineName could not download the images." 'Its output is above.' + Stop-Quickstart + } + Start-Gateway + # Delete the images this replaced. Removal fails, and they stay, while any + # container still uses them or another tag points at them. + $now = @(foreach ($image in $images) { "$((Invoke-Native $docker @('image', 'inspect', '--format', '{{.Id}}', $image)).Out | Select-Object -First 1)" }) + foreach ($id in $old) { + if ($id -and $now -notcontains $id) { $null = Invoke-Native $docker @('image', 'rm', $id) } + } + $after = Get-GatewayVersion + if ($before -and $before -eq $after) { + Step "LiteLLM $after is already the latest" + } elseif ($before -and $after) { + Step "Updated LiteLLM from $before to $after" + } elseif ($after) { + Step "LiteLLM $after is on the latest image" + } else { + Step 'LiteLLM is on the latest image' + } + } + + function Show-Logs([switch]$Follow) { + if (-not (Test-HasContainers -All)) { Step 'No logs yet: start LiteLLM first'; return } + if ($Follow) { + Say '{d:The gateway logs. Press Ctrl+C to stop following them.}' + $null = Invoke-Shown $docker ($cf + @('logs', '-f', '--tail', '100', 'litellm')) + return + } + $null = Invoke-Shown $docker ($cf + @('logs', '--tail', '100', 'litellm')) + } + + function Get-DesktopNote { + if ($engine -eq 'docker' -and ((Invoke-Native $docker @('info', '--format', '{{.OperatingSystem}}')).Out -match 'Docker Desktop')) { + return "listed as $($inst.Project) in Docker Desktop" + } + return '' + } + + function Show-Status { + $status = Get-GatewayStatus + $image = "$(Get-Images | Where-Object { $_ -match 'litellm' } | Select-Object -First 1)" + $desktop = Get-DesktopNote + $rows = New-Object System.Collections.ArrayList + if ($status -like 'Up*') { + $uptime = ($status -replace '^Up ', '') -replace ' \(.*$', '' + $version = Get-GatewayVersion + if (-not $version) { $version = 'starting' } + $requests = Get-RequestCount + $null = $rows.Add(@('Admin UI', $inst.Url)) + $null = $rows.Add(@('Version', $version)) + $null = $rows.Add(@('Uptime', $uptime)) + if ($requests -eq '0') { $requests = 'none yet' } elseif ($requests) { $requests = "$requests handled" } + if ($requests) { $null = $rows.Add(@('Requests', $requests)) } + $null = $rows.Add(@('Folder', (Tildify $inst.Dir))) + $null = $rows.Add(@('Image', $image)) + if ($desktop) { $null = $rows.Add(@('Desktop', $desktop)) } + $null = $rows.Add(@('Stop it', (Hint 'down'))) + Show-Box 'LiteLLM is running' $rows.ToArray() + } else { + $null = $rows.Add(@('Folder', (Tildify $inst.Dir))) + $null = $rows.Add(@('Image', $image)) + if ($desktop) { $null = $rows.Add(@('Desktop', $desktop)) } + $null = $rows.Add(@('Start it', (Hint 'up'))) + $null = $rows.Add(@('Remove it', (Hint 'uninstall'))) + Show-Box 'LiteLLM is stopped' $rows.ToArray() + } + } + + function Open-UI { + if (Test-Running) { Open-Url $inst.Url } else { Warn ('LiteLLM is stopped. Start it with: ' + (Lit (Hint 'up'))) } + } + + # Returns $false when the person keeps the install, so the menu can carry on. + function Remove-Gateway { + $dir = $inst.Dir + $volume = "$($inst.Project)_postgres_data" + if (-not (Test-Path -LiteralPath 'docker-compose.quickstart.yml')) { Save-Compose } + $images = Get-Images + $link = '' + if (-not $onWindows) { + $bin = Join-Path $HOME '.local/bin' + if ((Get-LinkTarget (Join-Path $bin 'litellm-gateway')) -eq (Join-Path $dir 'litellm-gateway')) { $link = Join-Path $bin 'litellm-gateway' } + } + if ($interactive) { + Write-Host '' + Say ('{b:This removes LiteLLM from ' + (Lit (Tildify $dir)) + ':}') + Say " - its containers ($engineName project $($inst.Project))" + Say " - its database, the $engineName volume $volume, with your models, keys, and spend logs" + Say (' - its files in ' + (Lit (Tildify $dir)) + ', including .env with the admin password') + if ($link) { Say (' - the litellm-gateway command in ' + (Lit (Tildify $bin))) } + if ((Menu 'Remove LiteLLM and delete its database?' 2 @(@('Yes, remove it'), @('No, keep it'))) -ne 1) { + Step 'Kept LiteLLM and its data' + return $false + } + } elseif (-not $noQuestions) { + Fail "Removing LiteLLM deletes its database ($engineName volume $volume)." 'To confirm, run it again with -Yes.' + Stop-Quickstart + } + + if (-not $style) { Say 'Removing the LiteLLM and Postgres containers and the database...' } + if ((Invoke-WithSpinner 'Removing the containers and the database' '' $docker ($cf + @('down', '--volumes', '--remove-orphans'))) -ne 0) { + Fail "$engineName could not remove them." 'Its output is above.' + Stop-Quickstart + } + Step "Removed the containers and the database {d:($volume)}" + + $kept = New-Object System.Collections.ArrayList + foreach ($image in $images) { + if ((Invoke-Native $docker @('image', 'inspect', $image)).Code -eq 0) { $null = $kept.Add($image) } + } + if ($kept.Count -and $interactive) { + # The size on disk, as the images list shows it (1.68GB, or 1.23 GB on Podman). + $bytes = 0.0 + foreach ($image in $kept) { + $size = "$((Invoke-Native $docker @('images', '--format', '{{.Size}}', $image)).Out | Select-Object -First 1)" + if ($size -match '^([\d.]+)\s*([kKMGT]?)') { + $unit = switch -CaseSensitive ($Matches[2]) { 'k' { 1e3 } 'K' { 1e3 } 'M' { 1e6 } 'G' { 1e9 } 'T' { 1e12 } default { 1 } } + $bytes += [double]::Parse($Matches[1], [Globalization.CultureInfo]::InvariantCulture) * $unit + } + } + $total = if ($bytes -gt 0) { ' (' + ($bytes / 1e9).ToString('0.0', [Globalization.CultureInfo]::InvariantCulture) + ' GB)' } else { '' } + $choice = Menu "Also delete the LiteLLM and Postgres images$($total)?" 2 @(@('Yes, delete them'), @('No, keep them', 'a reinstall starts faster')) + if ($choice -eq 1) { + foreach ($image in @($kept.ToArray())) { + if ((Invoke-Native $docker @('image', 'rm', $image)).Code -eq 0) { + Step "Deleted the image $image" + $kept.Remove($image) + } else { + Warn "Kept the image $image, which other containers use" + } + } + } + } + + foreach ($name in '.env', 'docker-compose.quickstart.yml', 'quickstart.ps1', 'quickstart.sh', 'litellm-gateway', 'litellm-gateway.cmd') { + Remove-Item -LiteralPath (Join-Path $dir $name) -Force -ErrorAction SilentlyContinue + } + $ignore = Join-Path $dir '.gitignore' + if ((Test-Path -LiteralPath $ignore) -and + -not ([IO.File]::ReadAllText($ignore) -split "`r?`n" | Where-Object { $_ -and $_ -ne '*' -and $_ -ne '.env' })) { + Remove-Item -LiteralPath $ignore -Force + } + if ($link) { Remove-Item -LiteralPath $link -Force -ErrorAction SilentlyContinue } + # Windows will not delete the folder a program is working in, this one included. + Set-Location -LiteralPath $HOME + [Environment]::CurrentDirectory = $HOME + try { + [IO.Directory]::Delete($dir, $false) + Step ('Deleted ' + (Lit (Tildify $dir))) + } catch { + $why = if (@(Get-ChildItem -LiteralPath $dir -Force -ErrorAction SilentlyContinue).Count) { 'it has other files' } else { 'a program is using it' } + Step ('Deleted the LiteLLM files in ' + (Lit (Tildify $dir)) + " {d:(the folder stays: $why)}") + } + + $left = New-Object System.Collections.ArrayList + if (Test-HasContainers -All) { $null = $left.Add('containers') } + if ((Invoke-Native $docker @('volume', 'inspect', $volume)).Code -eq 0) { $null = $left.Add("volume $volume") } + if (Test-Path -LiteralPath (Join-Path $dir '.env')) { $null = $left.Add((Join-Path (Tildify $dir) '.env')) } + if ($left.Count) { + Warn ('Still there: ' + (Lit ($left.ToArray() -join ', '))) + return $true + } + Step 'LiteLLM is removed. Nothing is left running, and its data is deleted' + if ($kept.Count) { + Say (" {d:The images stay, so a reinstall starts faster. To delete them:} $engine image rm " + ($kept.ToArray() -join ' ')) + } + return $true + } + + # Over SSH the clipboard would be the server's, not yours. On Linux, + # Set-Clipboard without xclip quietly keeps the text inside PowerShell, so + # only Windows and macOS, which always have a clipboard, get the option. + function Test-CanCopy { + return $interactive -and (-not $env:SSH_CONNECTION) -and (-not $env:SSH_TTY) -and + ($onWindows -or ((Test-Path variable:IsMacOS) -and $IsMacOS)) + } + + # The menu a rerun shows for an existing install. Its options follow whether + # LiteLLM is running, and come back after each choice until Done. + function Show-ManageMenu { + Show-Status + $canCopy = Test-CanCopy + $pick = 1 + while ($true) { + $options = New-Object System.Collections.ArrayList + if (Test-Running) { + $question = "LiteLLM is running at $($inst.Url). What next?" + $null = $options.Add(@('Open the admin UI in your browser')) + if ($canCopy) { $null = $options.Add(@('Copy the admin password', 'to the clipboard, ready to paste')) } + $null = $options.Add(@('Stop it', 'keeps your models, keys, and logs')) + $null = $options.Add(@('Restart it')) + $null = $options.Add(@('Show the logs', 'the last 100 lines')) + $null = $options.Add(@('Update it', 'download the latest image and restart')) + $null = $options.Add(@('Remove it', 'delete its containers, database, and files')) + $null = $options.Add(@('Done')) + } else { + $question = 'LiteLLM is stopped. What next?' + $null = $options.Add(@('Start it', "at $($inst.Url)")) + $null = $options.Add(@('Update it', 'download the latest image and start')) + $null = $options.Add(@('Show the logs', 'the last 100 lines')) + $null = $options.Add(@('Remove it', 'delete its containers, database, and files')) + $null = $options.Add(@('Done')) + } + if ($pick -gt $options.Count) { $pick = 1 } + $choice = Menu $question $pick $options.ToArray() + $picked = $options[$choice - 1][0] + $pick = 1 + switch -Wildcard ($picked) { + 'Open*' { Open-Url $inst.Url; $pick = $options.Count } + 'Copy*' { Copy-Password $inst.Dir } + 'Start*' { Start-Gateway } + 'Stop*' { Stop-Gateway } + 'Restart*' { Restart-Gateway } + 'Show*' { Show-Logs } + 'Update*' { Update-Gateway } + 'Remove*' { if (Remove-Gateway) { return } } + default { return } + } + } + } + # ------------------------------------------------------------ main $savedEncoding = $null + $savedLocation = (Get-Location).ProviderPath + $savedPort = $env:LITELLM_PORT try { if ($unicode -and $onWindows) { # Windows PowerShell writes the console in the OEM code page, which has no check marks. @@ -377,7 +1057,9 @@ function Invoke-LiteLLMQuickstart { [Net.ServicePointManager]::SecurityProtocol = [Net.ServicePointManager]::SecurityProtocol -bor [Net.SecurityProtocolType]::Tls12 } - if ($style) { + if ($action) { + # An action prints only its own result. + } elseif ($style) { Write-Host '' Say "{u:$($sym.Head) LiteLLM quickstart}" Say '{d: The gateway, Postgres, and the admin UI in one command}' @@ -416,6 +1098,32 @@ function Invoke-LiteLLMQuickstart { Fail "$engineName is installed but not running." $startHint return 1 } + + if ($action) { + $existing = Find-Install + if (-not $existing) { + $s = [IO.Path]::DirectorySeparatorChar + $places = if ($env:LITELLM_DIR) { $env:LITELLM_DIR } else { "this folder, .${s}litellm-gateway, or ~${s}litellm-gateway" } + Fail 'No LiteLLM install found' ('in ' + (Lit $places) + '.') + $runner = if ($onWindows) { 'powershell' } else { 'pwsh' } + Say " Install it with: $runner -c `"irm $scriptUrl | iex`"" -Err + return 1 + } + Use-Install $existing + $inst.Command = Find-Command + switch ($action) { + 'status' { Show-Status } + 'start' { if (Test-Running) { Step "LiteLLM is already running at {a:$($inst.Url)}" } else { Start-Gateway } } + 'stop' { Stop-Gateway } + 'restart' { Restart-Gateway } + 'logs' { if ($interactive) { Show-Logs -Follow } else { Show-Logs } } + 'open' { Open-UI } + 'update' { Update-Gateway } + 'uninstall' { $null = Remove-Gateway } + } + return 0 + } + $server = Invoke-Native $docker @('version', '--format', '{{.Server.Version}}') $engineVersion = if ($server.Code -eq 0 -and $server.Out.Count) { "$($server.Out[0]) " } else { '' } $composeVersion = "$($compose.Out[0])".TrimStart('v') @@ -423,13 +1131,12 @@ function Invoke-LiteLLMQuickstart { $homeDir = Join-Path $HOME 'litellm-gateway' $hereDir = Join-Path (Get-Location).ProviderPath 'litellm-gateway' - $found = $false - if ($env:LITELLM_DIR) { + $dir = Find-Install + $found = [bool]$dir + if ($found) { + # installed before + } elseif ($env:LITELLM_DIR) { $dir = $env:LITELLM_DIR - } elseif (Test-Path -LiteralPath (Join-Path $hereDir '.env')) { - $dir = $hereDir; $found = $true # installed in this folder before - } elseif (Test-Path -LiteralPath (Join-Path $homeDir '.env')) { - $dir = $homeDir; $found = $true # installed in the home folder before } elseif ($hereDir -eq $homeDir) { $dir = $homeDir } else { @@ -474,13 +1181,21 @@ function Invoke-LiteLLMQuickstart { # commits if the folder becomes a repository later. $ignore = Join-Path $dir '.gitignore' $current = if (Test-Path -LiteralPath $ignore) { [IO.File]::ReadAllText($ignore) } else { '' } - if (($current -split "`r?`n") -notcontains '.env') { + if (($current -split "`r?`n") -notcontains '.env' -and ($current -split "`r?`n") -notcontains '*') { # Start on a new line if the file does not end with one. $lead = if ($current.Length -gt 0 -and -not $current.EndsWith("`n")) { "`n" } else { '' } Write-PlainText $ignore "$lead.env`n" -Append } } + if ($found -and $interactive) { + Use-Install $dir + if (-not (Test-Path -LiteralPath 'quickstart.ps1') -or -not (Find-Command)) { Install-Command } + $inst.Command = Find-Command + Show-ManageMenu + return 0 + } + # The engine names containers and the database volume after the project, so # an install outside the home folder gets its own name and never shares a # database with another litellm-gateway folder. Windows paths ignore case, @@ -505,35 +1220,13 @@ function Invoke-LiteLLMQuickstart { } } - Invoke-WebRequest -UseBasicParsing -Uri $composeUrl -OutFile (Join-Path $dir 'docker-compose.quickstart.yml') + Save-Compose Step 'Downloaded docker-compose.quickstart.yml' - $saved = '' - if (Test-Path -LiteralPath $envFile) { - $m = Get-Content -LiteralPath $envFile | Where-Object { $_ -match '^LITELLM_PORT=(.*)$' } | Select-Object -Last 1 - if ($m) { $saved = ($m -split '=', 2)[1] } - } - if ($env:LITELLM_PORT) { - $port = [int]$env:LITELLM_PORT - } elseif ($saved) { - $port = [int]$saved - } elseif (Test-Path -LiteralPath $envFile) { - $port = 4000 # an existing install without a saved port runs on the compose default - } else { - $port = 4000 - while (-not (Test-PortFree $port)) { - $port++ - if ($port -gt 4099) { - Fail 'Ports 4000 to 4099 are all in use.' 'Set LITELLM_PORT to a free port and run this again.' - return 1 - } - } - if ($port -eq 4000) { Step 'Port {b:4000} is free' } else { Warn "Port 4000 is in use, so LiteLLM will use {b:$port}" } - } - if (Test-Path -LiteralPath $envFile) { Step 'Reusing .env, so existing keys and data keep working' } else { + $port = Select-Port $rng = [Security.Cryptography.RandomNumberGenerator]::Create() $hex = { param([int]$n) @@ -566,51 +1259,32 @@ function Invoke-LiteLLMQuickstart { Step 'Generated .env with your master key, salt key, and database password {d:(keep this file; only you can read it)}' } - # Compose prefers values already set in the environment over .env, so drop - # inherited ones: .env stays the only source for keys and the project name. - foreach ($name in 'LITELLM_MASTER_KEY', 'LITELLM_SALT_KEY', 'POSTGRES_PASSWORD', 'COMPOSE_PROJECT_NAME') { - Remove-Item "Env:$name" -ErrorAction SilentlyContinue - } - # The bind address follows .env when .env sets it. For an older .env without - # it, a value already in the environment is kept. - if (Get-Content -LiteralPath $envFile | Where-Object { $_ -match '^LITELLM_BIND=' }) { Remove-Item Env:LITELLM_BIND -ErrorAction SilentlyContinue } - $env:LITELLM_PORT = "$port" + Use-Install $dir + Install-Command + $inst.Command = Find-Command + Start-Gateway - $started = Get-Date - $detail = '' - foreach ($image in (Invoke-Native $docker @('compose', '-f', 'docker-compose.quickstart.yml', 'config', '--images')).Out) { - if ($image -and (Invoke-Native $docker @('image', 'inspect', "$image")).Code -ne 0) { $detail = 'downloading images, first run only' } - } - if (-not $style) { Say 'Starting LiteLLM and Postgres (the first run downloads the images)...' } - $code = Invoke-WithSpinner 'Starting LiteLLM and Postgres' $detail $docker @('compose', '-f', 'docker-compose.quickstart.yml', 'up', '-d') - if ($code -ne 0) { - Fail "$engineName could not start LiteLLM and Postgres." 'Its output is above.' - return 1 - } - if (-not (Wait-Ready $port)) { - Fail 'The gateway did not become ready in 3 minutes.' 'See what it logged:' - Say (' cd ' + (Lit (Tildify $dir)) + "; $engine compose -f docker-compose.quickstart.yml logs litellm") -Err - return 1 - } - Step ("LiteLLM and Postgres are up {d:$($sym.Dot) $(Elapsed $started)}") - - $url = "http://localhost:$port/ui" - $where = Tildify $dir - if ($where -match ' ') { $where = "`"$where`"" } + $url = $inst.Url Write-Host '' - # Over SSH the clipboard would be the server's, not yours. On Linux, - # Set-Clipboard without xclip quietly keeps the text inside PowerShell, so - # only Windows and macOS, which always have a clipboard, get the option. - $canCopy = $interactive -and (-not $env:SSH_CONNECTION) -and - ($onWindows -or ((Test-Path variable:IsMacOS) -and $IsMacOS)) + $canCopy = Test-CanCopy $password = "the LITELLM_MASTER_KEY value in $(Join-Path (Tildify $dir) '.env')" if ($canCopy) { $password = "copy it below, or LITELLM_MASTER_KEY in $(Join-Path (Tildify $dir) '.env')" } - Show-Box 'LiteLLM is running' @( - , @('Admin UI', $url) - , @('Username', 'admin') - , @('Password', $password) - , @('Next', 'in the UI, open Models + Endpoints > Add Model and paste a provider API key') - , @('Stop it', "cd $where; $engine compose -f docker-compose.quickstart.yml down")) + # A rerun finds ~\litellm-gateway on its own; any other folder needs naming. + $rerun = 'rerun the install command' + if ($dir -ne $homeDir) { $rerun += " with LITELLM_DIR=$(Tildify $dir)" } + $manage = if ($inst.Command) { "run $($inst.Command), or $rerun" } else { $rerun } + $rows = New-Object System.Collections.ArrayList + $null = $rows.Add(@('Admin UI', $url)) + $null = $rows.Add(@('Username', 'admin')) + $null = $rows.Add(@('Password', $password)) + $null = $rows.Add(@('Next', 'in the UI, open Models + Endpoints > Add Model and paste a provider API key')) + $null = $rows.Add(@('Stop it', (Hint 'down'))) + $null = $rows.Add(@('Start it', (Hint 'up'))) + $null = $rows.Add(@('Remove it', (Hint 'uninstall'))) + $desktop = Get-DesktopNote + if ($desktop) { $null = $rows.Add(@('Desktop', $desktop)) } + $null = $rows.Add(@('Manage', $manage)) + Show-Box 'LiteLLM is running' $rows.ToArray() if ($canCopy) { # Come back to the menu after each choice, with the next step as the default. @@ -625,15 +1299,20 @@ function Invoke-LiteLLMQuickstart { else { break } } } elseif ($interactive) { - if ((Menu 'Open the admin UI in your browser?' 1 @(, @('Yes'), @('No'))) -eq 1) { Open-Url $url } + if ((Menu 'Open the admin UI in your browser?' 1 @(@('Yes'), @('No'))) -eq 1) { Open-Url $url } } return 0 + } catch { + if ($_.Exception.Message -eq $exitSignal) { return 1 } + throw } finally { if ($savedEncoding) { [Console]::OutputEncoding = $savedEncoding } + if ($savedPort) { $env:LITELLM_PORT = $savedPort } else { Remove-Item Env:LITELLM_PORT -ErrorAction SilentlyContinue } + if (Test-Path -LiteralPath $savedLocation) { Set-Location -LiteralPath $savedLocation } } } -$code = Invoke-LiteLLMQuickstart -Yes:$Yes | Select-Object -Last 1 +$code = Invoke-LiteLLMQuickstart -Yes:$Yes -Rest $args | Select-Object -Last 1 $global:LASTEXITCODE = $code # Run from a file, or as `powershell -c "irm ... | iex"`, the exit code should # reach the caller. Pasted into an open window, `exit` would close that window, diff --git a/scripts/quickstart.sh b/scripts/quickstart.sh index 66f69f34ff3..0182b216788 100755 --- a/scripts/quickstart.sh +++ b/scripts/quickstart.sh @@ -10,10 +10,16 @@ # # Asks where to keep the files, then offers to copy the admin password and # open the admin UI; every question has a default you accept by pressing Enter. +# Run it again for a menu to open, stop, start, update, or remove the install. # It asks nothing when there is no terminal, under CI or Claude Code, or when # run with --yes. # -# --yes, -y no questions: install to ~/litellm-gateway, don't open a browser +# It also takes an action, which the litellm-gateway command it installs takes too: +# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh | sh -s -- stop +# status, start (up), stop (down), restart, logs, open, update, uninstall +# +# --yes, -y no questions: install to ~/litellm-gateway, don't open a browser, +# and let uninstall delete the database without asking # LITELLM_DIR folder to install into (skips the folder question) # LITELLM_PORT port for the gateway (default 4000, or the next free one) # NO_COLOR plain output, which agents, CI, and log files always get @@ -29,6 +35,9 @@ set -eu COMPOSE_URL="${LITELLM_COMPOSE_URL:-https://raw.githubusercontent.com/BerriAI/litellm/main/docker/docker-compose.quickstart.yml}" +SCRIPT_URL="${LITELLM_SCRIPT_URL:-https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh}" +# A port set by the caller is never moved; later steps export LITELLM_PORT too. +PORT_PINNED="${LITELLM_PORT:+1}" # ---------------------------------------------------------------- output @@ -296,6 +305,33 @@ menu() { # ---------------------------------------------------------------- steps +usage() { + cat <<'USAGE' +Usage: quickstart.sh [--yes] [action] + +With no action it installs the LiteLLM Gateway, or shows a menu to manage the +install it finds. Actions: + status whether it is running, and where + start start it (also: up) + stop stop it; models, keys, and logs are kept (also: down) + restart stop it and start it again + logs show the gateway's logs + open open the admin UI in your browser + update download the latest image and restart on it + uninstall remove its containers, its database, and its files + help show this help (also: -h, --help) + + --yes, -y ask nothing; without a terminal, uninstall needs it + +From the install command: + curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh | sh -s -- status +USAGE +} + +compose() { + "$ENGINE" compose -f docker-compose.quickstart.yml "$@" +} + port_free() { # curl exits 7 when nothing accepts the connection. rc=0 @@ -303,16 +339,64 @@ port_free() { [ "$rc" = 7 ] } +# What holds a port: sets OWNER to a description, FREE_IT to the command that +# frees it, and OWNER_PROJECT when it is another quickstart install. All stay +# empty when the holder cannot be named, such as another user's process. +port_owner() { + OWNER='' FREE_IT='' OWNER_PROJECT='' + name="$("$ENGINE" ps --format '{{.Names}} {{.Ports}}' 2>/dev/null | grep -E "[:]$1->" | head -n 1 | cut -d ' ' -f 1)" + if [ -n "$name" ]; then + OWNER_PROJECT="$("$ENGINE" inspect --format '{{index .Config.Labels "com.docker.compose.project"}}' "$name" 2>/dev/null)" || + OWNER_PROJECT='' + # Only another quickstart install keeps its project, since that is the one + # this script offers to stop. Whoever started the container set the label, + # so it must also be a name Compose would accept. + case "$OWNER_PROJECT" in + *[!a-z0-9_-]*) OWNER_PROJECT='' ;; + esac + case "$OWNER_PROJECT" in + litellm-gateway*) + OWNER="another LiteLLM quickstart install ($ENGINE_NAME project $OWNER_PROJECT)" + FREE_IT="$ENGINE compose -p $OWNER_PROJECT stop" + ;; + *) OWNER_PROJECT='' OWNER="the $ENGINE_NAME container $name" FREE_IT="$ENGINE stop $name" ;; + esac + return 0 + fi + command -v lsof >/dev/null 2>&1 || return 0 + pid="$(lsof -nP -iTCP:"$1" -sTCP:LISTEN -t 2>/dev/null | head -n 1)" + [ -n "$pid" ] || return 0 + FREE_IT="kill $pid" + case "$(ps -o args= -p "$pid" 2>/dev/null)" in + *litellm*) OWNER="a litellm server (pid $pid), maybe one started after pip install litellm" ;; + *) OWNER="$(basename "$(ps -o comm= -p "$pid" 2>/dev/null || echo a program)") (pid $pid)" ;; + esac +} + +# Sets DIR to an existing install: LITELLM_DIR, this folder, ./litellm-gateway, +# or ~/litellm-gateway, whichever has a .env first. +locate_install() { + DIR='' + if [ -n "${LITELLM_DIR:-}" ]; then + if [ -f "$LITELLM_DIR/.env" ]; then DIR="$LITELLM_DIR"; fi + elif [ -f .env ] && [ -f docker-compose.quickstart.yml ]; then + DIR="$(pwd)" + elif [ -f litellm-gateway/.env ]; then + DIR="$(pwd)/litellm-gateway" + elif [ -f "$HOME/litellm-gateway/.env" ]; then + DIR="$HOME/litellm-gateway" + fi + [ -n "$DIR" ] +} + pick_folder() { home_dir="$HOME/litellm-gateway" here_dir="$(pwd)/litellm-gateway" found=0 - if [ -n "${LITELLM_DIR:-}" ]; then + if locate_install; then + found=1 + elif [ -n "${LITELLM_DIR:-}" ]; then DIR="$LITELLM_DIR" - elif [ -f "$here_dir/.env" ]; then - DIR="$here_dir" found=1 # installed in this folder before - elif [ -f "$home_dir/.env" ]; then - DIR="$home_dir" found=1 # installed in the home folder before elif [ "$here_dir" = "$home_dir" ]; then DIR="$home_dir" else @@ -337,11 +421,11 @@ pick_folder() { # An ignore rule does not cover a tracked file, so new keys written here # would show up as a change to commit. fail "Git tracks a .env file in this folder," "so your keys could be committed." - cat >&2 <<'EOF' + cat >&2 <<'MSG' Install into another folder (set LITELLM_DIR), or stop tracking the file first with: git rm --cached .env -EOF +MSG exit 1 fi if [ "$created" = 1 ]; then @@ -360,7 +444,7 @@ EOF elif ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then # An existing folder outside git: ignore only .env, so it stays out of # commits if the folder becomes a repository later. - if ! grep -qxF '.env' .gitignore 2>/dev/null; then + if ! grep -qxE '\*|\.env' .gitignore 2>/dev/null; then # Start on a new line if the file does not end with one. if [ -s .gitignore ] && [ -n "$(tail -c 1 .gitignore)" ]; then printf '\n' >>.gitignore; fi printf '.env\n' >>.gitignore @@ -390,7 +474,8 @@ pick_port() { if [ "$PORT" = 4000 ]; then step "Port ${C_BOLD}4000${C_OFF} is free" else - warn "Port 4000 is in use, so LiteLLM will use ${C_BOLD}$PORT${C_OFF}" + port_owner 4000 + warn "Port 4000 is in use${OWNER:+ by $OWNER}, so LiteLLM will use ${C_BOLD}$PORT${C_OFF}" fi fi export LITELLM_PORT="$PORT" @@ -406,19 +491,210 @@ check_new_install() { # old database volume would lock the gateway out. Stop and explain instead. if "$ENGINE" volume inspect "${project}_postgres_data" >/dev/null 2>&1; then fail "Found a database from an earlier install" "($ENGINE_NAME volume ${project}_postgres_data)" - cat >&2 <&2 </dev/null && grep -q '^main "\$@"$' "$tmp"; then + mv -f "$tmp" quickstart.sh + else + rm -f "$tmp" + fi + fi + [ -f quickstart.sh ] || return 0 + quoted="$(printf '%s' "$DIR" | sed "s/'/'\\\\''/g")" + { + printf '#!/bin/sh\n' + printf '# Manages the LiteLLM Gateway in this folder. Run it for a menu, or with one of\n' + printf '# status, up, down, restart, logs, open, update, uninstall\n' + printf "export LITELLM_DIR='%s'\n" "$quoted" + # shellcheck disable=SC2016 # expanded when the command runs + printf 'exec sh "$LITELLM_DIR/quickstart.sh" "$@"\n' + } >litellm-gateway + chmod 755 litellm-gateway + # Linked into ~/.local/bin when that is on PATH. A link another install made + # stays, unless that install is gone. + link="$HOME/.local/bin/litellm-gateway" + case ":$PATH:" in *":$HOME/.local/bin:"*) ;; *) return 0 ;; esac + [ -d "$HOME/.local/bin" ] || return 0 + if [ -L "$link" ]; then + if [ "$(readlink "$link")" = "$DIR/litellm-gateway" ] || [ -e "$link" ]; then return 0; fi + rm -f "$link" + elif [ -e "$link" ]; then + return 0 + fi + ln -s "$DIR/litellm-gateway" "$link" +} + +# Sets COMMAND to how this install's litellm-gateway command runs from anywhere. +find_command() { + COMMAND='' + [ -x "$DIR/litellm-gateway" ] || return 0 + COMMAND="$(tildify "$DIR")/litellm-gateway" + case ":$PATH:" in + *":$HOME/.local/bin:"*) + if [ "$(readlink "$HOME/.local/bin/litellm-gateway" 2>/dev/null)" = "$DIR/litellm-gateway" ]; then + COMMAND=litellm-gateway + fi + ;; + esac +} + +hint() { + if [ -n "$COMMAND" ]; then + printf '%s %s' "$COMMAND" "$1" + return 0 + fi + case "$1" in + up) set -- up -d ;; + down) set -- stop ;; + uninstall) set -- down --volumes ;; + logs) set -- logs litellm ;; + esac + printf 'cd %s && %s compose -f docker-compose.quickstart.yml %s' "$(tildify "$DIR")" "$ENGINE" "$*" +} + +# The gateway container's status, such as "Up 2 hours", or nothing when there is none. +gateway_status() { + "$ENGINE" ps -a --filter "label=com.docker.compose.project=$project" \ + --filter label=com.docker.compose.service=litellm --format '{{.Status}}' 2>/dev/null | head -n 1 +} + +is_running() { + case "$(gateway_status)" in + Up*) return 0 ;; + *) return 1 ;; + esac +} + +project_containers() { + "$ENGINE" ps -q "$@" --filter "label=com.docker.compose.project=$project" 2>/dev/null +} + +# The running version, from the API description: the first "version" in it is +# the gateway's, and it needs no key. +gateway_version() { + curl -fsS --max-time 5 "http://127.0.0.1:$PORT/openapi.json" 2>/dev/null | + grep -o '"version": *"[^"]*"' | head -n 1 | sed 's/.*"\([^"]*\)"$/\1/' +} + +# How many requests the gateway has handled, failed ones included, counted from +# its request logs. It needs the admin key, which curl reads from stdin so it +# stays off the command line. Empty when the gateway cannot say. +gateway_requests() { + key="$(sed -n 's/^LITELLM_MASTER_KEY=//p' .env | tail -n 1)" + [ -n "$key" ] || return 0 + printf 'header = "Authorization: Bearer %s"\n' "$key" | + curl -fsS --max-time 5 -K - \ + "http://127.0.0.1:$PORT/global/activity?start_date=2000-01-01&end_date=$(date -u +%Y-%m-%d)" 2>/dev/null | + grep -o '"sum_api_requests": *[0-9]*' | grep -o '[0-9]*$' +} + +# Stops the other quickstart install port_owner found. +stop_owner() { + "$ENGINE" compose -p "$OWNER_PROJECT" stop +} + +# Sets NEXT_PORT to the first free port after $1, or empty when the next 100 +# are all taken. +next_free_port() { + NEXT_PORT=$(($1 + 1)) + while ! port_free "$NEXT_PORT"; do + NEXT_PORT=$((NEXT_PORT + 1)) + if [ "$NEXT_PORT" -gt $(($1 + 100)) ]; then + NEXT_PORT='' + return 0 + fi + done +} + +# Moves this install to NEXT_PORT and saves it in .env, through a fresh mktemp +# file (mode 600) and a rename so the keys are never half written. +move_port() { + ENV_TMP="$(mktemp .env.XXXXXX)" || return 1 + if ! awk '!/^LITELLM_PORT=/' .env >"$ENV_TMP" || + ! printf 'LITELLM_PORT=%s\n' "$NEXT_PORT" >>"$ENV_TMP" || ! mv -f "$ENV_TMP" .env; then + rm -f "$ENV_TMP" + ENV_TMP='' + return 1 + fi + ENV_TMP='' + PORT="$NEXT_PORT" + export LITELLM_PORT="$PORT" + url="http://localhost:$PORT/ui" +} + +# Before starting: if something else holds the port, stop it when it is another +# quickstart install, or move to a free port when it is anything else. +check_port() { + if is_running || port_free "$PORT"; then return 0; fi + port_owner "$PORT" + tab="$(printf '\t')" NEXT_PORT='' + if [ -n "$OWNER_PROJECT" ]; then + if [ "$INTERACTIVE" = 1 ]; then + menu "Port $PORT is in use by $OWNER. Stop it and start this one?" 1 \ + "Yes, stop it${tab}$FREE_IT" "No" + if [ "$CHOICE" = 1 ] && spin "Stopping it" "" stop_owner; then + step "Stopped $OWNER" + return 0 + fi + fi + elif [ -z "$PORT_PINNED" ]; then + next_free_port "$PORT" + if [ -n "$NEXT_PORT" ] && [ "$INTERACTIVE" = 1 ]; then + menu "Port $PORT is in use by ${OWNER:-another program}. Use port $NEXT_PORT instead?" 1 \ + "Yes, use $NEXT_PORT${tab}saved in $(tildify "$DIR")/.env" "No" + if [ "$CHOICE" = 1 ]; then + old="$PORT" + if move_port; then + step "LiteLLM moved from port $old to ${C_BOLD}$PORT${C_OFF}" + return 0 + fi + fail "Could not write $(tildify "$DIR")/.env." + exit 1 + fi + fi + fi + fail "Port $PORT is in use by ${OWNER:-another program}." \ + "Free it, or set LITELLM_PORT in $(tildify "$DIR")/.env to another port${NEXT_PORT:+, such as $NEXT_PORT}." + if [ -n "$FREE_IT" ]; then printf ' To free it: %s\n' "$FREE_IT" >&2; fi + exit 1 } wait_ready() { @@ -430,6 +706,214 @@ wait_ready() { done } +launch() { + check_port + [ -f docker-compose.quickstart.yml ] || fetch_compose + started="$(date +%s)" + detail='' + for image in $(compose config --images 2>/dev/null); do + "$ENGINE" image inspect "$image" >/dev/null 2>&1 || detail="downloading images, first run only" + done + [ "$STYLE" = 1 ] || echo "Starting LiteLLM and Postgres${detail:+ ($detail)}..." + if ! spin "Starting LiteLLM and Postgres" "$detail" compose up -d; then + fail "$ENGINE_NAME could not start LiteLLM and Postgres." "Its output is above." + exit 1 + fi + if ! spin "Waiting for the gateway to be ready" "" wait_ready; then + fail "The gateway did not become ready in 3 minutes." "See what it logged:" + echo " $(hint logs)" >&2 + exit 1 + fi + step "LiteLLM is up at ${C_ACC}$url${C_OFF} ${C_DIM}· $(elapsed_since "$started")" +} + +stop_stack() { + if [ -z "$(project_containers)" ]; then + step "LiteLLM is already stopped" + return 0 + fi + if ! spin "Stopping LiteLLM and Postgres" "" compose stop; then + fail "$ENGINE_NAME could not stop LiteLLM and Postgres." "Its output is above." + exit 1 + fi + step "Stopped LiteLLM. Your models, keys, and logs are kept ${C_DIM}· start it again: $(hint up)" +} + +restart_stack() { + if [ -n "$(project_containers)" ] && ! spin "Stopping LiteLLM and Postgres" "" compose stop; then + fail "$ENGINE_NAME could not stop LiteLLM and Postgres." "Its output is above." + exit 1 + fi + launch +} + +update_stack() { + before='' + if is_running; then before="$(gateway_version)"; fi + fetch_compose + install_command + find_command + images="$(compose config --images 2>/dev/null)" || images='' + old='' + for image in $images; do + old="$old $("$ENGINE" image inspect --format '{{.Id}}' "$image" 2>/dev/null)" + done + [ "$STYLE" = 1 ] || echo "Downloading the latest LiteLLM and Postgres images..." + if ! spin "Downloading the latest images" "" compose pull; then + fail "$ENGINE_NAME could not download the images." "Its output is above." + exit 1 + fi + launch + # Delete the images this replaced. Removal fails, and they stay, while any + # container still uses them or another tag points at them. + for id in $old; do + for image in $images; do + [ "$id" != "$("$ENGINE" image inspect --format '{{.Id}}' "$image" 2>/dev/null)" ] || continue 2 + done + "$ENGINE" image rm "$id" >/dev/null 2>&1 || true + done + after="$(gateway_version)" + if [ -n "$before" ] && [ "$before" = "$after" ]; then + step "LiteLLM $after is already the latest" + elif [ -n "$before" ] && [ -n "$after" ]; then + step "Updated LiteLLM from $before to $after" + else + step "LiteLLM ${after:+$after }is on the latest image" + fi +} + +show_logs() { + if [ -z "$(project_containers -a)" ]; then + step "No logs yet: start LiteLLM first" + return 0 + fi + if [ "${1:-}" = follow ]; then + printf '%sThe gateway logs. Press Ctrl+C to stop following them.%s\n' "${C_DIM}" "${C_OFF}" + exec "$ENGINE" compose -f docker-compose.quickstart.yml logs -f --tail 100 litellm + fi + compose logs --tail 100 litellm +} + +desktop_row() { + if [ "$ENGINE" = docker ] && docker info --format '{{.OperatingSystem}}' 2>/dev/null | grep -q 'Docker Desktop'; then + printf 'Desktop|listed as %s in Docker Desktop' "$project" + fi +} + +show_status() { + status="$(gateway_status)" + image="$(compose config --images 2>/dev/null | grep -m 1 litellm)" || image='' + desktop="$(desktop_row)" + case "$status" in + Up*) + uptime="${status#Up }" + version="$(gateway_version)" + requests="$(gateway_requests)" || requests='' + case "$requests" in + '') ;; + 0) requests='Requests|none yet' ;; + *) requests="Requests|$requests handled" ;; + esac + box "LiteLLM is running" "Admin UI|$url" "Version|${version:-starting}" "Uptime|${uptime%% (*}" \ + ${requests:+"$requests"} "Folder|$(tildify "$DIR")" "Image|$image" ${desktop:+"$desktop"} "Stop it|$(hint down)" + ;; + *) + box "LiteLLM is stopped" "Folder|$(tildify "$DIR")" "Image|$image" ${desktop:+"$desktop"} \ + "Start it|$(hint up)" "Remove it|$(hint uninstall)" + ;; + esac +} + +open_ui() { + if is_running; then + open_url "$url" + else + warn "LiteLLM is stopped. Start it with: $(hint up)" + fi +} + +# Returns 1 when the person keeps the install, so the menu can carry on. +uninstall() { + volume="${project}_postgres_data" + [ -f docker-compose.quickstart.yml ] || fetch_compose + images="$(compose config --images 2>/dev/null)" || images='' + link='' + if [ "$(readlink "$HOME/.local/bin/litellm-gateway" 2>/dev/null)" = "$DIR/litellm-gateway" ]; then + link="$HOME/.local/bin/litellm-gateway" + fi + if [ "$INTERACTIVE" = 1 ]; then + printf '\n%sThis removes LiteLLM from %s:%s\n' "${C_BOLD}" "$(tildify "$DIR")" "${C_OFF}" + printf ' - its containers (%s project %s)\n' "$ENGINE_NAME" "$project" + printf ' - its database, the %s volume %s, with your models, keys, and spend logs\n' "$ENGINE_NAME" "$volume" + printf ' - its files in %s, including .env with the admin password\n' "$(tildify "$DIR")" + if [ -n "$link" ]; then printf ' - the litellm-gateway command in %s\n' "$(tildify "$HOME/.local/bin")"; fi + menu "Remove LiteLLM and delete its database?" 2 "Yes, remove it" "No, keep it" + if [ "$CHOICE" != 1 ]; then + step "Kept LiteLLM and its data" + return 1 + fi + elif [ "$NO_QUESTIONS" != 1 ]; then + fail "Removing LiteLLM deletes its database ($ENGINE_NAME volume $volume)." "To confirm, run it again with --yes." + exit 1 + fi + + [ "$STYLE" = 1 ] || echo "Removing the LiteLLM and Postgres containers and the database..." + if ! spin "Removing the containers and the database" "" compose down --volumes --remove-orphans; then + fail "$ENGINE_NAME could not remove them." "Its output is above." + exit 1 + fi + step "Removed the containers and the database ${C_DIM}($volume)" + + kept='' + for image in $images; do + "$ENGINE" image inspect "$image" >/dev/null 2>&1 && kept="$kept $image" + done + if [ -n "$kept" ] && [ "$INTERACTIVE" = 1 ]; then + # The size on disk, as the images list shows it (1.68GB, or 1.23 GB on Podman). + size="$(for image in $kept; do "$ENGINE" images --format '{{.Size}}' "$image" 2>/dev/null | head -n 1; done | + awk '{ n = $1 + 0; u = $0; gsub(/[0-9. ]/, "", u); m = 1 + if (u ~ /^[kK]/) m = 1e3; else if (u ~ /^M/) m = 1e6; else if (u ~ /^G/) m = 1e9; else if (u ~ /^T/) m = 1e12 + s += n * m } END { if (s > 0) printf " (%.1f GB)", s / 1e9 }')" + tab="$(printf '\t')" + menu "Also delete the LiteLLM and Postgres images${size}?" 2 \ + "Yes, delete them" "No, keep them${tab}a reinstall starts faster" + if [ "$CHOICE" = 1 ]; then + for image in $kept; do + if "$ENGINE" image rm "$image" >/dev/null 2>&1; then + step "Deleted the image $image" + kept="$(printf '%s' "$kept" | sed "s| $image||")" + else + warn "Kept the image $image, which other containers use" + fi + done + fi + fi + + rm -f .env docker-compose.quickstart.yml quickstart.sh quickstart.ps1 litellm-gateway litellm-gateway.cmd + if [ -f .gitignore ] && ! grep -qvxE '\*|\.env' .gitignore; then rm -f .gitignore; fi + if [ -n "$link" ]; then rm -f "$link"; fi + cd "$HOME" + if rmdir "$DIR" 2>/dev/null; then + step "Deleted $(tildify "$DIR")" + else + step "Deleted the LiteLLM files in $(tildify "$DIR") ${C_DIM}(the folder stays: it has other files)" + fi + + left='' + if [ -n "$(project_containers -a)" ]; then left=" containers,"; fi + if "$ENGINE" volume inspect "$volume" >/dev/null 2>&1; then left="$left volume $volume,"; fi + if [ -f "$DIR/.env" ]; then left="$left $(tildify "$DIR")/.env,"; fi + if [ -n "$left" ]; then + warn "Still there:${left%,}" + return 0 + fi + step "LiteLLM is removed. Nothing is left running, and its data is deleted" + if [ -n "$kept" ]; then + printf ' %sThe images stay, so a reinstall starts faster. To delete them:%s %s image rm%s\n' \ + "${C_DIM}" "${C_OFF}" "$ENGINE" "$kept" + fi +} + open_url() { for opener in open xdg-open; do if command -v "$opener" >/dev/null 2>&1 && "$opener" "$1" >/dev/null 2>&1; then @@ -443,6 +927,7 @@ open_url() { # A clipboard to copy the admin password to, or nothing. Over SSH the # clipboard would be the server's, not yours. find_clipboard() { + CLIP='' [ -z "${SSH_CONNECTION:-}${SSH_TTY:-}" ] || return 0 if command -v pbcopy >/dev/null 2>&1; then CLIP=pbcopy elif [ -n "${WAYLAND_DISPLAY:-}" ] && command -v wl-copy >/dev/null 2>&1; then CLIP=wl-copy @@ -472,7 +957,6 @@ copy_password() { # After the summary: copy the password, open the admin UI, or finish. The # menu comes back after each choice with the next step as its default. next_steps() { - url="$1" [ "$INTERACTIVE" = 1 ] || return 0 if [ -z "$CLIP" ]; then menu "Open the admin UI in your browser?" 1 "Yes" "No" @@ -494,14 +978,65 @@ next_steps() { done } +# The menu a rerun shows for an existing install. Its options follow whether +# LiteLLM is running, and come back after each choice until Done. +manage() { + show_status + find_clipboard + tab="$(printf '\t')" + pick=1 + while :; do + if is_running; then + question="LiteLLM is running at $url. What next?" + set -- "Open the admin UI in your browser" + [ -z "$CLIP" ] || set -- "$@" "Copy the admin password${tab}to the clipboard, ready to paste" + set -- "$@" "Stop it${tab}keeps your models, keys, and logs" "Restart it" \ + "Show the logs${tab}the last 100 lines" "Update it${tab}download the latest image and restart" \ + "Remove it${tab}delete its containers, database, and files" "Done" + else + question="LiteLLM is stopped. What next?" + set -- "Start it${tab}at $url" "Update it${tab}download the latest image and start" \ + "Show the logs${tab}the last 100 lines" "Remove it${tab}delete its containers, database, and files" "Done" + fi + [ "$pick" -le $# ] || pick=1 + menu "$question" "$pick" "$@" + picked='' + eval "picked=\${$CHOICE}" + pick=1 + case "$picked" in + Open*) open_url "$url"; pick=$# ;; + Copy*) copy_password ;; + Start*) launch ;; + Stop*) stop_stack ;; + Restart*) restart_stack ;; + Show*) show_logs ;; + Update*) update_stack ;; + Remove*) if uninstall; then return 0; fi ;; + *) return 0 ;; + esac + done +} + main() { - NO_QUESTIONS=0 + NO_QUESTIONS=0 ACTION='' for arg in "$@"; do case "$arg" in -y | --yes) NO_QUESTIONS=1 ;; - *) echo "Unknown option: $arg" >&2; exit 1 ;; + -h | --help | help) usage; exit 0 ;; + status | start | up | stop | down | restart | logs | open | update | uninstall) + if [ -n "$ACTION" ]; then + echo "Choose one action, not both $ACTION and $arg" >&2 + exit 1 + fi + ACTION="$arg" + ;; + *) echo "Unknown option: $arg" >&2; usage >&2; exit 1 ;; esac done + case "$ACTION" in + up) ACTION=start ;; + down) ACTION=stop ;; + esac setup_output detect_terminal @@ -510,7 +1045,9 @@ main() { trap restore_terminal EXIT trap on_interrupt INT TERM - if [ "$STYLE" = 1 ]; then + if [ -n "$ACTION" ]; then + : + elif [ "$STYLE" = 1 ]; then printf '\n%s%s LiteLLM quickstart%s\n' "${C_ACC}${C_BOLD}" "$S_HEAD" "${C_OFF}" printf '%s The gateway, Postgres, and the admin UI in one command%s\n\n' "${C_DIM}" "${C_OFF}" else @@ -528,13 +1065,13 @@ main() { done if [ -z "$ENGINE" ]; then fail "No container engine found (docker, podman, or nerdctl)." "The LiteLLM Gateway runs in a container alongside a Postgres database." - cat >&2 <<'EOF' + cat >&2 <<'MSG' Install Docker (https://docs.docker.com/get-docker/), Podman (https://podman.io), or Rancher Desktop (https://rancherdesktop.io), then run this again. Or deploy in one click (Railway or Render): https://docs.litellm.ai/docs/proxy/docker_quick_start Only need to call models from Python? pip install litellm -EOF +MSG exit 1 fi case "$ENGINE" in @@ -552,19 +1089,50 @@ EOF fi exit 1 fi - command -v openssl >/dev/null 2>&1 || { fail "openssl is required to generate keys."; exit 1; } + + if [ -n "$ACTION" ]; then + if ! locate_install; then + fail "No LiteLLM install found" "in ${LITELLM_DIR:-this folder, ./litellm-gateway, or ~/litellm-gateway}." + printf ' Install it with: curl -fsSL %s | sh\n' "$SCRIPT_URL" >&2 + exit 1 + fi + use_install + find_command + case "$ACTION" in + status) show_status ;; + start) if is_running; then step "LiteLLM is already running at ${C_ACC}$url${C_OFF}"; else launch; fi ;; + stop) stop_stack ;; + restart) restart_stack ;; + logs) if [ "$INTERACTIVE" = 1 ]; then show_logs follow; else show_logs; fi ;; + open) open_ui ;; + update) update_stack ;; + uninstall) uninstall || true ;; + esac + return 0 + fi + engine_version="$("$ENGINE" version --format '{{.Server.Version}}' 2>/dev/null)" || engine_version='' step "$ENGINE_NAME ${engine_version:+$engine_version }with Compose ${compose_version#v} is running" pick_folder - [ -f .env ] || check_new_install - curl -fsSL -o docker-compose.quickstart.yml "$COMPOSE_URL" + if [ "$found" = 1 ] && [ "$INTERACTIVE" = 1 ]; then + use_install + if [ ! -x litellm-gateway ] || [ ! -f quickstart.sh ]; then install_command; fi + find_command + manage + return 0 + fi + if [ "$found" = 0 ]; then + command -v openssl >/dev/null 2>&1 || { fail "openssl is required to generate keys."; exit 1; } + check_new_install + fi + fetch_compose step "Downloaded docker-compose.quickstart.yml" - pick_port - if [ -f .env ]; then + if [ "$found" = 1 ]; then step "Reusing .env, so existing keys and data keep working" else + pick_port master="$(openssl rand -hex 32)" salt="$(openssl rand -hex 32)" db_password="$(openssl rand -hex 24)" @@ -583,47 +1151,34 @@ EOF step "Generated .env with your master key, salt key, and database password ${C_DIM}(keep this file; only you can read it)" fi - # Compose prefers values already set in the shell over .env, so drop any - # inherited ones: .env stays the only source for keys and the project name. - unset LITELLM_MASTER_KEY LITELLM_SALT_KEY POSTGRES_PASSWORD COMPOSE_PROJECT_NAME - # The bind address follows .env when .env sets it (every install this script - # creates does). For an older .env without it, a value exported in the shell - # is kept, so an intentional LITELLM_BIND=127.0.0.1: is not dropped. - if grep -q '^LITELLM_BIND=' .env; then unset LITELLM_BIND; fi + use_install + install_command + find_command + launch - started="$(date +%s)" - detail='' - for image in $("$ENGINE" compose -f docker-compose.quickstart.yml config --images 2>/dev/null); do - "$ENGINE" image inspect "$image" >/dev/null 2>&1 || detail="downloading images, first run only" - done - [ "$STYLE" = 1 ] || echo "Starting LiteLLM and Postgres (the first run downloads the images)..." - if ! spin "Starting LiteLLM and Postgres" "$detail" start_stack; then - fail "$ENGINE_NAME could not start LiteLLM and Postgres." "Its output is above." - exit 1 - fi - if ! spin "Waiting for the gateway to be ready" "" wait_ready; then - fail "The gateway did not become ready in 3 minutes." "See what it logged:" - echo " cd $(tildify "$DIR") && $ENGINE compose -f docker-compose.quickstart.yml logs litellm" >&2 - exit 1 - fi - step "LiteLLM and Postgres are up ${C_DIM}· $(elapsed_since "$started")" - - url="http://localhost:$PORT/ui" password="the LITELLM_MASTER_KEY value in $(tildify "$DIR")/.env" CLIP='' if [ "$INTERACTIVE" = 1 ]; then find_clipboard; fi if [ -n "$CLIP" ]; then password="copy it below, or LITELLM_MASTER_KEY in $(tildify "$DIR")/.env" fi + desktop="$(desktop_row)" + # A rerun finds ~/litellm-gateway on its own; any other folder needs naming. + rerun="rerun the install command" + [ "$DIR" = "$HOME/litellm-gateway" ] || rerun="$rerun with LITELLM_DIR=$(tildify "$DIR")" echo box "LiteLLM is running" \ "Admin UI|$url" \ "Username|admin" \ "Password|$password" \ "Next|in the UI, open Models + Endpoints > Add Model and paste a provider API key" \ - "Stop it|cd $(tildify "$DIR") && $ENGINE compose -f docker-compose.quickstart.yml down" + "Stop it|$(hint down)" \ + "Start it|$(hint up)" \ + "Remove it|$(hint uninstall)" \ + ${desktop:+"$desktop"} \ + "Manage|${COMMAND:+run $COMMAND, or }$rerun" - next_steps "$url" + next_steps } main "$@" diff --git a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py index 8c678cc1a92..cc1d5c35788 100644 --- a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py +++ b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py @@ -18,9 +18,12 @@ from redis import Redis from tests.integration._support.client import ( Gateway, + Scenario, delete_key_if_present, eventually, gateway_from_environment, + object_value, + string_value, ) from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import owned_proxy @@ -28,6 +31,7 @@ from tests.integration._support.wire import Reply, Request, Wire, wire_server _AUDIENCE: Final = "litellm-integration" _CLIENT_CLAIM: Final = "client_id" +_END_USER_CLAIM: Final = "eu" _KEY_ID: Final = "integration-signing-key" _UPSTREAM_KEY: Final = "synthetic-openai-key" _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" @@ -211,15 +215,23 @@ def _jwks_respond(private_key: rsa.RSAPrivateKey) -> Callable[[Request], Reply]: return respond -def _signed_jwt(private_key: rsa.RSAPrivateKey, client_id: str, issuer: str | None = None) -> str: +def _signed_jwt( + private_key: rsa.RSAPrivateKey, + client_id: str, + issuer: str | None = None, + *, + subject: str | None = None, + end_user: str | None = None, +) -> str: now: Final = int(time.time()) claims: Final[dict[str, JsonValue]] = { - "sub": "integration-user-" + uuid.uuid4().hex, + "sub": subject if subject is not None else "integration-user-" + uuid.uuid4().hex, "client_id": client_id, "aud": _AUDIENCE, "iat": now, "exp": now + 300, **({"iss": issuer} if issuer is not None else {}), + **({_END_USER_CLAIM: end_user} if end_user is not None else {}), } return jwt.encode(claims, private_key, algorithm="RS256", headers={"kid": _KEY_ID}) @@ -909,3 +921,287 @@ def test_repointing_a_warmed_mapping_to_a_new_issuer_moves_it_on_both_workers( upstream, tuple(_expected_chat_request(model, prompt) for prompt in prompts), ) + + +class _EndUserSpendRow(BaseModel): + api_key: str + end_user: str | None + + +@dataclass(frozen=True, slots=True) +class _EndUserClaimRuntime: + gateway: Gateway + model: str + private_key: rsa.RSAPrivateKey + + +@dataclass(frozen=True, slots=True) +class _MappedClient: + client_id: str + key: str + key_hash: str + + +def _end_user_claim_config(directory: Path, model: str, upstream_url: str) -> Path: + config: Final = directory / "jwt_end_user_claim_config.yaml" + config.write_text( + json.dumps( + { + "model_list": ( + { + "model_name": model, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + }, + ), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "disable_model_info_refresh": True, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "user_id_upsert": True, + "virtual_key_claim_field": _CLIENT_CLAIM, + "unregistered_jwt_client_behavior": "auto_register", + "end_user_id_jwt_field": _END_USER_CLAIM, + }, + }, + } + ) + ) + return config + + +@pytest.fixture(scope="module") +def end_user_claims(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_EndUserClaimRuntime]: + directory: Final = tmp_path_factory.mktemp("jwt-end-user-claim") + model: Final = "jwt-end-user-claim-" + uuid.uuid4().hex + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + with gateway_from_environment() as base, wire_server(_jwks_respond(private_key)) as jwks: + config: Final = _end_user_claim_config(directory, model, base.upstream_url) + overrides: Final = {"JWT_PUBLIC_KEY_URL": f"{jwks.url}/jwks", "JWT_AUDIENCE": _AUDIENCE} + with owned_proxy(base, directory / "proxy", overrides, config=config) as candidate: + eventually(lambda: _model_is_available(candidate, model), bool, seconds=30) + yield _EndUserClaimRuntime(gateway=candidate, model=model, private_key=private_key) + + +def _mapped_client(runtime: _EndUserClaimRuntime, scenario: Scenario) -> _MappedClient: + client_id: Final = "integration-client-" + uuid.uuid4().hex + key: Final = _create_key(runtime.gateway, runtime.model) + scenario.cleanups.callback(delete_key_if_present, runtime.gateway, key) + mapping: Final = _create_mapping(runtime.gateway, key, client_id) + scenario.cleanups.callback(_delete_mapping_if_present, runtime.gateway, mapping.id) + return _MappedClient(client_id=client_id, key=key, key_hash=sha256(key.encode()).hexdigest()) + + +def _end_user_chat( + runtime: _EndUserClaimRuntime, token: str, *, body_user: str | None = None, header_user: str | None = None +) -> httpx.Response: + return runtime.gateway.request( + "POST", + "/v1/chat/completions", + { + **_chat_body(runtime.model, "end user claim " + uuid.uuid4().hex), + **({"user": body_user} if body_user is not None else {}), + }, + key=token, + headers={"x-litellm-end-user-id": header_user} if header_user is not None else None, + ) + + +def _billed_end_user(response: httpx.Response) -> _EndUserSpendRow: + assert response.status_code == 200, response.text + identity: Final = _ChatResponse.model_validate_json(response.content).id + rows: Final = TypeAdapter(tuple[_EndUserSpendRow, ...]).validate_python( + eventually( + lambda: read_rows('SELECT api_key, end_user FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda values: len(values) == 1, + seconds=70, + ) + ) + return rows[0] + + +def _customer_ids(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT user_id FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_id,)) + + +def _delete_customer_if_present(gateway: Gateway, user_id: str) -> None: + if not _customer_ids(user_id): + return + response: Final = gateway.request("POST", "/customer/delete", {"user_ids": [user_id]}) + assert response.status_code == 200, response.text + + +def _forget_auto_registered_client(client_id: str, subject: str) -> None: + write_rows( + 'DELETE FROM "LiteLLM_VerificationToken" WHERE token IN ' + '(SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value = %s)', + (client_id,), + ) + write_rows('DELETE FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value = %s', (client_id,)) + write_rows('DELETE FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + + +def test_a_mapped_jwt_bills_and_registers_its_end_user_claim(end_user_claims: _EndUserClaimRuntime) -> None: + gateway: Final = end_user_claims.gateway + end_user: Final = "integration-end-user-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + client: Final = _mapped_client(end_user_claims, scenario) + scenario.cleanups.callback(_delete_customer_if_present, gateway, end_user) + token: Final = _signed_jwt(end_user_claims.private_key, client.client_id, end_user=end_user) + + billed: Final = _billed_end_user(_end_user_chat(end_user_claims, token)) + + assert billed == _EndUserSpendRow(api_key=client.key_hash, end_user=end_user) + customers: Final = eventually(lambda: _customer_ids(end_user), lambda rows: len(rows) == 1, seconds=70) + assert customers == [{"user_id": end_user}] + + +@pytest.mark.parametrize("spelling", ("body", "header")) +def test_a_mapped_jwt_end_user_claim_wins_over_the_caller_supplied_end_user( + end_user_claims: _EndUserClaimRuntime, spelling: str +) -> None: + gateway: Final = end_user_claims.gateway + end_user: Final = "integration-end-user-" + uuid.uuid4().hex + supplied: Final = "integration-supplied-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + client: Final = _mapped_client(end_user_claims, scenario) + scenario.cleanups.callback(_delete_customer_if_present, gateway, end_user) + scenario.cleanups.callback(_delete_customer_if_present, gateway, supplied) + token: Final = _signed_jwt(end_user_claims.private_key, client.client_id, end_user=end_user) + + response: Final = _end_user_chat( + end_user_claims, + token, + body_user=supplied if spelling == "body" else None, + header_user=supplied if spelling == "header" else None, + ) + + assert _billed_end_user(response) == _EndUserSpendRow(api_key=client.key_hash, end_user=end_user) + + +@pytest.mark.parametrize("claim", ("", None), ids=("empty", "missing")) +@pytest.mark.parametrize("spelling", ("body", "header")) +def test_a_mapped_jwt_without_an_end_user_claim_falls_back_to_the_caller_supplied_end_user( + end_user_claims: _EndUserClaimRuntime, spelling: str, claim: str | None +) -> None: + gateway: Final = end_user_claims.gateway + supplied: Final = "integration-supplied-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + client: Final = _mapped_client(end_user_claims, scenario) + scenario.cleanups.callback(_delete_customer_if_present, gateway, supplied) + token: Final = _signed_jwt(end_user_claims.private_key, client.client_id, end_user=claim) + + response: Final = _end_user_chat( + end_user_claims, + token, + body_user=supplied if spelling == "body" else None, + header_user=supplied if spelling == "header" else None, + ) + + assert _billed_end_user(response) == _EndUserSpendRow(api_key=client.key_hash, end_user=supplied) + + +def test_an_auto_registered_client_keeps_its_end_user_claim_after_the_first_request( + end_user_claims: _EndUserClaimRuntime, +) -> None: + gateway: Final = end_user_claims.gateway + client_id: Final = "integration-client-" + uuid.uuid4().hex + subject: Final = "integration-user-" + uuid.uuid4().hex + end_user: Final = "integration-end-user-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + scenario.cleanups.callback(_forget_auto_registered_client, client_id, subject) + scenario.cleanups.callback(_delete_customer_if_present, gateway, end_user) + token: Final = _signed_jwt(end_user_claims.private_key, client_id, subject=subject, end_user=end_user) + + registering: Final = _billed_end_user(_end_user_chat(end_user_claims, token)) + mapped: Final = _billed_end_user(_end_user_chat(end_user_claims, token)) + + mappings: Final = read_rows( + 'SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_name=%s AND jwt_claim_value=%s', + (_CLIENT_CLAIM, client_id), + ) + assert len(mappings) == 1, mappings + expected: Final = _EndUserSpendRow(api_key=string_value(mappings[0]["token"]), end_user=end_user) + assert (registering, mapped) == (expected, expected) + + +def test_a_mapped_jwt_end_user_claim_naming_an_over_budget_customer_is_refused_like_a_body_user( + end_user_claims: _EndUserClaimRuntime, +) -> None: + gateway: Final = end_user_claims.gateway + end_user: Final = "integration-end-user-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + client: Final = _mapped_client(end_user_claims, scenario) + created: Final = gateway.request( + "POST", "/customer/new", {"user_id": end_user, "max_budget": 0.05, "spend": 1.0} + ) + assert created.status_code == 200, created.text + scenario.cleanups.callback(_delete_customer_if_present, gateway, end_user) + token: Final = _signed_jwt(end_user_claims.private_key, client.client_id, end_user=end_user) + + by_body: Final = _end_user_chat(end_user_claims, client.key, body_user=end_user) + by_claim: Final = _end_user_chat(end_user_claims, token) + + assert by_body.status_code == 422, by_body.text + assert object_value(object_value(by_body.json())["error"])["type"] == "budget_exceeded", by_body.text + assert (by_claim.status_code, by_claim.json()) == (by_body.status_code, by_body.json()) + + +def test_a_mapped_jwt_end_user_claim_is_budget_checked_instead_of_an_over_budget_body_user( + end_user_claims: _EndUserClaimRuntime, +) -> None: + gateway: Final = end_user_claims.gateway + end_user: Final = "integration-end-user-" + uuid.uuid4().hex + over_budget: Final = "integration-over-budget-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + client: Final = _mapped_client(end_user_claims, scenario) + created: Final = gateway.request( + "POST", "/customer/new", {"user_id": over_budget, "max_budget": 0.05, "spend": 1.0} + ) + assert created.status_code == 200, created.text + scenario.cleanups.callback(_delete_customer_if_present, gateway, over_budget) + scenario.cleanups.callback(_delete_customer_if_present, gateway, end_user) + token: Final = _signed_jwt(end_user_claims.private_key, client.client_id, end_user=end_user) + + refused: Final = _end_user_chat(end_user_claims, client.key, body_user=over_budget) + served: Final = _end_user_chat(end_user_claims, token, body_user=over_budget) + + assert refused.status_code == 422, refused.text + assert _billed_end_user(served) == _EndUserSpendRow(api_key=client.key_hash, end_user=end_user) + + +@pytest.mark.parametrize("spelling", ("body", "header")) +def test_a_mapped_jwt_without_an_end_user_claim_budget_checks_the_caller_supplied_end_user( + end_user_claims: _EndUserClaimRuntime, spelling: str +) -> None: + gateway: Final = end_user_claims.gateway + over_budget: Final = "integration-over-budget-" + uuid.uuid4().hex + with gateway.scenario() as scenario: + client: Final = _mapped_client(end_user_claims, scenario) + created: Final = gateway.request( + "POST", "/customer/new", {"user_id": over_budget, "max_budget": 0.05, "spend": 1.0} + ) + assert created.status_code == 200, created.text + scenario.cleanups.callback(_delete_customer_if_present, gateway, over_budget) + token: Final = _signed_jwt(end_user_claims.private_key, client.client_id) + + refused: Final = _end_user_chat( + end_user_claims, + token, + body_user=over_budget if spelling == "body" else None, + header_user=over_budget if spelling == "header" else None, + ) + + assert refused.status_code == 422, refused.text + assert object_value(object_value(refused.json())["error"])["type"] == "budget_exceeded", refused.text diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index 4d0987ec1f9..c214352fdb0 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -4443,6 +4443,39 @@ CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( ), ) + +CLAUDE_SONNET_5_5_MESSAGES_STREAM: Final = CostTrackingTestCase( + name="claude-sonnet-5-5-messages_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-sonnet-5-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-sonnet-5-5","id":"msg_$REQUEST_ID","type":"message","role":"assistant","content":[],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"global"},"diagnostics":null}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}', + 'event: ping\ndata: {"type":"ping"}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig_$REQUEST_ID"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"scripted "}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":1}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"container":null},"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":412,"output_tokens_details":{"thinking_tokens":0}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.01287, input_cost=0.006072, output_cost=0.006798, prompt_tokens=1840, completion_tokens=412 + ), +) CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CostTrackingTestCase( name="claude-haiku-5-5-input_text", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -4806,6 +4839,39 @@ CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( ), ) + +CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = CostTrackingTestCase( + name="claude-haiku-5-5-messages_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-haiku-5-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-haiku-5-5","id":"msg_$REQUEST_ID","type":"message","role":"assistant","content":[],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"global"},"diagnostics":null}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}', + 'event: ping\ndata: {"type":"ping"}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig_$REQUEST_ID"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"scripted "}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":1}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"container":null},"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":412,"output_tokens_details":{"thinking_tokens":0}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412 + ), +) CLAUDE_FABLE_5_INPUT_TEXT: Final = CostTrackingTestCase( name="claude-fable-5-input_text", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5532,6 +5598,35 @@ CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( ), ) + +CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-1-messages_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5-1", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-fable-5-1","id":"msg_$REQUEST_ID","type":"message","role":"assistant","content":[],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"global"},"diagnostics":null}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: ping\ndata: {"type":"ping"}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"scripted "}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"container":null},"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":412,"output_tokens_details":{"thinking_tokens":0}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 + ), +) CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( name="claude-opus-5-5-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -5836,6 +5931,39 @@ CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( ), ) + +CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CostTrackingTestCase( + name="claude-opus-5-5-messages_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-opus-5-5", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-opus-5-5","id":"msg_$REQUEST_ID","type":"message","role":"assistant","content":[],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"global"},"diagnostics":null}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}', + 'event: ping\ndata: {"type":"ping"}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig_$REQUEST_ID"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"scripted "}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":1}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"container":null},"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":412,"output_tokens_details":{"thinking_tokens":0}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.0156, input_cost=0.00736, output_cost=0.00824, prompt_tokens=1840, completion_tokens=412 + ), +) CLAUDE_SONNET_5_5_TOOL_CALL: Final = CostTrackingTestCase( name="claude-sonnet-5-5-tool_call", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -6877,12 +7005,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_5_CACHE_WRITE_5M, CLAUDE_SONNET_5_5_CACHE_WRITE_1H, CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_STREAM, CLAUDE_HAIKU_5_5_INPUT_TEXT, CLAUDE_HAIKU_5_5_STREAM, CLAUDE_HAIKU_5_5_CACHE_READ, CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_FABLE_5_INPUT_TEXT, CLAUDE_FABLE_5_STREAM, CLAUDE_FABLE_5_CACHE_READ, @@ -6895,11 +7025,13 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_OPUS_5_5_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, CLAUDE_OPUS_5_5_CACHE_WRITE_5M, CLAUDE_OPUS_5_5_CACHE_WRITE_1H, CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, CLAUDE_SONNET_5_5_TOOL_CALL, CLAUDE_SONNET_5_5_THINKING, CLAUDE_HAIKU_5_5_TOOL_CALL, @@ -6938,6 +7070,10 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=CLAUDE_FABLE_5_INPUT_TEXT, streamed=CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_STREAM), StreamParityTestCase(plain=CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_STREAM), + StreamParityTestCase(plain=CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_MESSAGES_STREAM), + StreamParityTestCase(plain=CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_SONNET_5_5_MESSAGES_STREAM), + StreamParityTestCase(plain=CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_MESSAGES_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index 45ac7398d73..3c353a107f4 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -3097,7 +3097,16 @@ US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = CostTrackingTestCase( body={ "output": {"message": {"role": "assistant", "content": [{"text": "Hello."}]}}, "stopReason": "end_turn", - "usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252}, + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, "metrics": {"latencyMs": 42}, }, ), @@ -3508,11 +3517,14 @@ US_ANTHROPIC_CLAUDE_SONNET_5_5_CACHE_READ: Final = CostTrackingTestCase( "output": {"message": {"content": [{"text": "scripted answer 5eff7f685f14"}], "role": "assistant"}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 12288, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 640, "outputTokens": 380, + "serverToolUsage": {}, "totalTokens": 13308, - "cacheReadInputTokens": 12288, - "cacheWriteInputTokens": 0, }, }, ), @@ -3972,11 +3984,14 @@ US_ANTHROPIC_CLAUDE_OPUS_5_5_CACHE_READ: Final = CostTrackingTestCase( "output": {"message": {"content": [{"text": "scripted answer bbd4da2e326b"}], "role": "assistant"}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 12288, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 640, "outputTokens": 380, + "serverToolUsage": {}, "totalTokens": 13308, - "cacheReadInputTokens": 12288, - "cacheWriteInputTokens": 0, }, }, ), @@ -4230,11 +4245,14 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_CACHE_READ: Final = CostTrackingTestCase( "output": {"message": {"content": [{"text": "scripted answer cf757e3e24df"}], "role": "assistant"}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 12288, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 640, "outputTokens": 380, + "serverToolUsage": {}, "totalTokens": 13308, - "cacheReadInputTokens": 12288, - "cacheWriteInputTokens": 0, }, }, ), @@ -4746,11 +4764,14 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( "output": {"message": {"content": [{"text": "scripted answer 1a5a397f1134"}], "role": "assistant"}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 12288, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 640, "outputTokens": 380, + "serverToolUsage": {}, "totalTokens": 13308, - "cacheReadInputTokens": 12288, - "cacheWriteInputTokens": 0, }, }, ), @@ -5134,11 +5155,14 @@ US_ANTHROPIC_CLAUDE_SONNET_5_5_THINKING: Final = CostTrackingTestCase( "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer fe9c9714d04d"}]}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 1240, "outputTokens": 412, + "serverToolUsage": {}, "totalTokens": 1652, - "cacheReadInputTokens": 0, - "cacheWriteInputTokens": 0, }, "metrics": {"latencyMs": 2071}, }, @@ -5428,11 +5452,14 @@ US_ANTHROPIC_CLAUDE_OPUS_5_5_THINKING: Final = CostTrackingTestCase( }, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 1240, "outputTokens": 4040, + "serverToolUsage": {}, "totalTokens": 5280, - "cacheReadInputTokens": 0, - "cacheWriteInputTokens": 0, }, "metrics": {"latencyMs": 2071}, }, @@ -5560,11 +5587,14 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_THINKING: Final = CostTrackingTestCase( "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer 90a6bb1996a5"}]}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 1240, "outputTokens": 412, + "serverToolUsage": {}, "totalTokens": 1652, - "cacheReadInputTokens": 0, - "cacheWriteInputTokens": 0, }, "metrics": {"latencyMs": 2071}, }, @@ -5839,11 +5869,14 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer bca700fd52a2"}]}}, "stopReason": "end_turn", "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, "inputTokens": 1240, "outputTokens": 412, + "serverToolUsage": {}, "totalTokens": 1652, - "cacheReadInputTokens": 0, - "cacheWriteInputTokens": 0, }, "metrics": {"latencyMs": 2071}, }, diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index c4f56d4be5b..6ad0ff95615 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -8116,16 +8116,6 @@ GPT_6_1_SOL_CACHE_READ: Final = GPT_5_6_CACHE_READ.model_copy( "reasoning_tokens": 0, "rejected_prediction_tokens": 0, }, - "latency_checkpoint": { - "engine_tbt_ms": 23, - "engine_ttft_ms": 484, - "engine_ttlt_ms": 999, - "pre_inference_ms": 151, - "service_tbt_ms": 25, - "service_ttft_ms": 1085, - "service_ttlt_ms": 1610, - "user_visible_ttft_ms": 934, - }, "prompt_tokens": 12928, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 12288}, "total_tokens": 13308, @@ -8200,16 +8190,6 @@ GPT_6_LUNA_CACHE_READ: Final = GPT_5_6_CACHE_READ.model_copy( "reasoning_tokens": 0, "rejected_prediction_tokens": 0, }, - "latency_checkpoint": { - "engine_tbt_ms": 7, - "engine_ttft_ms": 110, - "engine_ttlt_ms": 706, - "pre_inference_ms": 254, - "service_tbt_ms": 7, - "service_ttft_ms": 647, - "service_ttlt_ms": 1238, - "user_visible_ttft_ms": 393, - }, "prompt_tokens": 12928, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 12288}, "total_tokens": 13308, @@ -8284,16 +8264,6 @@ GPT_6_ASTRA_CACHE_READ: Final = GPT_5_6_CACHE_READ.model_copy( "reasoning_tokens": 0, "rejected_prediction_tokens": 0, }, - "latency_checkpoint": { - "engine_tbt_ms": 25, - "engine_ttft_ms": 697, - "engine_ttlt_ms": 1373, - "pre_inference_ms": 142, - "service_tbt_ms": 26, - "service_ttft_ms": 1289, - "service_ttlt_ms": 1956, - "user_visible_ttft_ms": 1147, - }, "prompt_tokens": 12928, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 12288}, "total_tokens": 13308, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index bf5d57c2312..4cdd7fefd28 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -3237,6 +3237,7 @@ class TestJWTOAuth2Coexistence: jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "mapped-user"}) jwt_handler.get_user_email = MagicMock(return_value="mapped@example.com") jwt_handler.get_user_id = MagicMock(return_value="mapped-user") + jwt_handler.get_end_user_id = MagicMock(return_value=None) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( virtual_key_claim_field="sub", user_email_jwt_field="sub", @@ -3312,6 +3313,7 @@ class TestJWTOAuth2Coexistence: jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "jwt-principal"}) jwt_handler.get_user_email = MagicMock(return_value="principal@example.com") jwt_handler.get_user_id = MagicMock(return_value="jwt-principal") + jwt_handler.get_end_user_id = MagicMock(return_value=None) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( virtual_key_claim_field="sub", user_email_jwt_field="sub", @@ -3388,6 +3390,7 @@ class TestJWTOAuth2Coexistence: jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "mapped-user"}) jwt_handler.get_user_email = MagicMock(return_value="mapped@example.com") jwt_handler.get_user_id = MagicMock(return_value="mapped-user") + jwt_handler.get_end_user_id = MagicMock(return_value=None) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( virtual_key_claim_field="sub", user_email_jwt_field="sub", @@ -11288,6 +11291,7 @@ def _per_issuer_virtual_key_jwt_handler( ) -> MagicMock: jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True + jwt_handler.get_end_user_id.return_value = None jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( virtual_key_claim_field=global_claim_field, unregistered_jwt_client_behavior=global_behavior, diff --git a/tests/unit/proxy/proxy_server/test_routes_login_sso.py b/tests/unit/proxy/proxy_server/test_routes_login_sso.py index 8daea8d0ad2..75e14a89148 100644 --- a/tests/unit/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/unit/proxy/proxy_server/test_routes_login_sso.py @@ -670,6 +670,12 @@ def test_the_master_key_as_a_bearer_token_still_works_while_the_ui_password_is_b def test_a_database_users_correct_password_is_refused_while_blocked(client, monkeypatch, reset_login_throttle): """The block is hard: while it lasts, nothing from that source signs in as that user, right password or not, and the block is not extended by the refused attempts.""" + from types import SimpleNamespace + + from litellm.proxy.auth import login_throttle + + frozen_now = login_throttle.time.time() + monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: frozen_now)) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=64) _db_user(monkeypatch, "user@corp.com")