UI - Support 'batch' model health checks + make 'team-only' model concept clearer (#12770)

* fix(add_model_modes.tsx): add 'batch' mode to ui

* fix(main.py): support health checks on batches + support litellm_credentials on batches

* fix(add_model_tab.tsx): clarify what 'team' on add model means

litellm/batches/main.py

@@ -469,6 +469,7 @@ def retrieve_batch(
         raise e
 
 
+@client
 async def alist_batches(
     after: Optional[str] = None,
     limit: Optional[int] = None,
@@ -481,6 +482,7 @@ async def alist_batches(
     """
     Async: List your organization's batches.
     """
+
     try:
         loop = asyncio.get_event_loop()
         kwargs["alist_batches"] = True
@@ -510,6 +512,7 @@ async def alist_batches(
         raise e
 
 
+@client
 def list_batches(
     after: Optional[str] = None,
     limit: Optional[int] = None,

litellm/main.py

@@ -2925,7 +2925,7 @@ def completion(  # type: ignore # noqa: PLR0915
                     acompletion=acompletion,
                     client=client,
                     api_base=api_base,
-                    api_key=api_key
+                    api_key=api_key,
                 )
             elif bedrock_route == "converse_like":
                 model = model.replace("converse_like/", "")
@@ -3306,7 +3306,7 @@ def completion(  # type: ignore # noqa: PLR0915
                 custom_llm_provider=custom_llm_provider,
                 encoding=encoding,
                 stream=stream,
-                provider_config=bytez_transformation
+                provider_config=bytez_transformation,
             )
 
             pass
@@ -5563,6 +5563,9 @@ async def ahealth_check(
                 api_key=model_params.get("api_key", None),
                 api_version=model_params.get("api_version", None),
             ),
+            "batch": lambda: litellm.alist_batches(
+                **_filter_model_params(model_params),
+            ),
         }
 
         if mode in mode_handlers:

ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx

@@ -7,5 +7,6 @@ export const TEST_MODES = [
   { value: "audio_transcription", label: "Audio Transcription - /audio/transcriptions" },
   { value: "image_generation", label: "Image Generation - /images/generations" },
   { value: "rerank", label: "Rerank - /rerank" },
-  { value: "realtime", label: "Realtime - /realtime"}
+  { value: "realtime", label: "Realtime - /realtime"},
+  { value: "batch", label: "Batch - /batch"}
 ]; 

ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx

@@ -12,7 +12,7 @@ import { CredentialItem, modelAvailableCall } from "../networking";
 import ConnectionErrorDisplay from "./model_connection_test";
 import { TEST_MODES } from "./add_model_modes";
 import { Row, Col } from "antd";
-import { Text, TextInput } from "@tremor/react";
+import { Text, TextInput, Switch } from "@tremor/react";
 import TeamDropdown from "../common_components/team_dropdown";
 import { all_admin_roles } from "@/utils/roles";
 
@@ -31,6 +31,7 @@ interface AddModelTabProps {
   credentials: CredentialItem[];
   accessToken: string;
   userRole: string;
+  premiumUser: boolean;
 }
 
 const { Title, Link } = Typography;
@@ -50,6 +51,7 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
   credentials,
   accessToken,
   userRole,
+  premiumUser,
 }) => {
   // State for test mode and connection testing
   const [testMode, setTestMode] = useState<string>("chat");
@@ -68,6 +70,9 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
     setIsResultModalVisible(true);
   };
 
+  // State for team-only switch
+  const [isTeamOnly, setIsTeamOnly] = useState<boolean>(false);
+
   const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
 
   useEffect(() => {
@@ -238,20 +243,46 @@ const AddModelTab: React.FC<AddModelTabProps> = ({
               <span className="px-4 text-gray-500 text-sm">Additional Model Info Settings</span>
               <div className="flex-grow border-t border-gray-200"></div>
             </div>
+            {/* Team-only Model Switch */}
             <Form.Item
-              label="Team"
-              name="team_id"
+              label="Team-BYOK Model"
+              tooltip="Only use this model + credential combination for this team. Useful when teams want to onboard their own OpenAI keys."
               className="mb-4"
-              tooltip="Only keys for this team, will be able to call this model."
-              rules={[
-                {
-                  required: !isAdmin, // Required if not admin
-                  message: 'Please select a team.'
-                }
-              ]}
             >
-              <TeamDropdown teams={teams} />
+              <Tooltip 
+                title={!premiumUser ? "This is an enterprise-only feature. Upgrade to premium to restrict model+credential combinations to a specific team." : ""}
+                placement="top"
+              >
+                <Switch 
+                  checked={isTeamOnly}
+                  onChange={(checked) => {
+                    setIsTeamOnly(checked);
+                    if (!checked) {
+                      form.setFieldValue('team_id', undefined);
+                    }
+                  }}
+                  disabled={!premiumUser}
+                />
+              </Tooltip>
             </Form.Item>
+
+            {/* Conditional Team Selection */}
+            {isTeamOnly && (
+              <Form.Item
+                label="Select Team"
+                name="team_id"
+                className="mb-4"
+                tooltip="Only keys for this team will be able to call this model."
+                rules={[
+                  {
+                    required: isTeamOnly && !isAdmin,
+                    message: 'Please select a team.'
+                  }
+                ]}
+              >
+                <TeamDropdown teams={teams} disabled={!premiumUser} />
+              </Form.Item>
+            )}
             {
               isAdmin && (
                 <>

ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx

@@ -6,15 +6,18 @@ interface TeamDropdownProps {
   teams?: Team[] | null;
   value?: string;
   onChange?: (value: string) => void;
+  disabled?: boolean;
 }
 
-const TeamDropdown: React.FC<TeamDropdownProps> = ({ teams, value, onChange }) => {
+const TeamDropdown: React.FC<TeamDropdownProps> = ({ teams, value, onChange, disabled }) => {
+  console.log("disabled", disabled);
   return (
     <Select
       showSearch
       placeholder="Search or select a team"
       value={value}
       onChange={onChange}
+      disabled={disabled}
       filterOption={(input, option) => {
         if (!option) return false;
         // Get team data from the option key

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

@@ -1397,6 +1397,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
                 credentials={credentialsList}
                 accessToken={accessToken}
                 userRole={userRole}
+                premiumUser={premiumUser}
               />
             </TabPanel>
             <TabPanel>