LiteLLM Minor Fixes & Improvements (11/23/2024) (#6870)

* feat(pass_through_endpoints/): support logging anthropic/gemini pass through calls to langfuse/s3/etc.

* fix(utils.py): allow disabling end user cost tracking with new param

Allows proxy admin to disable cost tracking for end user - keeps prometheus metrics small

* docs(configs.md): add disable_end_user_cost_tracking reference to docs

* feat(key_management_endpoints.py): add support for restricting access to `/key/generate` by team/proxy level role

Enables admin to restrict key creation, and assign team admins to handle distributing keys

* test(test_key_management.py): add unit testing for personal / team key restriction checks

* docs: add docs on restricting key creation

* docs(finetuned_models.md): add new guide on calling finetuned models

* docs(input.md): cleanup anthropic supported params

Closes https://github.com/BerriAI/litellm/issues/6856

* test(test_embedding.py): add test for passing extra headers via embedding

* feat(cohere/embed): pass client to async embedding

* feat(rerank.py): add `/v1/rerank` if missing for cohere base url

Closes https://github.com/BerriAI/litellm/issues/6844

* fix(main.py): pass extra_headers param to openai

Fixes https://github.com/BerriAI/litellm/issues/6836

* fix(litellm_logging.py): don't disable global callbacks when dynamic callbacks are set

Fixes issue where global callbacks - e.g. prometheus were overriden when langfuse was set dynamically

* fix(handler.py): fix linting error

* fix: fix typing

* build: add conftest to proxy_admin_ui_tests/

* test: fix test

* fix: fix linting errors

* test: fix test

* fix: fix pass through testing

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -40,6 +40,77 @@ from litellm.proxy.utils import (
 )
 from litellm.secret_managers.main import get_secret
 
+
+def _is_team_key(data: GenerateKeyRequest):
+    return data.team_id is not None
+
+
+def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth):
+    if (
+        litellm.key_generation_settings is None
+        or litellm.key_generation_settings.get("team_key_generation") is None
+    ):
+        return True
+
+    if user_api_key_dict.team_member is None:
+        raise HTTPException(
+            status_code=400,
+            detail=f"User not assigned to team. Got team_member={user_api_key_dict.team_member}",
+        )
+
+    team_member_role = user_api_key_dict.team_member.role
+    if (
+        team_member_role
+        not in litellm.key_generation_settings["team_key_generation"][  # type: ignore
+            "allowed_team_member_roles"
+        ]
+    ):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Team member role {team_member_role} not in allowed_team_member_roles={litellm.key_generation_settings['team_key_generation']['allowed_team_member_roles']}",  # type: ignore
+        )
+    return True
+
+
+def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth):
+
+    if (
+        litellm.key_generation_settings is None
+        or litellm.key_generation_settings.get("personal_key_generation") is None
+    ):
+        return True
+
+    if (
+        user_api_key_dict.user_role
+        not in litellm.key_generation_settings["personal_key_generation"][  # type: ignore
+            "allowed_user_roles"
+        ]
+    ):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}",  # type: ignore
+        )
+    return True
+
+
+def key_generation_check(
+    user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest
+) -> bool:
+    """
+    Check if admin has restricted key creation to certain roles for teams or individuals
+    """
+    if litellm.key_generation_settings is None:
+        return True
+
+    ## check if key is for team or individual
+    is_team_key = _is_team_key(data=data)
+
+    if is_team_key:
+        return _team_key_generation_check(user_api_key_dict)
+    else:
+        return _personal_key_generation_check(user_api_key_dict=user_api_key_dict)
+
+
 router = APIRouter()
 
 
@@ -131,6 +202,8 @@ async def generate_key_fn(  # noqa: PLR0915
                 raise HTTPException(
                     status_code=status.HTTP_403_FORBIDDEN, detail=message
                 )
+        elif litellm.key_generation_settings is not None:
+            key_generation_check(user_api_key_dict=user_api_key_dict, data=data)
         # check if user set default key/generate params on config.yaml
         if litellm.default_key_generate_params is not None:
             for elem in data: