From c1dcc81fd500ef805df34c254e317c743eda77db Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Wed, 24 Dec 2025 14:06:59 -0800
Subject: [PATCH] Model Page sort by all models

---
 .../ModelsAndEndpointsView.tsx                |  1 -
 .../components/AllModelsTab.tsx               | 25 ++++++++-----------
 .../src/components/model_dashboard/table.tsx  | 21 +++++++++-------
 3 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
index 969c7bafa6b..bf001c62126 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
@@ -661,7 +661,6 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
                   setSelectedModelId={setSelectedModelId}
                   setSelectedTeamId={setSelectedTeamId}
                   setEditModel={setEditModel}
-                  modelData={modelData}
                 />
                 {!shouldHideAddModelTab && (
                   <TabPanel className="h-full">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
index 87fa0b1e3b6..04c05ede5c2 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
@@ -1,13 +1,14 @@
+import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import useTeams from "@/app/(dashboard)/hooks/useTeams";
 import { Team } from "@/components/key_team_helpers/key_list";
 import { ModelDataTable } from "@/components/model_dashboard/table";
 import { columns } from "@/components/molecules/models/columns";
 import { getDisplayModelName } from "@/components/view_model/model_name_display";
 import { InfoCircleOutlined } from "@ant-design/icons";
-import { PaginationState, Table as TableInstance } from "@tanstack/react-table";
+import { PaginationState } from "@tanstack/react-table";
 import { Grid, Select, SelectItem, TabPanel, Text } from "@tremor/react";
-import { useEffect, useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useState } from "react";
+import { useModelsInfo } from "../../hooks/models/useModels";
 
 type ModelViewMode = "all" | "current_team";
 
@@ -19,7 +20,6 @@ interface AllModelsTabProps {
   setSelectedModelId: (id: string) => void;
   setSelectedTeamId: (id: string) => void;
   setEditModel: (edit: boolean) => void;
-  modelData: any;
 }
 
 const AllModelsTab = ({
@@ -30,10 +30,10 @@ const AllModelsTab = ({
   setSelectedModelId,
   setSelectedTeamId,
   setEditModel,
-  modelData,
 }: AllModelsTabProps) => {
+  const { data: modelData } = useModelsInfo();
   const { userId, userRole, premiumUser } = useAuthorized();
-  const { teams } = useTeams();
+  const { data: teams } = useTeams();
 
   const [modelNameSearch, setModelNameSearch] = useState<string>("");
   const [modelViewMode, setModelViewMode] = useState<ModelViewMode>("current_team");
@@ -45,7 +45,6 @@ const AllModelsTab = ({
     pageIndex: 0,
     pageSize: 50,
   });
-  const tableRef = useRef<TableInstance<any>>(null);
 
   const filteredData = useMemo(() => {
     if (!modelData || !modelData.data || modelData.data.length === 0) {
@@ -88,12 +87,6 @@ const AllModelsTab = ({
     });
   }, [modelData, modelNameSearch, selectedModelGroup, selectedModelAccessGroupFilter, currentTeam, modelViewMode]);
 
-  const paginatedData = useMemo(() => {
-    const startIndex = pagination.pageIndex * pagination.pageSize;
-    const endIndex = startIndex + pagination.pageSize;
-    return filteredData.slice(startIndex, endIndex);
-  }, [filteredData, pagination.pageIndex, pagination.pageSize]);
-
   useEffect(() => {
     setPagination((prev: PaginationState) => ({ ...prev, pageIndex: 0 }));
   }, [modelNameSearch, selectedModelGroup, selectedModelAccessGroupFilter, currentTeam, modelViewMode]);
@@ -370,9 +363,11 @@ const AllModelsTab = ({
                 expandedRows,
                 setExpandedRows,
               )}
-              data={paginatedData}
+              data={filteredData}
               isLoading={false}
-              table={tableRef}
+              pagination={pagination}
+              onPaginationChange={setPagination}
+              enablePagination={true}
             />
           </div>
         </div>
diff --git a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx
index 344ff2e94f2..79224edba43 100644
--- a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx
+++ b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx
@@ -3,10 +3,13 @@ import {
   flexRender,
   getCoreRowModel,
   getSortedRowModel,
+  getPaginationRowModel,
   SortingState,
   useReactTable,
   ColumnResizeMode,
   VisibilityState,
+  PaginationState,
+  OnChangeFn,
 } from "@tanstack/react-table";
 import React from "react";
 import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react";
@@ -23,16 +26,20 @@ interface ModelDataTableProps<TData, TValue> {
   data: TData[];
   columns: ColumnDef<TData, TValue>[];
   isLoading?: boolean;
-  table: any; // Add table prop to access column visibility controls
   defaultSorting?: SortingState;
+  pagination?: PaginationState;
+  onPaginationChange?: OnChangeFn<PaginationState>;
+  enablePagination?: boolean;
 }
 
 export function ModelDataTable<TData, TValue>({
   data = [],
   columns,
   isLoading = false,
-  table,
   defaultSorting = [],
+  pagination,
+  onPaginationChange,
+  enablePagination = false,
 }: ModelDataTableProps<TData, TValue>) {
   const [sorting, setSorting] = React.useState<SortingState>(defaultSorting);
   const [columnResizeMode] = React.useState<ColumnResizeMode>("onChange");
@@ -46,13 +53,16 @@ export function ModelDataTable<TData, TValue>({
       sorting,
       columnSizing,
       columnVisibility,
+      ...(enablePagination && pagination ? { pagination } : {}),
     },
     columnResizeMode,
     onSortingChange: setSorting,
     onColumnSizingChange: setColumnSizing,
     onColumnVisibilityChange: setColumnVisibility,
+    ...(enablePagination && onPaginationChange ? { onPaginationChange } : {}),
     getCoreRowModel: getCoreRowModel(),
     getSortedRowModel: getSortedRowModel(),
+    ...(enablePagination ? { getPaginationRowModel: getPaginationRowModel() } : {}),
     enableSorting: true,
     enableColumnResizing: true,
     defaultColumn: {
@@ -61,13 +71,6 @@ export function ModelDataTable<TData, TValue>({
     },
   });
 
-  // Expose table instance to parent
-  React.useEffect(() => {
-    if (table) {
-      table.current = tableInstance;
-    }
-  }, [tableInstance, table]);
-
   const getHeaderText = (header: any): string => {
     if (typeof header === "string") {
       return header;