From dbbf6056535ac19120a9ca27dfd2bd307e043955 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 9 Apr 2024 14:50:20 -0700
Subject: [PATCH] fix(proxy_server.py): fix sso callback check to work for user
 email

---
 litellm/proxy/proxy_server.py | 40 ++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index fa571d4c640f..91cc57872991 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -7917,16 +7917,36 @@ def response_convertor(response, client):
                     "user_email": getattr(user_info, "user_id", user_email),
                 }
                 user_role = getattr(user_info, "user_role", None)
-            elif litellm.default_user_params is not None and isinstance(
-                litellm.default_user_params, dict
-            ):
-                user_defined_values = {
-                    "models": litellm.default_user_params.get("models", user_id_models),
-                    "user_id": litellm.default_user_params.get("user_id", user_id),
-                    "user_email": litellm.default_user_params.get(
-                        "user_email", user_email
-                    ),
-                }
+
+            else:
+                ## check if user-email in db ##
+                user_info = await prisma_client.db.litellm_usertable.find_first(
+                    where={"user_email": user_email}
+                )
+                if user_info is not None:
+                    user_defined_values = {
+                        "models": getattr(user_info, "models", user_id_models),
+                        "user_id": getattr(user_info, "user_id", user_id),
+                        "user_email": getattr(user_info, "user_id", user_email),
+                    }
+                    user_role = getattr(user_info, "user_role", None)
+
+                    # update id
+                    await prisma_client.db.litellm_usertable.update_many(
+                        where={"user_email": user_email}, data={"user_id": user_id}  # type: ignore
+                    )
+                elif litellm.default_user_params is not None and isinstance(
+                    litellm.default_user_params, dict
+                ):
+                    user_defined_values = {
+                        "models": litellm.default_user_params.get(
+                            "models", user_id_models
+                        ),
+                        "user_id": litellm.default_user_params.get("user_id", user_id),
+                        "user_email": litellm.default_user_params.get(
+                            "user_email", user_email
+                        ),
+                    }
     except Exception as e:
         pass