tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 13:10:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1991 from BerriAI/litellm_moderations_improvements

Browse Source 
feat(proxy_server.py): support key-level permissions for pii masking
This commit is contained in:
Krish Dholakia
2024-02-15 23:08:33 -08:00
committed by GitHub
parent 9b60ef9a3c 999fab82f7
commit b96c0e5ec4
9 changed files with 118 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pre-commit-config.yaml
+2 -2
View File
@@ -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/
litellm/llms/huggingface_restapi.py
+6 -6
View File
@@ -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
litellm/proxy/_types.py
+3 -1
View File
@@ -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(
litellm/proxy/db/dynamo_db.py
+6 -1
View File
@@ -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)
):
litellm/proxy/hooks/cache_control_check.py
+8 -6
View File
@@ -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.",
litellm/proxy/hooks/presidio_pii_masking.py
+16 -6
View File
@@ -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):
litellm/proxy/proxy_server.py
+14 -4
View File
@@ -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
litellm/tests/test_completion.py
+2
View File
@@ -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}")
litellm/utils.py
+61 -61
View File
@@ -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":