diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8978e0d1a0..2a84048e0b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,10 +1,10 @@ repos: - repo: https://github.com/psf/black - rev: stable + rev: 24.2.0 hooks: - id: black - repo: https://github.com/pycqa/flake8 - rev: 3.8.4 # The version of flake8 to use + rev: 7.0.0 # The version of flake8 to use hooks: - id: flake8 exclude: ^litellm/tests/|^litellm/proxy/proxy_cli.py|^litellm/integrations/|^litellm/proxy/tests/ diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index e66627cccf..d898ed8c7f 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -49,9 +49,9 @@ class HuggingfaceConfig: details: Optional[bool] = True # enables returning logprobs + best of max_new_tokens: Optional[int] = None repetition_penalty: Optional[float] = None - return_full_text: Optional[ - bool - ] = False # by default don't return the input as part of the output + return_full_text: Optional[bool] = ( + False # by default don't return the input as part of the output + ) seed: Optional[int] = None temperature: Optional[float] = None top_k: Optional[int] = None @@ -188,9 +188,9 @@ class Huggingface(BaseLLM): "content-type": "application/json", } if api_key and headers is None: - default_headers[ - "Authorization" - ] = f"Bearer {api_key}" # Huggingface Inference Endpoint default is to accept bearer tokens + default_headers["Authorization"] = ( + f"Bearer {api_key}" # Huggingface Inference Endpoint default is to accept bearer tokens + ) headers = default_headers elif headers: headers = headers diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7a6e400c89..3f8f1944ed 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -154,6 +154,7 @@ class GenerateKeyRequest(GenerateRequestBase): duration: Optional[str] = None aliases: Optional[dict] = {} config: Optional[dict] = {} + permissions: Optional[dict] = {} class GenerateKeyResponse(GenerateKeyRequest): @@ -166,7 +167,7 @@ class GenerateKeyResponse(GenerateKeyRequest): def set_model_info(cls, values): if values.get("token") is not None: values.update({"key": values.get("token")}) - dict_fields = ["metadata", "aliases", "config"] + dict_fields = ["metadata", "aliases", "config", "permissions"] for field in dict_fields: value = values.get(field) if value is not None and isinstance(value, str): @@ -381,6 +382,7 @@ class LiteLLM_VerificationToken(LiteLLMBase): budget_duration: Optional[str] = None budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] + permissions: Dict = {} class UserAPIKeyAuth( diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index a9461d9225..206fee7777 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -282,7 +282,12 @@ class DynamoDBWrapper(CustomDB): new_response = {} for k, v in response.items(): # handle json string if ( - (k == "aliases" or k == "config" or k == "metadata") + ( + k == "aliases" + or k == "config" + or k == "metadata" + or k == "permissions" + ) and v is not None and isinstance(v, str) ): diff --git a/litellm/proxy/hooks/cache_control_check.py b/litellm/proxy/hooks/cache_control_check.py index c50c4ec1fc..3160fe97ad 100644 --- a/litellm/proxy/hooks/cache_control_check.py +++ b/litellm/proxy/hooks/cache_control_check.py @@ -30,18 +30,20 @@ class _PROXY_CacheControlCheck(CustomLogger): self.print_verbose(f"Inside Cache Control Check Pre-Call Hook") allowed_cache_controls = user_api_key_dict.allowed_cache_controls - if (allowed_cache_controls is None) or ( - len(allowed_cache_controls) == 0 - ): # assume empty list to be nullable - https://github.com/prisma/prisma/issues/847#issuecomment-546895663 - return - if data.get("cache", None) is None: return cache_args = data.get("cache", None) if isinstance(cache_args, dict): for k, v in cache_args.items(): - if k not in allowed_cache_controls: + if ( + (allowed_cache_controls is not None) + and (isinstance(allowed_cache_controls, list)) + and ( + len(allowed_cache_controls) > 0 + ) # assume empty list to be nullable - https://github.com/prisma/prisma/issues/847#issuecomment-546895663 + and k not in allowed_cache_controls + ): raise HTTPException( status_code=403, detail=f"Not allowed to set {k} as a cache control. Contact admin to change permissions.", diff --git a/litellm/proxy/hooks/presidio_pii_masking.py b/litellm/proxy/hooks/presidio_pii_masking.py index 85e6260745..5152046bc5 100644 --- a/litellm/proxy/hooks/presidio_pii_masking.py +++ b/litellm/proxy/hooks/presidio_pii_masking.py @@ -61,7 +61,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): except: pass - async def check_pii(self, text: str) -> str: + async def check_pii(self, text: str, output_parse_pii: bool) -> str: """ [TODO] make this more performant for high-throughput scenario """ @@ -92,10 +92,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): start = item["start"] end = item["end"] replacement = item["text"] # replacement token - if ( - item["operator"] == "replace" - and litellm.output_parse_pii == True - ): + if item["operator"] == "replace" and output_parse_pii == True: # check if token in dict # if exists, add a uuid to the replacement token for swapping back to the original text in llm response output parsing if replacement in self.pii_tokens: @@ -125,13 +122,26 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): For multiple messages in /chat/completions, we'll need to call them in parallel. """ + permissions = user_api_key_dict.permissions + + if permissions.get("pii", True) == False: # allow key to turn off pii masking + return data + + output_parse_pii = permissions.get( + "output_parse_pii", litellm.output_parse_pii + ) # allow key to turn on/off output parsing for pii + if call_type == "completion": # /chat/completions requests messages = data["messages"] tasks = [] for m in messages: if isinstance(m["content"], str): - tasks.append(self.check_pii(text=m["content"])) + tasks.append( + self.check_pii( + text=m["content"], output_parse_pii=output_parse_pii + ) + ) responses = await asyncio.gather(*tasks) for index, r in enumerate(responses): if isinstance(messages[index]["content"], str): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7c522f7793..62e134cdc5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1016,7 +1016,10 @@ async def update_database( valid_token.spend = new_spend user_api_key_cache.set_cache(key=token, value=valid_token) except Exception as e: - verbose_proxy_logger.info(f"Update Key DB Call failed to execute") + traceback.print_exc() + verbose_proxy_logger.info( + f"Update Key DB Call failed to execute - {str(e)}" + ) ### UPDATE SPEND LOGS ### async def _insert_spend_log_to_db(): @@ -1631,6 +1634,7 @@ async def generate_key_helper_fn( update_key_values: Optional[dict] = None, key_alias: Optional[str] = None, allowed_cache_controls: Optional[list] = [], + permissions: Optional[dict] = {}, ): global prisma_client, custom_db_client, user_api_key_cache @@ -1662,12 +1666,14 @@ async def generate_key_helper_fn( aliases_json = json.dumps(aliases) config_json = json.dumps(config) + permissions_json = json.dumps(permissions) metadata_json = json.dumps(metadata) user_id = user_id or str(uuid.uuid4()) user_role = user_role or "app_user" tpm_limit = tpm_limit rpm_limit = rpm_limit allowed_cache_controls = allowed_cache_controls + try: # Create a new verification token (you may want to enhance this logic based on your needs) user_data = { @@ -1703,6 +1709,7 @@ async def generate_key_helper_fn( "budget_duration": key_budget_duration, "budget_reset_at": key_reset_at, "allowed_cache_controls": allowed_cache_controls, + "permissions": permissions_json, } if ( general_settings.get("allow_user_auth", False) == True @@ -1716,6 +1723,8 @@ async def generate_key_helper_fn( saved_token["config"] = json.loads(saved_token["config"]) if isinstance(saved_token["metadata"], str): saved_token["metadata"] = json.loads(saved_token["metadata"]) + if isinstance(saved_token["permissions"], str): + saved_token["permissions"] = json.loads(saved_token["permissions"]) if saved_token.get("expires", None) is not None and isinstance( saved_token["expires"], datetime ): @@ -1878,9 +1887,9 @@ async def initialize( user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param user_max_tokens = max_tokens dynamic_config[user_model]["max_tokens"] = max_tokens @@ -3044,6 +3053,7 @@ async def generate_key_fn( - max_budget: Optional[float] - Specify max budget for a given key. - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } + - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} Returns: - key: (str) The generated api key diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 283269001b..d102be25d3 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -446,6 +446,8 @@ def hf_test_completion_tgi(): ) # Add any assertions here to check the response print(response) + except litellm.ServiceUnavailableError as e: + pass except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/utils.py b/litellm/utils.py index aad2cfa53c..2e54b5e446 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1105,12 +1105,12 @@ class Logging: self.call_type == CallTypes.aimage_generation.value or self.call_type == CallTypes.image_generation.value ): - self.model_call_details[ - "response_cost" - ] = litellm.completion_cost( - completion_response=result, - model=self.model, - call_type=self.call_type, + self.model_call_details["response_cost"] = ( + litellm.completion_cost( + completion_response=result, + model=self.model, + call_type=self.call_type, + ) ) else: # check if base_model set on azure @@ -1118,12 +1118,12 @@ class Logging: model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details[ - "response_cost" - ] = litellm.completion_cost( - completion_response=result, - call_type=self.call_type, - model=base_model, + self.model_call_details["response_cost"] = ( + litellm.completion_cost( + completion_response=result, + call_type=self.call_type, + model=base_model, + ) ) verbose_logger.debug( f"Model={self.model}; cost={self.model_call_details['response_cost']}" @@ -1192,9 +1192,9 @@ class Logging: verbose_logger.debug( f"Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details[ - "complete_streaming_response" - ] = complete_streaming_response + self.model_call_details["complete_streaming_response"] = ( + complete_streaming_response + ) try: if self.model_call_details.get("cache_hit", False) == True: self.model_call_details["response_cost"] = 0.0 @@ -1204,11 +1204,11 @@ class Logging: model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details[ - "response_cost" - ] = litellm.completion_cost( - completion_response=complete_streaming_response, - model=base_model, + self.model_call_details["response_cost"] = ( + litellm.completion_cost( + completion_response=complete_streaming_response, + model=base_model, + ) ) verbose_logger.debug( f"Model={self.model}; cost={self.model_call_details['response_cost']}" @@ -1495,10 +1495,10 @@ class Logging: ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] callback.log_success_event( @@ -1578,9 +1578,9 @@ class Logging: verbose_logger.debug( "Async success callbacks: Got a complete streaming response" ) - self.model_call_details[ - "complete_streaming_response" - ] = complete_streaming_response + self.model_call_details["complete_streaming_response"] = ( + complete_streaming_response + ) try: if self.model_call_details.get("cache_hit", False) == True: self.model_call_details["response_cost"] = 0.0 @@ -2319,9 +2319,9 @@ def client(original_function): ): print_verbose(f"Checking Cache") preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) - kwargs[ - "preset_cache_key" - ] = preset_cache_key # for streaming calls, we need to pass the preset_cache_key + kwargs["preset_cache_key"] = ( + preset_cache_key # for streaming calls, we need to pass the preset_cache_key + ) cached_result = litellm.cache.get_cache(*args, **kwargs) if cached_result != None: if "detail" in cached_result: @@ -2619,17 +2619,17 @@ def client(original_function): cached_result = None elif isinstance(litellm.cache.cache, RedisSemanticCache): preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) - kwargs[ - "preset_cache_key" - ] = preset_cache_key # for streaming calls, we need to pass the preset_cache_key + kwargs["preset_cache_key"] = ( + preset_cache_key # for streaming calls, we need to pass the preset_cache_key + ) cached_result = await litellm.cache.async_get_cache( *args, **kwargs ) else: preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) - kwargs[ - "preset_cache_key" - ] = preset_cache_key # for streaming calls, we need to pass the preset_cache_key + kwargs["preset_cache_key"] = ( + preset_cache_key # for streaming calls, we need to pass the preset_cache_key + ) cached_result = litellm.cache.get_cache(*args, **kwargs) if cached_result is not None and not isinstance( @@ -3959,16 +3959,16 @@ def get_optional_params( True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params[ - "functions_unsupported_model" - ] = non_default_params.pop("tools") + optional_params["functions_unsupported_model"] = ( + non_default_params.pop("tools") + ) non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params[ - "functions_unsupported_model" - ] = non_default_params.pop("functions") + optional_params["functions_unsupported_model"] = ( + non_default_params.pop("functions") + ) elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -4148,9 +4148,9 @@ def get_optional_params( optional_params["top_p"] = top_p if n is not None: optional_params["best_of"] = n - optional_params[ - "do_sample" - ] = True # Need to sample if you want best of for hf inference endpoints + optional_params["do_sample"] = ( + True # Need to sample if you want best of for hf inference endpoints + ) if stream is not None: optional_params["stream"] = stream if stop is not None: @@ -4195,9 +4195,9 @@ def get_optional_params( if max_tokens is not None: optional_params["max_tokens"] = max_tokens if frequency_penalty is not None: - optional_params[ - "repetition_penalty" - ] = frequency_penalty # https://docs.together.ai/reference/inference + optional_params["repetition_penalty"] = ( + frequency_penalty # https://docs.together.ai/reference/inference + ) if stop is not None: optional_params["stop"] = stop if tools is not None: @@ -4313,9 +4313,9 @@ def get_optional_params( optional_params["top_p"] = top_p if n is not None: optional_params["best_of"] = n - optional_params[ - "do_sample" - ] = True # Need to sample if you want best of for hf inference endpoints + optional_params["do_sample"] = ( + True # Need to sample if you want best of for hf inference endpoints + ) if stream is not None: optional_params["stream"] = stream if stop is not None: @@ -4638,9 +4638,9 @@ def get_optional_params( extra_body["safe_mode"] = safe_mode if random_seed is not None: extra_body["random_seed"] = random_seed - optional_params[ - "extra_body" - ] = extra_body # openai client supports `extra_body` param + optional_params["extra_body"] = ( + extra_body # openai client supports `extra_body` param + ) elif custom_llm_provider == "openrouter": supported_params = [ "functions", @@ -4709,9 +4709,9 @@ def get_optional_params( extra_body["models"] = models if route is not None: extra_body["route"] = route - optional_params[ - "extra_body" - ] = extra_body # openai client supports `extra_body` param + optional_params["extra_body"] = ( + extra_body # openai client supports `extra_body` param + ) else: # assume passing in params for openai/azure openai supported_params = [ "functions", @@ -8475,10 +8475,10 @@ class CustomStreamWrapper: try: completion_obj["content"] = chunk.text if hasattr(chunk.candidates[0], "finish_reason"): - model_response.choices[ - 0 - ].finish_reason = map_finish_reason( - chunk.candidates[0].finish_reason.name + model_response.choices[0].finish_reason = ( + map_finish_reason( + chunk.candidates[0].finish_reason.name + ) ) except: if chunk.candidates[0].finish_reason.name == "SAFETY":