Org Flow Improvements (BerriAI#8549)

* refactor(organization.tsx): initial commit with orgs tab refactor make it similar to 'Teams' tab - simplifies org management actions * fix(page.tsx): pass user orgs to component * fix(organization_view.tsx): fix to pull info from org info endpoint * feat(organization_endpoints.py): return org members when calling /org/info * fix(organization_view.tsx): show org members on info page * feat(organization_view.tsx): allow adding user to org via user email Resolves BerriAI#8330 * fix(organization_endpoints.py): raise better error when duplicate user_email found in db * fix(organization_view.tsx): cleanup user_email for now not in initial org info - will need to prefetch * fix(page.tsx): fix getting user models on page load allows passing down the user models to org * fix(organizations.tsx): fix creating org on ui * fix(proxy/_types.py): include org created at and updated at cleans up ui * fix(navbar.tsx): cleanup * fix(organizations.tsx): fix tpm/rpm limits on org * fix(organizations.tsx): fix linting error * fix(organizations.tsx): fix linting \ * (Feat) - Add `/bedrock/meta.llama3-3-70b-instruct-v1:0` tool calling support + cost tracking + base llm unit test for tool calling (BerriAI#8545) * Add support for bedrock meta.llama3-3-70b-instruct-v1:0 tool calling (BerriAI#8512) * fix(converse_transformation.py): fixing bedrock meta.llama3-3-70b tool calling * test(test_bedrock_completion.py): adding llama3.3 tool compatibility check * add TestBedrockTestSuite * add bedrock llama 3.3 to base llm class * us.meta.llama3-3-70b-instruct-v1:0 * test_basic_tool_calling * TestAzureOpenAIO1 * test_basic_tool_calling * test_basic_tool_calling --------- Co-authored-by: miraclebakelaser <65143272+miraclebakelaser@users.noreply.github.com> * fix(general_settings.tsx): filter out empty dictionaries post fallback delete (BerriAI#8550) Fixes BerriAI#8331 * bump: version 1.61.3 → 1.61.4 * (perf) Fix memory leak on `/completions` route (BerriAI#8551) * initial mem util test * fix _cached_get_model_info_helper * test memory usage * fix tests * fix mem usage --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: miraclebakelaser <65143272+miraclebakelaser@users.noreply.github.com>
acuvity · Feb 20, 2025 · eae05e6 · eae05e6
1 parent 8e80408
commit eae05e6
Show file tree

Hide file tree

Showing 10 changed files with 938 additions and 299 deletions.
diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
@@ -1487,15 +1487,21 @@ class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase):
     organization_id: Optional[str] = None
     organization_alias: Optional[str] = None
     budget_id: str
+    spend: float = 0.0
     metadata: Optional[dict] = None
     models: List[str]
     created_by: str
     updated_by: str
 
 
 class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable):
-    members: List[LiteLLM_OrganizationMembershipTable]
-    teams: List[LiteLLM_TeamTable]
+    """Returned by the /organization/info endpoint and /organization/list endpoint"""
+
+    members: List[LiteLLM_OrganizationMembershipTable] = []
+    teams: List[LiteLLM_TeamTable] = []
+    litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
+    created_at: datetime
+    updated_at: datetime
 
 
 class NewOrganizationResponse(LiteLLM_OrganizationTable):

diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py
@@ -16,6 +16,7 @@
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 
+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
 from litellm.proxy.management_helpers.utils import (
@@ -250,17 +251,49 @@ async def list_organization(
     return response
 
 
-@router.post(
+@router.get(
     "/organization/info",
     tags=["organization management"],
     dependencies=[Depends(user_api_key_auth)],
+    response_model=LiteLLM_OrganizationTableWithMembers,
 )
-async def info_organization(data: OrganizationRequest):
+async def info_organization(organization_id: str):
     """
     Get the org specific information
     """
     from litellm.proxy.proxy_server import prisma_client
 
+    if prisma_client is None:
+        raise HTTPException(status_code=500, detail={"error": "No db connected"})
+
+    response: Optional[LiteLLM_OrganizationTableWithMembers] = (
+        await prisma_client.db.litellm_organizationtable.find_unique(
+            where={"organization_id": organization_id},
+            include={"litellm_budget_table": True, "members": True, "teams": True},
+        )
+    )
+
+    if response is None:
+        raise HTTPException(status_code=404, detail={"error": "Organization not found"})
+
+    response_pydantic_obj = LiteLLM_OrganizationTableWithMembers(
+        **response.model_dump()
+    )
+
+    return response_pydantic_obj
+
+
+@router.post(
+    "/organization/info",
+    tags=["organization management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def deprecated_info_organization(data: OrganizationRequest):
+    """
+    DEPRECATED: Use GET /organization/info instead
+    """
+    from litellm.proxy.proxy_server import prisma_client
+
     if prisma_client is None:
         raise HTTPException(status_code=500, detail={"error": "No db connected"})
 
@@ -378,6 +411,7 @@ async def organization_member_add(
             updated_organization_memberships=updated_organization_memberships,
         )
     except Exception as e:
+        verbose_proxy_logger.exception(f"Error adding member to organization: {e}")
         if isinstance(e, HTTPException):
             raise ProxyException(
                 message=getattr(e, "detail", f"Authentication Error({str(e)})"),
@@ -418,12 +452,17 @@ async def add_member_to_organization(
                 where={"user_id": member.user_id}
             )
 
-        if member.user_email is not None:
-            existing_user_email_row = (
-                await prisma_client.db.litellm_usertable.find_unique(
-                    where={"user_email": member.user_email}
+        if existing_user_id_row is None and member.user_email is not None:
+            try:
+                existing_user_email_row = (
+                    await prisma_client.db.litellm_usertable.find_unique(
+                        where={"user_email": member.user_email}
+                    )
+                )
+            except Exception as e:
+                raise ValueError(
+                    f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}"
                 )
-            )
 
         ## If user does not exist, create a new user
         if existing_user_id_row is None and existing_user_email_row is None:
@@ -477,4 +516,9 @@ async def add_member_to_organization(
         return user_object, organization_membership
 
     except Exception as e:
-        raise ValueError(f"Error adding member to organization: {e}")
+        import traceback
+
+        traceback.print_exc()
+        raise ValueError(
+            f"Error adding member={member} to organization={organization_id}: {e}"
+        )
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
@@ -27,6 +27,8 @@ import CacheDashboard from "@/components/cache_dashboard";
 import { setGlobalLitellmHeaderName } from "@/components/networking";
 import { Organization } from "@/components/networking";
 import GuardrailsPanel from "@/components/guardrails";
+import { fetchUserModels } from "@/components/create_key_button";
+
 function getCookie(name: string) {
   const cookieValue = document.cookie
     .split("; ")
@@ -81,6 +83,7 @@ export default function CreateKeyPage() {
   const [keys, setKeys] = useState<null | any[]>(null);
   const [currentOrg, setCurrentOrg] = useState<Organization>(defaultOrg);
   const [organizations, setOrganizations] = useState<Organization[]>([]);
+  const [userModels, setUserModels] = useState<string[]>([]);
   const [proxySettings, setProxySettings] = useState<ProxySettings>({
     PROXY_BASE_URL: "",
     PROXY_LOGOUT_URL: "",
@@ -171,6 +174,12 @@ export default function CreateKeyPage() {
       }
     }
   }, [token]);
+
+  useEffect(() => {
+    if (accessToken && userID && userRole) {
+      fetchUserModels(userID, userRole, accessToken, setUserModels);
+    }
+  }, [accessToken, userID, userRole]);
 
   const handleOrgChange = (org: Organization) => {
     setCurrentOrg(org);
@@ -283,11 +292,10 @@ export default function CreateKeyPage() {
                 />
               ) : page == "organizations" ? (
                 <Organizations
-                  teams={teams}
-                  setTeams={setTeams}
-                  searchParams={searchParams}
+                  organizations={organizations}
+                  setOrganizations={setOrganizations}
+                  userModels={userModels}
                   accessToken={accessToken}
-                  userID={userID}
                   userRole={userRole}
                   premiumUser={premiumUser}
                 />

diff --git a/ui/litellm-dashboard/src/components/common_components/all_view.tsx b/ui/litellm-dashboard/src/components/common_components/all_view.tsx
@@ -35,6 +35,7 @@ interface DeleteModalProps {
   onCancel: () => void;
   title: string;
   message: string;
+
 }
 
 interface DataTableProps {
@@ -43,14 +44,16 @@ interface DataTableProps {
   actions?: Action[];
   emptyMessage?: string;
   deleteModal?: DeleteModalProps;
+  onItemClick?: (item: any) => void;
 }
 
 const DataTable: React.FC<DataTableProps> = ({
   data,
   columns,
   actions,
   emptyMessage = "No data available",
-  deleteModal
+  deleteModal,
+  onItemClick
 }) => {
   const renderCell = (column: Column, row: any) => {
     const value = row[column.accessor];

diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx
@@ -1,8 +1,8 @@
 import { useState, useCallback } from 'react';
-import { Modal, Form, Button, Select } from 'antd';
+import { Modal, Form, Button, Select, Tooltip } from 'antd';
 import debounce from 'lodash/debounce';
 import { userFilterUICall } from "@/components/networking";
-
+import { InfoCircleOutlined } from '@ant-design/icons';
 interface User {
   user_id: string;
   user_email: string;
@@ -15,24 +15,40 @@ interface UserOption {
   user: User;
 }
 
+interface Role {
+  label: string;
+  value: string;
+  description: string;
+}
+
+
 interface FormValues {
   user_email: string;
   user_id: string;
-  role: 'admin' | 'user';
+  role: string;
 }
 
 interface UserSearchModalProps {
   isVisible: boolean;
   onCancel: () => void;
   onSubmit: (values: FormValues) => void;
   accessToken: string | null;
+  title?: string;
+  roles?: Role[];
+  defaultRole?: string;
 }
 
 const UserSearchModal: React.FC<UserSearchModalProps> = ({ 
   isVisible, 
   onCancel, 
   onSubmit,
-  accessToken
+  accessToken,
+  title = "Add Team Member",
+  roles = [
+    { label: "admin", value: "admin", description: "Admin role. Can create team keys, add members, and manage settings." },
+    { label: "user", value: "user", description: "User role. Can view team info, but not manage it." }
+  ],
+  defaultRole = "user"
 }) => {
   const [form] = Form.useForm<FormValues>();
   const [userOptions, setUserOptions] = useState<UserOption[]>([]);
@@ -97,7 +113,7 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
 
   return (
     <Modal
-      title="Add Team Member" 
+      title={title}
       open={isVisible}
       onCancel={handleClose}
       footer={null}
@@ -110,7 +126,7 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
         wrapperCol={{ span: 16 }}
         labelAlign="left"
         initialValues={{
-          role: "user",
+          role: defaultRole,
         }}
       >
         <Form.Item
@@ -156,9 +172,15 @@ const UserSearchModal: React.FC<UserSearchModalProps> = ({
           name="role"
           className="mb-4"
         >
-          <Select defaultValue="user">
-            <Select.Option value="admin">admin</Select.Option>
-            <Select.Option value="user">user</Select.Option>
+          <Select defaultValue={defaultRole}>
+            {roles.map(role => (
+              <Select.Option key={role.value} value={role.value}>
+                <Tooltip title={role.description}>
+                  <span className="font-medium">{role.label}</span>
+                  <span className="ml-2 text-gray-500 text-sm">- {role.description}</span>
+                </Tooltip>
+              </Select.Option>
+            ))}
           </Select>
         </Form.Item>
 

diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx
@@ -94,6 +94,29 @@ export const getTeamModels = (team: Team | null, allAvailableModels: string[]):
   return unfurlWildcardModelsInList(tempModelsToPick, allAvailableModels);
 };
 
+export const fetchUserModels = async (userID: string, userRole: string, accessToken: string, setUserModels: (models: string[]) => void) => {
+  try {
+    if (userID === null || userRole === null) {
+      return;
+    }
+
+    if (accessToken !== null) {
+      const model_available = await modelAvailableCall(
+        accessToken,
+        userID,
+        userRole
+      );
+      let available_model_names = model_available["data"].map(
+        (element: { id: string }) => element.id
+      );
+      console.log("available_model_names:", available_model_names);
+      setUserModels(available_model_names);
+    }
+  } catch (error) {
+    console.error("Error fetching user models:", error);
+  }
+};
+
 const CreateKey: React.FC<CreateKeyProps> = ({
   userID,
   team,
@@ -126,30 +149,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({
   };
 
   useEffect(() => {
-    const fetchUserModels = async () => {
-      try {
-        if (userID === null || userRole === null) {
-          return;
-        }
-
-        if (accessToken !== null) {
-          const model_available = await modelAvailableCall(
-            accessToken,
-            userID,
-            userRole
-          );
-          let available_model_names = model_available["data"].map(
-            (element: { id: string }) => element.id
-          );
-          console.log("available_model_names:", available_model_names);
-          setUserModels(available_model_names);
-        }
-      } catch (error) {
-        console.error("Error fetching user models:", error);
-      }
-    };
-
-    fetchUserModels();
+    if (userID && userRole && accessToken) {
+      fetchUserModels(userID, userRole, accessToken, setUserModels);
+    }
   }, [accessToken, userID, userRole]);
 
   useEffect(() => {