Merge pull request #4632 from BerriAI/litellm_set_ip_address_on_ui

ui - allow setting allowed ip addresses

litellm/proxy/proxy_server.py

@@ -186,6 +186,9 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import (
     router as spend_management_router,
 )
 from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
+from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
+    router as ui_crud_endpoints_router,
+)
 from litellm.proxy.utils import (
     DBClient,
     PrismaClient,
@@ -1281,7 +1284,7 @@ class ProxyConfig:
         return config
 
     async def save_config(self, new_config: dict):
-        global prisma_client, general_settings, user_config_file_path
+        global prisma_client, general_settings, user_config_file_path, store_model_in_db
         # Load existing config
         ## DB - writes valid config to db
         """
@@ -1290,6 +1293,7 @@ class ProxyConfig:
         """
         if prisma_client is not None and (
             general_settings.get("store_model_in_db", False) == True
+            or store_model_in_db
         ):
             # if using - db for config - models are in ModelTable
             new_config.pop("model_list", None)
@@ -9181,3 +9185,4 @@ app.include_router(spend_management_router)
 app.include_router(caching_router)
 app.include_router(analytics_router)
 app.include_router(debugging_endpoints_router)
+app.include_router(ui_crud_endpoints_router)

litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py

@@ -0,0 +1,113 @@
+#### CRUD ENDPOINTS for UI Settings #####
+from datetime import datetime, timedelta, timezone
+from typing import List, Optional
+
+import fastapi
+from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import *
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+router = APIRouter()
+
+
+class IPAddress(BaseModel):
+    ip: str
+
+
+@router.get(
+    "/get/allowed_ips",
+    tags=["Budget & Spend Tracking"],
+    dependencies=[Depends(user_api_key_auth)],
+    include_in_schema=False,
+)
+async def get_allowed_ips():
+    from litellm.proxy.proxy_server import general_settings
+
+    _allowed_ip = general_settings.get("allowed_ips")
+    return {"data": _allowed_ip}
+
+
+@router.post(
+    "/add/allowed_ip",
+    tags=["Budget & Spend Tracking"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def add_allowed_ip(ip_address: IPAddress):
+    from litellm.proxy.proxy_server import (
+        general_settings,
+        prisma_client,
+        proxy_config,
+        store_model_in_db,
+    )
+
+    _allowed_ips: List = general_settings.get("allowed_ips", [])
+    if ip_address.ip not in _allowed_ips:
+        _allowed_ips.append(ip_address.ip)
+        general_settings["allowed_ips"] = _allowed_ips
+    else:
+        raise HTTPException(status_code=400, detail="IP address already exists")
+
+    if prisma_client is None:
+        raise Exception("No DB Connected")
+
+    if store_model_in_db is not True:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."
+            },
+        )
+
+    # Load existing config
+    config = await proxy_config.get_config()
+    verbose_proxy_logger.debug("Loaded config: %s", config)
+    if "general_settings" not in config:
+        config["general_settings"] = {}
+
+    if "allowed_ips" not in config["general_settings"]:
+        config["general_settings"]["allowed_ips"] = []
+
+    if ip_address.ip not in config["general_settings"]["allowed_ips"]:
+        config["general_settings"]["allowed_ips"].append(ip_address.ip)
+
+    await proxy_config.save_config(new_config=config)
+
+    return {
+        "message": f"IP {ip_address.ip} address added successfully",
+        "status": "success",
+    }
+
+
+@router.post(
+    "/delete/allowed_ip",
+    tags=["Budget & Spend Tracking"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def delete_allowed_ip(ip_address: IPAddress):
+    from litellm.proxy.proxy_server import general_settings, proxy_config
+
+    _allowed_ips: List = general_settings.get("allowed_ips", [])
+    if ip_address.ip in _allowed_ips:
+        _allowed_ips.remove(ip_address.ip)
+        general_settings["allowed_ips"] = _allowed_ips
+    else:
+        raise HTTPException(status_code=404, detail="IP address not found")
+
+    # Load existing config
+    config = await proxy_config.get_config()
+    verbose_proxy_logger.debug("Loaded config: %s", config)
+    if "general_settings" not in config:
+        config["general_settings"] = {}
+
+    if "allowed_ips" not in config["general_settings"]:
+        config["general_settings"]["allowed_ips"] = []
+
+    if ip_address.ip in config["general_settings"]["allowed_ips"]:
+        config["general_settings"]["allowed_ips"].remove(ip_address.ip)
+
+    await proxy_config.save_config(new_config=config)
+
+    return {"message": f"IP {ip_address.ip} deleted successfully", "status": "success"}

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

@@ -221,6 +221,7 @@ const CreateKeyPage = () => {
               searchParams={searchParams}
               accessToken={accessToken}
               showSSOBanner={showSSOBanner}
+              premiumUser={premiumUser}
             />
           ) : page == "api_ref" ? (
             <APIRef 

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

@@ -15,7 +15,7 @@ import {
   message,
 } from "antd";
 import { CopyToClipboard } from "react-copy-to-clipboard";
-import { Select, SelectItem } from "@tremor/react";
+import { Select, SelectItem, Subtitle } from "@tremor/react";
 import {
   Table,
   TableBody,
@@ -40,6 +40,7 @@ interface AdminPanelProps {
   accessToken: string | null;
   setTeams: React.Dispatch<React.SetStateAction<Object[] | null>>;
   showSSOBanner: boolean;
+  premiumUser: boolean;
 }
 
 import {
@@ -50,12 +51,16 @@ import {
   setCallbacksCall,
   invitationCreateCall,
   getPossibleUserRoles,
+  addAllowedIP,
+  getAllowedIPs,
+  deleteAllowedIP,
 } from "./networking";
 
 const AdminPanel: React.FC<AdminPanelProps> = ({
   searchParams,
   accessToken,
   showSSOBanner,
+  premiumUser,
 }) => {
   const [form] = Form.useForm();
   const [memberForm] = Form.useForm();
@@ -73,6 +78,11 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
   const [isAddSSOModalVisible, setIsAddSSOModalVisible] = useState(false);
   const [isInstructionsModalVisible, setIsInstructionsModalVisible] =
     useState(false);
+  const [isAllowedIPModalVisible, setIsAllowedIPModalVisible] = useState(false);
+  const [isAddIPModalVisible, setIsAddIPModalVisible] = useState(false);
+  const [isDeleteIPModalVisible, setIsDeleteIPModalVisible] = useState(false);
+  const [allowedIPs, setAllowedIPs] = useState<string[]>([]);
+  const [ipToDelete, setIPToDelete] = useState<string | null>(null);
   const router = useRouter();
 
   const [possibleUIRoles, setPossibleUIRoles] = useState<null | Record<
@@ -85,6 +95,8 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
     isLocal ? "http://localhost:4000" : ""
   );
 
+  const all_ip_address_allowed = "All IP Addresses Allowed";
+
   let nonSssoUrl;
   try {
     nonSssoUrl = window.location.origin;
@@ -93,6 +105,72 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
   }
   nonSssoUrl += "/fallback/login";
 
+  const handleShowAllowedIPs = async () => {
+    try {
+      if (premiumUser !== true) {
+        message.error(
+          "This feature is only available for premium users. Please upgrade your account."
+        )
+        return
+      }
+      if (accessToken) {
+        const data = await getAllowedIPs(accessToken);
+        setAllowedIPs(data && data.length > 0 ? data : [all_ip_address_allowed]);
+      } else {
+        setAllowedIPs([all_ip_address_allowed]);
+      }
+    } catch (error) {
+      console.error("Error fetching allowed IPs:", error);
+      message.error(`Failed to fetch allowed IPs ${error}`);
+      setAllowedIPs([all_ip_address_allowed]);
+    } finally {
+      if (premiumUser === true) {
+        setIsAllowedIPModalVisible(true);
+      }
+    }
+  };
+  
+  const handleAddIP = async (values: { ip: string }) => {
+    try {
+      if (accessToken) {
+        await addAllowedIP(accessToken, values.ip);
+        // Fetch the updated list of IPs
+        const updatedIPs = await getAllowedIPs(accessToken);
+        setAllowedIPs(updatedIPs);
+        message.success('IP address added successfully');
+      }
+    } catch (error) {
+      console.error("Error adding IP:", error);
+      message.error(`Failed to add IP address ${error}`);
+    } finally {
+      setIsAddIPModalVisible(false);
+    }
+  };
+  
+  const handleDeleteIP = async (ip: string) => {
+    setIPToDelete(ip);
+    setIsDeleteIPModalVisible(true);
+  };
+  
+  const confirmDeleteIP = async () => {
+    if (ipToDelete && accessToken) {
+      try {
+        await deleteAllowedIP(accessToken, ipToDelete);
+        // Fetch the updated list of IPs
+        const updatedIPs = await getAllowedIPs(accessToken);
+        setAllowedIPs(updatedIPs.length > 0 ? updatedIPs : [all_ip_address_allowed]);
+        message.success('IP address deleted successfully');
+      } catch (error) {
+        console.error("Error deleting IP:", error);
+        message.error(`Failed to delete IP address ${error}`);
+      } finally {
+        setIsDeleteIPModalVisible(false);
+        setIPToDelete(null);
+      }
+    }
+  };
+
+
   const handleAddSSOOk = () => {
     setIsAddSSOModalVisible(false);
     form.resetFields();
@@ -532,10 +610,21 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
           </div>
         </Col>
       </Grid>
-      <Grid>
-        <Title level={4}>Add SSO</Title>
+      <Grid >
+        <Card>
+        <Title level={4}> ✨ Security Settings</Title>
+        <div style={{ display: 'flex', flexDirection: 'column', gap: '1rem', marginTop: '1rem' }}>
+    <div>
+      <Button onClick={() => premiumUser === true ? setIsAddSSOModalVisible(true) : message.error("Only premium users can add SSO")}>Add SSO</Button>
+    </div>
+    <div>
+      <Button onClick={handleShowAllowedIPs}>Allowed IPs</Button>
+    </div>
+  </div>
+        </Card>
+       
         <div className="flex justify-start mb-4">
-          <Button onClick={() => setIsAddSSOModalVisible(true)}>Add SSO</Button>
+         
           <Modal
             title="Add SSO"
             visible={isAddSSOModalVisible}
@@ -632,6 +721,81 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
               <Button2 onClick={handleInstructionsOk}>Done</Button2>
             </div>
           </Modal>
+          <Modal
+          title="Manage Allowed IP Addresses"
+          width={800}
+          visible={isAllowedIPModalVisible}
+          onCancel={() => setIsAllowedIPModalVisible(false)}
+          footer={[
+            <Button className="mx-1"key="add" onClick={() => setIsAddIPModalVisible(true)}>
+              Add IP Address
+            </Button>,
+            <Button key="close" onClick={() => setIsAllowedIPModalVisible(false)}>
+              Close
+            </Button>
+          ]}
+        >
+          <Table>
+  <TableHead>
+    <TableRow>
+      <TableHeaderCell>IP Address</TableHeaderCell>
+      <TableHeaderCell className="text-right">Action</TableHeaderCell>
+    </TableRow>
+  </TableHead>
+  <TableBody>
+  {allowedIPs.map((ip, index) => (
+  <TableRow key={index}>
+    <TableCell>{ip}</TableCell>
+    <TableCell className="text-right">
+      {ip !== all_ip_address_allowed && (
+        <Button onClick={() => handleDeleteIP(ip)} color="red" size="xs">
+          Delete
+        </Button>
+      )}
+    </TableCell>
+  </TableRow>
+))}
+  </TableBody>
+</Table>
+        </Modal>
+
+        <Modal
+          title="Add Allowed IP Address"
+          visible={isAddIPModalVisible}
+          onCancel={() => setIsAddIPModalVisible(false)}
+          footer={null}
+        >
+          <Form onFinish={handleAddIP}>
+            <Form.Item
+              name="ip"
+              rules={[{ required: true, message: 'Please enter an IP address' }]}
+            >
+              <Input placeholder="Enter IP address" />
+            </Form.Item>
+            <Form.Item>
+              <Button2 htmlType="submit">
+                Add IP Address
+              </Button2>
+            </Form.Item>
+          </Form>
+        </Modal>
+
+        <Modal
+          title="Confirm Delete"
+          visible={isDeleteIPModalVisible}
+          onCancel={() => setIsDeleteIPModalVisible(false)}
+          onOk={confirmDeleteIP}
+          footer={[
+            <Button className="mx-1"key="delete" onClick={() => confirmDeleteIP()}>
+              Yes
+            </Button>,
+            <Button key="close" onClick={() => setIsDeleteIPModalVisible(false)}>
+              Close
+            </Button>
+          ]}
+        >
+          <p>Are you sure you want to delete the IP address: {ipToDelete}?</p>
+        </Modal>
         </div>
         <Callout title="Login without SSO" color="teal">
           If you need to login without sso, you can access{" "}

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

@@ -98,7 +98,7 @@ const Sidebar: React.FC<SidebarProps> = ({
           ) : null}
           {userRole == "Admin" ? (
             <Menu.Item key="12" onClick={() => setPage("admin-panel")}>
-              <Text>Admin</Text>
+              <Text>Admin Settings</Text>
             </Menu.Item>
           ) : null}
           <Menu.Item key="13" onClick={() => setPage("api_ref")}>

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

@@ -788,6 +788,95 @@ export const modelHubCall = async (accessToken: String) => {
   }
 };
 
+// Function to get allowed IPs
+export const getAllowedIPs = async (accessToken: String) => {
+  try {
+    let url = proxyBaseUrl
+      ? `${proxyBaseUrl}/get/allowed_ips`
+      : `/get/allowed_ips`;
+
+    const response = await fetch(url, {
+      method: "GET",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+    });
+
+    if (!response.ok) {
+      const errorData = await response.text();
+      throw new Error(`Network response was not ok: ${errorData}`);
+    }
+
+    const data = await response.json();
+    console.log("getAllowedIPs:", data);
+    return data.data; // Assuming the API returns { data: [...] }
+  } catch (error) {
+    console.error("Failed to get allowed IPs:", error);
+    throw error;
+  }
+};
+
+// Function to add an allowed IP
+export const addAllowedIP = async (accessToken: String, ip: String) => {
+  try {
+    let url = proxyBaseUrl
+      ? `${proxyBaseUrl}/add/allowed_ip`
+      : `/add/allowed_ip`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({ ip: ip }),
+    });
+
+    if (!response.ok) {
+      const errorData = await response.text();
+      throw new Error(`Network response was not ok: ${errorData}`);
+    }
+
+    const data = await response.json();
+    console.log("addAllowedIP:", data);
+    return data;
+  } catch (error) {
+    console.error("Failed to add allowed IP:", error);
+    throw error;
+  }
+};
+
+// Function to delete an allowed IP
+export const deleteAllowedIP = async (accessToken: String, ip: String) => {
+  try {
+    let url = proxyBaseUrl
+      ? `${proxyBaseUrl}/delete/allowed_ip`
+      : `/delete/allowed_ip`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({ ip: ip }),
+    });
+
+    if (!response.ok) {
+      const errorData = await response.text();
+      throw new Error(`Network response was not ok: ${errorData}`);
+    }
+
+    const data = await response.json();
+    console.log("deleteAllowedIP:", data);
+    return data;
+  } catch (error) {
+    console.error("Failed to delete allowed IP:", error);
+    throw error;
+  }
+};
+
 export const modelMetricsCall = async (
   accessToken: String,
   userID: String,