LiteLLM Minor Fixes and Improvements (11/09/2024) (#5634)

* fix(caching.py): set ttl for async_increment cache

fixes issue where ttl for redis client was not being set on increment_cache

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

* fix(caching.py): fix increment cache w/ ttl for sync increment cache on redis

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

* fix(router.py): support adding retry policy + allowed fails policy via config.yaml

* fix(router.py): don't cooldown single deployments

No point, as there's no other deployment to loadbalance with.

* fix(user_api_key_auth.py): support setting allowed email domains on jwt tokens

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

* docs(token_auth.md): add user upsert + allowed email domain to jwt auth docs

* fix(litellm_pre_call_utils.py): fix dynamic key logging when team id is set

Fixes issue where key logging would not be set if team metadata was not none

* fix(secret_managers/main.py): load environment variables correctly

Fixes issue where os.environ/ was not being loaded correctly

* test(test_router.py): fix test

* feat(spend_tracking_utils.py): support logging additional usage params - e.g. prompt caching values for deepseek

* test: fix tests

* test: fix test

* test: fix test

* test: fix test

* test: fix test

litellm/proxy/auth/handle_jwt.py

@@ -78,6 +78,19 @@ class JWTHandler:
             return False
         return True
 
+    def is_enforced_email_domain(self) -> bool:
+        """
+        Returns:
+        - True: if 'user_allowed_email_domain' is set
+        - False: if 'user_allowed_email_domain' is None
+        """
+
+        if self.litellm_jwtauth.user_allowed_email_domain is not None and isinstance(
+            self.litellm_jwtauth.user_allowed_email_domain, str
+        ):
+            return True
+        return False
+
     def get_team_id(self, token: dict, default_value: Optional[str]) -> Optional[str]:
         try:
             if self.litellm_jwtauth.team_id_jwt_field is not None:
@@ -90,12 +103,14 @@ class JWTHandler:
             team_id = default_value
         return team_id
 
-    def is_upsert_user_id(self) -> bool:
+    def is_upsert_user_id(self, valid_user_email: Optional[bool] = None) -> bool:
         """
         Returns:
-        - True: if 'user_id_upsert' is set
+        - True: if 'user_id_upsert' is set AND valid_user_email is not False
         - False: if not
         """
+        if valid_user_email is False:
+            return False
         return self.litellm_jwtauth.user_id_upsert
 
     def get_user_id(self, token: dict, default_value: Optional[str]) -> Optional[str]:
@@ -103,11 +118,23 @@ class JWTHandler:
             if self.litellm_jwtauth.user_id_jwt_field is not None:
                 user_id = token[self.litellm_jwtauth.user_id_jwt_field]
             else:
-                user_id = None
+                user_id = default_value
         except KeyError:
             user_id = default_value
         return user_id
 
+    def get_user_email(
+        self, token: dict, default_value: Optional[str]
+    ) -> Optional[str]:
+        try:
+            if self.litellm_jwtauth.user_email_jwt_field is not None:
+                user_email = token[self.litellm_jwtauth.user_email_jwt_field]
+            else:
+                user_email = None
+        except KeyError:
+            user_email = default_value
+        return user_email
+
     def get_org_id(self, token: dict, default_value: Optional[str]) -> Optional[str]:
         try:
             if self.litellm_jwtauth.org_id_jwt_field is not None:
@@ -183,6 +210,16 @@ class JWTHandler:
 
         return public_key
 
+    def is_allowed_domain(self, user_email: str) -> bool:
+        if self.litellm_jwtauth.user_allowed_email_domain is None:
+            return True
+
+        email_domain = user_email.split("@")[-1]  # Extract domain from email
+        if email_domain == self.litellm_jwtauth.user_allowed_email_domain:
+            return True
+        else:
+            return False
+
     async def auth_jwt(self, token: str) -> dict:
         # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
         # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret