Merge pull request #4548 from BerriAI/ui_fix_onboarding_links

[Fix] Invite Links / Onboarding flow on admin ui

litellm/proxy/common_utils/openai_endpoint_utils.py

@@ -0,0 +1,21 @@
+"""
+Contains utils used by OpenAI compatible endpoints 
+"""
+
+
+def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict:
+    """
+    Removes sensitive information from a deployment dictionary.
+
+    Args:
+        deployment_dict (dict): The deployment dictionary to remove sensitive information from.
+
+    Returns:
+        dict: The modified deployment dictionary with sensitive information removed.
+    """
+    deployment_dict["litellm_params"].pop("api_key", None)
+    deployment_dict["litellm_params"].pop("vertex_credentials", None)
+    deployment_dict["litellm_params"].pop("aws_access_key_id", None)
+    deployment_dict["litellm_params"].pop("aws_secret_access_key", None)
+
+    return deployment_dict

litellm/proxy/proxy_server.py

@@ -143,6 +143,9 @@ from litellm.proxy.caching_routes import router as caching_router
 from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
 from litellm.proxy.common_utils.init_callbacks import initialize_callbacks_on_proxy
+from litellm.proxy.common_utils.openai_endpoint_utils import (
+    remove_sensitive_info_from_deployment,
+)
 from litellm.proxy.guardrails.init_guardrails import initialize_guardrails
 from litellm.proxy.health_check import perform_health_check
 from litellm.proxy.health_endpoints._health_endpoints import router as health_router
@@ -6645,26 +6648,81 @@ async def model_metrics_exceptions(
 
 @router.get(
     "/model/info",
-    description="Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",
     tags=["model management"],
     dependencies=[Depends(user_api_key_auth)],
 )
 @router.get(
     "/v1/model/info",
-    description="Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",
     tags=["model management"],
     dependencies=[Depends(user_api_key_auth)],
 )
 async def model_info_v1(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+    litellm_model_id: Optional[str] = None,
 ):
-    global llm_model_list, general_settings, user_config_file_path, proxy_config
+    """
+    Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)
+
+    Parameters:
+        litellm_model_id: Optional[str] = None (this is the value of `x-litellm-model-id` returned in response headers)
+
+        - When litellm_model_id is passed, it will return the info for that specific model
+        - When litellm_model_id is not passed, it will return the info for all models
+
+    Returns:
+        Returns a dictionary containing information about each model.
+
+    Example Response:
+    ```json
+    {
+        "data": [
+                    {
+                        "model_name": "fake-openai-endpoint",
+                        "litellm_params": {
+                            "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
+                            "model": "openai/fake"
+                        },
+                        "model_info": {
+                            "id": "112f74fab24a7a5245d2ced3536dd8f5f9192c57ee6e332af0f0512e08bed5af",
+                            "db_model": false
+                        }
+                    }
+                ]
+    }
+
+    ```
+    """
+    global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router
 
     if llm_model_list is None:
         raise HTTPException(
             status_code=500, detail={"error": "LLM Model List not loaded in"}
         )
 
+    if llm_router is None:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "LLM Router is not loaded in. Make sure you passed models in your config.yaml or on the LiteLLM Admin UI."
+            },
+        )
+
+    if litellm_model_id is not None:
+        # user is trying to get specific model from litellm router
+        deployment_info = llm_router.get_deployment(model_id=litellm_model_id)
+        if deployment_info is None:
+            raise HTTPException(
+                status_code=404,
+                detail={
+                    "error": f"Model id = {litellm_model_id} not found on litellm proxy"
+                },
+            )
+        _deployment_info_dict = deployment_info.model_dump()
+        _deployment_info_dict = remove_sensitive_info_from_deployment(
+            deployment_dict=_deployment_info_dict
+        )
+        return {"data": _deployment_info_dict}
+
     all_models: List[dict] = []
     ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ##
     if llm_model_list is None:
@@ -6726,10 +6784,7 @@ async def model_info_v1(
                 model_info[k] = v
         model["model_info"] = model_info
         # don't return the llm credentials
-        model["litellm_params"].pop("api_key", None)
-        model["litellm_params"].pop("vertex_credentials", None)
-        model["litellm_params"].pop("aws_access_key_id", None)
-        model["litellm_params"].pop("aws_secret_access_key", None)
+        model = remove_sensitive_info_from_deployment(deployment_dict=model)
 
     verbose_proxy_logger.debug("all_models: %s", all_models)
     return {"data": all_models}
@@ -7649,22 +7704,12 @@ async def claim_onboarding_link(data: InvitationClaim):
         )
 
     #### CHECK IF CLAIMED
-    ##### if claimed - check if within valid session (within 10 minutes of being claimed)
+    ##### if claimed - accept
     ##### if unclaimed - reject
 
-    current_time = litellm.utils.get_utc_datetime()
-
-    if invite_obj.is_accepted == True:
-        time_difference = current_time - invite_obj.updated_at
-
-        # Check if the difference is within 10 minutes
-        if time_difference > timedelta(minutes=10):
-            raise HTTPException(
-                status_code=401,
-                detail={
-                    "error": "The invitation link has already been claimed. Please ask your admin for a new invite link."
-                },
-            )
+    if invite_obj.is_accepted is True:
+        # this is a valid invite that was accepted
+        pass
     else:
         raise HTTPException(
             status_code=401,

ui/litellm-dashboard/src/app/onboarding/page.tsx

@@ -33,7 +33,7 @@ export default function Onboarding() {
   const [form] = Form.useForm();
   const searchParams = useSearchParams();
   const token = getCookie('token');
-  const inviteID = searchParams.get("id");
+  const inviteID = searchParams.get("invitation_id");
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [defaultUserEmail, setDefaultUserEmail] = useState<string>("");
   const [userEmail, setUserEmail] = useState<string>("");
@@ -92,7 +92,10 @@ export default function Onboarding() {
     ).then((data) => {
       let litellm_dashboard_ui = "/ui/";
       const user_id = data.data?.user_id || data.user_id;
-      litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwtToken;
+      litellm_dashboard_ui += "?userID=" + user_id;
+
+      // set cookie "token" to jwtToken
+      document.cookie = "token=" + jwtToken;
       console.log("redirecting to:", litellm_dashboard_ui);
 
       window.location.href = litellm_dashboard_ui;
@@ -101,7 +104,7 @@ export default function Onboarding() {
     // redirect to login page
   };
   return (
-    <div className="mx-auto max-w-md mt-10">
+    <div className="mx-auto w-full max-w-md mt-10">
       <Card>
         <Title className="text-sm mb-5 text-center">🚅 LiteLLM</Title>
         <Title className="text-xl">Sign up</Title>

ui/litellm-dashboard/src/app/page.tsx

@@ -77,6 +77,7 @@ const CreateKeyPage = () => {
   const searchParams = useSearchParams();
   const [modelData, setModelData] = useState<any>({ data: [] });
   const userID = searchParams.get("userID");
+  const invitation_id = searchParams.get("invitation_id");
   const token = getCookie('token');
 
   const [page, setPage] = useState("api-keys");
@@ -128,7 +129,23 @@ const CreateKeyPage = () => {
 
   return (
     <Suspense fallback={<div>Loading...</div>}>
-      <div className="flex flex-col min-h-screen">
+      {
+        invitation_id ? (
+          <UserDashboard
+              userID={userID}
+              userRole={userRole}
+              teams={teams}
+              keys={keys}
+              setUserRole={setUserRole}
+              userEmail={userEmail}
+              setUserEmail={setUserEmail}
+              setTeams={setTeams}
+              setKeys={setKeys}
+              setProxySettings={setProxySettings}
+              proxySettings={proxySettings}
+            />
+        ) : (
+        <div className="flex flex-col min-h-screen">
         <Navbar
           userID={userID}
           userRole={userRole}
@@ -140,11 +157,11 @@ const CreateKeyPage = () => {
         />
         <div className="flex flex-1 overflow-auto">
           <div className="mt-8">
-            <Sidebar
-              setPage={setPage}
-              userRole={userRole}
-              defaultSelectedKey={null}
-            />
+                <Sidebar
+                setPage={setPage}
+                userRole={userRole}
+                defaultSelectedKey={null}
+              />            
           </div>
 
           {page == "api-keys" ? (
@@ -250,7 +267,10 @@ const CreateKeyPage = () => {
             />
           )}
         </div>
-      </div>
+      </div>         
+        )
+      }
+      
     </Suspense>
   );
 };

ui/litellm-dashboard/src/components/onboarding_link.tsx

@@ -67,13 +67,13 @@ const OnboardingModal: React.FC<OnboardingProps> = ({
       <div className="flex justify-between pt-5 pb-2">
         <Text>Invitation Link</Text>
         <Text>
-          {baseUrl}/ui/onboarding?id={invitationLinkData?.id}
+          {baseUrl}/ui?invitation_id={invitationLinkData?.id}
         </Text>
       </div>
       <div className="flex justify-end mt-5">
         <div></div>
         <CopyToClipboard
-          text={`${baseUrl}/ui/onboarding?id=${invitationLinkData?.id}`}
+          text={`${baseUrl}/ui?invitation_id=${invitationLinkData?.id}`}
           onCopy={() => message.success("Copied!")}
         >
           <Button variant="primary">Copy invitation link</Button>

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -12,6 +12,7 @@ import ViewKeyTable from "./view_key_table";
 import ViewUserSpend from "./view_user_spend";
 import ViewUserTeam from "./view_user_team";
 import DashboardTeam from "./dashboard_default_team";
+import Onboarding from "../app/onboarding/page";
 import { useSearchParams, useRouter } from "next/navigation";
 import { jwtDecode } from "jwt-decode";
 import { Typography } from "antd";
@@ -76,6 +77,8 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
 
   const token = getCookie('token');
 
+  const invitation_id = searchParams.get("invitation_id");
+
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [teamSpend, setTeamSpend] = useState<number | null>(null);
   const [userModels, setUserModels] = useState<string[]>([]);
@@ -256,8 +259,15 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
     }
   }, [selectedTeam]);
 
+
+  if (invitation_id != null) {
+    return (
+      <Onboarding></Onboarding>
+    )
+  }
+
   if (userID == null || token == null) {
-    // Now you can construct the full URL
+    // user is not logged in as yet 
     const url = proxyBaseUrl
       ? `${proxyBaseUrl}/sso/key/generate`
       : `/sso/key/generate`;