ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/main' into add_helm_chart

Browse Source
This commit is contained in:
Shaun Maher 2024-01-30 09:44:56 +11:00
commit 9c4e101e32
79 changed files with 11161 additions and 492 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
.circleci/config.yml
View File

@ -147,12 +147,16 @@ jobs:
-e AZURE_API_KEY=$AZURE_API_KEY \
-e AZURE_FRANCE_API_KEY=$AZURE_FRANCE_API_KEY \
-e AZURE_EUROPE_API_KEY=$AZURE_EUROPE_API_KEY \
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
-e AWS_REGION_NAME=$AWS_REGION_NAME \
--name my-app \
-v $(pwd)/proxy_server_config.yaml:/app/config.yaml \
my-app:latest \
--config /app/config.yaml \
--port 4000 \
--num_workers 8
--num_workers 8 \
--debug
- run:
name: Install curl and dockerize
command: |

10
.github/workflows/ghcr_deploy.yml vendored
View File

@ -34,13 +34,6 @@ jobs:
with:
push: true
tags: litellm/litellm:${{ github.event.inputs.tag || 'latest' }}
-
name: Build and push litellm-ui image
uses: docker/build-push-action@v5
with:
push: true
file: ui/Dockerfile
tags: litellm/litellm-ui:${{ github.event.inputs.tag || 'latest' }}
-
name: Build and push litellm-database image
uses: docker/build-push-action@v5
@ -82,6 +75,8 @@ jobs:
push: true
tags: ${{ steps.meta.outputs.tags }}-${{ github.event.inputs.tag || 'latest' }}, ${{ steps.meta.outputs.tags }}-latest # if a tag is provided, use that, otherwise use the release tag, and if neither is available, use 'latest'
labels: ${{ steps.meta.outputs.labels }}
platform: local, linux/amd64,linux/arm64,linux/arm64/v8
build-and-push-image-ui:
runs-on: ubuntu-latest
permissions:
@ -112,6 +107,7 @@ jobs:
push: true
tags: ${{ steps.meta-ui.outputs.tags }}-${{ github.event.inputs.tag || 'latest' }}, ${{ steps.meta-ui.outputs.tags }}-latest
labels: ${{ steps.meta-ui.outputs.labels }}
platform: local, linux/amd64,linux/arm64,linux/arm64/v8
build-and-push-image-database:
runs-on: ubuntu-latest
permissions:

12
docker-compose.example.yml
View File

@ -1,12 +0,0 @@
version: "3.9"
services:
litellm:
image: ghcr.io/berriai/litellm:main
ports:
- "8000:8000" # Map the container port to the host, change the host port if necessary
volumes:
- ./litellm-config.yaml:/app/config.yaml # Mount the local configuration file
# You can change the port or number of workers as per your requirements or pass any new supported CLI augument. Make sure the port passed here matches with the container port defined above in `ports` value
command: [ "--config", "/app/config.yaml", "--port", "8000", "--num_workers", "8" ]
# ...rest of your docker-compose config if any

15
docker-compose.yml Normal file
View File

@ -0,0 +1,15 @@
version: "3.9"
services:
litellm:
image: ghcr.io/berriai/litellm:main-latest
volumes:
- ./proxy_server_config.yaml:/app/proxy_server_config.yaml # mount your litellm config.yaml
ports:
- "4000:4000"
environment:
- AZURE_API_KEY=sk-123
litellm-ui:
image: ghcr.io/berriai/litellm-ui:main-latest

19
docs/my-website/docs/embedding/supported_embedding.md
View File

@ -13,8 +13,8 @@ response = embedding(model='text-embedding-ada-002', input=["good morning from l
- `model`: *string* - ID of the model to use. `model='text-embedding-ada-002'`
- `input`: *array* - Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request, pass an array of strings or array of token arrays. The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002), cannot be an empty string, and any array must be 2048 dimensions or less.
```
- `input`: *string or array* - Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request, pass an array of strings or array of token arrays. The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002), cannot be an empty string, and any array must be 2048 dimensions or less.
```python
input=["good morning from litellm"]
```
@ -22,7 +22,11 @@ input=["good morning from litellm"]
- `user`: *string (optional)* A unique identifier representing your end-user,
- `timeout`: *integer* - The maximum time, in seconds, to wait for the API to respond. Defaults to 600 seconds (10 minutes).
- `dimensions`: *integer (Optional)* The number of dimensions the resulting output embeddings should have. Only supported in OpenAI/Azure text-embedding-3 and later models.
- `encoding_format`: *string (Optional)* The format to return the embeddings in. Can be either `"float"` or `"base64"`. Defaults to `encoding_format="float"`
- `timeout`: *integer (Optional)* - The maximum time, in seconds, to wait for the API to respond. Defaults to 600 seconds (10 minutes).
- `api_base`: *string (optional)* - The api endpoint you want to call the model with
@ -66,11 +70,18 @@ input=["good morning from litellm"]
from litellm import embedding
import os
os.environ['OPENAI_API_KEY'] = ""
response = embedding('text-embedding-ada-002', input=["good morning from litellm"])
response = embedding(
model="text-embedding-3-small",
input=["good morning from litellm", "this is another item"],
metadata={"anything": "good day"},
dimensions=5 # Only supported in text-embedding-3 and later models.
)
```
| Model Name | Function Call | Required OS Variables |
|----------------------|---------------------------------------------|--------------------------------------|
| text-embedding-3-small | `embedding('text-embedding-3-small', input)` | `os.environ['OPENAI_API_KEY']` |
| text-embedding-3-large | `embedding('text-embedding-3-large', input)` | `os.environ['OPENAI_API_KEY']` |
| text-embedding-ada-002 | `embedding('text-embedding-ada-002', input)` | `os.environ['OPENAI_API_KEY']` |
## Azure OpenAI Embedding Models

1
docs/my-website/docs/providers/openai.md
View File

@ -34,6 +34,7 @@ os.environ["OPENAI_API_BASE"] = "openaiai-api-base" # OPTIONAL
| Model Name | Function Call |
|-----------------------|-----------------------------------------------------------------|
| gpt-4-0125-preview | `response = completion(model="gpt-4-0125-preview", messages=messages)` |
| gpt-4-1106-preview | `response = completion(model="gpt-4-1106-preview", messages=messages)` |
| gpt-3.5-turbo-1106 | `response = completion(model="gpt-3.5-turbo-1106", messages=messages)` |
| gpt-3.5-turbo | `response = completion(model="gpt-3.5-turbo", messages=messages)` |

21
docs/my-website/docs/providers/vertex.md
View File

@ -20,6 +20,27 @@ litellm.vertex_location = "us-central1" # proj location
response = litellm.completion(model="gemini-pro", messages=[{"role": "user", "content": "write code for saying hi from LiteLLM"}])
```
## OpenAI Proxy Usage
1. Modify the config.yaml
```yaml
litellm_settings:
vertex_project: "hardy-device-38811" # Your Project ID
vertex_location: "us-central1" # proj location
model_list:
-model_name: team1-gemini-pro
litellm_params:
model: gemini-pro
```
2. Start the proxy
```bash
$ litellm --config /path/to/config.yaml
```
## Set Vertex Project & Vertex Location
All calls using Vertex AI require the following parameters:
* Your Project ID

9
docs/my-website/docs/proxy/alerting.md
View File

@ -1,6 +1,13 @@
# Slack Alerting
Get alerts for failed db read/writes, hanging api calls, failed api calls.
Get alerts for:
- hanging LLM api calls
- failed LLM api calls
- slow LLM api calls
- budget Tracking per key/user:
- When a User/Key crosses their Budget
- When a User/Key is 15% away from crossing their Budget
- failed db read/writes
## Quick Start

37
docs/my-website/docs/proxy/custom_pricing.md
View File

@ -54,7 +54,7 @@ model_list:
- model_name: sagemaker-embedding-model
litellm_params:
model: sagemaker/berri-benchmarking-gpt-j-6b-fp16
input_cost_per_second: 0.000420
input_cost_per_second: 0.000420
```
**Step 2: Start proxy**
@ -67,25 +67,28 @@ litellm /path/to/config.yaml
<Image img={require('../../img/spend_logs_table.png')} />
## Cost Per Token
## Cost Per Token (e.g. Azure)
```python
# !pip install boto3
from litellm import completion, completion_cost
os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""
## set ENV variables
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
def test_completion_sagemaker():
def test_completion_azure_model():
try:
print("testing sagemaker")
print("testing azure custom pricing")
# azure call
response = completion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
input_cost_per_token=0.005,
output_cost_per_token=1,
model = "azure/<your_deployment_name>",
messages = [{ "content": "Hello, how are you?","role": "user"}]
input_cost_per_token=0.005,
output_cost_per_token=1,
)
# Add any assertions here to check the response
print(response)
@ -94,15 +97,19 @@ def test_completion_sagemaker():
except Exception as e:
raise Exception(f"Error occurred: {e}")
test_completion_azure_model()
```
### Usage with OpenAI Proxy Server
```yaml
model_list:
- model_name: sagemaker-completion-model
- model_name: azure-model
litellm_params:
model: sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4
input_cost_per_token: 0.000420 # 👈 key change
output_cost_per_token: 0.000420 # 👈 key change
model: azure/<your_deployment_name>
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_API_BASE
api_version: os.envrion/AZURE_API_VERSION
input_cost_per_token: 0.000421 # 👈 ONLY to track cost per token
output_cost_per_token: 0.000520 # 👈 ONLY to track cost per token
```

63
docs/my-website/docs/proxy/ui.md
View File

@ -2,35 +2,23 @@ import Image from '@theme/IdealImage';
# [BETA] Admin UI
- Track Spend Per API Key, User
- Allow your users to create their own keys through a UI
:::info
This is in beta, so things may change. If you have feedback, [let us know](https://discord.com/invite/wuPM9dRgDw)
:::
Allow your users to create, view their own keys through a UI
<Image img={require('../../img/admin_ui_2.png')} />
## Quick Start
Requirements:
## 1. Changes to your config.yaml
- Need to a SMTP server connection to send emails (e.g. [Resend](https://resend.com/docs/send-with-smtp))
[**See code**](https://github.com/BerriAI/litellm/blob/61cd800b9ffbb02c286481d2056b65c7fb5447bf/litellm/proxy/proxy_server.py#L1782)
### Step 1. Save SMTP server credentials
```env
export SMTP_HOST="my-smtp-host"
export SMTP_USERNAME="my-smtp-password"
export SMTP_PASSWORD="my-smtp-password"
export SMTP_SENDER_EMAIL="krrish@berri.ai"
```
### Step 2. Enable user auth
In your config.yaml,
Set `allow_user_auth: true` on your config
```yaml
general_settings:
@ -38,13 +26,36 @@ general_settings:
allow_user_auth: true
```
This will enable:
* Users to create keys via `/key/generate` (by default, only admin can create keys)
* The `/user/auth` endpoint to send user's emails with their login credentials (key + user id)
## 2. Setup Google SSO - Use this to Authenticate Team Members to the UI
- Create an Oauth 2.0 Client
<Image img={require('../../img/google_oauth2.png')} />
### Step 3. Connect to UI
- Navigate to Google `Credenentials`
- Create a new Oauth client ID
- Set the `GOOGLE_CLIENT_ID` and `GOOGLE_CLIENT_SECRET` in your Proxy .env
- Set Redirect URL on your Oauth 2.0 Client
- Click on your Oauth 2.0 client on https://console.cloud.google.com/
- Set a redirect url = `<your proxy base url>/google-callback`
```
https://litellm-production-7002.up.railway.app/google-callback
```
<Image img={require('../../img/google_redirect.png')} />
## 3. Required env variables on your Proxy
You can use our hosted UI (https://dashboard.litellm.ai/) or [self-host your own](https://github.com/BerriAI/litellm/tree/main/ui).
```shell
PROXY_BASE_URL="<your deployed proxy endpoint>" example PROXY_BASE_URL=https://litellm-production-7002.up.railway.app/
# for Google SSO Login
GOOGLE_CLIENT_ID=
GOOGLE_CLIENT_SECRET=
```
## 4. Use UI
👉 Get Started here: https://litellm-dashboard.vercel.app/
<!-- You can use our hosted UI (https://dashboard.litellm.ai/) or [self-host your own](https://github.com/BerriAI/litellm/tree/main/ui).
If you self-host, you need to save the UI url in your proxy environment as `LITELLM_HOSTED_UI`.
@ -71,5 +82,5 @@ Connect your proxy to your UI, by entering:
### Spend Per User
<Image img={require('../../img/spend_per_user.png')} />
<Image img={require('../../img/spend_per_user.png')} /> -->

92
docs/my-website/docs/proxy/users.md
View File

@ -1,7 +1,7 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# 💰 Budgets, Rate Limits per user
# 💰 Budgets, Rate Limits
Requirements:
@ -10,22 +10,71 @@ Requirements:
## Set Budgets
You can set budgets at 3 levels:
- For the proxy
- For a user
- For a key
Set `max_budget` in (USD $) param in the `/user/new` or `/key/generate` request. By default the `max_budget` is set to `null` and is not checked for keys
<Tabs>
<TabItem value="per-user" label="Per User">
<TabItem value="proxy" label="For Proxy">
LiteLLM exposes a `/user/new` endpoint to create budgets for users, that persist across multiple keys.
Apply a budget across all calls on the proxy
**Step 1. Modify config.yaml**
```yaml
general_settings:
master_key: sk-1234
litellm_settings:
# other litellm settings
max_budget: 0 # (float) sets max budget as $0 USD
budget_duration: 30d # (str) frequency of reset - You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
```
**Step 2. Start proxy**
```bash
litellm /path/to/config.yaml
```
**Step 3. Send test call**
```bash
curl --location 'http://0.0.0.0:8000/chat/completions' \
--header 'Autherization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
<TabItem value="per-user" label="For User">
Apply a budget across multiple keys.
LiteLLM exposes a `/user/new` endpoint to create budgets for this.
You can:
- Add budgets to users [**Jump**](#add-budgets-to-users)
- Add budget durations, to reset spend [**Jump**](#add-budget-duration-to-users)
By default the `max_budget` is set to `null` and is not checked for keys
### **Add budgets to users**
```shell
curl --location 'http://localhost:8000/user/new' \
--header 'Authorization: Bearer <your-master-key>' \
--header 'Content-Type: application/json' \
--data-raw '{"models": ["azure-models"], "max_budget": 0, "user_id": "krrish3@berri.ai"}'
```
The request is a normal `/key/generate` request body + a `max_budget` field.
[**See Swagger**](https://litellm-api.up.railway.app/#/user%20management/new_user_user_new_post)
@ -40,9 +89,38 @@ The request is a normal `/key/generate` request body + a `max_budget` field.
}
```
### **Add budget duration to users**
`budget_duration`: Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
```
curl 'http://0.0.0.0:8000/user/new' \
--header 'Authorization: Bearer <your-master-key>' \
--header 'Content-Type: application/json' \
--data-raw '{
"team_id": "core-infra", # [OPTIONAL]
"max_budget": 10,
"budget_duration": 10s,
}'
```
### Create new keys for existing user
Now you can just call `/key/generate` with that user_id (i.e. krrish3@berri.ai) and:
- **Budget Check**: krrish3@berri.ai's budget (i.e. $10) will be checked for this key
- **Spend Tracking**: spend for this key will update krrish3@berri.ai's spend as well
```bash
curl --location 'http://0.0.0.0:8000/key/generate' \
--header 'Authorization: Bearer <your-master-key>' \
--header 'Content-Type: application/json' \
--data '{"models": ["azure-models"], "user_id": "krrish3@berri.ai"}'
```
</TabItem>
<TabItem value="per-key" label="Per Key">
<TabItem value="per-key" label="For Key">
Apply a budget on a key.
You can:
- Add budgets to keys [**Jump**](#add-budgets-to-keys)
@ -53,6 +131,8 @@ You can:
- After the key crosses it's `max_budget`, requests fail
- If duration set, spend is reset at the end of the duration
By default the `max_budget` is set to `null` and is not checked for keys
### **Add budgets to keys**
```bash

31
docs/my-website/docs/proxy/virtual_keys.md
View File

@ -81,15 +81,17 @@ curl 'http://0.0.0.0:8000/key/generate' \
Request Params:
- `models`: *list or null (optional)* - Specify the models a token has access too. If null, then token has access to all models on server.
- `duration`: *Optional[str]* - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
- `key_alias`: *Optional[str]* - User defined key alias
- `team_id`: *Optional[str]* - The team id of the user
- `models`: *Optional[list]* - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
- `aliases`: *Optional[dict]* - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
- `config`: *Optional[dict]* - any key-specific configs, overrides config in config.yaml
- `spend`: *Optional[int]* - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend
- `max_budget`: *Optional[float]* - Specify max budget for a given key.
- `max_parallel_requests`: *Optional[int]* - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
- `metadata`: *Optional[dict]* - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
- `duration`: *str or null (optional)* Specify the length of time the token is valid for. If null, default is set to 1 hour. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
- `metadata`: *dict or null (optional)* Pass metadata for the created token. If null defaults to {}
- `team_id`: *str or null (optional)* Specify team_id for the associated key
- `max_budget`: *float or null (optional)* Specify max budget (in Dollars $) for a given key. If no value is set, the key has no budget
### Response
@ -97,20 +99,11 @@ Request Params:
{
"key": "sk-kdEXbIqZRwEeEiHwdg7sFA", # Bearer token
"expires": "2023-11-19T01:38:25.838000+00:00" # datetime object
"key_name": "sk-...7sFA" # abbreviated key string, ONLY stored in db if `allow_user_auth: true` set - [see](./ui.md)
...
}
```
### Keys that don't expire
Just set duration to None.
```bash
curl --location 'http://0.0.0.0:8000/key/generate' \
--header 'Authorization: Bearer <your-master-key>' \
--header 'Content-Type: application/json' \
--data '{"models": ["azure-models"], "aliases": {"mistral-7b": "gpt-3.5-turbo"}, "duration": null}'
```
### Upgrade/Downgrade Models
If a user is expected to use a given model (i.e. gpt3-5), and you want to:

43
docs/my-website/docs/routing.md
View File

@ -605,6 +605,49 @@ response = router.completion(model="gpt-3.5-turbo", messages=messages)
print(f"response: {response}")
```
## Custom Callbacks - Track API Key, API Endpoint, Model Used
If you need to track the api_key, api endpoint, model, custom_llm_provider used for each completion call, you can setup a [custom callback](https://docs.litellm.ai/docs/observability/custom_callback)
### Usage
```python
import litellm
from litellm.integrations.custom_logger import CustomLogger
class MyCustomHandler(CustomLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Success")
print("kwargs=", kwargs)
litellm_params= kwargs.get("litellm_params")
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
custom_llm_provider= litellm_params.get("custom_llm_provider")
response_cost = kwargs.get("response_cost")
# print the values
print("api_key=", api_key)
print("api_base=", api_base)
print("custom_llm_provider=", custom_llm_provider)
print("response_cost=", response_cost)
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Failure")
print("kwargs=")
customHandler = MyCustomHandler()
litellm.callbacks = [customHandler]
# Init Router
router = Router(model_list=model_list, routing_strategy="simple-shuffle")
# router completion call
response = router.completion(
model="gpt-3.5-turbo",
messages=[{ "role": "user", "content": "Hi who are you"}]
)
```
## Deploy Router

BIN
docs/my-website/img/admin_ui_2.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 159 KiB

BIN
docs/my-website/img/google_oauth2.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 351 KiB

BIN
docs/my-website/img/google_redirect.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 297 KiB

11
litellm/budget_manager.py
View File

@ -1,3 +1,12 @@
# +-----------------------------------------------+
# | |
# | NOT PROXY BUDGET MANAGER |
# | proxy budget manager is in proxy_server.py |
# | |
# +-----------------------------------------------+
#
# Thank you users! We ❤️ you! - Krrish & Ishaan
import os, json, time
import litellm
from litellm.utils import ModelResponse
@ -16,7 +25,7 @@ class BudgetManager:
self.client_type = client_type
self.project_name = project_name
self.api_base = api_base or "https://api.litellm.ai"
self.headers = headers or {'Content-Type': 'application/json'}
self.headers = headers or {"Content-Type": "application/json"}
## load the data or init the initial dictionaries
self.load_data()

16
litellm/integrations/s3.py
View File

@ -8,7 +8,7 @@ dotenv.load_dotenv() # Loading env variables using dotenv
import traceback
import datetime, subprocess, sys
import litellm, uuid
from litellm._logging import print_verbose
from litellm._logging import print_verbose, verbose_logger
class S3Logger:
@ -31,7 +31,9 @@ class S3Logger:
import boto3
try:
print_verbose("in init s3 logger")
verbose_logger.debug(
f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}"
)
if litellm.s3_callback_params is not None:
# read in .env variables - example os.environ/AWS_BUCKET_NAME
@ -42,7 +44,7 @@ class S3Logger:
s3_bucket_name = litellm.s3_callback_params.get("s3_bucket_name")
s3_region_name = litellm.s3_callback_params.get("s3_region_name")
s3_api_version = litellm.s3_callback_params.get("s3_api_version")
s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl")
s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl", True)
s3_verify = litellm.s3_callback_params.get("s3_verify")
s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url")
s3_aws_access_key_id = litellm.s3_callback_params.get(
@ -59,6 +61,7 @@ class S3Logger:
self.bucket_name = s3_bucket_name
self.s3_path = s3_path
verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}")
# Create an S3 client with custom endpoint URL
self.s3_client = boto3.client(
"s3",
@ -84,7 +87,9 @@ class S3Logger:
def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
try:
print_verbose(f"s3 Logging - Enters logging function for model {kwargs}")
verbose_logger.debug(
f"s3 Logging - Enters logging function for model {kwargs}"
)
# construct payload to send to s3
# follows the same params as langfuse.py
@ -129,6 +134,7 @@ class S3Logger:
+ "-time="
+ str(start_time)
) # we need the s3 key to include the time, so we log cache hits too
s3_object_key += ".json"
import json
@ -151,5 +157,5 @@ class S3Logger:
return response
except Exception as e:
traceback.print_exc()
print_verbose(f"s3 Layer Error - {str(e)}\n{traceback.format_exc()}")
verbose_logger.debug(f"s3 Layer Error - {str(e)}\n{traceback.format_exc()}")
pass

15
litellm/llms/bedrock.py
View File

@ -659,9 +659,16 @@ def completion(
)
## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
prompt_tokens = len(encoding.encode(prompt))
completion_tokens = len(
encoding.encode(model_response["choices"][0]["message"].get("content", ""))
prompt_tokens = response_metadata.get(
"x-amzn-bedrock-input-token-count", len(encoding.encode(prompt))
)
completion_tokens = response_metadata.get(
"x-amzn-bedrock-output-token-count",
len(
encoding.encode(
model_response["choices"][0]["message"].get("content", "")
)
),
)
model_response["created"] = int(time.time())
@ -672,6 +679,8 @@ def completion(
total_tokens=prompt_tokens + completion_tokens,
)
model_response.usage = usage
model_response._hidden_params["region_name"] = client.meta.region_name
print_verbose(f"model_response._hidden_params: {model_response._hidden_params}")
return model_response
except BedrockError as e:
exception_mapping_worked = True

12
litellm/llms/ollama_chat.py
View File

@ -220,8 +220,10 @@ def get_ollama_response(
model_response["choices"][0]["message"] = response_json["message"]
model_response["created"] = int(time.time())
model_response["model"] = "ollama/" + model
prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore
completion_tokens = response_json["eval_count"]
prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=messages)) # type: ignore
completion_tokens = response_json.get(
"eval_count", litellm.token_counter(text=response_json["message"])
)
model_response["usage"] = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
@ -320,8 +322,10 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
model_response["choices"][0]["message"] = response_json["message"]
model_response["created"] = int(time.time())
model_response["model"] = "ollama/" + data["model"]
prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore
completion_tokens = response_json["eval_count"]
prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=data["messages"])) # type: ignore
completion_tokens = response_json.get(
"eval_count", litellm.token_counter(text=response_json["message"])
)
model_response["usage"] = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,

14
litellm/llms/openai.py
View File

@ -718,8 +718,22 @@ class OpenAIChatCompletion(BaseLLM):
return convert_to_model_response_object(response_object=response, model_response_object=model_response, response_type="image_generation") # type: ignore
except OpenAIError as e:
exception_mapping_worked = True
## LOGGING
logging_obj.post_call(
input=prompt,
api_key=api_key,
additional_args={"complete_input_dict": data},
original_response=str(e),
)
raise e
except Exception as e:
## LOGGING
logging_obj.post_call(
input=prompt,
api_key=api_key,
additional_args={"complete_input_dict": data},
original_response=str(e),
)
if hasattr(e, "status_code"):
raise OpenAIError(status_code=e.status_code, message=str(e))
else:

10
litellm/llms/prompt_templates/factory.py
View File

@ -99,12 +99,16 @@ def ollama_pt(
def mistral_instruct_pt(messages):
# Following the Mistral example's https://huggingface.co/docs/transformers/main/chat_templating
prompt = custom_prompt(
initial_prompt_value="<s>",
role_dict={
"system": {"pre_message": "[INST]", "post_message": "[/INST]"},
"user": {"pre_message": "[INST]", "post_message": "[/INST]"},
"assistant": {"pre_message": "[INST]", "post_message": "[/INST]"},
"system": {
"pre_message": "[INST] <<SYS>>\n",
"post_message": "<</SYS>> [/INST]\n",
},
"user": {"pre_message": "[INST] ", "post_message": " [/INST]\n"},
"assistant": {"pre_message": " ", "post_message": " "},
},
final_prompt_value="</s>",
messages=messages,

35
litellm/llms/sagemaker.py
View File

@ -34,22 +34,35 @@ class TokenIterator:
self.byte_iterator = iter(stream)
self.buffer = io.BytesIO()
self.read_pos = 0
self.end_of_data = False
def __iter__(self):
return self
def __next__(self):
while True:
self.buffer.seek(self.read_pos)
line = self.buffer.readline()
if line and line[-1] == ord("\n"):
self.read_pos += len(line) + 1
full_line = line[:-1].decode("utf-8")
line_data = json.loads(full_line.lstrip("data:").rstrip("/n"))
return line_data["token"]["text"]
chunk = next(self.byte_iterator)
self.buffer.seek(0, io.SEEK_END)
self.buffer.write(chunk["PayloadPart"]["Bytes"])
try:
while True:
self.buffer.seek(self.read_pos)
line = self.buffer.readline()
if line and line[-1] == ord("\n"):
response_obj = {"text": "", "is_finished": False}
self.read_pos += len(line) + 1
full_line = line[:-1].decode("utf-8")
line_data = json.loads(full_line.lstrip("data:").rstrip("/n"))
if line_data.get("generated_text", None) is not None:
self.end_of_data = True
response_obj["is_finished"] = True
response_obj["text"] = line_data["token"]["text"]
return response_obj
chunk = next(self.byte_iterator)
self.buffer.seek(0, io.SEEK_END)
self.buffer.write(chunk["PayloadPart"]["Bytes"])
except StopIteration as e:
if self.end_of_data == True:
raise e # Re-raise StopIteration
else:
self.end_of_data = True
return "data: [DONE]"
class SagemakerConfig:

62
litellm/main.py
View File

@ -10,12 +10,11 @@
import os, openai, sys, json, inspect, uuid, datetime, threading
from typing import Any, Literal, Union
from functools import partial
import dotenv, traceback, random, asyncio, time, contextvars
from copy import deepcopy
import httpx
import litellm
from ._logging import verbose_logger
from litellm import ( # type: ignore
client,
exception_type,
@ -274,14 +273,10 @@ async def acompletion(
else:
# Call the synchronous function using run_in_executor
response = await loop.run_in_executor(None, func_with_context) # type: ignore
# if kwargs.get("stream", False): # return an async generator
# return _async_streaming(
# response=response,
# model=model,
# custom_llm_provider=custom_llm_provider,
# args=args,
# )
# else:
if isinstance(response, CustomStreamWrapper):
response.set_logging_event_loop(
loop=loop
) # sets the logging event loop if the user does sync streaming (e.g. on proxy for sagemaker calls)
return response
except Exception as e:
custom_llm_provider = custom_llm_provider or "openai"
@ -590,6 +585,10 @@ def completion(
)
if model_response is not None and hasattr(model_response, "_hidden_params"):
model_response._hidden_params["custom_llm_provider"] = custom_llm_provider
model_response._hidden_params["region_name"] = kwargs.get(
"aws_region_name", None
) # support region-based pricing for bedrock
### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
if input_cost_per_token is not None and output_cost_per_token is not None:
litellm.register_model(
@ -1421,9 +1420,15 @@ def completion(
return response
response = model_response
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = litellm.vertex_project or get_secret("VERTEXAI_PROJECT")
vertex_ai_location = litellm.vertex_location or get_secret(
"VERTEXAI_LOCATION"
vertex_ai_project = (
optional_params.pop("vertex_ai_project", None)
or litellm.vertex_project
or get_secret("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.pop("vertex_ai_location", None)
or litellm.vertex_location
or get_secret("VERTEXAI_LOCATION")
)
model_response = vertex_ai.completion(
@ -1514,11 +1519,6 @@ def completion(
if (
"stream" in optional_params and optional_params["stream"] == True
): ## [BETA]
# sagemaker does not support streaming as of now so we're faking streaming:
# https://discuss.huggingface.co/t/streaming-output-text-when-deploying-on-sagemaker/39611
# "SageMaker is currently not supporting streaming responses."
# fake streaming for sagemaker
print_verbose(f"ENTERS SAGEMAKER CUSTOMSTREAMWRAPPER")
from .llms.sagemaker import TokenIterator
@ -1529,6 +1529,12 @@ def completion(
custom_llm_provider="sagemaker",
logging_obj=logging,
)
## LOGGING
logging.post_call(
input=messages,
api_key=None,
original_response=response,
)
return response
## RESPONSE OBJECT
@ -2221,6 +2227,7 @@ def embedding(
model,
input=[],
# Optional params
dimensions: Optional[int] = None,
timeout=600, # default to 10 minutes
# set api_base, api_version, api_key
api_base: Optional[str] = None,
@ -2241,6 +2248,7 @@ def embedding(
Parameters:
- model: The embedding model to use.
- input: The input for which embeddings are to be generated.
- dimensions: The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.
- timeout: The timeout value for the API call, default 10 mins
- litellm_call_id: The call ID for litellm logging.
- litellm_logging_obj: The litellm logging object.
@ -2274,6 +2282,7 @@ def embedding(
output_cost_per_second = kwargs.get("output_cost_per_second", None)
openai_params = [
"user",
"dimensions",
"request_timeout",
"api_base",
"api_version",
@ -2342,7 +2351,9 @@ def embedding(
api_key=api_key,
)
optional_params = get_optional_params_embeddings(
model=model,
user=user,
dimensions=dimensions,
encoding_format=encoding_format,
custom_llm_provider=custom_llm_provider,
**non_default_params,
@ -3064,7 +3075,7 @@ def image_generation(
custom_llm_provider=custom_llm_provider,
**non_default_params,
)
logging = litellm_logging_obj
logging: Logging = litellm_logging_obj
logging.update_environment_variables(
model=model,
user=user,
@ -3164,6 +3175,9 @@ async def ahealth_check(
if model is None:
raise Exception("model not set")
if model in litellm.model_cost and mode is None:
mode = litellm.model_cost[model]["mode"]
model, custom_llm_provider, _, _ = get_llm_provider(model=model)
mode = mode or "chat" # default to chat completion calls
@ -3342,6 +3356,16 @@ def stream_chunk_builder(
chunks: list, messages: Optional[list] = None, start_time=None, end_time=None
):
model_response = litellm.ModelResponse()
### SORT CHUNKS BASED ON CREATED ORDER ##
print_verbose("Goes into checking if chunk has hiddden created at param")
if chunks[0]._hidden_params.get("created_at", None):
print_verbose("Chunks have a created at hidden param")
# Sort chunks based on created_at in ascending order
chunks = sorted(
chunks, key=lambda x: x._hidden_params.get("created_at", float("inf"))
)
print_verbose("Chunks sorted")
# set hidden params from chunk to model_response
if model_response is not None and hasattr(model_response, "_hidden_params"):
model_response._hidden_params = chunks[0].get("_hidden_params", {})

74
litellm/proxy/_types.py
View File

@ -122,11 +122,12 @@ class ModelParams(LiteLLMBase):
return values
class GenerateKeyRequest(LiteLLMBase):
duration: Optional[str] = "1h"
class GenerateRequestBase(LiteLLMBase):
"""
Overlapping schema between key and user generate/update requests
"""
models: Optional[list] = []
aliases: Optional[dict] = {}
config: Optional[dict] = {}
spend: Optional[float] = 0
max_budget: Optional[float] = None
user_id: Optional[str] = None
@ -138,21 +139,42 @@ class GenerateKeyRequest(LiteLLMBase):
budget_duration: Optional[str] = None
class UpdateKeyRequest(LiteLLMBase):
class GenerateKeyRequest(GenerateRequestBase):
key_alias: Optional[str] = None
duration: Optional[str] = None
aliases: Optional[dict] = {}
config: Optional[dict] = {}
class GenerateKeyResponse(GenerateKeyRequest):
key: str
key_name: Optional[str] = None
expires: Optional[datetime]
user_id: str
@root_validator(pre=True)
def set_model_info(cls, values):
if values.get("token") is not None:
values.update({"key": values.get("token")})
dict_fields = ["metadata", "aliases", "config"]
for field in dict_fields:
value = values.get(field)
if value is not None and isinstance(value, str):
try:
values[field] = json.loads(value)
except json.JSONDecodeError:
raise ValueError(f"Field {field} should be a valid dictionary")
return values
class UpdateKeyRequest(GenerateKeyRequest):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
key: str
duration: Optional[str] = None
models: Optional[list] = None
aliases: Optional[dict] = None
config: Optional[dict] = None
spend: Optional[float] = None
max_budget: Optional[float] = None
user_id: Optional[str] = None
max_parallel_requests: Optional[int] = None
metadata: Optional[dict] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
class UserAPIKeyAuth(LiteLLMBase): # the expected response object for user api key auth
@ -174,12 +196,6 @@ class UserAPIKeyAuth(LiteLLMBase): # the expected response object for user api
rpm_limit: Optional[int] = None
class GenerateKeyResponse(LiteLLMBase):
key: str
expires: Optional[datetime]
user_id: str
class DeleteKeyRequest(LiteLLMBase):
keys: List
@ -192,6 +208,14 @@ class NewUserResponse(GenerateKeyResponse):
max_budget: Optional[float] = None
class UpdateUserRequest(GenerateRequestBase):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
user_id: str
spend: Optional[float] = None
metadata: Optional[dict] = None
class KeyManagementSystem(enum.Enum):
GOOGLE_KMS = "google_kms"
AZURE_KEY_VAULT = "azure_key_vault"
@ -297,6 +321,8 @@ class ConfigYAML(LiteLLMBase):
class LiteLLM_VerificationToken(LiteLLMBase):
token: str
key_name: Optional[str] = None
key_alias: Optional[str] = None
spend: float = 0.0
max_budget: Optional[float] = None
expires: Union[str, None]
@ -329,7 +355,7 @@ class LiteLLM_UserTable(LiteLLMBase):
if values.get("spend") is None:
values.update({"spend": 0.0})
if values.get("models") is None:
values.update({"models", []})
values.update({"models": []})
return values
@ -339,12 +365,12 @@ class LiteLLM_SpendLogs(LiteLLMBase):
model: Optional[str] = ""
call_type: str
spend: Optional[float] = 0.0
total_tokens: Optional[int] = 0
prompt_tokens: Optional[int] = 0
completion_tokens: Optional[int] = 0
startTime: Union[str, datetime, None]
endTime: Union[str, datetime, None]
user: Optional[str] = ""
modelParameters: Optional[Json] = {}
messages: Optional[Json] = []
response: Optional[Json] = {}
usage: Optional[Json] = {}
metadata: Optional[Json] = {}
cache_hit: Optional[str] = "False"
cache_key: Optional[str] = None

15
litellm/proxy/db/dynamo_db.py
View File

@ -5,6 +5,7 @@ from litellm.proxy._types import (
LiteLLM_Config,
LiteLLM_UserTable,
)
from litellm.proxy.utils import hash_token
from litellm import get_secret
from typing import Any, List, Literal, Optional, Union
import json
@ -186,7 +187,10 @@ class DynamoDBWrapper(CustomDB):
elif table_name == "spend":
table = client.table(self.database_arguments.spend_table_name)
value = value.copy()
for k, v in value.items():
if k == "token" and value[k].startswith("sk-"):
value[k] = hash_token(token=v)
if isinstance(v, datetime):
value[k] = v.isoformat()
@ -229,6 +233,10 @@ class DynamoDBWrapper(CustomDB):
table = client.table(self.database_arguments.config_table_name)
key_name = "param_name"
if key_name == "token" and key.startswith("sk-"):
# ensure it's hashed
key = hash_token(token=key)
response = await table.get_item({key_name: key})
new_response: Any = None
@ -243,6 +251,10 @@ class DynamoDBWrapper(CustomDB):
and isinstance(v, str)
):
new_response[k] = json.loads(v)
elif (k == "tpm_limit" or k == "rpm_limit") and isinstance(
v, float
):
new_response[k] = int(v)
else:
new_response[k] = v
new_response = LiteLLM_VerificationToken(**new_response)
@ -300,10 +312,13 @@ class DynamoDBWrapper(CustomDB):
# Initialize an empty UpdateExpression
actions: List = []
value = value.copy()
for k, v in value.items():
# Convert datetime object to ISO8601 string
if isinstance(v, datetime):
v = v.isoformat()
if k == "token" and value[k].startswith("sk-"):
value[k] = hash_token(token=v)
# Accumulate updates
actions.append((F(k), Value(value=v)))

26
litellm/proxy/proxy_cli.py
View File

@ -190,17 +190,27 @@ def run_server(
global feature_telemetry
args = locals()
if local:
from proxy_server import app, save_worker_config, usage_telemetry
from proxy_server import app, save_worker_config, usage_telemetry, ProxyConfig
else:
try:
from .proxy_server import app, save_worker_config, usage_telemetry
from .proxy_server import (
app,
save_worker_config,
usage_telemetry,
ProxyConfig,
)
except ImportError as e:
if "litellm[proxy]" in str(e):
# user is missing a proxy dependency, ask them to pip install litellm[proxy]
raise e
else:
# this is just a local/relative import error, user git cloned litellm
from proxy_server import app, save_worker_config, usage_telemetry
from proxy_server import (
app,
save_worker_config,
usage_telemetry,
ProxyConfig,
)
feature_telemetry = usage_telemetry
if version == True:
pkg_version = importlib.metadata.version("litellm")
@ -373,16 +383,16 @@ def run_server(
read from there and save it to os.env['DATABASE_URL']
"""
try:
import yaml
import yaml, asyncio
except:
raise ImportError(
"yaml needs to be imported. Run - `pip install 'litellm[proxy]'`"
)
if os.path.exists(config):
with open(config, "r") as config_file:
config = yaml.safe_load(config_file)
general_settings = config.get("general_settings", {})
proxy_config = ProxyConfig()
_, _, general_settings = asyncio.run(
proxy_config.load_config(router=None, config_file_path=config)
)
database_url = general_settings.get("database_url", None)
if database_url and database_url.startswith("os.environ/"):
original_dir = os.getcwd()

8
litellm/proxy/proxy_config.yaml
View File

@ -11,6 +11,12 @@ model_list:
output_cost_per_token: 0.00003
max_tokens: 4096
base_model: gpt-3.5-turbo
- model_name: gpt-4
litellm_params:
model: azure/chatgpt-v-2
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
- model_name: gpt-vision
litellm_params:
model: azure/gpt-4-vision
@ -61,6 +67,8 @@ model_list:
litellm_settings:
fallbacks: [{"openai-gpt-3.5": ["azure-gpt-3.5"]}]
success_callback: ['langfuse']
max_budget: 10 # global budget for proxy
budget_duration: 30d # global budget duration, will reset after 30d
# cache: True
# setting callback class
# callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance]

464
litellm/proxy/proxy_server.py
View File

@ -75,6 +75,7 @@ from litellm.proxy.utils import (
send_email,
get_logging_payload,
reset_budget,
hash_token,
)
from litellm.proxy.secret_managers.google_kms import load_google_kms
import pydantic
@ -103,6 +104,7 @@ from fastapi.responses import (
ORJSONResponse,
JSONResponse,
)
from fastapi.responses import RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security.api_key import APIKeyHeader
import json
@ -197,6 +199,7 @@ use_queue = False
health_check_interval = None
health_check_results = {}
queue: List = []
litellm_proxy_budget_name = "litellm-proxy-budget"
### INITIALIZE GLOBAL LOGGING OBJECT ###
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
### REDIS QUEUE ###
@ -242,6 +245,8 @@ async def user_api_key_auth(
response = await user_custom_auth(request=request, api_key=api_key)
return UserAPIKeyAuth.model_validate(response)
### LITELLM-DEFINED AUTH FUNCTION ###
if isinstance(api_key, str):
assert api_key.startswith("sk-") # prevent token hashes from being used
if master_key is None:
if isinstance(api_key, str):
return UserAPIKeyAuth(api_key=api_key)
@ -287,8 +292,9 @@ async def user_api_key_auth(
raise Exception("No connected db.")
## check for cache hit (In-Memory Cache)
if api_key.startswith("sk-"):
api_key = hash_token(token=api_key)
valid_token = user_api_key_cache.get_cache(key=api_key)
verbose_proxy_logger.debug(f"valid_token from cache: {valid_token}")
if valid_token is None:
## check db
verbose_proxy_logger.debug(f"api key: {api_key}")
@ -370,30 +376,81 @@ async def user_api_key_auth(
)
# Check 2. If user_id for this token is in budget
## Check 2.5 If global proxy is in budget
if valid_token.user_id is not None:
if prisma_client is not None:
user_id_information = await prisma_client.get_data(
user_id=valid_token.user_id, table_name="user"
user_id_list=[valid_token.user_id, litellm_proxy_budget_name],
table_name="user",
query_type="find_all",
)
if custom_db_client is not None:
user_id_information = await custom_db_client.get_data(
key=valid_token.user_id, table_name="user"
)
verbose_proxy_logger.debug(
f"user_id_information: {user_id_information}"
)
# Token exists, not expired now check if its in budget for the user
if valid_token.spend is not None and valid_token.user_id is not None:
user_max_budget = getattr(user_id_information, "max_budget", None)
user_current_spend = getattr(user_id_information, "spend", None)
if user_id_information is not None:
if isinstance(user_id_information, list):
## Check if user in budget
for _user in user_id_information:
if _user is None:
continue
assert isinstance(_user, dict)
# Token exists, not expired now check if its in budget for the user
user_max_budget = _user.get("max_budget", None)
user_current_spend = _user.get("spend", None)
if user_max_budget is not None and user_current_spend is not None:
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
verbose_proxy_logger.debug(
f"user_max_budget: {user_max_budget}; user_current_spend: {user_current_spend}"
)
if (
user_max_budget is not None
and user_current_spend is not None
):
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=user_max_budget,
user_current_spend=user_current_spend,
type="user_and_proxy_budget",
user_info=_user,
)
)
_user_id = _user.get("user_id", None)
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {_user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
)
else:
# Token exists, not expired now check if its in budget for the user
user_max_budget = getattr(
user_id_information, "max_budget", None
)
user_current_spend = getattr(user_id_information, "spend", None)
if (
user_max_budget is not None
and user_current_spend is not None
):
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=user_max_budget,
user_current_spend=user_current_spend,
type="user_budget",
user_info=user_id_information,
)
)
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
)
# Check 3. If token is expired
if valid_token.expires is not None:
current_time = datetime.now(timezone.utc)
@ -415,16 +472,25 @@ async def user_api_key_auth(
# Check 4. Token Spend is under budget
if valid_token.spend is not None and valid_token.max_budget is not None:
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=valid_token.max_budget,
user_current_spend=valid_token.spend,
type="token_budget",
user_info=valid_token,
)
)
if valid_token.spend > valid_token.max_budget:
raise Exception(
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Token: {valid_token.max_budget}"
)
# Token passed all checks
# Add token to cache
user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
api_key = valid_token.token
# Add hashed token to cache
user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
valid_token_dict = _get_pydantic_json_dict(valid_token)
valid_token_dict.pop("token", None)
"""
@ -450,16 +516,23 @@ async def user_api_key_auth(
)
if (
(route.startswith("/key/") or route.startswith("/user/"))
or route.startswith("/model/")
and not is_master_key_valid
and general_settings.get("allow_user_auth", False) != True
(
route.startswith("/key/")
or route.startswith("/user/")
or route.startswith("/model/")
)
and (not is_master_key_valid)
and (not general_settings.get("allow_user_auth", False))
):
assert not general_settings.get("allow_user_auth", False)
if route == "/key/info":
# check if user can access this route
query_params = request.query_params
key = query_params.get("key")
if prisma_client.hash_token(token=key) != api_key:
if (
key is not None
and prisma_client.hash_token(token=key) != api_key
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="user not allowed to access this key's info",
@ -478,7 +551,7 @@ async def user_api_key_auth(
pass
else:
raise Exception(
f"If master key is set, only master key can be used to generate, delete, update or get info for new keys/users"
f"only master key can be used to generate, delete, update or get info for new keys/users."
)
return UserAPIKeyAuth(api_key=api_key, **valid_token_dict)
@ -593,6 +666,12 @@ async def track_cost_callback(
"user_api_key_user_id", None
)
if kwargs.get("cache_hit", False) == True:
response_cost = 0.0
verbose_proxy_logger.info(
f"Cache Hit: response_cost {response_cost}, for user_id {user_id}"
)
verbose_proxy_logger.info(
f"response_cost {response_cost}, for user_id {user_id}"
)
@ -636,35 +715,54 @@ async def update_database(
### UPDATE USER SPEND ###
async def _update_user_db():
if user_id is None:
return
if prisma_client is not None:
existing_spend_obj = await prisma_client.get_data(user_id=user_id)
elif custom_db_client is not None:
existing_spend_obj = await custom_db_client.get_data(
key=user_id, table_name="user"
"""
- Update that user's row
- Update litellm-proxy-budget row (global proxy spend)
"""
user_ids = [user_id, litellm_proxy_budget_name]
data_list = []
for id in user_ids:
if id is None:
continue
if prisma_client is not None:
existing_spend_obj = await prisma_client.get_data(user_id=id)
elif custom_db_client is not None and id != litellm_proxy_budget_name:
existing_spend_obj = await custom_db_client.get_data(
key=id, table_name="user"
)
verbose_proxy_logger.debug(
f"Updating existing_spend_obj: {existing_spend_obj}"
)
if existing_spend_obj is None:
existing_spend = 0
else:
existing_spend = existing_spend_obj.spend
if existing_spend_obj is None:
existing_spend = 0
existing_spend_obj = LiteLLM_UserTable(
user_id=id, spend=0, max_budget=None, user_email=None
)
else:
existing_spend = existing_spend_obj.spend
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Calculate the new cost by adding the existing cost and response_cost
existing_spend_obj.spend = existing_spend + response_cost
verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}")
data_list.append(existing_spend_obj)
verbose_proxy_logger.debug(f"new cost: {new_spend}")
# Update the cost column for the given user id
if prisma_client is not None:
await prisma_client.update_data(
user_id=user_id, data={"spend": new_spend}
data_list=data_list, query_type="update_many", table_name="user"
)
elif custom_db_client is not None:
elif custom_db_client is not None and user_id is not None:
new_spend = data_list[0].spend
await custom_db_client.update_data(
key=user_id, value={"spend": new_spend}, table_name="user"
)
### UPDATE KEY SPEND ###
async def _update_key_db():
verbose_proxy_logger.debug(
f"adding spend to key db. Response cost: {response_cost}. Token: {token}."
)
if prisma_client is not None:
# Fetch the existing cost for the given token
existing_spend_obj = await prisma_client.get_data(token=token)
@ -916,7 +1014,7 @@ class ProxyConfig:
blue_color_code = "\033[94m"
reset_color_code = "\033[0m"
for key, value in litellm_settings.items():
if key == "cache":
if key == "cache" and value == True:
print(f"{blue_color_code}\nSetting Cache on Proxy") # noqa
from litellm.caching import Cache
@ -1155,6 +1253,8 @@ async def generate_key_helper_fn(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
query_type: Literal["insert_data", "update_data"] = "insert_data",
update_key_values: Optional[dict] = None,
key_alias: Optional[str] = None,
):
global prisma_client, custom_db_client
@ -1207,8 +1307,8 @@ async def generate_key_helper_fn(
config_json = json.dumps(config)
metadata_json = json.dumps(metadata)
user_id = user_id or str(uuid.uuid4())
tpm_limit = tpm_limit or sys.maxsize
rpm_limit = rpm_limit or sys.maxsize
tpm_limit = tpm_limit
rpm_limit = rpm_limit
if type(team_id) is not str:
team_id = str(team_id)
try:
@ -1228,6 +1328,7 @@ async def generate_key_helper_fn(
}
key_data = {
"token": token,
"key_alias": key_alias,
"expires": expires,
"models": models,
"aliases": aliases_json,
@ -1243,6 +1344,8 @@ async def generate_key_helper_fn(
"budget_duration": key_budget_duration,
"budget_reset_at": key_reset_at,
}
if general_settings.get("allow_user_auth", False) == True:
key_data["key_name"] = f"sk-...{token[-4:]}"
if prisma_client is not None:
## CREATE USER (If necessary)
verbose_proxy_logger.debug(f"prisma_client: Creating User={user_data}")
@ -1255,7 +1358,9 @@ async def generate_key_helper_fn(
key_data["models"] = user_row.models
elif query_type == "update_data":
user_row = await prisma_client.update_data(
data=user_data, table_name="user"
data=user_data,
table_name="user",
update_key_values=update_key_values,
)
## CREATE KEY
@ -1282,12 +1387,7 @@ async def generate_key_helper_fn(
except Exception as e:
traceback.print_exc()
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
return {
"token": token,
"expires": expires,
"user_id": user_id,
"max_budget": max_budget,
}
return key_data
async def delete_verification_token(tokens: List):
@ -1382,8 +1482,6 @@ async def initialize(
verbose_proxy_logger.setLevel(
level=logging.DEBUG
) # set proxy logs to debug
litellm.set_verbose = True
dynamic_config = {"general": {}, user_model: {}}
if config:
(
@ -1563,17 +1661,24 @@ async def startup_event():
if prisma_client is not None and master_key is not None:
# add master key to db
await generate_key_helper_fn(
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
duration=None,
models=[],
aliases={},
config={},
spend=0,
token=master_key,
user_id="default_user_id",
)
if (
prisma_client is not None
and litellm.max_budget > 0
and litellm.budget_duration is not None
):
if prisma_client is not None and litellm.max_budget > 0:
if litellm.budget_duration is None:
raise Exception(
"budget_duration not set on Proxy. budget_duration is required to use max_budget."
)
# add proxy budget to db in the user table
await generate_key_helper_fn(
user_id="litellm-proxy-budget",
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
aliases={},
@ -1582,6 +1687,10 @@ async def startup_event():
max_budget=litellm.max_budget,
budget_duration=litellm.budget_duration,
query_type="update_data",
update_key_values={
"max_budget": litellm.max_budget,
"budget_duration": litellm.budget_duration,
},
)
verbose_proxy_logger.debug(
@ -2248,7 +2357,8 @@ async def generate_key_fn(
Docs: https://docs.litellm.ai/docs/proxy/virtual_keys
Parameters:
- duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). **(Default is set to 1 hour.)**
- duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
- key_alias: Optional[str] - User defined key alias
- team_id: Optional[str] - The team id of the user
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
@ -2286,12 +2396,9 @@ async def generate_key_fn(
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
response = await generate_key_helper_fn(**data_json)
return GenerateKeyResponse(
key=response["token"],
expires=response["expires"],
user_id=response["user_id"],
)
return GenerateKeyResponse(**response)
except Exception as e:
traceback.print_exc()
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({str(e)})"),
@ -2356,10 +2463,10 @@ async def delete_key_fn(data: DeleteKeyRequest):
Delete a key from the key management system.
Parameters::
- keys (List[str]): A list of keys to delete. Example {"keys": ["sk-QWrxEynunsNpV1zT48HIrw"]}
- keys (List[str]): A list of keys or hashed keys to delete. Example {"keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]}
Returns:
- deleted_keys (List[str]): A list of deleted keys. Example {"deleted_keys": ["sk-QWrxEynunsNpV1zT48HIrw"]}
- deleted_keys (List[str]): A list of deleted keys. Example {"deleted_keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]}
Raises:
@ -2396,14 +2503,39 @@ async def delete_key_fn(data: DeleteKeyRequest):
"/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
)
async def info_key_fn(
key: str = fastapi.Query(..., description="Key in the request parameters"),
key: Optional[str] = fastapi.Query(
default=None, description="Key in the request parameters"
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Retrieve information about a key.
Parameters:
key: Optional[str] = Query parameter representing the key in the request
user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key
Returns:
Dict containing the key and its associated information
Example Curl:
```
curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \
-H "Authorization: Bearer sk-1234"
```
Example Curl - if no key is passed, it will use the Key Passed in Authorization Header
```
curl -X GET "http://0.0.0.0:8000/key/info" \
-H "Authorization: Bearer sk-02Wr4IAlN3NvPXvL5JVvDA"
```
"""
global prisma_client
try:
if prisma_client is None:
raise Exception(
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
if key == None:
key = user_api_key_dict.api_key
key_info = await prisma_client.get_data(token=key)
## REMOVE HASHED TOKEN INFO BEFORE RETURNING ##
try:
@ -2469,7 +2601,12 @@ async def spend_key_fn():
tags=["budget & spend Tracking"],
dependencies=[Depends(user_api_key_auth)],
)
async def spend_user_fn():
async def spend_user_fn(
user_id: Optional[str] = fastapi.Query(
default=None,
description="Get User Table row for user_id",
),
):
"""
View all users created, ordered by spend
@ -2478,6 +2615,12 @@ async def spend_user_fn():
curl -X GET "http://0.0.0.0:8000/spend/users" \
-H "Authorization: Bearer sk-1234"
```
View User Table row for user_id
```
curl -X GET "http://0.0.0.0:8000/spend/users?user_id=1234" \
-H "Authorization: Bearer sk-1234"
```
"""
global prisma_client
try:
@ -2486,9 +2629,15 @@ async def spend_user_fn():
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
user_info = await prisma_client.get_data(
table_name="user", query_type="find_all"
)
if user_id is not None:
user_info = await prisma_client.get_data(
table_name="user", query_type="find_unique", user_id=user_id
)
return [user_info]
else:
user_info = await prisma_client.get_data(
table_name="user", query_type="find_all"
)
return user_info
@ -2509,6 +2658,10 @@ async def view_spend_logs(
default=None,
description="Get spend logs based on api key",
),
user_id: Optional[str] = fastapi.Query(
default=None,
description="Get spend logs based on user_id",
),
request_id: Optional[str] = fastapi.Query(
default=None,
description="request_id to get spend logs for specific request_id. If none passed then pass spend logs for all requests",
@ -2534,6 +2687,12 @@ async def view_spend_logs(
curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=sk-Fn8Ej39NkBQmUagFEoUWPQ" \
-H "Authorization: Bearer sk-1234"
```
Example Request for specific user_id
```
curl -X GET "http://0.0.0.0:8000/spend/logs?user_id=ishaan@berri.ai" \
-H "Authorization: Bearer sk-1234"
```
"""
global prisma_client
try:
@ -2563,6 +2722,16 @@ async def view_spend_logs(
key_val={"key": "request_id", "value": request_id},
)
return [spend_log]
elif user_id is not None:
spend_log = await prisma_client.get_data(
table_name="spend",
query_type="find_all",
key_val={"key": "user", "value": user_id},
)
if isinstance(spend_log, list):
return spend_log
else:
return [spend_log]
else:
spend_logs = await prisma_client.get_data(
table_name="spend", query_type="find_all"
@ -2685,6 +2854,132 @@ async def user_auth(request: Request):
return "Email sent!"
@app.get("/google-login/key/generate", tags=["experimental"])
async def google_login(request: Request):
"""
Create Proxy API Keys using Google Workspace SSO. Requires setting GOOGLE_REDIRECT_URI in .env
GOOGLE_REDIRECT_URI should be the your deployed proxy endpoint, e.g. GOOGLE_REDIRECT_URI="https://litellm-production-7002.up.railway.app"
Example:
"""
GOOGLE_REDIRECT_URI = os.getenv("PROXY_BASE_URL")
if GOOGLE_REDIRECT_URI is None:
raise ProxyException(
message="PROXY_BASE_URL not set. Set it in .env file",
type="auth_error",
param="PROXY_BASE_URL",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
if GOOGLE_REDIRECT_URI.endswith("/"):
GOOGLE_REDIRECT_URI += "google-callback"
else:
GOOGLE_REDIRECT_URI += "/google-callback"
GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
if GOOGLE_CLIENT_ID is None:
GOOGLE_CLIENT_ID = (
"246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com"
)
verbose_proxy_logger.info(
f"In /google-login/key/generate, \nGOOGLE_REDIRECT_URI: {GOOGLE_REDIRECT_URI}\nGOOGLE_CLIENT_ID: {GOOGLE_CLIENT_ID}"
)
google_auth_url = f"https://accounts.google.com/o/oauth2/auth?client_id={GOOGLE_CLIENT_ID}&redirect_uri={GOOGLE_REDIRECT_URI}&response_type=code&scope=openid%20profile%20email"
return RedirectResponse(url=google_auth_url)
@app.get("/google-callback", tags=["experimental"], response_model=GenerateKeyResponse)
async def google_callback(code: str, request: Request):
import httpx
GOOGLE_REDIRECT_URI = os.getenv("PROXY_BASE_URL")
if GOOGLE_REDIRECT_URI is None:
raise ProxyException(
message="PROXY_BASE_URL not set. Set it in .env file",
type="auth_error",
param="PROXY_BASE_URL",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
# Add "/google-callback"" to your callback URL
if GOOGLE_REDIRECT_URI.endswith("/"):
GOOGLE_REDIRECT_URI += "google-callback"
else:
GOOGLE_REDIRECT_URI += "/google-callback"
GOOGLE_CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
if GOOGLE_CLIENT_ID is None:
GOOGLE_CLIENT_ID = (
"246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com"
)
GOOGLE_CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
if GOOGLE_CLIENT_SECRET is None:
GOOGLE_CLIENT_SECRET = "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"
verbose_proxy_logger.info(
f"/google-callback\n GOOGLE_REDIRECT_URI: {GOOGLE_REDIRECT_URI}\n GOOGLE_CLIENT_ID: {GOOGLE_CLIENT_ID}"
)
# Exchange code for access token
async with httpx.AsyncClient() as client:
token_url = f"https://oauth2.googleapis.com/token"
data = {
"code": code,
"client_id": GOOGLE_CLIENT_ID,
"client_secret": GOOGLE_CLIENT_SECRET,
"redirect_uri": GOOGLE_REDIRECT_URI,
"grant_type": "authorization_code",
}
response = await client.post(token_url, data=data)
# Process the response, extract user info, etc.
if response.status_code == 200:
access_token = response.json()["access_token"]
# Fetch user info using the access token
async with httpx.AsyncClient() as client:
user_info_url = "https://www.googleapis.com/oauth2/v1/userinfo"
headers = {"Authorization": f"Bearer {access_token}"}
user_info_response = await client.get(user_info_url, headers=headers)
# Process user info response
if user_info_response.status_code == 200:
user_info = user_info_response.json()
user_email = user_info.get("email")
user_name = user_info.get("name")
# we can use user_email on litellm proxy now
# TODO: Handle user info as needed, for example, store it in a database, authenticate the user, etc.
response = await generate_key_helper_fn(
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_email, "team_id": "litellm-dashboard"} # type: ignore
)
key = response["token"] # type: ignore
user_id = response["user_id"] # type: ignore
litellm_dashboard_ui = "https://litellm-dashboard.vercel.app/"
litellm_dashboard_ui += (
"?userID="
+ user_id
+ "&accessToken="
+ key
+ "&proxyBaseUrl="
+ os.getenv("PROXY_BASE_URL")
)
return RedirectResponse(url=litellm_dashboard_ui)
else:
# Handle user info retrieval error
raise HTTPException(
status_code=user_info_response.status_code,
detail=user_info_response.text,
)
else:
# Handle the error from the token exchange
raise HTTPException(status_code=response.status_code, detail=response.text)
@router.get(
"/user/info", tags=["user management"], dependencies=[Depends(user_api_key_auth)]
)
@ -2736,11 +3031,42 @@ async def user_info(
@router.post(
"/user/update", tags=["user management"], dependencies=[Depends(user_api_key_auth)]
)
async def user_update(request: Request):
async def user_update(data: UpdateUserRequest):
"""
[TODO]: Use this to update user budget
"""
pass
global prisma_client
try:
data_json: dict = data.json()
# get the row from db
if prisma_client is None:
raise Exception("Not connected to DB!")
non_default_values = {k: v for k, v in data_json.items() if v is not None}
response = await prisma_client.update_data(
user_id=data_json["user_id"],
data=non_default_values,
update_key_values=non_default_values,
)
return {"user_id": data_json["user_id"], **non_default_values}
# update based on remaining passed in values
except Exception as e:
traceback.print_exc()
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({str(e)})"),
type="auth_error",
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type="auth_error",
param=getattr(e, "param", "None"),
code=status.HTTP_400_BAD_REQUEST,
)
#### MODEL MANAGEMENT ####

13
litellm/proxy/schema.prisma
View File

@ -7,6 +7,7 @@ generator client {
provider = "prisma-client-py"
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @unique
team_id String?
@ -21,9 +22,11 @@ model LiteLLM_UserTable {
budget_reset_at DateTime?
}
// required for token gen
// Generate Tokens for Proxy
model LiteLLM_VerificationToken {
token String @unique
key_name String?
key_alias String?
spend Float @default(0.0)
expires DateTime?
models String[]
@ -40,22 +43,26 @@ model LiteLLM_VerificationToken {
budget_reset_at DateTime?
}
// store proxy config.yaml
model LiteLLM_Config {
param_name String @id
param_value Json?
}
// View spend, model, api_key per request
model LiteLLM_SpendLogs {
request_id String @unique
call_type String
api_key String @default ("")
spend Float @default(0.0)
total_tokens Int @default(0)
prompt_tokens Int @default(0)
completion_tokens Int @default(0)
startTime DateTime // Assuming start_time is a DateTime field
endTime DateTime // Assuming end_time is a DateTime field
model String @default("")
user String @default("")
modelParameters Json @default("{}")// Assuming optional_params is a JSON field
usage Json @default("{}")
metadata Json @default("{}")
cache_hit String @default("")
cache_key String @default("")
}

162
litellm/proxy/utils.py
View File

@ -181,6 +181,69 @@ class ProxyLogging:
level="Low",
)
async def budget_alerts(
self,
type: Literal["token_budget", "user_budget", "user_and_proxy_budget"],
user_max_budget: float,
user_current_spend: float,
user_info=None,
):
if self.alerting is None:
# do nothing if alerting is not switched on
return
if type == "user_and_proxy_budget":
user_info = dict(user_info)
user_id = user_info["user_id"]
max_budget = user_info["max_budget"]
spend = user_info["spend"]
user_email = user_info["user_email"]
user_info = f"""\nUser ID: {user_id}\nMax Budget: ${max_budget}\nSpend: ${spend}\nUser Email: {user_email}"""
elif type == "token_budget":
token_info = dict(user_info)
token = token_info["token"]
spend = token_info["spend"]
max_budget = token_info["max_budget"]
user_id = token_info["user_id"]
user_info = f"""\nToken: {token}\nSpend: ${spend}\nMax Budget: ${max_budget}\nUser ID: {user_id}"""
else:
user_info = str(user_info)
# percent of max_budget left to spend
percent_left = (user_max_budget - user_current_spend) / user_max_budget
verbose_proxy_logger.debug(
f"Budget Alerts: Percent left: {percent_left} for {user_info}"
)
# check if crossed budget
if user_current_spend >= user_max_budget:
verbose_proxy_logger.debug(f"Budget Crossed for {user_info}")
message = "Budget Crossed for" + user_info
await self.alerting_handler(
message=message,
level="High",
)
return
# check if 5% of max budget is left
if percent_left <= 0.05:
message = "5% budget left for" + user_info
await self.alerting_handler(
message=message,
level="Medium",
)
return
# check if 15% of max budget is left
if percent_left <= 0.15:
message = "15% budget left for" + user_info
await self.alerting_handler(
message=message,
level="Low",
)
return
return
async def alerting_handler(
self, message: str, level: Literal["Low", "Medium", "High"]
):
@ -191,6 +254,8 @@ class ProxyLogging:
- Requests are hanging
- Calls are failing
- DB Read/Writes are failing
- Proxy Close to max budget
- Key Close to max budget
Parameters:
level: str - Low|Medium|High - if calls might fail (Medium) or are failing (High); Currently, no alerts would be 'Low'.
@ -392,6 +457,7 @@ class PrismaClient:
self,
token: Optional[str] = None,
user_id: Optional[str] = None,
user_id_list: Optional[list] = None,
key_val: Optional[dict] = None,
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
query_type: Literal["find_unique", "find_all"] = "find_unique",
@ -408,7 +474,9 @@ class PrismaClient:
hashed_token = token
if token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
print_verbose("PrismaClient: find_unique")
verbose_proxy_logger.debug(
f"PrismaClient: find_unique for token: {hashed_token}"
)
if query_type == "find_unique":
response = await self.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
@ -473,11 +541,21 @@ class PrismaClient:
"budget_reset_at": {"lt": reset_at},
}
)
return response
elif table_name == "user" and query_type == "find_all":
response = await self.db.litellm_usertable.find_many( # type: ignore
order={"spend": "desc"},
)
elif query_type == "find_all" and user_id_list is not None:
user_id_values = str(tuple(user_id_list))
sql_query = f"""
SELECT *
FROM "LiteLLM_UserTable"
WHERE "user_id" IN {user_id_values}
"""
# Execute the raw query
# The asterisk before `user_id_list` unpacks the list into separate arguments
response = await self.db.query_raw(sql_query)
elif query_type == "find_all":
response = await self.db.litellm_usertable.find_many( # type: ignore
order={"spend": "desc"},
)
return response
elif table_name == "spend":
verbose_proxy_logger.debug(
@ -617,6 +695,7 @@ class PrismaClient:
user_id: Optional[str] = None,
query_type: Literal["update", "update_many"] = "update",
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
update_key_values: Optional[dict] = None,
):
"""
Update existing data
@ -649,23 +728,18 @@ class PrismaClient:
"""
if user_id is None:
user_id = db_data["user_id"]
update_user_row = await self.db.litellm_usertable.update(
if update_key_values is None:
update_key_values = db_data
update_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": user_id}, # type: ignore
data={**db_data}, # type: ignore
data={
"create": {**db_data}, # type: ignore
"update": {
**update_key_values # type: ignore
}, # just update user-specified values, if it already exists
},
)
if update_user_row is None:
# if the provided user does not exist, STILL Track this!
# make a new user with {"user_id": user_id, "spend": data['spend']}
db_data["user_id"] = user_id
update_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": user_id}, # type: ignore
data={
"create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists
},
)
print_verbose(
verbose_proxy_logger.info(
"\033[91m"
+ f"DB User Table - update succeeded {update_user_row}"
+ "\033[0m"
@ -714,13 +788,18 @@ class PrismaClient:
data_json = self.jsonify_object(data=user.model_dump())
except:
data_json = self.jsonify_object(data=user.dict())
batcher.litellm_usertable.update(
batcher.litellm_usertable.upsert(
where={"user_id": user.user_id}, # type: ignore
data={**data_json}, # type: ignore
data={
"create": {**data_json}, # type: ignore
"update": {
**data_json # type: ignore
}, # just update user-specified values, if it already exists
},
)
await batcher.commit()
print_verbose(
"\033[91m" + f"DB User Table update succeeded" + "\033[0m"
verbose_proxy_logger.info(
"\033[91m" + f"DB User Table Batch update succeeded" + "\033[0m"
)
except Exception as e:
asyncio.create_task(
@ -742,7 +821,13 @@ class PrismaClient:
Allow user to delete a key(s)
"""
try:
hashed_tokens = [self.hash_token(token=token) for token in tokens]
hashed_tokens = []
for token in tokens:
if isinstance(token, str) and token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
else:
hashed_token = token
hashed_tokens.append(hashed_token)
await self.db.litellm_verificationtoken.delete_many(
where={"token": {"in": hashed_tokens}}
)
@ -950,7 +1035,8 @@ async def send_email(sender_name, sender_email, receiver_email, subject, html):
print_verbose(f"SMTP Connection Init")
# Establish a secure connection with the SMTP server
with smtplib.SMTP(smtp_host, smtp_port) as server:
server.starttls()
if os.getenv("SMTP_TLS", 'True') != "False":
server.starttls()
# Login to your email account
server.login(smtp_username, smtp_password)
@ -959,7 +1045,7 @@ async def send_email(sender_name, sender_email, receiver_email, subject, html):
server.send_message(email_message)
except Exception as e:
print_verbose("An error occurred while sending the email:", str(e))
print_verbose("An error occurred while sending the email:" + str(e))
def hash_token(token: str):
@ -987,20 +1073,28 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time):
metadata = (
litellm_params.get("metadata", {}) or {}
) # if litellm_params['metadata'] == None
optional_params = kwargs.get("optional_params", {})
call_type = kwargs.get("call_type", "litellm.completion")
cache_hit = kwargs.get("cache_hit", False)
usage = response_obj["usage"]
if type(usage) == litellm.Usage:
usage = dict(usage)
id = response_obj.get("id", str(uuid.uuid4()))
api_key = metadata.get("user_api_key", "")
if api_key is not None and isinstance(api_key, str) and api_key.startswith("sk-"):
# hash the api_key
api_key = hash_token(api_key)
if "headers" in metadata and "authorization" in metadata["headers"]:
metadata["headers"].pop(
"authorization"
) # do not store the original `sk-..` api key in the db
if litellm.cache is not None:
cache_key = litellm.cache.get_cache_key(**kwargs)
else:
cache_key = "Cache OFF"
if cache_hit == True:
import time
id = f"{id}_cache_hit{time.time()}" # SpendLogs does not allow duplicate request_id
payload = {
"request_id": id,
@ -1011,9 +1105,11 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time):
"endTime": end_time,
"model": kwargs.get("model", ""),
"user": kwargs.get("user", ""),
"modelParameters": optional_params,
"usage": usage,
"metadata": metadata,
"cache_key": cache_key,
"total_tokens": usage.get("total_tokens", 0),
"prompt_tokens": usage.get("prompt_tokens", 0),
"completion_tokens": usage.get("completion_tokens", 0),
}
json_fields = [
@ -1038,8 +1134,6 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time):
payload[param] = payload[param].model_dump_json()
if type(payload[param]) == litellm.EmbeddingResponse:
payload[param] = payload[param].model_dump_json()
elif type(payload[param]) == litellm.Usage:
payload[param] = payload[param].model_dump_json()
else:
payload[param] = json.dumps(payload[param])

101
litellm/tests/test_s3_logs.py → litellm/tests/test_amazing_s3_logs.py
View File

@ -115,4 +115,103 @@ def test_s3_logging():
print("Passed! Testing async s3 logging")
test_s3_logging()
# test_s3_logging()
def test_s3_logging_r2():
# all s3 requests need to be in one test function
# since we are modifying stdout, and pytests runs tests in parallel
# on circle ci - we only test litellm.acompletion()
try:
# redirect stdout to log_file
# litellm.cache = litellm.Cache(
# type="s3", s3_bucket_name="litellm-r2-bucket", s3_region_name="us-west-2"
# )
litellm.set_verbose = True
from litellm._logging import verbose_logger
import logging
verbose_logger.setLevel(level=logging.DEBUG)
litellm.success_callback = ["s3"]
litellm.s3_callback_params = {
"s3_bucket_name": "litellm-r2-bucket",
"s3_aws_secret_access_key": "os.environ/R2_S3_ACCESS_KEY",
"s3_aws_access_key_id": "os.environ/R2_S3_ACCESS_ID",
"s3_endpoint_url": "os.environ/R2_S3_URL",
"s3_region_name": "os.environ/R2_S3_REGION_NAME",
}
print("Testing async s3 logging")
expected_keys = []
import time
curr_time = str(time.time())
async def _test():
return await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": f"This is a test {curr_time}"}],
max_tokens=10,
temperature=0.7,
user="ishaan-2",
)
response = asyncio.run(_test())
print(f"response: {response}")
expected_keys.append(response.id)
import boto3
s3 = boto3.client(
"s3",
endpoint_url=os.getenv("R2_S3_URL"),
region_name=os.getenv("R2_S3_REGION_NAME"),
aws_access_key_id=os.getenv("R2_S3_ACCESS_ID"),
aws_secret_access_key=os.getenv("R2_S3_ACCESS_KEY"),
)
bucket_name = "litellm-r2-bucket"
# List objects in the bucket
response = s3.list_objects(Bucket=bucket_name)
# # Sort the objects based on the LastModified timestamp
# objects = sorted(
# response["Contents"], key=lambda x: x["LastModified"], reverse=True
# )
# # Get the keys of the most recent objects
# most_recent_keys = [obj["Key"] for obj in objects]
# print(most_recent_keys)
# # for each key, get the part before "-" as the key. Do it safely
# cleaned_keys = []
# for key in most_recent_keys:
# split_key = key.split("-time=")
# cleaned_keys.append(split_key[0])
# print("\n most recent keys", most_recent_keys)
# print("\n cleaned keys", cleaned_keys)
# print("\n Expected keys: ", expected_keys)
# matches = 0
# for key in expected_keys:
# assert key in cleaned_keys
# if key in cleaned_keys:
# matches += 1
# # remove the match key
# cleaned_keys.remove(key)
# # this asserts we log, the first request + the 2nd cached request
# print("we had two matches ! passed ", matches)
# assert matches == 1
# try:
# # cleanup s3 bucket in test
# for key in most_recent_keys:
# s3.delete_object(Bucket=bucket_name, Key=key)
# except:
# # don't let cleanup fail a test
# pass
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
finally:
# post, close log file and verify
# Reset stdout to the original value
print("Passed! Testing async s3 logging")

4
litellm/tests/test_amazing_vertex_completion.py
View File

@ -95,7 +95,8 @@ def test_vertex_ai():
+ litellm.vertex_code_text_models
)
litellm.set_verbose = False
litellm.vertex_project = "reliablekeys"
vertex_ai_project = "reliablekeys"
# litellm.vertex_project = "reliablekeys"
test_models = random.sample(test_models, 1)
# test_models += litellm.vertex_language_models # always test gemini-pro
@ -117,6 +118,7 @@ def test_vertex_ai():
model=model,
messages=[{"role": "user", "content": "hi"}],
temperature=0.7,
vertex_ai_project=vertex_ai_project,
)
print("\nModel Response", response)
print(response)

9
litellm/tests/test_caching.py
View File

@ -127,9 +127,10 @@ def test_caching_with_models_v2():
]
litellm.cache = Cache()
print("test2 for caching")
litellm.set_verbose = True
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response3 = completion(model="command-nightly", messages=messages, caching=True)
response3 = completion(model="azure/chatgpt-v-2", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
@ -286,7 +287,7 @@ def test_redis_cache_completion():
response3 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, temperature=0.5
)
response4 = completion(model="command-nightly", messages=messages, caching=True)
response4 = completion(model="azure/chatgpt-v-2", messages=messages, caching=True)
print("\nresponse 1", response1)
print("\nresponse 2", response2)
@ -723,8 +724,8 @@ def test_cache_override():
print(f"Embedding 2 response time: {end_time - start_time} seconds")
assert (
end_time - start_time > 0.1
) # ensure 2nd response comes in over 0.1s. This should not be cached.
end_time - start_time > 0.05
) # ensure 2nd response comes in over 0.05s. This should not be cached.
# test_cache_override()

51
litellm/tests/test_completion.py
View File

@ -191,6 +191,21 @@ def test_completion_gpt4_turbo():
# test_completion_gpt4_turbo()
def test_completion_gpt4_turbo_0125():
try:
response = completion(
model="gpt-4-0125-preview",
messages=messages,
max_tokens=10,
)
print(response)
except openai.RateLimitError:
print("got a rate liimt error")
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="this test is flaky")
def test_completion_gpt4_vision():
try:
@ -224,7 +239,7 @@ def test_completion_gpt4_vision():
def test_completion_azure_gpt4_vision():
# azure/gpt-4 vision takes 5seconds to respond
# azure/gpt-4, vision takes 5seconds to respond
try:
litellm.set_verbose = True
response = completion(
@ -500,22 +515,22 @@ def hf_test_completion_tgi():
# hf_test_error_logs()
def test_completion_cohere(): # commenting for now as the cohere endpoint is being flaky
try:
litellm.CohereConfig(max_tokens=1000, stop_sequences=["a"])
response = completion(
model="command-nightly", messages=messages, logger_fn=logger_fn
)
# Add any assertions here to check the response
print(response)
response_str = response["choices"][0]["message"]["content"]
response_str_2 = response.choices[0].message.content
if type(response_str) != str:
pytest.fail(f"Error occurred: {e}")
if type(response_str_2) != str:
pytest.fail(f"Error occurred: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# def test_completion_cohere(): # commenting for now as the cohere endpoint is being flaky
# try:
# litellm.CohereConfig(max_tokens=10, stop_sequences=["a"])
# response = completion(
# model="command-nightly", messages=messages, logger_fn=logger_fn
# )
# # Add any assertions here to check the response
# print(response)
# response_str = response["choices"][0]["message"]["content"]
# response_str_2 = response.choices[0].message.content
# if type(response_str) != str:
# pytest.fail(f"Error occurred: {e}")
# if type(response_str_2) != str:
# pytest.fail(f"Error occurred: {e}")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_cohere()
@ -854,7 +869,7 @@ def test_completion_anyscale_with_functions():
def test_completion_azure_key_completion_arg():
# this tests if we can pass api_key to completion, when it's not in the env
# this tests if we can pass api_key to completion, when it's not in the env.
# DO NOT REMOVE THIS TEST. No MATTER WHAT Happens!
# If you want to remove it, speak to Ishaan!
# Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this

77
litellm/tests/test_completion_cost.py
View File

@ -124,7 +124,7 @@ def test_cost_azure_gpt_35():
)
test_cost_azure_gpt_35()
# test_cost_azure_gpt_35()
def test_cost_azure_embedding():
@ -158,3 +158,78 @@ def test_cost_azure_embedding():
# test_cost_azure_embedding()
def test_cost_openai_image_gen():
cost = litellm.completion_cost(
model="dall-e-2", size="1024-x-1024", quality="standard", n=1
)
assert cost == 0.019922944
def test_cost_bedrock_pricing():
"""
- get pricing specific to region for a model
"""
from litellm import ModelResponse, Choices, Message
from litellm.utils import Usage
litellm.set_verbose = True
input_tokens = litellm.token_counter(
model="bedrock/anthropic.claude-instant-v1",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
print(f"input_tokens: {input_tokens}")
output_tokens = litellm.token_counter(
model="bedrock/anthropic.claude-instant-v1",
text="It's all going well",
count_response_tokens=True,
)
print(f"output_tokens: {output_tokens}")
resp = ModelResponse(
id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content="It's all going well",
role="assistant",
),
)
],
created=1700775391,
model="anthropic.claude-instant-v1",
object="chat.completion",
system_fingerprint=None,
usage=Usage(
prompt_tokens=input_tokens,
completion_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
),
)
resp._hidden_params = {
"custom_llm_provider": "bedrock",
"region_name": "ap-northeast-1",
}
cost = litellm.completion_cost(
model="anthropic.claude-instant-v1",
completion_response=resp,
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
predicted_cost = input_tokens * 0.00000223 + 0.00000755 * output_tokens
assert cost == predicted_cost
def test_cost_bedrock_pricing_actual_calls():
litellm.set_verbose = True
model = "anthropic.claude-instant-v1"
messages = [{"role": "user", "content": "Hey, how's it going?"}]
response = litellm.completion(model=model, messages=messages)
assert response._hidden_params["region_name"] is not None
cost = litellm.completion_cost(
completion_response=response,
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
assert cost > 0

62
litellm/tests/test_configs/test_config_no_auth.yaml
View File

@ -53,9 +53,9 @@ model_list:
api_key: os.environ/AZURE_API_KEY
api_version: 2023-07-01-preview
model: azure/azure-embedding-model
model_name: azure-embedding-model
model_info:
mode: "embedding"
mode: embedding
model_name: azure-embedding-model
- litellm_params:
model: gpt-3.5-turbo
model_info:
@ -80,43 +80,49 @@ model_list:
description: this is a test openai model
id: 9b1ef341-322c-410a-8992-903987fef439
model_name: test_openai_models
- model_name: amazon-embeddings
litellm_params:
model: "bedrock/amazon.titan-embed-text-v1"
- litellm_params:
model: bedrock/amazon.titan-embed-text-v1
model_info:
mode: embedding
- model_name: "GPT-J 6B - Sagemaker Text Embedding (Internal)"
litellm_params:
model: "sagemaker/berri-benchmarking-gpt-j-6b-fp16"
model_name: amazon-embeddings
- litellm_params:
model: sagemaker/berri-benchmarking-gpt-j-6b-fp16
model_info:
mode: embedding
- model_name: dall-e-3
litellm_params:
model_name: GPT-J 6B - Sagemaker Text Embedding (Internal)
- litellm_params:
model: dall-e-3
model_info:
mode: image_generation
- model_name: dall-e-3
litellm_params:
model: "azure/dall-e-3-test"
api_version: "2023-12-01-preview"
api_base: "os.environ/AZURE_SWEDEN_API_BASE"
api_key: "os.environ/AZURE_SWEDEN_API_KEY"
model_name: dall-e-3
- litellm_params:
api_base: os.environ/AZURE_SWEDEN_API_BASE
api_key: os.environ/AZURE_SWEDEN_API_KEY
api_version: 2023-12-01-preview
model: azure/dall-e-3-test
model_info:
mode: image_generation
- model_name: dall-e-2
litellm_params:
model: "azure/"
api_version: "2023-06-01-preview"
api_base: "os.environ/AZURE_API_BASE"
api_key: "os.environ/AZURE_API_KEY"
model_name: dall-e-3
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_version: 2023-06-01-preview
model: azure/
model_info:
mode: image_generation
- model_name: text-embedding-ada-002
litellm_params:
model_name: dall-e-2
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_version: 2023-07-01-preview
model: azure/azure-embedding-model
api_base: "os.environ/AZURE_API_BASE"
api_key: "os.environ/AZURE_API_KEY"
api_version: "2023-07-01-preview"
model_info:
base_model: text-embedding-ada-002
mode: embedding
base_model: text-embedding-ada-002
model_name: text-embedding-ada-002
- litellm_params:
model: gpt-3.5-turbo
model_info:
description: this is a test openai model
id: 34cb2419-7c63-44ae-a189-53f1d1ce5953
model_name: test_openai_models

113
litellm/tests/test_custom_callback_input.py
View File

@ -556,6 +556,47 @@ async def test_async_chat_bedrock_stream():
# asyncio.run(test_async_chat_bedrock_stream())
## Test Sagemaker + Async
@pytest.mark.asyncio
async def test_async_chat_sagemaker_stream():
try:
customHandler = CompletionCustomHandler()
litellm.callbacks = [customHandler]
response = await litellm.acompletion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=[{"role": "user", "content": "Hi 👋 - i'm async sagemaker"}],
)
# test streaming
response = await litellm.acompletion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=[{"role": "user", "content": "Hi 👋 - i'm async sagemaker"}],
stream=True,
)
print(f"response: {response}")
async for chunk in response:
print(f"chunk: {chunk}")
continue
## test failure callback
try:
response = await litellm.acompletion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=[{"role": "user", "content": "Hi 👋 - i'm async sagemaker"}],
aws_region_name="my-bad-key",
stream=True,
)
async for chunk in response:
continue
except:
pass
time.sleep(1)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
except Exception as e:
pytest.fail(f"An exception occurred: {str(e)}")
# Text Completion
@ -778,47 +819,49 @@ async def test_async_embedding_azure_caching():
# Image Generation
# ## Test OpenAI + Sync
# def test_image_generation_openai():
# try:
# customHandler_success = CompletionCustomHandler()
# customHandler_failure = CompletionCustomHandler()
# litellm.callbacks = [customHandler_success]
## Test OpenAI + Sync
def test_image_generation_openai():
try:
customHandler_success = CompletionCustomHandler()
customHandler_failure = CompletionCustomHandler()
# litellm.callbacks = [customHandler_success]
# litellm.set_verbose = True
# litellm.set_verbose = True
# response = litellm.image_generation(
# prompt="A cute baby sea otter", model="dall-e-3"
# )
# response = litellm.image_generation(
# prompt="A cute baby sea otter", model="dall-e-3"
# )
# print(f"response: {response}")
# assert len(response.data) > 0
# print(f"response: {response}")
# assert len(response.data) > 0
# print(f"customHandler_success.errors: {customHandler_success.errors}")
# print(f"customHandler_success.states: {customHandler_success.states}")
# assert len(customHandler_success.errors) == 0
# assert len(customHandler_success.states) == 3 # pre, post, success
# # test failure callback
# litellm.callbacks = [customHandler_failure]
# try:
# response = litellm.image_generation(
# prompt="A cute baby sea otter", model="dall-e-4"
# )
# except:
# pass
# print(f"customHandler_failure.errors: {customHandler_failure.errors}")
# print(f"customHandler_failure.states: {customHandler_failure.states}")
# assert len(customHandler_failure.errors) == 0
# assert len(customHandler_failure.states) == 3 # pre, post, failure
# except litellm.RateLimitError as e:
# pass
# except litellm.ContentPolicyViolationError:
# pass # OpenAI randomly raises these errors - skip when they occur
# except Exception as e:
# pytest.fail(f"An exception occurred - {str(e)}")
# print(f"customHandler_success.errors: {customHandler_success.errors}")
# print(f"customHandler_success.states: {customHandler_success.states}")
# assert len(customHandler_success.errors) == 0
# assert len(customHandler_success.states) == 3 # pre, post, success
# test failure callback
litellm.callbacks = [customHandler_failure]
try:
response = litellm.image_generation(
prompt="A cute baby sea otter",
model="dall-e-2",
api_key="my-bad-api-key",
)
except:
pass
print(f"customHandler_failure.errors: {customHandler_failure.errors}")
print(f"customHandler_failure.states: {customHandler_failure.states}")
assert len(customHandler_failure.errors) == 0
assert len(customHandler_failure.states) == 3 # pre, post, failure
except litellm.RateLimitError as e:
pass
except litellm.ContentPolicyViolationError:
pass # OpenAI randomly raises these errors - skip when they occur
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_openai()
test_image_generation_openai()
## Test OpenAI + Async
## Test Azure + Sync

4
litellm/tests/test_custom_logger.py
View File

@ -206,12 +206,12 @@ def test_azure_completion_stream():
# checks if the model response available in the async + stream callbacks is equal to the received response
customHandler2 = MyCustomHandler()
litellm.callbacks = [customHandler2]
litellm.set_verbose = False
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "user",
"content": "write 1 sentence about litellm being amazing",
"content": f"write 1 sentence about litellm being amazing {time.time()}",
},
]
complete_streaming_response = ""

3
litellm/tests/test_dynamodb_logs.py
View File

@ -33,6 +33,7 @@ def pre_request():
import re
@pytest.mark.skip
def verify_log_file(log_file_path):
with open(log_file_path, "r") as log_file:
log_content = log_file.read()
@ -123,7 +124,7 @@ def test_dynamo_logging():
sys.stdout = original_stdout
# Close the file
log_file.close()
verify_log_file(file_name)
# verify_log_file(file_name)
print("Passed! Testing async dynamoDB logging")

44
litellm/tests/test_embedding.py
View File

@ -57,6 +57,48 @@ def test_openai_embedding():
# test_openai_embedding()
def test_openai_embedding_3():
try:
litellm.set_verbose = True
response = embedding(
model="text-embedding-3-small",
input=["good morning from litellm", "this is another item"],
metadata={"anything": "good day"},
dimensions=5,
)
print(f"response:", response)
litellm_response = dict(response)
litellm_response_keys = set(litellm_response.keys())
litellm_response_keys.discard("_response_ms")
print(litellm_response_keys)
print("LiteLLM Response\n")
# print(litellm_response)
# same request with OpenAI 1.0+
import openai
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
response = client.embeddings.create(
model="text-embedding-3-small",
input=["good morning from litellm", "this is another item"],
dimensions=5,
)
response = dict(response)
openai_response_keys = set(response.keys())
print(openai_response_keys)
assert (
litellm_response_keys == openai_response_keys
) # ENSURE the Keys in litellm response is exactly what the openai package returns
assert (
len(litellm_response["data"]) == 2
) # expect two embedding responses from litellm_response since input had two
print(openai_response_keys)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_openai_azure_embedding_simple():
try:
litellm.set_verbose = True
@ -186,7 +228,7 @@ def test_cohere_embedding3():
pytest.fail(f"Error occurred: {e}")
test_cohere_embedding3()
# test_cohere_embedding3()
def test_bedrock_embedding_titan():

14
litellm/tests/test_key_generate_dynamodb.py
View File

@ -33,7 +33,7 @@ from litellm.proxy.proxy_server import (
)
from litellm.proxy._types import NewUserRequest, DynamoDBArgs, GenerateKeyRequest
from litellm.proxy.utils import DBClient
from litellm.proxy.utils import DBClient, hash_token
from starlette.datastructures import URL
@ -116,6 +116,10 @@ def test_call_with_invalid_key(custom_db_client):
def test_call_with_invalid_model(custom_db_client):
# 3. Make a call to a key with an invalid model - expect to fail
from litellm._logging import verbose_proxy_logger
import logging
verbose_proxy_logger.setLevel(logging.DEBUG)
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
try:
@ -232,7 +236,7 @@ def test_call_with_user_over_budget(custom_db_client):
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},
@ -305,7 +309,7 @@ def test_call_with_user_over_budget_stream(custom_db_client):
"complete_streaming_response": resp,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},
@ -376,7 +380,7 @@ def test_call_with_user_key_budget(custom_db_client):
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},
@ -449,7 +453,7 @@ def test_call_with_key_over_budget_stream(custom_db_client):
"complete_streaming_response": resp,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},

309
litellm/tests/test_key_generate_prisma.py
View File

@ -12,6 +12,8 @@
# 11. Generate a Key, cal key/info, call key/update, call key/info
# 12. Make a call with key over budget, expect to fail
# 14. Make a streaming chat/completions call with key over budget, expect to fail
# 15. Generate key, when `allow_user_auth`=False - check if `/key/info` returns key_name=null
# 16. Generate key, when `allow_user_auth`=True - check if `/key/info` returns key_name=sk...<last-4-digits>
# function to call to generate key - async def new_user(data: NewUserRequest):
@ -24,7 +26,7 @@ from fastapi import Request
from datetime import datetime
load_dotenv()
import os, io
import os, io, time
# this file is to test litellm/proxy
@ -46,7 +48,7 @@ from litellm.proxy.proxy_server import (
spend_key_fn,
view_spend_logs,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
from litellm._logging import verbose_proxy_logger
verbose_proxy_logger.setLevel(level=logging.DEBUG)
@ -83,6 +85,10 @@ def prisma_client():
# Reset litellm.proxy.proxy_server.prisma_client to None
litellm.proxy.proxy_server.custom_db_client = None
litellm.proxy.proxy_server.litellm_proxy_budget_name = (
f"litellm-proxy-budget-{time.time()}"
)
litellm.proxy.proxy_server.user_custom_key_generate = None
return prisma_client
@ -282,6 +288,90 @@ def test_call_with_user_over_budget(prisma_client):
print(vars(e))
def test_call_with_proxy_over_budget(prisma_client):
# 5.1 Make a call with a proxy over budget, expect to fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}"
setattr(
litellm.proxy.proxy_server,
"litellm_proxy_budget_name",
litellm_proxy_budget_name,
)
try:
async def test():
await litellm.proxy.proxy_server.prisma_client.connect()
## CREATE PROXY + USER BUDGET ##
request = NewUserRequest(
max_budget=0.00001, user_id=litellm_proxy_budget_name
)
await new_user(request)
request = NewUserRequest()
key = await new_user(request)
print(key)
generated_key = key.key
user_id = key.user_id
bearer_token = "Bearer " + generated_key
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
# update spend using track_cost callback, make 2nd request, it should fail
from litellm.proxy.proxy_server import track_cost_callback
from litellm import ModelResponse, Choices, Message, Usage
resp = ModelResponse(
id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
role="assistant",
),
)
],
model="gpt-35-turbo", # azure always has model written like this
usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
)
await track_cost_callback(
kwargs={
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key_user_id": user_id,
}
},
"response_cost": 0.00002,
},
completion_response=resp,
start_time=datetime.now(),
end_time=datetime.now(),
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
pytest.fail(f"This should have failed!. They key crossed it's budget")
asyncio.run(test())
except Exception as e:
if hasattr(e, "message"):
error_detail = e.message
else:
error_detail = traceback.format_exc()
assert "Authentication Error, ExceededBudget:" in error_detail
print(vars(e))
def test_call_with_user_over_budget_stream(prisma_client):
# 6. Make a call with a key over budget, expect to fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
@ -358,6 +448,93 @@ def test_call_with_user_over_budget_stream(prisma_client):
print(vars(e))
def test_call_with_proxy_over_budget_stream(prisma_client):
# 6.1 Make a call with a global proxy over budget, expect to fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}"
setattr(
litellm.proxy.proxy_server,
"litellm_proxy_budget_name",
litellm_proxy_budget_name,
)
from litellm._logging import verbose_proxy_logger
import logging
litellm.set_verbose = True
verbose_proxy_logger.setLevel(logging.DEBUG)
try:
async def test():
await litellm.proxy.proxy_server.prisma_client.connect()
## CREATE PROXY + USER BUDGET ##
request = NewUserRequest(
max_budget=0.00001, user_id=litellm_proxy_budget_name
)
await new_user(request)
request = NewUserRequest()
key = await new_user(request)
print(key)
generated_key = key.key
user_id = key.user_id
bearer_token = "Bearer " + generated_key
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
# update spend using track_cost callback, make 2nd request, it should fail
from litellm.proxy.proxy_server import track_cost_callback
from litellm import ModelResponse, Choices, Message, Usage
resp = ModelResponse(
id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
role="assistant",
),
)
],
model="gpt-35-turbo", # azure always has model written like this
usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
)
await track_cost_callback(
kwargs={
"stream": True,
"complete_streaming_response": resp,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key_user_id": user_id,
}
},
"response_cost": 0.00002,
},
completion_response=ModelResponse(),
start_time=datetime.now(),
end_time=datetime.now(),
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
pytest.fail(f"This should have failed!. They key crossed it's budget")
asyncio.run(test())
except Exception as e:
error_detail = e.message
assert "Authentication Error, ExceededBudget:" in error_detail
print(vars(e))
def test_generate_and_call_with_valid_key_never_expires(prisma_client):
# 7. Make a call with an key that never expires, expect to pass
@ -716,6 +893,9 @@ def test_call_with_key_over_budget(prisma_client):
# update spend using track_cost callback, make 2nd request, it should fail
from litellm.proxy.proxy_server import track_cost_callback
from litellm import ModelResponse, Choices, Message, Usage
from litellm.caching import Cache
litellm.cache = Cache()
import time
request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"
@ -741,7 +921,7 @@ def test_call_with_key_over_budget(prisma_client):
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},
@ -763,6 +943,10 @@ def test_call_with_key_over_budget(prisma_client):
assert spend_log.request_id == request_id
assert spend_log.spend == float("2e-05")
assert spend_log.model == "chatgpt-v-2"
assert (
spend_log.cache_key
== "a61ae14fe4a8b8014a61e6ae01a100c8bc6770ac37c293242afed954bc69207d"
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
@ -776,6 +960,76 @@ def test_call_with_key_over_budget(prisma_client):
print(vars(e))
@pytest.mark.asyncio()
async def test_call_with_key_never_over_budget(prisma_client):
# Make a call with a key with budget=None, it should never fail
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
try:
await litellm.proxy.proxy_server.prisma_client.connect()
request = GenerateKeyRequest(max_budget=None)
key = await generate_key_fn(request)
print(key)
generated_key = key.key
user_id = key.user_id
bearer_token = "Bearer " + generated_key
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
# update spend using track_cost callback, make 2nd request, it should fail
from litellm.proxy.proxy_server import track_cost_callback
from litellm import ModelResponse, Choices, Message, Usage
import time
request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"
resp = ModelResponse(
id=request_id,
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
role="assistant",
),
)
],
model="gpt-35-turbo", # azure always has model written like this
usage=Usage(
prompt_tokens=210000, completion_tokens=200000, total_tokens=41000
),
)
await track_cost_callback(
kwargs={
"model": "chatgpt-v-2",
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},
"response_cost": 200000,
},
completion_response=resp,
start_time=datetime.now(),
end_time=datetime.now(),
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
except Exception as e:
pytest.fail(f"This should have not failed!. They key uses max_budget=None. {e}")
@pytest.mark.asyncio()
async def test_call_with_key_over_budget_stream(prisma_client):
# 14. Make a call with a key over budget, expect to fail
@ -832,7 +1086,7 @@ async def test_call_with_key_over_budget_stream(prisma_client):
"complete_streaming_response": resp,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
}
},
@ -861,7 +1115,7 @@ async def test_view_spend_per_user(prisma_client):
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
try:
user_by_spend = await spend_user_fn()
user_by_spend = await spend_user_fn(user_id=None)
assert type(user_by_spend) == list
assert len(user_by_spend) > 0
first_user = user_by_spend[0]
@ -889,3 +1143,48 @@ async def test_view_spend_per_key(prisma_client):
except Exception as e:
print("Got Exception", e)
pytest.fail(f"Got exception {e}")
@pytest.mark.asyncio()
async def test_key_name_null(prisma_client):
"""
- create key
- get key info
- assert key_name is null
"""
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
try:
request = GenerateKeyRequest()
key = await generate_key_fn(request)
generated_key = key.key
result = await info_key_fn(key=generated_key)
print("result from info_key_fn", result)
assert result["info"]["key_name"] is None
except Exception as e:
print("Got Exception", e)
pytest.fail(f"Got exception {e}")
@pytest.mark.asyncio()
async def test_key_name_set(prisma_client):
"""
- create key
- get key info
- assert key_name is not null
"""
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm.proxy.proxy_server, "general_settings", {"allow_user_auth": True})
await litellm.proxy.proxy_server.prisma_client.connect()
try:
request = GenerateKeyRequest()
key = await generate_key_fn(request)
generated_key = key.key
result = await info_key_fn(key=generated_key)
print("result from info_key_fn", result)
assert isinstance(result["info"]["key_name"], str)
except Exception as e:
print("Got Exception", e)
pytest.fail(f"Got exception {e}")

2
litellm/tests/test_proxy_pass_user_config.py
View File

@ -32,7 +32,7 @@ from litellm.proxy.proxy_server import (
) # Replace with the actual module where your FastAPI router is defined
# Your bearer token
token = ""
token = "sk-1234"
headers = {"Authorization": f"Bearer {token}"}

2
litellm/tests/test_proxy_server.py
View File

@ -31,7 +31,7 @@ from litellm.proxy.proxy_server import (
) # Replace with the actual module where your FastAPI router is defined
# Your bearer token
token = ""
token = "sk-1234"
headers = {"Authorization": f"Bearer {token}"}

2
litellm/tests/test_proxy_server_caching.py
View File

@ -33,7 +33,7 @@ from litellm.proxy.proxy_server import (
) # Replace with the actual module where your FastAPI router is defined
# Your bearer token
token = ""
token = "sk-1234"
headers = {"Authorization": f"Bearer {token}"}

80
litellm/tests/test_streaming.py
View File

@ -847,9 +847,13 @@ def test_sagemaker_weird_response():
logging_obj=logging_obj,
)
complete_response = ""
for chunk in response:
print(chunk)
complete_response += chunk["choices"][0]["delta"]["content"]
for idx, chunk in enumerate(response):
# print
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
complete_response += chunk
if finished:
break
assert len(complete_response) > 0
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
@ -872,41 +876,53 @@ async def test_sagemaker_streaming_async():
)
# Add any assertions here to check the response
print(response)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
idx = 0
async for chunk in response:
complete_response += chunk.choices[0].delta.content or ""
print(f"complete_response: {complete_response}")
assert len(complete_response) > 0
# print
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
complete_response += chunk
if finished:
break
idx += 1
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"completion_response: {complete_response}")
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# def test_completion_sagemaker_stream():
# try:
# response = completion(
# model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b",
# messages=messages,
# temperature=0.2,
# max_tokens=80,
# stream=True,
# )
# complete_response = ""
# has_finish_reason = False
# # Add any assertions here to check the response
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finish_reason = finished
# if finished:
# break
# complete_response += chunk
# if has_finish_reason is False:
# raise Exception("finish reason not set for last chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# except InvalidRequestError as e:
# pass
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
def test_completion_sagemaker_stream():
try:
response = completion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=messages,
temperature=0.2,
max_tokens=80,
stream=True,
)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
for idx, chunk in enumerate(response):
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
if finished:
break
complete_response += chunk
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_sagemaker_stream()

280
litellm/utils.py
View File

@ -10,6 +10,7 @@
import sys, re, binascii, struct
import litellm
import dotenv, json, traceback, threading, base64, ast
import subprocess, os
import litellm, openai
import itertools
@ -713,6 +714,7 @@ class ImageResponse(OpenAIObject):
############################################################
def print_verbose(print_statement):
try:
verbose_logger.debug(print_statement)
if litellm.set_verbose:
print(print_statement) # noqa
except:
@ -1090,7 +1092,12 @@ class Logging:
else: # streaming chunks + image gen.
self.model_call_details["response_cost"] = None
if litellm.max_budget and self.stream:
if (
litellm.max_budget
and self.stream
and result is not None
and "content" in result
):
time_diff = (end_time - start_time).total_seconds()
float_diff = float(time_diff)
litellm._current_cost += litellm.completion_cost(
@ -1412,7 +1419,9 @@ class Logging:
"""
Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
"""
print_verbose(f"Async success callbacks: {litellm._async_success_callback}")
verbose_logger.debug(
f"Async success callbacks: {litellm._async_success_callback}"
)
start_time, end_time, result = self._success_handler_helper_fn(
start_time=start_time, end_time=end_time, result=result, cache_hit=cache_hit
)
@ -1421,7 +1430,7 @@ class Logging:
if self.stream:
if result.choices[0].finish_reason is not None: # if it's the last chunk
self.streaming_chunks.append(result)
# print_verbose(f"final set of received chunks: {self.streaming_chunks}")
# verbose_logger.debug(f"final set of received chunks: {self.streaming_chunks}")
try:
complete_streaming_response = litellm.stream_chunk_builder(
self.streaming_chunks,
@ -1430,14 +1439,16 @@ class Logging:
end_time=end_time,
)
except Exception as e:
print_verbose(
verbose_logger.debug(
f"Error occurred building stream chunk: {traceback.format_exc()}"
)
complete_streaming_response = None
else:
self.streaming_chunks.append(result)
if complete_streaming_response is not None:
print_verbose("Async success callbacks: Got a complete streaming response")
verbose_logger.debug(
"Async success callbacks: Got a complete streaming response"
)
self.model_call_details[
"complete_streaming_response"
] = complete_streaming_response
@ -2019,14 +2030,15 @@ def client(original_function):
start_time=start_time,
)
## check if metadata is passed in
litellm_params = {}
if "metadata" in kwargs:
litellm_params = {"metadata": kwargs["metadata"]}
logging_obj.update_environment_variables(
model=model,
user="",
optional_params={},
litellm_params=litellm_params,
)
litellm_params["metadata"] = kwargs["metadata"]
logging_obj.update_environment_variables(
model=model,
user="",
optional_params={},
litellm_params=litellm_params,
)
return logging_obj
except Exception as e:
import logging
@ -2178,7 +2190,6 @@ def client(original_function):
result = original_function(*args, **kwargs)
end_time = datetime.datetime.now()
if "stream" in kwargs and kwargs["stream"] == True:
# TODO: Add to cache for streaming
if (
"complete_response" in kwargs
and kwargs["complete_response"] == True
@ -2872,8 +2883,13 @@ def token_counter(
print_verbose(
f"Token Counter - using generic token counter, for model={model}"
)
enc = tokenizer_json["tokenizer"].encode(text)
num_tokens = len(enc)
num_tokens = openai_token_counter(
text=text, # type: ignore
model="gpt-3.5-turbo",
messages=messages,
is_tool_call=is_tool_call,
count_response_tokens=count_response_tokens,
)
else:
num_tokens = len(encoding.encode(text)) # type: ignore
return num_tokens
@ -2885,6 +2901,7 @@ def cost_per_token(
completion_tokens=0,
response_time_ms=None,
custom_llm_provider=None,
region_name=None,
):
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -2901,16 +2918,46 @@ def cost_per_token(
prompt_tokens_cost_usd_dollar = 0
completion_tokens_cost_usd_dollar = 0
model_cost_ref = litellm.model_cost
model_with_provider = model
if custom_llm_provider is not None:
model_with_provider = custom_llm_provider + "/" + model
else:
model_with_provider = model
if region_name is not None:
model_with_provider_and_region = (
f"{custom_llm_provider}/{region_name}/{model}"
)
if (
model_with_provider_and_region in model_cost_ref
): # use region based pricing, if it's available
model_with_provider = model_with_provider_and_region
# see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models
verbose_logger.debug(f"Looking up model={model} in model_cost_map")
print_verbose(f"Looking up model={model} in model_cost_map")
if model_with_provider in model_cost_ref:
print_verbose(
f"Success: model={model_with_provider} in model_cost_map - {model_cost_ref[model_with_provider]}"
)
print_verbose(
f"applying cost={model_cost_ref[model_with_provider].get('input_cost_per_token', None)} for prompt_tokens={prompt_tokens}"
)
prompt_tokens_cost_usd_dollar = (
model_cost_ref[model_with_provider]["input_cost_per_token"] * prompt_tokens
)
print_verbose(
f"calculated prompt_tokens_cost_usd_dollar: {prompt_tokens_cost_usd_dollar}"
)
print_verbose(
f"applying cost={model_cost_ref[model_with_provider].get('output_cost_per_token', None)} for completion_tokens={completion_tokens}"
)
completion_tokens_cost_usd_dollar = (
model_cost_ref[model_with_provider]["output_cost_per_token"]
* completion_tokens
)
print_verbose(
f"calculated completion_tokens_cost_usd_dollar: {completion_tokens_cost_usd_dollar}"
)
return prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar
if model in model_cost_ref:
verbose_logger.debug(f"Success: model={model} in model_cost_map")
verbose_logger.debug(
print_verbose(f"Success: model={model} in model_cost_map")
print_verbose(
f"prompt_tokens={prompt_tokens}; completion_tokens={completion_tokens}"
)
if (
@ -2928,7 +2975,7 @@ def cost_per_token(
model_cost_ref[model].get("input_cost_per_second", None) is not None
and response_time_ms is not None
):
verbose_logger.debug(
print_verbose(
f"For model={model} - input_cost_per_second: {model_cost_ref[model].get('input_cost_per_second')}; response time: {response_time_ms}"
)
## COST PER SECOND ##
@ -2936,30 +2983,12 @@ def cost_per_token(
model_cost_ref[model]["input_cost_per_second"] * response_time_ms / 1000
)
completion_tokens_cost_usd_dollar = 0.0
verbose_logger.debug(
print_verbose(
f"Returned custom cost for model={model} - prompt_tokens_cost_usd_dollar: {prompt_tokens_cost_usd_dollar}, completion_tokens_cost_usd_dollar: {completion_tokens_cost_usd_dollar}"
)
return prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar
elif model_with_provider in model_cost_ref:
verbose_logger.debug(
f"Looking up model={model_with_provider} in model_cost_map"
)
verbose_logger.debug(
f"applying cost={model_cost_ref[model_with_provider]['input_cost_per_token']} for prompt_tokens={prompt_tokens}"
)
prompt_tokens_cost_usd_dollar = (
model_cost_ref[model_with_provider]["input_cost_per_token"] * prompt_tokens
)
verbose_logger.debug(
f"applying cost={model_cost_ref[model_with_provider]['output_cost_per_token']} for completion_tokens={completion_tokens}"
)
completion_tokens_cost_usd_dollar = (
model_cost_ref[model_with_provider]["output_cost_per_token"]
* completion_tokens
)
return prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar
elif "ft:gpt-3.5-turbo" in model:
verbose_logger.debug(f"Cost Tracking: {model} is an OpenAI FinteTuned LLM")
print_verbose(f"Cost Tracking: {model} is an OpenAI FinteTuned LLM")
# fuzzy match ft:gpt-3.5-turbo:abcd-id-cool-litellm
prompt_tokens_cost_usd_dollar = (
model_cost_ref["ft:gpt-3.5-turbo"]["input_cost_per_token"] * prompt_tokens
@ -3016,7 +3045,14 @@ def completion_cost(
prompt="",
messages: List = [],
completion="",
total_time=0.0, # used for replicate
total_time=0.0, # used for replicate, sagemaker
### REGION ###
custom_llm_provider=None,
region_name=None, # used for bedrock pricing
### IMAGE GEN ###
size=None,
quality=None,
n=None, # number of images
):
"""
Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm.
@ -3061,12 +3097,13 @@ def completion_cost(
model = (
model or completion_response["model"]
) # check if user passed an override for model, if it's none check completion_response['model']
if completion_response is not None and hasattr(
completion_response, "_hidden_params"
):
if hasattr(completion_response, "_hidden_params"):
custom_llm_provider = completion_response._hidden_params.get(
"custom_llm_provider", ""
)
region_name = completion_response._hidden_params.get(
"region_name", region_name
)
else:
if len(messages) > 0:
prompt_tokens = token_counter(model=model, messages=messages)
@ -3078,6 +3115,37 @@ def completion_cost(
f"Model is None and does not exist in passed completion_response. Passed completion_response={completion_response}, model={model}"
)
if size is not None and n is not None:
### IMAGE GENERATION COST CALCULATION ###
image_gen_model_name = f"{size}/{model}"
image_gen_model_name_with_quality = image_gen_model_name
if quality is not None:
image_gen_model_name_with_quality = f"{quality}/{image_gen_model_name}"
size = size.split("-x-")
height = int(size[0])
width = int(size[1])
verbose_logger.debug(f"image_gen_model_name: {image_gen_model_name}")
verbose_logger.debug(
f"image_gen_model_name_with_quality: {image_gen_model_name_with_quality}"
)
if image_gen_model_name in litellm.model_cost:
return (
litellm.model_cost[image_gen_model_name]["input_cost_per_pixel"]
* height
* width
* n
)
elif image_gen_model_name_with_quality in litellm.model_cost:
return (
litellm.model_cost[image_gen_model_name_with_quality][
"input_cost_per_pixel"
]
* height
* width
* n
)
else:
raise Exception(f"Model={model} not found in completion cost model map")
# Calculate cost based on prompt_tokens, completion_tokens
if "togethercomputer" in model or "together_ai" in model:
# together ai prices based on size of llm
@ -3096,8 +3164,13 @@ def completion_cost(
completion_tokens=completion_tokens,
custom_llm_provider=custom_llm_provider,
response_time_ms=total_time,
region_name=region_name,
)
return prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
_final_cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
print_verbose(
f"final cost: {_final_cost}; prompt_tokens_cost_usd_dollar: {prompt_tokens_cost_usd_dollar}; completion_tokens_cost_usd_dollar: {completion_tokens_cost_usd_dollar}"
)
return _final_cost
except Exception as e:
raise e
@ -3263,8 +3336,10 @@ def get_optional_params_image_gen(
def get_optional_params_embeddings(
# 2 optional params
model=None,
user=None,
encoding_format=None,
dimensions=None,
custom_llm_provider="",
**kwargs,
):
@ -3275,7 +3350,7 @@ def get_optional_params_embeddings(
for k, v in special_params.items():
passed_params[k] = v
default_params = {"user": None, "encoding_format": None}
default_params = {"user": None, "encoding_format": None, "dimensions": None}
non_default_params = {
k: v
@ -3283,6 +3358,19 @@ def get_optional_params_embeddings(
if (k in default_params and v != default_params[k])
}
## raise exception if non-default value passed for non-openai/azure embedding calls
if custom_llm_provider == "openai":
# 'dimensions` is only supported in `text-embedding-3` and later models
if (
model is not None
and "text-embedding-3" not in model
and "dimensions" in non_default_params.keys()
):
raise UnsupportedParamsError(
status_code=500,
message=f"Setting dimensions is not supported for OpenAI `text-embedding-3` and later models. To drop it from the call, set `litellm.drop_params = True`.",
)
if (
custom_llm_provider != "openai"
and custom_llm_provider != "azure"
@ -3337,6 +3425,10 @@ def get_optional_params(
custom_llm_provider != "bedrock" and custom_llm_provider != "sagemaker"
): # allow dynamically setting boto3 init logic
continue
elif (
k.startswith("vertex_") and custom_llm_provider != "vertex_ai"
): # allow dynamically setting vertex ai init logic
continue
passed_params[k] = v
default_params = {
"functions": None,
@ -7102,6 +7194,7 @@ class CustomStreamWrapper:
"model_id": (_model_info.get("id", None))
} # returned as x-litellm-model-id response header in proxy
self.response_id = None
self.logging_loop = None
def __iter__(self):
return self
@ -7672,6 +7765,27 @@ class CustomStreamWrapper:
}
return ""
def handle_sagemaker_stream(self, chunk):
if "data: [DONE]" in chunk:
text = ""
is_finished = True
finish_reason = "stop"
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
elif isinstance(chunk, dict):
if chunk["is_finished"] == True:
finish_reason = "stop"
else:
finish_reason = ""
return {
"text": chunk["text"],
"is_finished": chunk["is_finished"],
"finish_reason": finish_reason,
}
def chunk_creator(self, chunk):
model_response = ModelResponse(stream=True, model=self.model)
if self.response_id is not None:
@ -7679,6 +7793,7 @@ class CustomStreamWrapper:
else:
self.response_id = model_response.id
model_response._hidden_params["custom_llm_provider"] = self.custom_llm_provider
model_response._hidden_params["created_at"] = time.time()
model_response.choices = [StreamingChoices()]
model_response.choices[0].finish_reason = None
response_obj = {}
@ -7797,8 +7912,14 @@ class CustomStreamWrapper:
]
self.sent_last_chunk = True
elif self.custom_llm_provider == "sagemaker":
print_verbose(f"ENTERS SAGEMAKER STREAMING for chunk {chunk}")
completion_obj["content"] = chunk
verbose_logger.debug(f"ENTERS SAGEMAKER STREAMING for chunk {chunk}")
response_obj = self.handle_sagemaker_stream(chunk)
completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]:
model_response.choices[0].finish_reason = response_obj[
"finish_reason"
]
self.sent_last_chunk = True
elif self.custom_llm_provider == "petals":
if len(self.completion_stream) == 0:
if self.sent_last_chunk:
@ -7974,6 +8095,27 @@ class CustomStreamWrapper:
original_exception=e,
)
def set_logging_event_loop(self, loop):
self.logging_loop = loop
async def your_async_function(self):
# Your asynchronous code here
return "Your asynchronous code is running"
def run_success_logging_in_thread(self, processed_chunk):
# Create an event loop for the new thread
## ASYNC LOGGING
if self.logging_loop is not None:
future = asyncio.run_coroutine_threadsafe(
self.logging_obj.async_success_handler(processed_chunk),
loop=self.logging_loop,
)
result = future.result()
else:
asyncio.run(self.logging_obj.async_success_handler(processed_chunk))
## SYNC LOGGING
self.logging_obj.success_handler(processed_chunk)
## needs to handle the empty string case (even starting chunk can be an empty string)
def __next__(self):
try:
@ -7992,8 +8134,9 @@ class CustomStreamWrapper:
continue
## LOGGING
threading.Thread(
target=self.logging_obj.success_handler, args=(response,)
target=self.run_success_logging_in_thread, args=(response,)
).start() # log response
# RETURN RESULT
return response
except StopIteration:
@ -8049,13 +8192,34 @@ class CustomStreamWrapper:
raise StopAsyncIteration
else: # temporary patch for non-aiohttp async calls
# example - boto3 bedrock llms
processed_chunk = next(self)
asyncio.create_task(
self.logging_obj.async_success_handler(
processed_chunk,
)
)
return processed_chunk
while True:
if isinstance(self.completion_stream, str) or isinstance(
self.completion_stream, bytes
):
chunk = self.completion_stream
else:
chunk = next(self.completion_stream)
if chunk is not None and chunk != b"":
print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}")
processed_chunk = self.chunk_creator(chunk=chunk)
print_verbose(
f"PROCESSED CHUNK POST CHUNK CREATOR: {processed_chunk}"
)
if processed_chunk is None:
continue
## LOGGING
threading.Thread(
target=self.logging_obj.success_handler,
args=(processed_chunk,),
).start() # log processed_chunk
asyncio.create_task(
self.logging_obj.async_success_handler(
processed_chunk,
)
)
# RETURN RESULT
return processed_chunk
except StopAsyncIteration:
raise
except StopIteration:

23
model_prices_and_context_window.json
View File

@ -62,6 +62,15 @@
"litellm_provider": "openai",
"mode": "chat"
},
"gpt-4-0125-preview": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 4096,
"input_cost_per_token": 0.00001,
"output_cost_per_token": 0.00003,
"litellm_provider": "openai",
"mode": "chat"
},
"gpt-4-vision-preview": {
"max_tokens": 128000,
"max_input_tokens": 128000,
@ -143,6 +152,20 @@
"litellm_provider": "openai",
"mode": "chat"
},
"text-embedding-3-large": {
"max_tokens": 8191,
"input_cost_per_token": 0.00000013,
"output_cost_per_token": 0.000000,
"litellm_provider": "openai",
"mode": "embedding"
},
"text-embedding-3-small": {
"max_tokens": 8191,
"input_cost_per_token": 0.00000002,
"output_cost_per_token": 0.000000,
"litellm_provider": "openai",
"mode": "embedding"
},
"text-embedding-ada-002": {
"max_tokens": 8191,
"input_cost_per_token": 0.0000001,

6
poetry.lock generated
View File

@ -1158,13 +1158,13 @@ files = [
[[package]]
name = "openai"
version = "1.8.0"
version = "1.10.0"
description = "The official Python library for the openai API"
optional = false
python-versions = ">=3.7.1"
files = [
{file = "openai-1.8.0-py3-none-any.whl", hash = "sha256:0f8f53805826103fdd8adaf379ad3ec23f9d867e698cbc14caf34b778d150175"},
{file = "openai-1.8.0.tar.gz", hash = "sha256:93366be27802f517e89328801913d2a5ede45e3b86fdcab420385b8a1b88c767"},
{file = "openai-1.10.0-py3-none-any.whl", hash = "sha256:aa69e97d0223ace9835fbf9c997abe9ee95318f684fd2de6d02c870700c71ebc"},
{file = "openai-1.10.0.tar.gz", hash = "sha256:208886cb501b930dc63f48d51db9c15e5380380f80516d07332adad67c9f1053"},
]
[package.dependencies]

6
proxy_server_config.yaml
View File

@ -11,6 +11,10 @@ model_list:
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
- model_name: sagemaker-completion-model
litellm_params:
model: sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4
input_cost_per_second: 0.000420
- model_name: gpt-4
litellm_params:
model: azure/gpt-turbo
@ -41,6 +45,8 @@ model_list:
litellm_settings:
drop_params: True
max_budget: 100
budget_duration: 30d
general_settings:
master_key: sk-1234 # [OPTIONAL] Only use this if you to require all calls to contain this key (Authorization: Bearer sk-1234)
# database_url: "postgresql://<user>:<password>@<host>:<port>/<dbname>" # [OPTIONAL] use for token-based auth to proxy

4
pyproject.toml
View File

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.19.0"
version = "1.20.2"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@ -63,7 +63,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.19.0"
version = "1.20.2"
version_files = [
"pyproject.toml:^version"
]

13
schema.prisma
View File

@ -7,6 +7,7 @@ generator client {
provider = "prisma-client-py"
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @unique
team_id String?
@ -21,9 +22,11 @@ model LiteLLM_UserTable {
budget_reset_at DateTime?
}
// required for token gen
// Generate Tokens for Proxy
model LiteLLM_VerificationToken {
token String @unique
key_name String?
key_alias String?
spend Float @default(0.0)
expires DateTime?
models String[]
@ -40,22 +43,26 @@ model LiteLLM_VerificationToken {
budget_reset_at DateTime?
}
// store proxy config.yaml
model LiteLLM_Config {
param_name String @id
param_value Json?
}
// View spend, model, api_key per request
model LiteLLM_SpendLogs {
request_id String @unique
call_type String
api_key String @default ("")
spend Float @default(0.0)
total_tokens Int @default(0)
prompt_tokens Int @default(0)
completion_tokens Int @default(0)
startTime DateTime // Assuming start_time is a DateTime field
endTime DateTime // Assuming end_time is a DateTime field
model String @default("")
user String @default("")
modelParameters Json @default("{}")// Assuming optional_params is a JSON field
usage Json @default("{}")
metadata Json @default("{}")
cache_hit String @default("")
cache_key String @default("")
}

124
tests/test_keys.py
View File

@ -13,17 +13,21 @@ sys.path.insert(
import litellm
async def generate_key(session, i, budget=None, budget_duration=None):
async def generate_key(
session, i, budget=None, budget_duration=None, models=["azure-models", "gpt-4"]
):
url = "http://0.0.0.0:4000/key/generate"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {
"models": ["azure-models", "gpt-4"],
"models": models,
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
"duration": None,
"max_budget": budget,
"budget_duration": budget_duration,
}
print(f"data: {data}")
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
@ -67,6 +71,28 @@ async def update_key(session, get_key):
return await response.json()
async def update_proxy_budget(session):
"""
Make sure only models user has access to are returned
"""
url = "http://0.0.0.0:4000/user/update"
headers = {
"Authorization": f"Bearer sk-1234",
"Content-Type": "application/json",
}
data = {"user_id": "litellm-proxy-budget", "spend": 0}
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
print(response_text)
print()
if status != 200:
raise Exception(f"Request did not return a 200 status code: {status}")
return await response.json()
async def chat_completion(session, key, model="gpt-4"):
url = "http://0.0.0.0:4000/chat/completions"
headers = {
@ -89,7 +115,9 @@ async def chat_completion(session, key, model="gpt-4"):
print()
if status != 200:
raise Exception(f"Request did not return a 200 status code: {status}")
raise Exception(
f"Request did not return a 200 status code: {status}. Response: {response_text}"
)
return await response.json()
@ -135,6 +163,7 @@ async def test_key_update():
session=session,
get_key=key,
)
await update_proxy_budget(session=session) # resets proxy spend
await chat_completion(session=session, key=key)
@ -174,11 +203,14 @@ async def test_key_delete():
)
async def get_key_info(session, get_key, call_key):
async def get_key_info(session, call_key, get_key=None):
"""
Make sure only models user has access to are returned
"""
url = f"http://0.0.0.0:4000/key/info?key={get_key}"
if get_key is None:
url = "http://0.0.0.0:4000/key/info"
else:
url = f"http://0.0.0.0:4000/key/info?key={get_key}"
headers = {
"Authorization": f"Bearer {call_key}",
"Content-Type": "application/json",
@ -214,6 +246,9 @@ async def test_key_info():
await get_key_info(session=session, get_key=key, call_key="sk-1234")
# as key itself #
await get_key_info(session=session, get_key=key, call_key=key)
# as key itself, use the auth param, and no query key needed
await get_key_info(session=session, call_key=key)
# as random key #
key_gen = await generate_key(session=session, i=0)
random_key = key_gen["key"]
@ -254,14 +289,20 @@ async def test_key_info_spend_values():
await asyncio.sleep(5)
spend_logs = await get_spend_logs(session=session, request_id=response["id"])
print(f"spend_logs: {spend_logs}")
usage = spend_logs[0]["usage"]
completion_tokens = spend_logs[0]["completion_tokens"]
prompt_tokens = spend_logs[0]["prompt_tokens"]
print(f"prompt_tokens: {prompt_tokens}; completion_tokens: {completion_tokens}")
litellm.set_verbose = True
prompt_cost, completion_cost = litellm.cost_per_token(
model="gpt-35-turbo",
prompt_tokens=usage["prompt_tokens"],
completion_tokens=usage["completion_tokens"],
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
custom_llm_provider="azure",
)
print("prompt_cost: ", prompt_cost, "completion_cost: ", completion_cost)
response_cost = prompt_cost + completion_cost
print(f"response_cost: {response_cost}")
await asyncio.sleep(5) # allow db log to be updated
key_info = await get_key_info(session=session, get_key=key, call_key=key)
print(
@ -270,7 +311,7 @@ async def test_key_info_spend_values():
rounded_response_cost = round(response_cost, 8)
rounded_key_info_spend = round(key_info["info"]["spend"], 8)
assert rounded_response_cost == rounded_key_info_spend
## streaming
## streaming - azure
key_gen = await generate_key(session=session, i=0)
new_key = key_gen["key"]
prompt_tokens, completion_tokens = await chat_completion_streaming(
@ -295,6 +336,41 @@ async def test_key_info_spend_values():
assert rounded_response_cost == rounded_key_info_spend
@pytest.mark.asyncio
async def test_key_info_spend_values_sagemaker():
"""
Tests the sync streaming loop to ensure spend is correctly calculated.
- create key
- make completion call
- assert cost is expected value
"""
async with aiohttp.ClientSession() as session:
## streaming - sagemaker
key_gen = await generate_key(session=session, i=0, models=[])
new_key = key_gen["key"]
prompt_tokens, completion_tokens = await chat_completion_streaming(
session=session, key=new_key, model="sagemaker-completion-model"
)
# print(f"prompt_tokens: {prompt_tokens}, completion_tokens: {completion_tokens}")
# prompt_cost, completion_cost = litellm.cost_per_token(
# model="azure/gpt-35-turbo",
# prompt_tokens=prompt_tokens,
# completion_tokens=completion_tokens,
# )
# response_cost = prompt_cost + completion_cost
await asyncio.sleep(5) # allow db log to be updated
key_info = await get_key_info(
session=session, get_key=new_key, call_key=new_key
)
# print(
# f"response_cost: {response_cost}; key_info spend: {key_info['info']['spend']}"
# )
# rounded_response_cost = round(response_cost, 8)
rounded_key_info_spend = round(key_info["info"]["spend"], 8)
assert rounded_key_info_spend > 0
# assert rounded_response_cost == rounded_key_info_spend
@pytest.mark.asyncio
async def test_key_with_budgets():
"""
@ -314,7 +390,35 @@ async def test_key_with_budgets():
print(f"hashed_token: {hashed_token}")
key_info = await get_key_info(session=session, get_key=key, call_key=key)
reset_at_init_value = key_info["info"]["budget_reset_at"]
await asyncio.sleep(15)
await asyncio.sleep(30)
key_info = await get_key_info(session=session, get_key=key, call_key=key)
reset_at_new_value = key_info["info"]["budget_reset_at"]
assert reset_at_init_value != reset_at_new_value
@pytest.mark.asyncio
async def test_key_crossing_budget():
"""
- Create key with budget with budget=0.00000001
- make a /chat/completions call
- wait 5s
- make a /chat/completions call - should fail with key crossed it's budget
- Check if value updated
"""
from litellm.proxy.utils import hash_token
async with aiohttp.ClientSession() as session:
key_gen = await generate_key(session=session, i=0, budget=0.0000001)
key = key_gen["key"]
hashed_token = hash_token(token=key)
print(f"hashed_token: {hashed_token}")
response = await chat_completion(session=session, key=key)
print("response 1: ", response)
await asyncio.sleep(2)
try:
response = await chat_completion(session=session, key=key)
pytest.fail("Should have failed - Key crossed it's budget")
except Exception as e:
assert "ExceededTokenBudget: Current spend for token:" in str(e)

63
tests/test_users.py
View File

@ -4,6 +4,7 @@ import pytest
import asyncio
import aiohttp
import time
from openai import AsyncOpenAI
async def new_user(session, i, user_id=None, budget=None, budget_duration=None):
@ -105,7 +106,7 @@ async def test_user_update():
@pytest.mark.asyncio
async def test_users_with_budgets():
async def test_users_budgets_reset():
"""
- Create key with budget and 5s duration
- Get 'reset_at' value
@ -128,3 +129,63 @@ async def test_users_with_budgets():
)
reset_at_new_value = user_info["user_info"]["budget_reset_at"]
assert reset_at_init_value != reset_at_new_value
async def chat_completion(session, key, model="gpt-4"):
client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000")
messages = [
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": f"Hello! {time.time()}"},
]
data = {
"model": model,
"messages": messages,
}
response = await client.chat.completions.create(**data)
async def chat_completion_streaming(session, key, model="gpt-4"):
client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000")
messages = [
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": f"Hello! {time.time()}"},
]
data = {"model": model, "messages": messages, "stream": True}
response = await client.chat.completions.create(**data)
async for chunk in response:
continue
@pytest.mark.asyncio
async def test_global_proxy_budget_update():
"""
- Get proxy current spend
- Make chat completion call (normal)
- Assert spend increased
- Make chat completion call (streaming)
- Assert spend increased
"""
get_user = f"litellm-proxy-budget"
async with aiohttp.ClientSession() as session:
user_info = await get_user_info(
session=session, get_user=get_user, call_user="sk-1234"
)
original_spend = user_info["user_info"]["spend"]
await chat_completion(session=session, key="sk-1234")
await asyncio.sleep(5) # let db update
user_info = await get_user_info(
session=session, get_user=get_user, call_user="sk-1234"
)
new_spend = user_info["user_info"]["spend"]
print(f"new_spend: {new_spend}; original_spend: {original_spend}")
assert new_spend > original_spend
await chat_completion_streaming(session=session, key="sk-1234")
await asyncio.sleep(5) # let db update
user_info = await get_user_info(
session=session, get_user=get_user, call_user="sk-1234"
)
new_new_spend = user_info["user_info"]["spend"]
print(f"new_spend: {new_spend}; original_spend: {original_spend}")
assert new_new_spend > new_spend

9
ui/Dockerfile
View File

@ -1,5 +1,7 @@
# Use official Python base image
FROM python:3.9
FROM python:3.9.12
EXPOSE 8501
# Set the working directory in the container
WORKDIR /app
@ -11,7 +13,8 @@ COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Copy the entire project directory to the container
COPY . .
COPY admin.py .
# Set the entrypoint command to run admin.py with Streamlit
ENTRYPOINT ["streamlit", "run", "admin.py"]
ENTRYPOINT [ "streamlit", "run"]
CMD ["admin.py"]

3
ui/litellm-dashboard/.eslintrc.json Normal file
View File

@ -0,0 +1,3 @@
{
"extends": "next/core-web-vitals"
}

36
ui/litellm-dashboard/README.md Normal file
View File

@ -0,0 +1,36 @@
This is a [Next.js](https://nextjs.org/) project bootstrapped with [`create-next-app`](https://github.com/vercel/next.js/tree/canary/packages/create-next-app).
## Getting Started
First, run the development server:
```bash
npm run dev
# or
yarn dev
# or
pnpm dev
# or
bun dev
```
Open [http://localhost:3000](http://localhost:3000) with your browser to see the result.
You can start editing the page by modifying `app/page.tsx`. The page auto-updates as you edit the file.
This project uses [`next/font`](https://nextjs.org/docs/basic-features/font-optimization) to automatically optimize and load Inter, a custom Google Font.
## Learn More
To learn more about Next.js, take a look at the following resources:
- [Next.js Documentation](https://nextjs.org/docs) - learn about Next.js features and API.
- [Learn Next.js](https://nextjs.org/learn) - an interactive Next.js tutorial.
You can check out [the Next.js GitHub repository](https://github.com/vercel/next.js/) - your feedback and contributions are welcome!
## Deploy on Vercel
The easiest way to deploy your Next.js app is to use the [Vercel Platform](https://vercel.com/new?utm_medium=default-template&filter=next.js&utm_source=create-next-app&utm_campaign=create-next-app-readme) from the creators of Next.js.
Check out our [Next.js deployment documentation](https://nextjs.org/docs/deployment) for more details.

4
ui/litellm-dashboard/next.config.mjs Normal file
View File

@ -0,0 +1,4 @@
/** @type {import('next').NextConfig} */
const nextConfig = {};
export default nextConfig;

6071
ui/litellm-dashboard/package-lock.json generated Normal file
View File

File diff suppressed because it is too large Load Diff

34
ui/litellm-dashboard/package.json Normal file
View File

@ -0,0 +1,34 @@
{
"name": "litellm-dashboard",
"version": "0.1.0",
"private": true,
"scripts": {
"dev": "next dev",
"build": "next build",
"start": "next start",
"lint": "next lint"
},
"dependencies": {
"@headlessui/react": "^1.7.18",
"@headlessui/tailwindcss": "^0.2.0",
"@heroicons/react": "^1.0.6",
"@remixicon/react": "^4.1.1",
"@tremor/react": "^3.13.3",
"antd": "^5.13.2",
"next": "14.1.0",
"react": "^18",
"react-dom": "^18"
},
"devDependencies": {
"@tailwindcss/forms": "^0.5.7",
"@types/node": "^20",
"@types/react": "18.2.48",
"@types/react-dom": "^18",
"autoprefixer": "^10.4.17",
"eslint": "^8",
"eslint-config-next": "14.1.0",
"postcss": "^8.4.33",
"tailwindcss": "^3.4.1",
"typescript": "5.3.3"
}
}

6
ui/litellm-dashboard/postcss.config.js Normal file
View File

@ -0,0 +1,6 @@
module.exports = {
plugins: {
tailwindcss: {},
autoprefixer: {},
},
};

1
ui/litellm-dashboard/public/next.svg Normal file
View File

@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 394 80"><path fill="#000" d="M262 0h68.5v12.7h-27.2v66.6h-13.6V12.7H262V0ZM149 0v12.7H94v20.4h44.3v12.6H94v21h55v12.6H80.5V0h68.7zm34.3 0h-17.8l63.8 79.4h17.9l-32-39.7 32-39.6h-17.9l-23 28.6-23-28.6zm18.3 56.7-9-11-27.1 33.7h17.8l18.3-22.7z"/><path fill="#000" d="M81 79.3 17 0H0v79.3h13.6V17l50.2 62.3H81Zm252.6-.4c-1 0-1.8-.4-2.5-1s-1.1-1.6-1.1-2.6.3-1.8 1-2.5 1.6-1 2.6-1 1.8.3 2.5 1a3.4 3.4 0 0 1 .6 4.3 3.7 3.7 0 0 1-3 1.8zm23.2-33.5h6v23.3c0 2.1-.4 4-1.3 5.5a9.1 9.1 0 0 1-3.8 3.5c-1.6.8-3.5 1.3-5.7 1.3-2 0-3.7-.4-5.3-1s-2.8-1.8-3.7-3.2c-.9-1.3-1.4-3-1.4-5h6c.1.8.3 1.6.7 2.2s1 1.2 1.6 1.5c.7.4 1.5.5 2.4.5 1 0 1.8-.2 2.4-.6a4 4 0 0 0 1.6-1.8c.3-.8.5-1.8.5-3V45.5zm30.9 9.1a4.4 4.4 0 0 0-2-3.3 7.5 7.5 0 0 0-4.3-1.1c-1.3 0-2.4.2-3.3.5-.9.4-1.6 1-2 1.6a3.5 3.5 0 0 0-.3 4c.3.5.7.9 1.3 1.2l1.8 1 2 .5 3.2.8c1.3.3 2.5.7 3.7 1.2a13 13 0 0 1 3.2 1.8 8.1 8.1 0 0 1 3 6.5c0 2-.5 3.7-1.5 5.1a10 10 0 0 1-4.4 3.5c-1.8.8-4.1 1.2-6.8 1.2-2.6 0-4.9-.4-6.8-1.2-2-.8-3.4-2-4.5-3.5a10 10 0 0 1-1.7-5.6h6a5 5 0 0 0 3.5 4.6c1 .4 2.2.6 3.4.6 1.3 0 2.5-.2 3.5-.6 1-.4 1.8-1 2.4-1.7a4 4 0 0 0 .8-2.4c0-.9-.2-1.6-.7-2.2a11 11 0 0 0-2.1-1.4l-3.2-1-3.8-1c-2.8-.7-5-1.7-6.6-3.2a7.2 7.2 0 0 1-2.4-5.7 8 8 0 0 1 1.7-5 10 10 0 0 1 4.3-3.5c2-.8 4-1.2 6.4-1.2 2.3 0 4.4.4 6.2 1.2 1.8.8 3.2 2 4.3 3.4 1 1.4 1.5 3 1.5 5h-5.8z"/></svg>
Side by Side

After

Width:  |  Height:  |  Size: 1.3 KiB

1
ui/litellm-dashboard/public/vercel.svg Normal file
View File

@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 283 64"><path fill="black" d="M141 16c-11 0-19 7-19 18s9 18 20 18c7 0 13-3 16-7l-7-5c-2 3-6 4-9 4-5 0-9-3-10-7h28v-3c0-11-8-18-19-18zm-9 15c1-4 4-7 9-7s8 3 9 7h-18zm117-15c-11 0-19 7-19 18s9 18 20 18c6 0 12-3 16-7l-8-5c-2 3-5 4-8 4-5 0-9-3-11-7h28l1-3c0-11-8-18-19-18zm-10 15c2-4 5-7 10-7s8 3 9 7h-19zm-39 3c0 6 4 10 10 10 4 0 7-2 9-5l8 5c-3 5-9 8-17 8-11 0-19-7-19-18s8-18 19-18c8 0 14 3 17 8l-8 5c-2-3-5-5-9-5-6 0-10 4-10 10zm83-29v46h-9V5h9zM37 0l37 64H0L37 0zm92 5-27 48L74 5h10l18 30 17-30h10zm59 12v10l-3-1c-6 0-10 4-10 10v15h-9V17h9v9c0-5 6-9 13-9z"/></svg>
Side by Side

After

Width:  |  Height:  |  Size: 629 B

BIN
ui/litellm-dashboard/src/app/favicon.ico Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 25 KiB

33
ui/litellm-dashboard/src/app/globals.css Normal file
View File

@ -0,0 +1,33 @@
@tailwind base;
@tailwind components;
@tailwind utilities;
:root {
--foreground-rgb: 0, 0, 0;
--background-start-rgb: 214, 219, 220;
--background-end-rgb: 255, 255, 255;
}
/* @media (prefers-color-scheme: dark) {
:root {
--foreground-rgb: 255, 255, 255;
--background-start-rgb: 0, 0, 0;
--background-end-rgb: 0, 0, 0;
}
} */
body {
color: rgb(var(--foreground-rgb));
background: linear-gradient(
to bottom,
transparent,
rgb(var(--background-end-rgb))
)
rgb(var(--background-start-rgb));
}
@layer utilities {
.text-balance {
text-wrap: balance;
}
}

22
ui/litellm-dashboard/src/app/layout.tsx Normal file
View File

@ -0,0 +1,22 @@
import type { Metadata } from "next";
import { Inter } from "next/font/google";
import "./globals.css";
const inter = Inter({ subsets: ["latin"] });
export const metadata: Metadata = {
title: "Create Next App",
description: "Generated by create next app",
};
export default function RootLayout({
children,
}: Readonly<{
children: React.ReactNode;
}>) {
return (
<html lang="en">
<body className={inter.className}>{children}</body>
</html>
);
}

15
ui/litellm-dashboard/src/app/page.tsx Normal file
View File

@ -0,0 +1,15 @@
import React, { Suspense } from "react";
import Navbar from "../components/navbar";
import UserDashboard from "../components/user_dashboard";
const CreateKeyPage = () => {
return (
<Suspense fallback={<div>Loading...</div>}>
<div className="flex min-h-screen flex-col items-center">
<Navbar />
<UserDashboard />
</div>
</Suspense>
);
};
export default CreateKeyPage;

92
ui/litellm-dashboard/src/components/create_key_button.tsx Normal file
View File

@ -0,0 +1,92 @@
"use client";
import React, { useState, useEffect, useRef } from "react";
import { Button, TextInput, Grid, Col } from "@tremor/react";
import { message } from "antd";
import { Card, Metric, Text } from "@tremor/react";
import { keyCreateCall } from "./networking";
// Define the props type
interface CreateKeyProps {
userID: string;
accessToken: string;
proxyBaseUrl: string;
data: any[] | null;
setData: React.Dispatch<React.SetStateAction<any[] | null>>;
}
import { Modal, Button as Button2 } from "antd";
const CreateKey: React.FC<CreateKeyProps> = ({
userID,
accessToken,
proxyBaseUrl,
data,
setData,
}) => {
const [isModalVisible, setIsModalVisible] = useState(false);
const [apiKey, setApiKey] = useState(null);
const handleOk = () => {
// Handle the OK action
console.log("OK Clicked");
setIsModalVisible(false);
};
const handleCancel = () => {
// Handle the cancel action or closing the modal
console.log("Modal closed");
setIsModalVisible(false);
setApiKey(null);
};
const handleCreate = async () => {
if (data == null) {
return;
}
try {
message.info("Making API Call");
setIsModalVisible(true);
const response = await keyCreateCall(proxyBaseUrl, accessToken, userID);
// Successfully completed the deletion. Update the state to trigger a rerender.
setData([...data, response]);
setApiKey(response["key"]);
message.success("API Key Created");
} catch (error) {
console.error("Error deleting the key:", error);
// Handle any error situations, such as displaying an error message to the user.
}
};
return (
<div>
<Button className="mx-auto" onClick={handleCreate}>
+ Create New Key
</Button>
<Modal
title="Save your key"
open={isModalVisible}
onOk={handleOk}
onCancel={handleCancel}
>
<Grid numItems={1} className="gap-2 w-full">
<Col numColSpan={1}>
<p>
Please save this secret key somewhere safe and accessible. For
security reasons, <b>you will not be able to view it again</b>{" "}
through your LiteLLM account. If you lose this secret key, you will
need to generate a new one.
</p>
</Col>
<Col numColSpan={1}>
{apiKey != null ? (
<Text>API Key: {apiKey}</Text>
) : (
<Text>Key being created, this might take 30s</Text>
)}
</Col>
</Grid>
</Modal>
</div>
);
};
export default CreateKey;

64
ui/litellm-dashboard/src/components/enter_proxy_url.tsx Normal file
View File

@ -0,0 +1,64 @@
"use client";
import React, { useState, ChangeEvent } from "react";
import { Button, Col, Grid, TextInput } from "@tremor/react";
import { Card, Text } from "@tremor/react";
const EnterProxyUrl: React.FC = () => {
const [proxyUrl, setProxyUrl] = useState<string>("");
const [isUrlSaved, setIsUrlSaved] = useState<boolean>(false);
const handleUrlChange = (event: ChangeEvent<HTMLInputElement>) => {
setProxyUrl(event.target.value);
// Reset the saved status when the URL changes
setIsUrlSaved(false);
};
const handleSaveClick = () => {
// You can perform any additional validation or actions here
// For now, let's just display the message
setIsUrlSaved(true);
};
// Construct the URL for clicking
const clickableUrl = `${window.location.href}?proxyBaseUrl=${proxyUrl}`;
return (
<div>
<Card decoration="top" decorationColor="blue" style={{ width: '1000px' }}>
<Text>Admin Configuration</Text>
<label htmlFor="proxyUrl">Enter Proxy URL:</label>
<TextInput
type="text"
id="proxyUrl"
value={proxyUrl}
onChange={handleUrlChange}
placeholder="https://your-proxy-endpoint.com"
/>
<Button onClick={handleSaveClick} className="gap-2">Save</Button>
{/* Display message if the URL is saved */}
{isUrlSaved && (
<div>
<Grid numItems={1} className="gap-2">
<Col>
<p>
Proxy Admin UI (Save this URL): {clickableUrl}
</p>
</Col>
<Col>
<p>
Get Started with Proxy Admin UI 👉
<a href={clickableUrl} target="_blank" rel="noopener noreferrer" style={{ color: 'blue', textDecoration: 'underline' }}>
{clickableUrl}
</a>
</p>
</Col>
</Grid>
</div>
)}
</Card>
</div>
);
};
export default EnterProxyUrl;

26
ui/litellm-dashboard/src/components/navbar.tsx Normal file
View File

@ -0,0 +1,26 @@
import Link from 'next/link';
import Image from 'next/image'
import React, { useState } from 'react';
function Navbar() {
return (
<nav className="left-0 right-0 top-0 flex justify-between items-center h-12">
<div className="text-left mx-4 my-2 absolute top-0 left-0">
<div className="flex flex-col items-center">
<Link href="/">
<button className="text-gray-800 text-2xl px-4 py-1 rounded text-center">🚅 LiteLLM</button>
</Link>
</div>
</div>
<div className="text-right mx-4 my-2 absolute top-0 right-0">
<a href="https://docs.litellm.ai/docs/" target="_blank" rel="noopener noreferrer">
<button className="border border-gray-800 rounded-lg text-gray-800 text-xl px-4 py-1 rounded p-1 mr-2 text-center">Docs</button>
</a>
<a href="https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version" target="_blank" rel="noopener noreferrer">
<button className="border border-gray-800 rounded-lg text-gray-800 text-xl px-4 py-1 rounded p-1 text-center">Schedule Demo</button>
</a>
</div>
</nav>
)
}
export default Navbar;

98
ui/litellm-dashboard/src/components/networking.tsx Normal file
View File

@ -0,0 +1,98 @@
/**
* Helper file for calls being made to proxy
*/
export const keyCreateCall = async (
proxyBaseUrl: String,
accessToken: String,
userID: String
) => {
try {
const response = await fetch(`${proxyBaseUrl}/key/generate`, {
method: "POST",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
team_id: "core-infra-4",
max_budget: 10,
user_id: userID,
}),
});
if (!response.ok) {
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log(data);
return data;
// Handle success - you might want to update some state or UI based on the created key
} catch (error) {
console.error("Failed to create key:", error);
throw error;
}
};
export const keyDeleteCall = async (
proxyBaseUrl: String,
accessToken: String,
user_key: String
) => {
try {
const response = await fetch(`${proxyBaseUrl}/key/delete`, {
method: "POST",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
keys: [user_key],
}),
});
if (!response.ok) {
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log(data);
return data;
// Handle success - you might want to update some state or UI based on the created key
} catch (error) {
console.error("Failed to create key:", error);
throw error;
}
};
export const userInfoCall = async (
proxyBaseUrl: String,
accessToken: String,
userID: String
) => {
try {
const response = await fetch(
`${proxyBaseUrl}/user/info?user_id=${userID}`,
{
method: "GET",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
}
);
if (!response.ok) {
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log(data);
return data;
// Handle success - you might want to update some state or UI based on the created key
} catch (error) {
console.error("Failed to create key:", error);
throw error;
}
};

80
ui/litellm-dashboard/src/components/user_dashboard.tsx Normal file
View File

@ -0,0 +1,80 @@
"use client";
import React, { useState, useEffect } from "react";
import { userInfoCall } from "./networking";
import { Grid, Col, Card, Text } from "@tremor/react";
import CreateKey from "./create_key_button";
import ViewKeyTable from "./view_key_table";
import EnterProxyUrl from "./enter_proxy_url";
import { useSearchParams } from "next/navigation";
const UserDashboard = () => {
const [data, setData] = useState<null | any[]>(null); // Keep the initialization of state here
// Assuming useSearchParams() hook exists and works in your setup
const searchParams = useSearchParams();
const userID = searchParams.get("userID");
const accessToken = searchParams.get("accessToken");
const proxyBaseUrl = searchParams.get("proxyBaseUrl");
// Moved useEffect inside the component and used a condition to run fetch only if the params are available
useEffect(() => {
if (userID && accessToken && proxyBaseUrl && !data) {
const fetchData = async () => {
try {
const response = await userInfoCall(
proxyBaseUrl,
accessToken,
userID
);
setData(response["keys"]); // Assuming this is the correct path to your data
} catch (error) {
console.error("There was an error fetching the data", error);
// Optionally, update your UI to reflect the error state here as well
}
};
fetchData();
}
}, [userID, accessToken, proxyBaseUrl, data]);
if (proxyBaseUrl == null) {
return (
<div>
<EnterProxyUrl />
</div>
);
}
else if (userID == null || accessToken == null) {
// redirect to page: ProxyBaseUrl/google-login/key/generate
const baseUrl = proxyBaseUrl.endsWith('/') ? proxyBaseUrl : proxyBaseUrl + '/';
// Now you can construct the full URL
const url = `${baseUrl}google-login/key/generate`;
window.location.href = url;
return null;
}
return (
<Grid numItems={1} className="gap-0 p-10 h-[75vh] w-full">
<Col numColSpan={1}>
<ViewKeyTable
userID={userID}
accessToken={accessToken}
proxyBaseUrl={proxyBaseUrl}
data={data}
setData={setData}
/>
<CreateKey
userID={userID}
accessToken={accessToken}
proxyBaseUrl={proxyBaseUrl}
data={data}
setData={setData}
/>
</Col>
</Grid>
);
};
export default UserDashboard;

99
ui/litellm-dashboard/src/components/view_key_table.tsx Normal file
View File

@ -0,0 +1,99 @@
"use client";
import React, { useEffect } from "react";
import { keyDeleteCall } from "./networking";
import { StatusOnlineIcon, TrashIcon } from "@heroicons/react/outline";
import {
Badge,
Card,
Table,
TableBody,
TableCell,
TableHead,
TableHeaderCell,
TableRow,
Text,
Title,
Icon,
} from "@tremor/react";
// Define the props type
interface ViewKeyTableProps {
userID: string;
accessToken: string;
proxyBaseUrl: string;
data: any[] | null;
setData: React.Dispatch<React.SetStateAction<any[] | null>>;
}
const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
userID,
accessToken,
proxyBaseUrl,
data,
setData,
}) => {
const handleDelete = async (token: String) => {
if (data == null) {
return;
}
try {
await keyDeleteCall(proxyBaseUrl, accessToken, token);
// Successfully completed the deletion. Update the state to trigger a rerender.
const filteredData = data.filter((item) => item.token !== token);
setData(filteredData);
} catch (error) {
console.error("Error deleting the key:", error);
// Handle any error situations, such as displaying an error message to the user.
}
};
if (data == null) {
return;
}
console.log("RERENDER TRIGGERED");
return (
<Card className="w-full mx-auto flex-auto overflow-y-auto max-h-[50vh] mb-4">
<Title>API Keys</Title>
<Table className="mt-5">
<TableHead>
<TableRow>
<TableHeaderCell>Secret Key</TableHeaderCell>
<TableHeaderCell>Spend</TableHeaderCell>
<TableHeaderCell>Expires</TableHeaderCell>
</TableRow>
</TableHead>
<TableBody>
{data.map((item) => {
console.log(item);
return (
<TableRow key={item.token}>
<TableCell>
<Text>{item.key_name}</Text>
</TableCell>
<TableCell>
<Text>{item.spend}</Text>
</TableCell>
<TableCell>
{item.expires != null ? (
<Text>{item.expires}</Text>
) : (
<Text>Never expires</Text>
)}
</TableCell>
<TableCell>
<Icon
onClick={() => handleDelete(item.token)}
icon={TrashIcon}
size="xs"
/>
</TableCell>
</TableRow>
);
})}
</TableBody>
</Table>
</Card>
);
};
export default ViewKeyTable;

130
ui/litellm-dashboard/tailwind.config.js Normal file
View File

@ -0,0 +1,130 @@
/** @type {import('tailwindcss').Config} */
const colors = require("tailwindcss/colors");
module.exports = {
content: [
"./src/**/*.{js,ts,jsx,tsx}",
"./node_modules/@tremor/**/*.{js,ts,jsx,tsx}",
],
darkMode: "class",
theme: {
transparent: "transparent",
current: "currentColor",
extend: {
colors: {
// light mode
tremor: {
brand: {
faint: colors.blue[50],
muted: colors.blue[200],
subtle: colors.blue[400],
DEFAULT: colors.blue[500],
emphasis: colors.blue[700],
inverted: colors.white,
},
background: {
muted: colors.gray[50],
subtle: colors.gray[100],
DEFAULT: colors.white,
emphasis: colors.gray[700],
},
border: {
DEFAULT: colors.gray[200],
},
ring: {
DEFAULT: colors.gray[200],
},
content: {
subtle: colors.gray[400],
DEFAULT: colors.gray[500],
emphasis: colors.gray[700],
strong: colors.gray[900],
inverted: colors.white,
},
},
// dark mode
"dark-tremor": {
brand: {
faint: "#0B1229",
muted: colors.blue[950],
subtle: colors.blue[800],
DEFAULT: colors.blue[500],
emphasis: colors.blue[400],
inverted: colors.blue[950],
},
background: {
muted: "#131A2B",
subtle: colors.gray[800],
DEFAULT: colors.gray[900],
emphasis: colors.gray[300],
},
border: {
DEFAULT: colors.gray[700],
},
ring: {
DEFAULT: colors.gray[800],
},
content: {
subtle: colors.gray[600],
DEFAULT: colors.gray[500],
emphasis: colors.gray[200],
strong: colors.gray[50],
inverted: colors.gray[950],
},
},
},
boxShadow: {
// light
"tremor-input": "0 1px 2px 0 rgb(0 0 0 / 0.05)",
"tremor-card": "0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1)",
"tremor-dropdown": "0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1)",
// dark
"dark-tremor-input": "0 1px 2px 0 rgb(0 0 0 / 0.05)",
"dark-tremor-card": "0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1)",
"dark-tremor-dropdown": "0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1)",
},
borderRadius: {
"tremor-small": "0.375rem",
"tremor-default": "0.5rem",
"tremor-full": "9999px",
},
fontSize: {
"tremor-label": ["0.75rem", { lineHeight: "1rem" }],
"tremor-default": ["0.875rem", { lineHeight: "1.25rem" }],
"tremor-title": ["1.125rem", { lineHeight: "1.75rem" }],
"tremor-metric": ["1.875rem", { lineHeight: "2.25rem" }],
},
},
},
safelist: [
{
pattern:
/^(bg-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
variants: ["hover", "ui-selected"],
},
{
pattern:
/^(text-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
variants: ["hover", "ui-selected"],
},
{
pattern:
/^(border-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
variants: ["hover", "ui-selected"],
},
{
pattern:
/^(ring-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
},
{
pattern:
/^(stroke-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
},
{
pattern:
/^(fill-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
},
],
plugins: [require('@headlessui/tailwindcss'), require('@tailwindcss/forms')],
};

129
ui/litellm-dashboard/tailwind.config.ts Normal file
View File

@ -0,0 +1,129 @@
/** @type {import('tailwindcss').Config} */
const colors = require("tailwindcss/colors");
module.exports = {
content: [
"./src/**/*.{js,ts,jsx,tsx}",
"./node_modules/@tremor/**/*.{js,ts,jsx,tsx}",
],
theme: {
transparent: "transparent",
current: "currentColor",
extend: {
colors: {
// light mode
tremor: {
brand: {
faint: colors.blue[50],
muted: colors.blue[200],
subtle: colors.blue[400],
DEFAULT: colors.blue[500],
emphasis: colors.blue[700],
inverted: colors.white,
},
background: {
muted: colors.gray[50],
subtle: colors.gray[100],
DEFAULT: colors.white,
emphasis: colors.gray[700],
},
border: {
DEFAULT: colors.gray[200],
},
ring: {
DEFAULT: colors.gray[200],
},
content: {
subtle: colors.gray[400],
DEFAULT: colors.gray[500],
emphasis: colors.gray[700],
strong: colors.gray[900],
inverted: colors.white,
},
},
// dark mode
"dark-tremor": {
brand: {
faint: "#0B1229",
muted: colors.blue[950],
subtle: colors.blue[800],
DEFAULT: colors.blue[500],
emphasis: colors.blue[400],
inverted: colors.blue[950],
},
background: {
muted: "#131A2B",
subtle: colors.gray[800],
DEFAULT: colors.gray[900],
emphasis: colors.gray[300],
},
border: {
DEFAULT: colors.gray[700],
},
ring: {
DEFAULT: colors.gray[800],
},
content: {
subtle: colors.gray[600],
DEFAULT: colors.gray[500],
emphasis: colors.gray[200],
strong: colors.gray[50],
inverted: colors.gray[950],
},
},
},
boxShadow: {
// light
"tremor-input": "0 1px 2px 0 rgb(0 0 0 / 0.05)",
"tremor-card": "0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1)",
"tremor-dropdown": "0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1)",
// dark
"dark-tremor-input": "0 1px 2px 0 rgb(0 0 0 / 0.05)",
"dark-tremor-card": "0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1)",
"dark-tremor-dropdown": "0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1)",
},
borderRadius: {
"tremor-small": "0.375rem",
"tremor-default": "0.5rem",
"tremor-full": "9999px",
},
fontSize: {
"tremor-label": ["0.75rem", { lineHeight: "1rem" }],
"tremor-default": ["0.875rem", { lineHeight: "1.25rem" }],
"tremor-title": ["1.125rem", { lineHeight: "1.75rem" }],
"tremor-metric": ["1.875rem", { lineHeight: "2.25rem" }],
},
},
},
safelist: [
{
pattern:
/^(bg-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
variants: ["hover", "ui-selected"],
},
{
pattern:
/^(text-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
variants: ["hover", "ui-selected"],
},
{
pattern:
/^(border-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
variants: ["hover", "ui-selected"],
},
{
pattern:
/^(ring-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
},
{
pattern:
/^(stroke-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
},
{
pattern:
/^(fill-(?:slate|gray|zinc|neutral|stone|red|orange|amber|yellow|lime|green|emerald|teal|cyan|sky|blue|indigo|violet|purple|fuchsia|pink|rose)-(?:50|100|200|300|400|500|600|700|800|900|950))$/,
},
],
plugins: [require('@headlessui/tailwindcss'), require('@tailwindcss/forms')],
darkMode: "class"
};

26
ui/litellm-dashboard/tsconfig.json Normal file
View File

@ -0,0 +1,26 @@
{
"compilerOptions": {
"lib": ["dom", "dom.iterable", "esnext"],
"allowJs": true,
"skipLibCheck": true,
"strict": true,
"noEmit": true,
"esModuleInterop": true,
"module": "esnext",
"moduleResolution": "bundler",
"resolveJsonModule": true,
"isolatedModules": true,
"jsx": "preserve",
"incremental": true,
"plugins": [
{
"name": "next"
}
],
"paths": {
"@/*": ["./src/*"]
}
},
"include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],
"exclude": ["node_modules"]
}

1972
ui/package-lock.json generated Normal file
View File

File diff suppressed because it is too large Load Diff

10
ui/package.json Normal file
View File

@ -0,0 +1,10 @@
{
"dependencies": {
"@headlessui/react": "^1.7.18",
"@headlessui/tailwindcss": "^0.2.0",
"@tremor/react": "^3.13.3"
},
"devDependencies": {
"@tailwindcss/forms": "^0.5.7"
}
}