feat(extract models from all team models and all proxy models)

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -2799,6 +2799,88 @@ async def info_key_fn(
         raise handle_exception_on_proxy(e)
 
 
+@router.get(
+    "/key/{key_id}/models",
+    tags=["key management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def key_resolved_models_fn(
+    key_id: str,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    
+    """
+    from litellm.proxy.proxy_server import (
+            llm_router,
+            prisma_client,
+        )
+
+    try:
+        if prisma_client is None:
+            raise Exception(
+                "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
+            )
+
+        hashed_key = _hash_token_if_needed(token=key_id)
+        key_info = await prisma_client.db.litellm_verificationtoken.find_unique(
+            where={"token": hashed_key},  # type: ignore
+        )
+        if key_info is None:
+            raise ProxyException(
+                message="Key not found in database",
+                type=ProxyErrorTypes.not_found_error,
+                param="key",
+                code=status.HTTP_404_NOT_FOUND,
+            )
+
+        if (
+            await _can_user_query_key_info(
+                user_api_key_dict=user_api_key_dict,
+                key=key_id,
+                key_info=key_info,
+            )
+            is not True
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="You are not allowed to access this key's info. Your role={}".format(
+                    user_api_key_dict.user_role
+                ),
+            )
+
+        key_models = list[str](key_info.models or [])
+        all_models: List[str] = []
+
+        if llm_router is not None:
+            all_models = llm_router.get_model_names()
+
+        source: str = SpecialModelNames.no_default_models.value
+        resolved: List[str] = key_models
+
+        #Team Models
+        if (SpecialModelNames.all_team_models.value in key_models):
+            if key_info.team_id is not None:
+                source = SpecialModelNames.all_team_models.value
+                team_row = await prisma_client.db.litellm_teamtable.find_unique(
+                    where={"team_id": key_info.team_id},
+                )
+                if team_row is not None and team_row.models is not None:
+                    resolved = list[str](team_row.models)
+            else:
+                source = SpecialModelNames.all_team_models.value
+                resolved = all_models
+        
+        #Proxy Models
+        if SpecialModelNames.all_proxy_models.value in key_models or SpecialModelNames.all_proxy_models.value in resolved:
+            source = SpecialModelNames.all_proxy_models.value
+            resolved = all_models
+
+        return {"models": resolved, "source": source}
+    except Exception as e:
+        raise handle_exception_on_proxy(e)
+
+
 def _check_model_access_group(
     models: Optional[List[str]], llm_router: Optional[Router], premium_user: bool
 ) -> Literal[True]:

ui/litellm-dashboard/src/components/key_team_helpers/KeyModelList.tsx

@@ -0,0 +1,36 @@
+import React from "react";
+import {UseGetKeyModels} from "@/hooks/keys/useGetKeyModels";
+import { Card, Tag } from 'antd';
+
+interface KeyModelListProps {
+    key_id: string;
+}
+
+const extractDefaultTags = (source: string) => {
+    if (source === 'all-proxy-models') {
+        return <Tag className="ml-2">All proxy models</Tag>
+    } else if (source === 'all-team-models') {
+        return <Tag className="ml-2">All team models</Tag>
+    }
+    return ''
+}
+
+
+
+const KeyModelList: React.FC<KeyModelListProps> = ({ key_id }) => {
+    const { data: keyModels, isLoading} = UseGetKeyModels(key_id)
+
+    
+
+    const title = keyModels ? <>Model {extractDefaultTags(keyModels.source)}</> : 'Model'
+    return (
+        <Card title={title} loading={isLoading}>
+            <div className="mt-2 flex flex-wrap gap-2">
+                {keyModels && keyModels.models.map((item:string)=> {return <Tag>{item}</Tag>})}
+            </div>
+        </Card>
+    );
+};
+
+export default KeyModelList;
+

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

@@ -3308,6 +3308,41 @@ export const keyAliasesCall = async (
   }
 };
 
+export interface keyModelResponse {
+  source: string;
+  models: string[];
+}
+
+export const keyModelCall = async (
+  accessToken: string,
+  key_id: string
+): Promise<keyModelResponse> => {
+
+  try {
+    let url = proxyBaseUrl ? `${proxyBaseUrl}/key/${key_id}/models` : `/key/${key_id}/models`;
+
+    const response = await fetch(url, {
+      method: "GET",
+      headers: {
+        [globalLitellmHeaderName]: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+    });
+
+    if (!response.ok) {
+      const errorData = await response.json();
+      const errorMessage = deriveErrorMessage(errorData);
+      handleError(errorMessage);
+      throw new Error(errorMessage);
+    }
+
+    const data = await response.json();
+    return data;
+  } catch (error) {
+    throw error;
+  }
+};
+
 export const userDailyActivityAggregatedCall = async (accessToken: string, startTime: Date, endTime: Date, userId: string | null = null) => {
   /**
    * Get aggregated daily user activity (no pagination)

ui/litellm-dashboard/src/components/templates/key_info_view.tsx

@@ -23,6 +23,7 @@ import ObjectPermissionsView from "../object_permissions_view";
 import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal";
 import { parseErrorMessage } from "../shared/errorUtils";
 import { KeyEditView } from "./key_edit_view";
+import KeyModelList from "../key_team_helpers/KeyModelList";
 
 interface KeyInfoViewProps {
   keyId: string;
@@ -488,22 +489,16 @@ export default function KeyInfoView({
                   <Text>RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"}</Text>
                 </div>
               </Card>
-
-              <Card>
-                <Text>Models</Text>
-                <div className="mt-2 flex flex-wrap gap-2">
-                  {currentKeyData.models && currentKeyData.models.length > 0 ? (
-                    currentKeyData.models.map((model, index) => (
-                      <Badge key={index} color="red">
-                        {model}
-                      </Badge>
-                    ))
-                  ) : (
-                    <Text>No models specified</Text>
-                  )}
-                </div>
-              </Card>
-
+              {currentKeyData.models && currentKeyData.models.length > 0 ? (
+                <KeyModelList key_id={currentKeyData.token}/>
+              ) : (
+                <Card>
+                  <Text>Models</Text>
+                    <div className="mt-2 flex flex-wrap gap-2">
+                      <Text>No models specified</Text>
+                    </div>
+                </Card>
+              )}
               <Card>
                 <ObjectPermissionsView
                   objectPermission={currentKeyData.object_permission}

ui/litellm-dashboard/src/hooks/keys/useGetKeyModels.ts

@@ -0,0 +1,15 @@
+import { useQuery } from '@tanstack/react-query';
+import { keyModelCall } from '@/components/networking';
+import useAuthorized from '@/app/(dashboard)/hooks/useAuthorized';
+
+export const UseGetKeyModels = (key_id: string) => {
+  const { accessToken } = useAuthorized();
+  
+  return useQuery({
+    queryKey: ['keyModels', key_id],
+    queryFn: () => {
+        if (!accessToken) throw new Error("Access Token required");
+        return keyModelCall(accessToken, key_id);
+    },
+  });
+};