diff --git a/litellm/__init__.py b/litellm/__init__.py index 6d49c6ae96d..4e16319c64e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -146,6 +146,7 @@ suppress_debug_info = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None +default_team_settings: Optional[List] = None #### RELIABILITY #### request_timeout: Optional[float] = 6000 num_retries: Optional[int] = None # per model endpoint @@ -165,9 +166,6 @@ _key_management_system: Optional[KeyManagementSystem] = None def get_model_cost_map(url: str): - verbose_logger.debug( - f"os.getenv('LITELLM_LOCAL_MODEL_COST_MAP', False): {os.environ['LITELLM_LOCAL_MODEL_COST_MAP']}" - ) if ( os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", False) == True or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", False) == "True" @@ -175,7 +173,6 @@ def get_model_cost_map(url: str): import importlib.resources import json - verbose_logger.debug("RUNS LOCALLY") with importlib.resources.open_text( "litellm", "model_prices_and_context_window_backup.json" ) as f: diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py index 94619eb7341..70e8163bd5e 100644 --- a/litellm/integrations/langfuse.py +++ b/litellm/integrations/langfuse.py @@ -14,7 +14,7 @@ import litellm class LangFuseLogger: # Class variables or attributes - def __init__(self): + def __init__(self, langfuse_public_key=None, langfuse_secret=None): try: from langfuse import Langfuse except Exception as e: @@ -22,8 +22,8 @@ class LangFuseLogger: f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\033[0m" ) # Instance variables - self.secret_key = os.getenv("LANGFUSE_SECRET_KEY") - self.public_key = os.getenv("LANGFUSE_PUBLIC_KEY") + self.secret_key = langfuse_secret or os.getenv("LANGFUSE_SECRET_KEY") + self.public_key = langfuse_public_key or os.getenv("LANGFUSE_PUBLIC_KEY") self.langfuse_host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com") self.langfuse_release = os.getenv("LANGFUSE_RELEASE") self.langfuse_debug = os.getenv("LANGFUSE_DEBUG") diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4c88c0be97c..42bffdcad00 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -321,6 +321,7 @@ class LiteLLM_VerificationToken(LiteLLMBase): aliases: Dict = {} config: Dict = {} user_id: Optional[str] = None + team_id: Optional[str] = None max_parallel_requests: Optional[int] = None metadata: Dict = {} tpm_limit: Optional[int] = None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7c8e76062ba..c97ebc2b09a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1024,6 +1024,24 @@ class ProxyConfig: m["litellm_params"]["api_key"] = f"os.environ/{key_name}" await prisma_client.insert_data(data=new_config, table_name="config") + async def load_team_config(self, team_id: str): + """ + - for a given team id + - return the relevant completion() call params + """ + all_teams_config = litellm.default_team_settings + team_config: dict = {} + if all_teams_config is None: + return team_config + for team in all_teams_config: + 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 + async def load_config( self, router: Optional[litellm.Router], config_file_path: str ): @@ -2040,6 +2058,21 @@ async def chat_completion( 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) + 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: @@ -2215,6 +2248,21 @@ async def embeddings( data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id 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 = ( [m["model_name"] for m in llm_model_list] if llm_model_list is not None @@ -2361,6 +2409,21 @@ async def image_generation( data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id 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 = ( [m["model_name"] for m in llm_model_list] if llm_model_list is not None diff --git a/litellm/tests/test_team_config.py b/litellm/tests/test_team_config.py new file mode 100644 index 00000000000..c338307bcf5 --- /dev/null +++ b/litellm/tests/test_team_config.py @@ -0,0 +1,36 @@ +#### What this tests #### +# This tests if setting team_config actually works +import sys, os +import traceback +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm.proxy.proxy_server import ProxyConfig + + +@pytest.mark.asyncio +async def test_team_config(): + litellm.default_team_settings = [ + { + "team_id": "my-special-team", + "success_callback": ["langfuse"], + "langfuse_public_key": "os.environ/LANGFUSE_PUB_KEY_2", + "langfuse_secret": "os.environ/LANGFUSE_PRIVATE_KEY_2", + } + ] + proxyconfig = ProxyConfig() + + team_config = await proxyconfig.load_team_config(team_id="my-special-team") + assert len(team_config) > 0 + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hey, how's it going?"}], + } + team_config.pop("team_id") + response = litellm.completion(**{**data, **team_config}) + + print(f"response: {response}") diff --git a/litellm/utils.py b/litellm/utils.py index adadbb30836..03dfc40f22c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -752,6 +752,8 @@ class Logging: function_id, dynamic_success_callbacks=None, dynamic_async_success_callbacks=None, + langfuse_public_key=None, + langfuse_secret=None, ): if call_type not in [item.value for item in CallTypes]: allowed_values = ", ".join([item.value for item in CallTypes]) @@ -780,6 +782,9 @@ class Logging: self.dynamic_async_success_callbacks = ( dynamic_async_success_callbacks or [] ) # callbacks set for just that call + ## DYNAMIC LANGFUSE KEYS ## + self.langfuse_public_key = langfuse_public_key + self.langfuse_secret = langfuse_secret def update_environment_variables( self, model, user, optional_params, litellm_params, **additional_params @@ -1211,7 +1216,9 @@ class Logging: if "complete_streaming_response" not in kwargs: break else: - print_verbose("reaches langfuse for streaming logging!") + print_verbose( + "reaches langsmith for streaming logging!" + ) result = kwargs["complete_streaming_response"] langsmithLogger.log_event( kwargs=self.model_call_details, @@ -1279,7 +1286,10 @@ class Logging: print_verbose("reaches langfuse for streaming logging!") result = kwargs["complete_streaming_response"] if langFuseLogger is None: - langFuseLogger = LangFuseLogger() + langFuseLogger = LangFuseLogger( + langfuse_public_key=self.langfuse_public_key, + langfuse_secret=self.langfuse_secret, + ) langFuseLogger.log_event( kwargs=kwargs, response_obj=result, @@ -1965,7 +1975,7 @@ def client(original_function): # Pop the async items from success_callback in reverse order to avoid index issues for index in reversed(removed_async_items): kwargs["success_callback"].pop(index) - dynamic_success_callbacks = kwargs["success_callback"] + dynamic_success_callbacks = kwargs.pop("success_callback") if add_breadcrumb: add_breadcrumb( @@ -2030,6 +2040,8 @@ def client(original_function): start_time=start_time, dynamic_success_callbacks=dynamic_success_callbacks, dynamic_async_success_callbacks=dynamic_async_success_callbacks, + langfuse_public_key=kwargs.pop("langfuse_public_key", None), + langfuse_secret=kwargs.pop("langfuse_secret", None), ) ## check if metadata is passed in litellm_params = {} @@ -2041,7 +2053,7 @@ def client(original_function): optional_params={}, litellm_params=litellm_params, ) - return logging_obj + return logging_obj, kwargs except Exception as e: import logging @@ -2111,7 +2123,7 @@ def client(original_function): try: if logging_obj is None: - logging_obj = function_setup(start_time, *args, **kwargs) + logging_obj, kwargs = function_setup(start_time, *args, **kwargs) kwargs["litellm_logging_obj"] = logging_obj # CHECK FOR 'os.environ/' in kwargs @@ -2346,7 +2358,7 @@ def client(original_function): try: if logging_obj is None: - logging_obj = function_setup(start_time, *args, **kwargs) + logging_obj, kwargs = function_setup(start_time, *args, **kwargs) kwargs["litellm_logging_obj"] = logging_obj # [OPTIONAL] CHECK BUDGET