(feat)team-specific models

This commit is contained in:
ishaan-jaff 2024-02-10 21:13:39 -08:00
parent afd0d7cb41
commit 5ac5f6c53c
2 changed files with 18 additions and 0 deletions

View file

@ -92,6 +92,7 @@ from litellm.proxy.utils import (
hash_token, hash_token,
html_form, html_form,
_read_request_body, _read_request_body,
_is_valid_team_configs,
) )
from litellm.proxy.secret_managers.google_kms import load_google_kms from litellm.proxy.secret_managers.google_kms import load_google_kms
import pydantic import pydantic
@ -2294,6 +2295,9 @@ async def chat_completion(
pass pass
else: else:
team_id = team_config.pop("team_id", None) team_id = team_config.pop("team_id", None)
_is_valid_team_configs(
team_id=team_id, team_config=team_config, request_data=data
)
data["metadata"]["team_id"] = team_id data["metadata"]["team_id"] = team_id
data = { data = {
**team_config, **team_config,

View file

@ -1285,6 +1285,20 @@ async def _read_request_body(request):
return request_data return request_data
def _is_valid_team_configs(team_id=None, team_config=None, request_data=None):
if team_id is None or team_config is None or request_data is None:
return
# check if valid model called for team
if "models" in team_config:
valid_models = team_config["models"]
model_in_request = request_data["model"]
if model_in_request not in valid_models:
raise Exception(
f"Invalid model for team {team_id}: {model_in_request}. Valid models for team are: {valid_models}\n"
)
return
# LiteLLM Admin UI - Non SSO Login # LiteLLM Admin UI - Non SSO Login
html_form = """ html_form = """
<!DOCTYPE html> <!DOCTYPE html>