mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
988c37fda3
because that error can be thrown in several different places, so knowing the function it's being thrown from can be very useul for debugging.
9603 lines
367 KiB
Python
9603 lines
367 KiB
Python
import sys, os, platform, time, copy, re, asyncio, inspect
|
|
import threading, ast
|
|
import shutil, random, traceback, requests
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Optional, List, Callable
|
|
import secrets, subprocess
|
|
import hashlib, uuid
|
|
import warnings
|
|
import importlib
|
|
import warnings
|
|
|
|
|
|
def showwarning(message, category, filename, lineno, file=None, line=None):
|
|
traceback_info = f"{filename}:{lineno}: {category.__name__}: {message}\n"
|
|
if file is not None:
|
|
file.write(traceback_info)
|
|
|
|
|
|
warnings.showwarning = showwarning
|
|
warnings.filterwarnings("default", category=UserWarning)
|
|
|
|
# Your client code here
|
|
|
|
|
|
messages: list = []
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path - for litellm local dev
|
|
|
|
try:
|
|
import fastapi
|
|
import backoff
|
|
import yaml
|
|
import orjson
|
|
import logging
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
except ImportError as e:
|
|
raise ImportError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`")
|
|
|
|
import random
|
|
|
|
list_of_messages = [
|
|
"'The thing I wish you improved is...'",
|
|
"'A feature I really want is...'",
|
|
"'The worst thing about this product is...'",
|
|
"'This product would be better if...'",
|
|
"'I don't like how this works...'",
|
|
"'It would help me if you could add...'",
|
|
"'This feature doesn't meet my needs because...'",
|
|
"'I get frustrated when the product...'",
|
|
]
|
|
|
|
|
|
def generate_feedback_box():
|
|
box_width = 60
|
|
|
|
# Select a random message
|
|
message = random.choice(list_of_messages)
|
|
|
|
print() # noqa
|
|
print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa
|
|
print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa
|
|
print("\033[1;37m" + "# {:^59} #\033[0m".format(message)) # noqa
|
|
print( # noqa
|
|
"\033[1;37m"
|
|
+ "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new")
|
|
) # noqa
|
|
print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa
|
|
print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa
|
|
print() # noqa
|
|
print(" Thank you for using LiteLLM! - Krrish & Ishaan") # noqa
|
|
print() # noqa
|
|
print() # noqa
|
|
print() # noqa
|
|
print( # noqa
|
|
"\033[1;31mGive Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new\033[0m"
|
|
) # noqa
|
|
print() # noqa
|
|
print() # noqa
|
|
|
|
|
|
import litellm
|
|
from litellm.proxy.utils import (
|
|
PrismaClient,
|
|
DBClient,
|
|
get_instance_fn,
|
|
ProxyLogging,
|
|
_cache_user_row,
|
|
send_email,
|
|
get_logging_payload,
|
|
reset_budget,
|
|
hash_token,
|
|
html_form,
|
|
_read_request_body,
|
|
_is_valid_team_configs,
|
|
_is_user_proxy_admin,
|
|
_is_projected_spend_over_limit,
|
|
_get_projected_spend_over_limit,
|
|
update_spend,
|
|
encrypt_value,
|
|
decrypt_value,
|
|
)
|
|
from litellm.proxy.secret_managers.google_kms import load_google_kms
|
|
from litellm.proxy.secret_managers.aws_secret_manager import load_aws_secret_manager
|
|
import pydantic
|
|
from litellm.proxy._types import *
|
|
from litellm.caching import DualCache, RedisCache
|
|
from litellm.proxy.health_check import perform_health_check
|
|
from litellm.router import LiteLLM_Params, Deployment, updateDeployment
|
|
from litellm.router import ModelInfo as RouterModelInfo
|
|
from litellm._logging import verbose_router_logger, verbose_proxy_logger
|
|
from litellm.proxy.auth.handle_jwt import JWTHandler
|
|
from litellm.proxy.hooks.prompt_injection_detection import (
|
|
_OPTIONAL_PromptInjectionDetection,
|
|
)
|
|
from litellm.proxy.auth.auth_checks import (
|
|
common_checks,
|
|
get_end_user_object,
|
|
get_org_object,
|
|
get_team_object,
|
|
get_user_object,
|
|
allowed_routes_check,
|
|
get_actual_routes,
|
|
)
|
|
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
|
|
|
|
try:
|
|
from litellm._version import version
|
|
except:
|
|
version = "0.0.0"
|
|
litellm.suppress_debug_info = True
|
|
from fastapi import (
|
|
FastAPI,
|
|
Request,
|
|
HTTPException,
|
|
status,
|
|
Depends,
|
|
Header,
|
|
Response,
|
|
Form,
|
|
UploadFile,
|
|
File,
|
|
)
|
|
from fastapi.routing import APIRouter
|
|
from fastapi.security import OAuth2PasswordBearer
|
|
from fastapi.encoders import jsonable_encoder
|
|
from fastapi.responses import (
|
|
StreamingResponse,
|
|
FileResponse,
|
|
ORJSONResponse,
|
|
JSONResponse,
|
|
)
|
|
from fastapi.responses import RedirectResponse
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.security.api_key import APIKeyHeader
|
|
import json
|
|
import logging
|
|
from typing import Union
|
|
|
|
# import enterprise folder
|
|
try:
|
|
# when using litellm cli
|
|
import litellm.proxy.enterprise as enterprise
|
|
except Exception as e:
|
|
# when using litellm docker image
|
|
try:
|
|
import enterprise # type: ignore
|
|
except Exception as e:
|
|
pass
|
|
|
|
ui_link = f"/ui/"
|
|
ui_message = (
|
|
f"👉 [```LiteLLM Admin Panel on /ui```]({ui_link}). Create, Edit Keys with SSO"
|
|
)
|
|
|
|
_docs_url = None if os.getenv("NO_DOCS", "False") == "True" else "/"
|
|
|
|
app = FastAPI(
|
|
docs_url=_docs_url,
|
|
title="LiteLLM API",
|
|
description=f"Proxy Server to call 100+ LLMs in the OpenAI format\n\n{ui_message}",
|
|
version=version,
|
|
root_path=os.environ.get(
|
|
"SERVER_ROOT_PATH", ""
|
|
), # check if user passed root path, FastAPI defaults this value to ""
|
|
)
|
|
|
|
|
|
class ProxyException(Exception):
|
|
# NOTE: DO NOT MODIFY THIS
|
|
# This is used to map exactly to OPENAI Exceptions
|
|
def __init__(
|
|
self,
|
|
message: str,
|
|
type: str,
|
|
param: Optional[str],
|
|
code: Optional[int],
|
|
):
|
|
self.message = message
|
|
self.type = type
|
|
self.param = param
|
|
self.code = code
|
|
|
|
# rules for proxyExceptions
|
|
# Litellm router.py returns "No healthy deployment available" when there are no deployments available
|
|
# Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487
|
|
if (
|
|
"No healthy deployment available" in self.message
|
|
or "No deployments available" in self.message
|
|
):
|
|
self.code = 429
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Converts the ProxyException instance to a dictionary."""
|
|
return {
|
|
"message": self.message,
|
|
"type": self.type,
|
|
"param": self.param,
|
|
"code": self.code,
|
|
}
|
|
|
|
|
|
@app.exception_handler(ProxyException)
|
|
async def openai_exception_handler(request: Request, exc: ProxyException):
|
|
# NOTE: DO NOT MODIFY THIS, its crucial to map to Openai exceptions
|
|
return JSONResponse(
|
|
status_code=(
|
|
int(exc.code) if exc.code else status.HTTP_500_INTERNAL_SERVER_ERROR
|
|
),
|
|
content={
|
|
"error": {
|
|
"message": exc.message,
|
|
"type": exc.type,
|
|
"param": exc.param,
|
|
"code": exc.code,
|
|
}
|
|
},
|
|
)
|
|
|
|
|
|
router = APIRouter()
|
|
origins = ["*"]
|
|
|
|
# get current directory
|
|
try:
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
ui_path = os.path.join(current_dir, "_experimental", "out")
|
|
app.mount("/ui", StaticFiles(directory=ui_path, html=True), name="ui")
|
|
except:
|
|
pass
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins,
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
from typing import Dict
|
|
|
|
api_key_header = APIKeyHeader(
|
|
name="Authorization", auto_error=False, description="Bearer token"
|
|
)
|
|
user_api_base = None
|
|
user_model = None
|
|
user_debug = False
|
|
user_max_tokens = None
|
|
user_request_timeout = None
|
|
user_temperature = None
|
|
user_telemetry = True
|
|
user_config = None
|
|
user_headers = None
|
|
user_config_file_path = f"config_{int(time.time())}.yaml"
|
|
local_logging = True # writes logs to a local api_log.json file for debugging
|
|
experimental = False
|
|
#### GLOBAL VARIABLES ####
|
|
llm_router: Optional[litellm.Router] = None
|
|
llm_model_list: Optional[list] = None
|
|
general_settings: dict = {}
|
|
log_file = "api_log.json"
|
|
worker_config = None
|
|
master_key = None
|
|
otel_logging = False
|
|
prisma_client: Optional[PrismaClient] = None
|
|
custom_db_client: Optional[DBClient] = None
|
|
user_api_key_cache = DualCache()
|
|
redis_usage_cache: Optional[RedisCache] = (
|
|
None # redis cache used for tracking spend, tpm/rpm limits
|
|
)
|
|
user_custom_auth = None
|
|
user_custom_key_generate = None
|
|
use_background_health_checks = None
|
|
use_queue = False
|
|
health_check_interval = None
|
|
health_check_results = {}
|
|
queue: List = []
|
|
litellm_proxy_budget_name = "litellm-proxy-budget"
|
|
litellm_proxy_admin_name = "default_user_id"
|
|
ui_access_mode: Literal["admin", "all"] = "all"
|
|
proxy_budget_rescheduler_min_time = 597
|
|
proxy_budget_rescheduler_max_time = 605
|
|
proxy_batch_write_at = 10 # in seconds
|
|
litellm_master_key_hash = None
|
|
disable_spend_logs = False
|
|
jwt_handler = JWTHandler()
|
|
prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = None
|
|
store_model_in_db: bool = False
|
|
### INITIALIZE GLOBAL LOGGING OBJECT ###
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
|
|
### REDIS QUEUE ###
|
|
async_result = None
|
|
celery_app_conn = None
|
|
celery_fn = None # Redis Queue for handling requests
|
|
### DB WRITER ###
|
|
db_writer_client: Optional[HTTPHandler] = None
|
|
### logger ###
|
|
|
|
|
|
def _get_bearer_token(
|
|
api_key: str,
|
|
):
|
|
if api_key.startswith("Bearer "): # ensure Bearer token passed in
|
|
api_key = api_key.replace("Bearer ", "") # extract the token
|
|
else:
|
|
api_key = ""
|
|
return api_key
|
|
|
|
|
|
def _get_pydantic_json_dict(pydantic_obj: BaseModel) -> dict:
|
|
try:
|
|
return pydantic_obj.model_dump() # type: ignore
|
|
except:
|
|
# if using pydantic v1
|
|
return pydantic_obj.dict()
|
|
|
|
|
|
async def user_api_key_auth(
|
|
request: Request, api_key: str = fastapi.Security(api_key_header)
|
|
) -> UserAPIKeyAuth:
|
|
global master_key, prisma_client, llm_model_list, user_custom_auth, custom_db_client, general_settings
|
|
try:
|
|
if isinstance(api_key, str):
|
|
passed_in_key = api_key
|
|
api_key = _get_bearer_token(api_key=api_key)
|
|
|
|
### USER-DEFINED AUTH FUNCTION ###
|
|
if user_custom_auth is not None:
|
|
response = await user_custom_auth(request=request, api_key=api_key)
|
|
return UserAPIKeyAuth.model_validate(response)
|
|
|
|
### LITELLM-DEFINED AUTH FUNCTION ###
|
|
#### IF JWT ####
|
|
"""
|
|
LiteLLM supports using JWTs.
|
|
|
|
Enable this in proxy config, by setting
|
|
```
|
|
general_settings:
|
|
enable_jwt_auth: true
|
|
```
|
|
"""
|
|
route: str = request.url.path
|
|
|
|
if route in LiteLLMRoutes.public_routes.value:
|
|
# check if public endpoint
|
|
return UserAPIKeyAuth()
|
|
|
|
if general_settings.get("enable_jwt_auth", False) == True:
|
|
is_jwt = jwt_handler.is_jwt(token=api_key)
|
|
verbose_proxy_logger.debug("is_jwt: %s", is_jwt)
|
|
if is_jwt:
|
|
# check if valid token
|
|
valid_token = await jwt_handler.auth_jwt(token=api_key)
|
|
# get scopes
|
|
scopes = jwt_handler.get_scopes(token=valid_token)
|
|
|
|
# check if admin
|
|
is_admin = jwt_handler.is_admin(scopes=scopes)
|
|
# if admin return
|
|
if is_admin:
|
|
# check allowed admin routes
|
|
is_allowed = allowed_routes_check(
|
|
user_role="proxy_admin",
|
|
user_route=route,
|
|
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
|
|
)
|
|
if is_allowed:
|
|
return UserAPIKeyAuth()
|
|
else:
|
|
allowed_routes = (
|
|
jwt_handler.litellm_jwtauth.admin_allowed_routes
|
|
)
|
|
actual_routes = get_actual_routes(allowed_routes=allowed_routes)
|
|
raise Exception(
|
|
f"Admin not allowed to access this route. Route={route}, Allowed Routes={actual_routes}"
|
|
)
|
|
# get team id
|
|
team_id = jwt_handler.get_team_id(token=valid_token, default_value=None)
|
|
|
|
if team_id is None:
|
|
raise Exception(
|
|
f"No team id passed in. Field checked in jwt token - '{jwt_handler.litellm_jwtauth.team_id_jwt_field}'"
|
|
)
|
|
# check allowed team routes
|
|
is_allowed = allowed_routes_check(
|
|
user_role="team",
|
|
user_route=route,
|
|
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
|
|
)
|
|
if is_allowed == False:
|
|
allowed_routes = jwt_handler.litellm_jwtauth.team_allowed_routes
|
|
actual_routes = get_actual_routes(allowed_routes=allowed_routes)
|
|
raise Exception(
|
|
f"Team not allowed to access this route. Route={route}, Allowed Routes={actual_routes}"
|
|
)
|
|
|
|
# check if team in db
|
|
team_object = await get_team_object(
|
|
team_id=team_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
)
|
|
|
|
# [OPTIONAL] track spend for an org id - `LiteLLM_OrganizationTable`
|
|
org_id = jwt_handler.get_org_id(token=valid_token, default_value=None)
|
|
if org_id is not None:
|
|
_ = await get_org_object(
|
|
org_id=org_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
)
|
|
# [OPTIONAL] track spend against an internal employee - `LiteLLM_UserTable`
|
|
user_object = None
|
|
user_id = jwt_handler.get_user_id(token=valid_token, default_value=None)
|
|
if user_id is not None:
|
|
# get the user object
|
|
user_object = await get_user_object(
|
|
user_id=user_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
)
|
|
# save the user object to cache
|
|
await user_api_key_cache.async_set_cache(
|
|
key=user_id, value=user_object
|
|
)
|
|
# [OPTIONAL] track spend against an external user - `LiteLLM_EndUserTable`
|
|
end_user_object = None
|
|
end_user_id = jwt_handler.get_end_user_id(
|
|
token=valid_token, default_value=None
|
|
)
|
|
if end_user_id is not None:
|
|
# get the end-user object
|
|
end_user_object = await get_end_user_object(
|
|
end_user_id=end_user_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
)
|
|
# save the end-user object to cache
|
|
await user_api_key_cache.async_set_cache(
|
|
key=end_user_id, value=end_user_object
|
|
)
|
|
|
|
global_proxy_spend = None
|
|
if litellm.max_budget > 0: # user set proxy max budget
|
|
# check cache
|
|
global_proxy_spend = await user_api_key_cache.async_get_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name)
|
|
)
|
|
if global_proxy_spend is None and prisma_client is not None:
|
|
# get from db
|
|
sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
global_proxy_spend = response[0]["total_spend"]
|
|
|
|
await user_api_key_cache.async_set_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name),
|
|
value=global_proxy_spend,
|
|
ttl=60,
|
|
)
|
|
if global_proxy_spend is not None:
|
|
user_info = {
|
|
"user_id": litellm_proxy_admin_name,
|
|
"max_budget": litellm.max_budget,
|
|
"spend": global_proxy_spend,
|
|
"user_email": "",
|
|
}
|
|
asyncio.create_task(
|
|
proxy_logging_obj.budget_alerts(
|
|
user_max_budget=litellm.max_budget,
|
|
user_current_spend=global_proxy_spend,
|
|
type="user_and_proxy_budget",
|
|
user_info=user_info,
|
|
)
|
|
)
|
|
|
|
# get the request body
|
|
request_data = await _read_request_body(request=request)
|
|
|
|
# run through common checks
|
|
_ = common_checks(
|
|
request_body=request_data,
|
|
team_object=team_object,
|
|
user_object=user_object,
|
|
end_user_object=end_user_object,
|
|
general_settings=general_settings,
|
|
global_proxy_spend=global_proxy_spend,
|
|
route=route,
|
|
)
|
|
# save team object in cache
|
|
await user_api_key_cache.async_set_cache(
|
|
key=team_object.team_id, value=team_object
|
|
)
|
|
|
|
# return UserAPIKeyAuth object
|
|
return UserAPIKeyAuth(
|
|
api_key=None,
|
|
team_id=team_object.team_id,
|
|
team_tpm_limit=team_object.tpm_limit,
|
|
team_rpm_limit=team_object.rpm_limit,
|
|
team_models=team_object.models,
|
|
user_role="app_owner",
|
|
user_id=user_id,
|
|
org_id=org_id,
|
|
)
|
|
#### ELSE ####
|
|
if master_key is None:
|
|
if isinstance(api_key, str):
|
|
return UserAPIKeyAuth(api_key=api_key)
|
|
else:
|
|
return UserAPIKeyAuth()
|
|
elif api_key is None: # only require api key if master key is set
|
|
raise Exception("No api key passed in.")
|
|
elif api_key == "":
|
|
# missing 'Bearer ' prefix
|
|
raise Exception(
|
|
f"Malformed API Key passed in. Ensure Key has `Bearer ` prefix. Passed in: {passed_in_key}"
|
|
)
|
|
|
|
if route == "/user/auth":
|
|
if general_settings.get("allow_user_auth", False) == True:
|
|
return UserAPIKeyAuth()
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="'allow_user_auth' not set or set to False",
|
|
)
|
|
|
|
### CHECK IF ADMIN ###
|
|
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
|
|
### CHECK IF ADMIN ###
|
|
# note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
|
|
## Check CACHE
|
|
valid_token = user_api_key_cache.get_cache(key=hash_token(api_key))
|
|
if (
|
|
valid_token is not None
|
|
and isinstance(valid_token, UserAPIKeyAuth)
|
|
and valid_token.user_role == "proxy_admin"
|
|
):
|
|
return valid_token
|
|
|
|
try:
|
|
is_master_key_valid = secrets.compare_digest(api_key, master_key)
|
|
except Exception as e:
|
|
is_master_key_valid = False
|
|
|
|
if is_master_key_valid:
|
|
_user_api_key_obj = UserAPIKeyAuth(
|
|
api_key=master_key,
|
|
user_role="proxy_admin",
|
|
user_id=litellm_proxy_admin_name,
|
|
)
|
|
await user_api_key_cache.async_set_cache(
|
|
key=hash_token(master_key), value=_user_api_key_obj
|
|
)
|
|
|
|
return _user_api_key_obj
|
|
if isinstance(
|
|
api_key, str
|
|
): # if generated token, make sure it starts with sk-.
|
|
assert api_key.startswith("sk-") # prevent token hashes from being used
|
|
|
|
if (
|
|
prisma_client is None and custom_db_client is None
|
|
): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error
|
|
raise Exception("No connected db.")
|
|
|
|
## check for cache hit (In-Memory Cache)
|
|
original_api_key = api_key # (Patch: For DynamoDB Backwards Compatibility)
|
|
if api_key.startswith("sk-"):
|
|
api_key = hash_token(token=api_key)
|
|
valid_token = user_api_key_cache.get_cache(key=api_key)
|
|
if valid_token is None:
|
|
## check db
|
|
verbose_proxy_logger.debug("api key: %s", api_key)
|
|
if prisma_client is not None:
|
|
valid_token = await prisma_client.get_data(
|
|
token=api_key, table_name="combined_view"
|
|
)
|
|
elif custom_db_client is not None:
|
|
try:
|
|
valid_token = await custom_db_client.get_data(
|
|
key=api_key, table_name="key"
|
|
)
|
|
except:
|
|
# (Patch: For DynamoDB Backwards Compatibility)
|
|
valid_token = await custom_db_client.get_data(
|
|
key=original_api_key, table_name="key"
|
|
)
|
|
verbose_proxy_logger.debug("Token from db: %s", valid_token)
|
|
elif valid_token is not None:
|
|
verbose_proxy_logger.debug("API Key Cache Hit!")
|
|
user_id_information = None
|
|
if valid_token:
|
|
# Got Valid Token from Cache, DB
|
|
# Run checks for
|
|
# 1. If token can call model
|
|
# 2. If user_id for this token is in budget
|
|
# 3. If 'user' passed to /chat/completions, /embeddings endpoint is in budget
|
|
# 4. If token is expired
|
|
# 5. If token spend is under Budget for the token
|
|
# 6. If token spend per model is under budget per model
|
|
# 7. If token spend is under team budget
|
|
# 8. If team spend is under team budget
|
|
|
|
request_data = await _read_request_body(
|
|
request=request
|
|
) # request data, used across all checks. Making this easily available
|
|
|
|
# Check 1. If token can call model
|
|
_model_alias_map = {}
|
|
if (
|
|
hasattr(valid_token, "team_model_aliases")
|
|
and valid_token.team_model_aliases is not None
|
|
):
|
|
_model_alias_map = {
|
|
**valid_token.aliases,
|
|
**valid_token.team_model_aliases,
|
|
}
|
|
else:
|
|
_model_alias_map = {**valid_token.aliases}
|
|
litellm.model_alias_map = _model_alias_map
|
|
config = valid_token.config
|
|
if config != {}:
|
|
model_list = config.get("model_list", [])
|
|
llm_model_list = model_list
|
|
verbose_proxy_logger.debug(
|
|
f"\n new llm router model list {llm_model_list}"
|
|
)
|
|
if (
|
|
len(valid_token.models) == 0
|
|
): # assume an empty model list means all models are allowed to be called
|
|
pass
|
|
elif (
|
|
isinstance(valid_token.models, list)
|
|
and "all-team-models" in valid_token.models
|
|
):
|
|
# Do not do any validation at this step
|
|
# the validation will occur when checking the team has access to this model
|
|
pass
|
|
else:
|
|
try:
|
|
data = await request.json()
|
|
except json.JSONDecodeError:
|
|
data = {} # Provide a default value, such as an empty dictionary
|
|
model = data.get("model", None)
|
|
if model in litellm.model_alias_map:
|
|
model = litellm.model_alias_map[model]
|
|
|
|
## check if model in allowed model names
|
|
verbose_proxy_logger.debug(
|
|
f"LLM Model List pre access group check: {llm_model_list}"
|
|
)
|
|
from collections import defaultdict
|
|
|
|
access_groups = defaultdict(list)
|
|
if llm_model_list is not None:
|
|
for m in llm_model_list:
|
|
for group in m.get("model_info", {}).get("access_groups", []):
|
|
model_name = m["model_name"]
|
|
access_groups[group].append(model_name)
|
|
|
|
models_in_current_access_groups = []
|
|
if (
|
|
len(access_groups) > 0
|
|
): # check if token contains any model access groups
|
|
for idx, m in enumerate(
|
|
valid_token.models
|
|
): # loop token models, if any of them are an access group add the access group
|
|
if m in access_groups:
|
|
# if it is an access group we need to remove it from valid_token.models
|
|
models_in_group = access_groups[m]
|
|
models_in_current_access_groups.extend(models_in_group)
|
|
|
|
# Filter out models that are access_groups
|
|
filtered_models = [
|
|
m for m in valid_token.models if m not in access_groups
|
|
]
|
|
|
|
filtered_models += models_in_current_access_groups
|
|
verbose_proxy_logger.debug(
|
|
f"model: {model}; allowed_models: {filtered_models}"
|
|
)
|
|
if (
|
|
model is not None
|
|
and model not in filtered_models
|
|
and "*" not in filtered_models
|
|
):
|
|
raise ValueError(
|
|
f"API Key not allowed to access model. This token can only access models={valid_token.models}. Tried to access {model}"
|
|
)
|
|
valid_token.models = filtered_models
|
|
verbose_proxy_logger.debug(
|
|
f"filtered allowed_models: {filtered_models}; valid_token.models: {valid_token.models}"
|
|
)
|
|
|
|
# Check 2. If user_id for this token is in budget
|
|
if valid_token.user_id is not None:
|
|
user_id_list = [valid_token.user_id]
|
|
for id in user_id_list:
|
|
value = user_api_key_cache.get_cache(key=id)
|
|
if value is not None:
|
|
if user_id_information is None:
|
|
user_id_information = []
|
|
user_id_information.append(value)
|
|
if user_id_information is None or (
|
|
isinstance(user_id_information, list)
|
|
and len(user_id_information) < 1
|
|
):
|
|
if prisma_client is not None:
|
|
user_id_information = await prisma_client.get_data(
|
|
user_id_list=[
|
|
valid_token.user_id,
|
|
],
|
|
table_name="user",
|
|
query_type="find_all",
|
|
)
|
|
for _id in user_id_information:
|
|
await user_api_key_cache.async_set_cache(
|
|
key=_id["user_id"], value=_id, ttl=600
|
|
)
|
|
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}"
|
|
)
|
|
|
|
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)
|
|
# check if user is admin #
|
|
|
|
# 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)
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"user_id: {_user.get('user_id', None)}; 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)
|
|
expiry_time = datetime.fromisoformat(valid_token.expires)
|
|
if (
|
|
expiry_time.tzinfo is None
|
|
or expiry_time.tzinfo.utcoffset(expiry_time) is None
|
|
):
|
|
expiry_time = expiry_time.replace(tzinfo=timezone.utc)
|
|
verbose_proxy_logger.debug(
|
|
f"Checking if token expired, expiry time {expiry_time} and current time {current_time}"
|
|
)
|
|
if expiry_time < current_time:
|
|
# Token exists but is expired.
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
|
|
)
|
|
|
|
# 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}"
|
|
)
|
|
|
|
# Check 5. Token Model Spend is under Model budget
|
|
max_budget_per_model = valid_token.model_max_budget
|
|
|
|
if (
|
|
max_budget_per_model is not None
|
|
and isinstance(max_budget_per_model, dict)
|
|
and len(max_budget_per_model) > 0
|
|
):
|
|
current_model = request_data.get("model")
|
|
## GET THE SPEND FOR THIS MODEL
|
|
twenty_eight_days_ago = datetime.now() - timedelta(days=28)
|
|
model_spend = await prisma_client.db.litellm_spendlogs.group_by(
|
|
by=["model"],
|
|
sum={"spend": True},
|
|
where={
|
|
"AND": [
|
|
{"api_key": valid_token.token},
|
|
{"startTime": {"gt": twenty_eight_days_ago}},
|
|
{"model": current_model},
|
|
]
|
|
},
|
|
)
|
|
if (
|
|
len(model_spend) > 0
|
|
and max_budget_per_model.get(current_model, None) is not None
|
|
):
|
|
if (
|
|
model_spend[0]["model"] == current_model
|
|
and model_spend[0]["_sum"]["spend"]
|
|
>= max_budget_per_model[current_model]
|
|
):
|
|
current_model_spend = model_spend[0]["_sum"]["spend"]
|
|
current_model_budget = max_budget_per_model[current_model]
|
|
raise Exception(
|
|
f"ExceededModelBudget: Current spend for model: {current_model_spend}; Max Budget for Model: {current_model_budget}"
|
|
)
|
|
# Check 6. Token spend is under Team budget
|
|
if (
|
|
valid_token.spend is not None
|
|
and hasattr(valid_token, "team_max_budget")
|
|
and valid_token.team_max_budget is not None
|
|
):
|
|
asyncio.create_task(
|
|
proxy_logging_obj.budget_alerts(
|
|
user_max_budget=valid_token.team_max_budget,
|
|
user_current_spend=valid_token.spend,
|
|
type="token_budget",
|
|
user_info=valid_token,
|
|
)
|
|
)
|
|
|
|
if valid_token.spend >= valid_token.team_max_budget:
|
|
raise Exception(
|
|
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Team: {valid_token.team_max_budget}"
|
|
)
|
|
|
|
# Check 7. Team spend is under Team budget
|
|
if (
|
|
hasattr(valid_token, "team_spend")
|
|
and valid_token.team_spend is not None
|
|
and hasattr(valid_token, "team_max_budget")
|
|
and valid_token.team_max_budget is not None
|
|
):
|
|
asyncio.create_task(
|
|
proxy_logging_obj.budget_alerts(
|
|
user_max_budget=valid_token.team_max_budget,
|
|
user_current_spend=valid_token.team_spend,
|
|
type="token_budget",
|
|
user_info=valid_token,
|
|
)
|
|
)
|
|
|
|
if valid_token.team_spend >= valid_token.team_max_budget:
|
|
raise Exception(
|
|
f"ExceededTokenBudget: Current Team Spend: {valid_token.team_spend}; Max Budget for Team: {valid_token.team_max_budget}"
|
|
)
|
|
|
|
# Check 8: Additional Common Checks across jwt + key auth
|
|
_team_obj = LiteLLM_TeamTable(
|
|
team_id=valid_token.team_id,
|
|
max_budget=valid_token.team_max_budget,
|
|
spend=valid_token.team_spend,
|
|
tpm_limit=valid_token.team_tpm_limit,
|
|
rpm_limit=valid_token.team_rpm_limit,
|
|
blocked=valid_token.team_blocked,
|
|
models=valid_token.team_models,
|
|
)
|
|
|
|
user_api_key_cache.set_cache(
|
|
key=valid_token.team_id, value=_team_obj
|
|
) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
|
|
|
|
_end_user_object = None
|
|
if "user" in request_data:
|
|
_id = "end_user_id:{}".format(request_data["user"])
|
|
_end_user_object = await user_api_key_cache.async_get_cache(key=_id)
|
|
if _end_user_object is not None:
|
|
_end_user_object = LiteLLM_EndUserTable(**_end_user_object)
|
|
|
|
global_proxy_spend = None
|
|
if litellm.max_budget > 0: # user set proxy max budget
|
|
# check cache
|
|
global_proxy_spend = await user_api_key_cache.async_get_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name)
|
|
)
|
|
if global_proxy_spend is None:
|
|
# get from db
|
|
sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
global_proxy_spend = response[0]["total_spend"]
|
|
await user_api_key_cache.async_set_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name),
|
|
value=global_proxy_spend,
|
|
ttl=60,
|
|
)
|
|
|
|
if global_proxy_spend is not None:
|
|
user_info = {
|
|
"user_id": litellm_proxy_admin_name,
|
|
"max_budget": litellm.max_budget,
|
|
"spend": global_proxy_spend,
|
|
"user_email": "",
|
|
}
|
|
asyncio.create_task(
|
|
proxy_logging_obj.budget_alerts(
|
|
user_max_budget=litellm.max_budget,
|
|
user_current_spend=global_proxy_spend,
|
|
type="user_and_proxy_budget",
|
|
user_info=user_info,
|
|
)
|
|
)
|
|
_ = common_checks(
|
|
request_body=request_data,
|
|
team_object=_team_obj,
|
|
user_object=None,
|
|
end_user_object=_end_user_object,
|
|
general_settings=general_settings,
|
|
global_proxy_spend=global_proxy_spend,
|
|
route=route,
|
|
)
|
|
# Token passed all checks
|
|
api_key = valid_token.token
|
|
|
|
# Add hashed token to cache
|
|
await user_api_key_cache.async_set_cache(
|
|
key=api_key, value=valid_token, ttl=600
|
|
)
|
|
valid_token_dict = _get_pydantic_json_dict(valid_token)
|
|
valid_token_dict.pop("token", None)
|
|
"""
|
|
asyncio create task to update the user api key cache with the user db table as well
|
|
|
|
This makes the user row data accessible to pre-api call hooks.
|
|
"""
|
|
if custom_db_client is not None:
|
|
asyncio.create_task(
|
|
_cache_user_row(
|
|
user_id=valid_token.user_id,
|
|
cache=user_api_key_cache,
|
|
db=custom_db_client,
|
|
)
|
|
)
|
|
|
|
if not _is_user_proxy_admin(user_id_information): # if non-admin
|
|
if route in LiteLLMRoutes.openai_routes.value:
|
|
pass
|
|
elif (
|
|
route in LiteLLMRoutes.info_routes.value
|
|
): # check if user allowed to call an info route
|
|
if route == "/key/info":
|
|
# check if user can access this route
|
|
query_params = request.query_params
|
|
key = query_params.get("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",
|
|
)
|
|
elif route == "/user/info":
|
|
# check if user can access this route
|
|
query_params = request.query_params
|
|
user_id = query_params.get("user_id")
|
|
verbose_proxy_logger.debug(
|
|
f"user_id: {user_id} & valid_token.user_id: {valid_token.user_id}"
|
|
)
|
|
if user_id != valid_token.user_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="key not allowed to access this user's info",
|
|
)
|
|
elif route == "/model/info":
|
|
# /model/info just shows models user has access to
|
|
pass
|
|
elif route == "/team/info":
|
|
# check if key can access this team's info
|
|
query_params = request.query_params
|
|
team_id = query_params.get("team_id")
|
|
if team_id != valid_token.team_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="key not allowed to access this team's info",
|
|
)
|
|
elif (
|
|
_has_user_setup_sso()
|
|
and route in LiteLLMRoutes.sso_only_routes.value
|
|
):
|
|
pass
|
|
else:
|
|
user_role = "unknown"
|
|
user_id = "unknown"
|
|
if user_id_information is not None and isinstance(
|
|
user_id_information, list
|
|
):
|
|
_user = user_id_information[0]
|
|
user_role = _user.get("user_role", {}).get(
|
|
"user_role", "unknown"
|
|
)
|
|
user_id = _user.get("user_id", "unknown")
|
|
raise Exception(
|
|
f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={user_id}"
|
|
)
|
|
|
|
# check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions
|
|
# sso/login, ui/login, /key functions and /user functions
|
|
# this will never be allowed to call /chat/completions
|
|
token_team = getattr(valid_token, "team_id", None)
|
|
if token_team is not None and token_team == "litellm-dashboard":
|
|
# this token is only used for managing the ui
|
|
allowed_routes = [
|
|
"/sso",
|
|
"/login",
|
|
"/key/generate",
|
|
"/key/update",
|
|
"/key/info",
|
|
"/config",
|
|
"/spend",
|
|
"/user",
|
|
"/model/info",
|
|
"/v2/model/info",
|
|
"/v2/key/info",
|
|
"/models",
|
|
"/v1/models",
|
|
"/global/spend",
|
|
"/global/spend/logs",
|
|
"/global/spend/keys",
|
|
"/global/spend/models",
|
|
"/global/predict/spend/logs",
|
|
"/health/services",
|
|
]
|
|
# check if the current route startswith any of the allowed routes
|
|
if (
|
|
route is not None
|
|
and isinstance(route, str)
|
|
and any(
|
|
route.startswith(allowed_route) for allowed_route in allowed_routes
|
|
)
|
|
):
|
|
# Do something if the current route starts with any of the allowed routes
|
|
pass
|
|
else:
|
|
if user_id_information is not None and _is_user_proxy_admin(
|
|
user_id_information
|
|
):
|
|
return UserAPIKeyAuth(
|
|
api_key=api_key, user_role="proxy_admin", **valid_token_dict
|
|
)
|
|
elif (
|
|
_has_user_setup_sso()
|
|
and route in LiteLLMRoutes.sso_only_routes.value
|
|
):
|
|
return UserAPIKeyAuth(
|
|
api_key=api_key, user_role="app_owner", **valid_token_dict
|
|
)
|
|
else:
|
|
raise Exception(
|
|
f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
|
|
)
|
|
if valid_token is None:
|
|
# No token was found when looking up in the DB
|
|
raise Exception("Invalid token passed")
|
|
if valid_token_dict is not None:
|
|
return UserAPIKeyAuth(api_key=api_key, **valid_token_dict)
|
|
else:
|
|
raise Exception()
|
|
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_401_UNAUTHORIZED),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
raise e
|
|
raise ProxyException(
|
|
message="Authentication Error, " + str(e),
|
|
type="auth_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=status.HTTP_401_UNAUTHORIZED,
|
|
)
|
|
|
|
|
|
def prisma_setup(database_url: Optional[str]):
|
|
global prisma_client, proxy_logging_obj, user_api_key_cache
|
|
|
|
if database_url is not None:
|
|
try:
|
|
prisma_client = PrismaClient(
|
|
database_url=database_url, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
|
|
def load_from_azure_key_vault(use_azure_key_vault: bool = False):
|
|
if use_azure_key_vault is False:
|
|
return
|
|
|
|
try:
|
|
from azure.keyvault.secrets import SecretClient
|
|
from azure.identity import ClientSecretCredential
|
|
|
|
# Set your Azure Key Vault URI
|
|
KVUri = os.getenv("AZURE_KEY_VAULT_URI", None)
|
|
|
|
# Set your Azure AD application/client ID, client secret, and tenant ID
|
|
client_id = os.getenv("AZURE_CLIENT_ID", None)
|
|
client_secret = os.getenv("AZURE_CLIENT_SECRET", None)
|
|
tenant_id = os.getenv("AZURE_TENANT_ID", None)
|
|
|
|
if (
|
|
KVUri is not None
|
|
and client_id is not None
|
|
and client_secret is not None
|
|
and tenant_id is not None
|
|
):
|
|
# Initialize the ClientSecretCredential
|
|
credential = ClientSecretCredential(
|
|
client_id=client_id, client_secret=client_secret, tenant_id=tenant_id
|
|
)
|
|
|
|
# Create the SecretClient using the credential
|
|
client = SecretClient(vault_url=KVUri, credential=credential)
|
|
|
|
litellm.secret_manager_client = client
|
|
litellm._key_management_system = KeyManagementSystem.AZURE_KEY_VAULT
|
|
else:
|
|
raise Exception(
|
|
f"Missing KVUri or client_id or client_secret or tenant_id from environment"
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
"Error when loading keys from Azure Key Vault. Ensure you run `pip install azure-identity azure-keyvault-secrets`"
|
|
)
|
|
|
|
|
|
def cost_tracking():
|
|
global prisma_client, custom_db_client
|
|
if prisma_client is not None or custom_db_client is not None:
|
|
if isinstance(litellm.success_callback, list):
|
|
verbose_proxy_logger.debug("setting litellm success callback to track cost")
|
|
if (_PROXY_track_cost_callback) not in litellm.success_callback: # type: ignore
|
|
litellm.success_callback.append(_PROXY_track_cost_callback) # type: ignore
|
|
|
|
|
|
async def _PROXY_failure_handler(
|
|
kwargs, # kwargs to completion
|
|
completion_response: litellm.ModelResponse, # response from completion
|
|
start_time=None,
|
|
end_time=None, # start/end time for completion
|
|
):
|
|
global prisma_client
|
|
if prisma_client is not None:
|
|
verbose_proxy_logger.debug(
|
|
"inside _PROXY_failure_handler kwargs=", extra=kwargs
|
|
)
|
|
|
|
_exception = kwargs.get("exception")
|
|
_exception_type = _exception.__class__.__name__
|
|
_model = kwargs.get("model", None)
|
|
|
|
_optional_params = kwargs.get("optional_params", {})
|
|
_optional_params = copy.deepcopy(_optional_params)
|
|
|
|
for k, v in _optional_params.items():
|
|
v = str(v)
|
|
v = v[:100]
|
|
|
|
_status_code = "500"
|
|
try:
|
|
_status_code = str(_exception.status_code)
|
|
except:
|
|
# Don't let this fail logging the exception to the dB
|
|
pass
|
|
|
|
_litellm_params = kwargs.get("litellm_params", {}) or {}
|
|
_metadata = _litellm_params.get("metadata", {}) or {}
|
|
_model_id = _metadata.get("model_info", {}).get("id", "")
|
|
_model_group = _metadata.get("model_group", "")
|
|
api_base = litellm.get_api_base(model=_model, optional_params=_litellm_params)
|
|
_exception_string = str(_exception)[:500]
|
|
|
|
error_log = LiteLLM_ErrorLogs(
|
|
request_id=str(uuid.uuid4()),
|
|
model_group=_model_group,
|
|
model_id=_model_id,
|
|
litellm_model_name=kwargs.get("model"),
|
|
request_kwargs=_optional_params,
|
|
api_base=api_base,
|
|
exception_type=_exception_type,
|
|
status_code=_status_code,
|
|
exception_string=_exception_string,
|
|
startTime=kwargs.get("start_time"),
|
|
endTime=kwargs.get("end_time"),
|
|
)
|
|
|
|
# helper function to convert to dict on pydantic v2 & v1
|
|
error_log_dict = _get_pydantic_json_dict(error_log)
|
|
error_log_dict["request_kwargs"] = json.dumps(error_log_dict["request_kwargs"])
|
|
|
|
await prisma_client.db.litellm_errorlogs.create(
|
|
data=error_log_dict # type: ignore
|
|
)
|
|
|
|
pass
|
|
|
|
|
|
async def _PROXY_track_cost_callback(
|
|
kwargs, # kwargs to completion
|
|
completion_response: litellm.ModelResponse, # response from completion
|
|
start_time=None,
|
|
end_time=None, # start/end time for completion
|
|
):
|
|
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
|
|
global prisma_client, custom_db_client
|
|
try:
|
|
# check if it has collected an entire stream response
|
|
verbose_proxy_logger.debug("Proxy: In track_cost_callback for: %s", kwargs)
|
|
verbose_proxy_logger.debug(
|
|
f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}"
|
|
)
|
|
litellm_params = kwargs.get("litellm_params", {}) or {}
|
|
proxy_server_request = litellm_params.get("proxy_server_request") or {}
|
|
end_user_id = proxy_server_request.get("body", {}).get("user", None)
|
|
user_id = kwargs["litellm_params"]["metadata"].get("user_api_key_user_id", None)
|
|
team_id = kwargs["litellm_params"]["metadata"].get("user_api_key_team_id", None)
|
|
org_id = kwargs["litellm_params"]["metadata"].get("user_api_key_org_id", None)
|
|
if kwargs.get("response_cost", None) is not None:
|
|
response_cost = kwargs["response_cost"]
|
|
user_api_key = kwargs["litellm_params"]["metadata"].get(
|
|
"user_api_key", 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.debug(
|
|
f"user_api_key {user_api_key}, prisma_client: {prisma_client}, custom_db_client: {custom_db_client}"
|
|
)
|
|
if user_api_key is not None or user_id is not None or team_id is not None:
|
|
## UPDATE DATABASE
|
|
await update_database(
|
|
token=user_api_key,
|
|
response_cost=response_cost,
|
|
user_id=user_id,
|
|
end_user_id=end_user_id,
|
|
team_id=team_id,
|
|
kwargs=kwargs,
|
|
completion_response=completion_response,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
org_id=org_id,
|
|
)
|
|
|
|
await update_cache(
|
|
token=user_api_key,
|
|
user_id=user_id,
|
|
end_user_id=end_user_id,
|
|
response_cost=response_cost,
|
|
)
|
|
else:
|
|
raise Exception(
|
|
"User API key and team id and user id missing from custom callback."
|
|
)
|
|
else:
|
|
if kwargs["stream"] != True or (
|
|
kwargs["stream"] == True and "complete_streaming_response" in kwargs
|
|
):
|
|
raise Exception(
|
|
f"Model not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
|
|
)
|
|
except Exception as e:
|
|
error_msg = f"error in tracking cost callback - {traceback.format_exc()}"
|
|
model = kwargs.get("model", "")
|
|
metadata = kwargs.get("litellm_params", {}).get("metadata", {})
|
|
error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n metadata: {metadata}\n"
|
|
user_id = user_id or "not-found"
|
|
asyncio.create_task(
|
|
proxy_logging_obj.budget_alerts(
|
|
user_max_budget=0,
|
|
user_current_spend=0,
|
|
type="failed_tracking",
|
|
user_info=user_id,
|
|
error_message=error_msg,
|
|
)
|
|
)
|
|
verbose_proxy_logger.debug("error in tracking cost callback - %s", e)
|
|
|
|
|
|
def error_tracking():
|
|
global prisma_client, custom_db_client
|
|
if prisma_client is not None or custom_db_client is not None:
|
|
if isinstance(litellm.failure_callback, list):
|
|
verbose_proxy_logger.debug("setting litellm failure callback to track cost")
|
|
if (_PROXY_failure_handler) not in litellm.failure_callback: # type: ignore
|
|
litellm.failure_callback.append(_PROXY_failure_handler) # type: ignore
|
|
|
|
|
|
def _set_spend_logs_payload(
|
|
payload: dict, prisma_client: PrismaClient, spend_logs_url: Optional[str] = None
|
|
):
|
|
if prisma_client is not None and spend_logs_url is not None:
|
|
if isinstance(payload["startTime"], datetime):
|
|
payload["startTime"] = payload["startTime"].isoformat()
|
|
if isinstance(payload["endTime"], datetime):
|
|
payload["endTime"] = payload["endTime"].isoformat()
|
|
prisma_client.spend_log_transactions.append(payload)
|
|
elif prisma_client is not None:
|
|
prisma_client.spend_log_transactions.append(payload)
|
|
return prisma_client
|
|
|
|
|
|
async def update_database(
|
|
token,
|
|
response_cost,
|
|
user_id=None,
|
|
end_user_id=None,
|
|
team_id=None,
|
|
kwargs=None,
|
|
completion_response=None,
|
|
start_time=None,
|
|
end_time=None,
|
|
org_id=None,
|
|
):
|
|
try:
|
|
global prisma_client
|
|
verbose_proxy_logger.info(
|
|
f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}"
|
|
)
|
|
if token is not None and isinstance(token, str) and token.startswith("sk-"):
|
|
hashed_token = hash_token(token=token)
|
|
else:
|
|
hashed_token = token
|
|
|
|
### UPDATE USER SPEND ###
|
|
async def _update_user_db():
|
|
"""
|
|
- Update that user's row
|
|
- Update litellm-proxy-budget row (global proxy spend)
|
|
"""
|
|
## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db
|
|
existing_token_obj = await user_api_key_cache.async_get_cache(
|
|
key=hashed_token
|
|
)
|
|
existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
|
|
if existing_user_obj is not None and isinstance(existing_user_obj, dict):
|
|
existing_user_obj = LiteLLM_UserTable(**existing_user_obj)
|
|
data_list = []
|
|
try:
|
|
if prisma_client is not None: # update
|
|
user_ids = [user_id]
|
|
if (
|
|
litellm.max_budget > 0
|
|
): # track global proxy budget, if user set max budget
|
|
user_ids.append(litellm_proxy_budget_name)
|
|
### KEY CHANGE ###
|
|
for _id in user_ids:
|
|
if _id is not None:
|
|
prisma_client.user_list_transactons[_id] = (
|
|
response_cost
|
|
+ prisma_client.user_list_transactons.get(_id, 0)
|
|
)
|
|
if end_user_id is not None:
|
|
prisma_client.end_user_list_transactons[end_user_id] = (
|
|
response_cost
|
|
+ prisma_client.end_user_list_transactons.get(
|
|
end_user_id, 0
|
|
)
|
|
)
|
|
elif custom_db_client is not None:
|
|
for id in user_ids:
|
|
if id is None:
|
|
continue
|
|
if (
|
|
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:
|
|
# if user does not exist in LiteLLM_UserTable, create a new user
|
|
existing_spend = 0
|
|
max_user_budget = None
|
|
if litellm.max_user_budget is not None:
|
|
max_user_budget = litellm.max_user_budget
|
|
existing_spend_obj = LiteLLM_UserTable(
|
|
user_id=id,
|
|
spend=0,
|
|
max_budget=max_user_budget,
|
|
user_email=None,
|
|
)
|
|
else:
|
|
existing_spend = existing_spend_obj.spend
|
|
|
|
# Calculate the new cost by adding the existing cost and response_cost
|
|
existing_spend_obj.spend = existing_spend + response_cost
|
|
|
|
# track cost per model, for the given user
|
|
spend_per_model = existing_spend_obj.model_spend or {}
|
|
current_model = kwargs.get("model")
|
|
|
|
if current_model is not None and spend_per_model is not None:
|
|
if spend_per_model.get(current_model) is None:
|
|
spend_per_model[current_model] = response_cost
|
|
else:
|
|
spend_per_model[current_model] += response_cost
|
|
existing_spend_obj.model_spend = spend_per_model
|
|
|
|
valid_token = user_api_key_cache.get_cache(key=id)
|
|
if valid_token is not None and isinstance(valid_token, dict):
|
|
user_api_key_cache.set_cache(
|
|
key=id, value=existing_spend_obj.json()
|
|
)
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"user - new cost: {existing_spend_obj.spend}, user_id: {id}"
|
|
)
|
|
data_list.append(existing_spend_obj)
|
|
|
|
if 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",
|
|
)
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.info(
|
|
"\033[91m"
|
|
+ f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}"
|
|
)
|
|
|
|
### UPDATE KEY SPEND ###
|
|
async def _update_key_db():
|
|
try:
|
|
verbose_proxy_logger.debug(
|
|
f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}."
|
|
)
|
|
if hashed_token is None:
|
|
return
|
|
if prisma_client is not None:
|
|
prisma_client.key_list_transactons[hashed_token] = (
|
|
response_cost
|
|
+ prisma_client.key_list_transactons.get(hashed_token, 0)
|
|
)
|
|
elif custom_db_client is not None:
|
|
# Fetch the existing cost for the given token
|
|
existing_spend_obj = await custom_db_client.get_data(
|
|
key=token, table_name="key"
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"_update_key_db existing spend: {existing_spend_obj}"
|
|
)
|
|
if existing_spend_obj is None:
|
|
existing_spend = 0
|
|
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
|
|
|
|
verbose_proxy_logger.debug("new cost: %s", new_spend)
|
|
# Update the cost column for the given token
|
|
await custom_db_client.update_data(
|
|
key=token, value={"spend": new_spend}, table_name="key"
|
|
)
|
|
|
|
valid_token = user_api_key_cache.get_cache(key=token)
|
|
if valid_token is not None:
|
|
valid_token.spend = new_spend
|
|
user_api_key_cache.set_cache(key=token, value=valid_token)
|
|
except Exception as e:
|
|
verbose_proxy_logger.info(
|
|
f"Update Key DB Call failed to execute - {str(e)}\n{traceback.format_exc()}"
|
|
)
|
|
raise e
|
|
|
|
### UPDATE SPEND LOGS ###
|
|
async def _insert_spend_log_to_db():
|
|
try:
|
|
global prisma_client
|
|
if prisma_client is not None:
|
|
# Helper to generate payload to log
|
|
payload = get_logging_payload(
|
|
kwargs=kwargs,
|
|
response_obj=completion_response,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
|
|
payload["spend"] = response_cost
|
|
prisma_client = _set_spend_logs_payload(
|
|
payload=payload,
|
|
spend_logs_url=os.getenv("SPEND_LOGS_URL"),
|
|
prisma_client=prisma_client,
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}"
|
|
)
|
|
raise e
|
|
|
|
### UPDATE TEAM SPEND ###
|
|
async def _update_team_db():
|
|
try:
|
|
verbose_proxy_logger.debug(
|
|
f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}."
|
|
)
|
|
if team_id is None:
|
|
verbose_proxy_logger.debug(
|
|
"track_cost_callback: team_id is None. Not tracking spend for team"
|
|
)
|
|
return
|
|
if prisma_client is not None:
|
|
prisma_client.team_list_transactons[team_id] = (
|
|
response_cost
|
|
+ prisma_client.team_list_transactons.get(team_id, 0)
|
|
)
|
|
elif custom_db_client is not None:
|
|
# Fetch the existing cost for the given token
|
|
existing_spend_obj = await custom_db_client.get_data(
|
|
key=token, table_name="key"
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"_update_key_db existing spend: {existing_spend_obj}"
|
|
)
|
|
if existing_spend_obj is None:
|
|
existing_spend = 0
|
|
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
|
|
|
|
verbose_proxy_logger.debug("new cost: %s", new_spend)
|
|
# Update the cost column for the given token
|
|
await custom_db_client.update_data(
|
|
key=token, value={"spend": new_spend}, table_name="key"
|
|
)
|
|
|
|
valid_token = user_api_key_cache.get_cache(key=token)
|
|
if valid_token is not None:
|
|
valid_token.spend = new_spend
|
|
user_api_key_cache.set_cache(key=token, value=valid_token)
|
|
except Exception as e:
|
|
verbose_proxy_logger.info(
|
|
f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}"
|
|
)
|
|
raise e
|
|
|
|
### UPDATE ORG SPEND ###
|
|
async def _update_org_db():
|
|
try:
|
|
verbose_proxy_logger.debug(
|
|
"adding spend to org db. Response cost: {}. org_id: {}.".format(
|
|
response_cost, org_id
|
|
)
|
|
)
|
|
if org_id is None:
|
|
verbose_proxy_logger.debug(
|
|
"track_cost_callback: org_id is None. Not tracking spend for org"
|
|
)
|
|
return
|
|
if prisma_client is not None:
|
|
prisma_client.org_list_transactons[org_id] = (
|
|
response_cost
|
|
+ prisma_client.org_list_transactons.get(org_id, 0)
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.info(
|
|
f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}"
|
|
)
|
|
raise e
|
|
|
|
asyncio.create_task(_update_user_db())
|
|
asyncio.create_task(_update_key_db())
|
|
asyncio.create_task(_update_team_db())
|
|
asyncio.create_task(_update_org_db())
|
|
# asyncio.create_task(_insert_spend_log_to_db())
|
|
if disable_spend_logs == False:
|
|
await _insert_spend_log_to_db()
|
|
|
|
verbose_proxy_logger.debug("Runs spend update on all tables")
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
f"Error updating Prisma database: {traceback.format_exc()}"
|
|
)
|
|
|
|
|
|
async def update_cache(
|
|
token: Optional[str],
|
|
user_id: Optional[str],
|
|
end_user_id: Optional[str],
|
|
response_cost: Optional[float],
|
|
):
|
|
"""
|
|
Use this to update the cache with new user spend.
|
|
|
|
Put any alerting logic in here.
|
|
"""
|
|
|
|
### UPDATE KEY SPEND ###
|
|
async def _update_key_cache():
|
|
# Fetch the existing cost for the given token
|
|
if isinstance(token, str) and token.startswith("sk-"):
|
|
hashed_token = hash_token(token=token)
|
|
else:
|
|
hashed_token = token
|
|
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
|
|
existing_spend_obj = await user_api_key_cache.async_get_cache(key=hashed_token)
|
|
verbose_proxy_logger.debug(
|
|
f"_update_key_cache: existing_spend_obj={existing_spend_obj}"
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"_update_key_cache: existing spend: {existing_spend_obj}"
|
|
)
|
|
if existing_spend_obj is None:
|
|
existing_spend = 0
|
|
existing_spend_obj = LiteLLM_VerificationTokenView()
|
|
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
|
|
|
|
## CHECK IF USER PROJECTED SPEND > SOFT LIMIT
|
|
soft_budget_cooldown = existing_spend_obj.soft_budget_cooldown
|
|
if (
|
|
existing_spend_obj.soft_budget_cooldown == False
|
|
and existing_spend_obj.litellm_budget_table is not None
|
|
and (
|
|
_is_projected_spend_over_limit(
|
|
current_spend=new_spend,
|
|
soft_budget_limit=existing_spend_obj.litellm_budget_table.soft_budget,
|
|
)
|
|
== True
|
|
)
|
|
):
|
|
key_alias = existing_spend_obj.key_alias
|
|
projected_spend, projected_exceeded_date = _get_projected_spend_over_limit(
|
|
current_spend=new_spend,
|
|
soft_budget_limit=existing_spend_obj.litellm_budget_table.soft_budget,
|
|
)
|
|
soft_limit = existing_spend_obj.litellm_budget_table.soft_budget
|
|
user_info = {
|
|
"key_alias": key_alias,
|
|
"projected_spend": projected_spend,
|
|
"projected_exceeded_date": projected_exceeded_date,
|
|
}
|
|
# alert user
|
|
asyncio.create_task(
|
|
proxy_logging_obj.budget_alerts(
|
|
type="projected_limit_exceeded",
|
|
user_info=user_info,
|
|
user_max_budget=soft_limit,
|
|
user_current_spend=new_spend,
|
|
)
|
|
)
|
|
# set cooldown on alert
|
|
soft_budget_cooldown = True
|
|
|
|
if (
|
|
existing_spend_obj is not None
|
|
and getattr(existing_spend_obj, "team_spend", None) is not None
|
|
):
|
|
existing_team_spend = existing_spend_obj.team_spend
|
|
# Calculate the new cost by adding the existing cost and response_cost
|
|
existing_spend_obj.team_spend = existing_team_spend + response_cost
|
|
|
|
# Update the cost column for the given token
|
|
existing_spend_obj.spend = new_spend
|
|
user_api_key_cache.set_cache(key=hashed_token, value=existing_spend_obj)
|
|
|
|
async def _update_user_cache():
|
|
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
|
|
user_ids = [user_id]
|
|
try:
|
|
for _id in user_ids:
|
|
# Fetch the existing cost for the given user
|
|
if _id is None:
|
|
continue
|
|
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
|
|
if existing_spend_obj is None:
|
|
# if user does not exist in LiteLLM_UserTable, create a new user
|
|
existing_spend = 0
|
|
max_user_budget = None
|
|
if litellm.max_user_budget is not None:
|
|
max_user_budget = litellm.max_user_budget
|
|
existing_spend_obj = LiteLLM_UserTable(
|
|
user_id=_id,
|
|
spend=0,
|
|
max_budget=max_user_budget,
|
|
user_email=None,
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
|
|
)
|
|
if existing_spend_obj is None:
|
|
existing_spend = 0
|
|
else:
|
|
if isinstance(existing_spend_obj, dict):
|
|
existing_spend = existing_spend_obj["spend"]
|
|
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
|
|
|
|
# Update the cost column for the given user
|
|
if isinstance(existing_spend_obj, dict):
|
|
existing_spend_obj["spend"] = new_spend
|
|
user_api_key_cache.set_cache(key=_id, value=existing_spend_obj)
|
|
else:
|
|
existing_spend_obj.spend = new_spend
|
|
user_api_key_cache.set_cache(
|
|
key=_id, value=existing_spend_obj.json()
|
|
)
|
|
## UPDATE GLOBAL PROXY ##
|
|
global_proxy_spend = await user_api_key_cache.async_get_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name)
|
|
)
|
|
if global_proxy_spend is None:
|
|
await user_api_key_cache.async_set_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name), value=response_cost
|
|
)
|
|
elif response_cost is not None and global_proxy_spend is not None:
|
|
increment = global_proxy_spend + response_cost
|
|
await user_api_key_cache.async_set_cache(
|
|
key="{}:spend".format(litellm_proxy_admin_name), value=increment
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
f"An error occurred updating user cache: {str(e)}\n\n{traceback.format_exc()}"
|
|
)
|
|
|
|
async def _update_end_user_cache():
|
|
_id = "end_user_id:{}".format(end_user_id)
|
|
try:
|
|
# Fetch the existing cost for the given user
|
|
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
|
|
if existing_spend_obj is None:
|
|
# if user does not exist in LiteLLM_UserTable, create a new user
|
|
existing_spend = 0
|
|
max_user_budget = None
|
|
max_end_user_budget = None
|
|
if litellm.max_end_user_budget is not None:
|
|
max_end_user_budget = litellm.max_end_user_budget
|
|
existing_spend_obj = LiteLLM_EndUserTable(
|
|
user_id=_id,
|
|
spend=0,
|
|
blocked=False,
|
|
litellm_budget_table=LiteLLM_BudgetTable(
|
|
max_budget=max_end_user_budget
|
|
),
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
|
|
)
|
|
if existing_spend_obj is None:
|
|
existing_spend = 0
|
|
else:
|
|
if isinstance(existing_spend_obj, dict):
|
|
existing_spend = existing_spend_obj["spend"]
|
|
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
|
|
|
|
# Update the cost column for the given user
|
|
if isinstance(existing_spend_obj, dict):
|
|
existing_spend_obj["spend"] = new_spend
|
|
user_api_key_cache.set_cache(key=_id, value=existing_spend_obj)
|
|
else:
|
|
existing_spend_obj.spend = new_spend
|
|
user_api_key_cache.set_cache(key=_id, value=existing_spend_obj.json())
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
f"An error occurred updating end user cache: {str(e)}\n\n{traceback.format_exc()}"
|
|
)
|
|
|
|
if token is not None:
|
|
asyncio.create_task(_update_key_cache())
|
|
|
|
asyncio.create_task(_update_user_cache())
|
|
|
|
if end_user_id is not None:
|
|
asyncio.create_task(_update_end_user_cache())
|
|
|
|
|
|
def run_ollama_serve():
|
|
try:
|
|
command = ["ollama", "serve"]
|
|
|
|
with open(os.devnull, "w") as devnull:
|
|
process = subprocess.Popen(command, stdout=devnull, stderr=devnull)
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
f"""
|
|
LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve`
|
|
"""
|
|
)
|
|
|
|
|
|
async def _run_background_health_check():
|
|
"""
|
|
Periodically run health checks in the background on the endpoints.
|
|
|
|
Update health_check_results, based on this.
|
|
"""
|
|
global health_check_results, llm_model_list, health_check_interval
|
|
while True:
|
|
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
|
|
model_list=llm_model_list
|
|
)
|
|
|
|
# Update the global variable with the health check results
|
|
health_check_results["healthy_endpoints"] = healthy_endpoints
|
|
health_check_results["unhealthy_endpoints"] = unhealthy_endpoints
|
|
health_check_results["healthy_count"] = len(healthy_endpoints)
|
|
health_check_results["unhealthy_count"] = len(unhealthy_endpoints)
|
|
|
|
await asyncio.sleep(health_check_interval)
|
|
|
|
|
|
semaphore = asyncio.Semaphore(1)
|
|
|
|
|
|
class ProxyConfig:
|
|
"""
|
|
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
def is_yaml(self, config_file_path: str) -> bool:
|
|
if not os.path.isfile(config_file_path):
|
|
return False
|
|
|
|
_, file_extension = os.path.splitext(config_file_path)
|
|
return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml"
|
|
|
|
async def get_config(self, config_file_path: Optional[str] = None) -> dict:
|
|
global prisma_client, user_config_file_path
|
|
|
|
file_path = config_file_path or user_config_file_path
|
|
if config_file_path is not None:
|
|
user_config_file_path = config_file_path
|
|
# Load existing config
|
|
## Yaml
|
|
if os.path.exists(f"{file_path}"):
|
|
with open(f"{file_path}", "r") as config_file:
|
|
config = yaml.safe_load(config_file)
|
|
else:
|
|
config = {
|
|
"model_list": [],
|
|
"general_settings": {},
|
|
"router_settings": {},
|
|
"litellm_settings": {},
|
|
}
|
|
|
|
## DB
|
|
if prisma_client is not None and (
|
|
general_settings.get("store_model_in_db", False) == True
|
|
):
|
|
_tasks = []
|
|
keys = [
|
|
"general_settings",
|
|
"router_settings",
|
|
"litellm_settings",
|
|
"environment_variables",
|
|
]
|
|
for k in keys:
|
|
response = prisma_client.get_generic_data(
|
|
key="param_name", value=k, table_name="config"
|
|
)
|
|
_tasks.append(response)
|
|
|
|
responses = await asyncio.gather(*_tasks)
|
|
for response in responses:
|
|
if response is not None:
|
|
param_name = getattr(response, "param_name", None)
|
|
param_value = getattr(response, "param_value", None)
|
|
if param_name is not None and param_value is not None:
|
|
# check if param_name is already in the config
|
|
if param_name in config:
|
|
if isinstance(config[param_name], dict):
|
|
config[param_name].update(param_value)
|
|
else:
|
|
config[param_name] = param_value
|
|
else:
|
|
# if it's not in the config - then add it
|
|
config[param_name] = param_value
|
|
|
|
return config
|
|
|
|
async def save_config(self, new_config: dict):
|
|
global prisma_client, general_settings, user_config_file_path
|
|
# Load existing config
|
|
## DB - writes valid config to db
|
|
"""
|
|
- Do not write restricted params like 'api_key' to the database
|
|
- if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`)
|
|
"""
|
|
if prisma_client is not None and (
|
|
general_settings.get("store_model_in_db", False) == True
|
|
):
|
|
# if using - db for config - models are in ModelTable
|
|
new_config.pop("model_list", None)
|
|
await prisma_client.insert_data(data=new_config, table_name="config")
|
|
else:
|
|
# Save the updated config - if user is not using a dB
|
|
## YAML
|
|
with open(f"{user_config_file_path}", "w") as config_file:
|
|
yaml.dump(new_config, config_file, default_flow_style=False)
|
|
|
|
async def load_team_config(self, team_id: str):
|
|
"""
|
|
- for a given team id
|
|
- return the relevant completion() call params
|
|
"""
|
|
# load existing config
|
|
config = await self.get_config()
|
|
## LITELLM MODULE SETTINGS (e.g. litellm.drop_params=True,..)
|
|
litellm_settings = config.get("litellm_settings", {})
|
|
all_teams_config = litellm_settings.get("default_team_settings", None)
|
|
team_config: dict = {}
|
|
if all_teams_config is None:
|
|
return team_config
|
|
for team in all_teams_config:
|
|
if "team_id" not in team:
|
|
raise Exception(f"team_id missing from team: {team}")
|
|
if team_id == team["team_id"]:
|
|
team_config = team
|
|
break
|
|
for k, v in team_config.items():
|
|
if isinstance(v, str) and v.startswith("os.environ/"):
|
|
team_config[k] = litellm.get_secret(v)
|
|
return team_config
|
|
|
|
def _init_cache(
|
|
self,
|
|
cache_params: dict,
|
|
):
|
|
global redis_usage_cache
|
|
from litellm import Cache
|
|
|
|
if "default_in_memory_ttl" in cache_params:
|
|
litellm.default_in_memory_ttl = cache_params["default_in_memory_ttl"]
|
|
|
|
if "default_redis_ttl" in cache_params:
|
|
litellm.default_redis_ttl = cache_params["default_in_redis_ttl"]
|
|
|
|
litellm.cache = Cache(**cache_params)
|
|
|
|
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
|
|
## INIT PROXY REDIS USAGE CLIENT ##
|
|
redis_usage_cache = litellm.cache.cache
|
|
|
|
async def load_config(
|
|
self, router: Optional[litellm.Router], config_file_path: str
|
|
):
|
|
"""
|
|
Load config values into proxy global state
|
|
"""
|
|
global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db
|
|
|
|
# Load existing config
|
|
config = await self.get_config(config_file_path=config_file_path)
|
|
## PRINT YAML FOR CONFIRMING IT WORKS
|
|
printed_yaml = copy.deepcopy(config)
|
|
printed_yaml.pop("environment_variables", None)
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"Loaded config YAML (api_key and environment_variables are not shown):\n{json.dumps(printed_yaml, indent=2)}"
|
|
)
|
|
|
|
## ENVIRONMENT VARIABLES
|
|
environment_variables = config.get("environment_variables", None)
|
|
if environment_variables:
|
|
for key, value in environment_variables.items():
|
|
os.environ[key] = value
|
|
|
|
## LITELLM MODULE SETTINGS (e.g. litellm.drop_params=True,..)
|
|
litellm_settings = config.get("litellm_settings", None)
|
|
if litellm_settings is None:
|
|
litellm_settings = {}
|
|
if litellm_settings:
|
|
# ANSI escape code for blue text
|
|
blue_color_code = "\033[94m"
|
|
reset_color_code = "\033[0m"
|
|
for key, value in litellm_settings.items():
|
|
if key == "cache" and value == True:
|
|
print(f"{blue_color_code}\nSetting Cache on Proxy") # noqa
|
|
from litellm.caching import Cache
|
|
|
|
cache_params = {}
|
|
if "cache_params" in litellm_settings:
|
|
cache_params_in_config = litellm_settings["cache_params"]
|
|
# overwrie cache_params with cache_params_in_config
|
|
cache_params.update(cache_params_in_config)
|
|
|
|
cache_type = cache_params.get("type", "redis")
|
|
|
|
verbose_proxy_logger.debug("passed cache type=%s", cache_type)
|
|
|
|
if (
|
|
cache_type == "redis" or cache_type == "redis-semantic"
|
|
) and len(cache_params.keys()) == 0:
|
|
cache_host = litellm.get_secret("REDIS_HOST", None)
|
|
cache_port = litellm.get_secret("REDIS_PORT", None)
|
|
cache_password = None
|
|
cache_params.update(
|
|
{
|
|
"type": cache_type,
|
|
"host": cache_host,
|
|
"port": cache_port,
|
|
}
|
|
)
|
|
|
|
if litellm.get_secret("REDIS_PASSWORD", None) is not None:
|
|
cache_password = litellm.get_secret("REDIS_PASSWORD", None)
|
|
cache_params.update(
|
|
{
|
|
"password": cache_password,
|
|
}
|
|
)
|
|
|
|
# Assuming cache_type, cache_host, cache_port, and cache_password are strings
|
|
print( # noqa
|
|
f"{blue_color_code}Cache Type:{reset_color_code} {cache_type}"
|
|
) # noqa
|
|
print( # noqa
|
|
f"{blue_color_code}Cache Host:{reset_color_code} {cache_host}"
|
|
) # noqa
|
|
print( # noqa
|
|
f"{blue_color_code}Cache Port:{reset_color_code} {cache_port}"
|
|
) # noqa
|
|
print( # noqa
|
|
f"{blue_color_code}Cache Password:{reset_color_code} {cache_password}"
|
|
)
|
|
print() # noqa
|
|
if cache_type == "redis-semantic":
|
|
# by default this should always be async
|
|
cache_params.update({"redis_semantic_cache_use_async": True})
|
|
|
|
# users can pass os.environ/ variables on the proxy - we should read them from the env
|
|
for key, value in cache_params.items():
|
|
if type(value) is str and value.startswith("os.environ/"):
|
|
cache_params[key] = litellm.get_secret(value)
|
|
|
|
## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = <your-redis-url>`, _redis.py checks for REDIS specific environment variables
|
|
self._init_cache(cache_params=cache_params)
|
|
if litellm.cache is not None:
|
|
print( # noqa
|
|
f"{blue_color_code}Set Cache on LiteLLM Proxy: {vars(litellm.cache.cache)}{reset_color_code}"
|
|
)
|
|
elif key == "cache" and value == False:
|
|
pass
|
|
elif key == "callbacks":
|
|
if isinstance(value, list):
|
|
imported_list: List[Any] = []
|
|
for callback in value: # ["presidio", <my-custom-callback>]
|
|
if isinstance(callback, str) and callback == "presidio":
|
|
from litellm.proxy.hooks.presidio_pii_masking import (
|
|
_OPTIONAL_PresidioPIIMasking,
|
|
)
|
|
|
|
pii_masking_object = _OPTIONAL_PresidioPIIMasking()
|
|
imported_list.append(pii_masking_object)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "llamaguard_moderations"
|
|
):
|
|
from enterprise.enterprise_hooks.llama_guard import (
|
|
_ENTERPRISE_LlamaGuard,
|
|
)
|
|
|
|
llama_guard_object = _ENTERPRISE_LlamaGuard()
|
|
imported_list.append(llama_guard_object)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "google_text_moderation"
|
|
):
|
|
from enterprise.enterprise_hooks.google_text_moderation import (
|
|
_ENTERPRISE_GoogleTextModeration,
|
|
)
|
|
|
|
google_text_moderation_obj = (
|
|
_ENTERPRISE_GoogleTextModeration()
|
|
)
|
|
imported_list.append(google_text_moderation_obj)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "llmguard_moderations"
|
|
):
|
|
from enterprise.enterprise_hooks.llm_guard import (
|
|
_ENTERPRISE_LLMGuard,
|
|
)
|
|
|
|
llm_guard_moderation_obj = _ENTERPRISE_LLMGuard()
|
|
imported_list.append(llm_guard_moderation_obj)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "blocked_user_check"
|
|
):
|
|
from enterprise.enterprise_hooks.blocked_user_list import (
|
|
_ENTERPRISE_BlockedUserList,
|
|
)
|
|
|
|
blocked_user_list = _ENTERPRISE_BlockedUserList(
|
|
prisma_client=prisma_client
|
|
)
|
|
imported_list.append(blocked_user_list)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "banned_keywords"
|
|
):
|
|
from enterprise.enterprise_hooks.banned_keywords import (
|
|
_ENTERPRISE_BannedKeywords,
|
|
)
|
|
|
|
banned_keywords_obj = _ENTERPRISE_BannedKeywords()
|
|
imported_list.append(banned_keywords_obj)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "detect_prompt_injection"
|
|
):
|
|
from litellm.proxy.hooks.prompt_injection_detection import (
|
|
_OPTIONAL_PromptInjectionDetection,
|
|
)
|
|
|
|
prompt_injection_params = None
|
|
if "prompt_injection_params" in litellm_settings:
|
|
prompt_injection_params_in_config = (
|
|
litellm_settings["prompt_injection_params"]
|
|
)
|
|
prompt_injection_params = (
|
|
LiteLLMPromptInjectionParams(
|
|
**prompt_injection_params_in_config
|
|
)
|
|
)
|
|
|
|
prompt_injection_detection_obj = (
|
|
_OPTIONAL_PromptInjectionDetection(
|
|
prompt_injection_params=prompt_injection_params,
|
|
)
|
|
)
|
|
imported_list.append(prompt_injection_detection_obj)
|
|
elif (
|
|
isinstance(callback, str)
|
|
and callback == "batch_redis_requests"
|
|
):
|
|
from litellm.proxy.hooks.batch_redis_get import (
|
|
_PROXY_BatchRedisRequests,
|
|
)
|
|
|
|
batch_redis_obj = _PROXY_BatchRedisRequests()
|
|
imported_list.append(batch_redis_obj)
|
|
else:
|
|
imported_list.append(
|
|
get_instance_fn(
|
|
value=callback,
|
|
config_file_path=config_file_path,
|
|
)
|
|
)
|
|
litellm.callbacks = imported_list # type: ignore
|
|
else:
|
|
litellm.callbacks = [
|
|
get_instance_fn(
|
|
value=value,
|
|
config_file_path=config_file_path,
|
|
)
|
|
]
|
|
verbose_proxy_logger.debug(
|
|
f"{blue_color_code} Initialized Callbacks - {litellm.callbacks} {reset_color_code}"
|
|
)
|
|
elif key == "post_call_rules":
|
|
litellm.post_call_rules = [
|
|
get_instance_fn(value=value, config_file_path=config_file_path)
|
|
]
|
|
verbose_proxy_logger.debug(
|
|
f"litellm.post_call_rules: {litellm.post_call_rules}"
|
|
)
|
|
elif key == "success_callback":
|
|
litellm.success_callback = []
|
|
|
|
# initialize success callbacks
|
|
for callback in value:
|
|
# user passed custom_callbacks.async_on_succes_logger. They need us to import a function
|
|
if "." in callback:
|
|
litellm.success_callback.append(
|
|
get_instance_fn(value=callback)
|
|
)
|
|
# these are litellm callbacks - "langfuse", "sentry", "wandb"
|
|
else:
|
|
litellm.success_callback.append(callback)
|
|
if "prometheus" in callback:
|
|
verbose_proxy_logger.debug(
|
|
"Starting Prometheus Metrics on /metrics"
|
|
)
|
|
from prometheus_client import make_asgi_app
|
|
|
|
# Add prometheus asgi middleware to route /metrics requests
|
|
metrics_app = make_asgi_app()
|
|
app.mount("/metrics", metrics_app)
|
|
print( # noqa
|
|
f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}"
|
|
) # noqa
|
|
elif key == "failure_callback":
|
|
litellm.failure_callback = []
|
|
|
|
# initialize success callbacks
|
|
for callback in value:
|
|
# user passed custom_callbacks.async_on_succes_logger. They need us to import a function
|
|
if "." in callback:
|
|
litellm.failure_callback.append(
|
|
get_instance_fn(value=callback)
|
|
)
|
|
# these are litellm callbacks - "langfuse", "sentry", "wandb"
|
|
else:
|
|
litellm.failure_callback.append(callback)
|
|
print( # noqa
|
|
f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}"
|
|
) # noqa
|
|
elif key == "cache_params":
|
|
# this is set in the cache branch
|
|
# see usage here: https://docs.litellm.ai/docs/proxy/caching
|
|
pass
|
|
elif key == "default_team_settings":
|
|
for idx, team_setting in enumerate(
|
|
value
|
|
): # run through pydantic validation
|
|
try:
|
|
TeamDefaultSettings(**team_setting)
|
|
except:
|
|
raise Exception(
|
|
f"team_id missing from default_team_settings at index={idx}\npassed in value={team_setting}"
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"
|
|
)
|
|
setattr(litellm, key, value)
|
|
elif key == "upperbound_key_generate_params":
|
|
if value is not None and isinstance(value, dict):
|
|
for _k, _v in value.items():
|
|
if isinstance(_v, str) and _v.startswith("os.environ/"):
|
|
value[_k] = litellm.get_secret(_v)
|
|
litellm.upperbound_key_generate_params = (
|
|
LiteLLM_UpperboundKeyGenerateParams(**value)
|
|
)
|
|
else:
|
|
raise Exception(
|
|
f"Invalid value set for upperbound_key_generate_params - value={value}"
|
|
)
|
|
else:
|
|
verbose_proxy_logger.debug(
|
|
f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"
|
|
)
|
|
setattr(litellm, key, value)
|
|
|
|
## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging
|
|
general_settings = config.get("general_settings", {})
|
|
if general_settings is None:
|
|
general_settings = {}
|
|
if general_settings:
|
|
### LOAD SECRET MANAGER ###
|
|
key_management_system = general_settings.get("key_management_system", None)
|
|
if key_management_system is not None:
|
|
if key_management_system == KeyManagementSystem.AZURE_KEY_VAULT.value:
|
|
### LOAD FROM AZURE KEY VAULT ###
|
|
load_from_azure_key_vault(use_azure_key_vault=True)
|
|
elif key_management_system == KeyManagementSystem.GOOGLE_KMS.value:
|
|
### LOAD FROM GOOGLE KMS ###
|
|
load_google_kms(use_google_kms=True)
|
|
elif (
|
|
key_management_system
|
|
== KeyManagementSystem.AWS_SECRET_MANAGER.value
|
|
):
|
|
### LOAD FROM AWS SECRET MANAGER ###
|
|
load_aws_secret_manager(use_aws_secret_manager=True)
|
|
else:
|
|
raise ValueError("Invalid Key Management System selected")
|
|
key_management_settings = general_settings.get(
|
|
"key_management_settings", None
|
|
)
|
|
if key_management_settings is not None:
|
|
litellm._key_management_settings = KeyManagementSettings(
|
|
**key_management_settings
|
|
)
|
|
### [DEPRECATED] LOAD FROM GOOGLE KMS ### old way of loading from google kms
|
|
use_google_kms = general_settings.get("use_google_kms", False)
|
|
load_google_kms(use_google_kms=use_google_kms)
|
|
### [DEPRECATED] LOAD FROM AZURE KEY VAULT ### old way of loading from azure secret manager
|
|
use_azure_key_vault = general_settings.get("use_azure_key_vault", False)
|
|
load_from_azure_key_vault(use_azure_key_vault=use_azure_key_vault)
|
|
### ALERTING ###
|
|
proxy_logging_obj.update_values(
|
|
alerting=general_settings.get("alerting", None),
|
|
alerting_threshold=general_settings.get("alerting_threshold", 600),
|
|
alert_types=general_settings.get("alert_types", None),
|
|
redis_cache=redis_usage_cache,
|
|
)
|
|
### CONNECT TO DATABASE ###
|
|
database_url = general_settings.get("database_url", None)
|
|
if database_url and database_url.startswith("os.environ/"):
|
|
verbose_proxy_logger.debug("GOING INTO LITELLM.GET_SECRET!")
|
|
database_url = litellm.get_secret(database_url)
|
|
verbose_proxy_logger.debug("RETRIEVED DB URL: %s", database_url)
|
|
### MASTER KEY ###
|
|
master_key = general_settings.get(
|
|
"master_key", litellm.get_secret("LITELLM_MASTER_KEY", None)
|
|
)
|
|
if master_key and master_key.startswith("os.environ/"):
|
|
master_key = litellm.get_secret(master_key)
|
|
|
|
if master_key is not None and isinstance(master_key, str):
|
|
litellm_master_key_hash = hash_token(master_key)
|
|
### STORE MODEL IN DB ### feature flag for `/model/new`
|
|
store_model_in_db = general_settings.get("store_model_in_db", False)
|
|
if store_model_in_db is None:
|
|
store_model_in_db = False
|
|
### CUSTOM API KEY AUTH ###
|
|
## pass filepath
|
|
custom_auth = general_settings.get("custom_auth", None)
|
|
if custom_auth is not None:
|
|
user_custom_auth = get_instance_fn(
|
|
value=custom_auth, config_file_path=config_file_path
|
|
)
|
|
|
|
custom_key_generate = general_settings.get("custom_key_generate", None)
|
|
if custom_key_generate is not None:
|
|
user_custom_key_generate = get_instance_fn(
|
|
value=custom_key_generate, config_file_path=config_file_path
|
|
)
|
|
## dynamodb
|
|
database_type = general_settings.get("database_type", None)
|
|
if database_type is not None and (
|
|
database_type == "dynamo_db" or database_type == "dynamodb"
|
|
):
|
|
database_args = general_settings.get("database_args", None)
|
|
### LOAD FROM os.environ/ ###
|
|
for k, v in database_args.items():
|
|
if isinstance(v, str) and v.startswith("os.environ/"):
|
|
database_args[k] = litellm.get_secret(v)
|
|
if isinstance(k, str) and k == "aws_web_identity_token":
|
|
value = database_args[k]
|
|
verbose_proxy_logger.debug(
|
|
f"Loading AWS Web Identity Token from file: {value}"
|
|
)
|
|
if os.path.exists(value):
|
|
with open(value, "r") as file:
|
|
token_content = file.read()
|
|
database_args[k] = token_content
|
|
else:
|
|
verbose_proxy_logger.info(
|
|
f"DynamoDB Loading - {value} is not a valid file path"
|
|
)
|
|
verbose_proxy_logger.debug("database_args: %s", database_args)
|
|
custom_db_client = DBClient(
|
|
custom_db_args=database_args, custom_db_type=database_type
|
|
)
|
|
## ADMIN UI ACCESS ##
|
|
ui_access_mode = general_settings.get(
|
|
"ui_access_mode", "all"
|
|
) # can be either ["admin_only" or "all"]
|
|
## BUDGET RESCHEDULER ##
|
|
proxy_budget_rescheduler_min_time = general_settings.get(
|
|
"proxy_budget_rescheduler_min_time", proxy_budget_rescheduler_min_time
|
|
)
|
|
proxy_budget_rescheduler_max_time = general_settings.get(
|
|
"proxy_budget_rescheduler_max_time", proxy_budget_rescheduler_max_time
|
|
)
|
|
## BATCH WRITER ##
|
|
proxy_batch_write_at = general_settings.get(
|
|
"proxy_batch_write_at", proxy_batch_write_at
|
|
)
|
|
## DISABLE SPEND LOGS ## - gives a perf improvement
|
|
disable_spend_logs = general_settings.get(
|
|
"disable_spend_logs", disable_spend_logs
|
|
)
|
|
### BACKGROUND HEALTH CHECKS ###
|
|
# Enable background health checks
|
|
use_background_health_checks = general_settings.get(
|
|
"background_health_checks", False
|
|
)
|
|
health_check_interval = general_settings.get("health_check_interval", 300)
|
|
|
|
router_params: dict = {
|
|
"cache_responses": litellm.cache
|
|
!= None, # cache if user passed in cache values
|
|
}
|
|
## MODEL LIST
|
|
model_list = config.get("model_list", None)
|
|
if model_list:
|
|
router_params["model_list"] = model_list
|
|
print( # noqa
|
|
f"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m"
|
|
) # noqa
|
|
for model in model_list:
|
|
### LOAD FROM os.environ/ ###
|
|
for k, v in model["litellm_params"].items():
|
|
if isinstance(v, str) and v.startswith("os.environ/"):
|
|
model["litellm_params"][k] = litellm.get_secret(v)
|
|
print(f"\033[32m {model.get('model_name', '')}\033[0m") # noqa
|
|
litellm_model_name = model["litellm_params"]["model"]
|
|
litellm_model_api_base = model["litellm_params"].get("api_base", None)
|
|
if "ollama" in litellm_model_name and litellm_model_api_base is None:
|
|
run_ollama_serve()
|
|
|
|
## ROUTER SETTINGS (e.g. routing_strategy, ...)
|
|
router_settings = config.get("router_settings", None)
|
|
if router_settings and isinstance(router_settings, dict):
|
|
arg_spec = inspect.getfullargspec(litellm.Router)
|
|
# model list already set
|
|
exclude_args = {
|
|
"self",
|
|
"model_list",
|
|
}
|
|
|
|
available_args = [x for x in arg_spec.args if x not in exclude_args]
|
|
|
|
for k, v in router_settings.items():
|
|
if k in available_args:
|
|
router_params[k] = v
|
|
router = litellm.Router(**router_params, semaphore=semaphore) # type:ignore
|
|
return router, model_list, general_settings
|
|
|
|
def get_model_info_with_id(self, model) -> RouterModelInfo:
|
|
"""
|
|
Common logic across add + delete router models
|
|
Parameters:
|
|
- deployment
|
|
|
|
Return model info w/ id
|
|
"""
|
|
if model.model_info is not None and isinstance(model.model_info, dict):
|
|
if "id" not in model.model_info:
|
|
model.model_info["id"] = model.model_id
|
|
_model_info = RouterModelInfo(**model.model_info)
|
|
else:
|
|
_model_info = RouterModelInfo(id=model.model_id)
|
|
return _model_info
|
|
|
|
async def _delete_deployment(self, db_models: list) -> int:
|
|
"""
|
|
(Helper function of add deployment) -> combined to reduce prisma db calls
|
|
|
|
- Create all up list of model id's (db + config)
|
|
- Compare all up list to router model id's
|
|
- Remove any that are missing
|
|
|
|
Return:
|
|
- int - returns number of deleted deployments
|
|
"""
|
|
global user_config_file_path, llm_router
|
|
combined_id_list = []
|
|
if llm_router is None:
|
|
return 0
|
|
|
|
## DB MODELS ##
|
|
for m in db_models:
|
|
model_info = self.get_model_info_with_id(model=m)
|
|
if model_info.id is not None:
|
|
combined_id_list.append(model_info.id)
|
|
|
|
## CONFIG MODELS ##
|
|
config = await self.get_config(config_file_path=user_config_file_path)
|
|
model_list = config.get("model_list", None)
|
|
if model_list:
|
|
for model in model_list:
|
|
### LOAD FROM os.environ/ ###
|
|
for k, v in model["litellm_params"].items():
|
|
if isinstance(v, str) and v.startswith("os.environ/"):
|
|
model["litellm_params"][k] = litellm.get_secret(v)
|
|
|
|
## check if they have model-id's ##
|
|
model_id = model.get("model_info", {}).get("id", None)
|
|
if model_id is None:
|
|
## else - generate stable id's ##
|
|
model_id = llm_router._generate_model_id(
|
|
model_group=model["model_name"],
|
|
litellm_params=model["litellm_params"],
|
|
)
|
|
combined_id_list.append(model_id) # ADD CONFIG MODEL TO COMBINED LIST
|
|
|
|
router_model_ids = llm_router.get_model_ids()
|
|
# Check for model IDs in llm_router not present in combined_id_list and delete them
|
|
|
|
deleted_deployments = 0
|
|
for model_id in router_model_ids:
|
|
if model_id not in combined_id_list:
|
|
is_deleted = llm_router.delete_deployment(id=model_id)
|
|
if is_deleted is not None:
|
|
deleted_deployments += 1
|
|
return deleted_deployments
|
|
|
|
def _add_deployment(self, db_models: list) -> int:
|
|
"""
|
|
Iterate through db models
|
|
|
|
for any not in router - add them.
|
|
|
|
Return - number of deployments added
|
|
"""
|
|
import base64
|
|
|
|
if master_key is None or not isinstance(master_key, str):
|
|
raise Exception(
|
|
f"Master key is not initialized or formatted. master_key={master_key}"
|
|
)
|
|
|
|
if llm_router is None:
|
|
return 0
|
|
|
|
added_models = 0
|
|
## ADD MODEL LOGIC
|
|
for m in db_models:
|
|
_litellm_params = m.litellm_params
|
|
if isinstance(_litellm_params, dict):
|
|
# decrypt values
|
|
for k, v in _litellm_params.items():
|
|
if isinstance(v, str):
|
|
# decode base64
|
|
decoded_b64 = base64.b64decode(v)
|
|
# decrypt value
|
|
_value = decrypt_value(value=decoded_b64, master_key=master_key)
|
|
# sanity check if string > size 0
|
|
if len(_value) > 0:
|
|
_litellm_params[k] = _value
|
|
_litellm_params = LiteLLM_Params(**_litellm_params)
|
|
else:
|
|
verbose_proxy_logger.error(
|
|
f"Invalid model added to proxy db. Invalid litellm params. litellm_params={_litellm_params}"
|
|
)
|
|
continue # skip to next model
|
|
_model_info = self.get_model_info_with_id(model=m)
|
|
|
|
added = llm_router.add_deployment(
|
|
deployment=Deployment(
|
|
model_name=m.model_name,
|
|
litellm_params=_litellm_params,
|
|
model_info=_model_info,
|
|
)
|
|
)
|
|
|
|
if added is not None:
|
|
added_models += 1
|
|
return added_models
|
|
|
|
async def _update_llm_router(
|
|
self,
|
|
new_models: list,
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
global llm_router, llm_model_list, master_key, general_settings
|
|
import base64
|
|
|
|
if llm_router is None and master_key is not None:
|
|
verbose_proxy_logger.debug(f"len new_models: {len(new_models)}")
|
|
|
|
_model_list: list = []
|
|
for m in new_models:
|
|
_litellm_params = m.litellm_params
|
|
if isinstance(_litellm_params, dict):
|
|
# decrypt values
|
|
for k, v in _litellm_params.items():
|
|
if isinstance(v, str):
|
|
# decode base64
|
|
decoded_b64 = base64.b64decode(v)
|
|
# decrypt value
|
|
_litellm_params[k] = decrypt_value(
|
|
value=decoded_b64, master_key=master_key # type: ignore
|
|
)
|
|
_litellm_params = LiteLLM_Params(**_litellm_params)
|
|
else:
|
|
verbose_proxy_logger.error(
|
|
f"Invalid model added to proxy db. Invalid litellm params. litellm_params={_litellm_params}"
|
|
)
|
|
continue # skip to next model
|
|
|
|
_model_info = self.get_model_info_with_id(model=m)
|
|
_model_list.append(
|
|
Deployment(
|
|
model_name=m.model_name,
|
|
litellm_params=_litellm_params,
|
|
model_info=_model_info,
|
|
).to_json(exclude_none=True)
|
|
)
|
|
if len(_model_list) > 0:
|
|
verbose_proxy_logger.debug(f"_model_list: {_model_list}")
|
|
llm_router = litellm.Router(model_list=_model_list)
|
|
verbose_proxy_logger.debug(f"updated llm_router: {llm_router}")
|
|
else:
|
|
verbose_proxy_logger.debug(f"len new_models: {len(new_models)}")
|
|
## DELETE MODEL LOGIC
|
|
await self._delete_deployment(db_models=new_models)
|
|
|
|
## ADD MODEL LOGIC
|
|
self._add_deployment(db_models=new_models)
|
|
|
|
if llm_router is not None:
|
|
llm_model_list = llm_router.get_model_list()
|
|
|
|
# check if user set any callbacks in Config Table
|
|
config_data = await proxy_config.get_config()
|
|
litellm_settings = config_data.get("litellm_settings", {}) or {}
|
|
success_callbacks = litellm_settings.get("success_callback", None)
|
|
|
|
if success_callbacks is not None and isinstance(success_callbacks, list):
|
|
for success_callback in success_callbacks:
|
|
if success_callback not in litellm.success_callback:
|
|
litellm.success_callback.append(success_callback)
|
|
# we need to set env variables too
|
|
environment_variables = config_data.get("environment_variables", {})
|
|
for k, v in environment_variables.items():
|
|
try:
|
|
if v is not None:
|
|
decoded_b64 = base64.b64decode(v)
|
|
value = decrypt_value(value=decoded_b64, master_key=master_key) # type: ignore
|
|
os.environ[k] = value
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(
|
|
"Error setting env variable: %s - %s", k, str(e)
|
|
)
|
|
|
|
# general_settings
|
|
_general_settings = config_data.get("general_settings", {})
|
|
if "alerting" in _general_settings:
|
|
general_settings["alerting"] = _general_settings["alerting"]
|
|
proxy_logging_obj.alerting = general_settings["alerting"]
|
|
proxy_logging_obj.slack_alerting_instance.alerting = general_settings[
|
|
"alerting"
|
|
]
|
|
|
|
if "alert_types" in _general_settings:
|
|
general_settings["alert_types"] = _general_settings["alert_types"]
|
|
proxy_logging_obj.alert_types = general_settings["alert_types"]
|
|
proxy_logging_obj.slack_alerting_instance.update_values(
|
|
alert_types=general_settings["alert_types"]
|
|
)
|
|
|
|
if "alert_to_webhook_url" in _general_settings:
|
|
general_settings["alert_to_webhook_url"] = _general_settings[
|
|
"alert_to_webhook_url"
|
|
]
|
|
proxy_logging_obj.slack_alerting_instance.update_values(
|
|
alert_to_webhook_url=general_settings["alert_to_webhook_url"]
|
|
)
|
|
|
|
# router settings
|
|
if llm_router is not None and prisma_client is not None:
|
|
db_router_settings = await prisma_client.db.litellm_config.find_first(
|
|
where={"param_name": "router_settings"}
|
|
)
|
|
if (
|
|
db_router_settings is not None
|
|
and db_router_settings.param_value is not None
|
|
):
|
|
_router_settings = db_router_settings.param_value
|
|
llm_router.update_settings(**_router_settings)
|
|
|
|
async def add_deployment(
|
|
self,
|
|
prisma_client: PrismaClient,
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
"""
|
|
- Check db for new models (last 10 most recently updated)
|
|
- Check if model id's in router already
|
|
- If not, add to router
|
|
"""
|
|
global llm_router, llm_model_list, master_key, general_settings
|
|
|
|
try:
|
|
if master_key is None or not isinstance(master_key, str):
|
|
raise Exception(
|
|
f"Master key is not initialized or formatted. master_key={master_key}"
|
|
)
|
|
verbose_proxy_logger.debug(f"llm_router: {llm_router}")
|
|
new_models = await prisma_client.db.litellm_proxymodeltable.find_many()
|
|
await self._update_llm_router(
|
|
new_models=new_models, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(
|
|
"{}\nTraceback:{}".format(str(e), traceback.format_exc())
|
|
)
|
|
|
|
|
|
proxy_config = ProxyConfig()
|
|
|
|
|
|
def _duration_in_seconds(duration: str):
|
|
match = re.match(r"(\d+)([smhd]?)", duration)
|
|
if not match:
|
|
raise ValueError("Invalid duration format")
|
|
|
|
value, unit = match.groups()
|
|
value = int(value)
|
|
|
|
if unit == "s":
|
|
return value
|
|
elif unit == "m":
|
|
return value * 60
|
|
elif unit == "h":
|
|
return value * 3600
|
|
elif unit == "d":
|
|
return value * 86400
|
|
else:
|
|
raise ValueError("Unsupported duration unit")
|
|
|
|
|
|
async def generate_key_helper_fn(
|
|
duration: Optional[str],
|
|
models: list,
|
|
aliases: dict,
|
|
config: dict,
|
|
spend: float,
|
|
key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key
|
|
key_budget_duration: Optional[str] = None,
|
|
budget_id: Optional[float] = None, # budget id <-> LiteLLM_BudgetTable
|
|
soft_budget: Optional[
|
|
float
|
|
] = None, # soft_budget is used to set soft Budgets Per user
|
|
max_budget: Optional[float] = None, # max_budget is used to Budget Per user
|
|
budget_duration: Optional[str] = None, # max_budget is used to Budget Per user
|
|
token: Optional[str] = None,
|
|
user_id: Optional[str] = None,
|
|
team_id: Optional[str] = None,
|
|
user_email: Optional[str] = None,
|
|
user_role: Optional[str] = None,
|
|
max_parallel_requests: Optional[int] = None,
|
|
metadata: Optional[dict] = {},
|
|
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,
|
|
allowed_cache_controls: Optional[list] = [],
|
|
permissions: Optional[dict] = {},
|
|
model_max_budget: Optional[dict] = {},
|
|
teams: Optional[list] = None,
|
|
organization_id: Optional[str] = None,
|
|
table_name: Optional[Literal["key", "user"]] = None,
|
|
):
|
|
global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name
|
|
|
|
if prisma_client is None and custom_db_client is None:
|
|
raise Exception(
|
|
f"Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys "
|
|
)
|
|
|
|
if token is None:
|
|
token = f"sk-{secrets.token_urlsafe(16)}"
|
|
|
|
if duration is None: # allow tokens that never expire
|
|
expires = None
|
|
else:
|
|
duration_s = _duration_in_seconds(duration=duration)
|
|
expires = datetime.utcnow() + timedelta(seconds=duration_s)
|
|
|
|
if key_budget_duration is None: # one-time budget
|
|
key_reset_at = None
|
|
else:
|
|
duration_s = _duration_in_seconds(duration=key_budget_duration)
|
|
key_reset_at = datetime.utcnow() + timedelta(seconds=duration_s)
|
|
|
|
if budget_duration is None: # one-time budget
|
|
reset_at = None
|
|
else:
|
|
duration_s = _duration_in_seconds(duration=budget_duration)
|
|
reset_at = datetime.utcnow() + timedelta(seconds=duration_s)
|
|
|
|
aliases_json = json.dumps(aliases)
|
|
config_json = json.dumps(config)
|
|
permissions_json = json.dumps(permissions)
|
|
metadata_json = json.dumps(metadata)
|
|
model_max_budget_json = json.dumps(model_max_budget)
|
|
user_role = user_role or "app_user"
|
|
tpm_limit = tpm_limit
|
|
rpm_limit = rpm_limit
|
|
allowed_cache_controls = allowed_cache_controls
|
|
|
|
try:
|
|
# Create a new verification token (you may want to enhance this logic based on your needs)
|
|
user_data = {
|
|
"max_budget": max_budget,
|
|
"user_email": user_email,
|
|
"user_id": user_id,
|
|
"team_id": team_id,
|
|
"organization_id": organization_id,
|
|
"user_role": user_role,
|
|
"spend": spend,
|
|
"models": models,
|
|
"max_parallel_requests": max_parallel_requests,
|
|
"tpm_limit": tpm_limit,
|
|
"rpm_limit": rpm_limit,
|
|
"budget_duration": budget_duration,
|
|
"budget_reset_at": reset_at,
|
|
"allowed_cache_controls": allowed_cache_controls,
|
|
}
|
|
if teams is not None:
|
|
user_data["teams"] = teams
|
|
key_data = {
|
|
"token": token,
|
|
"key_alias": key_alias,
|
|
"expires": expires,
|
|
"models": models,
|
|
"aliases": aliases_json,
|
|
"config": config_json,
|
|
"spend": spend,
|
|
"max_budget": key_max_budget,
|
|
"user_id": user_id,
|
|
"team_id": team_id,
|
|
"max_parallel_requests": max_parallel_requests,
|
|
"metadata": metadata_json,
|
|
"tpm_limit": tpm_limit,
|
|
"rpm_limit": rpm_limit,
|
|
"budget_duration": key_budget_duration,
|
|
"budget_reset_at": key_reset_at,
|
|
"allowed_cache_controls": allowed_cache_controls,
|
|
"permissions": permissions_json,
|
|
"model_max_budget": model_max_budget_json,
|
|
"budget_id": budget_id,
|
|
}
|
|
|
|
if (
|
|
litellm.get_secret("DISABLE_KEY_NAME", False) == True
|
|
): # allow user to disable storing abbreviated key name (shown in UI, to help figure out which key spent how much)
|
|
pass
|
|
else:
|
|
key_data["key_name"] = f"sk-...{token[-4:]}"
|
|
saved_token = copy.deepcopy(key_data)
|
|
if isinstance(saved_token["aliases"], str):
|
|
saved_token["aliases"] = json.loads(saved_token["aliases"])
|
|
if isinstance(saved_token["config"], str):
|
|
saved_token["config"] = json.loads(saved_token["config"])
|
|
if isinstance(saved_token["metadata"], str):
|
|
saved_token["metadata"] = json.loads(saved_token["metadata"])
|
|
if isinstance(saved_token["permissions"], str):
|
|
saved_token["permissions"] = json.loads(saved_token["permissions"])
|
|
if isinstance(saved_token["model_max_budget"], str):
|
|
saved_token["model_max_budget"] = json.loads(
|
|
saved_token["model_max_budget"]
|
|
)
|
|
|
|
if saved_token.get("expires", None) is not None and isinstance(
|
|
saved_token["expires"], datetime
|
|
):
|
|
saved_token["expires"] = saved_token["expires"].isoformat()
|
|
if prisma_client is not None:
|
|
if (
|
|
table_name is None or table_name == "user"
|
|
): # do not auto-create users for `/key/generate`
|
|
## CREATE USER (If necessary)
|
|
if query_type == "insert_data":
|
|
user_row = await prisma_client.insert_data(
|
|
data=user_data, table_name="user"
|
|
)
|
|
## use default user model list if no key-specific model list provided
|
|
if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore
|
|
key_data["models"] = user_row.models
|
|
elif query_type == "update_data":
|
|
user_row = await prisma_client.update_data(
|
|
data=user_data,
|
|
table_name="user",
|
|
update_key_values=update_key_values,
|
|
)
|
|
if user_id == litellm_proxy_budget_name or (
|
|
table_name is not None and table_name == "user"
|
|
):
|
|
# do not create a key for litellm_proxy_budget_name or if table name is set to just 'user'
|
|
# we only need to ensure this exists in the user table
|
|
# the LiteLLM_VerificationToken table will increase in size if we don't do this check
|
|
return user_data
|
|
|
|
## CREATE KEY
|
|
verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data)
|
|
create_key_response = await prisma_client.insert_data(
|
|
data=key_data, table_name="key"
|
|
)
|
|
key_data["token_id"] = getattr(create_key_response, "token", None)
|
|
elif custom_db_client is not None:
|
|
if table_name is None or table_name == "user":
|
|
## CREATE USER (If necessary)
|
|
verbose_proxy_logger.debug(
|
|
"CustomDBClient: Creating User= %s", user_data
|
|
)
|
|
user_row = await custom_db_client.insert_data(
|
|
value=user_data, table_name="user"
|
|
)
|
|
if user_row is None:
|
|
# GET USER ROW
|
|
user_row = await custom_db_client.get_data(
|
|
key=user_id, table_name="user" # type: ignore
|
|
)
|
|
|
|
## use default user model list if no key-specific model list provided
|
|
if len(user_row.models) > 0 and len(key_data["models"]) == 0: # type: ignore
|
|
key_data["models"] = user_row.models
|
|
## CREATE KEY
|
|
verbose_proxy_logger.debug("CustomDBClient: Creating Key= %s", key_data)
|
|
await custom_db_client.insert_data(value=key_data, table_name="key")
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
if isinstance(e, HTTPException):
|
|
raise e
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail={"error": "Internal Server Error."},
|
|
)
|
|
|
|
# Add budget related info in key_data - this ensures it's returned
|
|
key_data["budget_id"] = budget_id
|
|
return key_data
|
|
|
|
|
|
async def delete_verification_token(tokens: List, user_id: Optional[str] = None):
|
|
global prisma_client
|
|
try:
|
|
if prisma_client:
|
|
# Assuming 'db' is your Prisma Client instance
|
|
# check if admin making request - don't filter by user-id
|
|
if user_id == litellm_proxy_admin_name:
|
|
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
|
|
# else
|
|
else:
|
|
deleted_tokens = await prisma_client.delete_data(
|
|
tokens=tokens, user_id=user_id
|
|
)
|
|
_num_deleted_tokens = deleted_tokens.get("deleted_keys", 0)
|
|
if _num_deleted_tokens != len(tokens):
|
|
raise Exception(
|
|
"Failed to delete all tokens. Tried to delete tokens that don't belong to user: "
|
|
+ str(user_id)
|
|
)
|
|
else:
|
|
raise Exception("DB not connected. prisma_client is None")
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
raise e
|
|
return deleted_tokens
|
|
|
|
|
|
def save_worker_config(**data):
|
|
import json
|
|
|
|
os.environ["WORKER_CONFIG"] = json.dumps(data)
|
|
|
|
|
|
async def initialize(
|
|
model=None,
|
|
alias=None,
|
|
api_base=None,
|
|
api_version=None,
|
|
debug=False,
|
|
detailed_debug=False,
|
|
temperature=None,
|
|
max_tokens=None,
|
|
request_timeout=600,
|
|
max_budget=None,
|
|
telemetry=False,
|
|
drop_params=True,
|
|
add_function_to_prompt=True,
|
|
headers=None,
|
|
save=False,
|
|
use_queue=False,
|
|
config=None,
|
|
):
|
|
global user_model, user_api_base, user_debug, user_detailed_debug, user_user_max_tokens, user_request_timeout, user_temperature, user_telemetry, user_headers, experimental, llm_model_list, llm_router, general_settings, master_key, user_custom_auth, prisma_client
|
|
generate_feedback_box()
|
|
user_model = model
|
|
user_debug = debug
|
|
if debug == True: # this needs to be first, so users can see Router init debugg
|
|
from litellm._logging import (
|
|
verbose_router_logger,
|
|
verbose_proxy_logger,
|
|
verbose_logger,
|
|
)
|
|
import logging
|
|
|
|
# this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS
|
|
verbose_logger.setLevel(level=logging.INFO) # sets package logs to info
|
|
verbose_router_logger.setLevel(level=logging.INFO) # set router logs to info
|
|
verbose_proxy_logger.setLevel(level=logging.INFO) # set proxy logs to info
|
|
if detailed_debug == True:
|
|
from litellm._logging import (
|
|
verbose_router_logger,
|
|
verbose_proxy_logger,
|
|
verbose_logger,
|
|
)
|
|
import logging
|
|
|
|
verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug
|
|
verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to debug
|
|
verbose_proxy_logger.setLevel(level=logging.DEBUG) # set proxy logs to debug
|
|
elif debug == False and detailed_debug == False:
|
|
# users can control proxy debugging using env variable = 'LITELLM_LOG'
|
|
litellm_log_setting = os.environ.get("LITELLM_LOG", "")
|
|
if litellm_log_setting != None:
|
|
if litellm_log_setting.upper() == "INFO":
|
|
from litellm._logging import verbose_router_logger, verbose_proxy_logger
|
|
import logging
|
|
|
|
# this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS
|
|
|
|
verbose_router_logger.setLevel(
|
|
level=logging.INFO
|
|
) # set router logs to info
|
|
verbose_proxy_logger.setLevel(
|
|
level=logging.INFO
|
|
) # set proxy logs to info
|
|
elif litellm_log_setting.upper() == "DEBUG":
|
|
from litellm._logging import verbose_router_logger, verbose_proxy_logger
|
|
import logging
|
|
|
|
verbose_router_logger.setLevel(
|
|
level=logging.DEBUG
|
|
) # set router logs to info
|
|
verbose_proxy_logger.setLevel(
|
|
level=logging.DEBUG
|
|
) # set proxy logs to debug
|
|
dynamic_config = {"general": {}, user_model: {}}
|
|
if config:
|
|
(
|
|
llm_router,
|
|
llm_model_list,
|
|
general_settings,
|
|
) = await proxy_config.load_config(router=llm_router, config_file_path=config)
|
|
if headers: # model-specific param
|
|
user_headers = headers
|
|
dynamic_config[user_model]["headers"] = headers
|
|
if api_base: # model-specific param
|
|
user_api_base = api_base
|
|
dynamic_config[user_model]["api_base"] = api_base
|
|
if api_version:
|
|
os.environ["AZURE_API_VERSION"] = (
|
|
api_version # set this for azure - litellm can read this from the env
|
|
)
|
|
if max_tokens: # model-specific param
|
|
user_max_tokens = max_tokens
|
|
dynamic_config[user_model]["max_tokens"] = max_tokens
|
|
if temperature: # model-specific param
|
|
user_temperature = temperature
|
|
dynamic_config[user_model]["temperature"] = temperature
|
|
if request_timeout:
|
|
user_request_timeout = request_timeout
|
|
dynamic_config[user_model]["request_timeout"] = request_timeout
|
|
if alias: # model-specific param
|
|
dynamic_config[user_model]["alias"] = alias
|
|
if drop_params == True: # litellm-specific param
|
|
litellm.drop_params = True
|
|
dynamic_config["general"]["drop_params"] = True
|
|
if add_function_to_prompt == True: # litellm-specific param
|
|
litellm.add_function_to_prompt = True
|
|
dynamic_config["general"]["add_function_to_prompt"] = True
|
|
if max_budget: # litellm-specific param
|
|
litellm.max_budget = max_budget
|
|
dynamic_config["general"]["max_budget"] = max_budget
|
|
if experimental:
|
|
pass
|
|
user_telemetry = telemetry
|
|
|
|
|
|
# for streaming
|
|
def data_generator(response):
|
|
verbose_proxy_logger.debug("inside generator")
|
|
for chunk in response:
|
|
verbose_proxy_logger.debug("returned chunk: %s", chunk)
|
|
try:
|
|
yield f"data: {json.dumps(chunk.dict())}\n\n"
|
|
except:
|
|
yield f"data: {json.dumps(chunk)}\n\n"
|
|
|
|
|
|
async def async_data_generator(response, user_api_key_dict):
|
|
verbose_proxy_logger.debug("inside generator")
|
|
try:
|
|
start_time = time.time()
|
|
async for chunk in response:
|
|
chunk = chunk.model_dump_json(exclude_none=True)
|
|
try:
|
|
yield f"data: {chunk}\n\n"
|
|
except Exception as e:
|
|
yield f"data: {str(e)}\n\n"
|
|
|
|
# Streaming is done, yield the [DONE] chunk
|
|
done_message = "[DONE]"
|
|
yield f"data: {done_message}\n\n"
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"\033[1;31mAn error occurred: {e}\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`"
|
|
)
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
if user_debug:
|
|
traceback.print_exc()
|
|
|
|
if isinstance(e, HTTPException):
|
|
raise e
|
|
else:
|
|
error_traceback = traceback.format_exc()
|
|
error_msg = f"{str(e)}\n\n{error_traceback}"
|
|
|
|
proxy_exception = ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
error_returned = json.dumps({"error": proxy_exception.to_dict()})
|
|
yield f"data: {error_returned}\n\n"
|
|
|
|
|
|
def select_data_generator(response, user_api_key_dict):
|
|
return async_data_generator(response=response, user_api_key_dict=user_api_key_dict)
|
|
|
|
|
|
def get_litellm_model_info(model: dict = {}):
|
|
model_info = model.get("model_info", {})
|
|
model_to_lookup = model.get("litellm_params", {}).get("model", None)
|
|
try:
|
|
if "azure" in model_to_lookup:
|
|
model_to_lookup = model_info.get("base_model", None)
|
|
litellm_model_info = litellm.get_model_info(model_to_lookup)
|
|
return litellm_model_info
|
|
except:
|
|
# this should not block returning on /model/info
|
|
# if litellm does not have info on the model it should return {}
|
|
return {}
|
|
|
|
|
|
def parse_cache_control(cache_control):
|
|
cache_dict = {}
|
|
directives = cache_control.split(", ")
|
|
|
|
for directive in directives:
|
|
if "=" in directive:
|
|
key, value = directive.split("=")
|
|
cache_dict[key] = value
|
|
else:
|
|
cache_dict[directive] = True
|
|
|
|
return cache_dict
|
|
|
|
|
|
def on_backoff(details):
|
|
# The 'tries' key in the details dictionary contains the number of completed tries
|
|
verbose_proxy_logger.debug("Backing off... this was attempt # %s", details["tries"])
|
|
|
|
|
|
@router.on_event("startup")
|
|
async def startup_event():
|
|
global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db
|
|
import json
|
|
|
|
### LOAD MASTER KEY ###
|
|
# check if master key set in environment - load from there
|
|
master_key = litellm.get_secret("LITELLM_MASTER_KEY", None)
|
|
# check if DATABASE_URL in environment - load from there
|
|
if prisma_client is None:
|
|
prisma_setup(database_url=os.getenv("DATABASE_URL"))
|
|
|
|
### LOAD CONFIG ###
|
|
worker_config = litellm.get_secret("WORKER_CONFIG")
|
|
verbose_proxy_logger.debug("worker_config: %s", worker_config)
|
|
# check if it's a valid file path
|
|
if os.path.isfile(worker_config):
|
|
if proxy_config.is_yaml(config_file_path=worker_config):
|
|
(
|
|
llm_router,
|
|
llm_model_list,
|
|
general_settings,
|
|
) = await proxy_config.load_config(
|
|
router=llm_router, config_file_path=worker_config
|
|
)
|
|
else:
|
|
await initialize(**worker_config)
|
|
else:
|
|
# if not, assume it's a json string
|
|
worker_config = json.loads(os.getenv("WORKER_CONFIG"))
|
|
await initialize(**worker_config)
|
|
|
|
## COST TRACKING ##
|
|
cost_tracking()
|
|
|
|
## Error Tracking ##
|
|
error_tracking()
|
|
|
|
db_writer_client = HTTPHandler()
|
|
|
|
proxy_logging_obj._init_litellm_callbacks() # INITIALIZE LITELLM CALLBACKS ON SERVER STARTUP <- do this to catch any logging errors on startup, not when calls are being made
|
|
|
|
## JWT AUTH ##
|
|
if general_settings.get("litellm_jwtauth", None) is not None:
|
|
for k, v in general_settings["litellm_jwtauth"].items():
|
|
if isinstance(v, str) and v.startswith("os.environ/"):
|
|
general_settings["litellm_jwtauth"][k] = litellm.get_secret(v)
|
|
litellm_jwtauth = LiteLLM_JWTAuth(**general_settings["litellm_jwtauth"])
|
|
else:
|
|
litellm_jwtauth = LiteLLM_JWTAuth()
|
|
jwt_handler.update_environment(
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
litellm_jwtauth=litellm_jwtauth,
|
|
)
|
|
|
|
if use_background_health_checks:
|
|
asyncio.create_task(
|
|
_run_background_health_check()
|
|
) # start the background health check coroutine.
|
|
|
|
if prompt_injection_detection_obj is not None:
|
|
prompt_injection_detection_obj.update_environment(router=llm_router)
|
|
|
|
verbose_proxy_logger.debug("prisma_client: %s", prisma_client)
|
|
if prisma_client is not None:
|
|
await prisma_client.connect()
|
|
|
|
verbose_proxy_logger.debug("custom_db_client client - %s", custom_db_client)
|
|
if custom_db_client is not None:
|
|
verbose_proxy_logger.debug("custom_db_client: connecting %s", custom_db_client)
|
|
await custom_db_client.connect()
|
|
|
|
if prisma_client is not None and master_key is not None:
|
|
# add master key to db
|
|
if os.getenv("PROXY_ADMIN_ID", None) is not None:
|
|
litellm_proxy_admin_name = os.getenv("PROXY_ADMIN_ID")
|
|
asyncio.create_task(
|
|
generate_key_helper_fn(
|
|
duration=None,
|
|
models=[],
|
|
aliases={},
|
|
config={},
|
|
spend=0,
|
|
token=master_key,
|
|
user_id=litellm_proxy_admin_name,
|
|
user_role="proxy_admin",
|
|
query_type="update_data",
|
|
update_key_values={
|
|
"user_role": "proxy_admin",
|
|
},
|
|
)
|
|
)
|
|
|
|
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
|
|
asyncio.create_task(
|
|
generate_key_helper_fn(
|
|
user_id=litellm_proxy_budget_name,
|
|
duration=None,
|
|
models=[],
|
|
aliases={},
|
|
config={},
|
|
spend=0,
|
|
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(
|
|
f"custom_db_client client {custom_db_client}. Master_key: {master_key}"
|
|
)
|
|
if custom_db_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
|
|
)
|
|
|
|
### CHECK IF VIEW EXISTS ###
|
|
if prisma_client is not None:
|
|
create_view_response = await prisma_client.check_view_exists()
|
|
|
|
### START BATCH WRITING DB + CHECKING NEW MODELS###
|
|
if prisma_client is not None:
|
|
scheduler = AsyncIOScheduler()
|
|
interval = random.randint(
|
|
proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time
|
|
) # random interval, so multiple workers avoid resetting budget at the same time
|
|
batch_writing_interval = random.randint(
|
|
proxy_batch_write_at - 3, proxy_batch_write_at + 3
|
|
) # random interval, so multiple workers avoid batch writing at the same time
|
|
|
|
### RESET BUDGET ###
|
|
if general_settings.get("disable_reset_budget", False) == False:
|
|
scheduler.add_job(
|
|
reset_budget, "interval", seconds=interval, args=[prisma_client]
|
|
)
|
|
|
|
### UPDATE SPEND ###
|
|
scheduler.add_job(
|
|
update_spend,
|
|
"interval",
|
|
seconds=batch_writing_interval,
|
|
args=[prisma_client, db_writer_client, proxy_logging_obj],
|
|
)
|
|
|
|
### ADD NEW MODELS ###
|
|
store_model_in_db = (
|
|
litellm.get_secret("STORE_MODEL_IN_DB", store_model_in_db)
|
|
or store_model_in_db
|
|
)
|
|
if store_model_in_db == True:
|
|
scheduler.add_job(
|
|
proxy_config.add_deployment,
|
|
"interval",
|
|
seconds=10,
|
|
args=[prisma_client, proxy_logging_obj],
|
|
)
|
|
|
|
# this will load all existing models on proxy startup
|
|
await proxy_config.add_deployment(
|
|
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
scheduler.start()
|
|
|
|
|
|
#### API ENDPOINTS ####
|
|
@router.get(
|
|
"/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]
|
|
)
|
|
@router.get(
|
|
"/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]
|
|
) # if project requires model list
|
|
def model_list(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global llm_model_list, general_settings
|
|
all_models = []
|
|
if len(user_api_key_dict.models) > 0:
|
|
all_models = user_api_key_dict.models
|
|
else:
|
|
## if no specific model access
|
|
if general_settings.get("infer_model_from_keys", False):
|
|
all_models = litellm.utils.get_valid_models()
|
|
if llm_model_list:
|
|
all_models = list(
|
|
set(all_models + [m["model_name"] for m in llm_model_list])
|
|
)
|
|
if user_model is not None:
|
|
all_models += [user_model]
|
|
verbose_proxy_logger.debug("all_models: %s", all_models)
|
|
return dict(
|
|
data=[
|
|
{
|
|
"id": model,
|
|
"object": "model",
|
|
"created": 1677610602,
|
|
"owned_by": "openai",
|
|
}
|
|
for model in all_models
|
|
],
|
|
object="list",
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v1/chat/completions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["chat/completions"],
|
|
)
|
|
@router.post(
|
|
"/chat/completions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["chat/completions"],
|
|
)
|
|
@router.post(
|
|
"/openai/deployments/{model:path}/chat/completions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["chat/completions"],
|
|
) # azure compatible endpoint
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=litellm.num_retries or 3, # maximum number of retries
|
|
max_time=litellm.request_timeout or 60, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
giveup=lambda e: not (
|
|
isinstance(e, ProxyException)
|
|
and getattr(e, "message", None) is not None
|
|
and isinstance(e.message, str)
|
|
and "Max parallel request limit reached" in e.message
|
|
), # the result of the logical expression is on the second position
|
|
)
|
|
async def chat_completion(
|
|
request: Request,
|
|
fastapi_response: Response,
|
|
model: Optional[str] = None,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global general_settings, user_debug, proxy_logging_obj, llm_model_list
|
|
try:
|
|
# async with llm_router.sem
|
|
data = {}
|
|
body = await request.body()
|
|
body_str = body.decode()
|
|
try:
|
|
data = ast.literal_eval(body_str)
|
|
except:
|
|
data = json.loads(body_str)
|
|
|
|
# Azure OpenAI only: check if user passed api-version
|
|
query_params = dict(request.query_params)
|
|
if "api-version" in query_params:
|
|
data["api_version"] = query_params["api-version"]
|
|
|
|
# Include original request and headers in the data
|
|
data["proxy_server_request"] = {
|
|
"url": str(request.url),
|
|
"method": request.method,
|
|
"headers": dict(request.headers),
|
|
"body": copy.copy(data), # use copy instead of deepcopy
|
|
}
|
|
|
|
## Cache Controls
|
|
headers = request.headers
|
|
verbose_proxy_logger.debug("Request Headers: %s", headers)
|
|
cache_control_header = headers.get("Cache-Control", None)
|
|
if cache_control_header:
|
|
cache_dict = parse_cache_control(cache_control_header)
|
|
data["ttl"] = cache_dict.get("s-maxage")
|
|
|
|
verbose_proxy_logger.debug("receiving data: %s", data)
|
|
data["model"] = (
|
|
general_settings.get("completion_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or model # for azure deployments
|
|
or data["model"] # default passed in http request
|
|
)
|
|
|
|
# users can pass in 'user' param to /chat/completions. Don't override it
|
|
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
|
|
# if users are using user_api_key_auth, set `user` in `data`
|
|
data["user"] = user_api_key_dict.user_id
|
|
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["user_api_key_team_alias"] = getattr(
|
|
user_api_key_dict, "team_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
|
|
### TEAM-SPECIFIC PARAMS ###
|
|
if user_api_key_dict.team_id is not None:
|
|
team_config = await proxy_config.load_team_config(
|
|
team_id=user_api_key_dict.team_id
|
|
)
|
|
if len(team_config) == 0:
|
|
pass
|
|
else:
|
|
team_id = team_config.pop("team_id", None)
|
|
_is_valid_team_configs(
|
|
team_id=team_id, team_config=team_config, request_data=data
|
|
)
|
|
data["metadata"]["team_id"] = team_id
|
|
data = {
|
|
**team_config,
|
|
**data,
|
|
} # add the team-specific configs to the completion call
|
|
|
|
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
|
|
# override with user settings, these are params passed via cli
|
|
if user_temperature:
|
|
data["temperature"] = user_temperature
|
|
if user_request_timeout:
|
|
data["request_timeout"] = user_request_timeout
|
|
if user_max_tokens:
|
|
data["max_tokens"] = user_max_tokens
|
|
if user_api_base:
|
|
data["api_base"] = user_api_base
|
|
|
|
### MODEL ALIAS MAPPING ###
|
|
# check if model name in model alias map
|
|
# get the actual model name
|
|
if data["model"] in litellm.model_alias_map:
|
|
data["model"] = litellm.model_alias_map[data["model"]]
|
|
|
|
## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call
|
|
data["litellm_call_id"] = str(uuid.uuid4())
|
|
logging_obj, data = litellm.utils.function_setup(
|
|
original_function="acompletion",
|
|
rules_obj=litellm.utils.Rules(),
|
|
start_time=datetime.now(),
|
|
**data,
|
|
)
|
|
|
|
data["litellm_logging_obj"] = logging_obj
|
|
|
|
### CALL HOOKS ### - modify incoming data before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, call_type="completion"
|
|
)
|
|
|
|
tasks = []
|
|
tasks.append(
|
|
proxy_logging_obj.during_call_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type="completion",
|
|
)
|
|
)
|
|
|
|
### ROUTE THE REQUEST ###
|
|
# Do not change this - it should be a constant time fetch - ALWAYS
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
# skip router if user passed their key
|
|
if "api_key" in data:
|
|
tasks.append(litellm.acompletion(**data))
|
|
elif "user_config" in data:
|
|
# initialize a new router instance. make request using this Router
|
|
router_config = data.pop("user_config")
|
|
user_router = litellm.Router(**router_config)
|
|
tasks.append(user_router.acompletion(**data))
|
|
elif (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
tasks.append(llm_router.acompletion(**data))
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
tasks.append(llm_router.acompletion(**data))
|
|
elif (
|
|
llm_router is not None and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
tasks.append(llm_router.acompletion(**data, specific_deployment=True))
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] not in router_model_names
|
|
and llm_router.default_deployment is not None
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
tasks.append(llm_router.acompletion(**data))
|
|
elif user_model is not None: # `litellm --model <your-model-name>`
|
|
tasks.append(litellm.acompletion(**data))
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={
|
|
"error": "chat_completion: Invalid model name passed in model="
|
|
+ data.get("model", "")
|
|
},
|
|
)
|
|
|
|
# wait for call to end
|
|
responses = await asyncio.gather(
|
|
*tasks
|
|
) # run the moderation check in parallel to the actual llm api call
|
|
response = responses[1]
|
|
|
|
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
|
model_id = hidden_params.get("model_id", None) or ""
|
|
cache_key = hidden_params.get("cache_key", None) or ""
|
|
|
|
# Post Call Processing
|
|
if llm_router is not None:
|
|
data["deployment"] = llm_router.get_deployment(model_id=model_id)
|
|
data["litellm_status"] = "success" # used for alerting
|
|
|
|
if (
|
|
"stream" in data and data["stream"] == True
|
|
): # use generate_responses to stream responses
|
|
custom_headers = {
|
|
"x-litellm-model-id": model_id,
|
|
"x-litellm-cache-key": cache_key,
|
|
}
|
|
selected_data_generator = select_data_generator(
|
|
response=response, user_api_key_dict=user_api_key_dict
|
|
)
|
|
return StreamingResponse(
|
|
selected_data_generator,
|
|
media_type="text/event-stream",
|
|
headers=custom_headers,
|
|
)
|
|
|
|
fastapi_response.headers["x-litellm-model-id"] = model_id
|
|
fastapi_response.headers["x-litellm-cache-key"] = cache_key
|
|
|
|
### CALL HOOKS ### - modify outgoing data
|
|
response = await proxy_logging_obj.post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict, response=response
|
|
)
|
|
|
|
return response
|
|
except Exception as e:
|
|
data["litellm_status"] = "fail" # used for alerting
|
|
traceback.print_exc()
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"\033[1;31mAn error occurred: {e}\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`"
|
|
)
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
if user_debug:
|
|
traceback.print_exc()
|
|
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "detail", str(e)),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
|
)
|
|
error_msg = f"{str(e)}"
|
|
raise ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v1/completions", dependencies=[Depends(user_api_key_auth)], tags=["completions"]
|
|
)
|
|
@router.post(
|
|
"/completions", dependencies=[Depends(user_api_key_auth)], tags=["completions"]
|
|
)
|
|
@router.post(
|
|
"/engines/{model:path}/completions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["completions"],
|
|
)
|
|
@router.post(
|
|
"/openai/deployments/{model:path}/completions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["completions"],
|
|
)
|
|
async def completion(
|
|
request: Request,
|
|
fastapi_response: Response,
|
|
model: Optional[str] = None,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
|
|
try:
|
|
body = await request.body()
|
|
body_str = body.decode()
|
|
try:
|
|
data = ast.literal_eval(body_str)
|
|
except:
|
|
data = json.loads(body_str)
|
|
|
|
data["user"] = data.get("user", user_api_key_dict.user_id)
|
|
data["model"] = (
|
|
general_settings.get("completion_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or model # for azure deployments
|
|
or data["model"] # default passed in http request
|
|
)
|
|
if user_model:
|
|
data["model"] = user_model
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["user_api_key_team_alias"] = getattr(
|
|
user_api_key_dict, "team_alias", None
|
|
)
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
|
|
# override with user settings, these are params passed via cli
|
|
if user_temperature:
|
|
data["temperature"] = user_temperature
|
|
if user_request_timeout:
|
|
data["request_timeout"] = user_request_timeout
|
|
if user_max_tokens:
|
|
data["max_tokens"] = user_max_tokens
|
|
if user_api_base:
|
|
data["api_base"] = user_api_base
|
|
|
|
### MODEL ALIAS MAPPING ###
|
|
# check if model name in model alias map
|
|
# get the actual model name
|
|
if data["model"] in litellm.model_alias_map:
|
|
data["model"] = litellm.model_alias_map[data["model"]]
|
|
|
|
### CALL HOOKS ### - modify incoming data before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, call_type="completion"
|
|
)
|
|
|
|
### ROUTE THE REQUESTs ###
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
# skip router if user passed their key
|
|
if "api_key" in data:
|
|
response = await litellm.atext_completion(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
response = await llm_router.atext_completion(**data)
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
response = await llm_router.atext_completion(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.atext_completion(
|
|
**data, specific_deployment=True
|
|
)
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] not in router_model_names
|
|
and llm_router.default_deployment is not None
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.atext_completion(**data)
|
|
elif user_model is not None: # `litellm --model <your-model-name>`
|
|
response = await litellm.atext_completion(**data)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={
|
|
"error": "completion: Invalid model name passed in model="
|
|
+ data.get("model", "")
|
|
},
|
|
)
|
|
|
|
if hasattr(response, "_hidden_params"):
|
|
model_id = response._hidden_params.get("model_id", None) or ""
|
|
original_response = (
|
|
response._hidden_params.get("original_response", None) or ""
|
|
)
|
|
else:
|
|
model_id = ""
|
|
original_response = ""
|
|
|
|
verbose_proxy_logger.debug("final response: %s", response)
|
|
if (
|
|
"stream" in data and data["stream"] == True
|
|
): # use generate_responses to stream responses
|
|
custom_headers = {
|
|
"x-litellm-model-id": model_id,
|
|
}
|
|
selected_data_generator = select_data_generator(
|
|
response=response, user_api_key_dict=user_api_key_dict
|
|
)
|
|
|
|
return StreamingResponse(
|
|
selected_data_generator,
|
|
media_type="text/event-stream",
|
|
headers=custom_headers,
|
|
)
|
|
|
|
fastapi_response.headers["x-litellm-model-id"] = model_id
|
|
return response
|
|
except Exception as e:
|
|
data["litellm_status"] = "fail" # used for alerting
|
|
verbose_proxy_logger.debug("EXCEPTION RAISED IN PROXY MAIN.PY")
|
|
verbose_proxy_logger.debug(
|
|
"\033[1;31mAn error occurred: %s\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`",
|
|
e,
|
|
)
|
|
traceback.print_exc()
|
|
error_traceback = traceback.format_exc()
|
|
error_msg = f"{str(e)}"
|
|
raise ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v1/embeddings",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["embeddings"],
|
|
)
|
|
@router.post(
|
|
"/embeddings",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["embeddings"],
|
|
)
|
|
@router.post(
|
|
"/openai/deployments/{model:path}/embeddings",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["embeddings"],
|
|
) # azure compatible endpoint
|
|
async def embeddings(
|
|
request: Request,
|
|
model: Optional[str] = None,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global proxy_logging_obj
|
|
try:
|
|
# Use orjson to parse JSON data, orjson speeds up requests significantly
|
|
body = await request.body()
|
|
data = orjson.loads(body)
|
|
|
|
# Include original request and headers in the data
|
|
data["proxy_server_request"] = {
|
|
"url": str(request.url),
|
|
"method": request.method,
|
|
"headers": dict(request.headers),
|
|
"body": copy.copy(data), # use copy instead of deepcopy
|
|
}
|
|
|
|
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
|
|
data["user"] = user_api_key_dict.user_id
|
|
|
|
data["model"] = (
|
|
general_settings.get("embedding_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or model # for azure deployments
|
|
or data["model"] # default passed in http request
|
|
)
|
|
if user_model:
|
|
data["model"] = user_model
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["user_api_key_team_alias"] = getattr(
|
|
user_api_key_dict, "team_alias", None
|
|
)
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
|
|
### TEAM-SPECIFIC PARAMS ###
|
|
if user_api_key_dict.team_id is not None:
|
|
team_config = await proxy_config.load_team_config(
|
|
team_id=user_api_key_dict.team_id
|
|
)
|
|
if len(team_config) == 0:
|
|
pass
|
|
else:
|
|
team_id = team_config.pop("team_id", None)
|
|
data["metadata"]["team_id"] = team_id
|
|
data = {
|
|
**team_config,
|
|
**data,
|
|
} # add the team-specific configs to the completion call
|
|
|
|
### MODEL ALIAS MAPPING ###
|
|
# check if model name in model alias map
|
|
# get the actual model name
|
|
if data["model"] in litellm.model_alias_map:
|
|
data["model"] = litellm.model_alias_map[data["model"]]
|
|
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
if (
|
|
"input" in data
|
|
and isinstance(data["input"], list)
|
|
and isinstance(data["input"][0], list)
|
|
and isinstance(data["input"][0][0], int)
|
|
): # check if array of tokens passed in
|
|
# check if non-openai/azure model called - e.g. for langchain integration
|
|
if llm_model_list is not None and data["model"] in router_model_names:
|
|
for m in llm_model_list:
|
|
if m["model_name"] == data["model"] and (
|
|
m["litellm_params"]["model"] in litellm.open_ai_embedding_models
|
|
or m["litellm_params"]["model"].startswith("azure/")
|
|
):
|
|
pass
|
|
else:
|
|
# non-openai/azure embedding model called with token input
|
|
input_list = []
|
|
for i in data["input"]:
|
|
input_list.append(
|
|
litellm.decode(model="gpt-3.5-turbo", tokens=i)
|
|
)
|
|
data["input"] = input_list
|
|
break
|
|
|
|
### CALL HOOKS ### - modify incoming data / reject request before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, call_type="embeddings"
|
|
)
|
|
|
|
## ROUTE TO CORRECT ENDPOINT ##
|
|
# skip router if user passed their key
|
|
if "api_key" in data:
|
|
response = await litellm.aembedding(**data)
|
|
elif "user_config" in data:
|
|
# initialize a new router instance. make request using this Router
|
|
router_config = data.pop("user_config")
|
|
user_router = litellm.Router(**router_config)
|
|
response = await user_router.aembedding(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
response = await llm_router.aembedding(**data)
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
response = await llm_router.aembedding(
|
|
**data
|
|
) # ensure this goes the llm_router, router will do the correct alias mapping
|
|
elif (
|
|
llm_router is not None and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.aembedding(**data, specific_deployment=True)
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] not in router_model_names
|
|
and llm_router.default_deployment is not None
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.aembedding(**data)
|
|
elif user_model is not None: # `litellm --model <your-model-name>`
|
|
response = await litellm.aembedding(**data)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={
|
|
"error": "embeddings: Invalid model name passed in model="
|
|
+ data.get("model", "")
|
|
},
|
|
)
|
|
|
|
### ALERTING ###
|
|
data["litellm_status"] = "success" # used for alerting
|
|
|
|
return response
|
|
except Exception as e:
|
|
data["litellm_status"] = "fail" # used for alerting
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
traceback.print_exc()
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "message", str(e)),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
|
)
|
|
else:
|
|
error_msg = f"{str(e)}"
|
|
raise ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v1/images/generations",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["images"],
|
|
)
|
|
@router.post(
|
|
"/images/generations",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["images"],
|
|
)
|
|
async def image_generation(
|
|
request: Request,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global proxy_logging_obj
|
|
try:
|
|
# Use orjson to parse JSON data, orjson speeds up requests significantly
|
|
body = await request.body()
|
|
data = orjson.loads(body)
|
|
|
|
# Include original request and headers in the data
|
|
data["proxy_server_request"] = {
|
|
"url": str(request.url),
|
|
"method": request.method,
|
|
"headers": dict(request.headers),
|
|
"body": copy.copy(data), # use copy instead of deepcopy
|
|
}
|
|
|
|
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
|
|
data["user"] = user_api_key_dict.user_id
|
|
|
|
data["model"] = (
|
|
general_settings.get("image_generation_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or data["model"] # default passed in http request
|
|
)
|
|
if user_model:
|
|
data["model"] = user_model
|
|
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["user_api_key_team_alias"] = getattr(
|
|
user_api_key_dict, "team_alias", None
|
|
)
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
|
|
### TEAM-SPECIFIC PARAMS ###
|
|
if user_api_key_dict.team_id is not None:
|
|
team_config = await proxy_config.load_team_config(
|
|
team_id=user_api_key_dict.team_id
|
|
)
|
|
if len(team_config) == 0:
|
|
pass
|
|
else:
|
|
team_id = team_config.pop("team_id", None)
|
|
data["metadata"]["team_id"] = team_id
|
|
data = {
|
|
**team_config,
|
|
**data,
|
|
} # add the team-specific configs to the completion call
|
|
|
|
### MODEL ALIAS MAPPING ###
|
|
# check if model name in model alias map
|
|
# get the actual model name
|
|
if data["model"] in litellm.model_alias_map:
|
|
data["model"] = litellm.model_alias_map[data["model"]]
|
|
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
|
|
### CALL HOOKS ### - modify incoming data / reject request before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, call_type="image_generation"
|
|
)
|
|
|
|
## ROUTE TO CORRECT ENDPOINT ##
|
|
# skip router if user passed their key
|
|
if "api_key" in data:
|
|
response = await litellm.aimage_generation(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
response = await llm_router.aimage_generation(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.aimage_generation(
|
|
**data, specific_deployment=True
|
|
)
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
response = await llm_router.aimage_generation(
|
|
**data
|
|
) # ensure this goes the llm_router, router will do the correct alias mapping
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] not in router_model_names
|
|
and llm_router.default_deployment is not None
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.aimage_generation(**data)
|
|
elif user_model is not None: # `litellm --model <your-model-name>`
|
|
response = await litellm.aimage_generation(**data)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={
|
|
"error": "image_generation: Invalid model name passed in model="
|
|
+ data.get("model", "")
|
|
},
|
|
)
|
|
|
|
### ALERTING ###
|
|
data["litellm_status"] = "success" # used for alerting
|
|
|
|
return response
|
|
except Exception as e:
|
|
data["litellm_status"] = "fail" # used for alerting
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
traceback.print_exc()
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "message", str(e)),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
|
)
|
|
else:
|
|
error_msg = f"{str(e)}"
|
|
raise ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v1/audio/transcriptions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["audio"],
|
|
)
|
|
@router.post(
|
|
"/audio/transcriptions",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
tags=["audio"],
|
|
)
|
|
async def audio_transcriptions(
|
|
request: Request,
|
|
file: UploadFile = File(...),
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Same params as:
|
|
|
|
https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl
|
|
"""
|
|
global proxy_logging_obj
|
|
try:
|
|
# Use orjson to parse JSON data, orjson speeds up requests significantly
|
|
form_data = await request.form()
|
|
data: Dict = {key: value for key, value in form_data.items() if key != "file"}
|
|
|
|
# Include original request and headers in the data
|
|
data["proxy_server_request"] = { # type: ignore
|
|
"url": str(request.url),
|
|
"method": request.method,
|
|
"headers": dict(request.headers),
|
|
"body": copy.copy(data), # use copy instead of deepcopy
|
|
}
|
|
|
|
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
|
|
data["user"] = user_api_key_dict.user_id
|
|
|
|
data["model"] = (
|
|
general_settings.get("moderation_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or data["model"] # default passed in http request
|
|
)
|
|
if user_model:
|
|
data["model"] = user_model
|
|
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["user_api_key_team_alias"] = getattr(
|
|
user_api_key_dict, "team_alias", None
|
|
)
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
data["metadata"]["file_name"] = file.filename
|
|
|
|
### TEAM-SPECIFIC PARAMS ###
|
|
if user_api_key_dict.team_id is not None:
|
|
team_config = await proxy_config.load_team_config(
|
|
team_id=user_api_key_dict.team_id
|
|
)
|
|
if len(team_config) == 0:
|
|
pass
|
|
else:
|
|
team_id = team_config.pop("team_id", None)
|
|
data["metadata"]["team_id"] = team_id
|
|
data = {
|
|
**team_config,
|
|
**data,
|
|
} # add the team-specific configs to the completion call
|
|
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
|
|
assert (
|
|
file.filename is not None
|
|
) # make sure filename passed in (needed for type)
|
|
|
|
_original_filename = file.filename
|
|
file_extension = os.path.splitext(file.filename)[1]
|
|
# rename the file to a random hash file name -> we eventuall remove the file and don't want to remove any local files
|
|
file.filename = f"tmp-request" + str(uuid.uuid4()) + file_extension
|
|
|
|
# IMP - Asserts that we've renamed the uploaded file, since we run os.remove(file.filename), we should rename the original file
|
|
assert file.filename != _original_filename
|
|
|
|
with open(file.filename, "wb+") as f:
|
|
f.write(await file.read())
|
|
try:
|
|
data["file"] = open(file.filename, "rb")
|
|
### CALL HOOKS ### - modify incoming data / reject request before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
call_type="audio_transcription",
|
|
)
|
|
|
|
## ROUTE TO CORRECT ENDPOINT ##
|
|
# skip router if user passed their key
|
|
if "api_key" in data:
|
|
response = await litellm.atranscription(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
response = await llm_router.atranscription(**data)
|
|
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.atranscription(
|
|
**data, specific_deployment=True
|
|
)
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
response = await llm_router.atranscription(
|
|
**data
|
|
) # ensure this goes the llm_router, router will do the correct alias mapping
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] not in router_model_names
|
|
and llm_router.default_deployment is not None
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.atranscription(**data)
|
|
elif user_model is not None: # `litellm --model <your-model-name>`
|
|
response = await litellm.atranscription(**data)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={
|
|
"error": "audio_transcriptions: Invalid model name passed in model="
|
|
+ data.get("model", "")
|
|
},
|
|
)
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
finally:
|
|
os.remove(file.filename) # Delete the saved file
|
|
|
|
### ALERTING ###
|
|
data["litellm_status"] = "success" # used for alerting
|
|
return response
|
|
except Exception as e:
|
|
data["litellm_status"] = "fail" # used for alerting
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
traceback.print_exc()
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "message", str(e.detail)),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
|
)
|
|
else:
|
|
error_traceback = traceback.format_exc()
|
|
error_msg = f"{str(e)}"
|
|
raise ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v1/moderations",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["moderations"],
|
|
)
|
|
@router.post(
|
|
"/moderations",
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_class=ORJSONResponse,
|
|
tags=["moderations"],
|
|
)
|
|
async def moderations(
|
|
request: Request,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
The moderations endpoint is a tool you can use to check whether content complies with an LLM Providers policies.
|
|
|
|
Quick Start
|
|
```
|
|
curl --location 'http://0.0.0.0:4000/moderations' \
|
|
--header 'Content-Type: application/json' \
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
--data '{"input": "Sample text goes here", "model": "text-moderation-stable"}'
|
|
```
|
|
"""
|
|
global proxy_logging_obj
|
|
try:
|
|
# Use orjson to parse JSON data, orjson speeds up requests significantly
|
|
body = await request.body()
|
|
data = orjson.loads(body)
|
|
|
|
# Include original request and headers in the data
|
|
data["proxy_server_request"] = {
|
|
"url": str(request.url),
|
|
"method": request.method,
|
|
"headers": dict(request.headers),
|
|
"body": copy.copy(data), # use copy instead of deepcopy
|
|
}
|
|
|
|
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
|
|
data["user"] = user_api_key_dict.user_id
|
|
|
|
data["model"] = (
|
|
general_settings.get("moderation_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or data["model"] # default passed in http request
|
|
)
|
|
if user_model:
|
|
data["model"] = user_model
|
|
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["user_api_key_team_alias"] = getattr(
|
|
user_api_key_dict, "team_alias", None
|
|
)
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
|
|
### TEAM-SPECIFIC PARAMS ###
|
|
if user_api_key_dict.team_id is not None:
|
|
team_config = await proxy_config.load_team_config(
|
|
team_id=user_api_key_dict.team_id
|
|
)
|
|
if len(team_config) == 0:
|
|
pass
|
|
else:
|
|
team_id = team_config.pop("team_id", None)
|
|
data["metadata"]["team_id"] = team_id
|
|
data = {
|
|
**team_config,
|
|
**data,
|
|
} # add the team-specific configs to the completion call
|
|
|
|
router_model_names = llm_router.model_names if llm_router is not None else []
|
|
|
|
### CALL HOOKS ### - modify incoming data / reject request before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, call_type="moderation"
|
|
)
|
|
|
|
start_time = time.time()
|
|
|
|
## ROUTE TO CORRECT ENDPOINT ##
|
|
# skip router if user passed their key
|
|
if "api_key" in data:
|
|
response = await litellm.amoderation(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
response = await llm_router.amoderation(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.amoderation(**data, specific_deployment=True)
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
response = await llm_router.amoderation(
|
|
**data
|
|
) # ensure this goes the llm_router, router will do the correct alias mapping
|
|
elif (
|
|
llm_router is not None
|
|
and data["model"] not in router_model_names
|
|
and llm_router.default_deployment is not None
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.amoderation(**data)
|
|
elif user_model is not None: # `litellm --model <your-model-name>`
|
|
response = await litellm.amoderation(**data)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={
|
|
"error": "moderations: Invalid model name passed in model="
|
|
+ data.get("model", "")
|
|
},
|
|
)
|
|
|
|
### ALERTING ###
|
|
data["litellm_status"] = "success" # used for alerting
|
|
|
|
return response
|
|
except Exception as e:
|
|
data["litellm_status"] = "fail" # used for alerting
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
traceback.print_exc()
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "message", str(e)),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
|
)
|
|
else:
|
|
error_traceback = traceback.format_exc()
|
|
error_msg = f"{str(e)}"
|
|
raise ProxyException(
|
|
message=getattr(e, "message", error_msg),
|
|
type=getattr(e, "type", "None"),
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", 500),
|
|
)
|
|
|
|
|
|
#### KEY MANAGEMENT ####
|
|
|
|
|
|
@router.post(
|
|
"/key/generate",
|
|
tags=["key management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_model=GenerateKeyResponse,
|
|
)
|
|
async def generate_key_fn(
|
|
data: GenerateKeyRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
Authorization: Optional[str] = Header(None),
|
|
):
|
|
"""
|
|
Generate an API key based on the provided data.
|
|
|
|
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").
|
|
- key_alias: Optional[str] - User defined key alias
|
|
- team_id: Optional[str] - The team id of the key
|
|
- user_id: Optional[str] - The user id of the key
|
|
- 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" }
|
|
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
|
- model_max_budget: Optional[dict] - key-specific model budget in USD. Example - {"text-davinci-002": 0.5, "gpt-3.5-turbo": 0.5}. IF null or {} then no model specific budget.
|
|
|
|
Examples:
|
|
|
|
1. Allow users to turn on/off pii masking
|
|
|
|
```bash
|
|
curl --location 'http://0.0.0.0:8000/key/generate' \
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
--header 'Content-Type: application/json' \
|
|
--data '{
|
|
"permissions": {"allow_pii_controls": true}
|
|
}'
|
|
```
|
|
|
|
Returns:
|
|
- key: (str) The generated api key
|
|
- expires: (datetime) Datetime object for when key expires.
|
|
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
|
|
"""
|
|
try:
|
|
global user_custom_key_generate
|
|
verbose_proxy_logger.debug("entered /key/generate")
|
|
|
|
if user_custom_key_generate is not None:
|
|
result = await user_custom_key_generate(data)
|
|
decision = result.get("decision", True)
|
|
message = result.get("message", "Authentication Failed - Custom Auth Rule")
|
|
if not decision:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN, detail=message
|
|
)
|
|
# check if user set default key/generate params on config.yaml
|
|
if litellm.default_key_generate_params is not None:
|
|
for elem in data:
|
|
key, value = elem
|
|
if value is None and key in [
|
|
"max_budget",
|
|
"user_id",
|
|
"team_id",
|
|
"max_parallel_requests",
|
|
"tpm_limit",
|
|
"rpm_limit",
|
|
"budget_duration",
|
|
]:
|
|
setattr(
|
|
data, key, litellm.default_key_generate_params.get(key, None)
|
|
)
|
|
elif key == "models" and value == []:
|
|
setattr(data, key, litellm.default_key_generate_params.get(key, []))
|
|
elif key == "metadata" and value == {}:
|
|
setattr(data, key, litellm.default_key_generate_params.get(key, {}))
|
|
|
|
# check if user set default key/generate params on config.yaml
|
|
if litellm.upperbound_key_generate_params is not None:
|
|
for elem in data:
|
|
# if key in litellm.upperbound_key_generate_params, use the min of value and litellm.upperbound_key_generate_params[key]
|
|
key, value = elem
|
|
if (
|
|
value is not None
|
|
and getattr(litellm.upperbound_key_generate_params, key, None)
|
|
is not None
|
|
):
|
|
# if value is float/int
|
|
if key in [
|
|
"max_budget",
|
|
"max_parallel_requests",
|
|
"tpm_limit",
|
|
"rpm_limit",
|
|
]:
|
|
if value > getattr(litellm.upperbound_key_generate_params, key):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"{key} is over max limit set in config - user_value={value}; max_value={getattr(litellm.upperbound_key_generate_params, key)}"
|
|
},
|
|
)
|
|
elif key == "budget_duration":
|
|
# budgets are in 1s, 1m, 1h, 1d, 1m (30s, 30m, 30h, 30d, 30m)
|
|
# compare the duration in seconds and max duration in seconds
|
|
upperbound_budget_duration = _duration_in_seconds(
|
|
duration=getattr(
|
|
litellm.upperbound_key_generate_params, key
|
|
)
|
|
)
|
|
user_set_budget_duration = _duration_in_seconds(duration=value)
|
|
if user_set_budget_duration > upperbound_budget_duration:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Budget duration is over max limit set in config - user_value={user_set_budget_duration}; max_value={upperbound_budget_duration}"
|
|
},
|
|
)
|
|
|
|
# TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable
|
|
_budget_id = None
|
|
if prisma_client is not None and data.soft_budget is not None:
|
|
# create the Budget Row for the LiteLLM Verification Token
|
|
budget_row = LiteLLM_BudgetTable(
|
|
soft_budget=data.soft_budget,
|
|
model_max_budget=data.model_max_budget or {},
|
|
)
|
|
new_budget = prisma_client.jsonify_object(
|
|
budget_row.json(exclude_none=True)
|
|
)
|
|
|
|
_budget = await prisma_client.db.litellm_budgettable.create(
|
|
data={
|
|
**new_budget, # type: ignore
|
|
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
}
|
|
)
|
|
_budget_id = getattr(_budget, "budget_id", None)
|
|
data_json = data.json() # type: ignore
|
|
# if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
|
|
if "max_budget" in data_json:
|
|
data_json["key_max_budget"] = data_json.pop("max_budget", None)
|
|
if _budget_id is not None:
|
|
data_json["budget_id"] = _budget_id
|
|
|
|
if "budget_duration" in data_json:
|
|
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
|
|
|
|
response = await generate_key_helper_fn(**data_json, table_name="key")
|
|
|
|
response["soft_budget"] = (
|
|
data.soft_budget
|
|
) # include the user-input soft budget in the response
|
|
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)})"),
|
|
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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def update_key_fn(request: Request, data: UpdateKeyRequest):
|
|
"""
|
|
Update an existing key
|
|
"""
|
|
global prisma_client
|
|
try:
|
|
data_json: dict = data.json()
|
|
key = data_json.pop("key")
|
|
# get the row from db
|
|
if prisma_client is None:
|
|
raise Exception("Not connected to DB!")
|
|
|
|
# get non default values for key
|
|
non_default_values = {}
|
|
for k, v in data_json.items():
|
|
if v is not None and v not in (
|
|
[],
|
|
{},
|
|
0,
|
|
): # models default to [], spend defaults to 0, we should not reset these values
|
|
non_default_values[k] = v
|
|
|
|
if "duration" in non_default_values:
|
|
duration = non_default_values.pop("duration")
|
|
duration_s = _duration_in_seconds(duration=duration)
|
|
expires = datetime.utcnow() + timedelta(seconds=duration_s)
|
|
non_default_values["expires"] = expires
|
|
|
|
response = await prisma_client.update_data(
|
|
token=key, data={**non_default_values, "token": key}
|
|
)
|
|
|
|
# Delete - key from cache, since it's been updated!
|
|
# key updated - a new model could have been added to this key. it should not block requests after this is done
|
|
user_api_key_cache.delete_cache(key)
|
|
hashed_token = hash_token(key)
|
|
user_api_key_cache.delete_cache(hashed_token)
|
|
|
|
return {"key": key, **response["data"]}
|
|
# update based on remaining passed in values
|
|
except Exception as e:
|
|
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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def delete_key_fn(
|
|
data: KeyRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Delete a key from the key management system.
|
|
|
|
Parameters::
|
|
- 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", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]}
|
|
|
|
|
|
Raises:
|
|
HTTPException: If an error occurs during key deletion.
|
|
"""
|
|
try:
|
|
global user_api_key_cache
|
|
keys = data.keys
|
|
if len(keys) == 0:
|
|
raise ProxyException(
|
|
message=f"No keys provided, passed in: keys={keys}",
|
|
type="auth_error",
|
|
param="keys",
|
|
code=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
## only allow user to delete keys they own
|
|
user_id = user_api_key_dict.user_id
|
|
verbose_proxy_logger.debug(
|
|
f"user_api_key_dict.user_role: {user_api_key_dict.user_role}"
|
|
)
|
|
if (
|
|
user_api_key_dict.user_role is not None
|
|
and user_api_key_dict.user_role == "proxy_admin"
|
|
):
|
|
user_id = None # unless they're admin
|
|
|
|
number_deleted_keys = await delete_verification_token(
|
|
tokens=keys, user_id=user_id
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"/key/delete - deleted_keys={number_deleted_keys['deleted_keys']}"
|
|
)
|
|
|
|
try:
|
|
assert len(keys) == number_deleted_keys["deleted_keys"]
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Not all keys passed in were deleted. This probably means you don't have access to delete all the keys passed in. Keys passed in={len(keys)}, Deleted keys ={number_deleted_keys['deleted_keys']}"
|
|
},
|
|
)
|
|
|
|
for key in keys:
|
|
user_api_key_cache.delete_cache(key)
|
|
# remove hash token from cache
|
|
hashed_token = hash_token(key)
|
|
user_api_key_cache.delete_cache(hashed_token)
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}"
|
|
)
|
|
|
|
return {"deleted_keys": keys}
|
|
except Exception as e:
|
|
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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/v2/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def info_key_fn_v2(
|
|
data: Optional[KeyRequest] = None,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Retrieve information about a list of keys.
|
|
|
|
**New endpoint**. Currently admin only.
|
|
Parameters:
|
|
keys: Optional[list] = body parameter representing the key(s) 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" \
|
|
-H "Authorization: Bearer sk-1234" \
|
|
-d {"keys": ["sk-1", "sk-2", "sk-3"]}
|
|
```
|
|
"""
|
|
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 data is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
|
detail={"message": "Malformed request. No keys passed in."},
|
|
)
|
|
|
|
key_info = await prisma_client.get_data(
|
|
token=data.keys, table_name="key", query_type="find_all"
|
|
)
|
|
filtered_key_info = []
|
|
for k in key_info:
|
|
try:
|
|
k = k.model_dump() # noqa
|
|
except:
|
|
# if using pydantic v1
|
|
k = k.dict()
|
|
filtered_key_info.append(k)
|
|
return {"key": data.keys, "info": filtered_key_info}
|
|
|
|
except Exception as e:
|
|
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,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def info_key_fn(
|
|
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:
|
|
key_info = key_info.model_dump() # noqa
|
|
except:
|
|
# if using pydantic v1
|
|
key_info = key_info.dict()
|
|
key_info.pop("token")
|
|
return {"key": key, "info": key_info}
|
|
except Exception as e:
|
|
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,
|
|
)
|
|
|
|
|
|
#### SPEND MANAGEMENT #####
|
|
|
|
|
|
@router.get(
|
|
"/spend/keys",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def spend_key_fn():
|
|
"""
|
|
View all keys created, ordered by spend
|
|
|
|
Example Request:
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/keys" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
"""
|
|
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"
|
|
)
|
|
|
|
key_info = await prisma_client.get_data(table_name="key", query_type="find_all")
|
|
return key_info
|
|
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={"error": str(e)},
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/spend/users",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
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
|
|
|
|
Example Request:
|
|
```
|
|
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:
|
|
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 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
|
|
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail={"error": str(e)},
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/spend/tags",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
responses={
|
|
200: {"model": List[LiteLLM_SpendLogs]},
|
|
},
|
|
)
|
|
async def view_spend_tags(
|
|
start_date: Optional[str] = fastapi.Query(
|
|
default=None,
|
|
description="Time from which to start viewing key spend",
|
|
),
|
|
end_date: Optional[str] = fastapi.Query(
|
|
default=None,
|
|
description="Time till which to view key spend",
|
|
),
|
|
):
|
|
"""
|
|
LiteLLM Enterprise - View Spend Per Request Tag
|
|
|
|
Example Request:
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/tags" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
|
|
Spend with Start Date and End Date
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/tags?start_date=2022-01-01&end_date=2022-02-01" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
"""
|
|
|
|
from enterprise.utils import get_spend_by_tags
|
|
|
|
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"
|
|
)
|
|
|
|
# run the following SQL query on prisma
|
|
"""
|
|
SELECT
|
|
jsonb_array_elements_text(request_tags) AS individual_request_tag,
|
|
COUNT(*) AS log_count,
|
|
SUM(spend) AS total_spend
|
|
FROM "LiteLLM_SpendLogs"
|
|
GROUP BY individual_request_tag;
|
|
"""
|
|
response = await get_spend_by_tags(
|
|
start_date=start_date, end_date=end_date, prisma_client=prisma_client
|
|
)
|
|
|
|
return response
|
|
except Exception as e:
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "detail", f"/spend/tags Error({str(e)})"),
|
|
type="internal_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
raise e
|
|
raise ProxyException(
|
|
message="/spend/tags Error" + str(e),
|
|
type="internal_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/global/spend/tags",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
include_in_schema=False,
|
|
responses={
|
|
200: {"model": List[LiteLLM_SpendLogs]},
|
|
},
|
|
)
|
|
async def global_view_spend_tags(
|
|
start_date: Optional[str] = fastapi.Query(
|
|
default=None,
|
|
description="Time from which to start viewing key spend",
|
|
),
|
|
end_date: Optional[str] = fastapi.Query(
|
|
default=None,
|
|
description="Time till which to view key spend",
|
|
),
|
|
):
|
|
"""
|
|
LiteLLM Enterprise - View Spend Per Request Tag. Used by LiteLLM UI
|
|
|
|
Example Request:
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/tags" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
|
|
Spend with Start Date and End Date
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/tags?start_date=2022-01-01&end_date=2022-02-01" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
"""
|
|
|
|
from enterprise.utils import ui_get_spend_by_tags
|
|
|
|
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"
|
|
)
|
|
|
|
response = await ui_get_spend_by_tags(
|
|
start_date=start_date, end_date=end_date, prisma_client=prisma_client
|
|
)
|
|
|
|
return response
|
|
except Exception as e:
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "detail", f"/spend/tags Error({str(e)})"),
|
|
type="internal_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
raise e
|
|
raise ProxyException(
|
|
message="/spend/tags Error" + str(e),
|
|
type="internal_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/spend/calculate",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
responses={
|
|
200: {
|
|
"cost": {
|
|
"description": "The calculated cost",
|
|
"example": 0.0,
|
|
"type": "float",
|
|
}
|
|
}
|
|
},
|
|
)
|
|
async def calculate_spend(request: Request):
|
|
"""
|
|
Accepts all the params of completion_cost.
|
|
|
|
Calculate spend **before** making call:
|
|
|
|
```
|
|
curl --location 'http://localhost:4000/spend/calculate'
|
|
--header 'Authorization: Bearer sk-1234'
|
|
--header 'Content-Type: application/json'
|
|
--data '{
|
|
"model": "anthropic.claude-v2",
|
|
"messages": [{"role": "user", "content": "Hey, how'''s it going?"}]
|
|
}'
|
|
```
|
|
|
|
Calculate spend **after** making call:
|
|
|
|
```
|
|
curl --location 'http://localhost:4000/spend/calculate'
|
|
--header 'Authorization: Bearer sk-1234'
|
|
--header 'Content-Type: application/json'
|
|
--data '{
|
|
"completion_response": {
|
|
"id": "chatcmpl-123",
|
|
"object": "chat.completion",
|
|
"created": 1677652288,
|
|
"model": "gpt-3.5-turbo-0125",
|
|
"system_fingerprint": "fp_44709d6fcb",
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Hello there, how may I assist you today?"
|
|
},
|
|
"logprobs": null,
|
|
"finish_reason": "stop"
|
|
}]
|
|
"usage": {
|
|
"prompt_tokens": 9,
|
|
"completion_tokens": 12,
|
|
"total_tokens": 21
|
|
}
|
|
}
|
|
}'
|
|
```
|
|
"""
|
|
from litellm import completion_cost
|
|
|
|
data = await request.json()
|
|
if "completion_response" in data:
|
|
data["completion_response"] = litellm.ModelResponse(
|
|
**data["completion_response"]
|
|
)
|
|
return {"cost": completion_cost(**data)}
|
|
|
|
|
|
@router.get(
|
|
"/spend/logs",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
responses={
|
|
200: {"model": List[LiteLLM_SpendLogs]},
|
|
},
|
|
)
|
|
async def view_spend_logs(
|
|
api_key: Optional[str] = fastapi.Query(
|
|
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",
|
|
),
|
|
start_date: Optional[str] = fastapi.Query(
|
|
default=None,
|
|
description="Time from which to start viewing key spend",
|
|
),
|
|
end_date: Optional[str] = fastapi.Query(
|
|
default=None,
|
|
description="Time till which to view key spend",
|
|
),
|
|
):
|
|
"""
|
|
View all spend logs, if request_id is provided, only logs for that request_id will be returned
|
|
|
|
Example Request for all logs
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/logs" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
|
|
Example Request for specific request_id
|
|
```
|
|
curl -X GET "http://0.0.0.0:8000/spend/logs?request_id=chatcmpl-6dcb2540-d3d7-4e49-bb27-291f863f112e" \
|
|
-H "Authorization: Bearer sk-1234"
|
|
```
|
|
|
|
Example Request for specific api_key
|
|
```
|
|
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:
|
|
verbose_proxy_logger.debug("inside view_spend_logs")
|
|
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"
|
|
)
|
|
spend_logs = []
|
|
if (
|
|
start_date is not None
|
|
and isinstance(start_date, str)
|
|
and end_date is not None
|
|
and isinstance(end_date, str)
|
|
):
|
|
# Convert the date strings to datetime objects
|
|
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
|
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
|
|
|
filter_query = {
|
|
"startTime": {
|
|
"gte": start_date_obj, # Greater than or equal to Start Date
|
|
"lte": end_date_obj, # Less than or equal to End Date
|
|
}
|
|
}
|
|
|
|
if api_key is not None and isinstance(api_key, str):
|
|
filter_query["api_key"] = api_key # type: ignore
|
|
elif request_id is not None and isinstance(request_id, str):
|
|
filter_query["request_id"] = request_id # type: ignore
|
|
elif user_id is not None and isinstance(user_id, str):
|
|
filter_query["user"] = user_id # type: ignore
|
|
|
|
# SQL query
|
|
response = await prisma_client.db.litellm_spendlogs.group_by(
|
|
by=["api_key", "user", "model", "startTime"],
|
|
where=filter_query, # type: ignore
|
|
sum={
|
|
"spend": True,
|
|
},
|
|
)
|
|
|
|
if (
|
|
isinstance(response, list)
|
|
and len(response) > 0
|
|
and isinstance(response[0], dict)
|
|
):
|
|
result: dict = {}
|
|
for record in response:
|
|
dt_object = datetime.strptime(
|
|
str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ"
|
|
) # type: ignore
|
|
date = dt_object.date()
|
|
if date not in result:
|
|
result[date] = {"users": {}, "models": {}}
|
|
api_key = record["api_key"]
|
|
user_id = record["user"]
|
|
model = record["model"]
|
|
result[date]["spend"] = (
|
|
result[date].get("spend", 0) + record["_sum"]["spend"]
|
|
)
|
|
result[date][api_key] = (
|
|
result[date].get(api_key, 0) + record["_sum"]["spend"]
|
|
)
|
|
result[date]["users"][user_id] = (
|
|
result[date]["users"].get(user_id, 0) + record["_sum"]["spend"]
|
|
)
|
|
result[date]["models"][model] = (
|
|
result[date]["models"].get(model, 0) + record["_sum"]["spend"]
|
|
)
|
|
return_list = []
|
|
final_date = None
|
|
for k, v in sorted(result.items()):
|
|
return_list.append({**v, "startTime": k})
|
|
final_date = k
|
|
|
|
end_date_date = end_date_obj.date()
|
|
if final_date is not None and final_date < end_date_date:
|
|
current_date = final_date + timedelta(days=1)
|
|
while current_date <= end_date_date:
|
|
# Represent current_date as string because original response has it this way
|
|
return_list.append(
|
|
{
|
|
"startTime": current_date,
|
|
"spend": 0,
|
|
"users": {},
|
|
"models": {},
|
|
}
|
|
) # If no data, will stay as zero
|
|
current_date += timedelta(days=1) # Move on to the next day
|
|
|
|
return return_list
|
|
|
|
return response
|
|
|
|
elif api_key is not None and isinstance(api_key, str):
|
|
if api_key.startswith("sk-"):
|
|
hashed_token = prisma_client.hash_token(token=api_key)
|
|
else:
|
|
hashed_token = api_key
|
|
spend_log = await prisma_client.get_data(
|
|
table_name="spend",
|
|
query_type="find_all",
|
|
key_val={"key": "api_key", "value": hashed_token},
|
|
)
|
|
if isinstance(spend_log, list):
|
|
return spend_log
|
|
else:
|
|
return [spend_log]
|
|
elif request_id is not None:
|
|
spend_log = await prisma_client.get_data(
|
|
table_name="spend",
|
|
query_type="find_unique",
|
|
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"
|
|
)
|
|
|
|
return spend_log
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
if isinstance(e, HTTPException):
|
|
raise ProxyException(
|
|
message=getattr(e, "detail", f"/spend/logs Error({str(e)})"),
|
|
type="internal_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
raise e
|
|
raise ProxyException(
|
|
message="/spend/logs Error" + str(e),
|
|
type="internal_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/global/spend/logs",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_spend_logs(
|
|
api_key: str = fastapi.Query(
|
|
default=None,
|
|
description="API Key to get global spend (spend per day for last 30d). Admin-only endpoint",
|
|
)
|
|
):
|
|
"""
|
|
[BETA] This is a beta endpoint. It will change.
|
|
|
|
Use this to get global spend (spend per day for last 30d). Admin-only endpoint
|
|
|
|
More efficient implementation of /spend/logs, by creating a view over the spend logs table.
|
|
"""
|
|
global prisma_client
|
|
if prisma_client is None:
|
|
raise ProxyException(
|
|
message="Prisma Client is not initialized",
|
|
type="internal_error",
|
|
param="None",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if api_key is None:
|
|
sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";"""
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
return response
|
|
else:
|
|
sql_query = (
|
|
"""
|
|
SELECT * FROM "MonthlyGlobalSpendPerKey"
|
|
WHERE "api_key" = '"""
|
|
+ api_key
|
|
+ """'
|
|
ORDER BY "date";
|
|
"""
|
|
)
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
return response
|
|
return
|
|
|
|
|
|
@router.get(
|
|
"/global/spend",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_spend():
|
|
"""
|
|
[BETA] This is a beta endpoint. It will change.
|
|
|
|
View total spend across all proxy keys
|
|
"""
|
|
global prisma_client
|
|
total_spend = 0.0
|
|
total_proxy_budget = 0.0
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
if response is not None:
|
|
if isinstance(response, list) and len(response) > 0:
|
|
total_spend = response[0].get("total_spend", 0.0)
|
|
|
|
return {"spend": total_spend, "max_budget": litellm.max_budget}
|
|
|
|
|
|
@router.get(
|
|
"/global/spend/keys",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_spend_keys(
|
|
limit: int = fastapi.Query(
|
|
default=None,
|
|
description="Number of keys to get. Will return Top 'n' keys.",
|
|
)
|
|
):
|
|
"""
|
|
[BETA] This is a beta endpoint. It will change.
|
|
|
|
Use this to get the top 'n' keys with the highest spend, ordered by spend.
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
sql_query = f"""SELECT * FROM "Last30dKeysBySpend" LIMIT {limit};"""
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
return response
|
|
|
|
|
|
@router.get(
|
|
"/global/spend/teams",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_spend_per_tea():
|
|
"""
|
|
[BETA] This is a beta endpoint. It will change.
|
|
|
|
Use this to get daily spend, grouped by `team_id` and `date`
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
sql_query = """
|
|
SELECT
|
|
t.team_alias as team_alias,
|
|
DATE(s."startTime") AS spend_date,
|
|
SUM(s.spend) AS total_spend
|
|
FROM
|
|
"LiteLLM_SpendLogs" s
|
|
LEFT JOIN
|
|
"LiteLLM_TeamTable" t ON s.team_id = t.team_id
|
|
WHERE
|
|
s."startTime" >= CURRENT_DATE - INTERVAL '30 days'
|
|
GROUP BY
|
|
t.team_alias,
|
|
DATE(s."startTime")
|
|
ORDER BY
|
|
spend_date;
|
|
"""
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
# transform the response for the Admin UI
|
|
spend_by_date = {}
|
|
team_aliases = set()
|
|
total_spend_per_team = {}
|
|
for row in response:
|
|
row_date = row["spend_date"]
|
|
if row_date is None:
|
|
continue
|
|
team_alias = row["team_alias"]
|
|
if team_alias is None:
|
|
team_alias = "Unassigned"
|
|
team_aliases.add(team_alias)
|
|
if row_date in spend_by_date:
|
|
# get the team_id for this entry
|
|
# get the spend for this entry
|
|
spend = row["total_spend"]
|
|
spend = round(spend, 2)
|
|
current_date_entries = spend_by_date[row_date]
|
|
current_date_entries[team_alias] = spend
|
|
else:
|
|
spend = row["total_spend"]
|
|
spend = round(spend, 2)
|
|
spend_by_date[row_date] = {team_alias: spend}
|
|
|
|
if team_alias in total_spend_per_team:
|
|
total_spend_per_team[team_alias] += spend
|
|
else:
|
|
total_spend_per_team[team_alias] = spend
|
|
|
|
total_spend_per_team_ui = []
|
|
# order the elements in total_spend_per_team by spend
|
|
total_spend_per_team = dict(
|
|
sorted(total_spend_per_team.items(), key=lambda item: item[1], reverse=True)
|
|
)
|
|
for team_id in total_spend_per_team:
|
|
# only add first 10 elements to total_spend_per_team_ui
|
|
if len(total_spend_per_team_ui) >= 10:
|
|
break
|
|
if team_id is None:
|
|
team_id = "Unassigned"
|
|
total_spend_per_team_ui.append(
|
|
{"team_id": team_id, "total_spend": total_spend_per_team[team_id]}
|
|
)
|
|
|
|
# sort spend_by_date by it's key (which is a date)
|
|
|
|
response_data = []
|
|
for key in spend_by_date:
|
|
value = spend_by_date[key]
|
|
response_data.append({"date": key, **value})
|
|
|
|
return {
|
|
"daily_spend": response_data,
|
|
"teams": list(team_aliases),
|
|
"total_spend_per_team": total_spend_per_team_ui,
|
|
}
|
|
|
|
|
|
@router.post(
|
|
"/global/spend/end_users",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_spend_end_users(data: Optional[GlobalEndUsersSpend] = None):
|
|
"""
|
|
[BETA] This is a beta endpoint. It will change.
|
|
|
|
Use this to get the top 'n' keys with the highest spend, ordered by spend.
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if data is None:
|
|
sql_query = f"""SELECT * FROM "Last30dTopEndUsersSpend";"""
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
else:
|
|
"""
|
|
Gets the top 100 end-users for a given api key
|
|
"""
|
|
current_date = datetime.now()
|
|
past_date = current_date - timedelta(days=30)
|
|
response = await prisma_client.db.litellm_spendlogs.group_by( # type: ignore
|
|
by=["end_user"],
|
|
where={
|
|
"AND": [{"startTime": {"gte": past_date}}, {"api_key": data.api_key}] # type: ignore
|
|
},
|
|
sum={"spend": True},
|
|
order={"_sum": {"spend": "desc"}}, # type: ignore
|
|
take=100,
|
|
count=True,
|
|
)
|
|
if response is not None and isinstance(response, list):
|
|
new_response = []
|
|
for r in response:
|
|
new_r = r
|
|
new_r["total_spend"] = r["_sum"]["spend"]
|
|
new_r["total_count"] = r["_count"]["_all"]
|
|
new_r.pop("_sum")
|
|
new_r.pop("_count")
|
|
new_response.append(new_r)
|
|
|
|
return response
|
|
|
|
|
|
@router.get(
|
|
"/global/spend/models",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_spend_models(
|
|
limit: int = fastapi.Query(
|
|
default=None,
|
|
description="Number of models to get. Will return Top 'n' models.",
|
|
)
|
|
):
|
|
"""
|
|
[BETA] This is a beta endpoint. It will change.
|
|
|
|
Use this to get the top 'n' keys with the highest spend, ordered by spend.
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
sql_query = f"""SELECT * FROM "Last30dModelsBySpend" LIMIT {limit};"""
|
|
|
|
response = await prisma_client.db.query_raw(query=sql_query)
|
|
|
|
return response
|
|
|
|
|
|
@router.post(
|
|
"/global/predict/spend/logs",
|
|
tags=["Budget & Spend Tracking"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def global_predict_spend_logs(request: Request):
|
|
from enterprise.utils import _forecast_daily_cost
|
|
|
|
data = await request.json()
|
|
data = data.get("data")
|
|
return _forecast_daily_cost(data)
|
|
|
|
|
|
#### USER MANAGEMENT ####
|
|
@router.post(
|
|
"/user/new",
|
|
tags=["user management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_model=NewUserResponse,
|
|
)
|
|
async def new_user(data: NewUserRequest):
|
|
"""
|
|
Use this to create a new user with a budget. This creates a new user and generates a new api key for the new user. The new api key is returned.
|
|
|
|
Returns user id, budget + new key.
|
|
|
|
Parameters:
|
|
- user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated.
|
|
- user_alias: Optional[str] - A descriptive name for you to know who this user id refers to.
|
|
- teams: Optional[list] - specify a list of team id's a user belongs to.
|
|
- organization_id: Optional[str] - specify the org a user belongs to.
|
|
- user_email: Optional[str] - Specify a user email.
|
|
- user_role: Optional[str] - Specify a user role - "admin", "app_owner", "app_user"
|
|
- max_budget: Optional[float] - Specify max budget for a given user.
|
|
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
|
|
- tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute)
|
|
- rpm_limit: Optional[int] - Specify rpm limit for a given user (Requests per minute)
|
|
- auto_create_key: bool - Default=True. Flag used for returning a key as part of the /user/new response
|
|
|
|
Returns:
|
|
- key: (str) The generated api key for the user
|
|
- expires: (datetime) Datetime object for when key expires.
|
|
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
|
|
- max_budget: (float|None) Max budget for given user.
|
|
"""
|
|
data_json = data.json() # type: ignore
|
|
if "user_role" in data_json:
|
|
user_role = data_json["user_role"]
|
|
if user_role is not None:
|
|
if user_role not in ["proxy_admin", "app_owner", "app_user"]:
|
|
raise ProxyException(
|
|
message=f"Invalid user role, passed in {user_role}. Must be one of 'admin', 'app_owner', 'app_user'",
|
|
type="invalid_user_role",
|
|
param="user_role",
|
|
code=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
if "user_id" in data_json and data_json["user_id"] is None:
|
|
data_json["user_id"] = str(uuid.uuid4())
|
|
auto_create_key = data_json.pop("auto_create_key", True)
|
|
if auto_create_key == False:
|
|
data_json["table_name"] = (
|
|
"user" # only create a user, don't create key if 'auto_create_key' set to False
|
|
)
|
|
response = await generate_key_helper_fn(**data_json)
|
|
|
|
# Admin UI Logic
|
|
# if team_id passed add this user to the team
|
|
if data_json.get("team_id", None) is not None:
|
|
await team_member_add(
|
|
data=TeamMemberAddRequest(
|
|
team_id=data_json.get("team_id", None),
|
|
member=Member(
|
|
user_id=data_json.get("user_id", None),
|
|
role="user",
|
|
user_email=data_json.get("user_email", None),
|
|
),
|
|
)
|
|
)
|
|
return NewUserResponse(
|
|
key=response.get("token", ""),
|
|
expires=response.get("expires", None),
|
|
max_budget=response["max_budget"],
|
|
user_id=response["user_id"],
|
|
team_id=response.get("team_id", None),
|
|
metadata=response.get("metadata", None),
|
|
models=response.get("models", None),
|
|
tpm_limit=response.get("tpm_limit", None),
|
|
rpm_limit=response.get("rpm_limit", None),
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/user/auth", tags=["user management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def user_auth(request: Request):
|
|
"""
|
|
Allows UI ("https://dashboard.litellm.ai/", or self-hosted - os.getenv("LITELLM_HOSTED_UI")) to request a magic link to be sent to user email, for auth to proxy.
|
|
|
|
Only allows emails from accepted email subdomains.
|
|
|
|
Rate limit: 1 request every 60s.
|
|
|
|
Only works, if you enable 'allow_user_auth' in general settings:
|
|
e.g.:
|
|
```yaml
|
|
general_settings:
|
|
allow_user_auth: true
|
|
```
|
|
|
|
Requirements:
|
|
SMTP server details saved in .env:
|
|
- os.environ["SMTP_HOST"]
|
|
- os.environ["SMTP_PORT"]
|
|
- os.environ["SMTP_USERNAME"]
|
|
- os.environ["SMTP_PASSWORD"]
|
|
- os.environ["SMTP_SENDER_EMAIL"]
|
|
"""
|
|
global prisma_client
|
|
|
|
data = await request.json() # type: ignore
|
|
user_email = data["user_email"]
|
|
page_params = data["page"]
|
|
if user_email is None:
|
|
raise HTTPException(status_code=400, detail="User email is none")
|
|
|
|
if prisma_client is None: # if no db connected, raise an error
|
|
raise Exception("No connected db.")
|
|
|
|
### Check if user email in user table
|
|
response = await prisma_client.get_generic_data(
|
|
key="user_email", value=user_email, table_name="users"
|
|
)
|
|
### if so - generate a 24 hr key with that user id
|
|
if response is not None:
|
|
user_id = response.user_id
|
|
response = await generate_key_helper_fn(
|
|
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id} # type: ignore
|
|
)
|
|
else: ### else - create new user
|
|
response = await generate_key_helper_fn(
|
|
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_email": user_email} # type: ignore
|
|
)
|
|
|
|
base_url = os.getenv("LITELLM_HOSTED_UI", "https://dashboard.litellm.ai/")
|
|
|
|
params = {
|
|
"sender_name": "LiteLLM Proxy",
|
|
"sender_email": os.getenv("SMTP_SENDER_EMAIL"),
|
|
"receiver_email": user_email,
|
|
"subject": "Your Magic Link",
|
|
"html": f"<strong> Follow this link, to login:\n\n{base_url}user/?token={response['token']}&user_id={response['user_id']}&page={page_params}</strong>",
|
|
}
|
|
|
|
await send_email(**params)
|
|
return "Email sent!"
|
|
|
|
|
|
@router.get(
|
|
"/user/info", tags=["user management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def user_info(
|
|
user_id: Optional[str] = fastapi.Query(
|
|
default=None, description="User ID in the request parameters"
|
|
),
|
|
view_all: bool = fastapi.Query(
|
|
default=False,
|
|
description="set to true to View all users. When using view_all, don't pass user_id",
|
|
),
|
|
page: Optional[int] = fastapi.Query(
|
|
default=0,
|
|
description="Page number for pagination. Only use when view_all is true",
|
|
),
|
|
page_size: Optional[int] = fastapi.Query(
|
|
default=25,
|
|
description="Number of items per page. Only use when view_all is true",
|
|
),
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Use this to get user information. (user row + all user key info)
|
|
|
|
Example request
|
|
```
|
|
curl -X GET 'http://localhost:8000/user/info?user_id=krrish7%40berri.ai' \
|
|
--header 'Authorization: Bearer sk-1234'
|
|
```
|
|
"""
|
|
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"
|
|
)
|
|
## GET USER ROW ##
|
|
if user_id is not None:
|
|
user_info = await prisma_client.get_data(user_id=user_id)
|
|
elif view_all == True:
|
|
if page is None:
|
|
page = 0
|
|
if page_size is None:
|
|
page_size = 25
|
|
offset = (page) * page_size # default is 0
|
|
limit = page_size # default is 10
|
|
user_info = await prisma_client.get_data(
|
|
table_name="user", query_type="find_all", offset=offset, limit=limit
|
|
)
|
|
return user_info
|
|
else:
|
|
user_info = None
|
|
## GET ALL TEAMS ##
|
|
team_list = []
|
|
team_id_list = []
|
|
# _DEPRECATED_ check if user in 'member' field
|
|
teams_1 = await prisma_client.get_data(
|
|
user_id=user_id, table_name="team", query_type="find_all"
|
|
)
|
|
|
|
if teams_1 is not None and isinstance(teams_1, list):
|
|
team_list = teams_1
|
|
for team in teams_1:
|
|
team_id_list.append(team.team_id)
|
|
|
|
if user_info is not None:
|
|
# *NEW* get all teams in user 'teams' field
|
|
teams_2 = await prisma_client.get_data(
|
|
team_id_list=user_info.teams, table_name="team", query_type="find_all"
|
|
)
|
|
|
|
if teams_2 is not None and isinstance(teams_2, list):
|
|
for team in teams_2:
|
|
if team.team_id not in team_id_list:
|
|
team_list.append(team)
|
|
team_id_list.append(team.team_id)
|
|
elif (
|
|
user_api_key_dict.user_id is not None and user_id is None
|
|
): # the key querying the endpoint is the one asking for it's teams
|
|
caller_user_info = await prisma_client.get_data(
|
|
user_id=user_api_key_dict.user_id
|
|
)
|
|
# *NEW* get all teams in user 'teams' field
|
|
if getattr(caller_user_info, "user_role", None) == "proxy_admin":
|
|
teams_2 = await prisma_client.get_data(
|
|
table_name="team",
|
|
query_type="find_all",
|
|
team_id_list=None,
|
|
)
|
|
else:
|
|
teams_2 = await prisma_client.get_data(
|
|
team_id_list=caller_user_info.teams,
|
|
table_name="team",
|
|
query_type="find_all",
|
|
)
|
|
|
|
if teams_2 is not None and isinstance(teams_2, list):
|
|
for team in teams_2:
|
|
if team.team_id not in team_id_list:
|
|
team_list.append(team)
|
|
team_id_list.append(team.team_id)
|
|
|
|
## GET ALL KEYS ##
|
|
keys = await prisma_client.get_data(
|
|
user_id=user_id,
|
|
table_name="key",
|
|
query_type="find_all",
|
|
expires=datetime.now(),
|
|
)
|
|
|
|
if user_info is None:
|
|
## make sure we still return a total spend ##
|
|
spend = 0
|
|
for k in keys:
|
|
spend += getattr(k, "spend", 0)
|
|
user_info = {"spend": spend}
|
|
|
|
## REMOVE HASHED TOKEN INFO before returning ##
|
|
returned_keys = []
|
|
for key in keys:
|
|
if (
|
|
key.token == litellm_master_key_hash
|
|
and general_settings.get("disable_master_key_return", False)
|
|
== True ## [IMPORTANT] used by hosted proxy-ui to prevent sharing master key on ui
|
|
):
|
|
continue
|
|
|
|
try:
|
|
key = key.model_dump() # noqa
|
|
except:
|
|
# if using pydantic v1
|
|
key = key.dict()
|
|
if (
|
|
"team_id" in key
|
|
and key["team_id"] is not None
|
|
and key["team_id"] != "litellm-dashboard"
|
|
):
|
|
team_info = await prisma_client.get_data(
|
|
team_id=key["team_id"], table_name="team"
|
|
)
|
|
team_alias = getattr(team_info, "team_alias", None)
|
|
key["team_alias"] = team_alias
|
|
else:
|
|
key["team_alias"] = "None"
|
|
returned_keys.append(key)
|
|
|
|
response_data = {
|
|
"user_id": user_id,
|
|
"user_info": user_info,
|
|
"keys": returned_keys,
|
|
"teams": team_list,
|
|
}
|
|
return response_data
|
|
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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/user/update", tags=["user management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def user_update(data: UpdateUserRequest):
|
|
"""
|
|
[TODO]: Use this to update user budget
|
|
"""
|
|
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!")
|
|
|
|
# get non default values for key
|
|
non_default_values = {}
|
|
for k, v in data_json.items():
|
|
if v is not None and v not in (
|
|
[],
|
|
{},
|
|
0,
|
|
): # models default to [], spend defaults to 0, we should not reset these values
|
|
non_default_values[k] = v
|
|
|
|
## ADD USER, IF NEW ##
|
|
verbose_proxy_logger.debug("/user/update: Received data = %s", data)
|
|
if data.user_id is not None and len(data.user_id) > 0:
|
|
non_default_values["user_id"] = data.user_id # type: ignore
|
|
verbose_proxy_logger.debug("In update user, user_id condition block.")
|
|
response = await prisma_client.update_data(
|
|
user_id=data.user_id,
|
|
data=non_default_values,
|
|
table_name="user",
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"received response from updating prisma client. response={response}"
|
|
)
|
|
elif data.user_email is not None:
|
|
non_default_values["user_id"] = str(uuid.uuid4())
|
|
non_default_values["user_email"] = data.user_email
|
|
## user email is not unique acc. to prisma schema -> future improvement
|
|
### for now: check if it exists in db, if not - insert it
|
|
existing_user_rows = await prisma_client.get_data(
|
|
key_val={"user_email": data.user_email},
|
|
table_name="user",
|
|
query_type="find_all",
|
|
)
|
|
if existing_user_rows is None or (
|
|
isinstance(existing_user_rows, list) and len(existing_user_rows) == 0
|
|
):
|
|
response = await prisma_client.insert_data(
|
|
data=non_default_values, table_name="user"
|
|
)
|
|
elif isinstance(existing_user_rows, list) and len(existing_user_rows) > 0:
|
|
for existing_user in existing_user_rows:
|
|
response = await prisma_client.update_data(
|
|
user_id=existing_user.user_id,
|
|
data=non_default_values,
|
|
table_name="user",
|
|
)
|
|
return response
|
|
# 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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/user/request_model",
|
|
tags=["user management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def user_request_model(request: Request):
|
|
"""
|
|
Allow a user to create a request to access a model
|
|
"""
|
|
global prisma_client
|
|
try:
|
|
data_json = await request.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}
|
|
new_models = non_default_values.get("models", None)
|
|
user_id = non_default_values.get("user_id", None)
|
|
justification = non_default_values.get("justification", None)
|
|
|
|
response = await prisma_client.insert_data(
|
|
data={
|
|
"models": new_models,
|
|
"justification": justification,
|
|
"user_id": user_id,
|
|
"status": "pending",
|
|
"request_id": str(uuid.uuid4()),
|
|
},
|
|
table_name="user_notification",
|
|
)
|
|
return {"status": "success"}
|
|
# 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,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/user/get_requests",
|
|
tags=["user management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def user_get_requests():
|
|
"""
|
|
Get all "Access" requests made by proxy users, access requests are requests for accessing models
|
|
"""
|
|
global prisma_client
|
|
try:
|
|
|
|
# get the row from db
|
|
if prisma_client is None:
|
|
raise Exception("Not connected to DB!")
|
|
|
|
# TODO: Optimize this so we don't read all the data here, eventually move to pagination
|
|
response = await prisma_client.get_data(
|
|
query_type="find_all",
|
|
table_name="user_notification",
|
|
)
|
|
return {"requests": response}
|
|
# 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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/end_user/block",
|
|
tags=["End User Management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def block_user(data: BlockUsers):
|
|
"""
|
|
[BETA] Reject calls with this end-user id
|
|
|
|
(any /chat/completion call with this user={end-user-id} param, will be rejected.)
|
|
|
|
```
|
|
curl -X POST "http://0.0.0.0:8000/user/block"
|
|
-H "Authorization: Bearer sk-1234"
|
|
-D '{
|
|
"user_ids": [<user_id>, ...]
|
|
}'
|
|
```
|
|
"""
|
|
try:
|
|
records = []
|
|
if prisma_client is not None:
|
|
for id in data.user_ids:
|
|
record = await prisma_client.db.litellm_endusertable.upsert(
|
|
where={"user_id": id}, # type: ignore
|
|
data={
|
|
"create": {"user_id": id, "blocked": True}, # type: ignore
|
|
"update": {"blocked": True},
|
|
},
|
|
)
|
|
records.append(record)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={"error": "Postgres DB Not connected"},
|
|
)
|
|
|
|
return {"blocked_users": records}
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"An error occurred - {str(e)}")
|
|
raise HTTPException(status_code=500, detail={"error": str(e)})
|
|
|
|
|
|
@router.post(
|
|
"/end_user/unblock",
|
|
tags=["End User Management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def unblock_user(data: BlockUsers):
|
|
"""
|
|
[BETA] Unblock calls with this user id
|
|
|
|
Example
|
|
```
|
|
curl -X POST "http://0.0.0.0:8000/user/unblock"
|
|
-H "Authorization: Bearer sk-1234"
|
|
-D '{
|
|
"user_ids": [<user_id>, ...]
|
|
}'
|
|
```
|
|
"""
|
|
from enterprise.enterprise_hooks.blocked_user_list import (
|
|
_ENTERPRISE_BlockedUserList,
|
|
)
|
|
|
|
if (
|
|
not any(isinstance(x, _ENTERPRISE_BlockedUserList) for x in litellm.callbacks)
|
|
or litellm.blocked_user_list is None
|
|
):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": "Blocked user check was never set. This call has no effect."
|
|
},
|
|
)
|
|
|
|
if isinstance(litellm.blocked_user_list, list):
|
|
for id in data.user_ids:
|
|
litellm.blocked_user_list.remove(id)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "`blocked_user_list` must be set as a list. Filepaths can't be updated."
|
|
},
|
|
)
|
|
|
|
return {"blocked_users": litellm.blocked_user_list}
|
|
|
|
|
|
@router.get(
|
|
"/user/get_users",
|
|
tags=["user management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def get_users(
|
|
role: str = fastapi.Query(
|
|
default=None,
|
|
description="Either 'proxy_admin', 'proxy_viewer', 'app_owner', 'app_user'",
|
|
)
|
|
):
|
|
"""
|
|
[BETA] This could change without notice. Give feedback - https://github.com/BerriAI/litellm/issues
|
|
|
|
Get all users who are a specific `user_role`.
|
|
|
|
Used by the UI to populate the user lists.
|
|
|
|
Currently - admin-only endpoint.
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={"error": f"No db connected. prisma client={prisma_client}"},
|
|
)
|
|
all_users = await prisma_client.get_data(
|
|
table_name="user", query_type="find_all", key_val={"user_role": role}
|
|
)
|
|
|
|
return all_users
|
|
|
|
|
|
#### TEAM MANAGEMENT ####
|
|
|
|
|
|
@router.post(
|
|
"/team/new",
|
|
tags=["team management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_model=LiteLLM_TeamTable,
|
|
)
|
|
async def new_team(
|
|
data: NewTeamRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Allow users to create a new team. Apply user permissions to their team.
|
|
|
|
[ASK FOR HELP](https://github.com/BerriAI/litellm/issues)
|
|
|
|
Parameters:
|
|
- team_alias: Optional[str] - User defined team alias
|
|
- team_id: Optional[str] - The team id of the user. If none passed, we'll generate it.
|
|
- members_with_roles: List[{"role": "admin" or "user", "user_id": "<user-id>"}] - A list of users and their roles in the team. Get user_id when making a new user via `/user/new`.
|
|
- metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
|
- tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit
|
|
- rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
|
|
- max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
|
|
- models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
|
|
- blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id.
|
|
|
|
Returns:
|
|
- team_id: (str) Unique team id - used for tracking spend across multiple keys for same team id.
|
|
|
|
_deprecated_params:
|
|
- admins: list - A list of user_id's for the admin role
|
|
- users: list - A list of user_id's for the user role
|
|
|
|
Example Request:
|
|
```
|
|
curl --location 'http://0.0.0.0:8000/team/new' \
|
|
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
|
|
--header 'Content-Type: application/json' \
|
|
|
|
--data '{
|
|
"team_alias": "my-new-team_2",
|
|
"members_with_roles": [{"role": "admin", "user_id": "user-1234"},
|
|
{"role": "user", "user_id": "user-2434"}]
|
|
}'
|
|
|
|
```
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if data.team_id is None:
|
|
data.team_id = str(uuid.uuid4())
|
|
|
|
if (
|
|
user_api_key_dict.user_role is None
|
|
or user_api_key_dict.user_role != "proxy_admin"
|
|
): # don't restrict proxy admin
|
|
if (
|
|
data.tpm_limit is not None
|
|
and user_api_key_dict.tpm_limit is not None
|
|
and data.tpm_limit > user_api_key_dict.tpm_limit
|
|
):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}. User role={user_api_key_dict.user_role}"
|
|
},
|
|
)
|
|
|
|
if (
|
|
data.rpm_limit is not None
|
|
and user_api_key_dict.rpm_limit is not None
|
|
and data.rpm_limit > user_api_key_dict.rpm_limit
|
|
):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}. User role={user_api_key_dict.user_role}"
|
|
},
|
|
)
|
|
|
|
if (
|
|
data.max_budget is not None
|
|
and user_api_key_dict.max_budget is not None
|
|
and data.max_budget > user_api_key_dict.max_budget
|
|
):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"max budget higher than user max. User max budget={user_api_key_dict.max_budget}. User role={user_api_key_dict.user_role}"
|
|
},
|
|
)
|
|
|
|
if data.models is not None and len(user_api_key_dict.models) > 0:
|
|
for m in data.models:
|
|
if m not in user_api_key_dict.models:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}. User id={user_api_key_dict.user_id}"
|
|
},
|
|
)
|
|
|
|
if user_api_key_dict.user_id is not None:
|
|
creating_user_in_list = False
|
|
for member in data.members_with_roles:
|
|
if member.user_id == user_api_key_dict.user_id:
|
|
creating_user_in_list = True
|
|
|
|
if creating_user_in_list == False:
|
|
data.members_with_roles.append(
|
|
Member(role="admin", user_id=user_api_key_dict.user_id)
|
|
)
|
|
|
|
## ADD TO MODEL TABLE
|
|
_model_id = None
|
|
if data.model_aliases is not None and isinstance(data.model_aliases, dict):
|
|
litellm_modeltable = LiteLLM_ModelTable(
|
|
model_aliases=json.dumps(data.model_aliases),
|
|
created_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
)
|
|
model_dict = await prisma_client.db.litellm_modeltable.create(
|
|
{**litellm_modeltable.json(exclude_none=True)} # type: ignore
|
|
) # type: ignore
|
|
|
|
_model_id = model_dict.id
|
|
|
|
## ADD TO TEAM TABLE
|
|
complete_team_data = LiteLLM_TeamTable(
|
|
**data.json(),
|
|
max_parallel_requests=user_api_key_dict.max_parallel_requests,
|
|
budget_duration=user_api_key_dict.budget_duration,
|
|
budget_reset_at=user_api_key_dict.budget_reset_at,
|
|
model_id=_model_id,
|
|
)
|
|
|
|
team_row = await prisma_client.insert_data(
|
|
data=complete_team_data.json(exclude_none=True), table_name="team"
|
|
)
|
|
|
|
## ADD TEAM ID TO USER TABLE ##
|
|
for user in complete_team_data.members_with_roles:
|
|
## add team id to user row ##
|
|
await prisma_client.update_data(
|
|
user_id=user.user_id,
|
|
data={"user_id": user.user_id, "teams": [team_row.team_id]},
|
|
update_key_values_custom_query={
|
|
"teams": {
|
|
"push ": [team_row.team_id],
|
|
}
|
|
},
|
|
)
|
|
try:
|
|
return team_row.model_dump()
|
|
except Exception as e:
|
|
return team_row.dict()
|
|
|
|
|
|
@router.post(
|
|
"/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def update_team(
|
|
data: UpdateTeamRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Use `/team/member_add` AND `/team/member/delete` to add/remove new team members
|
|
|
|
You can now update team budget / rate limits via /team/update
|
|
|
|
Parameters:
|
|
- team_id: str - The team id of the user. Required param.
|
|
- team_alias: Optional[str] - User defined team alias
|
|
- metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
|
- tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit
|
|
- rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
|
|
- max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
|
|
- models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
|
|
- blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id.
|
|
|
|
Example - update team TPM Limit
|
|
|
|
```
|
|
curl --location 'http://0.0.0.0:8000/team/update' \
|
|
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
|
|
--header 'Content-Type: application/json' \
|
|
|
|
--data-raw '{
|
|
"team_id": "litellm-test-client-id-new",
|
|
"tpm_limit": 100
|
|
}'
|
|
```
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if data.team_id is None:
|
|
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
|
|
verbose_proxy_logger.debug("/team/update - %s", data)
|
|
|
|
existing_team_row = await prisma_client.get_data(
|
|
team_id=data.team_id, table_name="team", query_type="find_unique"
|
|
)
|
|
|
|
updated_kv = data.json(exclude_none=True)
|
|
team_row = await prisma_client.update_data(
|
|
update_key_values=updated_kv,
|
|
data=updated_kv,
|
|
table_name="team",
|
|
team_id=data.team_id,
|
|
)
|
|
|
|
return team_row
|
|
|
|
|
|
@router.post(
|
|
"/team/member_add",
|
|
tags=["team management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def team_member_add(
|
|
data: TeamMemberAddRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
[BETA]
|
|
|
|
Add new members (either via user_email or user_id) to a team
|
|
|
|
If user doesn't exist, new user row will also be added to User Table
|
|
|
|
```
|
|
curl -X POST 'http://0.0.0.0:8000/team/update' \
|
|
|
|
-H 'Authorization: Bearer sk-1234' \
|
|
|
|
-H 'Content-Type: application/json' \
|
|
|
|
-D '{
|
|
"team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849",
|
|
"member": {"role": "user", "user_id": "krrish247652@berri.ai"}
|
|
}'
|
|
```
|
|
"""
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if data.team_id is None:
|
|
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
|
|
|
|
if data.member is None:
|
|
raise HTTPException(status_code=400, detail={"error": "No member passed in"})
|
|
|
|
existing_team_row = await prisma_client.get_data( # type: ignore
|
|
team_id=data.team_id, table_name="team", query_type="find_unique"
|
|
)
|
|
if existing_team_row is None:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail={
|
|
"error": f"Team not found for team_id={getattr(data, 'team_id', None)}"
|
|
},
|
|
)
|
|
|
|
new_member = data.member
|
|
|
|
existing_team_row.members_with_roles.append(new_member)
|
|
|
|
complete_team_data = LiteLLM_TeamTable(
|
|
**_get_pydantic_json_dict(existing_team_row),
|
|
)
|
|
|
|
team_row = await prisma_client.update_data(
|
|
update_key_values=complete_team_data.json(exclude_none=True),
|
|
data=complete_team_data.json(exclude_none=True),
|
|
table_name="team",
|
|
team_id=data.team_id,
|
|
)
|
|
|
|
## ADD USER, IF NEW ##
|
|
user_data = { # type: ignore
|
|
"teams": [team_row["team_id"]],
|
|
"models": team_row["data"].models,
|
|
}
|
|
if new_member.user_id is not None:
|
|
user_data["user_id"] = new_member.user_id # type: ignore
|
|
await prisma_client.update_data(
|
|
user_id=new_member.user_id,
|
|
data=user_data,
|
|
update_key_values_custom_query={
|
|
"teams": {
|
|
"push": [team_row["team_id"]],
|
|
}
|
|
},
|
|
table_name="user",
|
|
)
|
|
elif new_member.user_email is not None:
|
|
user_data["user_id"] = str(uuid.uuid4())
|
|
user_data["user_email"] = new_member.user_email
|
|
## user email is not unique acc. to prisma schema -> future improvement
|
|
### for now: check if it exists in db, if not - insert it
|
|
existing_user_row = await prisma_client.get_data(
|
|
key_val={"user_email": new_member.user_email},
|
|
table_name="user",
|
|
query_type="find_all",
|
|
)
|
|
if existing_user_row is None or (
|
|
isinstance(existing_user_row, list) and len(existing_user_row) == 0
|
|
):
|
|
|
|
await prisma_client.insert_data(data=user_data, table_name="user")
|
|
|
|
return team_row
|
|
|
|
|
|
@router.post(
|
|
"/team/member_delete",
|
|
tags=["team management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def team_member_delete(
|
|
data: TeamMemberDeleteRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
[BETA]
|
|
|
|
delete members (either via user_email or user_id) from a team
|
|
|
|
If user doesn't exist, an exception will be raised
|
|
```
|
|
curl -X POST 'http://0.0.0.0:8000/team/update' \
|
|
|
|
-H 'Authorization: Bearer sk-1234' \
|
|
|
|
-H 'Content-Type: application/json' \
|
|
|
|
-D '{
|
|
"team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849",
|
|
"user_id": "krrish247652@berri.ai"
|
|
}'
|
|
```
|
|
"""
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if data.team_id is None:
|
|
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
|
|
|
|
if data.user_id is None and data.user_email is None:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={"error": "Either user_id or user_email needs to be passed in"},
|
|
)
|
|
|
|
existing_team_row = await prisma_client.get_data( # type: ignore
|
|
team_id=data.team_id, table_name="team", query_type="find_unique"
|
|
)
|
|
|
|
## DELETE MEMBER FROM TEAM
|
|
new_team_members = []
|
|
for m in existing_team_row.members_with_roles:
|
|
if (
|
|
data.user_id is not None
|
|
and m["user_id"] is not None
|
|
and data.user_id == m["user_id"]
|
|
):
|
|
continue
|
|
elif (
|
|
data.user_email is not None
|
|
and m["user_email"] is not None
|
|
and data.user_email == m["user_email"]
|
|
):
|
|
continue
|
|
new_team_members.append(m)
|
|
existing_team_row.members_with_roles = new_team_members
|
|
complete_team_data = LiteLLM_TeamTable(
|
|
**existing_team_row.model_dump(),
|
|
)
|
|
|
|
team_row = await prisma_client.update_data(
|
|
update_key_values=complete_team_data.json(exclude_none=True),
|
|
data=complete_team_data.json(exclude_none=True),
|
|
table_name="team",
|
|
team_id=data.team_id,
|
|
)
|
|
|
|
## DELETE TEAM ID from USER ROW, IF EXISTS ##
|
|
# get user row
|
|
key_val = {}
|
|
if data.user_id is not None:
|
|
key_val["user_id"] = data.user_id
|
|
elif data.user_email is not None:
|
|
key_val["user_email"] = data.user_email
|
|
existing_user_rows = await prisma_client.get_data(
|
|
key_val=key_val,
|
|
table_name="user",
|
|
query_type="find_all",
|
|
)
|
|
user_data = { # type: ignore
|
|
"teams": [],
|
|
"models": team_row["data"].models,
|
|
}
|
|
if existing_user_rows is not None and (
|
|
isinstance(existing_user_rows, list) and len(existing_user_rows) > 0
|
|
):
|
|
for existing_user in existing_user_rows:
|
|
team_list = []
|
|
if hasattr(existing_user, "teams"):
|
|
team_list = existing_user.teams
|
|
team_list.remove(data.team_id)
|
|
user_data["user_id"] = existing_user.user_id
|
|
await prisma_client.update_data(
|
|
user_id=existing_user.user_id,
|
|
data=user_data,
|
|
update_key_values_custom_query={
|
|
"teams": {
|
|
"set": [team_row["team_id"]],
|
|
}
|
|
},
|
|
table_name="user",
|
|
)
|
|
|
|
return team_row["data"]
|
|
|
|
|
|
@router.post(
|
|
"/team/delete", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def delete_team(
|
|
data: DeleteTeamRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
delete team and associated team keys
|
|
|
|
```
|
|
curl --location 'http://0.0.0.0:8000/team/delete' \
|
|
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
|
|
--header 'Content-Type: application/json' \
|
|
|
|
--data-raw '{
|
|
"team_ids": ["45e3e396-ee08-4a61-a88e-16b3ce7e0849"]
|
|
}'
|
|
```
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if data.team_ids is None:
|
|
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
|
|
|
|
## DELETE ASSOCIATED KEYS
|
|
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
|
|
## DELETE TEAMS
|
|
deleted_teams = await prisma_client.delete_data(
|
|
team_id_list=data.team_ids, table_name="team"
|
|
)
|
|
return deleted_teams
|
|
|
|
|
|
@router.get(
|
|
"/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def team_info(
|
|
team_id: str = fastapi.Query(
|
|
default=None, description="Team ID in the request parameters"
|
|
)
|
|
):
|
|
"""
|
|
get info on team + related keys
|
|
|
|
```
|
|
curl --location 'http://localhost:4000/team/info' \
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
--header 'Content-Type: application/json' \
|
|
--data '{
|
|
"teams": ["<team-id>",..]
|
|
}'
|
|
```
|
|
"""
|
|
global prisma_client
|
|
try:
|
|
if prisma_client is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail={
|
|
"error": f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
|
},
|
|
)
|
|
if team_id is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
|
detail={"message": "Malformed request. No team id passed in."},
|
|
)
|
|
|
|
team_info = await prisma_client.get_data(
|
|
team_id=team_id, table_name="team", query_type="find_unique"
|
|
)
|
|
## GET ALL KEYS ##
|
|
keys = await prisma_client.get_data(
|
|
team_id=team_id,
|
|
table_name="key",
|
|
query_type="find_all",
|
|
expires=datetime.now(),
|
|
)
|
|
|
|
if team_info is None:
|
|
## make sure we still return a total spend ##
|
|
spend = 0
|
|
for k in keys:
|
|
spend += getattr(k, "spend", 0)
|
|
team_info = {"spend": spend}
|
|
|
|
## REMOVE HASHED TOKEN INFO before returning ##
|
|
for key in keys:
|
|
try:
|
|
key = key.model_dump() # noqa
|
|
except:
|
|
# if using pydantic v1
|
|
key = key.dict()
|
|
key.pop("token", None)
|
|
return {"team_id": team_id, "team_info": team_info, "keys": keys}
|
|
|
|
except Exception as e:
|
|
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,
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/team/block", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def block_team(
|
|
data: BlockTeamRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Blocks all calls from keys with this team id.
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise Exception("No DB Connected.")
|
|
|
|
record = await prisma_client.db.litellm_teamtable.update(
|
|
where={"team_id": data.team_id}, data={"blocked": True} # type: ignore
|
|
)
|
|
|
|
return record
|
|
|
|
|
|
@router.post(
|
|
"/team/unblock", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
|
)
|
|
async def unblock_team(
|
|
data: BlockTeamRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Blocks all calls from keys with this team id.
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise Exception("No DB Connected.")
|
|
|
|
record = await prisma_client.db.litellm_teamtable.update(
|
|
where={"team_id": data.team_id}, data={"blocked": False} # type: ignore
|
|
)
|
|
|
|
return record
|
|
|
|
|
|
#### ORGANIZATION MANAGEMENT ####
|
|
|
|
|
|
@router.post(
|
|
"/organization/new",
|
|
tags=["organization management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
response_model=NewOrganizationResponse,
|
|
)
|
|
async def new_organization(
|
|
data: NewOrganizationRequest,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
Allow orgs to own teams
|
|
|
|
Set org level budgets + model access.
|
|
|
|
Only admins can create orgs.
|
|
|
|
# Parameters
|
|
|
|
- `organization_alias`: *str* = The name of the organization.
|
|
- `models`: *List* = The models the organization has access to.
|
|
- `budget_id`: *Optional[str]* = The id for a budget (tpm/rpm/max budget) for the organization.
|
|
### IF NO BUDGET ID - CREATE ONE WITH THESE PARAMS ###
|
|
- `max_budget`: *Optional[float]* = Max budget for org
|
|
- `tpm_limit`: *Optional[int]* = Max tpm limit for org
|
|
- `rpm_limit`: *Optional[int]* = Max rpm limit for org
|
|
- `model_max_budget`: *Optional[dict]* = Max budget for a specific model
|
|
- `budget_duration`: *Optional[str]* = Frequency of reseting org budget
|
|
|
|
Case 1: Create new org **without** a budget_id
|
|
|
|
```bash
|
|
curl --location 'http://0.0.0.0:4000/organization/new' \
|
|
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
|
|
--header 'Content-Type: application/json' \
|
|
|
|
--data '{
|
|
"organization_alias": "my-secret-org",
|
|
"models": ["model1", "model2"],
|
|
"max_budget": 100
|
|
}'
|
|
|
|
|
|
```
|
|
|
|
Case 2: Create new org **with** a budget_id
|
|
|
|
```bash
|
|
curl --location 'http://0.0.0.0:4000/organization/new' \
|
|
|
|
--header 'Authorization: Bearer sk-1234' \
|
|
|
|
--header 'Content-Type: application/json' \
|
|
|
|
--data '{
|
|
"organization_alias": "my-secret-org",
|
|
"models": ["model1", "model2"],
|
|
"budget_id": "428eeaa8-f3ac-4e85-a8fb-7dc8d7aa8689"
|
|
}'
|
|
```
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if (
|
|
user_api_key_dict.user_role is None
|
|
or user_api_key_dict.user_role != "proxy_admin"
|
|
):
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail={
|
|
"error": f"Only admins can create orgs. Your role is = {user_api_key_dict.user_role}"
|
|
},
|
|
)
|
|
|
|
if data.budget_id is None:
|
|
"""
|
|
Every organization needs a budget attached.
|
|
|
|
If none provided, create one based on provided values
|
|
"""
|
|
budget_row = LiteLLM_BudgetTable(**data.json(exclude_none=True))
|
|
|
|
new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
|
|
|
|
_budget = await prisma_client.db.litellm_budgettable.create(
|
|
data={
|
|
**new_budget, # type: ignore
|
|
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
}
|
|
) # type: ignore
|
|
|
|
data.budget_id = _budget.budget_id
|
|
|
|
"""
|
|
Ensure only models that user has access to, are given to org
|
|
"""
|
|
if len(user_api_key_dict.models) == 0: # user has access to all models
|
|
pass
|
|
else:
|
|
if len(data.models) == 0:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"User not allowed to give access to all models. Select models you want org to have access to."
|
|
},
|
|
)
|
|
for m in data.models:
|
|
if m not in user_api_key_dict.models:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"User not allowed to give access to model={m}. Models you have access to = {user_api_key_dict.models}"
|
|
},
|
|
)
|
|
organization_row = LiteLLM_OrganizationTable(
|
|
**data.json(exclude_none=True),
|
|
created_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
)
|
|
new_organization_row = prisma_client.jsonify_object(
|
|
organization_row.json(exclude_none=True)
|
|
)
|
|
response = await prisma_client.db.litellm_organizationtable.create(
|
|
data={
|
|
**new_organization_row, # type: ignore
|
|
}
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
@router.post(
|
|
"/organization/update",
|
|
tags=["organization management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def update_organization():
|
|
"""[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues"""
|
|
pass
|
|
|
|
|
|
@router.post(
|
|
"/organization/delete",
|
|
tags=["organization management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def delete_organization():
|
|
"""[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues"""
|
|
pass
|
|
|
|
|
|
@router.post(
|
|
"/organization/info",
|
|
tags=["organization management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def info_organization(data: OrganizationRequest):
|
|
"""
|
|
Get the org specific information
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if len(data.organizations) == 0:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Specify list of organization id's to query. Passed in={data.organizations}"
|
|
},
|
|
)
|
|
response = await prisma_client.db.litellm_organizationtable.find_many(
|
|
where={"organization_id": {"in": data.organizations}},
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
#### BUDGET TABLE MANAGEMENT ####
|
|
|
|
|
|
@router.post(
|
|
"/budget/info",
|
|
tags=["organization management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def info_budget(data: BudgetRequest):
|
|
"""
|
|
Get the budget id specific information
|
|
"""
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
|
|
|
if len(data.budgets) == 0:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Specify list of budget id's to query. Passed in={data.budgets}"
|
|
},
|
|
)
|
|
response = await prisma_client.db.litellm_budgettable.find_many(
|
|
where={"budget_id": {"in": data.budgets}},
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
#### MODEL MANAGEMENT ####
|
|
|
|
|
|
#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
|
|
@router.post(
|
|
"/model/new",
|
|
description="Allows adding new models to the model list in the config.yaml",
|
|
tags=["model management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def add_new_model(
|
|
model_params: Deployment,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config, prisma_client, master_key, store_model_in_db, proxy_logging_obj
|
|
try:
|
|
import base64
|
|
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys"
|
|
},
|
|
)
|
|
|
|
model_response = None
|
|
# update DB
|
|
if store_model_in_db == True:
|
|
"""
|
|
- store model_list in db
|
|
- store keys separately
|
|
"""
|
|
# encrypt litellm params #
|
|
_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True)
|
|
for k, v in _litellm_params_dict.items():
|
|
if isinstance(v, str):
|
|
encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore
|
|
model_params.litellm_params[k] = base64.b64encode(
|
|
encrypted_value
|
|
).decode("utf-8")
|
|
_data: dict = {
|
|
"model_id": model_params.model_info.id,
|
|
"model_name": model_params.model_name,
|
|
"litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), # type: ignore
|
|
"model_info": model_params.model_info.model_dump_json( # type: ignore
|
|
exclude_none=True
|
|
),
|
|
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
}
|
|
if model_params.model_info.id is not None:
|
|
_data["model_id"] = model_params.model_info.id
|
|
model_response = await prisma_client.db.litellm_proxymodeltable.create(
|
|
data=_data # type: ignore
|
|
)
|
|
|
|
await proxy_config.add_deployment(
|
|
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."
|
|
},
|
|
)
|
|
|
|
return model_response
|
|
|
|
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 ####
|
|
@router.post(
|
|
"/model/update",
|
|
description="Edit existing model params",
|
|
tags=["model management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def update_model(
|
|
model_params: updateDeployment,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config, prisma_client, master_key, store_model_in_db, proxy_logging_obj
|
|
try:
|
|
import base64
|
|
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys"
|
|
},
|
|
)
|
|
# update DB
|
|
if store_model_in_db == True:
|
|
_model_id = None
|
|
_model_info = getattr(model_params, "model_info", None)
|
|
if _model_info is None:
|
|
raise Exception("model_info not provided")
|
|
|
|
_model_id = _model_info.id
|
|
if _model_id is None:
|
|
raise Exception("model_info.id not provided")
|
|
_existing_litellm_params = (
|
|
await prisma_client.db.litellm_proxymodeltable.find_unique(
|
|
where={"model_id": _model_id}
|
|
)
|
|
)
|
|
if _existing_litellm_params is None:
|
|
raise Exception("model not found")
|
|
_existing_litellm_params_dict = dict(
|
|
_existing_litellm_params.litellm_params
|
|
)
|
|
|
|
if model_params.litellm_params is None:
|
|
raise Exception("litellm_params not provided")
|
|
|
|
_new_litellm_params_dict = model_params.litellm_params.dict(
|
|
exclude_none=True
|
|
)
|
|
|
|
for key, value in _existing_litellm_params_dict.items():
|
|
if key in _new_litellm_params_dict:
|
|
_existing_litellm_params_dict[key] = _new_litellm_params_dict[key]
|
|
|
|
_data: dict = {
|
|
"litellm_params": json.dumps(_existing_litellm_params_dict), # type: ignore
|
|
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
|
}
|
|
model_response = await prisma_client.db.litellm_proxymodeltable.update(
|
|
where={"model_id": _model_id},
|
|
data=_data, # type: ignore
|
|
)
|
|
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,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/v2/model/info",
|
|
description="v2 - returns all the models set on the config.yaml, shows 'user_access' = True if the user has access to the model. Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",
|
|
tags=["model management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
include_in_schema=False,
|
|
)
|
|
async def model_info_v2(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
"""
|
|
BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now.
|
|
"""
|
|
global llm_model_list, general_settings, user_config_file_path, proxy_config
|
|
|
|
if llm_model_list is None or not isinstance(llm_model_list, list):
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": f"Invalid llm model list. llm_model_list={llm_model_list}"
|
|
},
|
|
)
|
|
|
|
# Load existing config
|
|
config = await proxy_config.get_config()
|
|
|
|
all_models = copy.deepcopy(llm_model_list)
|
|
if user_model is not None:
|
|
# if user does not use a config.yaml, https://github.com/BerriAI/litellm/issues/2061
|
|
all_models += [user_model]
|
|
|
|
# check all models user has access to in user_api_key_dict
|
|
user_models = []
|
|
if len(user_api_key_dict.models) > 0:
|
|
user_models = user_api_key_dict.models
|
|
|
|
# fill in model info based on config.yaml and litellm model_prices_and_context_window.json
|
|
for model in all_models:
|
|
# provided model_info in config.yaml
|
|
model_info = model.get("model_info", {})
|
|
|
|
# read litellm model_prices_and_context_window.json to get the following:
|
|
# input_cost_per_token, output_cost_per_token, max_tokens
|
|
litellm_model_info = get_litellm_model_info(model=model)
|
|
|
|
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
|
|
if litellm_model_info == {}:
|
|
# use litellm_param model_name to get model_info
|
|
litellm_params = model.get("litellm_params", {})
|
|
litellm_model = litellm_params.get("model", None)
|
|
try:
|
|
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
|
except:
|
|
litellm_model_info = {}
|
|
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
|
|
if litellm_model_info == {}:
|
|
# use litellm_param model_name to get model_info
|
|
litellm_params = model.get("litellm_params", {})
|
|
litellm_model = litellm_params.get("model", None)
|
|
split_model = litellm_model.split("/")
|
|
if len(split_model) > 0:
|
|
litellm_model = split_model[-1]
|
|
try:
|
|
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
|
except:
|
|
litellm_model_info = {}
|
|
for k, v in litellm_model_info.items():
|
|
if k not in model_info:
|
|
model_info[k] = v
|
|
model["model_info"] = model_info
|
|
# don't return the api key / vertex credentials
|
|
model["litellm_params"].pop("api_key", None)
|
|
model["litellm_params"].pop("vertex_credentials", None)
|
|
|
|
verbose_proxy_logger.debug("all_models: %s", all_models)
|
|
return {"data": all_models}
|
|
|
|
|
|
@router.get(
|
|
"/model/metrics",
|
|
description="View number of requests & avg latency per model on config.yaml",
|
|
tags=["model management"],
|
|
include_in_schema=False,
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def model_metrics(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
_selected_model_group: Optional[str] = "gpt-4-32k",
|
|
startTime: Optional[datetime] = None,
|
|
endTime: Optional[datetime] = None,
|
|
):
|
|
global prisma_client, llm_router
|
|
if prisma_client is None:
|
|
raise ProxyException(
|
|
message="Prisma Client is not initialized",
|
|
type="internal_error",
|
|
param="None",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
startTime = startTime or datetime.now() - timedelta(days=30)
|
|
endTime = endTime or datetime.now()
|
|
|
|
sql_query = """
|
|
SELECT
|
|
api_base,
|
|
model,
|
|
DATE_TRUNC('day', "startTime")::DATE AS day,
|
|
AVG(EXTRACT(epoch FROM ("endTime" - "startTime"))) / SUM(total_tokens) AS avg_latency_per_token
|
|
FROM
|
|
"LiteLLM_SpendLogs"
|
|
WHERE
|
|
"startTime" >= NOW() - INTERVAL '30 days'
|
|
AND "model" = $1 AND "cache_hit" != 'True'
|
|
GROUP BY
|
|
api_base,
|
|
model,
|
|
day
|
|
HAVING
|
|
SUM(total_tokens) > 0
|
|
ORDER BY
|
|
avg_latency_per_token DESC;
|
|
"""
|
|
_all_api_bases = set()
|
|
db_response = await prisma_client.db.query_raw(
|
|
sql_query, _selected_model_group, startTime, endTime
|
|
)
|
|
_daily_entries: dict = {} # {"Jun 23": {"model1": 0.002, "model2": 0.003}}
|
|
if db_response is not None:
|
|
for model_data in db_response:
|
|
_api_base = model_data["api_base"]
|
|
_model = model_data["model"]
|
|
_day = model_data["day"]
|
|
_avg_latency_per_token = model_data["avg_latency_per_token"]
|
|
if _day not in _daily_entries:
|
|
_daily_entries[_day] = {}
|
|
_combined_model_name = str(_model)
|
|
if "https://" in _api_base:
|
|
_combined_model_name = str(_api_base)
|
|
if "/openai/" in _combined_model_name:
|
|
_combined_model_name = _combined_model_name.split("/openai/")[0]
|
|
|
|
_all_api_bases.add(_combined_model_name)
|
|
_daily_entries[_day][_combined_model_name] = _avg_latency_per_token
|
|
|
|
"""
|
|
each entry needs to be like this:
|
|
{
|
|
date: 'Jun 23',
|
|
'gpt-4-https://api.openai.com/v1/': 0.002,
|
|
'gpt-43-https://api.openai.com-12/v1/': 0.002,
|
|
}
|
|
"""
|
|
# convert daily entries to list of dicts
|
|
|
|
response: List[dict] = []
|
|
|
|
# sort daily entries by date
|
|
_daily_entries = dict(sorted(_daily_entries.items(), key=lambda item: item[0]))
|
|
for day in _daily_entries:
|
|
entry = {"date": str(day)}
|
|
for model_key, latency in _daily_entries[day].items():
|
|
entry[model_key] = latency
|
|
response.append(entry)
|
|
|
|
return {
|
|
"data": response,
|
|
"all_api_bases": list(_all_api_bases),
|
|
}
|
|
|
|
|
|
@router.get(
|
|
"/model/metrics/slow_responses",
|
|
description="View number of hanging requests per model_group",
|
|
tags=["model management"],
|
|
include_in_schema=False,
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def model_metrics_slow_responses(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
_selected_model_group: Optional[str] = "gpt-4-32k",
|
|
startTime: Optional[datetime] = None,
|
|
endTime: Optional[datetime] = None,
|
|
):
|
|
global prisma_client, llm_router, proxy_logging_obj
|
|
if prisma_client is None:
|
|
raise ProxyException(
|
|
message="Prisma Client is not initialized",
|
|
type="internal_error",
|
|
param="None",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
startTime = startTime or datetime.now() - timedelta(days=30)
|
|
endTime = endTime or datetime.now()
|
|
|
|
alerting_threshold = (
|
|
proxy_logging_obj.slack_alerting_instance.alerting_threshold or 300
|
|
)
|
|
alerting_threshold = int(alerting_threshold)
|
|
|
|
sql_query = """
|
|
SELECT
|
|
api_base,
|
|
COUNT(*) AS total_count,
|
|
SUM(CASE
|
|
WHEN ("endTime" - "startTime") >= (INTERVAL '1 SECOND' * CAST($1 AS INTEGER)) THEN 1
|
|
ELSE 0
|
|
END) AS slow_count
|
|
FROM
|
|
"LiteLLM_SpendLogs"
|
|
WHERE
|
|
"model" = $2
|
|
AND "cache_hit" != 'True'
|
|
GROUP BY
|
|
api_base
|
|
ORDER BY
|
|
slow_count DESC;
|
|
"""
|
|
|
|
db_response = await prisma_client.db.query_raw(
|
|
sql_query, alerting_threshold, _selected_model_group
|
|
)
|
|
|
|
if db_response is not None:
|
|
for row in db_response:
|
|
_api_base = row.get("api_base") or ""
|
|
if "/openai/" in _api_base:
|
|
_api_base = _api_base.split("/openai/")[0]
|
|
row["api_base"] = _api_base
|
|
return db_response
|
|
|
|
|
|
@router.get(
|
|
"/model/metrics/exceptions",
|
|
description="View number of failed requests per model on config.yaml",
|
|
tags=["model management"],
|
|
include_in_schema=False,
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def model_metrics_exceptions(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
_selected_model_group: Optional[str] = None,
|
|
startTime: Optional[datetime] = None,
|
|
endTime: Optional[datetime] = None,
|
|
):
|
|
global prisma_client, llm_router
|
|
if prisma_client is None:
|
|
raise ProxyException(
|
|
message="Prisma Client is not initialized",
|
|
type="internal_error",
|
|
param="None",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
startTime = startTime or datetime.now() - timedelta(days=30)
|
|
endTime = endTime or datetime.now()
|
|
|
|
"""
|
|
"""
|
|
sql_query = """
|
|
WITH cte AS (
|
|
SELECT
|
|
CASE WHEN api_base = '' THEN litellm_model_name ELSE CONCAT(litellm_model_name, '-', api_base) END AS combined_model_api_base,
|
|
exception_type,
|
|
COUNT(*) AS num_exceptions
|
|
FROM "LiteLLM_ErrorLogs"
|
|
WHERE "startTime" >= $1::timestamp AND "endTime" <= $2::timestamp
|
|
GROUP BY combined_model_api_base, exception_type
|
|
)
|
|
SELECT
|
|
combined_model_api_base,
|
|
COUNT(*) AS total_exceptions,
|
|
json_object_agg(exception_type, num_exceptions) AS exception_counts
|
|
FROM cte
|
|
GROUP BY combined_model_api_base
|
|
ORDER BY total_exceptions DESC
|
|
LIMIT 200;
|
|
"""
|
|
db_response = await prisma_client.db.query_raw(sql_query, startTime, endTime)
|
|
response: List[dict] = []
|
|
exception_types = set()
|
|
|
|
"""
|
|
Return Data
|
|
{
|
|
"combined_model_api_base": "gpt-3.5-turbo-https://api.openai.com/v1/,
|
|
"total_exceptions": 5,
|
|
"BadRequestException": 5,
|
|
"TimeoutException": 2
|
|
}
|
|
"""
|
|
|
|
if db_response is not None:
|
|
# loop through all models
|
|
for model_data in db_response:
|
|
model = model_data.get("combined_model_api_base", "")
|
|
total_exceptions = model_data.get("total_exceptions", 0)
|
|
exception_counts = model_data.get("exception_counts", {})
|
|
curr_row = {
|
|
"model": model,
|
|
"total_exceptions": total_exceptions,
|
|
}
|
|
curr_row.update(exception_counts)
|
|
response.append(curr_row)
|
|
for k, v in exception_counts.items():
|
|
exception_types.add(k)
|
|
|
|
return {"data": response, "exception_types": list(exception_types)}
|
|
|
|
|
|
@router.get(
|
|
"/model/info",
|
|
description="Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",
|
|
tags=["model management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
@router.get(
|
|
"/v1/model/info",
|
|
description="Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",
|
|
tags=["model management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def model_info_v1(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global llm_model_list, general_settings, user_config_file_path, proxy_config
|
|
|
|
if llm_model_list is None:
|
|
raise HTTPException(
|
|
status_code=500, detail={"error": "LLM Model List not loaded in"}
|
|
)
|
|
|
|
if len(user_api_key_dict.models) > 0:
|
|
model_names = user_api_key_dict.models
|
|
_relevant_models = [m for m in llm_model_list if m["model_name"] in model_names]
|
|
all_models = copy.deepcopy(_relevant_models)
|
|
else:
|
|
all_models = copy.deepcopy(llm_model_list)
|
|
for model in all_models:
|
|
# provided model_info in config.yaml
|
|
model_info = model.get("model_info", {})
|
|
|
|
# read litellm model_prices_and_context_window.json to get the following:
|
|
# input_cost_per_token, output_cost_per_token, max_tokens
|
|
litellm_model_info = get_litellm_model_info(model=model)
|
|
|
|
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
|
|
if litellm_model_info == {}:
|
|
# use litellm_param model_name to get model_info
|
|
litellm_params = model.get("litellm_params", {})
|
|
litellm_model = litellm_params.get("model", None)
|
|
try:
|
|
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
|
except:
|
|
litellm_model_info = {}
|
|
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
|
|
if litellm_model_info == {}:
|
|
# use litellm_param model_name to get model_info
|
|
litellm_params = model.get("litellm_params", {})
|
|
litellm_model = litellm_params.get("model", None)
|
|
split_model = litellm_model.split("/")
|
|
if len(split_model) > 0:
|
|
litellm_model = split_model[-1]
|
|
try:
|
|
litellm_model_info = litellm.get_model_info(model=litellm_model)
|
|
except:
|
|
litellm_model_info = {}
|
|
for k, v in litellm_model_info.items():
|
|
if k not in model_info:
|
|
model_info[k] = v
|
|
model["model_info"] = model_info
|
|
# don't return the api key
|
|
model["litellm_params"].pop("api_key", None)
|
|
|
|
verbose_proxy_logger.debug("all_models: %s", all_models)
|
|
return {"data": all_models}
|
|
|
|
|
|
#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
|
|
@router.post(
|
|
"/model/delete",
|
|
description="Allows deleting models in the model list in the config.yaml",
|
|
tags=["model management"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def delete_model(model_info: ModelInfoDelete):
|
|
global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config
|
|
try:
|
|
"""
|
|
[BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
|
|
|
|
- Check if id in db
|
|
- Delete
|
|
"""
|
|
|
|
global prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys"
|
|
},
|
|
)
|
|
|
|
# update DB
|
|
if store_model_in_db == True:
|
|
"""
|
|
- store model_list in db
|
|
- store keys separately
|
|
"""
|
|
# encrypt litellm params #
|
|
result = await prisma_client.db.litellm_proxymodeltable.delete(
|
|
where={"model_id": model_info.id}
|
|
)
|
|
|
|
if result is None:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={"error": f"Model with id={model_info.id} not found in db"},
|
|
)
|
|
|
|
## DELETE FROM ROUTER ##
|
|
if llm_router is not None:
|
|
llm_router.delete_deployment(id=model_info.id)
|
|
|
|
return {"message": f"Model: {result.model_id} deleted successfully"}
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail={
|
|
"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."
|
|
},
|
|
)
|
|
|
|
except Exception as e:
|
|
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,
|
|
)
|
|
|
|
|
|
#### EXPERIMENTAL QUEUING ####
|
|
async def _litellm_chat_completions_worker(data, user_api_key_dict):
|
|
"""
|
|
worker to make litellm completions calls
|
|
"""
|
|
while True:
|
|
try:
|
|
### CALL HOOKS ### - modify incoming data before calling the model
|
|
data = await proxy_logging_obj.pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, call_type="completion"
|
|
)
|
|
|
|
verbose_proxy_logger.debug("_litellm_chat_completions_worker started")
|
|
### ROUTE THE REQUEST ###
|
|
router_model_names = (
|
|
[m["model_name"] for m in llm_model_list]
|
|
if llm_model_list is not None
|
|
else []
|
|
)
|
|
if (
|
|
llm_router is not None and data["model"] in router_model_names
|
|
): # model in router model list
|
|
response = await llm_router.acompletion(**data)
|
|
elif (
|
|
llm_router is not None and data["model"] in llm_router.deployment_names
|
|
): # model in router deployments, calling a specific deployment on the router
|
|
response = await llm_router.acompletion(
|
|
**data, specific_deployment=True
|
|
)
|
|
elif (
|
|
llm_router is not None
|
|
and llm_router.model_group_alias is not None
|
|
and data["model"] in llm_router.model_group_alias
|
|
): # model set in model_group_alias
|
|
response = await llm_router.acompletion(**data)
|
|
else: # router is not set
|
|
response = await litellm.acompletion(**data)
|
|
|
|
verbose_proxy_logger.debug("final response: {response}")
|
|
return response
|
|
except HTTPException as e:
|
|
verbose_proxy_logger.debug(
|
|
f"EXCEPTION RAISED IN _litellm_chat_completions_worker - {e.status_code}; {e.detail}"
|
|
)
|
|
if (
|
|
e.status_code == 429
|
|
and "Max parallel request limit reached" in e.detail
|
|
):
|
|
verbose_proxy_logger.debug("Max parallel request limit reached!")
|
|
timeout = litellm._calculate_retry_after(
|
|
remaining_retries=3, max_retries=3, min_timeout=1
|
|
)
|
|
await asyncio.sleep(timeout)
|
|
else:
|
|
raise e
|
|
|
|
|
|
@router.post(
|
|
"/queue/chat/completions",
|
|
tags=["experimental"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def async_queue_request(
|
|
request: Request,
|
|
model: Optional[str] = None,
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
):
|
|
global general_settings, user_debug, proxy_logging_obj
|
|
"""
|
|
v2 attempt at a background worker to handle queuing.
|
|
|
|
Just supports /chat/completion calls currently.
|
|
|
|
Now using a FastAPI background task + /chat/completions compatible endpoint
|
|
"""
|
|
try:
|
|
data = {}
|
|
data = await request.json() # type: ignore
|
|
|
|
# Include original request and headers in the data
|
|
data["proxy_server_request"] = {
|
|
"url": str(request.url),
|
|
"method": request.method,
|
|
"headers": dict(request.headers),
|
|
"body": copy.copy(data), # use copy instead of deepcopy
|
|
}
|
|
|
|
verbose_proxy_logger.debug("receiving data: %s", data)
|
|
data["model"] = (
|
|
general_settings.get("completion_model", None) # server default
|
|
or user_model # model name passed via cli args
|
|
or model # for azure deployments
|
|
or data["model"] # default passed in http request
|
|
)
|
|
|
|
# users can pass in 'user' param to /chat/completions. Don't override it
|
|
if data.get("user", None) is None and user_api_key_dict.user_id is not None:
|
|
# if users are using user_api_key_auth, set `user` in `data`
|
|
data["user"] = user_api_key_dict.user_id
|
|
|
|
if "metadata" not in data:
|
|
data["metadata"] = {}
|
|
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
|
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
|
|
_headers = dict(request.headers)
|
|
_headers.pop(
|
|
"authorization", None
|
|
) # do not store the original `sk-..` api key in the db
|
|
data["metadata"]["headers"] = _headers
|
|
data["metadata"]["user_api_key_alias"] = getattr(
|
|
user_api_key_dict, "key_alias", None
|
|
)
|
|
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
|
|
data["metadata"]["user_api_key_team_id"] = getattr(
|
|
user_api_key_dict, "team_id", None
|
|
)
|
|
data["metadata"]["endpoint"] = str(request.url)
|
|
|
|
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
|
|
# override with user settings, these are params passed via cli
|
|
if user_temperature:
|
|
data["temperature"] = user_temperature
|
|
if user_request_timeout:
|
|
data["request_timeout"] = user_request_timeout
|
|
if user_max_tokens:
|
|
data["max_tokens"] = user_max_tokens
|
|
if user_api_base:
|
|
data["api_base"] = user_api_base
|
|
|
|
response = await asyncio.wait_for(
|
|
_litellm_chat_completions_worker(
|
|
data=data, user_api_key_dict=user_api_key_dict
|
|
),
|
|
timeout=litellm.request_timeout,
|
|
)
|
|
|
|
if (
|
|
"stream" in data and data["stream"] == True
|
|
): # use generate_responses to stream responses
|
|
return StreamingResponse(
|
|
async_data_generator(
|
|
user_api_key_dict=user_api_key_dict, response=response
|
|
),
|
|
media_type="text/event-stream",
|
|
)
|
|
|
|
return response
|
|
except Exception as e:
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
user_api_key_dict=user_api_key_dict, original_exception=e
|
|
)
|
|
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,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/ollama_logs", dependencies=[Depends(user_api_key_auth)], tags=["experimental"]
|
|
)
|
|
async def retrieve_server_log(request: Request):
|
|
filepath = os.path.expanduser("~/.ollama/logs/server.log")
|
|
return FileResponse(filepath)
|
|
|
|
|
|
#### LOGIN ENDPOINTS ####
|
|
|
|
|
|
@app.get("/sso/key/generate", tags=["experimental"])
|
|
async def google_login(request: Request):
|
|
"""
|
|
Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
|
|
PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/"
|
|
Example:
|
|
"""
|
|
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
|
|
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
|
|
generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
|
|
# get url from request
|
|
redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url))
|
|
ui_username = os.getenv("UI_USERNAME")
|
|
if redirect_url.endswith("/"):
|
|
redirect_url += "sso/callback"
|
|
else:
|
|
redirect_url += "/sso/callback"
|
|
# Google SSO Auth
|
|
if google_client_id is not None:
|
|
from fastapi_sso.sso.google import GoogleSSO
|
|
|
|
google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
|
|
if google_client_secret is None:
|
|
raise ProxyException(
|
|
message="GOOGLE_CLIENT_SECRET not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GOOGLE_CLIENT_SECRET",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
google_sso = GoogleSSO(
|
|
client_id=google_client_id,
|
|
client_secret=google_client_secret,
|
|
redirect_uri=redirect_url,
|
|
)
|
|
verbose_proxy_logger.info(
|
|
f"In /google-login/key/generate, \nGOOGLE_REDIRECT_URI: {redirect_url}\nGOOGLE_CLIENT_ID: {google_client_id}"
|
|
)
|
|
with google_sso:
|
|
return await google_sso.get_login_redirect()
|
|
# Microsoft SSO Auth
|
|
elif microsoft_client_id is not None:
|
|
from fastapi_sso.sso.microsoft import MicrosoftSSO
|
|
|
|
microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
|
|
microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
|
|
if microsoft_client_secret is None:
|
|
raise ProxyException(
|
|
message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="MICROSOFT_CLIENT_SECRET",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
microsoft_sso = MicrosoftSSO(
|
|
client_id=microsoft_client_id,
|
|
client_secret=microsoft_client_secret,
|
|
tenant=microsoft_tenant,
|
|
redirect_uri=redirect_url,
|
|
allow_insecure_http=True,
|
|
)
|
|
with microsoft_sso:
|
|
return await microsoft_sso.get_login_redirect()
|
|
elif generic_client_id is not None:
|
|
from fastapi_sso.sso.generic import create_provider, DiscoveryDocument
|
|
|
|
generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
|
|
generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ")
|
|
generic_authorization_endpoint = os.getenv(
|
|
"GENERIC_AUTHORIZATION_ENDPOINT", None
|
|
)
|
|
generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
|
|
generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
|
|
if generic_client_secret is None:
|
|
raise ProxyException(
|
|
message="GENERIC_CLIENT_SECRET not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_CLIENT_SECRET",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if generic_authorization_endpoint is None:
|
|
raise ProxyException(
|
|
message="GENERIC_AUTHORIZATION_ENDPOINT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_AUTHORIZATION_ENDPOINT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if generic_token_endpoint is None:
|
|
raise ProxyException(
|
|
message="GENERIC_TOKEN_ENDPOINT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_TOKEN_ENDPOINT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if generic_userinfo_endpoint is None:
|
|
raise ProxyException(
|
|
message="GENERIC_USERINFO_ENDPOINT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_USERINFO_ENDPOINT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}"
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n"
|
|
)
|
|
discovery = DiscoveryDocument(
|
|
authorization_endpoint=generic_authorization_endpoint,
|
|
token_endpoint=generic_token_endpoint,
|
|
userinfo_endpoint=generic_userinfo_endpoint,
|
|
)
|
|
SSOProvider = create_provider(name="oidc", discovery_document=discovery)
|
|
generic_sso = SSOProvider(
|
|
client_id=generic_client_id,
|
|
client_secret=generic_client_secret,
|
|
redirect_uri=redirect_url,
|
|
allow_insecure_http=True,
|
|
scope=generic_scope,
|
|
)
|
|
with generic_sso:
|
|
# TODO: state should be a random string and added to the user session with cookie
|
|
# or a cryptographicly signed state that we can verify stateless
|
|
# For simplification we are using a static state, this is not perfect but some
|
|
# SSO providers do not allow stateless verification
|
|
redirect_params = {}
|
|
state = os.getenv("GENERIC_CLIENT_STATE", None)
|
|
if state:
|
|
redirect_params["state"] = state
|
|
return await generic_sso.get_login_redirect(**redirect_params) # type: ignore
|
|
elif ui_username is not None:
|
|
# No Google, Microsoft SSO
|
|
# Use UI Credentials set in .env
|
|
from fastapi.responses import HTMLResponse
|
|
|
|
return HTMLResponse(content=html_form, status_code=200)
|
|
else:
|
|
from fastapi.responses import HTMLResponse
|
|
|
|
return HTMLResponse(content=html_form, status_code=200)
|
|
|
|
|
|
@app.get("/fallback/login", tags=["experimental"], include_in_schema=False)
|
|
async def fallback_login(request: Request):
|
|
"""
|
|
Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
|
|
PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/"
|
|
Example:
|
|
"""
|
|
# get url from request
|
|
redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url))
|
|
ui_username = os.getenv("UI_USERNAME")
|
|
if redirect_url.endswith("/"):
|
|
redirect_url += "sso/callback"
|
|
else:
|
|
redirect_url += "/sso/callback"
|
|
|
|
if ui_username is not None:
|
|
# No Google, Microsoft SSO
|
|
# Use UI Credentials set in .env
|
|
from fastapi.responses import HTMLResponse
|
|
|
|
return HTMLResponse(content=html_form, status_code=200)
|
|
else:
|
|
from fastapi.responses import HTMLResponse
|
|
|
|
return HTMLResponse(content=html_form, status_code=200)
|
|
|
|
|
|
@router.post(
|
|
"/login", include_in_schema=False
|
|
) # hidden since this is a helper for UI sso login
|
|
async def login(request: Request):
|
|
try:
|
|
import multipart
|
|
except ImportError:
|
|
subprocess.run(["pip", "install", "python-multipart"])
|
|
global master_key
|
|
if master_key is None:
|
|
raise ProxyException(
|
|
message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.",
|
|
type="auth_error",
|
|
param="master_key",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
form = await request.form()
|
|
username = str(form.get("username"))
|
|
password = str(form.get("password"))
|
|
ui_username = os.getenv("UI_USERNAME", "admin")
|
|
ui_password = os.getenv("UI_PASSWORD", None)
|
|
if ui_password is None:
|
|
ui_password = str(master_key) if master_key is not None else None
|
|
if ui_password is None:
|
|
raise ProxyException(
|
|
message="set Proxy master key to use UI. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.",
|
|
type="auth_error",
|
|
param="UI_PASSWORD",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if secrets.compare_digest(username, ui_username) and secrets.compare_digest(
|
|
password, ui_password
|
|
):
|
|
user_role = "app_owner"
|
|
user_id = username
|
|
key_user_id = user_id
|
|
if (
|
|
os.getenv("PROXY_ADMIN_ID", None) is not None
|
|
and os.environ["PROXY_ADMIN_ID"] == user_id
|
|
) or user_id == "admin":
|
|
# checks if user is admin
|
|
user_role = "app_admin"
|
|
key_user_id = os.getenv("PROXY_ADMIN_ID", "default_user_id")
|
|
|
|
# Admin is Authe'd in - generate key for the UI to access Proxy
|
|
|
|
# ensure this user is set as the proxy admin, in this route there is no sso, we can assume this user is only the admin
|
|
await user_update(
|
|
data=UpdateUserRequest(
|
|
user_id=key_user_id,
|
|
user_role="proxy_admin",
|
|
)
|
|
)
|
|
if os.getenv("DATABASE_URL") is not None:
|
|
response = await generate_key_helper_fn(
|
|
**{"user_role": "proxy_admin", "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"} # type: ignore
|
|
)
|
|
else:
|
|
raise ProxyException(
|
|
message="No Database connected. Set DATABASE_URL in .env. If set, use `--detailed_debug` to debug issue.",
|
|
type="auth_error",
|
|
param="DATABASE_URL",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
key = response["token"] # type: ignore
|
|
litellm_dashboard_ui = os.getenv("PROXY_BASE_URL", "")
|
|
if litellm_dashboard_ui.endswith("/"):
|
|
litellm_dashboard_ui += "ui/"
|
|
else:
|
|
litellm_dashboard_ui += "/ui/"
|
|
import jwt
|
|
|
|
jwt_token = jwt.encode(
|
|
{
|
|
"user_id": user_id,
|
|
"key": key,
|
|
"user_email": user_id,
|
|
"user_role": "app_admin", # this is the path without sso - we can assume only admins will use this
|
|
"login_method": "username_password",
|
|
},
|
|
"secret",
|
|
algorithm="HS256",
|
|
)
|
|
litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token
|
|
return RedirectResponse(url=litellm_dashboard_ui, status_code=303)
|
|
else:
|
|
raise ProxyException(
|
|
message=f"Invalid credentials used to access UI. Passed in username: {username}, passed in password: {password}.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file",
|
|
type="auth_error",
|
|
param="invalid_credentials",
|
|
code=status.HTTP_401_UNAUTHORIZED,
|
|
)
|
|
|
|
|
|
@app.get("/get_image", include_in_schema=False)
|
|
def get_image():
|
|
"""Get logo to show on admin UI"""
|
|
from fastapi.responses import FileResponse
|
|
|
|
# get current_dir
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
default_logo = os.path.join(current_dir, "logo.jpg")
|
|
|
|
logo_path = os.getenv("UI_LOGO_PATH", default_logo)
|
|
verbose_proxy_logger.debug("Reading logo from path: %s", logo_path)
|
|
|
|
# Check if the logo path is an HTTP/HTTPS URL
|
|
if logo_path.startswith(("http://", "https://")):
|
|
# Download the image and cache it
|
|
response = requests.get(logo_path)
|
|
if response.status_code == 200:
|
|
# Save the image to a local file
|
|
cache_path = os.path.join(current_dir, "cached_logo.jpg")
|
|
with open(cache_path, "wb") as f:
|
|
f.write(response.content)
|
|
|
|
# Return the cached image as a FileResponse
|
|
return FileResponse(cache_path, media_type="image/jpeg")
|
|
else:
|
|
# Handle the case when the image cannot be downloaded
|
|
return FileResponse(default_logo, media_type="image/jpeg")
|
|
else:
|
|
# Return the local image file if the logo path is not an HTTP/HTTPS URL
|
|
return FileResponse(logo_path, media_type="image/jpeg")
|
|
|
|
|
|
@app.get("/sso/callback", tags=["experimental"])
|
|
async def auth_callback(request: Request):
|
|
"""Verify login"""
|
|
global general_settings, ui_access_mode
|
|
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
|
|
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
|
|
generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
|
|
# get url from request
|
|
redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url))
|
|
if redirect_url.endswith("/"):
|
|
redirect_url += "sso/callback"
|
|
else:
|
|
redirect_url += "/sso/callback"
|
|
if google_client_id is not None:
|
|
from fastapi_sso.sso.google import GoogleSSO
|
|
|
|
google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
|
|
if google_client_secret is None:
|
|
raise ProxyException(
|
|
message="GOOGLE_CLIENT_SECRET not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GOOGLE_CLIENT_SECRET",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
google_sso = GoogleSSO(
|
|
client_id=google_client_id,
|
|
redirect_uri=redirect_url,
|
|
client_secret=google_client_secret,
|
|
)
|
|
result = await google_sso.verify_and_process(request)
|
|
elif microsoft_client_id is not None:
|
|
from fastapi_sso.sso.microsoft import MicrosoftSSO
|
|
|
|
microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
|
|
microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
|
|
if microsoft_client_secret is None:
|
|
raise ProxyException(
|
|
message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="MICROSOFT_CLIENT_SECRET",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if microsoft_tenant is None:
|
|
raise ProxyException(
|
|
message="MICROSOFT_TENANT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="MICROSOFT_TENANT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
microsoft_sso = MicrosoftSSO(
|
|
client_id=microsoft_client_id,
|
|
client_secret=microsoft_client_secret,
|
|
tenant=microsoft_tenant,
|
|
redirect_uri=redirect_url,
|
|
allow_insecure_http=True,
|
|
)
|
|
result = await microsoft_sso.verify_and_process(request)
|
|
elif generic_client_id is not None:
|
|
# make generic sso provider
|
|
from fastapi_sso.sso.generic import create_provider, DiscoveryDocument, OpenID
|
|
|
|
generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
|
|
generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ")
|
|
generic_authorization_endpoint = os.getenv(
|
|
"GENERIC_AUTHORIZATION_ENDPOINT", None
|
|
)
|
|
generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
|
|
generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
|
|
generic_include_client_id = (
|
|
os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
|
|
)
|
|
if generic_client_secret is None:
|
|
raise ProxyException(
|
|
message="GENERIC_CLIENT_SECRET not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_CLIENT_SECRET",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if generic_authorization_endpoint is None:
|
|
raise ProxyException(
|
|
message="GENERIC_AUTHORIZATION_ENDPOINT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_AUTHORIZATION_ENDPOINT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if generic_token_endpoint is None:
|
|
raise ProxyException(
|
|
message="GENERIC_TOKEN_ENDPOINT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_TOKEN_ENDPOINT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
if generic_userinfo_endpoint is None:
|
|
raise ProxyException(
|
|
message="GENERIC_USERINFO_ENDPOINT not set. Set it in .env file",
|
|
type="auth_error",
|
|
param="GENERIC_USERINFO_ENDPOINT",
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}"
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n"
|
|
)
|
|
|
|
generic_user_id_attribute_name = os.getenv(
|
|
"GENERIC_USER_ID_ATTRIBUTE", "preferred_username"
|
|
)
|
|
generic_user_display_name_attribute_name = os.getenv(
|
|
"GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub"
|
|
)
|
|
generic_user_email_attribute_name = os.getenv(
|
|
"GENERIC_USER_EMAIL_ATTRIBUTE", "email"
|
|
)
|
|
generic_user_role_attribute_name = os.getenv(
|
|
"GENERIC_USER_ROLE_ATTRIBUTE", "role"
|
|
)
|
|
generic_user_first_name_attribute_name = os.getenv(
|
|
"GENERIC_USER_FIRST_NAME_ATTRIBUTE", "first_name"
|
|
)
|
|
generic_user_last_name_attribute_name = os.getenv(
|
|
"GENERIC_USER_LAST_NAME_ATTRIBUTE", "last_name"
|
|
)
|
|
|
|
verbose_proxy_logger.debug(
|
|
f" generic_user_id_attribute_name: {generic_user_id_attribute_name}\n generic_user_email_attribute_name: {generic_user_email_attribute_name}\n generic_user_role_attribute_name: {generic_user_role_attribute_name}"
|
|
)
|
|
|
|
discovery = DiscoveryDocument(
|
|
authorization_endpoint=generic_authorization_endpoint,
|
|
token_endpoint=generic_token_endpoint,
|
|
userinfo_endpoint=generic_userinfo_endpoint,
|
|
)
|
|
|
|
def response_convertor(response, client):
|
|
return OpenID(
|
|
id=response.get(generic_user_id_attribute_name),
|
|
display_name=response.get(generic_user_display_name_attribute_name),
|
|
email=response.get(generic_user_email_attribute_name),
|
|
first_name=response.get(generic_user_first_name_attribute_name),
|
|
last_name=response.get(generic_user_last_name_attribute_name),
|
|
)
|
|
|
|
SSOProvider = create_provider(
|
|
name="oidc",
|
|
discovery_document=discovery,
|
|
response_convertor=response_convertor,
|
|
)
|
|
generic_sso = SSOProvider(
|
|
client_id=generic_client_id,
|
|
client_secret=generic_client_secret,
|
|
redirect_uri=redirect_url,
|
|
allow_insecure_http=True,
|
|
scope=generic_scope,
|
|
)
|
|
verbose_proxy_logger.debug("calling generic_sso.verify_and_process")
|
|
result = await generic_sso.verify_and_process(
|
|
request, params={"include_client_id": generic_include_client_id}
|
|
)
|
|
verbose_proxy_logger.debug("generic result: %s", result)
|
|
|
|
# User is Authe'd in - generate key for the UI to access Proxy
|
|
user_email = getattr(result, "email", None)
|
|
user_id = getattr(result, "id", None)
|
|
|
|
# generic client id
|
|
if generic_client_id is not None:
|
|
user_id = getattr(result, "id", None)
|
|
user_email = getattr(result, "email", None)
|
|
user_role = getattr(result, generic_user_role_attribute_name, None)
|
|
|
|
if user_id is None:
|
|
user_id = getattr(result, "first_name", "") + getattr(result, "last_name", "")
|
|
|
|
user_info = None
|
|
user_id_models: List = []
|
|
|
|
# User might not be already created on first generation of key
|
|
# But if it is, we want their models preferences
|
|
default_ui_key_values = {
|
|
"duration": "2hr",
|
|
"key_max_budget": 0.01,
|
|
"aliases": {},
|
|
"config": {},
|
|
"spend": 0,
|
|
"team_id": "litellm-dashboard",
|
|
}
|
|
user_defined_values = {
|
|
"models": user_id_models,
|
|
"user_id": user_id,
|
|
"user_email": user_email,
|
|
}
|
|
_user_id_from_sso = user_id
|
|
try:
|
|
user_role = None
|
|
if prisma_client is not None:
|
|
user_info = await prisma_client.get_data(user_id=user_id, table_name="user")
|
|
verbose_proxy_logger.debug(
|
|
f"user_info: {user_info}; litellm.default_user_params: {litellm.default_user_params}"
|
|
)
|
|
if user_info is not None:
|
|
user_defined_values = {
|
|
"models": getattr(user_info, "models", []),
|
|
"user_id": getattr(user_info, "user_id", user_id),
|
|
"user_email": getattr(user_info, "user_id", user_email),
|
|
"user_role": getattr(user_info, "user_role", None),
|
|
}
|
|
user_role = getattr(user_info, "user_role", None)
|
|
|
|
## check if user-email in db ##
|
|
user_info = await prisma_client.db.litellm_usertable.find_first(
|
|
where={"user_email": user_email}
|
|
)
|
|
if user_info is not None:
|
|
user_defined_values = {
|
|
"models": getattr(user_info, "models", user_id_models),
|
|
"user_id": user_id,
|
|
"user_email": getattr(user_info, "user_id", user_email),
|
|
"user_role": getattr(user_info, "user_role", None),
|
|
}
|
|
user_role = getattr(user_info, "user_role", None)
|
|
|
|
# update id
|
|
await prisma_client.db.litellm_usertable.update_many(
|
|
where={"user_email": user_email}, data={"user_id": user_id} # type: ignore
|
|
)
|
|
elif litellm.default_user_params is not None and isinstance(
|
|
litellm.default_user_params, dict
|
|
):
|
|
user_defined_values = {
|
|
"models": litellm.default_user_params.get("models", user_id_models),
|
|
"user_id": litellm.default_user_params.get("user_id", user_id),
|
|
"user_email": litellm.default_user_params.get(
|
|
"user_email", user_email
|
|
),
|
|
}
|
|
except Exception as e:
|
|
pass
|
|
|
|
verbose_proxy_logger.info(
|
|
f"user_defined_values for creating ui key: {user_defined_values}"
|
|
)
|
|
response = await generate_key_helper_fn(
|
|
**default_ui_key_values, **user_defined_values # type: ignore
|
|
)
|
|
key = response["token"] # type: ignore
|
|
user_id = response["user_id"] # type: ignore
|
|
|
|
# This should always be true
|
|
# User_id on SSO == user_id in the LiteLLM_VerificationToken Table
|
|
assert user_id == _user_id_from_sso
|
|
litellm_dashboard_ui = "/ui/"
|
|
user_role = user_role or "app_owner"
|
|
if (
|
|
os.getenv("PROXY_ADMIN_ID", None) is not None
|
|
and os.environ["PROXY_ADMIN_ID"] == user_id
|
|
):
|
|
# checks if user is admin
|
|
user_role = "app_admin"
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"user_role: {user_role}; ui_access_mode: {ui_access_mode}"
|
|
)
|
|
## CHECK IF ROLE ALLOWED TO USE PROXY ##
|
|
if ui_access_mode == "admin_only" and "admin" not in user_role:
|
|
verbose_proxy_logger.debug("EXCEPTION RAISED")
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail={
|
|
"error": f"User not allowed to access proxy. User role={user_role}, proxy mode={ui_access_mode}"
|
|
},
|
|
)
|
|
|
|
import jwt
|
|
|
|
jwt_token = jwt.encode(
|
|
{
|
|
"user_id": user_id,
|
|
"key": key,
|
|
"user_email": user_email,
|
|
"user_role": user_role,
|
|
"login_method": "sso",
|
|
},
|
|
"secret",
|
|
algorithm="HS256",
|
|
)
|
|
litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token
|
|
return RedirectResponse(url=litellm_dashboard_ui)
|
|
|
|
|
|
#### BASIC ENDPOINTS ####
|
|
@router.post(
|
|
"/config/update",
|
|
tags=["config.yaml"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def update_config(config_info: ConfigYAML):
|
|
"""
|
|
For Admin UI - allows admin to update config via UI
|
|
|
|
Currently supports modifying General Settings + LiteLLM settings
|
|
"""
|
|
global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client
|
|
try:
|
|
import base64
|
|
|
|
"""
|
|
- Update the ConfigTable DB
|
|
- Run 'add_deployment'
|
|
"""
|
|
if prisma_client is None:
|
|
raise Exception("No DB Connected")
|
|
|
|
updated_settings = config_info.json(exclude_none=True)
|
|
updated_settings = prisma_client.jsonify_object(updated_settings)
|
|
for k, v in updated_settings.items():
|
|
if k == "router_settings":
|
|
await prisma_client.db.litellm_config.upsert(
|
|
where={"param_name": k},
|
|
data={
|
|
"create": {"param_name": k, "param_value": v},
|
|
"update": {"param_value": v},
|
|
},
|
|
)
|
|
|
|
### OLD LOGIC [TODO] MOVE TO DB ###
|
|
|
|
import base64
|
|
|
|
# Load existing config
|
|
config = await proxy_config.get_config()
|
|
verbose_proxy_logger.debug("Loaded config: %s", config)
|
|
|
|
# update the general settings
|
|
if config_info.general_settings is not None:
|
|
config.setdefault("general_settings", {})
|
|
updated_general_settings = config_info.general_settings.dict(
|
|
exclude_none=True
|
|
)
|
|
|
|
_existing_settings = config["general_settings"]
|
|
for k, v in updated_general_settings.items():
|
|
# overwrite existing settings with updated values
|
|
if k == "alert_to_webhook_url":
|
|
# check if slack is already enabled. if not, enable it
|
|
if "slack" not in _existing_settings:
|
|
if "alerting" not in _existing_settings:
|
|
_existing_settings["alerting"] = ["slack"]
|
|
elif isinstance(_existing_settings["alerting"], list):
|
|
_existing_settings["alerting"].append("slack")
|
|
_existing_settings[k] = v
|
|
config["general_settings"] = _existing_settings
|
|
|
|
if config_info.environment_variables is not None:
|
|
config.setdefault("environment_variables", {})
|
|
_updated_environment_variables = config_info.environment_variables
|
|
|
|
# encrypt updated_environment_variables #
|
|
for k, v in _updated_environment_variables.items():
|
|
if isinstance(v, str):
|
|
encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore
|
|
_updated_environment_variables[k] = base64.b64encode(
|
|
encrypted_value
|
|
).decode("utf-8")
|
|
|
|
_existing_env_variables = config["environment_variables"]
|
|
|
|
for k, v in _updated_environment_variables.items():
|
|
# overwrite existing env variables with updated values
|
|
_existing_env_variables[k] = _updated_environment_variables[k]
|
|
|
|
# update the litellm settings
|
|
if config_info.litellm_settings is not None:
|
|
config.setdefault("litellm_settings", {})
|
|
updated_litellm_settings = config_info.litellm_settings
|
|
config["litellm_settings"] = {
|
|
**updated_litellm_settings,
|
|
**config["litellm_settings"],
|
|
}
|
|
|
|
# if litellm.success_callback in updated_litellm_settings and config["litellm_settings"]
|
|
if (
|
|
"success_callback" in updated_litellm_settings
|
|
and "success_callback" in config["litellm_settings"]
|
|
):
|
|
|
|
# check both success callback are lists
|
|
if isinstance(
|
|
config["litellm_settings"]["success_callback"], list
|
|
) and isinstance(updated_litellm_settings["success_callback"], list):
|
|
combined_success_callback = (
|
|
config["litellm_settings"]["success_callback"]
|
|
+ updated_litellm_settings["success_callback"]
|
|
)
|
|
combined_success_callback = list(set(combined_success_callback))
|
|
config["litellm_settings"][
|
|
"success_callback"
|
|
] = combined_success_callback
|
|
|
|
# Save the updated config
|
|
await proxy_config.save_config(new_config=config)
|
|
|
|
await proxy_config.add_deployment(
|
|
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
|
|
return {"message": "Config updated successfully"}
|
|
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,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/get/config/callbacks",
|
|
tags=["config.yaml"],
|
|
include_in_schema=False,
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def get_config():
|
|
"""
|
|
For Admin UI - allows admin to view config via UI
|
|
# return the callbacks and the env variables for the callback
|
|
|
|
"""
|
|
global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key
|
|
try:
|
|
import base64
|
|
|
|
config_data = await proxy_config.get_config()
|
|
_litellm_settings = config_data.get("litellm_settings", {})
|
|
_general_settings = config_data.get("general_settings", {})
|
|
environment_variables = config_data.get("environment_variables", {})
|
|
|
|
# check if "langfuse" in litellm_settings
|
|
_success_callbacks = _litellm_settings.get("success_callback", [])
|
|
_data_to_return = []
|
|
"""
|
|
[
|
|
{
|
|
"name": "langfuse",
|
|
"variables": {
|
|
"LANGFUSE_PUB_KEY": "value",
|
|
"LANGFUSE_SECRET_KEY": "value",
|
|
"LANGFUSE_HOST": "value"
|
|
},
|
|
}
|
|
]
|
|
|
|
"""
|
|
for _callback in _success_callbacks:
|
|
if _callback == "openmeter":
|
|
env_vars = [
|
|
"OPENMETER_API_KEY",
|
|
]
|
|
env_vars_dict = {}
|
|
for _var in env_vars:
|
|
env_variable = environment_variables.get(_var, None)
|
|
if env_variable is None:
|
|
env_vars_dict[_var] = None
|
|
else:
|
|
# decode + decrypt the value
|
|
decoded_b64 = base64.b64decode(env_variable)
|
|
_decrypted_value = decrypt_value(
|
|
value=decoded_b64, master_key=master_key
|
|
)
|
|
env_vars_dict[_var] = _decrypted_value
|
|
|
|
_data_to_return.append({"name": _callback, "variables": env_vars_dict})
|
|
elif _callback == "langfuse":
|
|
_langfuse_vars = [
|
|
"LANGFUSE_PUBLIC_KEY",
|
|
"LANGFUSE_SECRET_KEY",
|
|
"LANGFUSE_HOST",
|
|
]
|
|
_langfuse_env_vars = {}
|
|
for _var in _langfuse_vars:
|
|
env_variable = environment_variables.get(_var, None)
|
|
if env_variable is None:
|
|
_langfuse_env_vars[_var] = None
|
|
else:
|
|
# decode + decrypt the value
|
|
decoded_b64 = base64.b64decode(env_variable)
|
|
_decrypted_value = decrypt_value(
|
|
value=decoded_b64, master_key=master_key
|
|
)
|
|
_langfuse_env_vars[_var] = _decrypted_value
|
|
|
|
_data_to_return.append(
|
|
{"name": _callback, "variables": _langfuse_env_vars}
|
|
)
|
|
|
|
# Check if slack alerting is on
|
|
_alerting = _general_settings.get("alerting", [])
|
|
alerting_data = []
|
|
if "slack" in _alerting:
|
|
_slack_vars = [
|
|
"SLACK_WEBHOOK_URL",
|
|
]
|
|
_slack_env_vars = {}
|
|
for _var in _slack_vars:
|
|
env_variable = environment_variables.get(_var, None)
|
|
if env_variable is None:
|
|
_value = os.getenv("SLACK_WEBHOOK_URL", None)
|
|
_slack_env_vars[_var] = _value
|
|
else:
|
|
# decode + decrypt the value
|
|
decoded_b64 = base64.b64decode(env_variable)
|
|
_decrypted_value = decrypt_value(
|
|
value=decoded_b64, master_key=master_key
|
|
)
|
|
_slack_env_vars[_var] = _decrypted_value
|
|
|
|
_alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types
|
|
_all_alert_types = (
|
|
proxy_logging_obj.slack_alerting_instance._all_possible_alert_types()
|
|
)
|
|
_alerts_to_webhook = (
|
|
proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url
|
|
)
|
|
alerting_data.append(
|
|
{
|
|
"name": "slack",
|
|
"variables": _slack_env_vars,
|
|
"active_alerts": _alerting_types,
|
|
"alerts_to_webhook": _alerts_to_webhook,
|
|
}
|
|
)
|
|
|
|
_router_settings = llm_router.get_settings()
|
|
return {
|
|
"status": "success",
|
|
"callbacks": _data_to_return,
|
|
"alerts": alerting_data,
|
|
"router_settings": _router_settings,
|
|
}
|
|
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,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/config/yaml",
|
|
tags=["config.yaml"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def config_yaml_endpoint(config_info: ConfigYAML):
|
|
"""
|
|
This is a mock endpoint, to show what you can set in config.yaml details in the Swagger UI.
|
|
|
|
Parameters:
|
|
|
|
The config.yaml object has the following attributes:
|
|
- **model_list**: *Optional[List[ModelParams]]* - A list of supported models on the server, along with model-specific configurations. ModelParams includes "model_name" (name of the model), "litellm_params" (litellm-specific parameters for the model), and "model_info" (additional info about the model such as id, mode, cost per token, etc).
|
|
|
|
- **litellm_settings**: *Optional[dict]*: Settings for the litellm module. You can specify multiple properties like "drop_params", "set_verbose", "api_base", "cache".
|
|
|
|
- **general_settings**: *Optional[ConfigGeneralSettings]*: General settings for the server like "completion_model" (default model for chat completion calls), "use_azure_key_vault" (option to load keys from azure key vault), "master_key" (key required for all calls to proxy), and others.
|
|
|
|
Please, refer to each class's description for a better understanding of the specific attributes within them.
|
|
|
|
Note: This is a mock endpoint primarily meant for demonstration purposes, and does not actually provide or change any configurations.
|
|
"""
|
|
return {"hello": "world"}
|
|
|
|
|
|
@router.get(
|
|
"/test",
|
|
tags=["health"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def test_endpoint(request: Request):
|
|
"""
|
|
[DEPRECATED] use `/health/liveliness` instead.
|
|
|
|
A test endpoint that pings the proxy server to check if it's healthy.
|
|
|
|
Parameters:
|
|
request (Request): The incoming request.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the route of the request URL.
|
|
"""
|
|
# ping the proxy server to check if its healthy
|
|
return {"route": request.url.path}
|
|
|
|
|
|
@router.get(
|
|
"/health/services",
|
|
tags=["health"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
include_in_schema=False,
|
|
)
|
|
async def health_services_endpoint(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
service: Literal[
|
|
"slack_budget_alerts", "langfuse", "slack", "openmeter"
|
|
] = fastapi.Query(description="Specify the service being hit."),
|
|
):
|
|
"""
|
|
Hidden endpoint.
|
|
|
|
Used by the UI to let user check if slack alerting is working as expected.
|
|
"""
|
|
try:
|
|
global general_settings, proxy_logging_obj
|
|
|
|
if service is None:
|
|
raise HTTPException(
|
|
status_code=400, detail={"error": "Service must be specified."}
|
|
)
|
|
if service not in ["slack_budget_alerts", "langfuse", "slack", "openmeter"]:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Service must be in list. Service={service}. List={['slack_budget_alerts']}"
|
|
},
|
|
)
|
|
|
|
if service == "openmeter":
|
|
_ = await litellm.acompletion(
|
|
model="openai/litellm-mock-response-model",
|
|
messages=[{"role": "user", "content": "Hey, how's it going?"}],
|
|
user="litellm:/health/services",
|
|
mock_response="This is a mock response",
|
|
)
|
|
return {
|
|
"status": "success",
|
|
"message": "Mock LLM request made - check openmeter.",
|
|
}
|
|
|
|
if service == "langfuse":
|
|
from litellm.integrations.langfuse import LangFuseLogger
|
|
|
|
langfuse_logger = LangFuseLogger()
|
|
langfuse_logger.Langfuse.auth_check()
|
|
_ = litellm.completion(
|
|
model="openai/litellm-mock-response-model",
|
|
messages=[{"role": "user", "content": "Hey, how's it going?"}],
|
|
user="litellm:/health/services",
|
|
mock_response="This is a mock response",
|
|
)
|
|
return {
|
|
"status": "success",
|
|
"message": "Mock LLM request made - check langfuse.",
|
|
}
|
|
|
|
if service == "slack" or service == "slack_budget_alerts":
|
|
if "slack" in general_settings.get("alerting", []):
|
|
# test_message = f"""\n🚨 `ProjectedLimitExceededError` 💸\n\n`Key Alias:` litellm-ui-test-alert \n`Expected Day of Error`: 28th March \n`Current Spend`: $100.00 \n`Projected Spend at end of month`: $1000.00 \n`Soft Limit`: $700"""
|
|
# check if user has opted into unique_alert_webhooks
|
|
if (
|
|
proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url
|
|
is not None
|
|
):
|
|
for (
|
|
alert_type
|
|
) in proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url:
|
|
"""
|
|
"llm_exceptions",
|
|
"llm_too_slow",
|
|
"llm_requests_hanging",
|
|
"budget_alerts",
|
|
"db_exceptions",
|
|
"""
|
|
# only test alert if it's in active alert types
|
|
if (
|
|
proxy_logging_obj.slack_alerting_instance.alert_types
|
|
is not None
|
|
and alert_type
|
|
not in proxy_logging_obj.slack_alerting_instance.alert_types
|
|
):
|
|
continue
|
|
test_message = "default test message"
|
|
if alert_type == "llm_exceptions":
|
|
test_message = f"LLM Exception test alert"
|
|
elif alert_type == "llm_too_slow":
|
|
test_message = f"LLM Too Slow test alert"
|
|
elif alert_type == "llm_requests_hanging":
|
|
test_message = f"LLM Requests Hanging test alert"
|
|
elif alert_type == "budget_alerts":
|
|
test_message = f"Budget Alert test alert"
|
|
elif alert_type == "db_exceptions":
|
|
test_message = f"DB Exception test alert"
|
|
|
|
await proxy_logging_obj.alerting_handler(
|
|
message=test_message, level="Low", alert_type=alert_type
|
|
)
|
|
else:
|
|
await proxy_logging_obj.alerting_handler(
|
|
message="This is a test slack alert message",
|
|
level="Low",
|
|
alert_type="budget_alerts",
|
|
)
|
|
return {
|
|
"status": "success",
|
|
"message": "Mock Slack Alert sent, verify Slack Alert Received on your channel",
|
|
}
|
|
else:
|
|
raise HTTPException(
|
|
status_code=422,
|
|
detail={
|
|
"error": '"{}" not in proxy config: general_settings. Unable to test this.'.format(
|
|
service
|
|
)
|
|
},
|
|
)
|
|
except Exception as e:
|
|
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_500_INTERNAL_SERVER_ERROR),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
raise e
|
|
raise ProxyException(
|
|
message="Authentication Error, " + str(e),
|
|
type="auth_error",
|
|
param=getattr(e, "param", "None"),
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
|
|
@router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)])
|
|
async def health_endpoint(
|
|
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
model: Optional[str] = fastapi.Query(
|
|
None, description="Specify the model name (optional)"
|
|
),
|
|
):
|
|
"""
|
|
🚨 USE `/health/liveliness` to health check the proxy 🚨
|
|
|
|
See more 👉 https://docs.litellm.ai/docs/proxy/health
|
|
|
|
|
|
Check the health of all the endpoints in config.yaml
|
|
|
|
To run health checks in the background, add this to config.yaml:
|
|
```
|
|
general_settings:
|
|
# ... other settings
|
|
background_health_checks: True
|
|
```
|
|
else, the health checks will be run on models when /health is called.
|
|
"""
|
|
global health_check_results, use_background_health_checks, user_model, llm_model_list
|
|
try:
|
|
if llm_model_list is None:
|
|
# if no router set, check if user set a model using litellm --model ollama/llama2
|
|
if user_model is not None:
|
|
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
|
|
model_list=[], cli_model=user_model
|
|
)
|
|
return {
|
|
"healthy_endpoints": healthy_endpoints,
|
|
"unhealthy_endpoints": unhealthy_endpoints,
|
|
"healthy_count": len(healthy_endpoints),
|
|
"unhealthy_count": len(unhealthy_endpoints),
|
|
}
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail={"error": "Model list not initialized"},
|
|
)
|
|
_llm_model_list = copy.deepcopy(llm_model_list)
|
|
### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ###
|
|
if len(user_api_key_dict.models) > 0:
|
|
allowed_model_names = user_api_key_dict.models
|
|
else:
|
|
allowed_model_names = [] #
|
|
if use_background_health_checks:
|
|
return health_check_results
|
|
else:
|
|
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
|
|
_llm_model_list, model
|
|
)
|
|
|
|
return {
|
|
"healthy_endpoints": healthy_endpoints,
|
|
"unhealthy_endpoints": unhealthy_endpoints,
|
|
"healthy_count": len(healthy_endpoints),
|
|
"unhealthy_count": len(unhealthy_endpoints),
|
|
}
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
raise e
|
|
|
|
|
|
db_health_cache = {"status": "unknown", "last_updated": datetime.now()}
|
|
|
|
|
|
def _db_health_readiness_check():
|
|
global db_health_cache, prisma_client
|
|
|
|
# Note - Intentionally don't try/except this so it raises an exception when it fails
|
|
|
|
# if timedelta is less than 2 minutes return DB Status
|
|
time_diff = datetime.now() - db_health_cache["last_updated"]
|
|
if db_health_cache["status"] != "unknown" and time_diff < timedelta(minutes=2):
|
|
return db_health_cache
|
|
prisma_client.health_check()
|
|
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
|
|
return db_health_cache
|
|
|
|
|
|
@router.get(
|
|
"/health/readiness",
|
|
tags=["health"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def health_readiness():
|
|
"""
|
|
Unprotected endpoint for checking if worker can receive requests
|
|
"""
|
|
try:
|
|
# get success callback
|
|
success_callback_names = []
|
|
try:
|
|
# this was returning a JSON of the values in some of the callbacks
|
|
# all we need is the callback name, hence we do str(callback)
|
|
success_callback_names = [str(x) for x in litellm.success_callback]
|
|
except:
|
|
# don't let this block the /health/readiness response, if we can't convert to str -> return litellm.success_callback
|
|
success_callback_names = litellm.success_callback
|
|
|
|
# check Cache
|
|
cache_type = None
|
|
if litellm.cache is not None:
|
|
from litellm.caching import RedisSemanticCache
|
|
|
|
cache_type = litellm.cache.type
|
|
|
|
if isinstance(litellm.cache.cache, RedisSemanticCache):
|
|
# ping the cache
|
|
# TODO: @ishaan-jaff - we should probably not ping the cache on every /health/readiness check
|
|
try:
|
|
index_info = await litellm.cache.cache._index_info()
|
|
except Exception as e:
|
|
index_info = "index does not exist - error: " + str(e)
|
|
cache_type = {"type": cache_type, "index_info": index_info}
|
|
|
|
# check DB
|
|
if prisma_client is not None: # if db passed in, check if it's connected
|
|
db_health_status = _db_health_readiness_check()
|
|
|
|
return {
|
|
"status": "healthy",
|
|
"db": "connected",
|
|
"cache": cache_type,
|
|
"litellm_version": version,
|
|
"success_callbacks": success_callback_names,
|
|
**db_health_status,
|
|
}
|
|
else:
|
|
return {
|
|
"status": "healthy",
|
|
"db": "Not connected",
|
|
"cache": cache_type,
|
|
"litellm_version": version,
|
|
"success_callbacks": success_callback_names,
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=503, detail=f"Service Unhealthy ({str(e)})")
|
|
|
|
|
|
@router.get(
|
|
"/health/liveliness",
|
|
tags=["health"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def health_liveliness():
|
|
"""
|
|
Unprotected endpoint for checking if worker is alive
|
|
"""
|
|
return "I'm alive!"
|
|
|
|
|
|
@router.get(
|
|
"/cache/ping",
|
|
tags=["caching"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def cache_ping():
|
|
"""
|
|
Endpoint for checking if cache can be pinged
|
|
"""
|
|
try:
|
|
litellm_cache_params = {}
|
|
specific_cache_params = {}
|
|
|
|
if litellm.cache is None:
|
|
raise HTTPException(
|
|
status_code=503, detail="Cache not initialized. litellm.cache is None"
|
|
)
|
|
|
|
for k, v in vars(litellm.cache).items():
|
|
try:
|
|
if k == "cache":
|
|
continue
|
|
litellm_cache_params[k] = str(copy.deepcopy(v))
|
|
except Exception:
|
|
litellm_cache_params[k] = "<unable to copy or convert>"
|
|
for k, v in vars(litellm.cache.cache).items():
|
|
try:
|
|
specific_cache_params[k] = str(v)
|
|
except Exception:
|
|
specific_cache_params[k] = "<unable to copy or convert>"
|
|
if litellm.cache.type == "redis":
|
|
# ping the redis cache
|
|
ping_response = await litellm.cache.ping()
|
|
verbose_proxy_logger.debug(
|
|
"/cache/ping: ping_response: " + str(ping_response)
|
|
)
|
|
# making a set cache call
|
|
# add cache does not return anything
|
|
await litellm.cache.async_add_cache(
|
|
result="test_key",
|
|
model="test-model",
|
|
messages=[{"role": "user", "content": "test from litellm"}],
|
|
)
|
|
verbose_proxy_logger.debug("/cache/ping: done with set_cache()")
|
|
return {
|
|
"status": "healthy",
|
|
"cache_type": litellm.cache.type,
|
|
"ping_response": True,
|
|
"set_cache_response": "success",
|
|
"litellm_cache_params": litellm_cache_params,
|
|
"redis_cache_params": specific_cache_params,
|
|
}
|
|
else:
|
|
return {
|
|
"status": "healthy",
|
|
"cache_type": litellm.cache.type,
|
|
"litellm_cache_params": litellm_cache_params,
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail=f"Service Unhealthy ({str(e)}).Cache parameters: {litellm_cache_params}.specific_cache_params: {specific_cache_params}",
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/cache/delete",
|
|
tags=["caching"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def cache_delete(request: Request):
|
|
"""
|
|
Endpoint for deleting a key from the cache. All responses from litellm proxy have `x-litellm-cache-key` in the headers
|
|
|
|
Parameters:
|
|
- **keys**: *Optional[List[str]]* - A list of keys to delete from the cache. Example {"keys": ["key1", "key2"]}
|
|
|
|
```shell
|
|
curl -X POST "http://0.0.0.0:4000/cache/delete" \
|
|
-H "Authorization: Bearer sk-1234" \
|
|
-d '{"keys": ["key1", "key2"]}'
|
|
```
|
|
|
|
"""
|
|
try:
|
|
if litellm.cache is None:
|
|
raise HTTPException(
|
|
status_code=503, detail="Cache not initialized. litellm.cache is None"
|
|
)
|
|
|
|
request_data = await request.json()
|
|
keys = request_data.get("keys", None)
|
|
|
|
if litellm.cache.type == "redis":
|
|
await litellm.cache.delete_cache_keys(keys=keys)
|
|
return {
|
|
"status": "success",
|
|
}
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Cache type {litellm.cache.type} does not support deleting a key. only `redis` is supported",
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Cache Delete Failed({str(e)})",
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/cache/redis/info",
|
|
tags=["caching"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def cache_redis_info():
|
|
"""
|
|
Endpoint for getting /redis/info
|
|
"""
|
|
try:
|
|
if litellm.cache is None:
|
|
raise HTTPException(
|
|
status_code=503, detail="Cache not initialized. litellm.cache is None"
|
|
)
|
|
if litellm.cache.type == "redis":
|
|
client_list = litellm.cache.cache.client_list()
|
|
redis_info = litellm.cache.cache.info()
|
|
num_clients = len(client_list)
|
|
return {
|
|
"num_clients": num_clients,
|
|
"clients": client_list,
|
|
"info": redis_info,
|
|
}
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Cache type {litellm.cache.type} does not support flushing",
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail=f"Service Unhealthy ({str(e)})",
|
|
)
|
|
|
|
|
|
@router.post(
|
|
"/cache/flushall",
|
|
tags=["caching"],
|
|
dependencies=[Depends(user_api_key_auth)],
|
|
)
|
|
async def cache_flushall():
|
|
"""
|
|
A function to flush all items from the cache. (All items will be deleted from the cache with this)
|
|
Raises HTTPException if the cache is not initialized or if the cache type does not support flushing.
|
|
Returns a dictionary with the status of the operation.
|
|
|
|
Usage:
|
|
```
|
|
curl -X POST http://0.0.0.0:4000/cache/flushall -H "Authorization: Bearer sk-1234"
|
|
```
|
|
"""
|
|
try:
|
|
if litellm.cache is None:
|
|
raise HTTPException(
|
|
status_code=503, detail="Cache not initialized. litellm.cache is None"
|
|
)
|
|
if litellm.cache.type == "redis":
|
|
litellm.cache.cache.flushall()
|
|
return {
|
|
"status": "success",
|
|
}
|
|
else:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Cache type {litellm.cache.type} does not support flushing",
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail=f"Service Unhealthy ({str(e)})",
|
|
)
|
|
|
|
|
|
@router.get("/", dependencies=[Depends(user_api_key_auth)])
|
|
async def home(request: Request):
|
|
return "LiteLLM: RUNNING"
|
|
|
|
|
|
@router.get("/routes", dependencies=[Depends(user_api_key_auth)])
|
|
async def get_routes():
|
|
"""
|
|
Get a list of available routes in the FastAPI application.
|
|
"""
|
|
routes = []
|
|
for route in app.routes:
|
|
route_info = {
|
|
"path": getattr(route, "path", None),
|
|
"methods": getattr(route, "methods", None),
|
|
"name": getattr(route, "name", None),
|
|
"endpoint": (
|
|
getattr(route, "endpoint", None).__name__
|
|
if getattr(route, "endpoint", None)
|
|
else None
|
|
),
|
|
}
|
|
routes.append(route_info)
|
|
|
|
return {"routes": routes}
|
|
|
|
|
|
#### TEST ENDPOINTS ####
|
|
@router.get("/token/generate", dependencies=[Depends(user_api_key_auth)])
|
|
async def token_generate():
|
|
"""
|
|
Test endpoint. Admin-only access. Meant for generating admin tokens with specific claims and testing if they work for creating keys, etc.
|
|
"""
|
|
# Initialize AuthJWTSSO with your OpenID Provider configuration
|
|
from fastapi_sso import AuthJWTSSO
|
|
|
|
auth_jwt_sso = AuthJWTSSO(
|
|
issuer=os.getenv("OPENID_BASE_URL"),
|
|
client_id=os.getenv("OPENID_CLIENT_ID"),
|
|
client_secret=os.getenv("OPENID_CLIENT_SECRET"),
|
|
scopes=["litellm_proxy_admin"],
|
|
)
|
|
|
|
token = auth_jwt_sso.create_access_token()
|
|
|
|
return {"token": token}
|
|
|
|
|
|
# @router.post("/update_database", dependencies=[Depends(user_api_key_auth)])
|
|
# async def update_database_endpoint(
|
|
# user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
|
# ):
|
|
# """
|
|
# Test endpoint. DO NOT MERGE IN PROD.
|
|
|
|
# Used for isolating and testing our prisma db update logic in high-traffic.
|
|
# """
|
|
# try:
|
|
# request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"
|
|
# resp = litellm.ModelResponse(
|
|
# id=request_id,
|
|
# choices=[
|
|
# litellm.Choices(
|
|
# finish_reason=None,
|
|
# index=0,
|
|
# message=litellm.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=litellm.Usage(
|
|
# prompt_tokens=210, completion_tokens=200, total_tokens=410
|
|
# ),
|
|
# )
|
|
# await _PROXY_track_cost_callback(
|
|
# kwargs={
|
|
# "model": "chatgpt-v-2",
|
|
# "stream": False,
|
|
# "litellm_params": {
|
|
# "metadata": {
|
|
# "user_api_key": user_api_key_dict.token,
|
|
# "user_api_key_user_id": user_api_key_dict.user_id,
|
|
# }
|
|
# },
|
|
# "response_cost": 0.00002,
|
|
# },
|
|
# completion_response=resp,
|
|
# start_time=datetime.now(),
|
|
# end_time=datetime.now(),
|
|
# )
|
|
# except Exception as e:
|
|
# raise e
|
|
|
|
|
|
def _has_user_setup_sso():
|
|
"""
|
|
Check if the user has set up single sign-on (SSO) by verifying the presence of Microsoft client ID, Google client ID, and UI username environment variables.
|
|
Returns a boolean indicating whether SSO has been set up.
|
|
"""
|
|
microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
|
|
google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
|
|
ui_username = os.getenv("UI_USERNAME", None)
|
|
|
|
sso_setup = (
|
|
(microsoft_client_id is not None)
|
|
or (google_client_id is not None)
|
|
or (ui_username is not None)
|
|
)
|
|
|
|
return sso_setup
|
|
|
|
|
|
@router.on_event("shutdown")
|
|
async def shutdown_event():
|
|
global prisma_client, master_key, user_custom_auth, user_custom_key_generate
|
|
if prisma_client:
|
|
verbose_proxy_logger.debug("Disconnecting from Prisma")
|
|
await prisma_client.disconnect()
|
|
|
|
if litellm.cache is not None:
|
|
await litellm.cache.disconnect()
|
|
|
|
await jwt_handler.close()
|
|
|
|
if db_writer_client is not None:
|
|
await db_writer_client.close()
|
|
## RESET CUSTOM VARIABLES ##
|
|
cleanup_router_config_variables()
|
|
|
|
|
|
def cleanup_router_config_variables():
|
|
global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, prisma_client, custom_db_client
|
|
|
|
# Set all variables to None
|
|
master_key = None
|
|
user_config_file_path = None
|
|
otel_logging = None
|
|
user_custom_auth = None
|
|
user_custom_auth_path = None
|
|
user_custom_key_generate = None
|
|
use_background_health_checks = None
|
|
health_check_interval = None
|
|
prisma_client = None
|
|
custom_db_client = None
|
|
|
|
|
|
app.include_router(router)
|