phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +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

33
litellm/router.py
View file

@ -1,5 +1,6 @@
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
@ -24,14 +25,13 @@ 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',
@ -63,10 +63,8 @@ class Router:
# 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
data["caching"] = self.cache_responses
# call via litellm.completion() # call via litellm.completion()
return litellm.completion(**{**data, **kwargs}) return litellm.completion(**{**data, "messages": messages, "caching": self.cache_responses, **kwargs})
async def acompletion(self, async def acompletion(self,
model: str, model: str,
@ -77,9 +75,7 @@ class Router:
# 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, def text_completion(self,
model: str, model: str,
@ -94,10 +90,8 @@ class Router:
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
data["caching"] = self.cache_responses
# call via litellm.completion() # call via litellm.completion()
return litellm.text_completion(**{**data, **kwargs}) return litellm.text_completion(**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs})
def embedding(self, def embedding(self,
model: str, model: str,
@ -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
data["caching"] = self.cache_responses
# call via litellm.embedding() # call via litellm.embedding()
return litellm.embedding(**{**data, **kwargs}) return litellm.embedding(**{**data, "input": input, "caching": self.cache_responses, **kwargs})
async def aembedding(self, async def aembedding(self,
model: str, model: str,
@ -122,12 +114,11 @@ 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
@ -148,8 +139,8 @@ class Router:
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.
""" """