From 4fdf1147d39f830267bf51d69c64f8853aa0a36c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 6 May 2024 17:44:55 -0700
Subject: [PATCH] ui - get/set model retry policy

---
 .../src/components/model_dashboard.tsx        | 30 +++++++++++++++++--
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx
index 1297db295..a62365474 100644
--- a/ui/litellm-dashboard/src/components/model_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx
@@ -18,7 +18,7 @@ import {
 } from "@tremor/react";
 import { TabPanel, TabPanels, TabGroup, TabList, Tab, TextInput, Icon, DateRangePicker } from "@tremor/react";
 import { Select, SelectItem, MultiSelect, MultiSelectItem, DateRangePickerValue } from "@tremor/react";
-import { modelInfoCall, userGetRequesedtModelsCall, modelCreateCall, Model, modelCostMap, modelDeleteCall, healthCheckCall, modelUpdateCall, modelMetricsCall, modelExceptionsCall, modelMetricsSlowResponsesCall, getCallbacksCall } from "./networking";
+import { modelInfoCall, userGetRequesedtModelsCall, modelCreateCall, Model, modelCostMap, modelDeleteCall, healthCheckCall, modelUpdateCall, modelMetricsCall, modelExceptionsCall, modelMetricsSlowResponsesCall, getCallbacksCall, setCallbacksCall } from "./networking";
 import { BarChart, AreaChart } from "@tremor/react";
 import {
   Button as Button2,
@@ -443,6 +443,29 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
     setLastRefreshed(currentDate.toLocaleString());
   };
 
+  const handleSaveRetrySettings = async () => {
+    if (!accessToken) {
+      console.error("Access token is missing");
+      return;
+    }
+
+    console.log("new modelGroupRetryPolicy:", modelGroupRetryPolicy);
+  
+    try {
+      const payload = {
+          router_settings: {
+            model_group_retry_policy: modelGroupRetryPolicy
+          }
+      };
+  
+      await setCallbacksCall(accessToken, payload);
+      message.success("Retry settings saved successfully");
+    } catch (error) {
+      console.error("Failed to save retry settings:", error);
+      message.error("Failed to save retry settings");
+    }
+  };
+
 
   useEffect(() => {
     if (!accessToken || !token || !userRole || !userID) {
@@ -1277,7 +1300,8 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
   <table>
   <tbody>
     {Object.entries(retry_policy_map).map(([exceptionType, retryPolicyKey], idx) => {
-      let retryCount = (modelGroupRetryPolicy[selectedModelGroup] ?? {})[retryPolicyKey];
+
+      let retryCount = modelGroupRetryPolicy?.[selectedModelGroup]?.[retryPolicyKey]
       if (retryCount == null) {
         retryCount = defaultRetry;
       }
@@ -1310,7 +1334,7 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
   </tbody>
 </table>
 }
-<Button className="mt-6 mr-8">
+<Button className="mt-6 mr-8" onClick={handleSaveRetrySettings}>
   Save
 </Button>