diff --git a/litellm/__init__.py b/litellm/__init__.py index 3e0a6de0ac..0021daef0b 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.getenv('LITELLM_LOCAL_MODEL_COST_MAP', False)}" - ) 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 94619eb734..70e8163bd5 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 4c88c0be97..42bffdcad0 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 7c8e76062b..c97ebc2b09 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_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 9d4318fe7f..d8ffcf022f 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -1266,9 +1266,7 @@ async def test_user_api_key_auth(prisma_client): pytest.fail(f"This should have failed!. IT's an invalid key") except ProxyException as exc: print(exc.message) - assert ( - exc.message == "Authentication Error, No API Key passed in. api_key is None" - ) + assert exc.message == "Authentication Error, No api key passed in." # Test case: Malformed API Key (missing 'Bearer ' prefix) try: diff --git a/litellm/tests/test_team_config.py b/litellm/tests/test_team_config.py new file mode 100644 index 0000000000..c338307bcf --- /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 8543858eec..2d428c26ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -750,6 +750,10 @@ class Logging: start_time, litellm_call_id, 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]) @@ -770,6 +774,17 @@ class Logging: self.streaming_chunks = [] # for generating complete stream response self.sync_streaming_chunks = [] # for generating complete stream response self.model_call_details = {} + self.dynamic_input_callbacks = [] # callbacks set for just that call + self.dynamic_failure_callbacks = [] # callbacks set for just that call + self.dynamic_success_callbacks = ( + dynamic_success_callbacks or [] + ) # callbacks set for just that call + 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 @@ -873,7 +888,8 @@ class Logging: ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - for callback in litellm.input_callback: + callbacks = litellm.input_callback + self.dynamic_input_callbacks + for callback in callbacks: try: if callback == "supabase": print_verbose("reaches supabase for logging!") @@ -946,43 +962,6 @@ class Logging: if capture_exception: # log this error to sentry for debugging capture_exception(e) - async def async_pre_call( - self, result=None, start_time=None, end_time=None, **kwargs - ): - """ - Â Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. - """ - start_time, end_time, result = self._success_handler_helper_fn( - start_time=start_time, end_time=end_time, result=result - ) - print_verbose(f"Async input callbacks: {litellm._async_input_callback}") - for callback in litellm._async_input_callback: - try: - if isinstance(callback, CustomLogger): # custom logger class - print_verbose(f"Async input callbacks: CustomLogger") - asyncio.create_task( - callback.async_log_input_event( - model=self.model, - messages=self.messages, - kwargs=self.model_call_details, - ) - ) - if callable(callback): # custom logger functions - print_verbose(f"Async success callbacks: async_log_event") - asyncio.create_task( - customLogger.async_log_input_event( - model=self.model, - messages=self.messages, - kwargs=self.model_call_details, - print_verbose=print_verbose, - callback_func=callback, - ) - ) - except: - print_verbose( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {traceback.format_exc()}" - ) - def post_call( self, original_response, input=None, api_key=None, additional_args={} ): @@ -1015,7 +994,9 @@ class Logging: ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - for callback in litellm.input_callback: + + callbacks = litellm.input_callback + self.dynamic_input_callbacks + for callback in callbacks: try: if callback == "lite_debugger": print_verbose("reaches litedebugger for post-call logging!") @@ -1164,8 +1145,8 @@ class Logging: f"Model={self.model} not found in completion cost map." ) self.model_call_details["response_cost"] = None - - for callback in litellm.success_callback: + callbacks = litellm.success_callback + self.dynamic_success_callbacks + for callback in callbacks: try: if callback == "lite_debugger": print_verbose("reaches lite_debugger for logging!") @@ -1235,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, @@ -1303,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, @@ -1466,7 +1452,10 @@ class Logging: ) self.model_call_details["response_cost"] = None - for callback in litellm._async_success_callback: + callbacks = ( + litellm._async_success_callback + self.dynamic_async_success_callbacks + ) + for callback in callbacks: try: if callback == "cache" and litellm.cache is not None: # set_cache once complete streaming response is built @@ -1968,6 +1957,26 @@ def client(original_function): # Pop the async items from failure_callback in reverse order to avoid index issues for index in reversed(removed_async_items): litellm.failure_callback.pop(index) + ### DYNAMIC CALLBACKS ### + dynamic_success_callbacks = [] + dynamic_async_success_callbacks = [] + if kwargs.get("success_callback", None) is not None and isinstance( + kwargs["success_callback"], list + ): + removed_async_items = [] + for index, callback in enumerate(kwargs["success_callback"]): + if ( + inspect.iscoroutinefunction(callback) + or callback == "dynamodb" + or callback == "s3" + ): + dynamic_async_success_callbacks.append(callback) + removed_async_items.append(index) + # 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.pop("success_callback") + if add_breadcrumb: add_breadcrumb( category="litellm.llm_call", @@ -2029,6 +2038,10 @@ def client(original_function): function_id=function_id, call_type=call_type, 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 = {} @@ -2040,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 @@ -2110,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 @@ -2345,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