From 57a78703481f3bd80274ee5befbcf2e0bb54b114 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 6 May 2024 17:23:51 -0700
Subject: [PATCH] ui - show retry policy per model group

---
 .../src/components/model_dashboard.tsx        | 52 ++++++++++++-------
 1 file changed, 34 insertions(+), 18 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx
index 43ca9fdff34..1297db29593 100644
--- a/ui/litellm-dashboard/src/components/model_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx
@@ -1275,24 +1275,40 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
               <Text className="mb-6">How many retries should be attempted based on the Exception</Text>
               {retry_policy_map &&
   <table>
-    <tbody>
-  {Object.keys(retry_policy_map).map((key, idx) => (
-    <tr key={idx} className="flex justify-between items-center mt-2">
-      <td>
-        <Text>{key}</Text>
-      </td>
-      <td>
-        <InputNumber
-          className="ml-5"
-          defaultValue={defaultRetry}
-          min={0}
-          step={1}
-        />
-      </td>
-    </tr>
-  ))}
-</tbody>
-  </table>
+  <tbody>
+    {Object.entries(retry_policy_map).map(([exceptionType, retryPolicyKey], idx) => {
+      let retryCount = (modelGroupRetryPolicy[selectedModelGroup] ?? {})[retryPolicyKey];
+      if (retryCount == null) {
+        retryCount = defaultRetry;
+      }
+      
+      return (
+        <tr key={idx} className="flex justify-between items-center mt-2">
+          <td>
+            <Text>{exceptionType}</Text>
+          </td>
+          <td>
+            <InputNumber
+              className="ml-5"
+              value={retryCount}
+              min={0}
+              step={1}
+              onChange={(value) => {
+                setModelGroupRetryPolicy({
+                  ...modelGroupRetryPolicy,
+                  [selectedModelGroup]: {
+                    ...modelGroupRetryPolicy[selectedModelGroup],
+                    [retryPolicyKey]: value,
+                  },
+                });
+              }}
+            />
+          </td>
+        </tr>
+      );
+    })}
+  </tbody>
+</table>
 }
 <Button className="mt-6 mr-8">
   Save