phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

avoid overwriting litellm_params

Browse source
This commit is contained in:
mc-marcocheng 2023-10-27 15:30:34 +08:00
parent 895cb5d0f9
commit f43d59fff8
1 changed files with 60 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

129
litellm/router.py
View file

@ -1,16 +1,17 @@
from typing import Union, List, Dict, Optional
from datetime import datetime from datetime import datetime
from typing import Dict, List, Optional, Union
import litellm import litellm
class Router: class Router:
""" """
Example usage: Example usage:
from litellm import Router from litellm import Router
model_list = [{ model_list = [{
"model_name": "gpt-3.5-turbo", # openai model name "model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call "litellm_params": { # params for litellm completion/embedding call
"model": "azure/<your-deployment-name>", "model": "azure/<your-deployment-name>",
"api_key": <your-api-key>, "api_key": <your-api-key>,
"api_version": <your-api-version>, "api_version": <your-api-version>,
"api_base": <your-api-base> "api_base": <your-api-base>
@ -23,16 +24,15 @@ class Router:
""" """
model_names: List = [] model_names: List = []
cache_responses: bool = False cache_responses: bool = False
def __init__(self, def __init__(self,
model_list: Optional[list]=None, model_list: Optional[list] = None,
redis_host: Optional[str] = None, redis_host: Optional[str] = None,
redis_port: Optional[int] = None, redis_port: Optional[int] = None,
redis_password: Optional[str] = None, redis_password: Optional[str] = None,
cache_responses: bool = False) -> None: cache_responses: bool = False) -> None:
if model_list: if model_list:
self.model_list = model_list self.set_model_list(model_list)
self.model_names = [m["model_name"] for m in model_list] if redis_host is not None and redis_port is not None and redis_password is not None:
if redis_host is not None and redis_port is not None and redis_password is not None:
cache_config = { cache_config = {
'type': 'redis', 'type': 'redis',
'host': redis_host, 'host': redis_host,
@ -45,61 +45,55 @@ class Router:
} }
self.cache = litellm.Cache(cache_config) # use Redis for tracking load balancing self.cache = litellm.Cache(cache_config) # use Redis for tracking load balancing
if cache_responses: if cache_responses:
litellm.cache = litellm.Cache(**cache_config) # use Redis for caching completion requests litellm.cache = litellm.Cache(**cache_config) # use Redis for caching completion requests
self.cache_responses = cache_responses self.cache_responses = cache_responses
litellm.success_callback = [self.deployment_callback] litellm.success_callback = [self.deployment_callback]
def completion(self, def completion(self,
model: str, model: str,
messages: List[Dict[str, str]], messages: List[Dict[str, str]],
is_retry: Optional[bool] = False, is_retry: Optional[bool] = False,
is_fallback: Optional[bool] = False, is_fallback: Optional[bool] = False,
**kwargs): **kwargs):
""" """
Example usage: Example usage:
response = router.completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}] response = router.completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}]
""" """
# pick the one that is available (lowest TPM/RPM) # pick the one that is available (lowest TPM/RPM)
deployment = self.get_available_deployment(model=model, messages=messages) deployment = self.get_available_deployment(model=model, messages=messages)
data = deployment["litellm_params"] data = deployment["litellm_params"]
data["messages"] = messages # call via litellm.completion()
data["caching"] = self.cache_responses return litellm.completion(**{**data, "messages": messages, "caching": self.cache_responses, **kwargs})
# call via litellm.completion()
return litellm.completion(**{**data, **kwargs})
async def acompletion(self, async def acompletion(self,
model: str, model: str,
messages: List[Dict[str, str]], messages: List[Dict[str, str]],
is_retry: Optional[bool] = False, is_retry: Optional[bool] = False,
is_fallback: Optional[bool] = False, is_fallback: Optional[bool] = False,
**kwargs): **kwargs):
# pick the one that is available (lowest TPM/RPM) # pick the one that is available (lowest TPM/RPM)
deployment = self.get_available_deployment(model=model, messages=messages) deployment = self.get_available_deployment(model=model, messages=messages)
data = deployment["litellm_params"] data = deployment["litellm_params"]
data["messages"] = messages return await litellm.acompletion(**{**data, "messages": messages, "caching": self.cache_responses, **kwargs})
data["caching"] = self.cache_responses
return await litellm.acompletion(**{**data, **kwargs}) def text_completion(self,
model: str,
def text_completion(self, prompt: str,
model: str,
prompt: str,
is_retry: Optional[bool] = False, is_retry: Optional[bool] = False,
is_fallback: Optional[bool] = False, is_fallback: Optional[bool] = False,
is_async: Optional[bool] = False, is_async: Optional[bool] = False,
**kwargs): **kwargs):
messages=[{"role": "user", "content": prompt}] messages=[{"role": "user", "content": prompt}]
# pick the one that is available (lowest TPM/RPM) # pick the one that is available (lowest TPM/RPM)
deployment = self.get_available_deployment(model=model, messages=messages) deployment = self.get_available_deployment(model=model, messages=messages)
data = deployment["litellm_params"] data = deployment["litellm_params"]
data["prompt"] = prompt # call via litellm.completion()
data["caching"] = self.cache_responses return litellm.text_completion(**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs})
# call via litellm.completion()
return litellm.text_completion(**{**data, **kwargs})
def embedding(self, def embedding(self,
model: str, model: str,
input: Union[str, List], input: Union[str, List],
is_async: Optional[bool] = False, is_async: Optional[bool] = False,
@ -108,10 +102,8 @@ class Router:
deployment = self.get_available_deployment(model=model, input=input) deployment = self.get_available_deployment(model=model, input=input)
data = deployment["litellm_params"] data = deployment["litellm_params"]
data["input"] = input # call via litellm.embedding()
data["caching"] = self.cache_responses return litellm.embedding(**{**data, "input": input, "caching": self.cache_responses, **kwargs})
# call via litellm.embedding()
return litellm.embedding(**{**data, **kwargs})
async def aembedding(self, async def aembedding(self,
model: str, model: str,
@ -122,14 +114,13 @@ class Router:
deployment = self.get_available_deployment(model=model, input=input) deployment = self.get_available_deployment(model=model, input=input)
data = deployment["litellm_params"] data = deployment["litellm_params"]
data["input"] = input return await litellm.aembedding(**{**data, "input": input, "caching": self.cache_responses, **kwargs})
data["caching"] = self.cache_responses
return await litellm.aembedding(**{**data, **kwargs})
def set_model_list(self, model_list: list): def set_model_list(self, model_list: list):
self.model_list = model_list self.model_list = model_list
self.model_names = [m["model_name"] for m in model_list]
def get_model_names(self): def get_model_names(self):
return self.model_names return self.model_names
def deployment_callback( def deployment_callback(
@ -146,21 +137,21 @@ class Router:
total_tokens = completion_response['usage']['total_tokens'] total_tokens = completion_response['usage']['total_tokens']
self._set_deployment_usage(model_name, total_tokens) self._set_deployment_usage(model_name, total_tokens)
def get_available_deployment(self, def get_available_deployment(self,
model: str, model: str,
messages: Optional[List[Dict[str, str]]]=None, messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]]=None): input: Optional[Union[str, List]] = None):
""" """
Returns a deployment with the lowest TPM/RPM usage. Returns a deployment with the lowest TPM/RPM usage.
""" """
# get list of potential deployments # get list of potential deployments
potential_deployments = [] potential_deployments = []
for item in self.model_list: for item in self.model_list:
if item["model_name"] == model: if item["model_name"] == model:
potential_deployments.append(item) potential_deployments.append(item)
# set first model as current model # set first model as current model
deployment = potential_deployments[0] deployment = potential_deployments[0]
# get model tpm, rpm limits # get model tpm, rpm limits
@ -170,7 +161,7 @@ class Router:
# get deployment current usage # get deployment current usage
current_tpm, current_rpm = self._get_deployment_usage(deployment_name=deployment["litellm_params"]["model"]) current_tpm, current_rpm = self._get_deployment_usage(deployment_name=deployment["litellm_params"]["model"])
# get encoding # get encoding
if messages: if messages:
token_count = litellm.token_counter(model=deployment["model_name"], messages=messages) token_count = litellm.token_counter(model=deployment["model_name"], messages=messages)
elif input: elif input:
@ -179,9 +170,9 @@ class Router:
else: else:
input_text = input input_text = input
token_count = litellm.token_counter(model=deployment["model_name"], text=input_text) token_count = litellm.token_counter(model=deployment["model_name"], text=input_text)
# if at model limit, return lowest used # if at model limit, return lowest used
if current_tpm + token_count > tpm or current_rpm + 1 >= rpm: if current_tpm + token_count > tpm or current_rpm + 1 >= rpm:
# ----------------------- # -----------------------
# Find lowest used model # Find lowest used model
# ---------------------- # ----------------------
@ -194,17 +185,17 @@ class Router:
if item_tpm == 0: if item_tpm == 0:
return item return item
elif item_tpm + token_count > item["tpm"] or item_rpm + 1 >= item["rpm"]: elif item_tpm + token_count > item["tpm"] or item_rpm + 1 >= item["rpm"]:
continue continue
elif item_tpm < lowest_tpm: elif item_tpm < lowest_tpm:
lowest_tpm = item_tpm lowest_tpm = item_tpm
deployment = item deployment = item
# if none, raise exception # if none, raise exception
if deployment is None: if deployment is None:
raise ValueError(f"No models available.") raise ValueError(f"No models available.")
# return model # return model
return deployment return deployment
def _get_deployment_usage( def _get_deployment_usage(
@ -224,24 +215,24 @@ class Router:
tpm = self.cache.get_cache(tpm_key) tpm = self.cache.get_cache(tpm_key)
rpm = self.cache.get_cache(rpm_key) rpm = self.cache.get_cache(rpm_key)
if tpm is None: if tpm is None:
tpm = 0 tpm = 0
if rpm is None: if rpm is None:
rpm = 0 rpm = 0
return int(tpm), int(rpm) return int(tpm), int(rpm)
def increment(self, key: str, increment_value: int): def increment(self, key: str, increment_value: int):
# get value # get value
cached_value = self.cache.get_cache(key) cached_value = self.cache.get_cache(key)
# update value # update value
try: try:
cached_value = cached_value + increment_value cached_value = cached_value + increment_value
except: except:
cached_value = increment_value cached_value = increment_value
# save updated value # save updated value
self.cache.add_cache(result=cached_value, cache_key=key) self.cache.add_cache(result=cached_value, cache_key=key)
def _set_deployment_usage( def _set_deployment_usage(
self, self,
model_name: str, model_name: str,