From e753f61231be60544b5de180ca0c9e0f2a310bf7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 7 Mar 2025 17:01:11 -0800
Subject: [PATCH] fix logs page selected log (#9061)

---
 .../src/components/view_logs/columns.tsx      | 16 +++++-
 .../src/components/view_logs/index.tsx        | 19 +++++++
 .../src/components/view_logs/table.tsx        | 51 ++++++++++++++++++-
 3 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx
index 0652e3a8bfc..724732ec6ae 100644
--- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx
@@ -42,9 +42,23 @@ export const columns: ColumnDef<LogEntry>[] = [
           {...{
             onClick: row.getToggleExpandedHandler(),
             style: { cursor: "pointer" },
+            "aria-label": row.getIsExpanded() ? "Collapse row" : "Expand row",
           }}
+          className="w-6 h-6 flex items-center justify-center"
         >
-          {row.getIsExpanded() ? "▼" : "▶"}
+          <svg
+            className={`w-4 h-4 transition-transform ${row.getIsExpanded() ? 'transform rotate-90' : ''}`}
+            fill="none"
+            stroke="currentColor"
+            viewBox="0 0 24 24"
+          >
+            <path
+              strokeLinecap="round"
+              strokeLinejoin="round"
+              strokeWidth={2}
+              d="M9 5l7 7-7 7"
+            />
+          </svg>
         </button>
       ) : (
         "●"
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx
index db88e2acaaa..ab38c3fabe7 100644
--- a/ui/litellm-dashboard/src/components/view_logs/index.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx
@@ -67,6 +67,7 @@ export default function SpendLogsTable({
   const [filterByCurrentUser, setFilterByCurrentUser] = useState(
     userRole && internalUserRoles.includes(userRole)
   );
+  const [expandedRequestId, setExpandedRequestId] = useState<string | null>(null);
 
   const queryClient = useQueryClient();
 
@@ -176,6 +177,18 @@ export default function SpendLogsTable({
     refetchIntervalInBackground: true,
   });
 
+  // Add this effect to preserve expanded state when data refreshes
+  useEffect(() => {
+    if (logs.data?.data && expandedRequestId) {
+      // Check if the expanded request ID still exists in the new data
+      const stillExists = logs.data.data.some(log => log.request_id === expandedRequestId);
+      if (!stillExists) {
+        // If the request ID no longer exists in the data, clear the expanded state
+        setExpandedRequestId(null);
+      }
+    }
+  }, [logs.data?.data, expandedRequestId]);
+
   if (!accessToken || !token || !userRole || !userID) {
     console.log(
       "got None values for one of accessToken, token, userRole, userID",
@@ -220,6 +233,10 @@ export default function SpendLogsTable({
     return `${start.format('MMM D')} - ${now.format('MMM D')}`;
   };
 
+  const handleRowExpand = (requestId: string | null) => {
+    setExpandedRequestId(requestId);
+  };
+
   return (
     <div className="w-full p-6">
       <div className="flex items-center justify-between mb-4">
@@ -573,6 +590,8 @@ export default function SpendLogsTable({
           data={filteredData}
           renderSubComponent={RequestViewer}
           getRowCanExpand={() => true}
+          onRowExpand={handleRowExpand}
+          expandedRequestId={expandedRequestId}
         />
       </div>
     </div>
diff --git a/ui/litellm-dashboard/src/components/view_logs/table.tsx b/ui/litellm-dashboard/src/components/view_logs/table.tsx
index c2e8da9dea3..af9bc5afd25 100644
--- a/ui/litellm-dashboard/src/components/view_logs/table.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/table.tsx
@@ -1,4 +1,4 @@
-import { Fragment } from "react";
+import { Fragment, useEffect } from "react";
 import {
   ColumnDef,
   flexRender,
@@ -23,14 +23,18 @@ interface DataTableProps<TData, TValue> {
   renderSubComponent: (props: { row: Row<TData> }) => React.ReactElement;
   getRowCanExpand: (row: Row<TData>) => boolean;
   isLoading?: boolean;
+  expandedRequestId?: string | null;
+  onRowExpand?: (requestId: string | null) => void;
 }
 
-export function DataTable<TData, TValue>({
+export function DataTable<TData extends { request_id: string }, TValue>({
   data = [],
   columns,
   getRowCanExpand,
   renderSubComponent,
   isLoading = false,
+  expandedRequestId,
+  onRowExpand,
 }: DataTableProps<TData, TValue>) {
   const table = useReactTable({
     data,
@@ -38,8 +42,51 @@ export function DataTable<TData, TValue>({
     getRowCanExpand,
     getCoreRowModel: getCoreRowModel(),
     getExpandedRowModel: getExpandedRowModel(),
+    state: {
+      expanded: expandedRequestId 
+        ? data.reduce((acc, row, index) => {
+            if (row.request_id === expandedRequestId) {
+              acc[index] = true;
+            }
+            return acc;
+          }, {} as Record<string, boolean>)
+        : {},
+    },
+    onExpandedChange: (updater) => {
+      if (!onRowExpand) return;
+      
+      // Get current expanded state
+      const currentExpanded = expandedRequestId 
+        ? data.reduce((acc, row, index) => {
+            if (row.request_id === expandedRequestId) {
+              acc[index] = true;
+            }
+            return acc;
+          }, {} as Record<string, boolean>)
+        : {};
+      
+      // Calculate new expanded state
+      const newExpanded = typeof updater === 'function' 
+        ? updater(currentExpanded) 
+        : updater;
+      
+      // If empty, it means we're closing the expanded row
+      if (Object.keys(newExpanded).length === 0) {
+        onRowExpand(null);
+        return;
+      }
+      
+      // Find the request_id of the expanded row
+      const expandedIndex = Object.keys(newExpanded)[0];
+      const expandedRow = expandedIndex !== undefined ? data[parseInt(expandedIndex)] : null;
+      
+      // Call the onRowExpand callback with the request_id
+      onRowExpand(expandedRow ? expandedRow.request_id : null);
+    },
   });
 
+  // No need for the useEffect here as we're handling everything in onExpandedChange
+  
   return (
     <div className="rounded-lg custom-border">
       <Table className="[&_td]:py-0.5 [&_th]:py-1">