kenken999/litellmlope
0
1import sys, os, platform, time, copy, re, asyncio, inspect2import threading, ast3import shutil, random, traceback, requests4from datetime import datetime, timedelta, timezone5from typing import Optional, List6import secrets, subprocess7import hashlib, uuid8import warnings9import importlib10 11messages: list = []12sys.path.insert(13 0, os.path.abspath("../..")14) # Adds the parent directory to the system path - for litellm local dev15 16sample = """17 from openai import OpenAI18 import json19 20 base_url = "https://ka1kuk-litellm.hf.space"21 api_key = "hf_xxxx"22 23 client = OpenAI(base_url=base_url, api_key=api_key)24 25 messages = [{"role": "user", "content": "What's the capital of France?"}]26 27 response = client.chat.completions.create(28 model="huggingface/mistralai/Mixtral-8x7B-Instruct-v0.1",29 response_format={ "type": "json_object" },30 messages=messages,31 stream=False,32 )33 34 print(response.choices[0].message.content)35"""36 37description = f"Proxy Server to call 100+ LLMs in the OpenAI format\n\nSample with openai library:\n\n{sample}" 38 39try:40 import fastapi41 import backoff42 import yaml43 import orjson44 import logging45except ImportError as e:46 raise ImportError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`")47 48import litellm49from litellm.proxy.utils import (50 PrismaClient,51 DBClient,52 get_instance_fn,53 ProxyLogging,54 _cache_user_row,55 send_email,56)57from litellm.proxy.secret_managers.google_kms import load_google_kms58import pydantic59from litellm.proxy._types import *60from litellm.caching import DualCache61from litellm.proxy.health_check import perform_health_check62from litellm._logging import verbose_router_logger, verbose_proxy_logger63 64litellm.suppress_debug_info = True65from fastapi import (66 FastAPI,67 Request,68 HTTPException,69 status,70 Depends,71 BackgroundTasks,72 Header,73 Response,74)75from fastapi.routing import APIRouter76from fastapi.security import OAuth2PasswordBearer77from fastapi.encoders import jsonable_encoder78from fastapi.responses import StreamingResponse, FileResponse, ORJSONResponse79from fastapi.middleware.cors import CORSMiddleware80from fastapi.security.api_key import APIKeyHeader81import json82import logging83from typing import Union84 85app = FastAPI(86 docs_url="/",87 title="LiteLLM API",88 description= description,89)90router = APIRouter()91origins = ["*"]92 93app.add_middleware(94 CORSMiddleware,95 allow_origins=origins,96 allow_credentials=True,97 allow_methods=["*"],98 allow_headers=["*"],99)100 101 102from typing import Dict103 104api_key_header = APIKeyHeader(name="Authorization", auto_error=False)105user_api_base = None106user_model = None107user_debug = False108user_max_tokens = None109user_request_timeout = None110user_temperature = None111user_telemetry = True112user_config = None113user_headers = None114user_config_file_path = f"config_{int(time.time())}.yaml"115local_logging = True # writes logs to a local api_log.json file for debugging116experimental = False117#### GLOBAL VARIABLES ####118llm_router: Optional[litellm.Router] = None119llm_model_list: Optional[list] = None120general_settings: dict = {}121log_file = "api_log.json"122worker_config = None123master_key = None124otel_logging = False125prisma_client: Optional[PrismaClient] = None126custom_db_client: Optional[DBClient] = None127user_api_key_cache = DualCache()128user_custom_auth = None129use_background_health_checks = None130use_queue = False131health_check_interval = None132health_check_results = {}133queue: List = []134### INITIALIZE GLOBAL LOGGING OBJECT ###135proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)136### REDIS QUEUE ###137async_result = None138celery_app_conn = None139celery_fn = None # Redis Queue for handling requests140### logger ###141 142 143def usage_telemetry(144 feature: str,145): # helps us know if people are using this feature. Set `litellm --telemetry False` to your cli call to turn this off146 if user_telemetry:147 data = {"feature": feature} # "local_proxy_server"148 threading.Thread(149 target=litellm.utils.litellm_telemetry, args=(data,), daemon=True150 ).start()151 152 153def _get_bearer_token(api_key: str):154 assert api_key.startswith("Bearer ") # ensure Bearer token passed in155 api_key = api_key.replace("Bearer ", "") # extract the token156 return api_key157 158 159def _get_pydantic_json_dict(pydantic_obj: BaseModel) -> dict:160 try:161 return pydantic_obj.model_dump() # type: ignore162 except:163 # if using pydantic v1164 return pydantic_obj.dict()165 166 167async def user_api_key_auth(168 request: Request, api_key: str = fastapi.Security(api_key_header)169) -> UserAPIKeyAuth:170 global master_key, prisma_client, llm_model_list, user_custom_auth, custom_db_client171 try:172 if isinstance(api_key, str):173 api_key = _get_bearer_token(api_key=api_key)174 ### USER-DEFINED AUTH FUNCTION ###175 if user_custom_auth is not None:176 response = await user_custom_auth(request=request, api_key=api_key)177 return UserAPIKeyAuth.model_validate(response)178 ### LITELLM-DEFINED AUTH FUNCTION ###179 if master_key is None:180 if isinstance(api_key, str):181 return UserAPIKeyAuth(api_key=api_key)182 else:183 return UserAPIKeyAuth()184 185 route: str = request.url.path186 if route == "/user/auth":187 if general_settings.get("allow_user_auth", False) == True:188 return UserAPIKeyAuth()189 else:190 raise HTTPException(191 status_code=status.HTTP_403_FORBIDDEN,192 detail="'allow_user_auth' not set or set to False",193 )194 195 if api_key is None: # only require api key if master key is set196 raise Exception(f"No api key passed in.")197 198 # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead199 is_master_key_valid = secrets.compare_digest(api_key, master_key)200 if is_master_key_valid:201 return UserAPIKeyAuth(api_key=master_key)202 203 if route.startswith("/config/") and not is_master_key_valid:204 raise Exception(f"Only admin can modify config")205 206 if (207 (route.startswith("/key/") or route.startswith("/user/"))208 or route.startswith("/model/")209 and not is_master_key_valid210 and general_settings.get("allow_user_auth", False) != True211 ):212 raise Exception(213 f"If master key is set, only master key can be used to generate, delete, update or get info for new keys/users"214 )215 216 if (217 prisma_client is None and custom_db_client is None218 ): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error219 raise Exception("No connected db.")220 221 ## check for cache hit (In-Memory Cache)222 valid_token = user_api_key_cache.get_cache(key=api_key)223 verbose_proxy_logger.debug(f"valid_token from cache: {valid_token}")224 if valid_token is None:225 ## check db226 verbose_proxy_logger.debug(f"api key: {api_key}")227 if prisma_client is not None:228 valid_token = await prisma_client.get_data(229 token=api_key,230 )231 232 expires = datetime.utcnow().replace(tzinfo=timezone.utc)233 elif custom_db_client is not None:234 valid_token = await custom_db_client.get_data(235 key=api_key, table_name="key"236 )237 # Token exists, now check expiration.238 if valid_token.expires is not None:239 expiry_time = datetime.fromisoformat(valid_token.expires)240 if expiry_time >= datetime.utcnow():241 # Token exists and is not expired.242 return response243 else:244 # Token exists but is expired.245 raise HTTPException(246 status_code=status.HTTP_403_FORBIDDEN,247 detail="expired user key",248 )249 verbose_proxy_logger.debug(f"valid token from prisma: {valid_token}")250 user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)251 elif valid_token is not None:252 verbose_proxy_logger.debug(f"API Key Cache Hit!")253 if valid_token:254 litellm.model_alias_map = valid_token.aliases255 config = valid_token.config256 if config != {}:257 model_list = config.get("model_list", [])258 llm_model_list = model_list259 verbose_proxy_logger.debug(260 f"\n new llm router model list {llm_model_list}"261 )262 if (263 len(valid_token.models) == 0264 ): # assume an empty model list means all models are allowed to be called265 pass266 else:267 try:268 data = await request.json()269 except json.JSONDecodeError:270 data = {} # Provide a default value, such as an empty dictionary271 model = data.get("model", None)272 if model in litellm.model_alias_map:273 model = litellm.model_alias_map[model]274 if model and model not in valid_token.models:275 raise Exception(f"Token not allowed to access model")276 api_key = valid_token.token277 valid_token_dict = _get_pydantic_json_dict(valid_token)278 valid_token_dict.pop("token", None)279 """280 asyncio create task to update the user api key cache with the user db table as well281 282 This makes the user row data accessible to pre-api call hooks.283 """284 if prisma_client is not None:285 asyncio.create_task(286 _cache_user_row(287 user_id=valid_token.user_id,288 cache=user_api_key_cache,289 db=prisma_client,290 )291 )292 elif custom_db_client is not None:293 asyncio.create_task(294 _cache_user_row(295 user_id=valid_token.user_id,296 cache=user_api_key_cache,297 db=custom_db_client,298 )299 )300 return UserAPIKeyAuth(api_key=api_key, **valid_token_dict)301 else:302 raise Exception(f"Invalid token")303 except Exception as e:304 # verbose_proxy_logger.debug(f"An exception occurred - {traceback.format_exc()}")305 traceback.print_exc()306 if isinstance(e, HTTPException):307 raise e308 else:309 raise HTTPException(310 status_code=status.HTTP_401_UNAUTHORIZED,311 detail="invalid user key",312 )313 314 315def prisma_setup(database_url: Optional[str]):316 global prisma_client, proxy_logging_obj, user_api_key_cache317 318 if database_url is not None:319 try:320 prisma_client = PrismaClient(321 database_url=database_url, proxy_logging_obj=proxy_logging_obj322 )323 except Exception as e:324 raise e325 326 327def load_from_azure_key_vault(use_azure_key_vault: bool = False):328 if use_azure_key_vault is False:329 return330 331 try:332 from azure.keyvault.secrets import SecretClient333 from azure.identity import ClientSecretCredential334 335 # Set your Azure Key Vault URI336 KVUri = os.getenv("AZURE_KEY_VAULT_URI", None)337 338 # Set your Azure AD application/client ID, client secret, and tenant ID339 client_id = os.getenv("AZURE_CLIENT_ID", None)340 client_secret = os.getenv("AZURE_CLIENT_SECRET", None)341 tenant_id = os.getenv("AZURE_TENANT_ID", None)342 343 if (344 KVUri is not None345 and client_id is not None346 and client_secret is not None347 and tenant_id is not None348 ):349 # Initialize the ClientSecretCredential350 credential = ClientSecretCredential(351 client_id=client_id, client_secret=client_secret, tenant_id=tenant_id352 )353 354 # Create the SecretClient using the credential355 client = SecretClient(vault_url=KVUri, credential=credential)356 357 litellm.secret_manager_client = client358 litellm._key_management_system = KeyManagementSystem.AZURE_KEY_VAULT359 else:360 raise Exception(361 f"Missing KVUri or client_id or client_secret or tenant_id from environment"362 )363 except Exception as e:364 verbose_proxy_logger.debug(365 "Error when loading keys from Azure Key Vault. Ensure you run `pip install azure-identity azure-keyvault-secrets`"366 )367 368 369def cost_tracking():370 global prisma_client, custom_db_client371 if prisma_client is not None or custom_db_client is not None:372 if isinstance(litellm.success_callback, list):373 verbose_proxy_logger.debug("setting litellm success callback to track cost")374 if (track_cost_callback) not in litellm.success_callback: # type: ignore375 litellm.success_callback.append(track_cost_callback) # type: ignore376 377 378async def track_cost_callback(379 kwargs, # kwargs to completion380 completion_response: litellm.ModelResponse, # response from completion381 start_time=None,382 end_time=None, # start/end time for completion383):384 global prisma_client, custom_db_client385 try:386 # check if it has collected an entire stream response387 verbose_proxy_logger.debug(388 f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}"389 )390 if "complete_streaming_response" in kwargs:391 # for tracking streaming cost we pass the "messages" and the output_text to litellm.completion_cost392 completion_response = kwargs["complete_streaming_response"]393 response_cost = litellm.completion_cost(394 completion_response=completion_response395 )396 verbose_proxy_logger.debug(f"streaming response_cost {response_cost}")397 user_api_key = kwargs["litellm_params"]["metadata"].get(398 "user_api_key", None399 )400 user_id = kwargs["litellm_params"]["metadata"].get(401 "user_api_key_user_id", None402 )403 if user_api_key and (404 prisma_client is not None or custom_db_client is not None405 ):406 await update_database(token=user_api_key, response_cost=response_cost)407 elif kwargs["stream"] == False: # for non streaming responses408 response_cost = litellm.completion_cost(409 completion_response=completion_response410 )411 user_api_key = kwargs["litellm_params"]["metadata"].get(412 "user_api_key", None413 )414 user_id = kwargs["litellm_params"]["metadata"].get(415 "user_api_key_user_id", None416 )417 if user_api_key and (418 prisma_client is not None or custom_db_client is not None419 ):420 await update_database(421 token=user_api_key, response_cost=response_cost, user_id=user_id422 )423 except Exception as e:424 verbose_proxy_logger.debug(f"error in tracking cost callback - {str(e)}")425 426 427async def update_database(token, response_cost, user_id=None):428 try:429 verbose_proxy_logger.debug(430 f"Enters prisma db call, token: {token}; user_id: {user_id}"431 )432 433 ### UPDATE USER SPEND ###434 async def _update_user_db():435 if user_id is None:436 return437 if prisma_client is not None:438 existing_spend_obj = await prisma_client.get_data(user_id=user_id)439 elif custom_db_client is not None:440 existing_spend_obj = await custom_db_client.get_data(441 key=user_id, table_name="user"442 )443 if existing_spend_obj is None:444 existing_spend = 0445 else:446 existing_spend = existing_spend_obj.spend447 448 # Calculate the new cost by adding the existing cost and response_cost449 new_spend = existing_spend + response_cost450 451 verbose_proxy_logger.debug(f"new cost: {new_spend}")452 # Update the cost column for the given user id453 if prisma_client is not None:454 await prisma_client.update_data(455 user_id=user_id, data={"spend": new_spend}456 )457 elif custom_db_client is not None:458 await custom_db_client.update_data(459 key=user_id, value={"spend": new_spend}, table_name="user"460 )461 462 ### UPDATE KEY SPEND ###463 async def _update_key_db():464 if prisma_client is not None:465 # Fetch the existing cost for the given token466 existing_spend_obj = await prisma_client.get_data(token=token)467 verbose_proxy_logger.debug(f"existing spend: {existing_spend_obj}")468 if existing_spend_obj is None:469 existing_spend = 0470 else:471 existing_spend = existing_spend_obj.spend472 # Calculate the new cost by adding the existing cost and response_cost473 new_spend = existing_spend + response_cost474 475 verbose_proxy_logger.debug(f"new cost: {new_spend}")476 # Update the cost column for the given token477 await prisma_client.update_data(token=token, data={"spend": new_spend})478 elif custom_db_client is not None:479 # Fetch the existing cost for the given token480 existing_spend_obj = await custom_db_client.get_data(481 key=token, table_name="key"482 )483 verbose_proxy_logger.debug(f"existing spend: {existing_spend_obj}")484 if existing_spend_obj is None:485 existing_spend = 0486 else:487 existing_spend = existing_spend_obj.spend488 # Calculate the new cost by adding the existing cost and response_cost489 new_spend = existing_spend + response_cost490 491 verbose_proxy_logger.debug(f"new cost: {new_spend}")492 # Update the cost column for the given token493 await custom_db_client.update_data(494 key=token, value={"spend": new_spend}, table_name="key"495 )496 497 tasks = []498 tasks.append(_update_user_db())499 tasks.append(_update_key_db())500 await asyncio.gather(*tasks)501 except Exception as e:502 verbose_proxy_logger.debug(503 f"Error updating Prisma database: {traceback.format_exc()}"504 )505 pass506 507 508def run_ollama_serve():509 try:510 command = ["ollama", "serve"]511 512 with open(os.devnull, "w") as devnull:513 process = subprocess.Popen(command, stdout=devnull, stderr=devnull)514 except Exception as e:515 verbose_proxy_logger.debug(516 f"""517 LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve`518 """519 )520 521 522async def _run_background_health_check():523 """524 Periodically run health checks in the background on the endpoints.525 526 Update health_check_results, based on this.527 """528 global health_check_results, llm_model_list, health_check_interval529 while True:530 healthy_endpoints, unhealthy_endpoints = await perform_health_check(531 model_list=llm_model_list532 )533 534 # Update the global variable with the health check results535 health_check_results["healthy_endpoints"] = healthy_endpoints536 health_check_results["unhealthy_endpoints"] = unhealthy_endpoints537 health_check_results["healthy_count"] = len(healthy_endpoints)538 health_check_results["unhealthy_count"] = len(unhealthy_endpoints)539 540 await asyncio.sleep(health_check_interval)541 542 543class ProxyConfig:544 """545 Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.546 """547 548 def __init__(self) -> None:549 pass550 551 def is_yaml(self, config_file_path: str) -> bool:552 if not os.path.isfile(config_file_path):553 return False554 555 _, file_extension = os.path.splitext(config_file_path)556 return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml"557 558 async def get_config(self, config_file_path: Optional[str] = None) -> dict:559 global prisma_client, user_config_file_path560 561 file_path = config_file_path or user_config_file_path562 if config_file_path is not None:563 user_config_file_path = config_file_path564 # Load existing config565 ## Yaml566 if os.path.exists(f"{file_path}"):567 with open(f"{file_path}", "r") as config_file:568 config = yaml.safe_load(config_file)569 else:570 config = {571 "model_list": [],572 "general_settings": {},573 "router_settings": {},574 "litellm_settings": {},575 }576 577 ## DB578 if (579 prisma_client is not None580 and litellm.get_secret("SAVE_CONFIG_TO_DB", False) == True581 ):582 prisma_setup(database_url=None) # in case it's not been connected yet583 _tasks = []584 keys = [585 "model_list",586 "general_settings",587 "router_settings",588 "litellm_settings",589 ]590 for k in keys:591 response = prisma_client.get_generic_data(592 key="param_name", value=k, table_name="config"593 )594 _tasks.append(response)595 596 responses = await asyncio.gather(*_tasks)597 598 return config599 600 async def save_config(self, new_config: dict):601 global prisma_client, llm_router, user_config_file_path, llm_model_list, general_settings602 # Load existing config603 backup_config = await self.get_config()604 605 # Save the updated config606 ## YAML607 with open(f"{user_config_file_path}", "w") as config_file:608 yaml.dump(new_config, config_file, default_flow_style=False)609 610 # update Router - verifies if this is a valid config611 try:612 (613 llm_router,614 llm_model_list,615 general_settings,616 ) = await proxy_config.load_config(617 router=llm_router, config_file_path=user_config_file_path618 )619 except Exception as e:620 traceback.print_exc()621 # Revert to old config instead622 with open(f"{user_config_file_path}", "w") as config_file:623 yaml.dump(backup_config, config_file, default_flow_style=False)624 raise HTTPException(status_code=400, detail="Invalid config passed in")625 626 ## DB - writes valid config to db627 """628 - Do not write restricted params like 'api_key' to the database629 - if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`)630 """631 if (632 prisma_client is not None633 and litellm.get_secret("SAVE_CONFIG_TO_DB", default_value=False) == True634 ):635 ### KEY REMOVAL ###636 models = new_config.get("model_list", [])637 for m in models:638 if m.get("litellm_params", {}).get("api_key", None) is not None:639 # pop the key640 api_key = m["litellm_params"].pop("api_key")641 # store in local env642 key_name = f"LITELLM_MODEL_KEY_{uuid.uuid4()}"643 os.environ[key_name] = api_key644 # save the key name (not the value)645 m["litellm_params"]["api_key"] = f"os.environ/{key_name}"646 await prisma_client.insert_data(data=new_config, table_name="config")647 648 async def load_config(649 self, router: Optional[litellm.Router], config_file_path: str650 ):651 """652 Load config values into proxy global state653 """654 global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, use_background_health_checks, health_check_interval, use_queue, custom_db_client655 656 # Load existing config657 config = await self.get_config(config_file_path=config_file_path)658 ## PRINT YAML FOR CONFIRMING IT WORKS659 printed_yaml = copy.deepcopy(config)660 printed_yaml.pop("environment_variables", None)661 662 verbose_proxy_logger.debug(663 f"Loaded config YAML (api_key and environment_variables are not shown):\n{json.dumps(printed_yaml, indent=2)}"664 )665 666 ## ENVIRONMENT VARIABLES667 environment_variables = config.get("environment_variables", None)668 if environment_variables:669 for key, value in environment_variables.items():670 os.environ[key] = value671 672 ## LITELLM MODULE SETTINGS (e.g. litellm.drop_params=True,..)673 litellm_settings = config.get("litellm_settings", None)674 if litellm_settings is None:675 litellm_settings = {}676 if litellm_settings:677 # ANSI escape code for blue text678 blue_color_code = "\033[94m"679 reset_color_code = "\033[0m"680 for key, value in litellm_settings.items():681 if key == "cache":682 print(f"{blue_color_code}\nSetting Cache on Proxy") # noqa683 from litellm.caching import Cache684 685 cache_params = {}686 if "cache_params" in litellm_settings:687 cache_params_in_config = litellm_settings["cache_params"]688 # overwrie cache_params with cache_params_in_config689 cache_params.update(cache_params_in_config)690 691 cache_type = cache_params.get("type", "redis")692 693 verbose_proxy_logger.debug(f"passed cache type={cache_type}")694 695 if cache_type == "redis":696 cache_host = litellm.get_secret("REDIS_HOST", None)697 cache_port = litellm.get_secret("REDIS_PORT", None)698 cache_password = litellm.get_secret("REDIS_PASSWORD", None)699 700 cache_params.update(701 {702 "type": cache_type,703 "host": cache_host,704 "port": cache_port,705 "password": cache_password,706 }707 )708 # Assuming cache_type, cache_host, cache_port, and cache_password are strings709 print( # noqa710 f"{blue_color_code}Cache Type:{reset_color_code} {cache_type}"711 ) # noqa712 print( # noqa713 f"{blue_color_code}Cache Host:{reset_color_code} {cache_host}"714 ) # noqa715 print( # noqa716 f"{blue_color_code}Cache Port:{reset_color_code} {cache_port}"717 ) # noqa718 print( # noqa719 f"{blue_color_code}Cache Password:{reset_color_code} {cache_password}"720 )721 print() # noqa722 723 # users can pass os.environ/ variables on the proxy - we should read them from the env724 for key, value in cache_params.items():725 if type(value) is str and value.startswith("os.environ/"):726 cache_params[key] = litellm.get_secret(value)727 728 ## 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 variables729 litellm.cache = Cache(**cache_params)730 print( # noqa731 f"{blue_color_code}Set Cache on LiteLLM Proxy: {vars(litellm.cache.cache)}{reset_color_code}"732 )733 elif key == "callbacks":734 litellm.callbacks = [735 get_instance_fn(value=value, config_file_path=config_file_path)736 ]737 verbose_proxy_logger.debug(738 f"{blue_color_code} Initialized Callbacks - {litellm.callbacks} {reset_color_code}"739 )740 elif key == "post_call_rules":741 litellm.post_call_rules = [742 get_instance_fn(value=value, config_file_path=config_file_path)743 ]744 verbose_proxy_logger.debug(745 f"litellm.post_call_rules: {litellm.post_call_rules}"746 )747 elif key == "success_callback":748 litellm.success_callback = []749 750 # intialize success callbacks751 for callback in value:752 # user passed custom_callbacks.async_on_succes_logger. They need us to import a function753 if "." in callback:754 litellm.success_callback.append(755 get_instance_fn(value=callback)756 )757 # these are litellm callbacks - "langfuse", "sentry", "wandb"758 else:759 litellm.success_callback.append(callback)760 verbose_proxy_logger.debug(761 f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}"762 )763 elif key == "failure_callback":764 litellm.failure_callback = []765 766 # intialize success callbacks767 for callback in value:768 # user passed custom_callbacks.async_on_succes_logger. They need us to import a function769 if "." in callback:770 litellm.failure_callback.append(771 get_instance_fn(value=callback)772 )773 # these are litellm callbacks - "langfuse", "sentry", "wandb"774 else:775 litellm.failure_callback.append(callback)776 verbose_proxy_logger.debug(777 f"{blue_color_code} Initialized Success Callbacks - {litellm.failure_callback} {reset_color_code}"778 )779 elif key == "cache_params":780 # this is set in the cache branch781 # see usage here: https://docs.litellm.ai/docs/proxy/caching782 pass783 else:784 setattr(litellm, key, value)785 786 ## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging787 general_settings = config.get("general_settings", {})788 if general_settings is None:789 general_settings = {}790 if general_settings:791 ### LOAD SECRET MANAGER ###792 key_management_system = general_settings.get("key_management_system", None)793 if key_management_system is not None:794 if key_management_system == KeyManagementSystem.AZURE_KEY_VAULT.value:795 ### LOAD FROM AZURE KEY VAULT ###796 load_from_azure_key_vault(use_azure_key_vault=True)797 elif key_management_system == KeyManagementSystem.GOOGLE_KMS.value:798 ### LOAD FROM GOOGLE KMS ###799 load_google_kms(use_google_kms=True)800 else:801 raise ValueError("Invalid Key Management System selected")802 ### [DEPRECATED] LOAD FROM GOOGLE KMS ### old way of loading from google kms803 use_google_kms = general_settings.get("use_google_kms", False)804 load_google_kms(use_google_kms=use_google_kms)805 ### [DEPRECATED] LOAD FROM AZURE KEY VAULT ### old way of loading from azure secret manager806 use_azure_key_vault = general_settings.get("use_azure_key_vault", False)807 load_from_azure_key_vault(use_azure_key_vault=use_azure_key_vault)808 ### ALERTING ###809 proxy_logging_obj.update_values(810 alerting=general_settings.get("alerting", None),811 alerting_threshold=general_settings.get("alerting_threshold", 600),812 )813 ### CONNECT TO DATABASE ###814 database_url = general_settings.get("database_url", None)815 if database_url and database_url.startswith("os.environ/"):816 verbose_proxy_logger.debug(f"GOING INTO LITELLM.GET_SECRET!")817 database_url = litellm.get_secret(database_url)818 verbose_proxy_logger.debug(f"RETRIEVED DB URL: {database_url}")819 ### MASTER KEY ###820 master_key = general_settings.get(821 "master_key", litellm.get_secret("LITELLM_MASTER_KEY", None)822 )823 if master_key and master_key.startswith("os.environ/"):824 master_key = litellm.get_secret(master_key)825 ### CUSTOM API KEY AUTH ###826 ## pass filepath827 custom_auth = general_settings.get("custom_auth", None)828 if custom_auth is not None:829 user_custom_auth = get_instance_fn(830 value=custom_auth, config_file_path=config_file_path831 )832 ## dynamodb833 database_type = general_settings.get("database_type", None)834 if database_type is not None and (835 database_type == "dynamo_db" or database_type == "dynamodb"836 ):837 database_args = general_settings.get("database_args", None)838 custom_db_client = DBClient(839 custom_db_args=database_args, custom_db_type=database_type840 )841 ## COST TRACKING ##842 cost_tracking()843 ### BACKGROUND HEALTH CHECKS ###844 # Enable background health checks845 use_background_health_checks = general_settings.get(846 "background_health_checks", False847 )848 health_check_interval = general_settings.get("health_check_interval", 300)849 850 router_params: dict = {851 "num_retries": 3,852 "cache_responses": litellm.cache853 != None, # cache if user passed in cache values854 }855 ## MODEL LIST856 model_list = config.get("model_list", None)857 if model_list:858 router_params["model_list"] = model_list859 print( # noqa860 f"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m"861 ) # noqa862 for model in model_list:863 ### LOAD FROM os.environ/ ###864 for k, v in model["litellm_params"].items():865 if isinstance(v, str) and v.startswith("os.environ/"):866 model["litellm_params"][k] = litellm.get_secret(v)867 print(f"\033[32m {model.get('model_name', '')}\033[0m") # noqa868 litellm_model_name = model["litellm_params"]["model"]869 litellm_model_api_base = model["litellm_params"].get("api_base", None)870 if "ollama" in litellm_model_name and litellm_model_api_base is None:871 run_ollama_serve()872 873 ## ROUTER SETTINGS (e.g. routing_strategy, ...)874 router_settings = config.get("router_settings", None)875 if router_settings and isinstance(router_settings, dict):876 arg_spec = inspect.getfullargspec(litellm.Router)877 # model list already set878 exclude_args = {879 "self",880 "model_list",881 }882 883 available_args = [x for x in arg_spec.args if x not in exclude_args]884 885 for k, v in router_settings.items():886 if k in available_args:887 router_params[k] = v888 889 router = litellm.Router(**router_params) # type:ignore890 return router, model_list, general_settings891 892 893proxy_config = ProxyConfig()894 895 896async def generate_key_helper_fn(897 duration: Optional[str],898 models: list,899 aliases: dict,900 config: dict,901 spend: float,902 max_budget: Optional[float] = None,903 token: Optional[str] = None,904 user_id: Optional[str] = None,905 user_email: Optional[str] = None,906 max_parallel_requests: Optional[int] = None,907 metadata: Optional[dict] = {},908):909 global prisma_client, custom_db_client910 911 if prisma_client is None and custom_db_client is None:912 raise Exception(913 f"Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys "914 )915 916 if token is None:917 token = f"sk-{secrets.token_urlsafe(16)}"918 919 def _duration_in_seconds(duration: str):920 match = re.match(r"(\d+)([smhd]?)", duration)921 if not match:922 raise ValueError("Invalid duration format")923 924 value, unit = match.groups()925 value = int(value)926 927 if unit == "s":928 return value929 elif unit == "m":930 return value * 60931 elif unit == "h":932 return value * 3600933 elif unit == "d":934 return value * 86400935 else:936 raise ValueError("Unsupported duration unit")937 938 if duration is None: # allow tokens that never expire939 expires = None940 else:941 duration_s = _duration_in_seconds(duration=duration)942 expires = datetime.utcnow() + timedelta(seconds=duration_s)943 944 aliases_json = json.dumps(aliases)945 config_json = json.dumps(config)946 metadata_json = json.dumps(metadata)947 user_id = user_id or str(uuid.uuid4())948 try:949 # Create a new verification token (you may want to enhance this logic based on your needs)950 user_data = {951 "max_budget": max_budget,952 "user_email": user_email,953 "user_id": user_id,954 "spend": spend,955 }956 key_data = {957 "token": token,958 "expires": expires,959 "models": models,960 "aliases": aliases_json,961 "config": config_json,962 "spend": spend,963 "user_id": user_id,964 "max_parallel_requests": max_parallel_requests,965 "metadata": metadata_json,966 }967 if prisma_client is not None:968 verification_token_data = dict(key_data)969 verification_token_data.update(user_data)970 verbose_proxy_logger.debug("PrismaClient: Before Insert Data")971 await prisma_client.insert_data(data=verification_token_data)972 elif custom_db_client is not None:973 ## CREATE USER (If necessary)974 await custom_db_client.insert_data(value=user_data, table_name="user")975 ## CREATE KEY976 await custom_db_client.insert_data(value=key_data, table_name="key")977 except Exception as e:978 traceback.print_exc()979 raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)980 return {981 "token": token,982 "expires": expires,983 "user_id": user_id,984 "max_budget": max_budget,985 }986 987 988async def delete_verification_token(tokens: List):989 global prisma_client990 try:991 if prisma_client:992 # Assuming 'db' is your Prisma Client instance993 deleted_tokens = await prisma_client.delete_data(tokens=tokens)994 else:995 raise Exception996 except Exception as e:997 traceback.print_exc()998 raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)999 return deleted_tokens1000 1001 1002def save_worker_config(**data):1003 import json1004 1005 os.environ["WORKER_CONFIG"] = json.dumps(data)1006 1007 1008async def initialize(1009 model=None,1010 alias=None,1011 api_base=None,1012 api_version=None,1013 debug=False,1014 detailed_debug=False,1015 temperature=None,1016 max_tokens=None,1017 request_timeout=600,1018 max_budget=None,1019 telemetry=False,1020 drop_params=True,1021 add_function_to_prompt=True,1022 headers=None,1023 save=False,1024 use_queue=False,1025 config=None,1026):1027 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_client1028 user_model = model1029 user_debug = debug1030 if debug == True: # this needs to be first, so users can see Router init debugg1031 from litellm._logging import verbose_router_logger, verbose_proxy_logger1032 import logging1033 1034 # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS1035 1036 verbose_router_logger.setLevel(level=logging.INFO) # set router logs to info1037 verbose_proxy_logger.setLevel(level=logging.INFO) # set proxy logs to info1038 if detailed_debug == True:1039 from litellm._logging import verbose_router_logger, verbose_proxy_logger1040 import logging1041 1042 verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to info1043 verbose_proxy_logger.setLevel(level=logging.DEBUG) # set proxy logs to debug1044 litellm.set_verbose = True1045 elif debug == False and detailed_debug == False:1046 # users can control proxy debugging using env variable = 'LITELLM_LOG'1047 litellm_log_setting = os.environ.get("LITELLM_LOG", "")1048 if litellm_log_setting != None:1049 if litellm_log_setting.upper() == "INFO":1050 from litellm._logging import verbose_router_logger, verbose_proxy_logger1051 import logging1052 1053 # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS1054 1055 verbose_router_logger.setLevel(1056 level=logging.INFO1057 ) # set router logs to info1058 verbose_proxy_logger.setLevel(1059 level=logging.INFO1060 ) # set proxy logs to info1061 elif litellm_log_setting.upper() == "DEBUG":1062 from litellm._logging import verbose_router_logger, verbose_proxy_logger1063 import logging1064 1065 verbose_router_logger.setLevel(1066 level=logging.DEBUG1067 ) # set router logs to info1068 verbose_proxy_logger.setLevel(1069 level=logging.DEBUG1070 ) # set proxy logs to debug1071 litellm.set_verbose = True1072 1073 dynamic_config = {"general": {}, user_model: {}}1074 if config:1075 (1076 llm_router,1077 llm_model_list,1078 general_settings,1079 ) = await proxy_config.load_config(router=llm_router, config_file_path=config)1080 if headers: # model-specific param1081 user_headers = headers1082 dynamic_config[user_model]["headers"] = headers1083 if api_base: # model-specific param1084 user_api_base = api_base1085 dynamic_config[user_model]["api_base"] = api_base1086 if api_version:1087 os.environ[1088 "AZURE_API_VERSION"1089 ] = api_version # set this for azure - litellm can read this from the env1090 if max_tokens: # model-specific param1091 user_max_tokens = max_tokens1092 dynamic_config[user_model]["max_tokens"] = max_tokens1093 if temperature: # model-specific param1094 user_temperature = temperature1095 dynamic_config[user_model]["temperature"] = temperature1096 if request_timeout:1097 user_request_timeout = request_timeout1098 dynamic_config[user_model]["request_timeout"] = request_timeout1099 if alias: # model-specific param1100 dynamic_config[user_model]["alias"] = alias1101 if drop_params == True: # litellm-specific param1102 litellm.drop_params = True1103 dynamic_config["general"]["drop_params"] = True1104 if add_function_to_prompt == True: # litellm-specific param1105 litellm.add_function_to_prompt = True1106 dynamic_config["general"]["add_function_to_prompt"] = True1107 if max_budget: # litellm-specific param1108 litellm.max_budget = max_budget1109 dynamic_config["general"]["max_budget"] = max_budget1110 if experimental:1111 pass1112 user_telemetry = telemetry1113 usage_telemetry(feature="local_proxy_server")1114 1115 1116# for streaming1117def data_generator(response):1118 verbose_proxy_logger.debug("inside generator")1119 for chunk in response:1120 verbose_proxy_logger.debug(f"returned chunk: {chunk}")1121 try:1122 yield f"data: {json.dumps(chunk.dict())}\n\n"1123 except:1124 yield f"data: {json.dumps(chunk)}\n\n"1125 1126 1127async def async_data_generator(response, user_api_key_dict):1128 verbose_proxy_logger.debug("inside generator")1129 try:1130 start_time = time.time()1131 async for chunk in response:1132 verbose_proxy_logger.debug(f"returned chunk: {chunk}")1133 try:1134 yield f"data: {json.dumps(chunk.dict())}\n\n"1135 except Exception as e:1136 yield f"data: {str(e)}\n\n"1137 1138 ### ALERTING ###1139 end_time = time.time()1140 asyncio.create_task(1141 proxy_logging_obj.response_taking_too_long(1142 start_time=start_time, end_time=end_time, type="slow_response"1143 )1144 )1145 1146 # Streaming is done, yield the [DONE] chunk1147 done_message = "[DONE]"1148 yield f"data: {done_message}\n\n"1149 except Exception as e:1150 yield f"data: {str(e)}\n\n"1151 1152 1153def get_litellm_model_info(model: dict = {}):1154 model_info = model.get("model_info", {})1155 model_to_lookup = model.get("litellm_params", {}).get("model", None)1156 try:1157 if "azure" in model_to_lookup:1158 model_to_lookup = model_info.get("base_model", None)1159 litellm_model_info = litellm.get_model_info(model_to_lookup)1160 return litellm_model_info1161 except:1162 # this should not block returning on /model/info1163 # if litellm does not have info on the model it should return {}1164 return {}1165 1166 1167def parse_cache_control(cache_control):1168 cache_dict = {}1169 directives = cache_control.split(", ")1170 1171 for directive in directives:1172 if "=" in directive:1173 key, value = directive.split("=")1174 cache_dict[key] = value1175 else:1176 cache_dict[directive] = True1177 1178 return cache_dict1179 1180 1181@router.on_event("startup")1182async def startup_event():1183 global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings1184 import json1185 1186 ### LOAD MASTER KEY ###1187 # check if master key set in environment - load from there1188 master_key = litellm.get_secret("LITELLM_MASTER_KEY", None)1189 # check if DATABASE_URL in environment - load from there1190 if prisma_client is None:1191 prisma_setup(database_url=os.getenv("DATABASE_URL"))1192 1193 ### LOAD CONFIG ###1194 worker_config = litellm.get_secret("WORKER_CONFIG")1195 verbose_proxy_logger.debug(f"worker_config: {worker_config}")1196 # check if it's a valid file path1197 if os.path.isfile(worker_config):1198 if proxy_config.is_yaml(config_file_path=worker_config):1199 (1200 llm_router,