From 3ac0461d43c4ffd426539c889dd0381a532a9130 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 24 Feb 2025 23:42:04 -0800
Subject: [PATCH] (UI) Allow adding MSFT SSO on UI  (#8779)

* ui - allow adding SSO on ui

* add msft sso to ui

* sso on admin ui
---
 .../src/components/SSOModals.tsx              | 215 ++++++++++++++++++
 .../src/components/admins.tsx                 | 136 +++--------
 2 files changed, 248 insertions(+), 103 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/SSOModals.tsx

diff --git a/ui/litellm-dashboard/src/components/SSOModals.tsx b/ui/litellm-dashboard/src/components/SSOModals.tsx
new file mode 100644
index 00000000000..6e0c5740161
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/SSOModals.tsx
@@ -0,0 +1,215 @@
+import React from "react";
+import { Modal, Form, Input, Button as Button2, Select } from "antd";
+import { Text, TextInput } from "@tremor/react";
+
+interface SSOModalsProps {
+  isAddSSOModalVisible: boolean;
+  isInstructionsModalVisible: boolean;
+  handleAddSSOOk: () => void;
+  handleAddSSOCancel: () => void;
+  handleShowInstructions: (formValues: Record<string, any>) => void;
+  handleInstructionsOk: () => void;
+  handleInstructionsCancel: () => void;
+  form: any; // Replace with proper Form type if available
+}
+
+const ssoProviderLogoMap: Record<string, string> = {
+  google: "https://artificialanalysis.ai/img/logos/google_small.svg",
+  microsoft: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg",
+  okta: "https://www.okta.com/sites/default/files/Okta_Logo_BrightBlue_Medium.png",
+  generic: "",
+};
+
+// Define the SSO provider configuration type
+interface SSOProviderConfig {
+  envVarMap: Record<string, string>;
+  fields: Array<{
+    label: string;
+    name: string;
+    placeholder?: string;
+  }>;
+}
+
+// Define configurations for each SSO provider
+const ssoProviderConfigs: Record<string, SSOProviderConfig> = {
+  google: {
+    envVarMap: {
+      google_client_id: 'GOOGLE_CLIENT_ID',
+      google_client_secret: 'GOOGLE_CLIENT_SECRET',
+    },
+    fields: [
+      { label: 'GOOGLE CLIENT ID', name: 'google_client_id' },
+      { label: 'GOOGLE CLIENT SECRET', name: 'google_client_secret' },
+    ],
+  },
+  microsoft: {
+    envVarMap: {
+      microsoft_client_id: 'MICROSOFT_CLIENT_ID',
+      microsoft_client_secret: 'MICROSOFT_CLIENT_SECRET',
+      microsoft_tenant: 'MICROSOFT_TENANT',
+    },
+    fields: [
+      { label: 'MICROSOFT CLIENT ID', name: 'microsoft_client_id' },
+      { label: 'MICROSOFT CLIENT SECRET', name: 'microsoft_client_secret' },
+      { label: 'MICROSOFT TENANT', name: 'microsoft_tenant' },
+    ],
+  },
+  okta: {
+    envVarMap: {
+      generic_client_id: 'GENERIC_CLIENT_ID',
+      generic_client_secret: 'GENERIC_CLIENT_SECRET',
+      generic_authorization_endpoint: 'GENERIC_AUTHORIZATION_ENDPOINT',
+      generic_token_endpoint: 'GENERIC_TOKEN_ENDPOINT',
+      generic_userinfo_endpoint: 'GENERIC_USERINFO_ENDPOINT',
+    },
+    fields: [
+      { label: 'GENERIC CLIENT ID', name: 'generic_client_id' },
+      { label: 'GENERIC CLIENT SECRET', name: 'generic_client_secret' },
+      { label: 'AUTHORIZATION ENDPOINT', name: 'generic_authorization_endpoint', placeholder: 'https://your-okta-domain/authorize' },
+      { label: 'TOKEN ENDPOINT', name: 'generic_token_endpoint', placeholder: 'https://your-okta-domain/token' },
+      { label: 'USERINFO ENDPOINT', name: 'generic_userinfo_endpoint', placeholder: 'https://your-okta-domain/userinfo' },
+    ],
+  },
+  generic: {
+    envVarMap: {
+      generic_client_id: 'GENERIC_CLIENT_ID',
+      generic_client_secret: 'GENERIC_CLIENT_SECRET',
+      generic_authorization_endpoint: 'GENERIC_AUTHORIZATION_ENDPOINT',
+      generic_token_endpoint: 'GENERIC_TOKEN_ENDPOINT',
+      generic_userinfo_endpoint: 'GENERIC_USERINFO_ENDPOINT',
+    },
+    fields: [
+      { label: 'GENERIC CLIENT ID', name: 'generic_client_id' },
+      { label: 'GENERIC CLIENT SECRET', name: 'generic_client_secret' },
+      { label: 'AUTHORIZATION ENDPOINT', name: 'generic_authorization_endpoint' },
+      { label: 'TOKEN ENDPOINT', name: 'generic_token_endpoint' },
+      { label: 'USERINFO ENDPOINT', name: 'generic_userinfo_endpoint' },
+    ],
+  },
+};
+
+const SSOModals: React.FC<SSOModalsProps> = ({
+  isAddSSOModalVisible,
+  isInstructionsModalVisible,
+  handleAddSSOOk,
+  handleAddSSOCancel,
+  handleShowInstructions,
+  handleInstructionsOk,
+  handleInstructionsCancel,
+  form,
+}) => {
+  // Helper function to render provider fields
+  const renderProviderFields = (provider: string) => {
+    const config = ssoProviderConfigs[provider];
+    if (!config) return null;
+
+    return config.fields.map((field) => (
+      <Form.Item
+        key={field.name}
+        label={field.label}
+        name={field.name}
+        rules={[{ required: true, message: `Please enter the ${field.label.toLowerCase()}` }]}
+      >
+        {field.name.includes('client') ? (
+          <Input.Password />
+        ) : (
+          <TextInput placeholder={field.placeholder} />
+        )}
+      </Form.Item>
+    ));
+  };
+
+  return (
+    <>
+      <Modal
+        title="Add SSO"
+        visible={isAddSSOModalVisible}
+        width={800}
+        footer={null}
+        onOk={handleAddSSOOk}
+        onCancel={handleAddSSOCancel}
+      >
+        <Form
+          form={form}
+          onFinish={handleShowInstructions}
+          labelCol={{ span: 8 }}
+          wrapperCol={{ span: 16 }}
+          labelAlign="left"
+        >
+          <>
+            <Form.Item
+              label="SSO Provider"
+              name="sso_provider"
+              rules={[{ required: true, message: "Please select an SSO provider" }]}
+            >
+              <Select>
+                {Object.entries(ssoProviderLogoMap).map(([value, logo]) => (
+                  <Select.Option key={value} value={value}>
+                    <div style={{ display: 'flex', alignItems: 'center', padding: '4px 0' }}>
+                      {logo && <img src={logo} alt={value} style={{ height: 24, width: 24, marginRight: 12, objectFit: 'contain' }} />}
+                      <span>{value.charAt(0).toUpperCase() + value.slice(1)} SSO</span>
+                    </div>
+                  </Select.Option>
+                ))}
+              </Select>
+            </Form.Item>
+
+            <Form.Item
+              noStyle
+              shouldUpdate={(prevValues, currentValues) => prevValues.sso_provider !== currentValues.sso_provider}
+            >
+              {({ getFieldValue }) => {
+                const provider = getFieldValue('sso_provider');
+                return provider ? renderProviderFields(provider) : null;
+              }}
+            </Form.Item>
+
+            <Form.Item
+              label="Proxy Admin Email"
+              name="user_email"
+              rules={[{ required: true, message: "Please enter the email of the proxy admin" }]}
+            >
+              <TextInput />
+            </Form.Item>
+            <Form.Item
+              label="PROXY BASE URL"
+              name="proxy_base_url"
+              rules={[{ required: true, message: "Please enter the proxy base url" }]}
+            >
+              <TextInput />
+            </Form.Item>
+          </>
+          <div style={{ textAlign: "right", marginTop: "10px" }}>
+            <Button2 htmlType="submit">Save</Button2>
+          </div>
+        </Form>
+      </Modal>
+
+      <Modal
+        title="SSO Setup Instructions"
+        visible={isInstructionsModalVisible}
+        width={800}
+        footer={null}
+        onOk={handleInstructionsOk}
+        onCancel={handleInstructionsCancel}
+      >
+        <p>Follow these steps to complete the SSO setup:</p>
+        <Text className="mt-2">1. DO NOT Exit this TAB</Text>
+        <Text className="mt-2">2. Open a new tab, visit your proxy base url</Text>
+        <Text className="mt-2">
+          3. Confirm your SSO is configured correctly and you can login on the new
+          Tab
+        </Text>
+        <Text className="mt-2">
+          4. If Step 3 is successful, you can close this tab
+        </Text>
+        <div style={{ textAlign: "right", marginTop: "10px" }}>
+          <Button2 onClick={handleInstructionsOk}>Done</Button2>
+        </div>
+      </Modal>
+    </>
+  );
+};
+
+export { ssoProviderConfigs };  // Export for use in other components
+export default SSOModals; 
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx
index 5d5b8c2ccbc..cd0f34b7060 100644
--- a/ui/litellm-dashboard/src/components/admins.tsx
+++ b/ui/litellm-dashboard/src/components/admins.tsx
@@ -36,6 +36,9 @@ import {
 import { PencilAltIcon } from "@heroicons/react/outline";
 import OnboardingModal from "./onboarding_link";
 import { InvitationLink } from "./onboarding_link";
+import SSOModals from "./SSOModals";
+import { ssoProviderConfigs } from './SSOModals';
+
 interface AdminPanelProps {
   searchParams: any;
   accessToken: string | null;
@@ -484,13 +487,27 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
     if (accessToken == null) {
       return;
     }
-    let payload = {
-      environment_variables: {
-        PROXY_BASE_URL: formValues.proxy_base_url,
-        GOOGLE_CLIENT_ID: formValues.google_client_id,
-        GOOGLE_CLIENT_SECRET: formValues.google_client_secret,
-      },
+
+    const provider = formValues.sso_provider;
+    const config = ssoProviderConfigs[provider];
+    
+    const envVars: Record<string, string> = {
+      PROXY_BASE_URL: formValues.proxy_base_url,
     };
+
+    // Add provider-specific environment variables using the configuration
+    if (config) {
+      Object.entries(config.envVarMap).forEach(([formKey, envKey]) => {
+        if (formValues[formKey]) {
+          envVars[envKey] = formValues[formKey];
+        }
+      });
+    }
+
+    const payload = {
+      environment_variables: envVars,
+    };
+    
     setCallbacksCall(accessToken, payload);
   };
   console.log(`admins: ${admins?.length}`);
@@ -512,103 +529,16 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
         </Card>
        
         <div className="flex justify-start mb-4">
-         
-          <Modal
-            title="Add SSO"
-            visible={isAddSSOModalVisible}
-            width={800}
-            footer={null}
-            onOk={handleAddSSOOk}
-            onCancel={handleAddSSOCancel}
-          >
-            <Form
-              form={form}
-              onFinish={handleShowInstructions}
-              labelCol={{ span: 8 }}
-              wrapperCol={{ span: 16 }}
-              labelAlign="left"
-            >
-              <>
-                <Form.Item
-                  label="Admin Email"
-                  name="user_email"
-                  rules={[
-                    {
-                      required: true,
-                      message: "Please enter the email of the proxy admin",
-                    },
-                  ]}
-                >
-                  <Input />
-                </Form.Item>
-                <Form.Item
-                  label="PROXY BASE URL"
-                  name="proxy_base_url"
-                  rules={[
-                    {
-                      required: true,
-                      message: "Please enter the proxy base url",
-                    },
-                  ]}
-                >
-                  <Input />
-                </Form.Item>
-
-                <Form.Item
-                  label="GOOGLE CLIENT ID"
-                  name="google_client_id"
-                  rules={[
-                    {
-                      required: true,
-                      message: "Please enter the google client id",
-                    },
-                  ]}
-                >
-                  <Input.Password />
-                </Form.Item>
-
-                <Form.Item
-                  label="GOOGLE CLIENT SECRET"
-                  name="google_client_secret"
-                  rules={[
-                    {
-                      required: true,
-                      message: "Please enter the google client secret",
-                    },
-                  ]}
-                >
-                  <Input.Password />
-                </Form.Item>
-              </>
-              <div style={{ textAlign: "right", marginTop: "10px" }}>
-                <Button2 htmlType="submit">Save</Button2>
-              </div>
-            </Form>
-          </Modal>
-          <Modal
-            title="SSO Setup Instructions"
-            visible={isInstructionsModalVisible}
-            width={800}
-            footer={null}
-            onOk={handleInstructionsOk}
-            onCancel={handleInstructionsCancel}
-          >
-            <p>Follow these steps to complete the SSO setup:</p>
-            <Text className="mt-2">1. DO NOT Exit this TAB</Text>
-            <Text className="mt-2">
-              2. Open a new tab, visit your proxy base url
-            </Text>
-            <Text className="mt-2">
-              3. Confirm your SSO is configured correctly and you can login on
-              the new Tab
-            </Text>
-            <Text className="mt-2">
-              4. If Step 3 is successful, you can close this tab
-            </Text>
-            <div style={{ textAlign: "right", marginTop: "10px" }}>
-              <Button2 onClick={handleInstructionsOk}>Done</Button2>
-            </div>
-          </Modal>
+          <SSOModals
+            isAddSSOModalVisible={isAddSSOModalVisible}
+            isInstructionsModalVisible={isInstructionsModalVisible}
+            handleAddSSOOk={handleAddSSOOk}
+            handleAddSSOCancel={handleAddSSOCancel}
+            handleShowInstructions={handleShowInstructions}
+            handleInstructionsOk={handleInstructionsOk}
+            handleInstructionsCancel={handleInstructionsCancel}
+            form={form}
+          />
           <Modal
           title="Manage Allowed IP Addresses"
           width={800}