Merge pull request #3492 from BerriAI/litellm_fixes_retry_policy

[UI] Fixes for getting/setting Retry Policy by Model Group

litellm/router.py

@@ -3268,6 +3268,8 @@ class Router:
 
         if retry_policy is None:
             return None
+        if isinstance(retry_policy, dict):
+            retry_policy = RetryPolicy(**retry_policy)
         if (
             isinstance(exception, litellm.BadRequestError)
             and retry_policy.BadRequestErrorRetries is not None

litellm/types/router.py

@@ -55,6 +55,7 @@ class UpdateRouterConfig(BaseModel):
 
     routing_strategy_args: Optional[dict] = None
     routing_strategy: Optional[str] = None
+    model_group_retry_policy: Optional[dict] = None
     allowed_fails: Optional[int] = None
     cooldown_time: Optional[float] = None
     num_retries: Optional[int] = None
@@ -344,3 +345,4 @@ class RetryPolicy(BaseModel):
     TimeoutErrorRetries: Optional[int] = None
     RateLimitErrorRetries: Optional[int] = None
     ContentPolicyViolationErrorRetries: Optional[int] = None
+    InternalServerErrorRetries: Optional[int] = None

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,
@@ -60,6 +60,10 @@ interface EditModelModalProps {
   onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted
 }
 
+interface RetryPolicyObject {
+  [key: string]: { [retryPolicyKey: string]: number } | undefined;
+}
+
 //["OpenAI", "Azure OpenAI", "Anthropic", "Gemini (Google AI Studio)", "Amazon Bedrock", "OpenAI-Compatible Endpoints (Groq, Together AI, Mistral AI, etc.)"]
 
 enum Providers {
@@ -89,7 +93,8 @@ const retry_policy_map: Record <string, string> = {
   "AuthenticationError  (401)": "AuthenticationErrorRetries",
   "TimeoutError (408)": "TimeoutErrorRetries",
   "RateLimitError (429)": "RateLimitErrorRetries",
-  "ContentPolicyViolationError (400)": "ContentPolicyViolationErrorRetries"
+  "ContentPolicyViolationError (400)": "ContentPolicyViolationErrorRetries",
+  "InternalServerError (500)": "InternalServerErrorRetries"
 };
 
 
@@ -222,7 +227,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
     to: new Date(),
   });
 
-  const [modelGroupRetryPolicy, setModelGroupRetryPolicy] = useState<Record<string, number>>({});
+  const [modelGroupRetryPolicy, setModelGroupRetryPolicy] = useState<RetryPolicyObject | null>(null);
   const [defaultRetry, setDefaultRetry] = useState<number>(0);
 
 
@@ -443,6 +448,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) {
@@ -1275,26 +1303,46 @@ 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>
+  <tbody>
-  {Object.keys(retry_policy_map).map((key, idx) => (
+    {Object.entries(retry_policy_map).map(([exceptionType, retryPolicyKey], idx) => {
-    <tr key={idx} className="flex justify-between items-center mt-2">
+
-      <td>
+      let retryCount = modelGroupRetryPolicy?.[selectedModelGroup!]?.[retryPolicyKey]
-        <Text>{key}</Text>
+      if (retryCount == null) {
-      </td>
+        retryCount = defaultRetry;
-      <td>
+      }
-        <InputNumber
+      
-          className="ml-5"
+      return (
-          defaultValue={defaultRetry}
+        <tr key={idx} className="flex justify-between items-center mt-2">
-          min={0}
+          <td>
-          step={1}
+            <Text>{exceptionType}</Text>
-        />
+          </td>
-      </td>
+          <td>
-    </tr>
+            <InputNumber
-  ))}
+              className="ml-5"
-</tbody>
+              value={retryCount}
-  </table>
+              min={0}
+              step={1}
+              onChange={(value) => {
+                setModelGroupRetryPolicy(prevModelGroupRetryPolicy => {
+                  const prevRetryPolicy = prevModelGroupRetryPolicy?.[selectedModelGroup!] ?? {};
+                  return {
+                    ...prevModelGroupRetryPolicy ?? {},
+                    [selectedModelGroup!]: {
+                      ...prevRetryPolicy,
+                      [retryPolicyKey!]: value,
+                    },
+                  } as RetryPolicyObject;
+                });
+              }}
+            />
+          </td>
+        </tr>
+      );
+    })}
+  </tbody>
+</table>
 }
-<Button className="mt-6 mr-8">
+<Button className="mt-6 mr-8" onClick={handleSaveRetrySettings}>
   Save
 </Button>